diff --git a/src/unilab/envs/mdp/commands/velocity_command.py b/src/unilab/envs/mdp/commands/velocity_command.py index 3932d77d3..66e49f701 100644 --- a/src/unilab/envs/mdp/commands/velocity_command.py +++ b/src/unilab/envs/mdp/commands/velocity_command.py @@ -117,7 +117,8 @@ def _validate_cfg(cfg: UniformVelocityCommandCfg) -> None: def command(self) -> np.ndarray: return self.vel_command_b - def _update_metrics(self) -> None: + def _update_metrics(self, env_ids: np.ndarray | None = None) -> None: + del env_ids # Metrics accumulate over all rows on every compute. max_command_steps = self.cfg.resampling_time_range[1] / self._env.step_dt self.metrics["error_vel_xy"] += ( np.linalg.norm( diff --git a/src/unilab/managers/command_manager.py b/src/unilab/managers/command_manager.py index 81451bb2a..4f67ff774 100644 --- a/src/unilab/managers/command_manager.py +++ b/src/unilab/managers/command_manager.py @@ -87,7 +87,8 @@ def compute(self, dt: float | np.ndarray, env_ids: np.ndarray | None = None) -> With env_ids=None (the per-step path) all envs are updated; with env_ids (the reset path) timers and the command update are scoped to those envs. - Metrics are always refreshed. + Metrics are refreshed each call; terms may scope per-row metric work to + env_ids since other rows are unchanged since the per-step update. dt may be a scalar (all envs) or a per-env tensor (auto-reset path, where freshly reset envs get zero to keep their timers full). A tensor @@ -104,7 +105,7 @@ def compute(self, dt: float | np.ndarray, env_ids: np.ndarray | None = None) -> dt_is_finite = np.isfinite(dt).all() if isinstance(dt, np.ndarray) else np.isfinite(dt) if not dt_is_finite: raise ValueError(f"CommandTerm '{self.name}' received non-finite dt.") - self._update_metrics() + self._update_metrics(env_ids) self._validate_metrics() if env_ids is None: self.time_left -= dt @@ -156,8 +157,13 @@ def _resample(self, env_ids: np.ndarray) -> None: self.command_counter[env_ids] += 1 @abc.abstractmethod - def _update_metrics(self) -> None: - """Update the metrics based on the current state.""" + def _update_metrics(self, env_ids: np.ndarray | None = None) -> None: + """Update the metrics based on the current state. + + env_ids is None on the per-step update (all envs) and the reset env ids on + the reset path. Terms may scope per-row metric work to env_ids; rows outside + env_ids are unchanged since the last per-step update and stay valid. + """ raise NotImplementedError @abc.abstractmethod diff --git a/src/unilab/tasks/motion_tracking/common/manager_terms.py b/src/unilab/tasks/motion_tracking/common/manager_terms.py index 9dc3b6b13..2d7788398 100644 --- a/src/unilab/tasks/motion_tracking/common/manager_terms.py +++ b/src/unilab/tasks/motion_tracking/common/manager_terms.py @@ -2,8 +2,10 @@ from __future__ import annotations +import dataclasses import math from dataclasses import dataclass, field +from types import SimpleNamespace from typing import TYPE_CHECKING, Any, Literal, cast import numpy as np @@ -185,6 +187,10 @@ def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv): self._env_error = np.empty(self.num_envs, dtype=dtype) self._reward_term = np.empty(self.num_envs, dtype=dtype) self._robot_cache_step = -1 + # Env ids of the most recent scoped (reset-path) compute; None after a + # per-step compute. Written by `_update_command`, consumed by + # `post_compute` to restrict refresh work to the reset rows. + self._post_compute_env_ids: np.ndarray | None = None self._robot_body_pos_w = np.empty_like(self._body_pos_w) self._robot_body_quat_w = np.empty((self.num_envs, num_bodies, 4), dtype=dtype) self._robot_body_lin_vel_w = np.empty_like(self._body_pos_w) @@ -341,29 +347,55 @@ def reset(self, env_ids: np.ndarray | slice | None) -> dict[str, float]: ) return super().reset(ids) - def _refresh_motion(self) -> None: - self.motion.get_motion_at_frame(self.time_steps, out=self._motion_data) - np.add( - self._motion_data.body_pos_w, - self._env.scene.env_origins[:, None, :], - out=self._body_pos_w, - ) + def _refresh_motion(self, env_ids: np.ndarray | None = None) -> None: + """Refresh motion-reference buffers from the current frame indices. + + With env_ids=None all rows are refreshed in place; with env_ids only + those rows are gathered and scattered (partial-reset path). Rows outside + env_ids keep the values produced by the last per-step refresh, which are + still valid because untouched envs did not advance or resample frames. + """ + if env_ids is None: + self.motion.get_motion_at_frame(self.time_steps, out=self._motion_data) + np.add( + self._motion_data.body_pos_w, + self._env.scene.env_origins[:, None, :], + out=self._body_pos_w, + ) + width = self.motion.num_joints + self._command[:, :width] = self._motion_data.joint_pos + self._command[:, width:] = self._motion_data.joint_vel + return + data = self.motion.get_motion_at_frame(self.time_steps[env_ids]) + for motion_field in dataclasses.fields(data): + value = getattr(data, motion_field.name) + target = getattr(self._motion_data, motion_field.name) + if value is None or target is None: + continue + target[env_ids] = value + self._body_pos_w[env_ids] = data.body_pos_w + self._env.scene.env_origins[env_ids, None, :] width = self.motion.num_joints - self._command[:, :width] = self._motion_data.joint_pos - self._command[:, width:] = self._motion_data.joint_vel + self._command[env_ids, :width] = data.joint_pos + self._command[env_ids, width:] = data.joint_vel - def _refresh_robot_state(self, *, force: bool = False) -> None: + def _refresh_robot_state( + self, *, force: bool = False, env_ids: np.ndarray | None = None + ) -> None: step = self._env.common_step_counter if not force and self._robot_cache_step == step: return + sel: np.ndarray | slice = slice(None) if env_ids is None else env_ids body_index = self._robot_body_ids - self._robot_body_pos_w[:] = self.robot.data.body_link_pos_w[:, body_index] - self._robot_body_quat_w[:] = self.robot.data.body_link_quat_w[:, body_index] - self._robot_body_lin_vel_w[:] = self.robot.data.body_link_lin_vel_w[:, body_index] - self._robot_body_ang_vel_w[:] = self.robot.data.body_link_ang_vel_w[:, body_index] + self._robot_body_pos_w[sel] = self.robot.data.body_link_pos_w[sel][:, body_index] + self._robot_body_quat_w[sel] = self.robot.data.body_link_quat_w[sel][:, body_index] + self._robot_body_lin_vel_w[sel] = self.robot.data.body_link_lin_vel_w[sel][:, body_index] + self._robot_body_ang_vel_w[sel] = self.robot.data.body_link_ang_vel_w[sel][:, body_index] self._robot_cache_step = step - def _refresh_relative_state(self) -> None: + def _refresh_relative_state(self, env_ids: np.ndarray | None = None) -> None: + if env_ids is not None: + self._refresh_relative_state_rows(env_ids) + return update_relative_transforms( self, self._motion_data, @@ -391,39 +423,111 @@ def _refresh_relative_state(self) -> None: self.robot_body_ori_b, ) - def _update_metrics(self) -> None: - self.metrics["error_anchor_pos"][:] = np.linalg.norm( - self.anchor_pos_w - self.robot_anchor_pos_w, axis=-1 + def _refresh_relative_state_rows(self, env_ids: np.ndarray) -> None: + """Row-scoped variant of `_refresh_relative_state` for partial resets. + + Computes the same transforms on the gathered reset rows and scatters the + results back; untouched rows keep their per-step values. + """ + num_rows = len(env_ids) + num_bodies = self._robot_body_pos_w.shape[1] + dtype = self._body_pos_w.dtype + robot_pos_rows = self._robot_body_pos_w[env_ids] + robot_quat_rows = self._robot_body_quat_w[env_ids] + motion_rows = SimpleNamespace( + body_pos_w=self._motion_data.body_pos_w[env_ids], + body_quat_w=self._motion_data.body_quat_w[env_ids], + ) + # `update_relative_transforms` reads/writes these attributes on the env; + # a namespace with row-sized scratch keeps the shared buffers untouched. + scratch = SimpleNamespace( + anchor_body_idx=self.anchor_body_idx, + _delta_pos_w=np.empty((num_rows, 3), dtype=dtype), + _delta_ori_w=np.empty((num_rows, 4), dtype=dtype), + body_quat_relative_w=np.empty((num_rows, num_bodies, 4), dtype=dtype), + body_pos_relative_w=np.empty((num_rows, num_bodies, 3), dtype=dtype), + _body_vec_error=np.empty((num_rows, num_bodies, 3), dtype=dtype), + _env_error=np.empty(num_rows, dtype=dtype), + _reward_term=np.empty(num_rows, dtype=dtype), + ) + update_relative_transforms(scratch, motion_rows, robot_pos_rows, robot_quat_rows) + self.body_pos_relative_w[env_ids] = scratch.body_pos_relative_w + self.body_quat_relative_w[env_ids] = scratch.body_quat_relative_w + + anchor_idx = self.anchor_body_idx + robot_anchor_pos_rows = robot_pos_rows[:, anchor_idx] + robot_anchor_quat_rows = robot_quat_rows[:, anchor_idx] + anchor_pos_rows = self._body_pos_w[env_ids][:, anchor_idx] + anchor_quat_rows = motion_rows.body_quat_w[:, anchor_idx] + motion_anchor_pos_b = np.empty((num_rows, 3), dtype=dtype) + motion_anchor_ori_b = np.empty((num_rows, 6), dtype=dtype) + np_write_relative_anchor_transform_pos_rot6d( + robot_anchor_pos_rows, + robot_anchor_quat_rows, + anchor_pos_rows, + anchor_quat_rows, + motion_anchor_pos_b, + motion_anchor_ori_b, ) - self.metrics["error_anchor_rot"][:] = np.sqrt( - np_quat_error_magnitude_squared_batched(self.anchor_quat_w, self.robot_anchor_quat_w) + self.motion_anchor_pos_b[env_ids] = motion_anchor_pos_b + self.motion_anchor_ori_b[env_ids] = motion_anchor_ori_b + robot_body_pos_b = np.empty((num_rows, num_bodies, 3), dtype=dtype) + write_body_pos_in_anchor_frame( + robot_anchor_pos_rows, + robot_anchor_quat_rows, + robot_pos_rows, + robot_body_pos_b, + body_vec_error=scratch._body_vec_error, + ) + self.robot_body_pos_b[env_ids] = robot_body_pos_b + robot_body_ori_b = np.empty((num_rows, num_bodies, 6), dtype=dtype) + write_body_ori6_in_anchor_frame( + robot_anchor_quat_rows, + robot_quat_rows, + robot_body_ori_b, ) - self.metrics["error_anchor_lin_vel"][:] = np.linalg.norm( - self.anchor_lin_vel_w - self.robot_anchor_lin_vel_w, axis=-1 + self.robot_body_ori_b[env_ids] = robot_body_ori_b + + def _update_metrics(self, env_ids: np.ndarray | None = None) -> None: + # All error metrics are row-wise functions of the motion/robot buffers; + # on the reset path only the reset rows are recomputed since other rows + # are unchanged since the per-step update. + sel: np.ndarray | slice = slice(None) if env_ids is None else env_ids + self.metrics["error_anchor_pos"][sel] = np.linalg.norm( + self.anchor_pos_w[sel] - self.robot_anchor_pos_w[sel], axis=-1 ) - self.metrics["error_anchor_ang_vel"][:] = np.linalg.norm( - self.anchor_ang_vel_w - self.robot_anchor_ang_vel_w, axis=-1 + self.metrics["error_anchor_rot"][sel] = np.sqrt( + np_quat_error_magnitude_squared_batched( + self.anchor_quat_w[sel], self.robot_anchor_quat_w[sel] + ) ) - self.metrics["error_body_pos"][:] = np.linalg.norm( - self.body_pos_relative_w - self.robot_body_pos_w, axis=-1 + self.metrics["error_anchor_lin_vel"][sel] = np.linalg.norm( + self.anchor_lin_vel_w[sel] - self.robot_anchor_lin_vel_w[sel], axis=-1 + ) + self.metrics["error_anchor_ang_vel"][sel] = np.linalg.norm( + self.anchor_ang_vel_w[sel] - self.robot_anchor_ang_vel_w[sel], axis=-1 + ) + self.metrics["error_body_pos"][sel] = np.linalg.norm( + self.body_pos_relative_w[sel] - self.robot_body_pos_w[sel], axis=-1 ).mean(axis=-1) - self.metrics["error_body_rot"][:] = np.sqrt( + self.metrics["error_body_rot"][sel] = np.sqrt( np_quat_error_magnitude_squared_batched( - self.body_quat_relative_w, self.robot_body_quat_w + self.body_quat_relative_w[sel], self.robot_body_quat_w[sel] ) ).mean(axis=-1) - self.metrics["error_body_lin_vel"][:] = np.linalg.norm( - self.body_lin_vel_w - self.robot_body_lin_vel_w, axis=-1 + self.metrics["error_body_lin_vel"][sel] = np.linalg.norm( + self.body_lin_vel_w[sel] - self.robot_body_lin_vel_w[sel], axis=-1 ).mean(axis=-1) - self.metrics["error_body_ang_vel"][:] = np.linalg.norm( - self.body_ang_vel_w - self.robot_body_ang_vel_w, axis=-1 + self.metrics["error_body_ang_vel"][sel] = np.linalg.norm( + self.body_ang_vel_w[sel] - self.robot_body_ang_vel_w[sel], axis=-1 ).mean(axis=-1) - self.metrics["error_joint_pos"][:] = np.linalg.norm( - self.joint_pos - self.robot_joint_pos, axis=-1 + self.metrics["error_joint_pos"][sel] = np.linalg.norm( + self.joint_pos[sel] - self.robot_joint_pos[sel], axis=-1 ) - self.metrics["error_joint_vel"][:] = np.linalg.norm( - self.joint_vel - self.robot_joint_vel, axis=-1 + self.metrics["error_joint_vel"][sel] = np.linalg.norm( + self.joint_vel[sel] - self.robot_joint_vel[sel], axis=-1 ) + # Sampler statistics are global scalars, so every row tracks them. self.metrics["sampling_entropy"].fill(self.sampler.sampling_entropy) self.metrics["sampling_top1_prob"].fill(self.sampler.sampling_top1_prob) self.metrics["sampling_top1_bin"].fill(self.sampler.sampling_top1_bin) @@ -459,8 +563,9 @@ def _resample_command(self, env_ids: np.ndarray) -> None: self.robot.write_root_state_to_sim(root_state, env_ids=env_ids) def _update_command(self, env_ids: np.ndarray | None) -> None: + self._post_compute_env_ids = env_ids if env_ids is not None: - self._refresh_motion() + self._refresh_motion(env_ids) return self.sampler.update_failure_stats(self._env.termination_manager.terminated) active_ids = np.flatnonzero(~self._env.reset_buf).astype(np.int32, copy=False) @@ -470,8 +575,11 @@ def _update_command(self, env_ids: np.ndarray | None) -> None: self._refresh_motion() def post_compute(self) -> None: - self._refresh_robot_state(force=True) - self._refresh_relative_state() + # On the reset path only the reset rows changed (via the committed + # set_state writes and the motion resample), so refresh just those rows. + env_ids = self._post_compute_env_ids + self._refresh_robot_state(force=True, env_ids=env_ids) + self._refresh_relative_state(env_ids) @dataclass(kw_only=True) diff --git a/src/unilab/tasks/motion_tracking/g1/manager_terms.py b/src/unilab/tasks/motion_tracking/g1/manager_terms.py index 09ddc07f9..32e34cbb8 100644 --- a/src/unilab/tasks/motion_tracking/g1/manager_terms.py +++ b/src/unilab/tasks/motion_tracking/g1/manager_terms.py @@ -80,12 +80,15 @@ def object_quat_w(self) -> np.ndarray: def object_state_b(self) -> np.ndarray: return self._object_obs_b - def _refresh_motion(self) -> None: - super()._refresh_motion() + def _refresh_motion(self, env_ids: np.ndarray | None = None) -> None: + super()._refresh_motion(env_ids) value = self.box_motion.object_pos_w if value is None: raise RuntimeError("Box motion object position was not materialized") - np.add(value, self._env.scene.env_origins, out=self._object_pos_w) + if env_ids is None: + np.add(value, self._env.scene.env_origins, out=self._object_pos_w) + else: + self._object_pos_w[env_ids] = value[env_ids] + self._env.scene.env_origins[env_ids] def _resample_command(self, env_ids: np.ndarray) -> None: super()._resample_command(env_ids) @@ -111,23 +114,43 @@ def _resample_command(self, env_ids: np.ndarray) -> None: ) self.object.write_root_state_to_sim(object_state, env_ids=env_ids) - def _refresh_object_state(self) -> None: + def _refresh_object_state(self, env_ids: np.ndarray | None = None) -> None: + if env_ids is None: + np_write_relative_anchor_transform_pos_rot6d( + self.robot_anchor_pos_w, + self.robot_anchor_quat_w, + self.object.data.root_link_pos_w, + self.object.data.root_link_quat_w, + self._object_obs_b[:, :3], + self._object_obs_b[:, 3:9], + ) + self._object_obs_b[:, 9:12] = np_quat_apply_inverse( + self.robot_anchor_quat_w, + self.object.data.root_link_lin_vel_w, + ) + return + num_rows = len(env_ids) + dtype = self._object_obs_b.dtype + pos_b = np.empty((num_rows, 3), dtype=dtype) + rot_b = np.empty((num_rows, 6), dtype=dtype) np_write_relative_anchor_transform_pos_rot6d( - self.robot_anchor_pos_w, - self.robot_anchor_quat_w, - self.object.data.root_link_pos_w, - self.object.data.root_link_quat_w, - self._object_obs_b[:, :3], - self._object_obs_b[:, 3:9], + self.robot_anchor_pos_w[env_ids], + self.robot_anchor_quat_w[env_ids], + self.object.data.root_link_pos_w[env_ids], + self.object.data.root_link_quat_w[env_ids], + pos_b, + rot_b, ) - self._object_obs_b[:, 9:12] = np_quat_apply_inverse( - self.robot_anchor_quat_w, - self.object.data.root_link_lin_vel_w, + self._object_obs_b[env_ids, :3] = pos_b + self._object_obs_b[env_ids, 3:9] = rot_b + self._object_obs_b[env_ids, 9:12] = np_quat_apply_inverse( + self.robot_anchor_quat_w[env_ids], + self.object.data.root_link_lin_vel_w[env_ids], ) def post_compute(self) -> None: super().post_compute() - self._refresh_object_state() + self._refresh_object_state(self._post_compute_env_ids) def _box_command(env: ManagerBasedRlEnv, command_name: str) -> BoxMotionCommand: diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 886f90dbf..c3fb9299c 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -226,7 +226,7 @@ def __init__(self, cfg: _CommandCfg, env) -> None: def command(self) -> np.ndarray: return self._command - def _update_metrics(self) -> None: + def _update_metrics(self, env_ids: np.ndarray | None = None) -> None: return None def _resample_command(self, env_ids: np.ndarray) -> None: diff --git a/tests/envs/test_motion_command_partial_reset.py b/tests/envs/test_motion_command_partial_reset.py new file mode 100644 index 000000000..8f979f9ad --- /dev/null +++ b/tests/envs/test_motion_command_partial_reset.py @@ -0,0 +1,188 @@ +"""Row-scoped partial-reset parity tests for MotionCommand (issue #1261). + +On the partial-reset path the command manager recomputes only the reset rows; +untouched rows reuse the values produced by the per-step compute of the same +control step. These tests pin that contract: + +- untouched rows keep their command/motion/relative-state/metric values + bit-identically across a partial reset (except the sampler-stat metrics, + which track global sampler scalars); +- reset rows match a full recompute of the same post-reset state, which is + what the pre-#1261 full-batch implementation produced. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Callable + +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 +from unilab.tasks.motion_tracking.common.manager_terms import MotionCommand +from unilab.tasks.motion_tracking.g1.manager_terms import BoxMotionCommand + +_ROOT = Path(__file__).parents[2] + +_CASES = ( + ("ppo", "g1_flip_tracking", "G1FlipTracking"), + ("ppo", "g1_box_tracking", "G1BoxTracking"), +) + +_SAMPLER_STAT_METRICS = ("sampling_entropy", "sampling_top1_prob", "sampling_top1_bin") + +_ACCESSORS: dict[str, Callable[[MotionCommand], np.ndarray]] = { + "command": lambda term: term.command, + "time_steps": lambda term: term.time_steps, + "body_pos_w": lambda term: term.body_pos_w, + "body_pos_relative_w": lambda term: term.body_pos_relative_w, + "body_quat_relative_w": lambda term: term.body_quat_relative_w, + "motion_anchor_pos_b": lambda term: term.motion_anchor_pos_b, + "motion_anchor_ori_b": lambda term: term.motion_anchor_ori_b, + "robot_body_pos_b": lambda term: term.robot_body_pos_b, + "robot_body_ori_b": lambda term: term.robot_body_ori_b, + "joint_default_bias": lambda term: term.joint_default_bias, + "robot_body_pos_w": lambda term: term._robot_body_pos_w, + "robot_body_quat_w": lambda term: term._robot_body_quat_w, + "robot_body_lin_vel_w": lambda term: term._robot_body_lin_vel_w, + "robot_body_ang_vel_w": lambda term: term._robot_body_ang_vel_w, +} + +_MOTION_DATA_FIELDS = ( + "joint_pos", + "joint_vel", + "body_pos_w", + "body_quat_w", + "body_lin_vel_w", + "body_ang_vel_w", + "object_pos_w", + "object_quat_w", + "object_lin_vel_w", + "object_ang_vel_w", +) + + +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 _buffers(term: MotionCommand) -> dict[str, np.ndarray]: + buffers = {name: accessor(term) for name, accessor in _ACCESSORS.items()} + for field_name in _MOTION_DATA_FIELDS: + value = getattr(term._motion_data, field_name, None) + if value is not None: + buffers[f"motion_data.{field_name}"] = value + if isinstance(term, BoxMotionCommand): + buffers["object_pos_w"] = term.object_pos_w + buffers["object_state_b"] = term.object_state_b + return buffers + + +def _current(term: MotionCommand, name: str) -> np.ndarray: + if name in _ACCESSORS: + return _ACCESSORS[name](term) + if name.startswith("motion_data."): + return getattr(term._motion_data, name.split(".", 1)[1]) + if name == "object_pos_w": + return term.object_pos_w # type: ignore[attr-defined] + if name == "object_state_b": + return term.object_state_b # type: ignore[attr-defined] + raise KeyError(name) + + +@pytest.mark.parametrize(("config_root", "task", "identity"), _CASES) +def test_motion_command_partial_reset_row_parity( + config_root: str, task: str, identity: str +) -> 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(config_root, task, "mujoco", identity, num_envs) + try: + env.init_state() + term = env.command_manager.get_term("motion") + assert isinstance(term, MotionCommand) + action_dim = term.motion.num_joints + rng = np.random.default_rng(7) + for _ in range(5): + env.step((0.1 * rng.standard_normal((num_envs, action_dim))).astype(np.float32)) + + reset_ids = np.array([0, 2], dtype=np.int32) + keep_ids = np.array([1, 3], dtype=np.int32) + before = {name: value.copy() for name, value in _buffers(term).items()} + metrics_before = {name: value.copy() for name, value in term.metrics.items()} + reset_obs, _ = env.reset(env_ids=reset_ids) + after = _buffers(term) + + # Untouched rows keep their per-step values bit-identically, except the + # sampler-stat metrics which track global sampler scalars. + for name, value in before.items(): + np.testing.assert_array_equal( + after[name][keep_ids], + value[keep_ids], + err_msg=f"untouched rows changed for {name}", + ) + for name, value in metrics_before.items(): + if name in _SAMPLER_STAT_METRICS: + expected = np.full( + num_envs, getattr(term.sampler, name), dtype=term.metrics[name].dtype + ) + np.testing.assert_array_equal(term.metrics[name], expected) + else: + np.testing.assert_array_equal( + term.metrics[name][keep_ids], + value[keep_ids], + err_msg=f"untouched rows changed for metric {name}", + ) + + # Reset observations are finite and scattered into the env state. + assert env.state is not None + for group, values in reset_obs.items(): + assert values.shape[0] == len(reset_ids) + assert np.isfinite(values).all() + np.testing.assert_array_equal(env.state.obs[group][reset_ids], values) + + # Reference: a full recompute of the post-reset state (the pre-#1261 + # behavior) must agree with the row-scoped results on every row. + reference = {name: value.copy() for name, value in _buffers(term).items()} + term._refresh_motion() + term._refresh_robot_state(force=True) + term._refresh_relative_state() + if isinstance(term, BoxMotionCommand): + term._refresh_object_state() + for name, value in reference.items(): + # The default bias is intentionally resampled from RNG on reset rows. + if name == "joint_default_bias": + np.testing.assert_array_equal(_current(term, name), value) + continue + np.testing.assert_allclose( + _current(term, name), + value, + rtol=1e-6, + atol=1e-7, + err_msg=f"row-scoped refresh disagrees with full refresh for {name}", + ) + finally: + env.close() diff --git a/tests/managers/test_event_command_metrics_recorder.py b/tests/managers/test_event_command_metrics_recorder.py index 9efc46a5a..e2bfb87d7 100644 --- a/tests/managers/test_event_command_metrics_recorder.py +++ b/tests/managers/test_event_command_metrics_recorder.py @@ -126,7 +126,7 @@ def __init__(self, cfg: DummyCommandCfg, env: FakeEnv): def command(self) -> np.ndarray: return self._command - def _update_metrics(self) -> None: + def _update_metrics(self, env_ids: np.ndarray | None = None) -> None: self.metrics["error"] += 1.0 def _resample_command(self, env_ids: np.ndarray) -> None: