diff --git a/scripts/play_interactive.py b/scripts/play_interactive.py index 1ea9377bf..cf8fe4bef 100644 --- a/scripts/play_interactive.py +++ b/scripts/play_interactive.py @@ -52,10 +52,9 @@ def ensure_registries(): from unilab.base import registry from unilab.config.structured_configs import PPOConfig -from unilab.utils.obs_utils import flatten_obs_dict from unilab.utils.rsl_rl_compat import convert_config_v3_to_v4, is_rsl_rl_v4 +from unilab.utils.rsl_rl_vec_env_wrapper import RslRlVecEnvWrapper from unilab.utils.run_utils import get_latest_run -from unilab.utils.torch_utils import to_torch try: from rsl_rl.runners import OnPolicyRunner @@ -65,87 +64,6 @@ def ensure_registries(): from tensordict import TensorDict -# --------------------------------------------------------------------------- -# Minimal env wrapper (mirrors RslRlVecEnvWrapper in train_rsl_rl.py) -# --------------------------------------------------------------------------- - - -class RslRlVecEnvWrapper: - def __init__(self, env, device="cuda", policy_obs_mode: str = "flat"): - self.env = env - self.cfg = env.cfg - self.device = device - self.policy_obs_mode = policy_obs_mode - self.num_envs = env.num_envs - self.observation_space = env.observation_space - self.action_space = env.action_space - self._actor_obs_dim = int(env.obs_groups_spec.get("obs", sum(env.obs_groups_spec.values()))) - self._flat_obs_dim = int(sum(env.obs_groups_spec.values())) - self.num_obs = self._flat_obs_dim if policy_obs_mode == "flat" else self._actor_obs_dim - self.num_privileged_obs = self.num_obs - self.num_actions = env.action_space.shape[0] - self.episode_returns = torch.zeros(self.num_envs, device=device) - self.episode_lengths = torch.zeros(self.num_envs, device=device) - self.episode_length_buf = self.episode_lengths - self.max_episode_length = int(env.cfg.max_episode_seconds / env.cfg.ctrl_dt) - self.reset() - - def _obs_to_tensordict(self, obs: dict[str, np.ndarray]) -> TensorDict: - actor = to_torch(obs["obs"], self.device) - if self.policy_obs_mode == "actor": - policy = actor - else: - policy = to_torch(flatten_obs_dict(obs), self.device) - td = {"policy": policy, "actor": actor} - if "privileged" in obs: - td["privileged"] = to_torch(obs["privileged"], self.device) - return TensorDict(td, batch_size=self.num_envs, device=self.device) - - def step(self, actions): - actions_np = ( - actions.detach().cpu().numpy() if isinstance(actions, torch.Tensor) else actions - ) - state = self.env.step(actions_np) - rewards = to_torch(state.reward, self.device) - dones = to_torch(state.done, self.device).bool() - self.episode_returns += rewards - self.episode_lengths += 1 - infos = {} - done_idx = torch.nonzero(dones).flatten() - if len(done_idx) > 0: - if hasattr(state, "truncated"): - infos["time_outs"] = to_torch(state.truncated, self.device).bool() - self.episode_returns[done_idx] = 0 - self.episode_lengths[done_idx] = 0 - if hasattr(state, "info") and "log" in state.info: - infos["log"] = state.info["log"] - return self._obs_to_tensordict(state.obs), rewards, dones, infos - - def reset(self): - if self.env.state is None: - self.env.init_state() - env_indices = np.arange(self.num_envs, dtype=np.int32) - reset_out = self.env.reset(env_indices) - if isinstance(reset_out, tuple) and len(reset_out) == 2: - obs_out, _ = reset_out - elif isinstance(reset_out, tuple) and len(reset_out) == 3: - _, obs_out, _ = reset_out - else: - raise ValueError( - "Unexpected env.reset return format. Expected (obs, info) or (_, obs, _)." - ) - self.episode_returns[:] = 0 - self.episode_lengths[:] = 0 - return self._obs_to_tensordict(obs_out), {} - - def get_observations(self): - return self._obs_to_tensordict(self.env.state.obs) - - def get_privileged_observations(self): - if self.policy_obs_mode == "actor": - return to_torch(self.env.state.obs["obs"], self.device) - return to_torch(flatten_obs_dict(self.env.state.obs), self.device) - def _infer_checkpoint_actor_input_dim(ckpt_path: str) -> int | None: loaded = torch.load(ckpt_path, map_location="cpu", weights_only=True) diff --git a/scripts/train_rsl_rl.py b/scripts/train_rsl_rl.py index 04a5585bf..cdbd307ea 100644 --- a/scripts/train_rsl_rl.py +++ b/scripts/train_rsl_rl.py @@ -24,6 +24,7 @@ from unilab.utils.experiment_tracking import ExperimentTracker, patch_rsl_rl_wandb_writer from unilab.utils.obs_utils import flatten_obs_dict from unilab.utils.reward_utils import resolve_reward_dict +from unilab.utils.rsl_rl_vec_env_wrapper import RslRlVecEnvWrapper from unilab.utils.torch_utils import to_numpy, to_torch from unilab.utils.xml_utils import materialize_scene_visual_override @@ -225,90 +226,6 @@ def run_motrix_rsl_play_loop( steps_run += 1 -class RslRlVecEnvWrapper: - """Wrapper to adapt NpEnv to RSL-RL OnPolicyRunner interface.""" - - def __init__(self, env, device="cuda"): - self.env = env - self.cfg = env.cfg - self.device = device - self.num_envs = env.num_envs - self.observation_space = env.observation_space - self.action_space = env.action_space - self.num_obs = sum(env.obs_groups_spec.values()) - self.num_privileged_obs = self.num_obs - self.num_actions = env.action_space.shape[0] - - self.episode_returns = torch.zeros(self.num_envs, device=self.device) - self.episode_lengths = torch.zeros(self.num_envs, device=self.device) - - self.episode_length_buf = self.episode_lengths - self.max_episode_length = np.ceil(env.cfg.max_episode_seconds / env.cfg.ctrl_dt) - - self.reset() - - def _obs_to_tensordict(self, obs: dict[str, np.ndarray]) -> TensorDict: - actor = to_torch(obs["obs"], self.device) - td = {"actor": actor} - if "privileged" in obs: - td["privileged"] = to_torch(obs["privileged"], self.device) - if hasattr(self.env, "get_policy_obs"): - policy_obs = self.env.get_policy_obs(obs) - else: - policy_obs = flatten_obs_dict(obs) - td["policy"] = to_torch(policy_obs, self.device) - else: - td["policy"] = actor - return TensorDict(td, batch_size=self.num_envs, device=self.device) - - def step(self, actions): - if isinstance(actions, torch.Tensor): - actions_np = actions.detach().cpu().numpy() - else: - actions_np = actions - - state = self.env.step(actions_np) - - # Convert output to torch tensors on target device - rewards = to_torch(state.reward, self.device) - dones = to_torch(state.done, self.device).bool() - - self.episode_returns += rewards - self.episode_lengths += 1 - - infos = {} - done_indices = torch.nonzero(dones).flatten() - if len(done_indices) > 0: - if hasattr(state, "truncated"): - infos["time_outs"] = to_torch(state.truncated, self.device).bool() - self.episode_returns[done_indices] = 0 - self.episode_lengths[done_indices] = 0 - - if hasattr(state, "info") and "log" in state.info: - infos["log"] = state.info["log"] - - obs_dict = self._obs_to_tensordict(state.obs) - return obs_dict, rewards, dones, infos - - def reset(self): - if self.env.state is None: - self.env.init_state() - env_indices = np.arange(self.num_envs, dtype=np.int32) - obs_out, _ = self.env.reset(env_indices) - - self.episode_returns[:] = 0 - self.episode_lengths[:] = 0 - - return self._obs_to_tensordict(obs_out), {} - - def get_observations(self): - return self._obs_to_tensordict(self.env.state.obs) - - def get_privileged_observations(self): - obs = to_torch(flatten_obs_dict(self.env.state.obs), self.device) - return obs - - def play_rsl_rl(cfg: DictConfig, device: str) -> str | None: """Play mode for RSL-RL.""" diff --git a/src/unilab/utils/__init__.py b/src/unilab/utils/__init__.py index 5954b6b21..0eec1a8d7 100644 --- a/src/unilab/utils/__init__.py +++ b/src/unilab/utils/__init__.py @@ -2,6 +2,7 @@ from unilab.utils.algo_utils import build_actor, ensure_registries from unilab.utils.offpolicy_logger import OffPolicyLogger from unilab.utils.onpolicy_logger import OnPolicyLogger +from unilab.utils.rsl_rl_vec_env_wrapper import RslRlVecEnvWrapper from unilab.utils.torch_utils import to_numpy, to_torch __all__ = [ @@ -11,4 +12,5 @@ "OnPolicyLogger", "ensure_registries", "build_actor", + "RslRlVecEnvWrapper", ] diff --git a/src/unilab/utils/rsl_rl_vec_env_wrapper.py b/src/unilab/utils/rsl_rl_vec_env_wrapper.py new file mode 100644 index 000000000..e9ddd926f --- /dev/null +++ b/src/unilab/utils/rsl_rl_vec_env_wrapper.py @@ -0,0 +1,155 @@ +"""Shared RSL-RL vectorized environment wrapper. + +This module provides a unified RslRlVecEnvWrapper that aligns with the current +env contract (obs, info) reset format and is used by both training and play scripts. +""" + +import numpy as np +import torch +from tensordict import TensorDict + +from unilab.utils.obs_utils import flatten_obs_dict +from unilab.utils.torch_utils import to_torch + + +class RslRlVecEnvWrapper: + """Wrapper to adapt NpEnv to RSL-RL OnPolicyRunner interface. + + This wrapper aligns with the current env contract: + - reset() returns (obs_dict, info_dict) + - step() returns state with .obs, .reward, .done, .truncated, .info attributes + + Args: + env: The environment to wrap (must follow NpEnv contract). + device: Device to place tensors on ("cuda", "mps", or "cpu"). + policy_obs_mode: Observation mode for policy ("flat" or "actor"). + "flat" uses flattened obs dict, "actor" uses only the "obs" key. + Default is "flat" for backward compatibility with training scripts. + """ + + def __init__(self, env, device: str = "cuda", policy_obs_mode: str = "flat"): + self.env = env + self.cfg = env.cfg + self.device = device + self.policy_obs_mode = policy_obs_mode + self.num_envs = env.num_envs + self.observation_space = env.observation_space + self.action_space = env.action_space + + # Compute observation dimensions + self._actor_obs_dim = int(env.obs_groups_spec.get("obs", sum(env.obs_groups_spec.values()))) + self._flat_obs_dim = int(sum(env.obs_groups_spec.values())) + self.num_obs = self._flat_obs_dim if policy_obs_mode == "flat" else self._actor_obs_dim + self.num_privileged_obs = self.num_obs + self.num_actions = env.action_space.shape[0] + + # Episode tracking + self.episode_returns = torch.zeros(self.num_envs, device=device) + self.episode_lengths = torch.zeros(self.num_envs, device=device) + self.episode_length_buf = self.episode_lengths + self.max_episode_length = int(env.cfg.max_episode_seconds / env.cfg.ctrl_dt) + + # Initialize + self.reset() + + def _obs_to_tensordict(self, obs: dict[str, np.ndarray]) -> TensorDict: + """Convert observation dict to TensorDict for RSL-RL. + + Args: + obs: Observation dictionary with "obs" key and optional "privileged" key. + + Returns: + TensorDict with "policy" and "actor" keys (and optional "privileged"). + """ + actor = to_torch(obs["obs"], self.device) + + if self.policy_obs_mode == "actor": + policy = actor + else: + policy = to_torch(flatten_obs_dict(obs), self.device) + + td: dict[str, torch.Tensor] = {"policy": policy, "actor": actor} + + if "privileged" in obs: + td["privileged"] = to_torch(obs["privileged"], self.device) + + return TensorDict(td, batch_size=self.num_envs, device=self.device) + + def step(self, actions): + """Execute one step in the environment. + + Args: + actions: Actions to execute (torch.Tensor or numpy array). + + Returns: + Tuple of (obs_tensordict, rewards, dones, infos). + """ + # Convert actions to numpy + if isinstance(actions, torch.Tensor): + actions_np = actions.detach().cpu().numpy() + else: + actions_np = actions + + # Step the environment + state = self.env.step(actions_np) + + # Convert outputs to torch tensors + rewards = to_torch(state.reward, self.device) + dones = to_torch(state.done, self.device).bool() + + # Update episode statistics + self.episode_returns += rewards + self.episode_lengths += 1 + + # Build info dict + infos = {} + done_indices = torch.nonzero(dones).flatten() + if len(done_indices) > 0: + if hasattr(state, "truncated"): + infos["time_outs"] = to_torch(state.truncated, self.device).bool() + self.episode_returns[done_indices] = 0 + self.episode_lengths[done_indices] = 0 + + if hasattr(state, "info") and "log" in state.info: + infos["log"] = state.info["log"] + + obs_dict = self._obs_to_tensordict(state.obs) + return obs_dict, rewards, dones, infos + + def reset(self): + """Reset the environment. + + Returns: + Tuple of (obs_tensordict, empty_info_dict). + """ + # Ensure state is initialized + if self.env.state is None: + self.env.init_state() + + # Reset all environments + env_indices = np.arange(self.num_envs, dtype=np.int32) + obs_out, _ = self.env.reset(env_indices) + + # Reset episode statistics + self.episode_returns[:] = 0 + self.episode_lengths[:] = 0 + + return self._obs_to_tensordict(obs_out), {} + + def get_observations(self): + """Get current observations without stepping. + + Returns: + TensorDict with current observations. + """ + return self._obs_to_tensordict(self.env.state.obs) + + def get_privileged_observations(self): + """Get current privileged observations. + + Returns: + Torch tensor with privileged observations (or policy obs if unavailable). + """ + if self.policy_obs_mode == "actor": + return to_torch(self.env.state.obs["obs"], self.device) + return to_torch(flatten_obs_dict(self.env.state.obs), self.device) diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 4aaaa9a6e..48b51ca9d 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -848,3 +848,96 @@ def test_play_resolve_checkpoint_empty_dir(tmp_path): run_dir.mkdir() result = mod.resolve_checkpoint("MyTask", str(run_dir)) assert result is None + + +# --------------------------------------------------------------------------- +# play_interactive.py — RslRlVecEnvWrapper contract behavior +# --------------------------------------------------------------------------- + + +def _play_interactive(): + """Load play_interactive.py as a module.""" + return _load_script("play_interactive") + + +def test_play_wrapper_imports_shared_implementation(): + """Verify play_interactive.py uses shared RslRlVecEnvWrapper.""" + from unilab.utils.rsl_rl_vec_env_wrapper import RslRlVecEnvWrapper as SharedWrapper + + mod = _play_interactive() + # The wrapper class in play_interactive should be the shared one + assert mod.RslRlVecEnvWrapper is SharedWrapper + + +def test_play_wrapper_uses_current_reset_contract(): + """Verify wrapper reset() uses current (obs, info) contract, not old (_, obs, _).""" + import numpy as np + from tensordict import TensorDict + + from unilab.utils.rsl_rl_vec_env_wrapper import RslRlVecEnvWrapper + + # Create a fake environment that returns (obs, info) tuple + class FakeEnv: + def __init__(self): + self.num_envs = 2 + self.state = type("State", (), {"obs": {"obs": np.ones((2, 5), dtype=np.float32)}})() + self.cfg = type("Cfg", (), {"max_episode_seconds": 10.0, "ctrl_dt": 0.02})() + self.observation_space = type("Space", (), {"shape": (5,)})() + self.action_space = type("Space", (), {"shape": (3,)})() + self.obs_groups_spec = {"obs": 5} + + def init_state(self): + pass + + def reset(self, env_indices): + # Returns current contract: (obs, info) + return {"obs": np.ones((2, 5), dtype=np.float32)}, {} + + env = FakeEnv() + wrapper = RslRlVecEnvWrapper(env, device="cpu", policy_obs_mode="flat") + + # Reset should work with current contract + obs_td, info = wrapper.reset() + + assert isinstance(obs_td, TensorDict) + assert "policy" in obs_td + assert "actor" in obs_td + assert obs_td.batch_size == (2,) + + +def test_play_wrapper_policy_obs_mode_actor(): + """Verify wrapper supports policy_obs_mode='actor'.""" + import numpy as np + + from unilab.utils.rsl_rl_vec_env_wrapper import RslRlVecEnvWrapper + + class FakeEnv: + def __init__(self): + self.num_envs = 1 + self.state = type("State", (), {"obs": {"obs": np.ones((1, 3), dtype=np.float32)}})() + self.cfg = type("Cfg", (), {"max_episode_seconds": 10.0, "ctrl_dt": 0.02})() + self.observation_space = type("Space", (), {"shape": (3,)})() + self.action_space = type("Space", (), {"shape": (2,)})() + self.obs_groups_spec = {"obs": 3, "privileged": 2} + + def init_state(self): + pass + + def reset(self, env_indices): + return { + "obs": np.ones((1, 3), dtype=np.float32), + "privileged": np.zeros((1, 2), dtype=np.float32), + }, {} + + env = FakeEnv() + + # Test actor mode - num_obs should match actor obs dim only + wrapper_actor = RslRlVecEnvWrapper(env, device="cpu", policy_obs_mode="actor") + assert wrapper_actor.num_obs == 3 # Only "obs" group + assert wrapper_actor._actor_obs_dim == 3 + assert wrapper_actor._flat_obs_dim == 5 # obs + privileged + + obs_td, _ = wrapper_actor.reset() + # In actor mode, policy obs should equal actor obs + assert obs_td["policy"].shape == (1, 3) + assert obs_td["actor"].shape == (1, 3)