refactor: replace flat obs array with structured dict[str, np.ndarray]

CLAUDE.md

@@ -71,9 +71,15 @@ tests/
 │   ├── test_shared_obs_stats.py
 │   └── test_async_runner.py
 ├── base/
-│   └── test_registry.py
+│   ├── test_registry.py
+│   └── test_np_env.py             # NpEnvState + NpEnv dict-obs contract
 ├── config/
-│   └── test_locomotion_params.py
+│   ├── test_locomotion_params.py
+│   └── test_manipulation_params.py
+├── envs/
+│   └── test_env_configs.py        # obs_groups_spec dims + env instantiation
+├── utils/
+│   └── test_obs_utils.py          # flatten_obs_dict
 ├── scripts/
 │   └── test_train_scripts.py
 └── algos/
@@ -85,6 +91,15 @@ tests/
 Tests marked `@pytest.mark.slow` require a real MuJoCo environment and are excluded from CI
 by default. Run them locally when working on runner/learner code.
 
+## Testing
+
+**New features must ship with tests.** When developing a new feature or refactoring, design and write comprehensive unit tests alongside the feature code — not as an afterthought. Tests should cover:
+
+- Normal behaviour and edge cases
+- Error paths and invalid inputs
+- Contract verification (e.g. interface shapes, types, key presence)
+- Integration with neighbouring modules when relevant (`@pytest.mark.slow` for MuJoCo-dependent tests)
+
 ## Git Commits
 
 Use Conventional Commits:

scripts/play_interactive.py

@@ -48,6 +48,7 @@ def ensure_registries():
 
 from unilab.base import registry
 from unilab.config import locomotion_params
+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.run_utils import get_latest_run
 from unilab.utils.torch_utils import to_torch
@@ -73,7 +74,7 @@ def __init__(self, env, device="cuda"):
         self.num_envs = env.num_envs
         self.observation_space = env.observation_space
         self.action_space = env.action_space
-        self.num_obs = env.observation_space.shape[0]
+        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=device)
@@ -82,12 +83,15 @@ def __init__(self, env, device="cuda"):
         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:
+        flat = to_torch(flatten_obs_dict(obs), self.device)
+        return TensorDict({"policy": flat}, 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)
-        obs = to_torch(state.obs, self.device)
         rewards = to_torch(state.reward, self.device)
         dones = to_torch(state.done, self.device).bool()
         self.episode_returns += rewards
@@ -101,25 +105,22 @@ def step(self, actions):
             self.episode_lengths[done_idx] = 0
         if hasattr(state, "info") and "log" in state.info:
             infos["log"] = state.info["log"]
-        obs_dict = TensorDict({"policy": obs}, batch_size=self.num_envs, device=self.device)
-        return obs_dict, rewards, dones, infos
+        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)
         _, obs_out, _ = self.env.reset(env_indices)
-        obs = to_torch(obs_out, self.device)
         self.episode_returns[:] = 0
         self.episode_lengths[:] = 0
-        return TensorDict({"policy": obs}, batch_size=self.num_envs, device=self.device), {}
+        return self._obs_to_tensordict(obs_out), {}
 
     def get_observations(self):
-        obs = to_torch(self.env.state.obs, self.device)
-        return TensorDict({"policy": obs}, batch_size=self.num_envs, device=self.device)
+        return self._obs_to_tensordict(self.env.state.obs)
 
     def get_privileged_observations(self):
-        return to_torch(self.env.state.obs, self.device)
+        return to_torch(flatten_obs_dict(self.env.state.obs), self.device)
 
 
 # ---------------------------------------------------------------------------

scripts/train_appo.py

@@ -57,7 +57,7 @@ def play_appo(args, rl_cfg):
     print(f"Using device for play: {device}")
 
     env = registry.make(args.task, num_envs=args.play_env_num, sim_backend="mujoco")
-    obs_dim = env.observation_space.shape[0]
+    obs_dim = sum(env.obs_groups_spec.values())
     action_dim = env.action_space.shape[0]
 
     rl_cfg_dict = dict(rl_cfg)
@@ -142,8 +142,10 @@ def play_appo(args, rl_cfg):
     if env.state is None:
         env.init_state()
     env_indices = np.arange(args.play_env_num, dtype=np.int32)
+    from unilab.utils.obs_utils import flatten_obs_dict
+
     _, obs_out, _ = env.reset(env_indices)
-    obs_np = np.asarray(obs_out, dtype=np.float32)
+    obs_np = np.asarray(flatten_obs_dict(obs_out), dtype=np.float32)
 
     state_list = []
     num_steps = 150
@@ -156,11 +158,7 @@ def play_appo(args, rl_cfg):
             actions_torch = actor(td)
             actions_np = actions_torch.cpu().numpy().astype(np.float32)
             state = env.step(actions_np)
-            if hasattr(state, "obs"):
-                next_obs_raw = state.obs
-            else:
-                next_obs_raw = state[0]
-            obs_np = np.asarray(next_obs_raw, dtype=np.float32)
+            obs_np = np.asarray(flatten_obs_dict(state.obs), dtype=np.float32)
             state_list.append(np.asarray(env._backend.get_physics_state(), dtype=np.float32).copy())
 
     print("Rendering frames...")

scripts/train_mlx_ppo.py

@@ -50,6 +50,7 @@ def ensure_registries() -> None:
 from unilab.base import registry
 from unilab.config import locomotion_params
 from unilab.utils import render_many
+from unilab.utils.obs_utils import flatten_obs_dict
 from unilab.utils.onpolicy_logger import OnPolicyLogger
 
 TASK_STEP_TUNING = {
@@ -165,7 +166,7 @@ def play_mlx_ppo(args, cfg, dtype, use_fp16, resolved_sim_backend, task_log_root
     play_model_dtype = mx.float32 if use_fp16 else dtype
     play_env_num = args.play_env_num
     env = registry.make(args.task, num_envs=play_env_num, sim_backend=resolved_sim_backend)
-    obs_dim = env.observation_space.shape[0]
+    obs_dim = sum(env.obs_groups_spec.values())
     action_dim = env.action_space.shape[0]
     model = build_model(cfg, obs_dim, action_dim, dtype=play_model_dtype)
 
@@ -202,8 +203,8 @@ def play_mlx_ppo(args, cfg, dtype, use_fp16, resolved_sim_backend, task_log_root
     if env.state is None:
         env.init_state()
     play_reset_indices = np.arange(env.num_envs, dtype=np.int32)
-    _, obs, _ = env.reset(play_reset_indices)
-    obs = mx.array(obs)
+    _, obs_dict_play, _ = env.reset(play_reset_indices)
+    obs = mx.array(flatten_obs_dict(obs_dict_play))
 
     if resolved_sim_backend == "motrix":
         print("[MLX PPO] Starting interactive visualization (motrix native renderer)...")
@@ -227,7 +228,7 @@ def play_mlx_ppo(args, cfg, dtype, use_fp16, resolved_sim_backend, task_log_root
                 )
                 env_actions = np.asarray(actions)
                 state = env.step(env_actions)
-                raw_obs = state.obs
+                raw_obs = flatten_obs_dict(state.obs)
                 obs = mx.nan_to_num(raw_obs, nan=0.0, posinf=0.0, neginf=0.0)
 
                 current_time = time.perf_counter()
@@ -259,7 +260,7 @@ def play_mlx_ppo(args, cfg, dtype, use_fp16, resolved_sim_backend, task_log_root
         actions = actions.astype(dtype) if getattr(actions, "dtype", None) != dtype else actions
         env_actions = np.asarray(actions)
         state = env.step(env_actions)
-        raw_obs = state.obs
+        raw_obs = flatten_obs_dict(state.obs)
         obs = mx.nan_to_num(raw_obs, nan=0.0, posinf=0.0, neginf=0.0)
         state_list.append(_get_physics_state_snapshot(env))
 
@@ -428,10 +429,10 @@ def log(msg: str) -> None:
     if env.state is None:
         env.init_state()
     reset_indices = np.arange(env.num_envs, dtype=np.int32)
-    _, obs, _ = env.reset(reset_indices)
-    obs = mx.array(obs)
+    _, obs_dict, _ = env.reset(reset_indices)
+    obs = mx.array(flatten_obs_dict(obs_dict))
 
-    obs_dim = env.observation_space.shape[0]
+    obs_dim = sum(env.obs_groups_spec.values())
     action_dim = env.action_space.shape[0]
 
     model = build_model(cfg, obs_dim, action_dim, dtype=model_dtype)
@@ -599,7 +600,7 @@ def log(msg: str) -> None:
             # Conversion boundary: env output → rollout; sanitize Nan/Inf only here (no forced reset).
             raw_rewards = mx.array(state.reward)
             raw_dones = mx.array(state.done)
-            raw_obs = mx.array(state.obs)
+            raw_obs = mx.array(flatten_obs_dict(state.obs))
             rewards = mx.nan_to_num(raw_rewards, nan=0.0, posinf=0.0, neginf=0.0)
             dones = mx.where(mx.isfinite(raw_dones), raw_dones, mx.ones_like(raw_dones)).astype(
                 dtype

scripts/train_offpolicy.py

@@ -158,8 +158,8 @@ def build_runner(algo_name: str, args, cfg):
             ensure_registries()
             device = args.device or get_default_device()
             env = registry.make(args.task, num_envs=1, sim_backend=args.sim_backend)
-            assert env.observation_space.shape and env.action_space.shape
-            obs_dim = env.observation_space.shape[0]
+            assert env.action_space.shape
+            obs_dim = sum(env.obs_groups_spec.values())
             action_dim = env.action_space.shape[0]
             env.close()
 
@@ -286,8 +286,10 @@ def play_offpolicy(algo_name: str, args, cfg) -> None:
     device = default_device(torch, args.device)
     print(f"Using device for play: {device}")
 
+    from unilab.utils.obs_utils import flatten_obs_dict
+
     env = registry.make(args.task, num_envs=args.play_env_num, sim_backend=args.sim_backend)
-    obs_dim = env.observation_space.shape[0]
+    obs_dim = sum(env.obs_groups_spec.values())
     action_dim = env.action_space.shape[0]
 
     normalizer = None
@@ -339,7 +341,7 @@ def play_offpolicy(algo_name: str, args, cfg) -> None:
         env.init_state()
     env_indices = np.arange(args.play_env_num, dtype=np.int32)
     _, obs_out, _ = env.reset(env_indices)
-    obs_np = np.asarray(obs_out, dtype=np.float32)
+    obs_np = np.asarray(flatten_obs_dict(obs_out), dtype=np.float32)
 
     # Use Motrix native rendering
     if args.sim_backend == "motrix":
@@ -364,7 +366,7 @@ def play_offpolicy(algo_name: str, args, cfg) -> None:
                         else actor(obs_torch).cpu().numpy()
                     )
                     state = env.step(actions_np)
-                    obs_np = np.asarray(state.obs, dtype=np.float32)
+                    obs_np = np.asarray(flatten_obs_dict(state.obs), dtype=np.float32)
 
                     current_time = time.perf_counter()
                     elapsed = current_time - last_render_time
@@ -397,7 +399,7 @@ def play_offpolicy(algo_name: str, args, cfg) -> None:
                 else actor(obs_torch).cpu().numpy()
             )
             state = env.step(actions_np)
-            obs_np = np.asarray(state.obs, dtype=np.float32)
+            obs_np = np.asarray(flatten_obs_dict(state.obs), dtype=np.float32)
             state_list.append(np.asarray(env._backend.get_physics_state(), dtype=np.float32).copy())
 
     print("Rendering frames...")

scripts/train_rsl_rl.py

@@ -34,6 +34,7 @@ def ensure_registries():
 
 from unilab.base import registry
 from unilab.config import locomotion_params, manipulation_params
+from unilab.utils.obs_utils import flatten_obs_dict
 from unilab.utils.torch_utils import to_numpy, to_torch
 
 # Try importing rsl_rl
@@ -62,9 +63,8 @@ def __init__(self, env, device="cuda"):
         self.device = device
         self.num_envs = env.num_envs
         self.observation_space = env.observation_space
-        self.obs_dict = self.cfg.obs_config.obs_dict
         self.action_space = env.action_space
-        self.num_obs = env.observation_space.shape[0]
+        self.num_obs = sum(env.obs_groups_spec.values())
         self.num_privileged_obs = self.num_obs
         self.num_actions = env.action_space.shape[0]
 
@@ -78,6 +78,18 @@ def __init__(self, env, device="cuda"):
         # RSL-RL runner calls get_observations() in __init__, so we need to ensure env is reset
         self.reset()
 
+    def _obs_to_tensordict(self, obs: dict[str, np.ndarray]) -> TensorDict:
+        actor = to_torch(obs["actor"], self.device)
+        if "privileged" in obs:
+            policy = to_torch(flatten_obs_dict(obs), self.device)
+        else:
+            policy = actor
+        return TensorDict(
+            {"policy": policy, "actor": actor},
+            batch_size=self.num_envs,
+            device=self.device,
+        )
+
     def step(self, actions):
         # Convert actions to numpy (CPU)
         if isinstance(actions, torch.Tensor):
@@ -89,7 +101,6 @@ def step(self, actions):
         state = self.env.step(actions_np)
 
         # Convert output to torch tensors on target device
-        obs = to_torch(state.obs, self.device)
         rewards = to_torch(state.reward, self.device)
         dones = to_torch(state.done, self.device).bool()
 
@@ -114,12 +125,7 @@ def step(self, actions):
         if hasattr(state, "info") and "log" in state.info:
             infos["log"] = state.info["log"]
 
-        # obs_dict = TensorDict({"policy": obs}, batch_size=self.num_envs, device=self.device)
-        obs_dict = TensorDict(
-            {"policy": obs, "actor": obs[:, self.env.actor_indices]},
-            batch_size=self.num_envs,
-            device=self.device,
-        )
+        obs_dict = self._obs_to_tensordict(state.obs)
         return obs_dict, rewards, dones, infos
 
     def reset(self):
@@ -128,79 +134,22 @@ def reset(self):
             self.env.init_state()
         env_indices = np.arange(self.num_envs, dtype=np.int32)
         _, obs_out, _ = self.env.reset(env_indices)
-        obs = to_torch(obs_out, self.device)
 
         self.episode_returns[:] = 0
         self.episode_lengths[:] = 0
 
-        return TensorDict(
-            {"policy": obs, "actor": obs[:, self.env.actor_indices]},
-            batch_size=self.num_envs,
-            device=self.device,
-        ), {}
+        return self._obs_to_tensordict(obs_out), {}
 
     def get_observations(self):
-        obs = to_torch(self.env.state.obs, self.device)
-        return TensorDict(
-            {"policy": obs, "actor": obs[:, self.env.actor_indices]},
-            batch_size=self.num_envs,
-            device=self.device,
-        )
+        return self._obs_to_tensordict(self.env.state.obs)
 
     def get_privileged_observations(self):
-        obs = to_torch(self.env.state.obs, self.device)
+        obs = to_torch(flatten_obs_dict(self.env.state.obs), self.device)
         return obs
 
 
-def RslRlAacVecEnvWrapper(RslRlVecEnvWrapper):  # Asymmetric Actor-Critic
-    def step(self, actions):
-        # Convert actions to numpy (CPU)
-        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 output to torch tensors on target device
-        obs = to_torch(state.obs, self.device)
-        rewards = to_torch(state.reward, self.device)
-        dones = to_torch(state.done, self.device).bool()
-
-        # Update logging info
-        self.episode_returns += rewards
-        self.episode_lengths += 1
-
-        infos = {}
-        # Check for dones
-        done_indices = torch.nonzero(dones).flatten()
-        if len(done_indices) > 0:
-            # Handle limits and timeouts (RSL-RL expects 'time_outs' in extras/infos)
-            if hasattr(state, "truncated"):
-                infos["time_outs"] = to_torch(state.truncated, self.device).bool()
-
-            # Reset buffers for done envs
-            self.episode_returns[done_indices] = 0
-            self.episode_lengths[done_indices] = 0
-
-        # Pass per-step logs if available (gs_playground style)
-        # prioritizing 'log' over 'episode' allows per-step metric logging
-        if hasattr(state, "info") and "log" in state.info:
-            infos["log"] = state.info["log"]
-
-        obs_dict = TensorDict(
-            {"policy": obs, "actor": obs[:, self.env.actor_indices]},
-            batch_size=self.num_envs,
-            device=self.device,
-        )
-
-        return obs_dict, rewards, dones, infos
-
-
 def play_rsl_rl(args, cfg, device):
     """Play mode for RSL-RL."""
-    import torch
 
     from unilab.base import registry