From 113c108ddcf0e9c49ab4548b36a961929f2f0d50 Mon Sep 17 00:00:00 2001
From: tatp-yf <jyf23@mails.tsinghua.edu.cn>
Date: Wed, 18 Mar 2026 00:01:01 +0800
Subject: [PATCH] refactor: replace flat obs array with structured dict[str,
 np.ndarray]
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Environments now emit dict obs (e.g. {"actor": ..., "privileged": ...})
instead of concatenating then slicing flat arrays. This eliminates fragile
index alignment across env, config, and wrapper layers.

- NpEnvState.obs: np.ndarray → dict[str, np.ndarray]
- Add obs_groups_spec abstract property to ABEnv/NpEnv
- NpEnv.observation_space computed from obs_groups_spec
- Add flatten_obs_dict utility for downstream consumers
- Delete ObsConfig and actor_indices from all envs
- Update RSL-RL, APPO, off-policy, FastSAC, MLX wrappers
- Update config obs_groups keys to "actor"/"critic"
- Add comprehensive tests (61 new) and testing principle to CLAUDE.md
---
 CLAUDE.md                                     |  19 +-
 scripts/play_interactive.py                   |  19 +-
 scripts/train_appo.py                         |  12 +-
 scripts/train_mlx_ppo.py                      |  19 +-
 scripts/train_offpolicy.py                    |  14 +-
 scripts/train_rsl_rl.py                       |  87 +---
 src/unilab/algos/torch/appo/runner.py         |   3 +-
 src/unilab/algos/torch/appo/worker.py         |  10 +-
 src/unilab/algos/torch/fast_sac/runner.py     |   5 +-
 src/unilab/algos/torch/offpolicy/worker.py    |  12 +-
 src/unilab/base/base.py                       |   5 +
 src/unilab/base/np_env.py                     |  39 +-
 src/unilab/config/locomotion_params.py        |  10 +-
 src/unilab/config/manipulation_params.py      |   2 +-
 src/unilab/envs/locomotion/g1/joystick.py     |  45 +-
 src/unilab/envs/locomotion/g1/joystick_sac.py |  25 --
 .../locomotion/g1/motion_tracking/tracking.py |  57 +--
 src/unilab/envs/locomotion/go1/joystick.py    |  43 +-
 src/unilab/envs/locomotion/go2/joystick.py    |  45 +-
 src/unilab/envs/locomotion/obs_config.py      |  37 --
 .../inhand_rot_allegro/rotation.py            |  60 +--
 .../inhand_rot_allegro/rotation_sac.py        |  29 +-
 src/unilab/utils/device_utils.py              |   5 +-
 src/unilab/utils/obs_utils.py                 |   8 +
 tests/base/test_np_env.py                     | 424 ++++++++++++++++++
 tests/base/test_registry.py                   |   4 +
 tests/config/test_locomotion_params.py        |  11 +-
 tests/config/test_manipulation_params.py      | 146 ++++++
 tests/conftest.py                             |   4 +
 tests/envs/test_env_configs.py                |  51 +--
 tests/utils/__init__.py                       |   0
 tests/utils/test_obs_utils.py                 |  92 ++++
 32 files changed, 911 insertions(+), 431 deletions(-)
 delete mode 100644 src/unilab/envs/locomotion/obs_config.py
 create mode 100644 src/unilab/utils/obs_utils.py
 create mode 100644 tests/base/test_np_env.py
 create mode 100644 tests/config/test_manipulation_params.py
 create mode 100644 tests/utils/__init__.py
 create mode 100644 tests/utils/test_obs_utils.py

diff --git a/CLAUDE.md b/CLAUDE.md
index b5b1cad7b..f4cef6a47 100644
--- a/CLAUDE.md
+++ b/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:
diff --git a/scripts/play_interactive.py b/scripts/play_interactive.py
index 0297803ea..a8ce1fa1f 100644
--- a/scripts/play_interactive.py
+++ b/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)
 
 
 # ---------------------------------------------------------------------------
diff --git a/scripts/train_appo.py b/scripts/train_appo.py
index f7f870aa3..c5787a8e7 100644
--- a/scripts/train_appo.py
+++ b/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...")
diff --git a/scripts/train_mlx_ppo.py b/scripts/train_mlx_ppo.py
index 20c341dea..86b0764fb 100644
--- a/scripts/train_mlx_ppo.py
+++ b/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
diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py
index c029ea829..848f862ca 100644
--- a/scripts/train_offpolicy.py
+++ b/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...")
diff --git a/scripts/train_rsl_rl.py b/scripts/train_rsl_rl.py
index f380ca86b..a6287f804 100644
--- a/scripts/train_rsl_rl.py
+++ b/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
 
diff --git a/src/unilab/algos/torch/appo/runner.py b/src/unilab/algos/torch/appo/runner.py
index 686ccc954..878fae6dc 100644
--- a/src/unilab/algos/torch/appo/runner.py
+++ b/src/unilab/algos/torch/appo/runner.py
@@ -86,9 +86,8 @@ def _detect_dims(self):
         ensure_registries()
 
         env = registry.make(self.env_name, num_envs=1, sim_backend="mujoco")
-        assert env.observation_space.shape is not None
+        obs_dim = sum(env.obs_groups_spec.values())
         assert env.action_space.shape is not None
-        obs_dim = env.observation_space.shape[0]
         action_dim = env.action_space.shape[0]
         env.close()
 
diff --git a/src/unilab/algos/torch/appo/worker.py b/src/unilab/algos/torch/appo/worker.py
index 4298ff202..8ebc983a3 100644
--- a/src/unilab/algos/torch/appo/worker.py
+++ b/src/unilab/algos/torch/appo/worker.py
@@ -16,6 +16,7 @@
 from rsl_rl.utils import resolve_callable
 
 from unilab.utils.algo_utils import ensure_registries
+from unilab.utils.obs_utils import flatten_obs_dict
 
 
 def appo_collector_fn(
@@ -105,7 +106,7 @@ def to_float32_np(x):
             x = x.cpu().numpy()
         return np.asarray(x, dtype=np.float32)
 
-    obs_np = to_float32_np(obs_out)
+    obs_np = to_float32_np(flatten_obs_dict(obs_out))
 
     # Pre-allocate obs TensorDict once; update in-place each step to avoid
     # repeated TensorDict construction overhead in the hot loop.
@@ -170,15 +171,14 @@ def to_float32_np(x):
                 write_buf["dones"][:, step] = done_raw
                 write_buf["truncated"][:, step] = truncated_raw
 
-                obs_np = to_float32_np(next_obs_raw)
+                obs_np = to_float32_np(flatten_obs_dict(next_obs_raw))
 
                 # Bootstrap from true terminal obs so next rollout starts cleanly
                 if "_final_observation" in state.info:
                     has_final = np.asarray(state.info["_final_observation"], dtype=bool)
                     if np.any(has_final):
-                        obs_np[has_final] = to_float32_np(state.info["final_observation"])[
-                            has_final
-                        ]
+                        final_obs_flat = flatten_obs_dict(state.info["final_observation"])
+                        obs_np[has_final] = to_float32_np(final_obs_flat)[has_final]
 
                 # Episode tracking (vectorized)
                 total_steps += num_envs
diff --git a/src/unilab/algos/torch/fast_sac/runner.py b/src/unilab/algos/torch/fast_sac/runner.py
index e4c214ccd..4a25cfff8 100644
--- a/src/unilab/algos/torch/fast_sac/runner.py
+++ b/src/unilab/algos/torch/fast_sac/runner.py
@@ -45,10 +45,9 @@ def __init__(
 
         ensure_registries()
         env: Any = registry.make(env_name, num_envs=1, sim_backend=sim_backend)
-        obs_space_shape = env.observation_space.shape
+        obs_dim = sum(env.obs_groups_spec.values())
         act_space_shape = env.action_space.shape
-        assert obs_space_shape is not None and act_space_shape is not None
-        obs_dim = obs_space_shape[0]
+        assert act_space_shape is not None
         action_dim = act_space_shape[0]
         mujoco_model = getattr(env, "_backend", None)
         if mujoco_model is not None:
diff --git a/src/unilab/algos/torch/offpolicy/worker.py b/src/unilab/algos/torch/offpolicy/worker.py
index b1a4497e8..62cf4f788 100644
--- a/src/unilab/algos/torch/offpolicy/worker.py
+++ b/src/unilab/algos/torch/offpolicy/worker.py
@@ -11,6 +11,7 @@
 import torch
 
 from unilab.utils.algo_utils import build_actor, ensure_registries
+from unilab.utils.obs_utils import flatten_obs_dict
 
 
 def off_policy_collector_fn(
@@ -109,9 +110,8 @@ def _run_collector(
     )
 
     # Build actor (always on CPU for env interaction)
-    assert env.observation_space.shape is not None
+    obs_dim = sum(env.obs_groups_spec.values())
     assert env.action_space.shape is not None
-    obs_dim = env.observation_space.shape[0]
     action_dim = env.action_space.shape[0]
     actor = build_actor(
         algo_type, obs_dim, action_dim, actor_hidden_dim, use_layer_norm, "cpu", num_envs
@@ -141,7 +141,7 @@ def _run_collector(
     # Initial step to get first observation
     actions_np = np.zeros((num_envs, action_dim), dtype=np.float32)
     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)
     max_episode_steps = getattr(getattr(env, "cfg", None), "max_episode_steps", None)
     if max_episode_steps is not None and int(max_episode_steps) > 0:
         step_offsets = np.random.randint(
@@ -214,7 +214,7 @@ def _run_collector(
             timing_count += 1
 
         # Extract data as numpy
-        next_obs_np = np.asarray(state.obs, dtype=np.float32)
+        next_obs_np = np.asarray(flatten_obs_dict(state.obs), dtype=np.float32)
         rewards_np = np.asarray(state.reward, dtype=np.float32).ravel()
 
         terminated_np = (
@@ -241,7 +241,9 @@ def _run_collector(
             has_final = state.info["_final_observation"]
             has_final_np = np.asarray(has_final, dtype=bool)
             if np.any(has_final_np):
-                final_obs_np = np.asarray(state.info["final_observation"], dtype=np.float32)
+                final_obs_np = np.asarray(
+                    flatten_obs_dict(state.info["final_observation"]), dtype=np.float32
+                )
                 next_obs_np[has_final_np] = final_obs_np[has_final_np]
 
         # Write to replay buffer
diff --git a/src/unilab/base/base.py b/src/unilab/base/base.py
index 59ad799d7..0e92ea8fc 100644
--- a/src/unilab/base/base.py
+++ b/src/unilab/base/base.py
@@ -68,6 +68,11 @@ def observation_space(self) -> gym.Space:
     def action_space(self) -> gym.Space:
         """Action space"""
 
+    @property
+    @abc.abstractmethod
+    def obs_groups_spec(self) -> dict[str, int]:
+        """Map from observation group name to its dimension."""
+
     @property
     @abc.abstractmethod
     def state(self) -> Any:
diff --git a/src/unilab/base/np_env.py b/src/unilab/base/np_env.py
index 1e1f88762..393b090b0 100644
--- a/src/unilab/base/np_env.py
+++ b/src/unilab/base/np_env.py
@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 from typing import Any, Optional, Tuple
 
+import gymnasium as gym
 import numpy as np
 
 from unilab.base.backend import SimBackend
@@ -14,12 +15,11 @@
 
 @dataclass
 class NpEnvState:
-    obs: np.ndarray
+    obs: dict[str, np.ndarray]
     reward: np.ndarray
     terminated: np.ndarray
     truncated: np.ndarray
     info: dict
-    # critic_obs: Optional[np.ndarray] = None
 
     @property
     def done(self) -> np.ndarray:
@@ -57,11 +57,24 @@ def num_envs(self) -> int:
     def state(self) -> Optional[NpEnvState]:
         return self._state
 
+    @property
+    def obs_groups_spec(self) -> dict[str, int]:
+        """Return observation group dimensions, e.g. {"actor": 98, "privileged": 3}.
+
+        Subclasses MUST override this property.
+        """
+        raise NotImplementedError
+
+    @property
+    def observation_space(self) -> gym.Space:
+        total = sum(self.obs_groups_spec.values())
+        return gym.spaces.Box(-np.inf, np.inf, shape=(total,), dtype=np.float64)
+
     def init_state(self) -> NpEnvState:
         dtype = get_global_dtype()
-        obs_shape = self.observation_space.shape
-        assert obs_shape is not None
-        obs = np.zeros((self._num_envs, obs_shape[0]), dtype=dtype)
+        obs = {
+            k: np.zeros((self._num_envs, d), dtype=dtype) for k, d in self.obs_groups_spec.items()
+        }
         reward = np.zeros((self._num_envs,), dtype=dtype)
         terminated = np.ones((self._num_envs,), dtype=bool)
         truncated = np.zeros((self._num_envs,), dtype=bool)
@@ -123,15 +136,21 @@ def _reset_done_envs(self) -> None:
         self._state.info["steps"][env_indices] = 0
 
         if "final_observation" not in self._state.info:
-            self._state.info["final_observation"] = np.zeros_like(self._state.obs)
+            self._state.info["final_observation"] = {
+                k: np.zeros_like(v) for k, v in self._state.obs.items()
+            }
             self._state.info["_final_observation"] = np.zeros((self._num_envs,), dtype=bool)
 
         self._state.info["_final_observation"][:] = False
         self._state.info["_final_observation"][env_indices] = True
-        self._state.info["final_observation"][env_indices] = self._state.obs[env_indices]
+        for key in self._state.obs:
+            self._state.info["final_observation"][key][env_indices] = self._state.obs[key][
+                env_indices
+            ]
 
         new_obs, _, info1 = self.reset(env_indices)
-        self._state.obs[env_indices] = new_obs
+        for key in self._state.obs:
+            self._state.obs[key][env_indices] = new_obs[key]
 
         if info1:
             for key, value in info1.items():
@@ -160,7 +179,9 @@ def update_state(self, state: NpEnvState) -> NpEnvState:
         """子类实现：计算 obs/reward/terminated"""
 
     @abc.abstractmethod
-    def reset(self, env_indices: np.ndarray) -> Tuple[np.ndarray, np.ndarray, dict]:
+    def reset(
+        self, env_indices: np.ndarray
+    ) -> Tuple[dict[str, np.ndarray], dict[str, np.ndarray], dict]:
         """子类实现：重置指定环境"""
 
     def close(self) -> None:
diff --git a/src/unilab/config/locomotion_params.py b/src/unilab/config/locomotion_params.py
index 0aa3a9925..5f1fe2083 100644
--- a/src/unilab/config/locomotion_params.py
+++ b/src/unilab/config/locomotion_params.py
@@ -29,19 +29,19 @@ def ppo_config(env_name: str) -> config_dict.ConfigDict:
     # elif env_name in ("G1JoystickFlatTerrain", "G1WalkTaskMjSAC"):
     #     cfg.num_envs = 2048
     #     cfg.max_iterations = 220
-    obs_group_set = {"default": ["policy"]}
+    obs_group_set = {"actor": ["actor"], "critic": ["actor"]}
     if env_name == "Go1JoystickFlatTerrain":
         cfg.max_iterations = 151
-        obs_group_set = {"default": ["policy"], "actor": ["actor"]}
+        obs_group_set = {"actor": ["actor"], "critic": ["actor", "privileged"]}
     elif env_name == "G1JoystickFlatTerrain":
-        obs_group_set = {"default": ["policy"], "actor": ["actor"]}
+        obs_group_set = {"actor": ["actor"], "critic": ["actor", "privileged"]}
         cfg.num_envs = 2048
         cfg.max_iterations = 220
     elif env_name == "Go2JoystickFlatTerrain":
-        obs_group_set = {"default": ["policy"], "actor": ["actor"]}
+        obs_group_set = {"actor": ["actor"], "critic": ["actor", "privileged"]}
         pass
     elif env_name == "G1MotionTracking":
-        obs_group_set = {"default": ["policy"], "actor": ["actor"]}
+        obs_group_set = {"actor": ["actor"], "critic": ["actor", "privileged"]}
         cfg.num_envs = 4096
         cfg.max_iterations = 30000
         cfg.save_interval = 500
diff --git a/src/unilab/config/manipulation_params.py b/src/unilab/config/manipulation_params.py
index 030e615db..4006c1227 100644
--- a/src/unilab/config/manipulation_params.py
+++ b/src/unilab/config/manipulation_params.py
@@ -41,7 +41,7 @@ def ppo_config(env_name: str) -> config_dict.ConfigDict:
         seed=1,
         num_envs=cfg.num_envs,
         runner_class_name="OnPolicyRunner",
-        obs_groups={"default": ["policy"]},
+        obs_groups={"actor": ["actor"], "critic": ["actor"]},
         num_steps_per_env=cfg.num_steps_per_env,
         max_iterations=cfg.max_iterations,
         empirical_normalization=cfg.empirical_normalization,
diff --git a/src/unilab/envs/locomotion/g1/joystick.py b/src/unilab/envs/locomotion/g1/joystick.py
index 836c9513d..7b5148260 100644
--- a/src/unilab/envs/locomotion/g1/joystick.py
+++ b/src/unilab/envs/locomotion/g1/joystick.py
@@ -5,7 +5,6 @@
 import math
 from dataclasses import dataclass, field
 
-import gymnasium as gym
 import numpy as np
 from etils import epath
 
@@ -14,26 +13,9 @@
 from unilab.base.dtype_config import get_global_dtype
 from unilab.base.np_env import NpEnvState
 from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv
-from unilab.envs.locomotion.obs_config import ObsConfig
 from unilab.utils.math_utils import np_quat_mul, np_yaw_to_quat
 
 
-def _g1_obs_config() -> ObsConfig:
-    return ObsConfig(
-        obs_dict={
-            "vel": 3,
-            "gyro": 3,
-            "gravity": 3,
-            "diff": 29,
-            "dof_vel": 29,
-            "action": 29,
-            "cmd": 3,
-            "phase": 2,
-        },
-        actor_obs=["gyro", "gravity", "diff", "dof_vel", "action", "cmd", "phase"],
-    )
-
-
 @dataclass
 class InitState:
     pos = [0.0, 0.0, 0.754]
@@ -113,7 +95,6 @@ class G1JoystickPPOCfg(G1BaseCfg):
     init_state: InitState = field(default_factory=InitState)
     commands: Commands = field(default_factory=Commands)
     reward_config: RewardConfigPPO = field(default_factory=RewardConfigPPO)
-    obs_config: ObsConfig = field(default_factory=_g1_obs_config)
 
 
 @registry.env("G1JoystickFlatTerrain", sim_backend="mujoco")
@@ -135,9 +116,13 @@ def __init__(self, cfg: G1JoystickPPOCfg, num_envs=1, backend_type="mujoco"):
         if self._pose_weights.shape[0] != self._num_action:
             raise ValueError("pose_weights length mismatch")
 
-        self._init_obs_space()
         self._init_reward_functions()
 
+    @property
+    def obs_groups_spec(self) -> dict[str, int]:
+        # gyro(3) + gravity(3) + diff(29) + dof_vel(29) + action(29) + cmd(3) + phase(2) = 98
+        return {"actor": 98, "privileged": 3}
+
     def _init_reward_functions(self):
         self._reward_fns = {
             "tracking_lin_vel": self._reward_tracking_lin_vel,
@@ -151,17 +136,6 @@ def _init_reward_functions(self):
             "pose": self._reward_pose,
         }
 
-    def _init_obs_space(self):
-        obs_cfg = self._cfg.obs_config
-        self._observation_space = gym.spaces.Box(
-            low=-float("inf"), high=float("inf"), shape=(obs_cfg.total_dim,), dtype=float
-        )
-        self.actor_indices = obs_cfg.actor_indices
-
-    @property
-    def observation_space(self) -> gym.spaces.Box:
-        return self._observation_space  # type: ignore[no-any-return]
-
     def update_state(self, state: NpEnvState) -> NpEnvState:
         linvel = self.get_local_linvel()
         gyro = self.get_gyro()
@@ -180,16 +154,19 @@ def update_state(self, state: NpEnvState) -> NpEnvState:
         obs = self._compute_obs(state.info, linvel, gyro, gravity, dof_pos, dof_vel)
         return state.replace(obs=obs, reward=reward, terminated=terminated)
 
-    def _compute_obs(self, info: dict, linvel, gyro, gravity, dof_pos, dof_vel) -> np.ndarray:
+    def _compute_obs(
+        self, info: dict, linvel, gyro, gravity, dof_pos, dof_vel
+    ) -> dict[str, np.ndarray]:
         diff = dof_pos - self.default_angles
         command = info["commands"]
         last_actions = info.get("current_actions", np.zeros_like(diff))
         gait_phase = info.get("gait_phase", np.zeros((self._num_envs, 2), dtype=get_global_dtype()))
-        return np.concatenate(
-            [linvel, gyro, -gravity, diff, dof_vel, last_actions, command, gait_phase],
+        actor = np.concatenate(
+            [gyro, -gravity, diff, dof_vel, last_actions, command, gait_phase],
             axis=1,
             dtype=get_global_dtype(),
         )
+        return {"actor": actor, "privileged": linvel}
 
     def get_obs_structure(self) -> dict:
         """Return observation structure for symmetry augmentation."""
diff --git a/src/unilab/envs/locomotion/g1/joystick_sac.py b/src/unilab/envs/locomotion/g1/joystick_sac.py
index e332eaf56..702ca7f25 100644
--- a/src/unilab/envs/locomotion/g1/joystick_sac.py
+++ b/src/unilab/envs/locomotion/g1/joystick_sac.py
@@ -13,29 +13,6 @@
 from unilab.base.dtype_config import get_global_dtype
 from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv
 from unilab.envs.locomotion.g1.joystick import G1JoystickPPO, InitState
-from unilab.envs.locomotion.obs_config import ObsConfig
-
-
-def _g1_sac_obs_config() -> ObsConfig:
-    """ObsConfig for G1 (29 DoF) SAC.
-
-    Must stay consistent with G1JoystickPPO._compute_obs() which outputs:
-        vel(3) + gyro(3) + gravity(3) + diff(29) + dof_vel(29)
-        + action(29) + cmd(3) + phase(2) = 101 total
-    """
-    return ObsConfig(
-        obs_dict={
-            "vel": 3,
-            "gyro": 3,
-            "gravity": 3,
-            "diff": 29,
-            "dof_vel": 29,
-            "action": 29,
-            "cmd": 3,
-            "phase": 2,
-        },
-        actor_obs=["gyro", "gravity", "diff", "dof_vel", "action", "cmd", "phase"],
-    )
 
 
 @dataclass
@@ -128,7 +105,6 @@ class G1JoystickSACCfg(G1BaseCfg):
     commands: Commands = field(default_factory=Commands)
     reward_config: RewardConfigSAC = field(default_factory=RewardConfigSAC)
     control_config: ControlConfigSAC = field(default_factory=ControlConfigSAC)  # type: ignore[assignment]
-    obs_config: ObsConfig = field(default_factory=_g1_sac_obs_config)
 
 
 @registry.env("G1WalkTaskMjSAC", sim_backend="mujoco")
@@ -158,7 +134,6 @@ def __init__(self, cfg: G1JoystickSACCfg, num_envs=1, backend_type="mujoco"):
             degree=0.002,  # 调整速度加快
         )
 
-        self._init_obs_space()
         self._init_reward_functions()
 
     def _init_reward_functions(self):
diff --git a/src/unilab/envs/locomotion/g1/motion_tracking/tracking.py b/src/unilab/envs/locomotion/g1/motion_tracking/tracking.py
index a1eb664cd..b2711e0c0 100644
--- a/src/unilab/envs/locomotion/g1/motion_tracking/tracking.py
+++ b/src/unilab/envs/locomotion/g1/motion_tracking/tracking.py
@@ -4,7 +4,6 @@
 
 from dataclasses import dataclass, field
 
-import gymnasium as gym
 import mujoco
 import numpy as np
 from etils import epath
@@ -173,21 +172,17 @@ def __init__(self, cfg: G1MotionTrackingCfg, num_envs=1, backend_type="mujoco"):
         self.body_quat_relative_w[:, :, 0] = 1.0  # Initialize to identity quaternion
 
         self._enable_reward_log = True
-        self._init_obs_space()
         self._init_reward_functions()
 
-    def _init_obs_space(self):
-        # Actor observations: motion_anchor_pos_b (3), motion_anchor_ori_b (6),
-        # base_lin_vel (3), base_ang_vel (3), joint_pos (num_joints), joint_vel (num_joints), actions (num_joints)
-        num_actor_obs = 3 + 6 + 3 + 3 + self._num_action * 3
-
-        # Critic observations: actor + body_pos_b (num_bodies * 3) + body_ori_b (num_bodies * 6)
-        num_critic_obs = num_actor_obs + len(self._cfg.body_names) * 9
-
-        self._observation_space = gym.spaces.Box(
-            low=-float("inf"), high=float("inf"), shape=(num_critic_obs,), dtype=float
-        )
-        self.num_actor_obs = num_actor_obs
+    @property
+    def obs_groups_spec(self) -> dict[str, int]:
+        # Actor: motion_anchor_pos_b(3) + motion_anchor_ori_b(6) + linvel(3) + gyro(3)
+        #        + joint_pos(n) + joint_vel(n) + actions(n)
+        n = self._num_action
+        actor_dim = 3 + 6 + 3 + 3 + n * 3
+        # Privileged: body_pos_b(num_bodies*3) + body_ori_b(num_bodies*6)
+        privileged_dim = len(self._cfg.body_names) * 9
+        return {"actor": actor_dim, "privileged": privileged_dim}
 
     def _init_reward_functions(self):
         self._reward_fns = {
@@ -201,10 +196,6 @@ def _init_reward_functions(self):
             "joint_limit": self._reward_joint_limit,
         }
 
-    @property
-    def observation_space(self) -> gym.spaces.Box:
-        return self._observation_space  # type: ignore[no-any-return]
-
     def update_state(self, state: NpEnvState) -> NpEnvState:
         # Get current motion data
         motion_data = self.motion_sampler.get_current_motion()
@@ -329,8 +320,9 @@ def _compute_obs(
         dof_vel: np.ndarray,
         robot_body_pos_w: np.ndarray,
         robot_body_quat_w: np.ndarray,
-    ) -> np.ndarray:
-        """Compute observations (critic observations, actor is subset)."""
+    ) -> dict[str, np.ndarray]:
+        """Compute observations as dict with actor and privileged groups."""
+        num_envs = linvel.shape[0]
         # Get anchor states
         anchor_pos_w = motion_data.body_pos_w[:, self.anchor_body_idx]
         anchor_quat_w = motion_data.body_quat_w[:, self.anchor_body_idx]
@@ -347,7 +339,7 @@ def _compute_obs(
             robot_anchor_pos_w, robot_anchor_quat_w, anchor_pos_w, anchor_quat_w
         )
         motion_anchor_ori_mat = matrix_from_quat(motion_anchor_ori_rel)
-        motion_anchor_ori_b = motion_anchor_ori_mat[:, :, :2].reshape(self._num_envs, 6)
+        motion_anchor_ori_b = motion_anchor_ori_mat[:, :, :2].reshape(num_envs, 6)
 
         # Joint positions and velocities
         joint_pos_rel = dof_pos - self.default_angles
@@ -368,12 +360,12 @@ def _compute_obs(
             dtype=get_global_dtype(),
         )
 
-        # Robot body positions in robot anchor frame
+        # Robot body positions in robot anchor frame (privileged)
         robot_body_pos_b = np.zeros(
-            (self._num_envs, len(self._cfg.body_names), 3), dtype=get_global_dtype()
+            (num_envs, len(self._cfg.body_names), 3), dtype=get_global_dtype()
         )
         robot_body_ori_b = np.zeros(
-            (self._num_envs, len(self._cfg.body_names), 6), dtype=get_global_dtype()
+            (num_envs, len(self._cfg.body_names), 6), dtype=get_global_dtype()
         )
         for i in range(len(self._cfg.body_names)):
             pos_b, ori_b = subtract_frame_transforms(
@@ -384,20 +376,18 @@ def _compute_obs(
             )
             robot_body_pos_b[:, i] = pos_b
             ori_mat = matrix_from_quat(ori_b)
-            robot_body_ori_b[:, i] = ori_mat[:, :, :2].reshape(self._num_envs, 6)
+            robot_body_ori_b[:, i] = ori_mat[:, :, :2].reshape(num_envs, 6)
 
-        # Critic observations (privileged)
-        critic_obs = np.concatenate(
+        privileged_obs = np.concatenate(
             [
-                actor_obs,
-                robot_body_pos_b.reshape(self._num_envs, -1),
-                robot_body_ori_b.reshape(self._num_envs, -1),
+                robot_body_pos_b.reshape(num_envs, -1),
+                robot_body_ori_b.reshape(num_envs, -1),
             ],
             axis=1,
             dtype=get_global_dtype(),
         )
 
-        return critic_obs
+        return {"actor": actor_obs, "privileged": privileged_obs}
 
     def _compute_reward(
         self,
@@ -594,7 +584,4 @@ def reset(self, env_indices: np.ndarray):
             info, motion_data, linvel, gyro, dof_pos, dof_vel, robot_body_pos_w, robot_body_quat_w
         )
 
-        # Actor observations are first num_actor_obs elements
-        actor_obs = obs[:, : self.num_actor_obs]
-
-        return actor_obs, obs, info
+        return obs, obs, info
diff --git a/src/unilab/envs/locomotion/go1/joystick.py b/src/unilab/envs/locomotion/go1/joystick.py
index 8042c1502..402e1e172 100644
--- a/src/unilab/envs/locomotion/go1/joystick.py
+++ b/src/unilab/envs/locomotion/go1/joystick.py
@@ -2,7 +2,6 @@
 
 from dataclasses import dataclass, field
 
-import gymnasium as gym
 import numpy as np
 from etils import epath
 
@@ -11,7 +10,6 @@
 from unilab.base.dtype_config import get_global_dtype
 from unilab.base.np_env import NpEnvState
 from unilab.envs.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv
-from unilab.envs.locomotion.obs_config import ObsConfig
 from unilab.utils.math_utils import np_quat_mul, np_yaw_to_quat
 
 
@@ -68,22 +66,6 @@ class Domain_Rand:
     max_force = [1, 1, 0.5]
 
 
-def _go1_obs_config() -> ObsConfig:
-    return ObsConfig(
-        obs_dict={
-            "vel": 3,
-            "gyro": 3,
-            "gravity": 3,
-            "diff": 12,
-            "dof_vel": 12,
-            "action": 12,
-            "cmd": 3,
-            "phase": 4,
-        },
-        actor_obs=["gyro", "gravity", "diff", "dof_vel", "action", "cmd", "phase"],
-    )
-
-
 @registry.envcfg("Go1JoystickFlatTerrain")
 @dataclass
 class Go1JoystickCfg(Go1BaseCfg):
@@ -94,7 +76,6 @@ class Go1JoystickCfg(Go1BaseCfg):
     reward_config: RewardConfig = field(default_factory=RewardConfig)
     sensor: JoystickSensor = field(default_factory=JoystickSensor)  # type: ignore[assignment]
     domain_rand: Domain_Rand = field(default_factory=Domain_Rand)
-    obs_config: ObsConfig = field(default_factory=_go1_obs_config)
 
 
 @registry.env("Go1JoystickFlatTerrain", sim_backend="mujoco")
@@ -108,13 +89,17 @@ def __init__(self, cfg: Go1JoystickCfg, num_envs=1, backend_type="mujoco"):
         )
         super().__init__(cfg, backend, num_envs)
         self._enable_reward_log = True
-        self._init_obs_space()
         self._init_reward_functions()
         self.phase = np.zeros((num_envs,), dtype=np.float32)
         self.feet_phase = np.zeros((num_envs, len(cfg.sensor.feet_force)), dtype=np.float32)
         self.gait_frequency = 2
         self.feet_force = np.zeros((num_envs, len(cfg.sensor.feet_force), 3), dtype=np.float32)
 
+    @property
+    def obs_groups_spec(self) -> dict[str, int]:
+        # gyro(3) + gravity(3) + diff(12) + dof_vel(12) + action(12) + cmd(3) + phase(4) = 49
+        return {"actor": 49, "privileged": 3}
+
     def _init_reward_functions(self):
         self._reward_fns = {
             "tracking_lin_vel": self._reward_tracking_lin_vel,
@@ -127,17 +112,6 @@ def _init_reward_functions(self):
             "contact": self._reward_contact,
         }
 
-    def _init_obs_space(self):
-        obs_cfg = self._cfg.obs_config
-        self._observation_space = gym.spaces.Box(
-            low=-float("inf"), high=float("inf"), shape=(obs_cfg.total_dim,), dtype=float
-        )
-        self.actor_indices = obs_cfg.actor_indices
-
-    @property
-    def observation_space(self) -> gym.spaces.Box:
-        return self._observation_space  # type: ignore[no-any-return]
-
     def update_state(self, state: NpEnvState) -> NpEnvState:
         self.phase = np.fmod(self.phase + self._cfg.ctrl_dt * self.gait_frequency, 1.0)
         self.feet_phase[:, 0] = self.phase
@@ -163,15 +137,16 @@ def update_state(self, state: NpEnvState) -> NpEnvState:
 
     def _compute_obs(
         self, info: dict, linvel, gyro, gravity, dof_pos, dof_vel, feet_phase
-    ) -> np.ndarray:
+    ) -> dict[str, np.ndarray]:
         diff = dof_pos - self.default_angles
         command = info["commands"]
         last_actions = info.get("current_actions", np.zeros_like(diff))
-        return np.concatenate(
-            [linvel, gyro, -gravity, diff, dof_vel, last_actions, command, feet_phase],
+        actor = np.concatenate(
+            [gyro, -gravity, diff, dof_vel, last_actions, command, feet_phase],
             axis=1,
             dtype=get_global_dtype(),
         )
+        return {"actor": actor, "privileged": linvel}
 
     def _compute_reward(self, info: dict, linvel, gyro, dof_pos) -> np.ndarray:
         dtype = get_global_dtype()
diff --git a/src/unilab/envs/locomotion/go2/joystick.py b/src/unilab/envs/locomotion/go2/joystick.py
index b6413de19..3ca9ac473 100644
--- a/src/unilab/envs/locomotion/go2/joystick.py
+++ b/src/unilab/envs/locomotion/go2/joystick.py
@@ -2,7 +2,6 @@
 
 from dataclasses import dataclass, field
 
-import gymnasium as gym
 import numpy as np
 from etils import epath
 
@@ -10,7 +9,6 @@
 from unilab.base.backend import create_backend
 from unilab.base.np_env import NpEnvState
 from unilab.envs.locomotion.go2.base import Go2BaseCfg, Go2BaseEnv
-from unilab.envs.locomotion.obs_config import ObsConfig
 from unilab.utils.math_utils import np_quat_mul, np_yaw_to_quat
 
 
@@ -19,21 +17,6 @@ class InitState:
     pos = [0.0, 0.0, 0.42]
 
 
-def _go2_obs_config() -> ObsConfig:
-    return ObsConfig(
-        obs_dict={
-            "vel": 3,
-            "gyro": 3,
-            "gravity": 3,
-            "diff": 12,
-            "dof_vel": 12,
-            "action": 12,
-            "cmd": 3,
-        },
-        actor_obs=["gyro", "gravity", "diff", "dof_vel", "action", "cmd"],
-    )
-
-
 @dataclass
 class Commands:
     vel_limit = [
@@ -88,7 +71,6 @@ class Go2JoystickCfg(Go2BaseCfg):
     commands: Commands = field(default_factory=Commands)
     reward_config: RewardConfig = field(default_factory=RewardConfig)
     domain_rand: Domain_Rand = field(default_factory=Domain_Rand)
-    obs_config: ObsConfig = field(default_factory=_go2_obs_config)
 
 
 @registry.env("Go2JoystickFlatTerrain", sim_backend="mujoco")
@@ -101,9 +83,13 @@ def __init__(self, cfg: Go2JoystickCfg, num_envs=1, backend_type="mujoco"):
             backend_type, cfg.model_file, num_envs, cfg.sim_dt, base_name=cfg.asset.base_name
         )
         super().__init__(cfg, backend, num_envs)
-        self._init_obs_space()
         self._init_reward_functions()
 
+    @property
+    def obs_groups_spec(self) -> dict[str, int]:
+        # gyro(3) + gravity(3) + diff(12) + dof_vel(12) + action(12) + cmd(3) = 45
+        return {"actor": 45, "privileged": 3}
+
     def _init_reward_functions(self):
         self._reward_fns = {
             "tracking_lin_vel": self._reward_tracking_lin_vel,
@@ -118,24 +104,13 @@ def _init_reward_functions(self):
             "foot_drag_penalty": self._reward_foot_drag,
         }
 
-    def _init_obs_space(self):
-        obs_cfg = self._cfg.obs_config
-        self._observation_space = gym.spaces.Box(
-            low=-float("inf"), high=float("inf"), shape=(obs_cfg.total_dim,), dtype=float
-        )
-        self.actor_indices = obs_cfg.actor_indices
-
-    @property
-    def observation_space(self) -> gym.spaces.Box:
-        return self._observation_space  # type: ignore[no-any-return]
-
     def update_state(self, state: NpEnvState) -> NpEnvState:
         terminated = self._compute_termination()
         reward = self._compute_reward(state.info)
         obs = self._compute_obs(state.info)
         return state.replace(obs=obs, reward=reward, terminated=terminated)
 
-    def _compute_obs(self, info: dict) -> np.ndarray:
+    def _compute_obs(self, info: dict) -> dict[str, np.ndarray]:
         linvel = self.get_local_linvel()
         gyro = self.get_gyro()
         gravity = self._backend.get_sensor_data("upvector")
@@ -146,9 +121,8 @@ def _compute_obs(self, info: dict) -> np.ndarray:
         command = info["commands"]
         last_actions = info.get("current_actions", np.zeros_like(diff))
 
-        return np.concatenate(
-            [linvel, gyro, -gravity, diff, dof_vel, last_actions, command], axis=1
-        )
+        actor = np.concatenate([gyro, -gravity, diff, dof_vel, last_actions, command], axis=1)
+        return {"actor": actor, "privileged": linvel}
 
     def _compute_reward(self, info: dict) -> np.ndarray:
         reward = np.zeros((self._num_envs,), dtype=np.float32)
@@ -249,4 +223,5 @@ def reset(self, env_indices: np.ndarray):
         info["commands"][env_indices] = commands
 
         obs = self._compute_obs(info)
-        return obs[env_indices], obs[env_indices], {k: v[env_indices] for k, v in info.items()}
+        obs_subset = {k: v[env_indices] for k, v in obs.items()}
+        return obs_subset, obs_subset, {k: v[env_indices] for k, v in info.items()}
diff --git a/src/unilab/envs/locomotion/obs_config.py b/src/unilab/envs/locomotion/obs_config.py
deleted file mode 100644
index 37d9478bb..000000000
--- a/src/unilab/envs/locomotion/obs_config.py
+++ /dev/null
@@ -1,37 +0,0 @@
-from __future__ import annotations
-
-from dataclasses import dataclass, field
-
-
-@dataclass
-class ObsConfig:
-    """Observation layout for locomotion environments.
-
-    ``obs_dict`` maps observation name → dimension in the order the
-    fields are concatenated into the full observation vector.
-    ``actor_obs`` lists the subset of keys exposed to the policy.
-    ``actor_indices`` derives the corresponding flat index slice
-    automatically, so callers never have to compute offsets by hand.
-    """
-
-    obs_dict: dict[str, int] = field(default_factory=dict)
-    actor_obs: list[str] = field(default_factory=list)
-
-    @property
-    def total_dim(self) -> int:
-        return sum(self.obs_dict.values())
-
-    @property
-    def actor_indices(self) -> list[int]:
-        """Flat indices into the full obs vector for the ``actor_obs`` keys."""
-        offsets: dict[str, int] = {}
-        s = 0
-        for k, v in self.obs_dict.items():
-            offsets[k] = s
-            s += v
-        indices: list[int] = []
-        for key in self.actor_obs:
-            if key in offsets:
-                start = offsets[key]
-                indices.extend(range(start, start + self.obs_dict[key]))
-        return indices
diff --git a/src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py b/src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py
index a30021bed..7614773a8 100644
--- a/src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py
+++ b/src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py
@@ -3,7 +3,6 @@
 from dataclasses import dataclass, field
 from typing import Optional, Tuple
 
-import gymnasium as gym
 import numpy as np
 from etils import epath
 
@@ -36,17 +35,6 @@ class RewardConfig:
     reset_z_threshold: float = 0.125  # metres
 
 
-@dataclass
-class obs_cfg:
-    obs_dict = {
-        "dof_pos_norm": 16,
-        "targets": 16,
-        "ball_pos": 3,
-    }  # 'obs_name': dim
-    actor_obs = ["dof_pos_norm", "targets", "ball_pos"]
-    num_lag_steps: int = 3
-
-
 @dataclass
 class DomainRandConfig:
     # Small random perturbation added to each joint at reset (rad).
@@ -64,7 +52,6 @@ class AllegroRotationCfg(AllegroBaseCfg):
     max_episode_seconds: float = 20.0  # same as HORA
     reward_config: RewardConfig = field(default_factory=RewardConfig)
     domain_rand: DomainRandConfig = field(default_factory=DomainRandConfig)
-    obs_config: obs_cfg = field(default_factory=obs_cfg)
 
     # World-frame unit vector around which to reward rotation.
     # Default (+z) = counterclockwise when viewed from above.
@@ -103,36 +90,15 @@ def __init__(self, cfg: AllegroRotationCfg, num_envs: int = 1, backend_type: str
         self._grasp_cache_loaded = False
 
         self._init_reward_functions()
-        self._init_obs_space()
 
     # ── Spaces ──────────────────────────────────────────────────────
 
-    def _init_obs_space(self):
-        obs_cfg = self._cfg.obs_config
-        num_obs_per_step = sum(obs_cfg.obs_dict.values())
-        num_obs = obs_cfg.num_lag_steps * num_obs_per_step
-        self._observation_space = gym.spaces.Box(
-            low=-np.inf, high=np.inf, shape=(num_obs,), dtype=float
-        )
-        self.actor_indices = self._get_actor_indices()
-
-    def _get_actor_indices(self):
-        s = 0
-        indices = []
-        for i in self._cfg.obs_config.actor_obs:
-            s = 0
-            for k, v in self._cfg.obs_config.obs_dict.items():
-                if k == i:
-                    for q in range(s, s + v):
-                        indices.append(q)
-                    s += v
-                    break
-                s += v
-        return indices
+    _NUM_OBS_PER_STEP = 35  # dof_pos_norm(16) + targets(16) + ball_pos(3)
+    _NUM_LAG_STEPS = 3
 
     @property
-    def observation_space(self) -> gym.spaces.Box:
-        return self._observation_space  # type: ignore[no-any-return]
+    def obs_groups_spec(self) -> dict[str, int]:
+        return {"actor": self._NUM_OBS_PER_STEP * self._NUM_LAG_STEPS}
 
     # ── Reward functions ─────────────────────────────────────────────
 
@@ -245,8 +211,8 @@ def update_state(self, state: NpEnvState) -> NpEnvState:
             np.zeros(
                 (
                     self._num_envs,
-                    self._cfg.obs_config.num_lag_steps,
-                    sum(self._cfg.obs_config.obs_dict.values()),
+                    self._NUM_LAG_STEPS,
+                    self._NUM_OBS_PER_STEP,
                 ),
                 dtype=self._np_dtype,
             ),
@@ -284,16 +250,16 @@ def _compute_rewards(self, state: NpEnvState) -> tuple:
 
         return total * self._cfg.ctrl_dt, log
 
-    def _get_obs(self, info: dict) -> np.ndarray:
+    def _get_obs(self, info: dict) -> dict[str, np.ndarray]:
         """Extract observation from info. Pure function - does not modify state."""
-        num_obs_per_step = sum(self._cfg.obs_config.obs_dict.values())
-        num_lag_steps = self._cfg.obs_config.num_lag_steps
         obs_lag_history = info.get(
             "obs_lag_history",
-            np.zeros((self._num_envs, num_lag_steps, num_obs_per_step), dtype=self._np_dtype),
+            np.zeros(
+                (self._num_envs, self._NUM_LAG_STEPS, self._NUM_OBS_PER_STEP), dtype=self._np_dtype
+            ),
         )
         num_envs = obs_lag_history.shape[0]
-        return np.asarray(obs_lag_history.reshape(num_envs, -1))
+        return {"actor": np.asarray(obs_lag_history.reshape(num_envs, -1))}
 
     # ── Reset ────────────────────────────────────────────────────────
 
@@ -366,8 +332,8 @@ def reset(self, env_indices: np.ndarray) -> Tuple[np.ndarray, np.ndarray, dict]:
             init_obs[:, None, :],
             (
                 num_reset,
-                self._cfg.obs_config.num_lag_steps,
-                sum(self._cfg.obs_config.obs_dict.values()),
+                self._NUM_LAG_STEPS,
+                self._NUM_OBS_PER_STEP,
             ),
         ).copy()  # (num_reset, num_lag_steps, num_obs_per_step)
 
diff --git a/src/unilab/envs/manipulation/inhand_rot_allegro/rotation_sac.py b/src/unilab/envs/manipulation/inhand_rot_allegro/rotation_sac.py
index efb98a226..f8b63fdd4 100644
--- a/src/unilab/envs/manipulation/inhand_rot_allegro/rotation_sac.py
+++ b/src/unilab/envs/manipulation/inhand_rot_allegro/rotation_sac.py
@@ -2,13 +2,12 @@
 
 from dataclasses import dataclass, field
 
-import gymnasium as gym
 import numpy as np
 from etils import epath
-from unilab.envs.backend import create_backend
-from unilab.envs.np_env import NpEnvState
 
 from unilab.base import registry
+from unilab.base.backend import create_backend
+from unilab.base.np_env import NpEnvState
 from unilab.envs.manipulation.inhand_rot_allegro.base import AllegroBaseCfg, AllegroBaseMjEnv
 from unilab.utils.math_utils import np_quat_conjugate, np_quat_mul, np_quat_to_axis_angle
 
@@ -89,19 +88,15 @@ def __init__(self, cfg: AllegroRotationSacCfg, num_envs: int = 1, backend_type:
         self._grasp_cache_loaded = False
 
         self._init_reward_functions()
-        self._init_obs_space()
 
     # ── Spaces ──────────────────────────────────────────────────────
 
-    def _init_obs_space(self):
-        # Last 3 timesteps of [dof_pos_norm(16) + targets(16) + ball_pos(3)] = 3 * 35 = 105
-        self._observation_space = gym.spaces.Box(
-            low=-np.inf, high=np.inf, shape=(105,), dtype=float
-        )
+    _NUM_OBS_PER_STEP = 35  # dof_pos_norm(16) + targets(16) + ball_pos(3)
+    _NUM_LAG_STEPS = 3
 
     @property
-    def observation_space(self) -> gym.spaces.Box:
-        return self._observation_space  # type: ignore[return-value, no-any-return]
+    def obs_groups_spec(self) -> dict[str, int]:
+        return {"actor": self._NUM_OBS_PER_STEP * self._NUM_LAG_STEPS}
 
     # ── Reward functions ─────────────────────────────────────────────
 
@@ -252,13 +247,16 @@ def _compute_rewards(self, state: NpEnvState) -> tuple:
 
         return total * self._cfg.ctrl_dt, log
 
-    def _get_obs(self, info: dict) -> np.ndarray:
+    def _get_obs(self, info: dict) -> dict[str, np.ndarray]:
         """Extract observation from info. Pure function - does not modify state."""
         obs_lag_history = info.get(
-            "obs_lag_history", np.zeros((self._num_envs, 3, 35), dtype=self._np_dtype)
+            "obs_lag_history",
+            np.zeros(
+                (self._num_envs, self._NUM_LAG_STEPS, self._NUM_OBS_PER_STEP), dtype=self._np_dtype
+            ),
         )
         num_envs = obs_lag_history.shape[0]
-        return np.asarray(obs_lag_history.reshape(num_envs, -1))
+        return {"actor": np.asarray(obs_lag_history.reshape(num_envs, -1))}
 
     # ── Reset ────────────────────────────────────────────────────────
 
@@ -327,7 +325,8 @@ def reset(self, env_indices: np.ndarray) -> tuple:
             [dof_pos_norm, init_ctrl, ball_pos_f32], axis=1
         )  # (num_reset, 35)
         obs_lag_history = np.broadcast_to(
-            init_obs[:, None, :], (num_reset, 3, 35)
+            init_obs[:, None, :],
+            (num_reset, self._NUM_LAG_STEPS, self._NUM_OBS_PER_STEP),
         ).copy()  # (num_reset, 3, 35)
 
         info = {
diff --git a/src/unilab/utils/device_utils.py b/src/unilab/utils/device_utils.py
index 0c54b7df4..031b5cdbe 100644
--- a/src/unilab/utils/device_utils.py
+++ b/src/unilab/utils/device_utils.py
@@ -15,10 +15,9 @@ def get_default_device() -> str:
 def get_env_dims(env_name: str, sim_backend: str = "mujoco") -> tuple[int, int]:
     """Get observation and action dimensions from environment."""
     env = registry.make(env_name, num_envs=1, sim_backend=sim_backend)
-    obs_shape = env.observation_space.shape
+    obs_dim = sum(env.obs_groups_spec.values())
     action_shape = env.action_space.shape
-    assert obs_shape is not None and action_shape is not None
-    obs_dim = obs_shape[0]
+    assert action_shape is not None
     action_dim = action_shape[0]
     env.close()  # type: ignore[attr-defined]
     return obs_dim, action_dim
diff --git a/src/unilab/utils/obs_utils.py b/src/unilab/utils/obs_utils.py
new file mode 100644
index 000000000..0e74da4c0
--- /dev/null
+++ b/src/unilab/utils/obs_utils.py
@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import numpy as np
+
+
+def flatten_obs_dict(obs: dict[str, np.ndarray]) -> np.ndarray:
+    """Concatenate obs groups in insertion order -> flat (N, total_dim) array."""
+    return np.concatenate(list(obs.values()), axis=1)
diff --git a/tests/base/test_np_env.py b/tests/base/test_np_env.py
new file mode 100644
index 000000000..983fc18e5
--- /dev/null
+++ b/tests/base/test_np_env.py
@@ -0,0 +1,424 @@
+"""Tests for NpEnvState and NpEnv dict-obs contract.
+
+These tests use a minimal concrete NpEnv stub — no MuJoCo required.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Tuple
+from unittest.mock import MagicMock
+
+import gymnasium as gym
+import numpy as np
+import pytest
+
+from unilab.base.base import EnvCfg
+from unilab.base.np_env import NpEnv, NpEnvState
+
+# ---------------------------------------------------------------------------
+# Fixtures: minimal concrete NpEnv
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _StubCfg(EnvCfg):
+    max_episode_seconds: float = 1.0
+    ctrl_dt: float = 0.1
+    sim_dt: float = 0.01
+
+
+class _StubNpEnv(NpEnv):
+    """Concrete NpEnv for testing — no physics, deterministic outputs."""
+
+    OBS_SPEC = {"actor": 5, "privileged": 2}
+
+    def __init__(self, num_envs: int = 4):
+        cfg = _StubCfg()
+        backend = MagicMock()
+        backend.backend_type = "mujoco"
+        backend.step = MagicMock()
+        super().__init__(cfg, backend, num_envs)
+        self._reset_count = 0
+
+    @property
+    def obs_groups_spec(self) -> dict[str, int]:
+        return self.OBS_SPEC
+
+    @property
+    def action_space(self) -> gym.Space:
+        return gym.spaces.Box(low=-1.0, high=1.0, shape=(3,), dtype=np.float32)
+
+    def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray:
+        return actions
+
+    def update_state(self, state: NpEnvState) -> NpEnvState:
+        obs = {
+            "actor": np.ones((self._num_envs, 5), dtype=np.float32),
+            "privileged": np.full((self._num_envs, 2), 0.5, dtype=np.float32),
+        }
+        return state.replace(
+            obs=obs,
+            reward=np.ones((self._num_envs,), dtype=np.float32),
+            terminated=np.zeros((self._num_envs,), dtype=bool),
+            truncated=np.zeros((self._num_envs,), dtype=bool),
+        )
+
+    def reset(
+        self, env_indices: np.ndarray
+    ) -> Tuple[dict[str, np.ndarray], dict[str, np.ndarray], dict]:
+        self._reset_count += 1
+        n = len(env_indices)
+        obs = {
+            "actor": np.zeros((n, 5), dtype=np.float32),
+            "privileged": np.zeros((n, 2), dtype=np.float32),
+        }
+        return obs, obs, {}
+
+
+class _TerminatingStubEnv(_StubNpEnv):
+    """Like _StubNpEnv but terminates specified envs each step."""
+
+    def __init__(self, num_envs: int = 4, terminate_indices: list[int] | None = None):
+        super().__init__(num_envs)
+        self._terminate_indices = terminate_indices or []
+
+    def update_state(self, state: NpEnvState) -> NpEnvState:
+        state = super().update_state(state)
+        terminated = np.zeros((self._num_envs,), dtype=bool)
+        terminated[self._terminate_indices] = True
+        return state.replace(terminated=terminated)
+
+
+# ---------------------------------------------------------------------------
+# NpEnvState tests
+# ---------------------------------------------------------------------------
+
+
+class TestNpEnvState:
+    def test_obs_is_dict(self):
+        obs = {"actor": np.zeros((2, 4)), "privileged": np.zeros((2, 1))}
+        state = NpEnvState(
+            obs=obs,
+            reward=np.zeros(2),
+            terminated=np.zeros(2, dtype=bool),
+            truncated=np.zeros(2, dtype=bool),
+            info={},
+        )
+        assert isinstance(state.obs, dict)
+        assert "actor" in state.obs
+        assert "privileged" in state.obs
+
+    def test_done_combines_terminated_and_truncated(self):
+        state = NpEnvState(
+            obs={"a": np.zeros((3, 1))},
+            reward=np.zeros(3),
+            terminated=np.array([True, False, False]),
+            truncated=np.array([False, False, True]),
+            info={},
+        )
+        done = state.done
+        np.testing.assert_array_equal(done, [True, False, True])
+
+    def test_done_both_true(self):
+        state = NpEnvState(
+            obs={"a": np.zeros((1, 1))},
+            reward=np.zeros(1),
+            terminated=np.array([True]),
+            truncated=np.array([True]),
+            info={},
+        )
+        assert state.done[0] is np.True_
+
+    def test_replace_preserves_type(self):
+        obs = {"actor": np.zeros((2, 3))}
+        state = NpEnvState(
+            obs=obs,
+            reward=np.zeros(2),
+            terminated=np.zeros(2, dtype=bool),
+            truncated=np.zeros(2, dtype=bool),
+            info={},
+        )
+        new_obs = {"actor": np.ones((2, 3))}
+        state2 = state.replace(obs=new_obs)
+        assert isinstance(state2, NpEnvState)
+        assert isinstance(state2.obs, dict)
+        np.testing.assert_array_equal(state2.obs["actor"], 1.0)
+        # Original unchanged
+        np.testing.assert_array_equal(state.obs["actor"], 0.0)
+
+    def test_replace_partial_update(self):
+        state = NpEnvState(
+            obs={"a": np.zeros((1, 1))},
+            reward=np.array([0.0]),
+            terminated=np.array([False]),
+            truncated=np.array([False]),
+            info={"k": "v"},
+        )
+        state2 = state.replace(reward=np.array([5.0]))
+        assert state2.reward[0] == 5.0
+        assert state2.obs is state.obs  # not copied
+
+
+# ---------------------------------------------------------------------------
+# NpEnv — obs_groups_spec / observation_space
+# ---------------------------------------------------------------------------
+
+
+class TestNpEnvObsSpec:
+    def test_obs_groups_spec_returns_dict(self):
+        env = _StubNpEnv(num_envs=2)
+        spec = env.obs_groups_spec
+        assert isinstance(spec, dict)
+        assert spec == {"actor": 5, "privileged": 2}
+
+    def test_observation_space_total_dim(self):
+        env = _StubNpEnv(num_envs=2)
+        space = env.observation_space
+        assert isinstance(space, gym.spaces.Box)
+        assert space.shape == (7,)  # 5 + 2
+
+    def test_observation_space_bounds(self):
+        env = _StubNpEnv(num_envs=1)
+        space = env.observation_space
+        assert np.all(space.low == -np.inf)
+        assert np.all(space.high == np.inf)
+
+    def test_bare_npenv_obs_groups_spec_raises(self):
+        """NpEnv.obs_groups_spec raises NotImplementedError if not overridden."""
+        # Access the property via the NpEnv class descriptor directly
+        with pytest.raises(NotImplementedError):
+            NpEnv.obs_groups_spec.fget(MagicMock(spec=NpEnv))  # type: ignore[union-attr]
+
+
+# ---------------------------------------------------------------------------
+# NpEnv.init_state — dict obs allocation
+# ---------------------------------------------------------------------------
+
+
+class TestNpEnvInitState:
+    def test_init_state_returns_np_env_state(self):
+        env = _StubNpEnv(num_envs=4)
+        state = env.init_state()
+        assert isinstance(state, NpEnvState)
+
+    def test_init_state_obs_is_dict(self):
+        env = _StubNpEnv(num_envs=4)
+        env.init_state()
+        assert isinstance(env.state.obs, dict)
+
+    def test_init_state_obs_keys_match_spec(self):
+        env = _StubNpEnv(num_envs=4)
+        env.init_state()
+        assert set(env.state.obs.keys()) == {"actor", "privileged"}
+
+    def test_init_state_obs_shapes(self):
+        env = _StubNpEnv(num_envs=4)
+        env.init_state()
+        assert env.state.obs["actor"].shape == (4, 5)
+        assert env.state.obs["privileged"].shape == (4, 2)
+
+    def test_init_state_obs_zeros(self):
+        """Initial obs should be zeros (before reset fills them)."""
+        env = _StubNpEnv(num_envs=2)
+        # init_state internally calls _reset_done_envs which calls reset(),
+        # so the obs after init_state will be the reset values
+        env.init_state()
+        # After init_state, reset was called (terminated=True initially)
+        # Our stub resets to zeros, so should be zeros
+        np.testing.assert_array_equal(env.state.obs["actor"], 0.0)
+
+    def test_init_state_triggers_reset(self):
+        env = _StubNpEnv(num_envs=2)
+        env.init_state()
+        # init_state sets terminated=True initially, so _reset_done_envs triggers
+        assert env._reset_count > 0
+
+    def test_init_state_reward_shape(self):
+        env = _StubNpEnv(num_envs=3)
+        env.init_state()
+        assert env.state.reward.shape == (3,)
+
+    def test_init_state_steps_initialized(self):
+        env = _StubNpEnv(num_envs=2)
+        env.init_state()
+        assert "steps" in env.state.info
+        np.testing.assert_array_equal(env.state.info["steps"], 0)
+
+
+# ---------------------------------------------------------------------------
+# NpEnv.step — dict obs flow
+# ---------------------------------------------------------------------------
+
+
+class TestNpEnvStep:
+    def test_step_returns_state_with_dict_obs(self):
+        env = _StubNpEnv(num_envs=2)
+        env.init_state()
+        actions = np.zeros((2, 3))
+        state = env.step(actions)
+        assert isinstance(state.obs, dict)
+        assert "actor" in state.obs
+        assert "privileged" in state.obs
+
+    def test_step_obs_values_from_update_state(self):
+        env = _StubNpEnv(num_envs=2)
+        env.init_state()
+        state = env.step(np.zeros((2, 3)))
+        # _StubNpEnv.update_state fills actor=1.0, privileged=0.5
+        np.testing.assert_array_equal(state.obs["actor"], 1.0)
+        np.testing.assert_array_equal(state.obs["privileged"], 0.5)
+
+    def test_step_increments_counter(self):
+        env = _StubNpEnv(num_envs=1)
+        env.init_state()
+        env.step(np.zeros((1, 3)))
+        assert env.step_counter == 1
+        env.step(np.zeros((1, 3)))
+        assert env.step_counter == 2
+
+    def test_step_increments_steps_info(self):
+        env = _StubNpEnv(num_envs=2)
+        env.init_state()
+        env.step(np.zeros((2, 3)))
+        np.testing.assert_array_equal(env.state.info["steps"], 1)
+
+    def test_step_auto_inits_if_no_state(self):
+        env = _StubNpEnv(num_envs=1)
+        # Don't call init_state
+        state = env.step(np.zeros((1, 3)))
+        assert isinstance(state.obs, dict)
+
+
+# ---------------------------------------------------------------------------
+# NpEnv._reset_done_envs — dict final_observation handling
+# ---------------------------------------------------------------------------
+
+
+class TestResetDoneEnvs:
+    def test_done_envs_get_reset_obs(self):
+        env = _TerminatingStubEnv(num_envs=4, terminate_indices=[0, 2])
+        env.init_state()
+        state = env.step(np.zeros((4, 3)))
+        # Envs 0, 2 terminated — their obs should be reset values (zeros from stub)
+        np.testing.assert_array_equal(state.obs["actor"][0], 0.0)
+        np.testing.assert_array_equal(state.obs["actor"][2], 0.0)
+        # Envs 1, 3 not terminated — their obs should be update_state values (ones)
+        np.testing.assert_array_equal(state.obs["actor"][1], 1.0)
+        np.testing.assert_array_equal(state.obs["actor"][3], 1.0)
+
+    def test_final_observation_is_dict(self):
+        env = _TerminatingStubEnv(num_envs=2, terminate_indices=[0])
+        env.init_state()
+        env.step(np.zeros((2, 3)))
+        assert isinstance(env.state.info["final_observation"], dict)
+        assert "actor" in env.state.info["final_observation"]
+        assert "privileged" in env.state.info["final_observation"]
+
+    def test_final_observation_captures_pre_reset_obs(self):
+        env = _TerminatingStubEnv(num_envs=2, terminate_indices=[0])
+        env.init_state()
+        env.step(np.zeros((2, 3)))
+        # update_state fills actor=1.0 before reset replaces it with 0.0
+        # final_observation should capture the pre-reset obs (1.0)
+        np.testing.assert_array_equal(env.state.info["final_observation"]["actor"][0], 1.0)
+
+    def test_final_observation_mask(self):
+        env = _TerminatingStubEnv(num_envs=3, terminate_indices=[1])
+        env.init_state()
+        env.step(np.zeros((3, 3)))
+        mask = env.state.info["_final_observation"]
+        np.testing.assert_array_equal(mask, [False, True, False])
+
+    def test_no_termination_skips_reset(self):
+        env = _TerminatingStubEnv(num_envs=2, terminate_indices=[])
+        env.init_state()
+        reset_count_after_init = env._reset_count
+        # Step without any terminations
+        env.step(np.zeros((2, 3)))
+        # No additional resets should have occurred
+        assert env._reset_count == reset_count_after_init
+
+    def test_steps_reset_for_done_envs(self):
+        env = _TerminatingStubEnv(num_envs=3, terminate_indices=[0, 2])
+        env.init_state()
+        env.step(np.zeros((3, 3)))
+        # After step, terminated envs had their steps reset to 0
+        # then +1 from step, but _reset_done_envs is called AFTER info["steps"] += 1,
+        # and it sets steps[env_indices] = 0
+        assert env.state.info["steps"][0] == 0
+        assert env.state.info["steps"][2] == 0
+        assert env.state.info["steps"][1] == 1
+
+    def test_truncation_triggers_reset(self):
+        """max_episode_steps triggers truncation → reset."""
+        env = _StubNpEnv(num_envs=1)
+        env.init_state()
+        # max_episode_steps = 1.0 / 0.1 = 10
+        for _ in range(10):
+            env.step(np.zeros((1, 3)))
+        # After 10 steps, env should have been truncated and reset
+        assert env.state.info["steps"][0] == 0
+
+
+# ---------------------------------------------------------------------------
+# Edge cases
+# ---------------------------------------------------------------------------
+
+
+class TestEdgeCases:
+    def test_single_obs_group_symmetric(self):
+        """Env with only "actor" (no privileged) works correctly."""
+
+        class _SymmetricEnv(_StubNpEnv):
+            OBS_SPEC = {"actor": 10}
+
+            def update_state(self, state):
+                obs = {"actor": np.ones((self._num_envs, 10), dtype=np.float32)}
+                return state.replace(
+                    obs=obs,
+                    reward=np.zeros(self._num_envs, dtype=np.float32),
+                    terminated=np.zeros(self._num_envs, dtype=bool),
+                    truncated=np.zeros(self._num_envs, dtype=bool),
+                )
+
+            def reset(self, env_indices):
+                n = len(env_indices)
+                obs = {"actor": np.zeros((n, 10), dtype=np.float32)}
+                return obs, obs, {}
+
+        env = _SymmetricEnv(num_envs=2)
+        env.init_state()
+        state = env.step(np.zeros((2, 3)))
+        assert set(state.obs.keys()) == {"actor"}
+        assert state.obs["actor"].shape == (2, 10)
+        assert env.observation_space.shape == (10,)
+
+    def test_many_obs_groups(self):
+        """Env with 3+ obs groups allocates correctly."""
+
+        class _MultiGroupEnv(_StubNpEnv):
+            OBS_SPEC = {"actor": 4, "privileged": 2, "history": 8}
+
+            def update_state(self, state):
+                obs = {
+                    k: np.ones((self._num_envs, d), dtype=np.float32)
+                    for k, d in self.OBS_SPEC.items()
+                }
+                return state.replace(
+                    obs=obs,
+                    reward=np.zeros(self._num_envs, dtype=np.float32),
+                    terminated=np.zeros(self._num_envs, dtype=bool),
+                    truncated=np.zeros(self._num_envs, dtype=bool),
+                )
+
+            def reset(self, env_indices):
+                n = len(env_indices)
+                obs = {k: np.zeros((n, d), dtype=np.float32) for k, d in self.OBS_SPEC.items()}
+                return obs, obs, {}
+
+        env = _MultiGroupEnv(num_envs=1)
+        env.init_state()
+        state = env.step(np.zeros((1, 3)))
+        assert set(state.obs.keys()) == {"actor", "privileged", "history"}
+        assert env.observation_space.shape == (14,)  # 4+2+8
diff --git a/tests/base/test_registry.py b/tests/base/test_registry.py
index 97310c518..59f855447 100644
--- a/tests/base/test_registry.py
+++ b/tests/base/test_registry.py
@@ -50,6 +50,10 @@ def observation_space(self):
     def action_space(self):
         return gym.spaces.Box(low=-1.0, high=1.0, shape=(2,), dtype=np.float32)
 
+    @property
+    def obs_groups_spec(self) -> dict[str, int]:
+        return {"actor": 4}
+
     @property
     def state(self):
         return None
diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py
index 4bd7abe2a..09bf04552 100644
--- a/tests/config/test_locomotion_params.py
+++ b/tests/config/test_locomotion_params.py
@@ -100,12 +100,12 @@ def test_offpolicy_invalid_algo_raises():
 
 
 def test_ppo_config_go1_has_actor_obs_group():
-    """Go1 branch sets actor obs group in addition to default."""
+    """Go1 branch sets actor and critic obs groups."""
     cfg = ppo_config("Go1JoystickFlatTerrain")
     d = cfg.to_dict()
     obs = d["obs_groups"]
     assert "actor" in obs, "Go1 ppo_config must have actor obs group"
-    assert "default" in obs
+    assert "critic" in obs
 
 
 def test_ppo_config_go1_max_iterations():
@@ -135,9 +135,10 @@ def test_ppo_config_unknown_env_uses_defaults():
     """Unknown env_name falls through all branches → uses default obs_groups."""
     cfg = ppo_config("SomeOtherEnv")
     d = cfg.to_dict()
-    assert "default" in d["obs_groups"]
-    # no actor obs group for unknown env
-    assert "actor" not in d["obs_groups"]
+    assert "actor" in d["obs_groups"]
+    assert "critic" in d["obs_groups"]
+    # no privileged obs for unknown env
+    assert "privileged" not in d["obs_groups"].get("critic", [])
 
 
 # ---------------------------------------------------------------------------
diff --git a/tests/config/test_manipulation_params.py b/tests/config/test_manipulation_params.py
new file mode 100644
index 000000000..860bb70ef
--- /dev/null
+++ b/tests/config/test_manipulation_params.py
@@ -0,0 +1,146 @@
+"""Tests for manipulation_params config builders."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+
+import pytest
+
+from unilab.config.manipulation_params import appo_config, offpolicy_config, ppo_config
+
+# ---------------------------------------------------------------------------
+# ppo_config
+# ---------------------------------------------------------------------------
+
+
+def test_ppo_config_required_keys():
+    cfg = ppo_config("AllegroInhandRotation")
+    d = cfg.to_dict()
+    for key in ("seed", "num_envs", "obs_groups", "policy", "algorithm"):
+        assert key in d, f"ppo_config missing key: {key}"
+
+
+def test_ppo_config_obs_groups_has_actor():
+    cfg = ppo_config("AllegroInhandRotation")
+    d = cfg.to_dict()
+    assert "actor" in d["obs_groups"], "obs_groups must contain 'actor'"
+    assert "critic" in d["obs_groups"], "obs_groups must contain 'critic'"
+
+
+def test_ppo_config_actor_obs_group_contains_actor():
+    """actor obs group lists 'actor' key."""
+    cfg = ppo_config("AllegroInhandRotation")
+    d = cfg.to_dict()
+    assert "actor" in d["obs_groups"]["actor"]
+
+
+def test_ppo_config_allegro_num_envs():
+    cfg = ppo_config("AllegroInhandRotation")
+    assert cfg.num_envs == 16384
+
+
+def test_ppo_config_allegro_max_iterations():
+    cfg = ppo_config("AllegroInhandRotation")
+    assert cfg.max_iterations == 501
+
+
+def test_ppo_config_allegro_sac_raises():
+    with pytest.raises(NotImplementedError):
+        ppo_config("AllegroInhandRotationSac")
+
+
+def test_ppo_config_to_dict_no_raise():
+    cfg = ppo_config("AllegroInhandRotation")
+    d = cfg.to_dict()
+    assert isinstance(d, dict)
+
+
+# ---------------------------------------------------------------------------
+# appo_config
+# ---------------------------------------------------------------------------
+
+
+def test_appo_config_required_keys():
+    cfg = appo_config("AllegroInhandRotation")
+    d = cfg.to_dict()
+    for key in ("algo", "actor", "critic", "algorithm"):
+        assert key in d, f"appo_config missing key: {key}"
+
+
+def test_appo_config_algo_value():
+    cfg = appo_config("AllegroInhandRotation")
+    assert cfg.algo == "appo"
+
+
+def test_appo_config_allegro_num_envs():
+    cfg = appo_config("AllegroInhandRotation")
+    assert cfg.num_envs == 16384
+
+
+def test_appo_config_allegro_max_iterations():
+    cfg = appo_config("AllegroInhandRotation")
+    assert cfg.max_iterations == 501
+
+
+def test_appo_config_allegro_empirical_normalization():
+    cfg = appo_config("AllegroInhandRotation")
+    assert cfg.empirical_normalization is True
+
+
+def test_appo_config_sac_raises():
+    with pytest.raises(NotImplementedError):
+        appo_config("AllegroInhandRotationSac")
+
+
+def test_appo_config_to_dict_no_raise():
+    cfg = appo_config("AllegroInhandRotation")
+    d = cfg.to_dict()
+    assert isinstance(d, dict)
+
+
+def test_appo_config_deepcopy_safety():
+    """deepcopy + pop('class_name') must not affect the original config."""
+    cfg = appo_config("AllegroInhandRotation")
+    d = cfg.to_dict()
+    actor_copy = deepcopy(d["actor"])
+    actor_copy.pop("class_name", None)
+    d2 = cfg.to_dict()
+    assert "class_name" in d2["actor"]
+
+
+# ---------------------------------------------------------------------------
+# offpolicy_config — SAC
+# ---------------------------------------------------------------------------
+
+
+def test_offpolicy_sac_required_keys():
+    cfg = offpolicy_config("sac", "AllegroInhandRotationSac")
+    d = cfg.to_dict()
+    for key in ("algo", "num_envs", "batch_size", "replay_buffer_n"):
+        assert key in d, f"offpolicy_config(sac) missing key: {key}"
+    assert d["algo"] == "sac"
+
+
+def test_offpolicy_sac_allegro_rotation_raises():
+    with pytest.raises(NotImplementedError):
+        offpolicy_config("sac", "AllegroInhandRotation")
+
+
+def test_offpolicy_sac_allegro_rotation_sac_num_envs():
+    cfg = offpolicy_config("sac", "AllegroInhandRotationSac")
+    assert cfg.num_envs == 2048
+
+
+def test_offpolicy_sac_allegro_rotation_sac_max_iterations():
+    cfg = offpolicy_config("sac", "AllegroInhandRotationSac")
+    assert cfg.max_iterations == 25000
+
+
+# ---------------------------------------------------------------------------
+# offpolicy_config — invalid algo
+# ---------------------------------------------------------------------------
+
+
+def test_offpolicy_invalid_algo_raises():
+    with pytest.raises(ValueError, match="Unsupported"):
+        offpolicy_config("dqn", "AllegroInhandRotation")
diff --git a/tests/conftest.py b/tests/conftest.py
index 170c6f0fe..dfc3dd71b 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -58,6 +58,10 @@ def observation_space(self) -> gym.Space:
     def action_space(self) -> gym.Space:
         return self._act_space
 
+    @property
+    def obs_groups_spec(self) -> dict[str, int]:
+        return {"actor": _DUMMY_OBS_DIM}
+
     def close(self) -> None:
         pass
 
diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py
index 471862c21..ce89c945b 100644
--- a/tests/envs/test_env_configs.py
+++ b/tests/envs/test_env_configs.py
@@ -33,49 +33,42 @@
 # ---------------------------------------------------------------------------
 
 
-def test_g1_joystick_ppo_cfg_has_obs_config():
-    """G1JoystickPPOCfg must have obs_config (baseline sanity check)."""
-    from unilab.envs.locomotion.g1.joystick import G1JoystickPPOCfg
+def test_g1_joystick_ppo_cfg_obs_groups_spec():
+    """G1JoystickPPO must declare obs_groups_spec with actor and privileged groups."""
+    from unilab.envs.locomotion.g1.joystick import G1JoystickPPO, G1JoystickPPOCfg
 
     cfg = G1JoystickPPOCfg()
-    assert hasattr(cfg, "obs_config"), "G1JoystickPPOCfg missing obs_config"
-    assert hasattr(cfg.obs_config, "obs_dict")
-    assert hasattr(cfg.obs_config, "actor_obs")
+    # Access obs_groups_spec from the class (it's a class-level property)
+    # We need an instance to check, but can't instantiate without MuJoCo sim,
+    # so just verify the cfg no longer has obs_config (removed in dict obs refactor).
+    assert not hasattr(cfg, "obs_config"), "obs_config should have been removed"
 
 
-def test_g1_joystick_sac_cfg_has_obs_config():
-    """G1JoystickSACCfg must have obs_config because G1WalkTaskMjSAC inherits
-    _init_obs_space() from G1JoystickPPO which reads self._cfg.obs_config.
-
-    Regression test for:
-        AttributeError: 'G1JoystickSACCfg' object has no attribute 'obs_config'
-    """
+def test_g1_joystick_sac_cfg_no_obs_config():
+    """G1JoystickSACCfg should no longer have obs_config after dict obs refactor."""
     from unilab.envs.locomotion.g1.joystick_sac import G1JoystickSACCfg
 
     cfg = G1JoystickSACCfg()
-    assert hasattr(cfg, "obs_config"), (
-        "G1JoystickSACCfg must declare obs_config so that the inherited "
-        "G1JoystickPPO._init_obs_space() can resolve self._cfg.obs_config"
+    assert not hasattr(cfg, "obs_config"), (
+        "obs_config should have been removed in the dict obs refactor"
     )
-    assert hasattr(cfg.obs_config, "obs_dict")
-    assert hasattr(cfg.obs_config, "actor_obs")
 
 
-def test_g1_joystick_sac_obs_config_total_dim():
-    """obs_config total dim must match what _compute_obs actually concatenates.
+def test_g1_joystick_ppo_obs_groups_spec_dims():
+    """obs_groups_spec total dim must match what _compute_obs actually produces.
 
     G1JoystickPPO._compute_obs outputs (G1 has 29 DoF):
-        linvel(3) + gyro(3) + gravity(3) + diff(29) + dof_vel(29)
-        + last_actions(29) + command(3) + gait_phase(2) = 101
+        actor: gyro(3) + gravity(3) + diff(29) + dof_vel(29)
+            + last_actions(29) + command(3) + gait_phase(2) = 98
+        privileged: linvel(3)
     """
-    from unilab.envs.locomotion.g1.joystick_sac import G1JoystickSACCfg
+    from unilab.envs.locomotion.g1.joystick import G1JoystickPPO
 
-    cfg = G1JoystickSACCfg()
-    total = sum(cfg.obs_config.obs_dict.values())
-    assert total == 101, (
-        f"obs_config total dim is {total}, expected 101. "
-        "obs_config.obs_dict does not match _compute_obs output."
-    )
+    # obs_groups_spec is a @property; access via descriptor protocol
+    spec = G1JoystickPPO.obs_groups_spec.fget(None)  # type: ignore[union-attr]
+    assert spec is not None
+    assert spec["actor"] == 98
+    assert spec["privileged"] == 3
 
 
 # ---------------------------------------------------------------------------
diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/utils/test_obs_utils.py b/tests/utils/test_obs_utils.py
new file mode 100644
index 000000000..dc2312b9f
--- /dev/null
+++ b/tests/utils/test_obs_utils.py
@@ -0,0 +1,92 @@
+"""Tests for unilab.utils.obs_utils."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from unilab.utils.obs_utils import flatten_obs_dict
+
+# ---------------------------------------------------------------------------
+# flatten_obs_dict — basic behaviour
+# ---------------------------------------------------------------------------
+
+
+class TestFlattenObsDict:
+    """Unit tests for flatten_obs_dict."""
+
+    def test_single_group(self):
+        obs = {"actor": np.ones((4, 8))}
+        flat = flatten_obs_dict(obs)
+        assert flat.shape == (4, 8)
+        np.testing.assert_array_equal(flat, np.ones((4, 8)))
+
+    def test_two_groups_concatenated(self):
+        obs = {
+            "actor": np.ones((2, 5)),
+            "privileged": np.full((2, 3), 2.0),
+        }
+        flat = flatten_obs_dict(obs)
+        assert flat.shape == (2, 8)
+        np.testing.assert_array_equal(flat[:, :5], 1.0)
+        np.testing.assert_array_equal(flat[:, 5:], 2.0)
+
+    def test_insertion_order_preserved(self):
+        """Groups are concatenated in dict insertion order."""
+        a = np.array([[1.0, 2.0]])
+        b = np.array([[3.0]])
+        c = np.array([[4.0, 5.0, 6.0]])
+        obs = {"first": a, "second": b, "third": c}
+        flat = flatten_obs_dict(obs)
+        expected = np.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]])
+        np.testing.assert_array_equal(flat, expected)
+
+    def test_three_groups(self):
+        obs = {
+            "actor": np.zeros((3, 10)),
+            "privileged": np.ones((3, 4)),
+            "extra": np.full((3, 2), 9.0),
+        }
+        flat = flatten_obs_dict(obs)
+        assert flat.shape == (3, 16)
+
+    def test_dtype_preserved_float32(self):
+        obs = {"actor": np.zeros((1, 5), dtype=np.float32)}
+        flat = flatten_obs_dict(obs)
+        assert flat.dtype == np.float32
+
+    def test_dtype_preserved_float64(self):
+        obs = {"actor": np.zeros((1, 5), dtype=np.float64)}
+        flat = flatten_obs_dict(obs)
+        assert flat.dtype == np.float64
+
+    def test_mixed_dtype_upcasts(self):
+        """NumPy concatenation rules apply — float64 wins over float32."""
+        obs = {
+            "actor": np.zeros((2, 3), dtype=np.float32),
+            "privileged": np.zeros((2, 1), dtype=np.float64),
+        }
+        flat = flatten_obs_dict(obs)
+        assert flat.dtype == np.float64
+
+    def test_large_batch(self):
+        n = 4096
+        obs = {"actor": np.random.randn(n, 98), "privileged": np.random.randn(n, 3)}
+        flat = flatten_obs_dict(obs)
+        assert flat.shape == (n, 101)
+        np.testing.assert_array_equal(flat[:, :98], obs["actor"])
+        np.testing.assert_array_equal(flat[:, 98:], obs["privileged"])
+
+    def test_single_env(self):
+        obs = {"actor": np.array([[1.0, 2.0, 3.0]])}
+        flat = flatten_obs_dict(obs)
+        assert flat.shape == (1, 3)
+
+    def test_values_roundtrip(self):
+        """Flatten then slice recovers original groups."""
+        actor = np.random.randn(8, 45)
+        priv = np.random.randn(8, 3)
+        obs = {"actor": actor, "privileged": priv}
+        flat = flatten_obs_dict(obs)
+        np.testing.assert_array_equal(flat[:, :45], actor)
+        np.testing.assert_array_equal(flat[:, 45:], priv)