Add WrapperConfig for environment wrappers and update SACConfig properties

- Introduced `WrapperConfig` dataclass for environment wrapper configurations.
- Updated `ManiskillEnvConfig` to include a `wrapper` field for enhanced environment management.
- Modified `SACConfig` to return `None` for `observation_delta_indices` and `action_delta_indices` properties.
- Refactored `make_robot_env` function to improve readability and maintainability.
This commit is contained in:
AdilZouitine
2025-03-27 17:07:06 +00:00
parent d0b7690bc0
commit 79e0f6e06c
4 changed files with 114 additions and 55 deletions

View File

@@ -163,6 +163,12 @@ class VideoRecordConfig:
record_dir: str = "videos"
trajectory_name: str = "trajectory"
@dataclass
class WrapperConfig:
"""Configuration for environment wrappers."""
delta_action: float | None = None
joint_masking_action_space: list[bool] | None = None
@EnvConfig.register_subclass("maniskill_push")
@dataclass
class ManiskillEnvConfig(EnvConfig):
@@ -181,6 +187,7 @@ class ManiskillEnvConfig(EnvConfig):
device: str = "cuda"
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)
features: dict[str, PolicyFeature] = field(
default_factory=lambda: {
"action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),

View File

@@ -233,11 +233,11 @@ class SACConfig(PreTrainedConfig):
@property
def observation_delta_indices(self) -> list:
return list(range(1 - self.n_obs_steps, 1))
return None
@property
def action_delta_indices(self) -> list:
return [0] # SAC typically predicts one action at a time
return None # SAC typically predicts one action at a time
@property
def reward_delta_indices(self) -> None: