huggingface · AdilZouitine · Apr 16, 2025 · Mar 31, 2025 · Mar 31, 2025 · Mar 31, 2025
diff --git a/lerobot/common/envs/configs.py b/lerobot/common/envs/configs.py
@@ -171,7 +171,6 @@ class VideoRecordConfig:
 class WrapperConfig:
     """Configuration for environment wrappers."""
 
-    delta_action: float | None = None
     joint_masking_action_space: list[bool] | None = None
 
 
@@ -191,7 +190,6 @@ class EnvWrapperConfig:
     """Configuration for environment wrappers."""
 
     display_cameras: bool = False
-    delta_action: float = 0.1
     use_relative_joint_positions: bool = True
     add_joint_velocity_to_observation: bool = False
     add_ee_pose_to_observation: bool = False
@@ -203,6 +201,10 @@ class EnvWrapperConfig:
     joint_masking_action_space: Optional[Any] = None
     ee_action_space_params: Optional[EEActionSpaceConfig] = None
     use_gripper: bool = False
+    gripper_quantization_threshold: float | None = 0.8
+    gripper_penalty: float = 0.0
+    gripper_penalty_in_reward: bool = False
+    open_gripper_on_reset: bool = False
 
 
 @EnvConfig.register_subclass(name="gym_manipulator")
@@ -254,6 +256,7 @@ class ManiskillEnvConfig(EnvConfig):
     robot: str = "so100"  # This is a hack to make the robot config work
     video_record: VideoRecordConfig = field(default_factory=VideoRecordConfig)
     wrapper: WrapperConfig = field(default_factory=WrapperConfig)
+    mock_gripper: bool = False
     features: dict[str, PolicyFeature] = field(
         default_factory=lambda: {
             "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),

diff --git a/lerobot/common/policies/sac/configuration_sac.py b/lerobot/common/policies/sac/configuration_sac.py
@@ -51,8 +51,8 @@ class ActorNetworkConfig:
 @dataclass
 class PolicyConfig:
     use_tanh_squash: bool = True
-    log_std_min: int = -5
-    log_std_max: int = 2
+    log_std_min: float = 1e-5
+    log_std_max: float = 10.0
     init_final: float = 0.05
 
 
@@ -85,12 +85,15 @@ class SACConfig(PreTrainedConfig):
         freeze_vision_encoder: Whether to freeze the vision encoder during training.
         image_encoder_hidden_dim: Hidden dimension size for the image encoder.
         shared_encoder: Whether to use a shared encoder for actor and critic.
+        num_discrete_actions: Number of discrete actions, eg for gripper actions.
+        image_embedding_pooling_dim: Dimension of the image embedding pooling.
         concurrency: Configuration for concurrency settings.
         actor_learner: Configuration for actor-learner architecture.
         online_steps: Number of steps for online training.
         online_env_seed: Seed for the online environment.
         online_buffer_capacity: Capacity of the online replay buffer.
         offline_buffer_capacity: Capacity of the offline replay buffer.
+        async_prefetch: Whether to use asynchronous prefetching for the buffers.
         online_step_before_learning: Number of steps before learning starts.
         policy_update_freq: Frequency of policy updates.
         discount: Discount factor for the SAC algorithm.
@@ -118,7 +121,7 @@ class SACConfig(PreTrainedConfig):
         }
     )
 
-    dataset_stats: dict[str, dict[str, list[float]]] = field(
+    dataset_stats: dict[str, dict[str, list[float]]] | None = field(
         default_factory=lambda: {
             "observation.image": {
                 "mean": [0.485, 0.456, 0.406],
@@ -144,12 +147,15 @@ class SACConfig(PreTrainedConfig):
    freeze_vision_encoder: bool = True
     image_encoder_hidden_dim: int = 32
     shared_encoder: bool = True
+    num_discrete_actions: int | None = None
+    image_embedding_pooling_dim: int = 8
 
     # Training parameter
     online_steps: int = 1000000
     online_env_seed: int = 10000
     online_buffer_capacity: int = 100000
     offline_buffer_capacity: int = 100000
+    async_prefetch: bool = False
     online_step_before_learning: int = 100
     policy_update_freq: int = 1
 
@@ -173,7 +179,7 @@ class SACConfig(PreTrainedConfig):
     critic_network_kwargs: CriticNetworkConfig = field(default_factory=CriticNetworkConfig)
     actor_network_kwargs: ActorNetworkConfig = field(default_factory=ActorNetworkConfig)
     policy_kwargs: PolicyConfig = field(default_factory=PolicyConfig)
-
+    grasp_critic_network_kwargs: CriticNetworkConfig = field(default_factory=CriticNetworkConfig)
     actor_learner_config: ActorLearnerConfig = field(default_factory=ActorLearnerConfig)
     concurrency: ConcurrencyConfig = field(default_factory=ConcurrencyConfig)