From 11bb0ac336d4b65427c505f6eea354f614ee4659 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Sun, 23 Aug 2026 22:30:47 +0800 Subject: [PATCH] perf: scope partial-reset observation rebuild to reset rows Issue #1269 (roadmap #1259 R2). On the partial-reset path the ObservationManager computed the full batch through the whole noise/clip/scale/NaN/delay/history/concat pipeline and the env then sliced out the reset rows, discarding the rest. compute(env_ids=...) now returns only the reset rows: groups without delay/history terms run the row-independent pipeline stages on the reset rows only, while term calls and noise draws stay full-batch so the (num_envs, ...) term contract and the shared RNG stream are bit-identical to the full-batch path. Groups with temporal buffers keep the full pipeline and only slice the final output. --- src/unilab/envs/manager_based_rl_env.py | 12 +- src/unilab/managers/observation_manager.py | 82 +++++++-- tests/envs/test_observation_partial_reset.py | 84 +++++++++ .../test_observation_partial_reset.py | 163 ++++++++++++++++++ 4 files changed, 324 insertions(+), 17 deletions(-) create mode 100644 tests/envs/test_observation_partial_reset.py create mode 100644 tests/managers/test_observation_partial_reset.py diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 2890957d4..9c82b40db 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/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 " diff --git a/src/unilab/managers/observation_manager.py b/src/unilab/managers/observation_manager.py index a21877b0c..9a649a800 100644 --- a/src/unilab/managers/observation_manager.py +++ b/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] + ) diff --git a/tests/envs/test_observation_partial_reset.py b/tests/envs/test_observation_partial_reset.py new file mode 100644 index 000000000..d5d1c0f06 --- /dev/null +++ b/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() diff --git a/tests/managers/test_observation_partial_reset.py b/tests/managers/test_observation_partial_reset.py new file mode 100644 index 000000000..046581b01 --- /dev/null +++ b/tests/managers/test_observation_partial_reset.py @@ -0,0 +1,163 @@ +"""Row-scoped partial-reset parity tests for ObservationManager (issue #1259 R2). + +On the partial-reset path (compute(update_history=True, env_ids=...)) the +manager returns only the reset rows and processes them row-scoped whenever the +group has no delay/history terms. These tests pin that contract: + +- reset rows are bit-identical to a full-batch compute sliced to those rows, + and the shared RNG stream is consumed identically (noise stays full-batch); +- groups with delay/history terms fall back to the full-batch pipeline and + only slice the final output; untouched rows' buffers are not advanced; +- NaN diagnostics on the row-scoped path report real env indices. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from unilab.managers import ObservationGroupCfg, ObservationManager, ObservationTermCfg +from unilab.managers._noise import UniformNoiseCfg + +from .conftest import FakeEnv + + +def _noisy_cfg() -> dict[str, ObservationGroupCfg]: + return { + "policy": ObservationGroupCfg( + terms={ + "state": ObservationTermCfg( + func=lambda env: env.obs, + noise=UniformNoiseCfg(n_min=-0.1, n_max=0.1), + clip=(-10.0, 10.0), + scale=2.0, + ), + "bias": ObservationTermCfg( + func=lambda env: np.ones((env.num_envs, 1), dtype=np.float32) + ), + }, + enable_corruption=True, + ), + } + + +def test_partial_reset_rows_match_full_compute_bitwise() -> None: + env = FakeEnv(seed=11) + manager = ObservationManager(_noisy_cfg(), env) + manager.compute(update_history=True) # populate caches like a step would + + ids = np.array([0, 2], dtype=np.int32) + rng_state = env.rng.bit_generator.state + rows = manager.compute(update_history=True, env_ids=ids) + rng_after_rows = env.rng.bit_generator.state + + env.rng.bit_generator.state = rng_state + full = manager.compute(update_history=True) + rng_after_full = env.rng.bit_generator.state + + assert rows["policy"].shape == (len(ids), full["policy"].shape[1]) + np.testing.assert_array_equal(rows["policy"], full["policy"][ids]) + assert rng_after_rows == rng_after_full + + +def test_partial_reset_does_not_populate_obs_cache() -> None: + env = FakeEnv(seed=3) + manager = ObservationManager(_noisy_cfg(), env) + manager.compute(update_history=True) + assert manager._obs_buffer is not None + + manager.reset(np.array([1], dtype=np.int32)) + assert manager._obs_buffer is None + manager.compute(update_history=True, env_ids=np.array([1], dtype=np.int32)) + # The reset path leaves the cache invalidated; the next per-step compute + # refreshes it with a full-batch entry. + assert manager._obs_buffer is None + manager.compute(update_history=True) + assert manager._obs_buffer is not None + assert manager._obs_buffer["policy"].shape[0] == env.num_envs + + +def test_partial_reset_temporal_group_falls_back_and_preserves_rows() -> None: + env = FakeEnv(seed=5) + cfg = { + "policy": ObservationGroupCfg( + terms={ + "state": ObservationTermCfg(func=lambda env: env.obs, history_length=3), + "delayed": ObservationTermCfg( + func=lambda env: env.obs, delay_min_lag=1, delay_max_lag=1 + ), + }, + ), + } + manager = ObservationManager(cfg, env) + for step in range(4): + env.obs = env.obs + 1 + manager.compute(update_history=True) + + ids = np.array([1], dtype=np.int32) + keep_ids = np.array([0, 2, 3], dtype=np.int32) + history_before = manager._group_obs_term_history_buffer["policy"]["state"].buffer.copy() + delay_before = manager._group_obs_term_delay_buffer["policy"]["delayed"].peek().copy() + + rng_state = env.rng.bit_generator.state + rows = manager.compute(update_history=True, env_ids=ids) + + env.rng.bit_generator.state = rng_state + full = manager.compute(update_history=True, env_ids=ids) + np.testing.assert_array_equal(rows["policy"], full["policy"]) + assert rows["policy"].shape[0] == len(ids) + + # Untouched rows keep their history and delayed values bit-identically; + # the reset row is backfilled with its post-reset frame in every slot. + history_after = manager._group_obs_term_history_buffer["policy"]["state"].buffer + np.testing.assert_array_equal(history_after[keep_ids], history_before[keep_ids]) + reset_slots = history_after[ids[0]] + np.testing.assert_array_equal(reset_slots, np.broadcast_to(reset_slots[-1], reset_slots.shape)) + np.testing.assert_array_equal( + manager._group_obs_term_delay_buffer["policy"]["delayed"].peek()[keep_ids], + delay_before[keep_ids], + ) + + # Reference: the full-batch pipeline sliced to the reset rows agrees. + env.rng.bit_generator.state = rng_state + reference = manager.compute(update_history=True)["policy"][ids] + env.rng.bit_generator.state = rng_state + np.testing.assert_array_equal(rows["policy"], reference) + + +@pytest.mark.parametrize("bad", [np.nan, np.inf]) +def test_partial_reset_nan_error_reports_env_ids(bad: float) -> None: + def invalid(env: FakeEnv) -> np.ndarray: + result = env.obs.copy() + result[2, 0] = bad + return result + + env = FakeEnv(seed=7) + manager = ObservationManager( + {"policy": ObservationGroupCfg(terms={"bad": ObservationTermCfg(func=invalid)})}, + env, + ) + with pytest.raises(ValueError, match=r"for environments: \[2\]"): + manager.compute(update_history=True, env_ids=np.array([0, 2], dtype=np.int32)) + + +def test_partial_reset_nan_on_untouched_row_is_not_rechecked() -> None: + def invalid(env: FakeEnv) -> np.ndarray: + result = env.obs.copy() + result[1, 0] = np.nan + return result + + env = FakeEnv(seed=7) + manager = ObservationManager( + { + "policy": ObservationGroupCfg( + terms={"bad": ObservationTermCfg(func=invalid)}, nan_policy="error" + ) + }, + env, + ) + # Row-scoped NaN checks cover only the reset rows; untouched rows were + # already checked by the per-step compute of their control step. + rows = manager.compute(update_history=True, env_ids=np.array([0, 2], dtype=np.int32)) + assert rows["policy"].shape == (2, env.obs.shape[1]) + assert np.isfinite(rows["policy"]).all()