fix(issue-167): align play_interactive.py with current env contract

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)

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."""
 

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",
 ]

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)