perf: scope partial-reset observation rebuild to reset rows

src/unilab/envs/manager_based_rl_env.py

@@ -629,9 +629,11 @@ def reset(
         self.command_manager.post_compute()
         command_compute_ms = (time.perf_counter() - t0) * 1000.0
         t0 = time.perf_counter()
+        # Row-scoped reset rebuild (issue #1259 R2): the observation manager
+        # returns only the reset rows, so no full-batch slice is needed here.
         manager_obs = self.observation_manager.compute(update_history=True, env_ids=ids)
-        mapped_obs = self._map_observations(manager_obs)
-        reset_obs = {name: values[ids].copy() for name, values in mapped_obs.items()}
+        mapped_obs = self._map_observations(manager_obs, num_rows=len(ids))
+        reset_obs = {name: values.copy() for name, values in mapped_obs.items()}
         obs_build_ms = (time.perf_counter() - t0) * 1000.0
 
         if self._state is not None:
@@ -711,6 +713,7 @@ def _normalize_reset_ids(
     def _map_observations(
         self,
         manager_obs: dict[str, np.ndarray | dict[str, np.ndarray]],
+        num_rows: int | None = None,
     ) -> dict[str, np.ndarray]:
         mapping = {"obs": self._cfg.policy_observation_group}
         if self._cfg.critic_observation_group is not None:
@@ -723,7 +726,10 @@ def _map_observations(
                     f"ManagerBasedRlEnv observation group '{group_name}' returned "
                     f"{type(value).__name__}, expected np.ndarray"
                 )
-            expected = (self.num_envs, self._mapped_obs_dims[output_name])
+            expected = (
+                self.num_envs if num_rows is None else num_rows,
+                self._mapped_obs_dims[output_name],
+            )
             if value.shape != expected:
                 raise ValueError(
                     f"ManagerBasedRlEnv observation group '{group_name}' returned shape "

src/unilab/managers/observation_manager.py

@@ -278,13 +278,21 @@ def reset(self, env_ids: np.ndarray | slice | None = None) -> dict[str, float]:
                 mod.reset(env_ids=env_ids)
         return {}
 
-    def _check_and_handle_nans(self, tensor: np.ndarray, context: str, policy: str) -> np.ndarray:
+    def _check_and_handle_nans(
+        self,
+        tensor: np.ndarray,
+        context: str,
+        policy: str,
+        env_ids: np.ndarray | None = None,
+    ) -> np.ndarray:
         """Check for NaN/Inf and handle according to policy.
 
         Args:
-          tensor: Observation tensor to check.
+          tensor: Observation tensor to check. On the reset path this holds only
+            the reset rows; pass env_ids so diagnostics report real env indices.
           context: Context string for error/warning messages (e.g., "actor/base_lin_vel").
           policy: NaN handling policy ("disabled", "warn", "sanitize", "error").
+          env_ids: Optional mapping from tensor rows to env indices (reset path).
 
         Returns:
           The tensor, potentially sanitized depending on policy.
@@ -301,10 +309,17 @@ def _check_and_handle_nans(self, tensor: np.ndarray, context: str, policy: str)
         if not (has_nan or has_inf):
             return tensor
 
+        def _row_env_ids(mask: np.ndarray) -> list[int]:
+            rows = np.flatnonzero(np.asarray(mask, dtype=bool))
+            if env_ids is not None:
+                rows = np.asarray(env_ids)[rows]
+            result: list[int] = rows.tolist()
+            return result
+
         if policy == "error":
             invalid = ~np.isfinite(tensor)
-            nan_mask = invalid.reshape(self.num_envs, -1).any(axis=1)
-            nan_env_ids = np.flatnonzero(nan_mask).tolist()
+            nan_mask = np.asarray(invalid.reshape(tensor.shape[0], -1).any(axis=1))
+            nan_env_ids = _row_env_ids(nan_mask)
             invalid_kind = (
                 "NaN"
                 if has_nan and not has_inf
@@ -319,8 +334,8 @@ def _check_and_handle_nans(self, tensor: np.ndarray, context: str, policy: str)
 
         if policy == "warn":
             invalid = ~np.isfinite(tensor)
-            nan_mask = invalid.reshape(self.num_envs, -1).any(axis=1)
-            nan_env_ids = np.flatnonzero(nan_mask).tolist()
+            nan_mask = np.asarray(invalid.reshape(tensor.shape[0], -1).any(axis=1))
+            nan_env_ids = _row_env_ids(nan_mask)
             print(
                 f"[ObservationManager] NaN/Inf in '{context}' "
                 f"(envs: {nan_env_ids[:5]}). Sanitizing to 0."
@@ -337,10 +352,14 @@ def compute(
         """Compute observations for all groups.
 
         With env_ids=None (the per-step path), history and delay buffers advance
-        for all envs. With env_ids (the reset path), only the reset envs' buffers
-        receive their post-reset frame (a backfill); other envs' buffers, delay
-        schedules, and lag draws are untouched, so a partial reset does not
-        advance their observation timelines.
+        for all envs and the returned arrays cover the full batch. With env_ids
+        (the reset path), only the reset envs' buffers receive their post-reset
+        frame (a backfill); other envs' buffers, delay schedules, and lag draws
+        are untouched, so a partial reset does not advance their observation
+        timelines. The returned arrays then hold only the reset rows, in env_ids
+        order, and the observation cache is left invalidated (the next per-step
+        compute refreshes it); noise is still drawn with full-batch shapes so
+        the shared RNG stream matches the full-batch implementation exactly.
         """
         if env_ids is not None and not update_history:
             raise ValueError("env_ids is only meaningful with update_history=True.")
@@ -354,7 +373,8 @@ def compute(
         self._last_compute_timing_ms = {}
         for group_name in self._group_obs_term_names:
             obs_buffer[group_name] = self.compute_group(group_name, update_history, env_ids)
-        self._obs_buffer = obs_buffer
+        if env_ids is None:
+            self._obs_buffer = obs_buffer
         return obs_buffer
 
     def compute_group(
@@ -381,6 +401,13 @@ def compute_group(
         }
         term_timing: dict[str, float] = {}
         obs_terms = zip(group_term_names, self._group_obs_term_cfgs[group_name], strict=False)
+        # Reset path (issue #1259 R2): when no term in this group uses delay or
+        # history buffers, everything downstream of the term call is row
+        # independent, so only the reset rows are processed. Term calls and
+        # noise stay full-batch: term funcs are contracted to return
+        # (num_envs, ...) and full-shape noise draws keep the shared RNG stream
+        # and per-row noise values identical to the full-batch path.
+        row_scoped = env_ids is not None and not self._group_obs_temporal[group_name]
         for term_name, term_cfg in obs_terms:
             term_t0 = time.perf_counter()
             term_getter0 = self._env._mba_getter_total_ms()
@@ -395,13 +422,17 @@ def compute_group(
                     f"ObservationManager term '{group_name}/{term_name}' returned shape "
                     f"{obs.shape}, expected (num_envs, ...) with num_envs={self.num_envs}."
                 )
-            obs = obs.copy()
+            if not row_scoped:
+                obs = obs.copy()
             t0 = time.perf_counter()
             if isinstance(term_cfg.noise, noise_cfg.NoiseCfg):
                 obs = term_cfg.noise.apply(obs, rng=self._env.rng)
             elif isinstance(term_cfg.noise, noise_cfg.NoiseModelCfg):
                 obs = self._group_obs_class_instances[group_name][term_name](obs)
             phase_ms["noise"] += time.perf_counter() - t0
+            if row_scoped:
+                # Fresh row copy; safe for the in-place clip/scale below.
+                obs = obs[env_ids]
             t0 = time.perf_counter()
             if term_cfg.clip:
                 np.clip(obs, term_cfg.clip[0], term_cfg.clip[1], out=obs)
@@ -415,7 +446,10 @@ def compute_group(
             t0 = time.perf_counter()
             if group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled":
                 obs = self._check_and_handle_nans(
-                    obs, context=f"{group_name}/{term_name}", policy=group_cfg.nan_policy
+                    obs,
+                    context=f"{group_name}/{term_name}",
+                    policy=group_cfg.nan_policy,
+                    env_ids=env_ids if row_scoped else None,
                 )
             phase_ms["nan_check"] += time.perf_counter() - t0
 
@@ -464,6 +498,7 @@ def compute_group(
                         group_obs[term_name],
                         context=f"{group_name}/{term_name}",
                         policy=group_cfg.nan_policy,
+                        env_ids=env_ids if row_scoped else None,
                     )
         phase_ms["nan_check"] += time.perf_counter() - t0
 
@@ -477,13 +512,25 @@ def compute_group(
             t0 = time.perf_counter()
             if not group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled":
                 result = self._check_and_handle_nans(
-                    result, context=group_name, policy=group_cfg.nan_policy
+                    result,
+                    context=group_name,
+                    policy=group_cfg.nan_policy,
+                    env_ids=env_ids if row_scoped else None,
                 )
             phase_ms["nan_check"] += time.perf_counter() - t0
         else:
             phase_ms["concat"] += time.perf_counter() - t0
             result = group_obs
 
+        if env_ids is not None and not row_scoped:
+            # Groups with delay/history terms ran the full-batch pipeline above
+            # (buffer readout stays full-batch); slice the reset rows to match
+            # the reset-path return contract.
+            if isinstance(result, dict):
+                result = {name: values[env_ids] for name, values in result.items()}
+            else:
+                result = result[env_ids]
+
         timing = self._last_compute_timing_ms
         for phase_name, elapsed_s in phase_ms.items():
             key = f"mba_obs_{phase_name}_ms"
@@ -501,6 +548,9 @@ def _prepare_terms(self) -> None:
         self._group_obs_class_instances: dict[str, dict[str, noise_model.NoiseModel]] = {}
         self._group_obs_term_delay_buffer: dict[str, dict[str, DelayBuffer]] = dict()
         self._group_obs_term_history_buffer: dict[str, dict[str, CircularBuffer]] = dict()
+        # Whether any term in the group uses delay/history buffers. Groups
+        # without temporal terms can be row-scoped on the reset path.
+        self._group_obs_temporal: dict[str, bool] = dict()
 
         for group_name, group_cfg in self.cfg.items():
             if group_cfg is None:
@@ -630,3 +680,7 @@ def _prepare_terms(self) -> None:
 
             self._group_obs_term_delay_buffer[group_name] = group_entry_delay_buffer
             self._group_obs_term_history_buffer[group_name] = group_entry_history_buffer
+            self._group_obs_temporal[group_name] = any(
+                term_cfg.delay_max_lag > 0 or term_cfg.history_length > 0
+                for term_cfg in self._group_obs_term_cfgs[group_name]
+            )

tests/envs/test_observation_partial_reset.py

@@ -0,0 +1,84 @@
+"""Row-scoped partial-reset observation rebuild tests (issue #1259 R2).
+
+On the partial-reset path the observation manager returns only the reset rows
+instead of a full batch that the env then slices. These tests pin the env-level
+contract:
+
+- reset observations arrive row-shaped, are finite, and are scattered into the
+  env state at exactly the reset rows;
+- untouched rows of ``state.obs`` keep their per-step values bit-identically;
+- a subsequent step recomputes full-batch observations normally.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+from hydra import compose, initialize_config_dir
+from hydra.core.global_hydra import GlobalHydra
+
+from unilab.base import registry
+from unilab.base.config_adapter import BackendAdapter
+from unilab.base.config_materialization import apply_cfg_overrides
+from unilab.envs import ManagerBasedRlEnvCfg
+
+_ROOT = Path(__file__).parents[2]
+
+
+def _make_env(config_root: str, task: str, backend: str, identity: str, num_envs: int):
+    registry.ensure_registries()
+    GlobalHydra.instance().clear()
+    with initialize_config_dir(config_dir=str(_ROOT / "conf" / config_root), version_base="1.3"):
+        owner = compose("config", overrides=[f"task={task}/{backend}"])
+    cfg = registry.materialize_env_config(identity)
+    assert isinstance(cfg, ManagerBasedRlEnvCfg)
+    override = BackendAdapter(
+        owner, root_dir=_ROOT, algo_name=config_root
+    ).build_task_env_cfg_override()
+    apply_cfg_overrides(cfg, override)
+    return registry.make(
+        identity, num_envs=num_envs, sim_backend=backend, env_cfg_override=override
+    )
+
+
+def test_observation_partial_reset_row_contract() -> None:
+    pytest.importorskip("mujoco")
+    try:
+        from mujoco_uni.batch_env import BatchEnvPool as _  # noqa: F401
+    except Exception:
+        pytest.skip("mujoco_uni.batch_env not available")
+
+    num_envs = 4
+    env = _make_env("sac", "g1_motion_tracking", "mujoco", "G1MotionTrackingSAC", num_envs)
+    try:
+        env.init_state()
+        action_dim = 29
+        rng = np.random.default_rng(7)
+        for _ in range(5):
+            env.step((0.1 * rng.standard_normal((num_envs, action_dim))).astype(np.float32))
+
+        assert env.state is not None
+        obs_before = {name: values.copy() for name, values in env.state.obs.items()}
+        reset_ids = np.array([0, 2], dtype=np.int32)
+        keep_ids = np.array([1, 3], dtype=np.int32)
+        reset_obs, _ = env.reset(env_ids=reset_ids)
+
+        for group, values in reset_obs.items():
+            assert values.shape == (len(reset_ids), obs_before[group].shape[1])
+            assert np.isfinite(values).all()
+            np.testing.assert_array_equal(env.state.obs[group][reset_ids], values)
+            np.testing.assert_array_equal(
+                env.state.obs[group][keep_ids],
+                obs_before[group][keep_ids],
+                err_msg=f"untouched rows changed for group {group}",
+            )
+
+        # The next per-step compute rebuilds full-batch observations normally.
+        env.step((0.1 * rng.standard_normal((num_envs, action_dim))).astype(np.float32))
+        for group, values in env.state.obs.items():
+            assert values.shape == obs_before[group].shape
+            assert np.isfinite(values).all()
+    finally:
+        env.close()