From 53a782393b8cbe65e71a4b21ac3a9e854b7b6beb Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Sun, 5 Jul 2026 16:22:44 +0800 Subject: [PATCH 1/6] perf: optimize reset done body pose lookup --- .../benchmark_offpolicy_collector_active.py | 25 +++++- src/unilab/base/backend/base.py | 5 ++ src/unilab/base/backend/motrix/backend.py | 45 +++++++++++ src/unilab/base/np_env.py | 8 ++ src/unilab/dr/manager.py | 1 + .../envs/motion_tracking/g1/tracking.py | 8 +- tests/dr/test_manager.py | 7 ++ tests/envs/test_env_configs.py | 79 +++++++++++++++++++ 8 files changed, 174 insertions(+), 4 deletions(-) diff --git a/benchmark/benchmark_offpolicy_collector_active.py b/benchmark/benchmark_offpolicy_collector_active.py index 588d11b95..6f23d0d4c 100644 --- a/benchmark/benchmark_offpolicy_collector_active.py +++ b/benchmark/benchmark_offpolicy_collector_active.py @@ -91,6 +91,14 @@ "dr_reset_plan_ms", "dr_reset_payload_filter_ms", "dr_reset_set_state_ms", + "dr_reset_set_state_qpos_convert_ms", + "dr_reset_set_state_slice_ms", + "dr_reset_set_state_data_write_ms", + "dr_reset_set_state_ctrl_ms", + "dr_reset_set_state_forward_kinematic_ms", + "dr_reset_set_state_refresh_cache_ms", + "dr_reset_set_state_invalidate_cache_ms", + "dr_reset_set_state_backend_internal_gap_ms", "dr_reset_build_observation_ms", "dr_reset_internal_gap_ms", "dr_reset_observation_getters_ms", @@ -122,6 +130,17 @@ ("dr_reset_plan_ms", "dr_reset_plan_ms"), ("dr_reset_payload_filter_ms", "dr_reset_payload_filter_ms"), ("dr_reset_set_state_ms", "dr_reset_set_state_ms"), + ("dr_reset_set_state_qpos_convert_ms", "dr_reset_set_state_qpos_convert_ms"), + ("dr_reset_set_state_slice_ms", "dr_reset_set_state_slice_ms"), + ("dr_reset_set_state_data_write_ms", "dr_reset_set_state_data_write_ms"), + ("dr_reset_set_state_ctrl_ms", "dr_reset_set_state_ctrl_ms"), + ("dr_reset_set_state_forward_kinematic_ms", "dr_reset_set_state_forward_kinematic_ms"), + ("dr_reset_set_state_refresh_cache_ms", "dr_reset_set_state_refresh_cache_ms"), + ("dr_reset_set_state_invalidate_cache_ms", "dr_reset_set_state_invalidate_cache_ms"), + ( + "dr_reset_set_state_backend_internal_gap_ms", + "dr_reset_set_state_backend_internal_gap_ms", + ), ("dr_reset_build_observation_ms", "dr_reset_build_observation_ms"), ("dr_reset_internal_gap_ms", "dr_reset_internal_gap_ms"), ("dr_reset_observation_getters_ms", "dr_reset_observation_getters_ms"), @@ -1395,6 +1414,10 @@ def _print_result(result: CollectorResult) -> None: "reset_done_info_scatter_ms", "dr_reset_plan_ms", "dr_reset_set_state_ms", + "dr_reset_set_state_slice_ms", + "dr_reset_set_state_data_write_ms", + "dr_reset_set_state_forward_kinematic_ms", + "dr_reset_set_state_refresh_cache_ms", "dr_reset_build_observation_ms", "dr_reset_obs_get_motion_ms", "dr_reset_observation_getters_ms", @@ -1563,7 +1586,7 @@ def main() -> int: print("\nReset done timing (subparts of NpEnv reset_done_ms):") print(_format_reset_done_timing_table(results)) print( - "\nDR reset timing (subparts of reset call; reset obs getters currently read full batch):" + "\nDR reset timing (subparts of reset call; includes reset observation getters):" ) print(_format_dr_reset_timing_table(results)) else: diff --git a/src/unilab/base/backend/base.py b/src/unilab/base/backend/base.py index 96cc56f53..2ef20e403 100644 --- a/src/unilab/base/backend/base.py +++ b/src/unilab/base/backend/base.py @@ -281,6 +281,11 @@ def set_state( randomization: Optional backend randomization payload. """ + @property + def last_set_state_timing_ms(self) -> dict[str, float]: + """Return backend-internal timing from the most recent ``set_state`` call.""" + return {} + @abc.abstractmethod def get_dr_capabilities(self) -> DomainRandomizationCapabilities: """Return supported domain-randomization capabilities for this backend.""" diff --git a/src/unilab/base/backend/motrix/backend.py b/src/unilab/base/backend/motrix/backend.py index daa1f5bd8..8e3a062a5 100644 --- a/src/unilab/base/backend/motrix/backend.py +++ b/src/unilab/base/backend/motrix/backend.py @@ -278,6 +278,7 @@ def __init__( self._render_tracking_camera: MotrixTrackingCamera | None = None self.backend_type = "motrix" self._link_velocity_cache: np.ndarray | None = None + self._last_set_state_timing_ms: dict[str, float] = {} # Pre-cache link objects to avoid repeated get_link() lookups. self._link_cache: dict[int, "mtx.Link"] = {} @@ -317,6 +318,10 @@ def model(self): def data(self): return self._data + @property + def last_set_state_timing_ms(self) -> dict[str, float]: + return dict(self._last_set_state_timing_ms) + # ------------------------------------------------------------------ # # Model properties # # ------------------------------------------------------------------ # @@ -580,21 +585,31 @@ def set_state( qvel: np.ndarray, randomization: ResetRandomizationPayload | None = None, ) -> None: + total_t0 = time.perf_counter() + self._last_set_state_timing_ms = {} + + t0 = time.perf_counter() qpos_motrix = self._mujoco_qpos_to_motrix(qpos) + qpos_convert_ms = (time.perf_counter() - t0) * 1000.0 + t0 = time.perf_counter() # Create mask for batch operation mask = np.zeros(self._num_envs, dtype=bool) mask[env_indices] = True data_slice = self._data[mask] + slice_ms = (time.perf_counter() - t0) * 1000.0 # Batch set state + t0 = time.perf_counter() data_slice.reset(self._model) self._clear_applied_body_forces(env_indices) self._apply_init_geom_size_overrides(data_slice, env_indices) self._apply_reset_randomization(data_slice, env_indices, randomization) data_slice.set_dof_vel(qvel) data_slice.set_dof_pos(qpos_motrix, self._model) + data_write_ms = (time.perf_counter() - t0) * 1000.0 + t0 = time.perf_counter() if self._supports_position_actuator_gains and len(self._joint_dof_pos_indices) == int( self.num_actuators ): @@ -606,10 +621,40 @@ def set_state( else: ctrl = np.zeros((len(env_indices), self.num_actuators), dtype=self._np_dtype) data_slice.actuator_ctrls = np.ascontiguousarray(ctrl) + ctrl_ms = (time.perf_counter() - t0) * 1000.0 + t0 = time.perf_counter() self._model.forward_kinematic(data_slice) + forward_kinematic_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() self._refresh_link_pose_cache(env_indices) + refresh_cache_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() self._invalidate_link_velocity_cache() + invalidate_cache_ms = (time.perf_counter() - t0) * 1000.0 + + total_ms = (time.perf_counter() - total_t0) * 1000.0 + measured_ms = ( + qpos_convert_ms + + slice_ms + + data_write_ms + + ctrl_ms + + forward_kinematic_ms + + refresh_cache_ms + + invalidate_cache_ms + ) + self._last_set_state_timing_ms = { + "dr_reset_set_state_qpos_convert_ms": qpos_convert_ms, + "dr_reset_set_state_slice_ms": slice_ms, + "dr_reset_set_state_data_write_ms": data_write_ms, + "dr_reset_set_state_ctrl_ms": ctrl_ms, + "dr_reset_set_state_forward_kinematic_ms": forward_kinematic_ms, + "dr_reset_set_state_refresh_cache_ms": refresh_cache_ms, + "dr_reset_set_state_invalidate_cache_ms": invalidate_cache_ms, + "dr_reset_set_state_backend_internal_gap_ms": total_ms - measured_ms, + } def get_dr_capabilities(self) -> DomainRandomizationCapabilities: supported_reset_terms = { diff --git a/src/unilab/base/np_env.py b/src/unilab/base/np_env.py index 65408d330..82ec1aa5a 100644 --- a/src/unilab/base/np_env.py +++ b/src/unilab/base/np_env.py @@ -34,6 +34,14 @@ "dr_reset_plan_ms", "dr_reset_payload_filter_ms", "dr_reset_set_state_ms", + "dr_reset_set_state_qpos_convert_ms", + "dr_reset_set_state_slice_ms", + "dr_reset_set_state_data_write_ms", + "dr_reset_set_state_ctrl_ms", + "dr_reset_set_state_forward_kinematic_ms", + "dr_reset_set_state_refresh_cache_ms", + "dr_reset_set_state_invalidate_cache_ms", + "dr_reset_set_state_backend_internal_gap_ms", "dr_reset_build_observation_ms", "dr_reset_internal_gap_ms", "dr_reset_observation_getters_ms", diff --git a/src/unilab/dr/manager.py b/src/unilab/dr/manager.py index 6beddb281..f6469d6ca 100644 --- a/src/unilab/dr/manager.py +++ b/src/unilab/dr/manager.py @@ -71,6 +71,7 @@ def reset(self, env_ids: np.ndarray) -> tuple[dict[str, np.ndarray], dict]: "dr_reset_build_observation_ms": build_observation_ms, "dr_reset_internal_gap_ms": total_ms - measured_ms, } + timing.update(self._env._backend.last_set_state_timing_ms) provider_timing = getattr(self._provider, "last_reset_observation_timing_ms", {}) if isinstance(provider_timing, dict): timing.update(provider_timing) diff --git a/src/unilab/envs/motion_tracking/g1/tracking.py b/src/unilab/envs/motion_tracking/g1/tracking.py index 5e86a1bba..eac620b60 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking.py +++ b/src/unilab/envs/motion_tracking/g1/tracking.py @@ -446,7 +446,9 @@ def build_reset_observation( dof_vel_ms = (time.perf_counter() - t0) * 1000.0 t0 = time.perf_counter() - all_pos_w, all_quat_w = env._get_body_pose_w() + robot_body_pos_w, robot_body_quat_w = env._backend.get_body_pose_w_rows( + env_ids, env.body_ids + ) body_pose_ms = (time.perf_counter() - t0) * 1000.0 obs_info = dict(info_updates) @@ -465,8 +467,8 @@ def build_reset_observation( gyro, dof_pos, dof_vel, - all_pos_w[env_ids], - all_quat_w[env_ids], + robot_body_pos_w, + robot_body_quat_w, ), ) compute_obs_ms = (time.perf_counter() - t0) * 1000.0 diff --git a/tests/dr/test_manager.py b/tests/dr/test_manager.py index c8150b30a..1e3640fcd 100644 --- a/tests/dr/test_manager.py +++ b/tests/dr/test_manager.py @@ -6,6 +6,7 @@ from typing import Any import numpy as np +import pytest from unilab.dr import ( DomainRandomizationCapabilities, @@ -122,6 +123,10 @@ class _FakeBackend: def __post_init__(self) -> None: self.last_randomization: ResetRandomizationPayload | None = None + @property + def last_set_state_timing_ms(self) -> dict[str, float]: + return {"dr_reset_set_state_forward_kinematic_ms": 1.25} + def get_dr_capabilities(self) -> DomainRandomizationCapabilities: return self.capabilities @@ -179,6 +184,8 @@ def test_manager_skips_unsupported_reset_terms_with_warning(caplog): "motrix backend does not support reset randomization terms: kp; skipping them." in caplog.text ) + timing = manager.last_reset_timing_ms + assert timing["dr_reset_set_state_forward_kinematic_ms"] == pytest.approx(1.25) def test_manager_keeps_supported_reset_terms_without_warning(caplog): diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index a570094f0..be0eaa9a8 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -477,6 +477,85 @@ def get_body_pose_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray] np.testing.assert_array_equal(env._backend.calls[0], np.array([1, 3], dtype=np.int32)) +def test_g1_motion_tracking_reset_observation_uses_sparse_body_pose_rows(): + from unilab.envs.motion_tracking.g1.motion_loader import MotionData + from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingDomainRandomizationProvider + + class FakeBackend: + def __init__(self) -> None: + self.row_calls: list[tuple[np.ndarray, np.ndarray]] = [] + + def get_body_pose_w_rows( + self, env_ids: np.ndarray, body_ids: np.ndarray + ) -> tuple[np.ndarray, np.ndarray]: + self.row_calls.append((env_ids.copy(), body_ids.copy())) + rows = len(env_ids) + bodies = len(body_ids) + return ( + np.full((rows, bodies, 3), 2.0, dtype=np.float32), + np.full((rows, bodies, 4), 3.0, dtype=np.float32), + ) + + def get_body_pose_w(self, body_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + raise AssertionError("reset observation should use sparse body pose rows") + + class FakeMotionLoader: + def get_motion_at_frame(self, frames: np.ndarray) -> MotionData: + rows = len(frames) + return MotionData( + joint_pos=np.zeros((rows, 2), dtype=np.float32), + joint_vel=np.zeros((rows, 2), dtype=np.float32), + body_pos_w=np.zeros((rows, 2, 3), dtype=np.float32), + body_quat_w=np.zeros((rows, 2, 4), dtype=np.float32), + body_lin_vel_w=np.zeros((rows, 2, 3), dtype=np.float32), + body_ang_vel_w=np.zeros((rows, 2, 3), dtype=np.float32), + ) + + class FakeMotionSampler: + current_frames = np.array([10, 11, 12, 13], dtype=np.int32) + + env = SimpleNamespace( + _backend=FakeBackend(), + body_ids=np.array([1, 3], dtype=np.int32), + motion_loader=FakeMotionLoader(), + motion_sampler=FakeMotionSampler(), + get_local_linvel=lambda: np.zeros((4, 3), dtype=np.float32), + get_gyro=lambda: np.zeros((4, 3), dtype=np.float32), + get_dof_pos=lambda: np.zeros((4, 2), dtype=np.float32), + get_dof_vel=lambda: np.zeros((4, 2), dtype=np.float32), + ) + + captured: dict[str, np.ndarray] = {} + + def compute_obs( + obs_info, + motion_data, + linvel, + gyro, + dof_pos, + dof_vel, + robot_body_pos_w, + robot_body_quat_w, + ): + del obs_info, motion_data, linvel, gyro, dof_pos, dof_vel + captured["robot_body_pos_w"] = robot_body_pos_w + captured["robot_body_quat_w"] = robot_body_quat_w + return {"obs": np.zeros((2, 1), dtype=np.float32)} + + env._compute_obs = compute_obs + provider = G1MotionTrackingDomainRandomizationProvider() + env_ids = np.array([1, 3], dtype=np.int32) + + obs = provider.build_reset_observation(env, env_ids, {}) + + assert obs["obs"].shape == (2, 1) + assert len(env._backend.row_calls) == 1 + np.testing.assert_array_equal(env._backend.row_calls[0][0], env_ids) + np.testing.assert_array_equal(env._backend.row_calls[0][1], env.body_ids) + assert captured["robot_body_pos_w"].shape == (2, 2, 3) + assert captured["robot_body_quat_w"].shape == (2, 2, 4) + + def _compute_g1_motion_tracking_obs_stub(env_cls: type): from unilab.envs.motion_tracking.g1.motion_loader import MotionData From 66546c4b4bdbe37def83af956b8b47c7d1033f77 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Sun, 5 Jul 2026 16:32:32 +0800 Subject: [PATCH 2/6] perf: reuse motrix reset data slice --- src/unilab/base/backend/motrix/backend.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/src/unilab/base/backend/motrix/backend.py b/src/unilab/base/backend/motrix/backend.py index 8e3a062a5..f44648443 100644 --- a/src/unilab/base/backend/motrix/backend.py +++ b/src/unilab/base/backend/motrix/backend.py @@ -628,7 +628,7 @@ def set_state( forward_kinematic_ms = (time.perf_counter() - t0) * 1000.0 t0 = time.perf_counter() - self._refresh_link_pose_cache(env_indices) + self._refresh_link_pose_cache(env_indices, data_slice=data_slice) refresh_cache_ms = (time.perf_counter() - t0) * 1000.0 t0 = time.perf_counter() @@ -1030,13 +1030,17 @@ def _motrix_qpos_to_mujoco(self, qpos: np.ndarray) -> np.ndarray: qpos_mujoco[..., quat_indices] = qpos[..., quat_indices[[3, 0, 1, 2]]] return qpos_mujoco - def _refresh_link_pose_cache(self, env_indices: np.ndarray | None = None) -> None: + def _refresh_link_pose_cache( + self, env_indices: np.ndarray | None = None, data_slice: Any | None = None + ) -> None: if env_indices is None: self._link_poses = self._model.get_link_poses(self._data) else: - mask = np.zeros(self._num_envs, dtype=bool) - mask[env_indices] = True - self._link_poses[env_indices] = self._model.get_link_poses(self._data[mask]) + if data_slice is None: + mask = np.zeros(self._num_envs, dtype=bool) + mask[env_indices] = True + data_slice = self._data[mask] + self._link_poses[env_indices] = self._model.get_link_poses(data_slice) def _refresh_link_velocity_cache(self, env_indices: np.ndarray | None = None) -> None: if env_indices is None: From b210ea6f245e68b3f25f3c29ea8d040c953736ec Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Sun, 5 Jul 2026 17:10:59 +0800 Subject: [PATCH 3/6] bench: add reset done attribution harness --- benchmark/benchmark_reset_done_attribution.py | 384 ++++++++++++++++++ .../test_reset_done_attribution_benchmark.py | 40 ++ 2 files changed, 424 insertions(+) create mode 100644 benchmark/benchmark_reset_done_attribution.py create mode 100644 tests/benchmark/test_reset_done_attribution_benchmark.py diff --git a/benchmark/benchmark_reset_done_attribution.py b/benchmark/benchmark_reset_done_attribution.py new file mode 100644 index 000000000..eb4d8e794 --- /dev/null +++ b/benchmark/benchmark_reset_done_attribution.py @@ -0,0 +1,384 @@ +#!/usr/bin/env python3 +"""Attribution harness for reset_done changes from issue #673. + +This benchmark avoids full collector throughput attribution. It constructs one +env per case and alternates old/new implementations in the same process, using +the same env_ids and reset payloads. + +Usage: + uv run benchmark/benchmark_reset_done_attribution.py + uv run benchmark/benchmark_reset_done_attribution.py --num-envs 8192 --reset-count 256 + uv run benchmark/benchmark_reset_done_attribution.py --out-json /tmp/reset_done_attr.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +import types +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Callable + +import numpy as np +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra + +ROOT_DIR = Path(__file__).resolve().parents[1] +if str(ROOT_DIR) not in sys.path: + sys.path.append(str(ROOT_DIR)) + +from benchmark.core.device_info import get_device_info_dict, get_device_info_line + + +@dataclass(frozen=True) +class TimingStats: + mean_ms: float + median_ms: float + min_ms: float + max_ms: float + samples_ms: list[float] + + +@dataclass(frozen=True) +class BodyPoseAttributionResult: + backend: str + num_envs: int + reset_count: int + body_count: int + repeats: int + old_build_observation_ms: TimingStats + new_build_observation_ms: TimingStats + old_body_pose_ms: TimingStats + new_body_pose_ms: TimingStats + build_observation_speedup: float + body_pose_speedup: float + + +@dataclass(frozen=True) +class SetStateAttributionResult: + backend: str + num_envs: int + reset_count: int + repeats: int + old_set_state_ms: TimingStats + new_set_state_ms: TimingStats + old_refresh_cache_ms: TimingStats + new_refresh_cache_ms: TimingStats + set_state_speedup: float + refresh_cache_speedup: float + + +def _stats(samples_ms: list[float]) -> TimingStats: + if not samples_ms: + raise ValueError("no samples") + arr = np.asarray(samples_ms, dtype=np.float64) + return TimingStats( + mean_ms=float(arr.mean()), + median_ms=float(np.median(arr)), + min_ms=float(arr.min()), + max_ms=float(arr.max()), + samples_ms=[float(v) for v in samples_ms], + ) + + +def _speedup(old: TimingStats, new: TimingStats) -> float: + return old.mean_ms / new.mean_ms if new.mean_ms > 0.0 else float("inf") + + +def _sample_env_ids(num_envs: int, reset_count: int, seed: int) -> np.ndarray: + if reset_count <= 0: + raise ValueError("reset-count must be > 0") + if reset_count > num_envs: + raise ValueError("reset-count must be <= num-envs") + rng = np.random.default_rng(seed) + return np.sort(rng.choice(num_envs, size=reset_count, replace=False).astype(np.int32)) + + +def _compose_env(algo: str, task: str, backend: str, num_envs: int): + from unilab.training import BackendAdapter, create_env, ensure_registries + + ensure_registries() + GlobalHydra.instance().clear() + with initialize_config_dir(version_base=None, config_dir=str(ROOT_DIR / "conf" / "offpolicy")): + cfg = compose( + "config", + overrides=[ + f"algo={algo}", + f"task={algo}/{task}/{backend}", + f"algo.num_envs={num_envs}", + ], + ) + env_cfg_override = BackendAdapter( + cfg, + root_dir=ROOT_DIR, + algo_name=algo, + ).build_task_env_cfg_override() + env = create_env(cfg, num_envs=num_envs, env_cfg_override=env_cfg_override) + env.init_state() + return env + + +def _body_pose_old_rows(backend: Any, rows: np.ndarray, body_ids: np.ndarray): + pos, quat = backend.get_body_pose_w(body_ids) + return pos[rows], quat[rows] + + +def _run_body_pose_once(provider: Any, env: Any, env_ids: np.ndarray) -> tuple[float, float]: + t0 = time.perf_counter() + provider.build_reset_observation(env, env_ids, {}) + total_ms = (time.perf_counter() - t0) * 1000.0 + timing = provider.last_reset_observation_timing_ms + return total_ms, float(timing["dr_reset_obs_get_body_pose_ms"]) + + +def benchmark_body_pose_attribution( + *, + backend: str, + num_envs: int, + reset_count: int, + warmup_repeats: int, + measure_repeats: int, + seed: int, +) -> BodyPoseAttributionResult: + env = _compose_env("sac", "g1_motion_tracking", backend, num_envs) + try: + provider = env._dr_manager._provider + env_ids = _sample_env_ids(num_envs, reset_count, seed) + original_rows = env._backend.get_body_pose_w_rows + + def old_rows(self, rows: np.ndarray, body_ids: np.ndarray): + return _body_pose_old_rows(self, rows, body_ids) + + old_pos, old_quat = _body_pose_old_rows(env._backend, env_ids, env.body_ids) + new_pos, new_quat = original_rows(env_ids, env.body_ids) + np.testing.assert_allclose(old_pos, new_pos, rtol=0, atol=0) + np.testing.assert_allclose(old_quat, new_quat, rtol=0, atol=0) + + old_total: list[float] = [] + new_total: list[float] = [] + old_body: list[float] = [] + new_body: list[float] = [] + + for repeat_idx in range(warmup_repeats + measure_repeats): + record = repeat_idx >= warmup_repeats + + env._backend.get_body_pose_w_rows = types.MethodType(old_rows, env._backend) + total_ms, body_ms = _run_body_pose_once(provider, env, env_ids) + if record: + old_total.append(total_ms) + old_body.append(body_ms) + + env._backend.get_body_pose_w_rows = original_rows + total_ms, body_ms = _run_body_pose_once(provider, env, env_ids) + if record: + new_total.append(total_ms) + new_body.append(body_ms) + + env._backend.get_body_pose_w_rows = original_rows + old_total_stats = _stats(old_total) + new_total_stats = _stats(new_total) + old_body_stats = _stats(old_body) + new_body_stats = _stats(new_body) + return BodyPoseAttributionResult( + backend=backend, + num_envs=num_envs, + reset_count=reset_count, + body_count=int(len(env.body_ids)), + repeats=measure_repeats, + old_build_observation_ms=old_total_stats, + new_build_observation_ms=new_total_stats, + old_body_pose_ms=old_body_stats, + new_body_pose_ms=new_body_stats, + build_observation_speedup=_speedup(old_total_stats, new_total_stats), + body_pose_speedup=_speedup(old_body_stats, new_body_stats), + ) + finally: + env.close() + + +def _old_refresh_link_pose_cache(self, env_indices: np.ndarray | None = None, data_slice=None): + del data_slice + if env_indices is None: + self._link_poses = self._model.get_link_poses(self._data) + else: + mask = np.zeros(self._num_envs, dtype=bool) + mask[env_indices] = True + self._link_poses[env_indices] = self._model.get_link_poses(self._data[mask]) + + +def _run_set_state_once(env: Any, env_ids: np.ndarray, qpos: np.ndarray, qvel: np.ndarray) -> dict: + env._backend.set_state(env_ids, qpos, qvel) + return env._backend.last_set_state_timing_ms + + +def benchmark_set_state_attribution( + *, + num_envs: int, + reset_count: int, + warmup_repeats: int, + measure_repeats: int, + seed: int, +) -> SetStateAttributionResult: + env = _compose_env("flashsac", "g1_walk_flat", "motrix", num_envs) + try: + env_ids = _sample_env_ids(num_envs, reset_count, seed) + plan = env._dr_manager._provider.build_reset_plan(env, env_ids) + qpos = np.asarray(plan.qpos) + qvel = np.asarray(plan.qvel) + original_refresh = env._backend._refresh_link_pose_cache + + old_set_state: list[float] = [] + new_set_state: list[float] = [] + old_refresh: list[float] = [] + new_refresh: list[float] = [] + + for repeat_idx in range(warmup_repeats + measure_repeats): + record = repeat_idx >= warmup_repeats + + env._backend._refresh_link_pose_cache = types.MethodType( + _old_refresh_link_pose_cache, env._backend + ) + timing = _run_set_state_once(env, env_ids, qpos, qvel) + if record: + old_set_state.append(sum(v for k, v in timing.items() if k.startswith("dr_reset_set_state_"))) + old_refresh.append(float(timing["dr_reset_set_state_refresh_cache_ms"])) + + env._backend._refresh_link_pose_cache = original_refresh + timing = _run_set_state_once(env, env_ids, qpos, qvel) + if record: + new_set_state.append(sum(v for k, v in timing.items() if k.startswith("dr_reset_set_state_"))) + new_refresh.append(float(timing["dr_reset_set_state_refresh_cache_ms"])) + + env._backend._refresh_link_pose_cache = original_refresh + old_set_state_stats = _stats(old_set_state) + new_set_state_stats = _stats(new_set_state) + old_refresh_stats = _stats(old_refresh) + new_refresh_stats = _stats(new_refresh) + return SetStateAttributionResult( + backend="motrix", + num_envs=num_envs, + reset_count=reset_count, + repeats=measure_repeats, + old_set_state_ms=old_set_state_stats, + new_set_state_ms=new_set_state_stats, + old_refresh_cache_ms=old_refresh_stats, + new_refresh_cache_ms=new_refresh_stats, + set_state_speedup=_speedup(old_set_state_stats, new_set_state_stats), + refresh_cache_speedup=_speedup(old_refresh_stats, new_refresh_stats), + ) + finally: + env.close() + + +def _print_stats(label: str, old: TimingStats, new: TimingStats, speedup: float) -> None: + print( + f" {label}: old={old.mean_ms:.6f} ms new={new.mean_ms:.6f} ms " + f"speedup={speedup:.2f}x" + ) + + +def _run_safely(label: str, fn: Callable[[], Any]) -> Any | None: + try: + return fn() + except Exception as exc: + print(f"{label}: ERROR {type(exc).__name__}: {exc}") + return None + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--num-envs", type=int, default=8192) + parser.add_argument("--reset-count", type=int, default=256) + parser.add_argument("--warmup-repeats", type=int, default=5) + parser.add_argument("--measure-repeats", type=int, default=30) + parser.add_argument("--seed", type=int, default=673) + parser.add_argument("--body-pose-backends", default="motrix,mujoco") + parser.add_argument("--skip-set-state", action="store_true") + parser.add_argument("--out-json", type=Path, default=None) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + print(f"Device: {get_device_info_line()}") + + body_pose_results: list[BodyPoseAttributionResult] = [] + for backend in [part.strip() for part in args.body_pose_backends.split(",") if part.strip()]: + result = _run_safely( + f"body_pose/{backend}", + lambda backend=backend: benchmark_body_pose_attribution( + backend=backend, + num_envs=args.num_envs, + reset_count=args.reset_count, + warmup_repeats=args.warmup_repeats, + measure_repeats=args.measure_repeats, + seed=args.seed, + ), + ) + if result is None: + continue + body_pose_results.append(result) + print(f"body_pose/{backend}:") + _print_stats( + "build_observation_ms", + result.old_build_observation_ms, + result.new_build_observation_ms, + result.build_observation_speedup, + ) + _print_stats( + "body_pose_ms", + result.old_body_pose_ms, + result.new_body_pose_ms, + result.body_pose_speedup, + ) + + set_state_result = None + if not args.skip_set_state: + set_state_result = _run_safely( + "set_state/motrix", + lambda: benchmark_set_state_attribution( + num_envs=args.num_envs, + reset_count=args.reset_count, + warmup_repeats=args.warmup_repeats, + measure_repeats=args.measure_repeats, + seed=args.seed, + ), + ) + if set_state_result is not None: + print("set_state/motrix:") + _print_stats( + "set_state_ms", + set_state_result.old_set_state_ms, + set_state_result.new_set_state_ms, + set_state_result.set_state_speedup, + ) + _print_stats( + "refresh_cache_ms", + set_state_result.old_refresh_cache_ms, + set_state_result.new_refresh_cache_ms, + set_state_result.refresh_cache_speedup, + ) + + if args.out_json is not None: + payload = { + "device": get_device_info_dict(), + "num_envs": args.num_envs, + "reset_count": args.reset_count, + "warmup_repeats": args.warmup_repeats, + "measure_repeats": args.measure_repeats, + "seed": args.seed, + "body_pose": [asdict(result) for result in body_pose_results], + "set_state": asdict(set_state_result) if set_state_result is not None else None, + } + args.out_json.parent.mkdir(parents=True, exist_ok=True) + args.out_json.write_text(json.dumps(payload, indent=2), encoding="utf-8") + print(f"Wrote JSON: {args.out_json}") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/benchmark/test_reset_done_attribution_benchmark.py b/tests/benchmark/test_reset_done_attribution_benchmark.py new file mode 100644 index 000000000..6fcc6e630 --- /dev/null +++ b/tests/benchmark/test_reset_done_attribution_benchmark.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +import pytest +from benchmark import benchmark_reset_done_attribution as bench + + +def test_stats_reports_distribution() -> None: + stats = bench._stats([3.0, 1.0, 2.0]) + + assert stats.mean_ms == pytest.approx(2.0) + assert stats.median_ms == pytest.approx(2.0) + assert stats.min_ms == pytest.approx(1.0) + assert stats.max_ms == pytest.approx(3.0) + assert stats.samples_ms == [3.0, 1.0, 2.0] + + +def test_sample_env_ids_is_sorted_unique_and_seeded() -> None: + env_ids = bench._sample_env_ids(num_envs=16, reset_count=5, seed=673) + + assert env_ids.tolist() == sorted(env_ids.tolist()) + assert len(set(env_ids.tolist())) == 5 + assert env_ids.tolist() == bench._sample_env_ids(16, 5, 673).tolist() + + +def test_sample_env_ids_rejects_invalid_count() -> None: + with pytest.raises(ValueError, match="reset-count must be > 0"): + bench._sample_env_ids(num_envs=16, reset_count=0, seed=673) + + with pytest.raises(ValueError, match="reset-count must be <= num-envs"): + bench._sample_env_ids(num_envs=16, reset_count=17, seed=673) + + +def test_parse_args_defaults_to_issue_scale() -> None: + args = bench.parse_args([]) + + assert args.num_envs == 8192 + assert args.reset_count == 256 + assert args.measure_repeats == 30 + assert args.body_pose_backends == "motrix,mujoco" + assert not args.skip_set_state From 8f59a4ca2da841cc0224e384104b7fd5fd215c3d Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Sun, 5 Jul 2026 17:32:49 +0800 Subject: [PATCH 4/6] perf: refine motrix reset state attribution --- .../benchmark_offpolicy_collector_active.py | 12 ++++ benchmark/benchmark_reset_done_attribution.py | 41 ++++++++++++- src/unilab/base/backend/motrix/backend.py | 58 +++++++++++++++++-- src/unilab/base/np_env.py | 6 ++ 4 files changed, 110 insertions(+), 7 deletions(-) diff --git a/benchmark/benchmark_offpolicy_collector_active.py b/benchmark/benchmark_offpolicy_collector_active.py index 6f23d0d4c..fdaaad57e 100644 --- a/benchmark/benchmark_offpolicy_collector_active.py +++ b/benchmark/benchmark_offpolicy_collector_active.py @@ -94,6 +94,12 @@ "dr_reset_set_state_qpos_convert_ms", "dr_reset_set_state_slice_ms", "dr_reset_set_state_data_write_ms", + "dr_reset_set_state_data_reset_ms", + "dr_reset_set_state_clear_forces_ms", + "dr_reset_set_state_geom_overrides_ms", + "dr_reset_set_state_randomization_ms", + "dr_reset_set_state_set_dof_vel_ms", + "dr_reset_set_state_set_dof_pos_ms", "dr_reset_set_state_ctrl_ms", "dr_reset_set_state_forward_kinematic_ms", "dr_reset_set_state_refresh_cache_ms", @@ -133,6 +139,12 @@ ("dr_reset_set_state_qpos_convert_ms", "dr_reset_set_state_qpos_convert_ms"), ("dr_reset_set_state_slice_ms", "dr_reset_set_state_slice_ms"), ("dr_reset_set_state_data_write_ms", "dr_reset_set_state_data_write_ms"), + ("dr_reset_set_state_data_reset_ms", "dr_reset_set_state_data_reset_ms"), + ("dr_reset_set_state_clear_forces_ms", "dr_reset_set_state_clear_forces_ms"), + ("dr_reset_set_state_geom_overrides_ms", "dr_reset_set_state_geom_overrides_ms"), + ("dr_reset_set_state_randomization_ms", "dr_reset_set_state_randomization_ms"), + ("dr_reset_set_state_set_dof_vel_ms", "dr_reset_set_state_set_dof_vel_ms"), + ("dr_reset_set_state_set_dof_pos_ms", "dr_reset_set_state_set_dof_pos_ms"), ("dr_reset_set_state_ctrl_ms", "dr_reset_set_state_ctrl_ms"), ("dr_reset_set_state_forward_kinematic_ms", "dr_reset_set_state_forward_kinematic_ms"), ("dr_reset_set_state_refresh_cache_ms", "dr_reset_set_state_refresh_cache_ms"), diff --git a/benchmark/benchmark_reset_done_attribution.py b/benchmark/benchmark_reset_done_attribution.py index eb4d8e794..ead1f9209 100644 --- a/benchmark/benchmark_reset_done_attribution.py +++ b/benchmark/benchmark_reset_done_attribution.py @@ -67,6 +67,8 @@ class SetStateAttributionResult: new_set_state_ms: TimingStats old_refresh_cache_ms: TimingStats new_refresh_cache_ms: TimingStats + old_component_ms: dict[str, TimingStats] + new_component_ms: dict[str, TimingStats] set_state_speedup: float refresh_cache_speedup: float @@ -209,9 +211,30 @@ def _old_refresh_link_pose_cache(self, env_indices: np.ndarray | None = None, da self._link_poses[env_indices] = self._model.get_link_poses(self._data[mask]) +SET_STATE_COMPONENT_KEYS = ( + "dr_reset_set_state_qpos_convert_ms", + "dr_reset_set_state_slice_ms", + "dr_reset_set_state_data_write_ms", + "dr_reset_set_state_data_reset_ms", + "dr_reset_set_state_clear_forces_ms", + "dr_reset_set_state_geom_overrides_ms", + "dr_reset_set_state_randomization_ms", + "dr_reset_set_state_set_dof_vel_ms", + "dr_reset_set_state_set_dof_pos_ms", + "dr_reset_set_state_ctrl_ms", + "dr_reset_set_state_forward_kinematic_ms", + "dr_reset_set_state_refresh_cache_ms", + "dr_reset_set_state_invalidate_cache_ms", + "dr_reset_set_state_backend_internal_gap_ms", +) + + def _run_set_state_once(env: Any, env_ids: np.ndarray, qpos: np.ndarray, qvel: np.ndarray) -> dict: + t0 = time.perf_counter() env._backend.set_state(env_ids, qpos, qvel) - return env._backend.last_set_state_timing_ms + wall_ms = (time.perf_counter() - t0) * 1000.0 + timing = env._backend.last_set_state_timing_ms + return {"dr_reset_set_state_wall_ms": wall_ms, **timing} def benchmark_set_state_attribution( @@ -234,6 +257,8 @@ def benchmark_set_state_attribution( new_set_state: list[float] = [] old_refresh: list[float] = [] new_refresh: list[float] = [] + old_components: dict[str, list[float]] = {key: [] for key in SET_STATE_COMPONENT_KEYS} + new_components: dict[str, list[float]] = {key: [] for key in SET_STATE_COMPONENT_KEYS} for repeat_idx in range(warmup_repeats + measure_repeats): record = repeat_idx >= warmup_repeats @@ -243,14 +268,18 @@ def benchmark_set_state_attribution( ) timing = _run_set_state_once(env, env_ids, qpos, qvel) if record: - old_set_state.append(sum(v for k, v in timing.items() if k.startswith("dr_reset_set_state_"))) + old_set_state.append(float(timing["dr_reset_set_state_wall_ms"])) old_refresh.append(float(timing["dr_reset_set_state_refresh_cache_ms"])) + for key, values in old_components.items(): + values.append(float(timing.get(key, 0.0))) env._backend._refresh_link_pose_cache = original_refresh timing = _run_set_state_once(env, env_ids, qpos, qvel) if record: - new_set_state.append(sum(v for k, v in timing.items() if k.startswith("dr_reset_set_state_"))) + new_set_state.append(float(timing["dr_reset_set_state_wall_ms"])) new_refresh.append(float(timing["dr_reset_set_state_refresh_cache_ms"])) + for key, values in new_components.items(): + values.append(float(timing.get(key, 0.0))) env._backend._refresh_link_pose_cache = original_refresh old_set_state_stats = _stats(old_set_state) @@ -266,6 +295,8 @@ def benchmark_set_state_attribution( new_set_state_ms=new_set_state_stats, old_refresh_cache_ms=old_refresh_stats, new_refresh_cache_ms=new_refresh_stats, + old_component_ms={key: _stats(values) for key, values in old_components.items()}, + new_component_ms={key: _stats(values) for key, values in new_components.items()}, set_state_speedup=_speedup(old_set_state_stats, new_set_state_stats), refresh_cache_speedup=_speedup(old_refresh_stats, new_refresh_stats), ) @@ -361,6 +392,10 @@ def main(argv: list[str] | None = None) -> int: set_state_result.new_refresh_cache_ms, set_state_result.refresh_cache_speedup, ) + for key in SET_STATE_COMPONENT_KEYS: + old = set_state_result.old_component_ms[key] + new = set_state_result.new_component_ms[key] + print(f" {key}: old={old.mean_ms:.6f} ms new={new.mean_ms:.6f} ms") if args.out_json is not None: payload = { diff --git a/src/unilab/base/backend/motrix/backend.py b/src/unilab/base/backend/motrix/backend.py index f44648443..130eac1c7 100644 --- a/src/unilab/base/backend/motrix/backend.py +++ b/src/unilab/base/backend/motrix/backend.py @@ -60,6 +60,16 @@ def _first_scalar(value: Any) -> float: return float(arr.reshape(-1)[0]) +def _contiguous_slice(indices: np.ndarray) -> slice | None: + if indices.size == 0: + return None + start = int(indices[0]) + stop = start + int(indices.size) + if np.array_equal(indices, np.arange(start, stop, dtype=indices.dtype)): + return slice(start, stop) + return None + + @dataclass class _MotrixSceneContext: model: "mtx.SceneModel" @@ -185,6 +195,7 @@ def __init__( self._body_floatingbase = self._body.floatingbase self._joint_dof_pos_indices = np.asarray(self._model.joint_dof_pos_indices, dtype=np.intp) self._joint_dof_vel_indices = np.asarray(self._model.joint_dof_vel_indices, dtype=np.intp) + self._joint_dof_pos_slice = _contiguous_slice(self._joint_dof_pos_indices) position_actuators: list["mtx.PositionActuator"] = [] for actuator in self._model.actuators: if actuator.typ == "position": @@ -212,6 +223,11 @@ def __init__( joint_pos_idx.append(int(joint.dof_pos_index)) if len(joint_pos_idx) == int(self._model.num_actuators): self._actuator_joint_pos_indices = np.asarray(joint_pos_idx, dtype=np.intp) + self._actuator_joint_pos_slice = ( + _contiguous_slice(self._actuator_joint_pos_indices) + if self._actuator_joint_pos_indices is not None + else None + ) self._default_actuator_kp = np.zeros((self.num_actuators,), dtype=np.float32) self._default_actuator_kd = np.zeros((self.num_actuators,), dtype=np.float32) for actuator in self._position_actuators: @@ -599,25 +615,53 @@ def set_state( data_slice = self._data[mask] slice_ms = (time.perf_counter() - t0) * 1000.0 - # Batch set state t0 = time.perf_counter() data_slice.reset(self._model) + data_reset_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() self._clear_applied_body_forces(env_indices) + clear_forces_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() self._apply_init_geom_size_overrides(data_slice, env_indices) + geom_overrides_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() self._apply_reset_randomization(data_slice, env_indices, randomization) + randomization_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() data_slice.set_dof_vel(qvel) + set_dof_vel_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() data_slice.set_dof_pos(qpos_motrix, self._model) - data_write_ms = (time.perf_counter() - t0) * 1000.0 + set_dof_pos_ms = (time.perf_counter() - t0) * 1000.0 + data_write_ms = ( + data_reset_ms + + clear_forces_ms + + geom_overrides_ms + + randomization_ms + + set_dof_vel_ms + + set_dof_pos_ms + ) t0 = time.perf_counter() if self._supports_position_actuator_gains and len(self._joint_dof_pos_indices) == int( self.num_actuators ): # Fully-actuated model: hold every joint at its reset position (unchanged). - ctrl = qpos_motrix[:, self._joint_dof_pos_indices] + if self._joint_dof_pos_slice is not None: + ctrl = qpos_motrix[:, self._joint_dof_pos_slice] + else: + ctrl = qpos_motrix[:, self._joint_dof_pos_indices] elif self._actuator_joint_pos_indices is not None: # Under-actuated / parallel model: hold only the actuated joints. - ctrl = qpos_motrix[:, self._actuator_joint_pos_indices] + if self._actuator_joint_pos_slice is not None: + ctrl = qpos_motrix[:, self._actuator_joint_pos_slice] + else: + ctrl = qpos_motrix[:, self._actuator_joint_pos_indices] else: ctrl = np.zeros((len(env_indices), self.num_actuators), dtype=self._np_dtype) data_slice.actuator_ctrls = np.ascontiguousarray(ctrl) @@ -649,6 +693,12 @@ def set_state( "dr_reset_set_state_qpos_convert_ms": qpos_convert_ms, "dr_reset_set_state_slice_ms": slice_ms, "dr_reset_set_state_data_write_ms": data_write_ms, + "dr_reset_set_state_data_reset_ms": data_reset_ms, + "dr_reset_set_state_clear_forces_ms": clear_forces_ms, + "dr_reset_set_state_geom_overrides_ms": geom_overrides_ms, + "dr_reset_set_state_randomization_ms": randomization_ms, + "dr_reset_set_state_set_dof_vel_ms": set_dof_vel_ms, + "dr_reset_set_state_set_dof_pos_ms": set_dof_pos_ms, "dr_reset_set_state_ctrl_ms": ctrl_ms, "dr_reset_set_state_forward_kinematic_ms": forward_kinematic_ms, "dr_reset_set_state_refresh_cache_ms": refresh_cache_ms, diff --git a/src/unilab/base/np_env.py b/src/unilab/base/np_env.py index 82ec1aa5a..17a9037d8 100644 --- a/src/unilab/base/np_env.py +++ b/src/unilab/base/np_env.py @@ -37,6 +37,12 @@ "dr_reset_set_state_qpos_convert_ms", "dr_reset_set_state_slice_ms", "dr_reset_set_state_data_write_ms", + "dr_reset_set_state_data_reset_ms", + "dr_reset_set_state_clear_forces_ms", + "dr_reset_set_state_geom_overrides_ms", + "dr_reset_set_state_randomization_ms", + "dr_reset_set_state_set_dof_vel_ms", + "dr_reset_set_state_set_dof_pos_ms", "dr_reset_set_state_ctrl_ms", "dr_reset_set_state_forward_kinematic_ms", "dr_reset_set_state_refresh_cache_ms", From 27ba24b3243089daedb37fe9e6c588b9ad66fabe Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Sun, 5 Jul 2026 17:52:05 +0800 Subject: [PATCH 5/6] bench: profile mujoco reset state attribution --- .../benchmark_offpolicy_collector_active.py | 17 ++++ benchmark/benchmark_reset_done_attribution.py | 87 +++++++++++++++++++ src/unilab/base/backend/mujoco/backend.py | 58 ++++++++++++- src/unilab/base/np_env.py | 7 ++ .../test_reset_done_attribution_benchmark.py | 1 + 5 files changed, 168 insertions(+), 2 deletions(-) diff --git a/benchmark/benchmark_offpolicy_collector_active.py b/benchmark/benchmark_offpolicy_collector_active.py index fdaaad57e..2419910ad 100644 --- a/benchmark/benchmark_offpolicy_collector_active.py +++ b/benchmark/benchmark_offpolicy_collector_active.py @@ -91,6 +91,10 @@ "dr_reset_plan_ms", "dr_reset_payload_filter_ms", "dr_reset_set_state_ms", + "dr_reset_set_state_alloc_state_ms", + "dr_reset_set_state_fill_qpos_ms", + "dr_reset_set_state_fill_qvel_ms", + "dr_reset_set_state_env_ids_ms", "dr_reset_set_state_qpos_convert_ms", "dr_reset_set_state_slice_ms", "dr_reset_set_state_data_write_ms", @@ -102,6 +106,9 @@ "dr_reset_set_state_set_dof_pos_ms", "dr_reset_set_state_ctrl_ms", "dr_reset_set_state_forward_kinematic_ms", + "dr_reset_set_state_pool_reset_ms", + "dr_reset_set_state_physics_state_scatter_ms", + "dr_reset_set_state_sensor_scatter_ms", "dr_reset_set_state_refresh_cache_ms", "dr_reset_set_state_invalidate_cache_ms", "dr_reset_set_state_backend_internal_gap_ms", @@ -136,6 +143,10 @@ ("dr_reset_plan_ms", "dr_reset_plan_ms"), ("dr_reset_payload_filter_ms", "dr_reset_payload_filter_ms"), ("dr_reset_set_state_ms", "dr_reset_set_state_ms"), + ("dr_reset_set_state_alloc_state_ms", "dr_reset_set_state_alloc_state_ms"), + ("dr_reset_set_state_fill_qpos_ms", "dr_reset_set_state_fill_qpos_ms"), + ("dr_reset_set_state_fill_qvel_ms", "dr_reset_set_state_fill_qvel_ms"), + ("dr_reset_set_state_env_ids_ms", "dr_reset_set_state_env_ids_ms"), ("dr_reset_set_state_qpos_convert_ms", "dr_reset_set_state_qpos_convert_ms"), ("dr_reset_set_state_slice_ms", "dr_reset_set_state_slice_ms"), ("dr_reset_set_state_data_write_ms", "dr_reset_set_state_data_write_ms"), @@ -147,6 +158,12 @@ ("dr_reset_set_state_set_dof_pos_ms", "dr_reset_set_state_set_dof_pos_ms"), ("dr_reset_set_state_ctrl_ms", "dr_reset_set_state_ctrl_ms"), ("dr_reset_set_state_forward_kinematic_ms", "dr_reset_set_state_forward_kinematic_ms"), + ("dr_reset_set_state_pool_reset_ms", "dr_reset_set_state_pool_reset_ms"), + ( + "dr_reset_set_state_physics_state_scatter_ms", + "dr_reset_set_state_physics_state_scatter_ms", + ), + ("dr_reset_set_state_sensor_scatter_ms", "dr_reset_set_state_sensor_scatter_ms"), ("dr_reset_set_state_refresh_cache_ms", "dr_reset_set_state_refresh_cache_ms"), ("dr_reset_set_state_invalidate_cache_ms", "dr_reset_set_state_invalidate_cache_ms"), ( diff --git a/benchmark/benchmark_reset_done_attribution.py b/benchmark/benchmark_reset_done_attribution.py index ead1f9209..048dc3c9d 100644 --- a/benchmark/benchmark_reset_done_attribution.py +++ b/benchmark/benchmark_reset_done_attribution.py @@ -73,6 +73,16 @@ class SetStateAttributionResult: refresh_cache_speedup: float +@dataclass(frozen=True) +class SetStateProfileResult: + backend: str + num_envs: int + reset_count: int + repeats: int + set_state_ms: TimingStats + component_ms: dict[str, TimingStats] + + def _stats(samples_ms: list[float]) -> TimingStats: if not samples_ms: raise ValueError("no samples") @@ -212,6 +222,10 @@ def _old_refresh_link_pose_cache(self, env_indices: np.ndarray | None = None, da SET_STATE_COMPONENT_KEYS = ( + "dr_reset_set_state_alloc_state_ms", + "dr_reset_set_state_fill_qpos_ms", + "dr_reset_set_state_fill_qvel_ms", + "dr_reset_set_state_env_ids_ms", "dr_reset_set_state_qpos_convert_ms", "dr_reset_set_state_slice_ms", "dr_reset_set_state_data_write_ms", @@ -223,6 +237,9 @@ def _old_refresh_link_pose_cache(self, env_indices: np.ndarray | None = None, da "dr_reset_set_state_set_dof_pos_ms", "dr_reset_set_state_ctrl_ms", "dr_reset_set_state_forward_kinematic_ms", + "dr_reset_set_state_pool_reset_ms", + "dr_reset_set_state_physics_state_scatter_ms", + "dr_reset_set_state_sensor_scatter_ms", "dr_reset_set_state_refresh_cache_ms", "dr_reset_set_state_invalidate_cache_ms", "dr_reset_set_state_backend_internal_gap_ms", @@ -304,6 +321,49 @@ def benchmark_set_state_attribution( env.close() +def benchmark_set_state_profile( + *, + backend: str, + num_envs: int, + reset_count: int, + warmup_repeats: int, + measure_repeats: int, + seed: int, +) -> SetStateProfileResult: + env = _compose_env("flashsac", "g1_walk_flat", backend, num_envs) + try: + env_ids = _sample_env_ids(num_envs, reset_count, seed) + plan = env._dr_manager._provider.build_reset_plan(env, env_ids) + qpos = np.asarray(plan.qpos) + qvel = np.asarray(plan.qvel) + randomization = plan.randomization + + set_state: list[float] = [] + components: dict[str, list[float]] = {key: [] for key in SET_STATE_COMPONENT_KEYS} + + for repeat_idx in range(warmup_repeats + measure_repeats): + record = repeat_idx >= warmup_repeats + t0 = time.perf_counter() + env._backend.set_state(env_ids, qpos, qvel, randomization=randomization) + wall_ms = (time.perf_counter() - t0) * 1000.0 + timing = env._backend.last_set_state_timing_ms + if record: + set_state.append(wall_ms) + for key, values in components.items(): + values.append(float(timing.get(key, 0.0))) + + return SetStateProfileResult( + backend=backend, + num_envs=num_envs, + reset_count=reset_count, + repeats=measure_repeats, + set_state_ms=_stats(set_state), + component_ms={key: _stats(values) for key, values in components.items()}, + ) + finally: + env.close() + + def _print_stats(label: str, old: TimingStats, new: TimingStats, speedup: float) -> None: print( f" {label}: old={old.mean_ms:.6f} ms new={new.mean_ms:.6f} ms " @@ -327,6 +387,7 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: parser.add_argument("--measure-repeats", type=int, default=30) parser.add_argument("--seed", type=int, default=673) parser.add_argument("--body-pose-backends", default="motrix,mujoco") + parser.add_argument("--set-state-profile-backends", default="motrix,mujoco") parser.add_argument("--skip-set-state", action="store_true") parser.add_argument("--out-json", type=Path, default=None) return parser.parse_args(argv) @@ -367,6 +428,7 @@ def main(argv: list[str] | None = None) -> int: ) set_state_result = None + set_state_profiles: list[SetStateProfileResult] = [] if not args.skip_set_state: set_state_result = _run_safely( "set_state/motrix", @@ -397,6 +459,30 @@ def main(argv: list[str] | None = None) -> int: new = set_state_result.new_component_ms[key] print(f" {key}: old={old.mean_ms:.6f} ms new={new.mean_ms:.6f} ms") + for backend in [ + part.strip() for part in args.set_state_profile_backends.split(",") if part.strip() + ]: + profile = _run_safely( + f"set_state_profile/{backend}", + lambda backend=backend: benchmark_set_state_profile( + backend=backend, + num_envs=args.num_envs, + reset_count=args.reset_count, + warmup_repeats=args.warmup_repeats, + measure_repeats=args.measure_repeats, + seed=args.seed, + ), + ) + if profile is None: + continue + set_state_profiles.append(profile) + print(f"set_state_profile/{backend}:") + print(f" set_state_ms: current={profile.set_state_ms.mean_ms:.6f} ms") + for key in SET_STATE_COMPONENT_KEYS: + value = profile.component_ms[key] + if value.mean_ms > 0.0: + print(f" {key}: current={value.mean_ms:.6f} ms") + if args.out_json is not None: payload = { "device": get_device_info_dict(), @@ -407,6 +493,7 @@ def main(argv: list[str] | None = None) -> int: "seed": args.seed, "body_pose": [asdict(result) for result in body_pose_results], "set_state": asdict(set_state_result) if set_state_result is not None else None, + "set_state_profiles": [asdict(result) for result in set_state_profiles], } args.out_json.parent.mkdir(parents=True, exist_ok=True) args.out_json.write_text(json.dumps(payload, indent=2), encoding="utf-8") diff --git a/src/unilab/base/backend/mujoco/backend.py b/src/unilab/base/backend/mujoco/backend.py index d2c9e4013..5ea6b76ee 100644 --- a/src/unilab/base/backend/mujoco/backend.py +++ b/src/unilab/base/backend/mujoco/backend.py @@ -319,6 +319,7 @@ def __init__( self._np_dtype = np_dtype if np_dtype is not None else get_global_dtype() self.backend_type = "mujoco" self._pending_xfrc_applied = np.zeros((num_envs, 6 * self._model.nbody), dtype=np.float64) + self._last_set_state_timing_ms: dict[str, float] = {} # Thread configuration. self._n_threads = min(num_envs, cpu_count() * 2) @@ -409,6 +410,10 @@ def _get_sensor_view(prefix, dim): self._tracked_linvel_b_all = _get_sensor_view("track_linvel_b", 3) self._tracked_angvel_b_all = _get_sensor_view("track_angvel_b", 3) + @property + def last_set_state_timing_ms(self) -> dict[str, float]: + return dict(self._last_set_state_timing_ms) + def _load_base_model(self) -> mujoco.MjModel: if isinstance(self._model_file, mujoco.MjModel): if self.add_body_sensors: @@ -878,22 +883,71 @@ def set_state( qvel: np.ndarray, randomization: ResetRandomizationPayload | None = None, ) -> None: + total_t0 = time.perf_counter() + self._last_set_state_timing_ms = {} if len(env_indices) == 0: return num_reset = len(env_indices) + + t0 = time.perf_counter() state_np = np.zeros((num_reset, self._physics_state.shape[1]), dtype=np.float64) + alloc_state_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() state_np[:, self._idx_qpos : self._idx_qpos + self.nq] = qpos + fill_qpos_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() state_np[:, self._idx_qvel : self._idx_qvel + self.nv] = qvel + fill_qvel_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + env_ids = np.asarray(env_indices, dtype=np.int32) + env_ids_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + reset_randomization = self._translate_reset_randomization(randomization, num_reset) + randomization_ms = (time.perf_counter() - t0) * 1000.0 + t0 = time.perf_counter() state_out, sensor_np = self._pool.reset( # type: ignore[union-attr] - env_ids=np.asarray(env_indices, dtype=np.int32), + env_ids=env_ids, initial_state=state_np, - randomization=self._translate_reset_randomization(randomization, num_reset), + randomization=reset_randomization, ) + pool_reset_ms = (time.perf_counter() - t0) * 1000.0 + t0 = time.perf_counter() self._physics_state[env_indices] = state_out.astype(self._np_dtype) + physics_state_scatter_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() self._sensor_data[env_indices] = sensor_np.astype(self._np_dtype) + sensor_scatter_ms = (time.perf_counter() - t0) * 1000.0 + + total_ms = (time.perf_counter() - total_t0) * 1000.0 + measured_ms = ( + alloc_state_ms + + fill_qpos_ms + + fill_qvel_ms + + env_ids_ms + + randomization_ms + + pool_reset_ms + + physics_state_scatter_ms + + sensor_scatter_ms + ) + self._last_set_state_timing_ms = { + "dr_reset_set_state_alloc_state_ms": alloc_state_ms, + "dr_reset_set_state_fill_qpos_ms": fill_qpos_ms, + "dr_reset_set_state_fill_qvel_ms": fill_qvel_ms, + "dr_reset_set_state_env_ids_ms": env_ids_ms, + "dr_reset_set_state_randomization_ms": randomization_ms, + "dr_reset_set_state_pool_reset_ms": pool_reset_ms, + "dr_reset_set_state_physics_state_scatter_ms": physics_state_scatter_ms, + "dr_reset_set_state_sensor_scatter_ms": sensor_scatter_ms, + "dr_reset_set_state_backend_internal_gap_ms": total_ms - measured_ms, + } def get_dr_capabilities(self) -> DomainRandomizationCapabilities: return DomainRandomizationCapabilities( diff --git a/src/unilab/base/np_env.py b/src/unilab/base/np_env.py index 17a9037d8..018df89e8 100644 --- a/src/unilab/base/np_env.py +++ b/src/unilab/base/np_env.py @@ -34,6 +34,10 @@ "dr_reset_plan_ms", "dr_reset_payload_filter_ms", "dr_reset_set_state_ms", + "dr_reset_set_state_alloc_state_ms", + "dr_reset_set_state_fill_qpos_ms", + "dr_reset_set_state_fill_qvel_ms", + "dr_reset_set_state_env_ids_ms", "dr_reset_set_state_qpos_convert_ms", "dr_reset_set_state_slice_ms", "dr_reset_set_state_data_write_ms", @@ -45,6 +49,9 @@ "dr_reset_set_state_set_dof_pos_ms", "dr_reset_set_state_ctrl_ms", "dr_reset_set_state_forward_kinematic_ms", + "dr_reset_set_state_pool_reset_ms", + "dr_reset_set_state_physics_state_scatter_ms", + "dr_reset_set_state_sensor_scatter_ms", "dr_reset_set_state_refresh_cache_ms", "dr_reset_set_state_invalidate_cache_ms", "dr_reset_set_state_backend_internal_gap_ms", diff --git a/tests/benchmark/test_reset_done_attribution_benchmark.py b/tests/benchmark/test_reset_done_attribution_benchmark.py index 6fcc6e630..0aa776f8e 100644 --- a/tests/benchmark/test_reset_done_attribution_benchmark.py +++ b/tests/benchmark/test_reset_done_attribution_benchmark.py @@ -37,4 +37,5 @@ def test_parse_args_defaults_to_issue_scale() -> None: assert args.reset_count == 256 assert args.measure_repeats == 30 assert args.body_pose_backends == "motrix,mujoco" + assert args.set_state_profile_backends == "motrix,mujoco" assert not args.skip_set_state From 6ae8f8d7e9114648830fd927b69fa161fb1928cf Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Sun, 5 Jul 2026 18:16:48 +0800 Subject: [PATCH 6/6] chore: remove reset attribution scaffolding --- .../benchmark_offpolicy_collector_active.py | 54 +- benchmark/benchmark_reset_done_attribution.py | 506 ------------------ src/unilab/base/backend/base.py | 5 - src/unilab/base/backend/motrix/backend.py | 75 +-- src/unilab/base/backend/mujoco/backend.py | 58 +- src/unilab/base/np_env.py | 21 - src/unilab/dr/manager.py | 1 - .../test_reset_done_attribution_benchmark.py | 41 -- tests/dr/test_manager.py | 7 - 9 files changed, 4 insertions(+), 764 deletions(-) delete mode 100644 benchmark/benchmark_reset_done_attribution.py delete mode 100644 tests/benchmark/test_reset_done_attribution_benchmark.py diff --git a/benchmark/benchmark_offpolicy_collector_active.py b/benchmark/benchmark_offpolicy_collector_active.py index 2419910ad..588d11b95 100644 --- a/benchmark/benchmark_offpolicy_collector_active.py +++ b/benchmark/benchmark_offpolicy_collector_active.py @@ -91,27 +91,6 @@ "dr_reset_plan_ms", "dr_reset_payload_filter_ms", "dr_reset_set_state_ms", - "dr_reset_set_state_alloc_state_ms", - "dr_reset_set_state_fill_qpos_ms", - "dr_reset_set_state_fill_qvel_ms", - "dr_reset_set_state_env_ids_ms", - "dr_reset_set_state_qpos_convert_ms", - "dr_reset_set_state_slice_ms", - "dr_reset_set_state_data_write_ms", - "dr_reset_set_state_data_reset_ms", - "dr_reset_set_state_clear_forces_ms", - "dr_reset_set_state_geom_overrides_ms", - "dr_reset_set_state_randomization_ms", - "dr_reset_set_state_set_dof_vel_ms", - "dr_reset_set_state_set_dof_pos_ms", - "dr_reset_set_state_ctrl_ms", - "dr_reset_set_state_forward_kinematic_ms", - "dr_reset_set_state_pool_reset_ms", - "dr_reset_set_state_physics_state_scatter_ms", - "dr_reset_set_state_sensor_scatter_ms", - "dr_reset_set_state_refresh_cache_ms", - "dr_reset_set_state_invalidate_cache_ms", - "dr_reset_set_state_backend_internal_gap_ms", "dr_reset_build_observation_ms", "dr_reset_internal_gap_ms", "dr_reset_observation_getters_ms", @@ -143,33 +122,6 @@ ("dr_reset_plan_ms", "dr_reset_plan_ms"), ("dr_reset_payload_filter_ms", "dr_reset_payload_filter_ms"), ("dr_reset_set_state_ms", "dr_reset_set_state_ms"), - ("dr_reset_set_state_alloc_state_ms", "dr_reset_set_state_alloc_state_ms"), - ("dr_reset_set_state_fill_qpos_ms", "dr_reset_set_state_fill_qpos_ms"), - ("dr_reset_set_state_fill_qvel_ms", "dr_reset_set_state_fill_qvel_ms"), - ("dr_reset_set_state_env_ids_ms", "dr_reset_set_state_env_ids_ms"), - ("dr_reset_set_state_qpos_convert_ms", "dr_reset_set_state_qpos_convert_ms"), - ("dr_reset_set_state_slice_ms", "dr_reset_set_state_slice_ms"), - ("dr_reset_set_state_data_write_ms", "dr_reset_set_state_data_write_ms"), - ("dr_reset_set_state_data_reset_ms", "dr_reset_set_state_data_reset_ms"), - ("dr_reset_set_state_clear_forces_ms", "dr_reset_set_state_clear_forces_ms"), - ("dr_reset_set_state_geom_overrides_ms", "dr_reset_set_state_geom_overrides_ms"), - ("dr_reset_set_state_randomization_ms", "dr_reset_set_state_randomization_ms"), - ("dr_reset_set_state_set_dof_vel_ms", "dr_reset_set_state_set_dof_vel_ms"), - ("dr_reset_set_state_set_dof_pos_ms", "dr_reset_set_state_set_dof_pos_ms"), - ("dr_reset_set_state_ctrl_ms", "dr_reset_set_state_ctrl_ms"), - ("dr_reset_set_state_forward_kinematic_ms", "dr_reset_set_state_forward_kinematic_ms"), - ("dr_reset_set_state_pool_reset_ms", "dr_reset_set_state_pool_reset_ms"), - ( - "dr_reset_set_state_physics_state_scatter_ms", - "dr_reset_set_state_physics_state_scatter_ms", - ), - ("dr_reset_set_state_sensor_scatter_ms", "dr_reset_set_state_sensor_scatter_ms"), - ("dr_reset_set_state_refresh_cache_ms", "dr_reset_set_state_refresh_cache_ms"), - ("dr_reset_set_state_invalidate_cache_ms", "dr_reset_set_state_invalidate_cache_ms"), - ( - "dr_reset_set_state_backend_internal_gap_ms", - "dr_reset_set_state_backend_internal_gap_ms", - ), ("dr_reset_build_observation_ms", "dr_reset_build_observation_ms"), ("dr_reset_internal_gap_ms", "dr_reset_internal_gap_ms"), ("dr_reset_observation_getters_ms", "dr_reset_observation_getters_ms"), @@ -1443,10 +1395,6 @@ def _print_result(result: CollectorResult) -> None: "reset_done_info_scatter_ms", "dr_reset_plan_ms", "dr_reset_set_state_ms", - "dr_reset_set_state_slice_ms", - "dr_reset_set_state_data_write_ms", - "dr_reset_set_state_forward_kinematic_ms", - "dr_reset_set_state_refresh_cache_ms", "dr_reset_build_observation_ms", "dr_reset_obs_get_motion_ms", "dr_reset_observation_getters_ms", @@ -1615,7 +1563,7 @@ def main() -> int: print("\nReset done timing (subparts of NpEnv reset_done_ms):") print(_format_reset_done_timing_table(results)) print( - "\nDR reset timing (subparts of reset call; includes reset observation getters):" + "\nDR reset timing (subparts of reset call; reset obs getters currently read full batch):" ) print(_format_dr_reset_timing_table(results)) else: diff --git a/benchmark/benchmark_reset_done_attribution.py b/benchmark/benchmark_reset_done_attribution.py deleted file mode 100644 index 048dc3c9d..000000000 --- a/benchmark/benchmark_reset_done_attribution.py +++ /dev/null @@ -1,506 +0,0 @@ -#!/usr/bin/env python3 -"""Attribution harness for reset_done changes from issue #673. - -This benchmark avoids full collector throughput attribution. It constructs one -env per case and alternates old/new implementations in the same process, using -the same env_ids and reset payloads. - -Usage: - uv run benchmark/benchmark_reset_done_attribution.py - uv run benchmark/benchmark_reset_done_attribution.py --num-envs 8192 --reset-count 256 - uv run benchmark/benchmark_reset_done_attribution.py --out-json /tmp/reset_done_attr.json -""" - -from __future__ import annotations - -import argparse -import json -import sys -import time -import types -from dataclasses import asdict, dataclass -from pathlib import Path -from typing import Any, Callable - -import numpy as np -from hydra import compose, initialize_config_dir -from hydra.core.global_hydra import GlobalHydra - -ROOT_DIR = Path(__file__).resolve().parents[1] -if str(ROOT_DIR) not in sys.path: - sys.path.append(str(ROOT_DIR)) - -from benchmark.core.device_info import get_device_info_dict, get_device_info_line - - -@dataclass(frozen=True) -class TimingStats: - mean_ms: float - median_ms: float - min_ms: float - max_ms: float - samples_ms: list[float] - - -@dataclass(frozen=True) -class BodyPoseAttributionResult: - backend: str - num_envs: int - reset_count: int - body_count: int - repeats: int - old_build_observation_ms: TimingStats - new_build_observation_ms: TimingStats - old_body_pose_ms: TimingStats - new_body_pose_ms: TimingStats - build_observation_speedup: float - body_pose_speedup: float - - -@dataclass(frozen=True) -class SetStateAttributionResult: - backend: str - num_envs: int - reset_count: int - repeats: int - old_set_state_ms: TimingStats - new_set_state_ms: TimingStats - old_refresh_cache_ms: TimingStats - new_refresh_cache_ms: TimingStats - old_component_ms: dict[str, TimingStats] - new_component_ms: dict[str, TimingStats] - set_state_speedup: float - refresh_cache_speedup: float - - -@dataclass(frozen=True) -class SetStateProfileResult: - backend: str - num_envs: int - reset_count: int - repeats: int - set_state_ms: TimingStats - component_ms: dict[str, TimingStats] - - -def _stats(samples_ms: list[float]) -> TimingStats: - if not samples_ms: - raise ValueError("no samples") - arr = np.asarray(samples_ms, dtype=np.float64) - return TimingStats( - mean_ms=float(arr.mean()), - median_ms=float(np.median(arr)), - min_ms=float(arr.min()), - max_ms=float(arr.max()), - samples_ms=[float(v) for v in samples_ms], - ) - - -def _speedup(old: TimingStats, new: TimingStats) -> float: - return old.mean_ms / new.mean_ms if new.mean_ms > 0.0 else float("inf") - - -def _sample_env_ids(num_envs: int, reset_count: int, seed: int) -> np.ndarray: - if reset_count <= 0: - raise ValueError("reset-count must be > 0") - if reset_count > num_envs: - raise ValueError("reset-count must be <= num-envs") - rng = np.random.default_rng(seed) - return np.sort(rng.choice(num_envs, size=reset_count, replace=False).astype(np.int32)) - - -def _compose_env(algo: str, task: str, backend: str, num_envs: int): - from unilab.training import BackendAdapter, create_env, ensure_registries - - ensure_registries() - GlobalHydra.instance().clear() - with initialize_config_dir(version_base=None, config_dir=str(ROOT_DIR / "conf" / "offpolicy")): - cfg = compose( - "config", - overrides=[ - f"algo={algo}", - f"task={algo}/{task}/{backend}", - f"algo.num_envs={num_envs}", - ], - ) - env_cfg_override = BackendAdapter( - cfg, - root_dir=ROOT_DIR, - algo_name=algo, - ).build_task_env_cfg_override() - env = create_env(cfg, num_envs=num_envs, env_cfg_override=env_cfg_override) - env.init_state() - return env - - -def _body_pose_old_rows(backend: Any, rows: np.ndarray, body_ids: np.ndarray): - pos, quat = backend.get_body_pose_w(body_ids) - return pos[rows], quat[rows] - - -def _run_body_pose_once(provider: Any, env: Any, env_ids: np.ndarray) -> tuple[float, float]: - t0 = time.perf_counter() - provider.build_reset_observation(env, env_ids, {}) - total_ms = (time.perf_counter() - t0) * 1000.0 - timing = provider.last_reset_observation_timing_ms - return total_ms, float(timing["dr_reset_obs_get_body_pose_ms"]) - - -def benchmark_body_pose_attribution( - *, - backend: str, - num_envs: int, - reset_count: int, - warmup_repeats: int, - measure_repeats: int, - seed: int, -) -> BodyPoseAttributionResult: - env = _compose_env("sac", "g1_motion_tracking", backend, num_envs) - try: - provider = env._dr_manager._provider - env_ids = _sample_env_ids(num_envs, reset_count, seed) - original_rows = env._backend.get_body_pose_w_rows - - def old_rows(self, rows: np.ndarray, body_ids: np.ndarray): - return _body_pose_old_rows(self, rows, body_ids) - - old_pos, old_quat = _body_pose_old_rows(env._backend, env_ids, env.body_ids) - new_pos, new_quat = original_rows(env_ids, env.body_ids) - np.testing.assert_allclose(old_pos, new_pos, rtol=0, atol=0) - np.testing.assert_allclose(old_quat, new_quat, rtol=0, atol=0) - - old_total: list[float] = [] - new_total: list[float] = [] - old_body: list[float] = [] - new_body: list[float] = [] - - for repeat_idx in range(warmup_repeats + measure_repeats): - record = repeat_idx >= warmup_repeats - - env._backend.get_body_pose_w_rows = types.MethodType(old_rows, env._backend) - total_ms, body_ms = _run_body_pose_once(provider, env, env_ids) - if record: - old_total.append(total_ms) - old_body.append(body_ms) - - env._backend.get_body_pose_w_rows = original_rows - total_ms, body_ms = _run_body_pose_once(provider, env, env_ids) - if record: - new_total.append(total_ms) - new_body.append(body_ms) - - env._backend.get_body_pose_w_rows = original_rows - old_total_stats = _stats(old_total) - new_total_stats = _stats(new_total) - old_body_stats = _stats(old_body) - new_body_stats = _stats(new_body) - return BodyPoseAttributionResult( - backend=backend, - num_envs=num_envs, - reset_count=reset_count, - body_count=int(len(env.body_ids)), - repeats=measure_repeats, - old_build_observation_ms=old_total_stats, - new_build_observation_ms=new_total_stats, - old_body_pose_ms=old_body_stats, - new_body_pose_ms=new_body_stats, - build_observation_speedup=_speedup(old_total_stats, new_total_stats), - body_pose_speedup=_speedup(old_body_stats, new_body_stats), - ) - finally: - env.close() - - -def _old_refresh_link_pose_cache(self, env_indices: np.ndarray | None = None, data_slice=None): - del data_slice - if env_indices is None: - self._link_poses = self._model.get_link_poses(self._data) - else: - mask = np.zeros(self._num_envs, dtype=bool) - mask[env_indices] = True - self._link_poses[env_indices] = self._model.get_link_poses(self._data[mask]) - - -SET_STATE_COMPONENT_KEYS = ( - "dr_reset_set_state_alloc_state_ms", - "dr_reset_set_state_fill_qpos_ms", - "dr_reset_set_state_fill_qvel_ms", - "dr_reset_set_state_env_ids_ms", - "dr_reset_set_state_qpos_convert_ms", - "dr_reset_set_state_slice_ms", - "dr_reset_set_state_data_write_ms", - "dr_reset_set_state_data_reset_ms", - "dr_reset_set_state_clear_forces_ms", - "dr_reset_set_state_geom_overrides_ms", - "dr_reset_set_state_randomization_ms", - "dr_reset_set_state_set_dof_vel_ms", - "dr_reset_set_state_set_dof_pos_ms", - "dr_reset_set_state_ctrl_ms", - "dr_reset_set_state_forward_kinematic_ms", - "dr_reset_set_state_pool_reset_ms", - "dr_reset_set_state_physics_state_scatter_ms", - "dr_reset_set_state_sensor_scatter_ms", - "dr_reset_set_state_refresh_cache_ms", - "dr_reset_set_state_invalidate_cache_ms", - "dr_reset_set_state_backend_internal_gap_ms", -) - - -def _run_set_state_once(env: Any, env_ids: np.ndarray, qpos: np.ndarray, qvel: np.ndarray) -> dict: - t0 = time.perf_counter() - env._backend.set_state(env_ids, qpos, qvel) - wall_ms = (time.perf_counter() - t0) * 1000.0 - timing = env._backend.last_set_state_timing_ms - return {"dr_reset_set_state_wall_ms": wall_ms, **timing} - - -def benchmark_set_state_attribution( - *, - num_envs: int, - reset_count: int, - warmup_repeats: int, - measure_repeats: int, - seed: int, -) -> SetStateAttributionResult: - env = _compose_env("flashsac", "g1_walk_flat", "motrix", num_envs) - try: - env_ids = _sample_env_ids(num_envs, reset_count, seed) - plan = env._dr_manager._provider.build_reset_plan(env, env_ids) - qpos = np.asarray(plan.qpos) - qvel = np.asarray(plan.qvel) - original_refresh = env._backend._refresh_link_pose_cache - - old_set_state: list[float] = [] - new_set_state: list[float] = [] - old_refresh: list[float] = [] - new_refresh: list[float] = [] - old_components: dict[str, list[float]] = {key: [] for key in SET_STATE_COMPONENT_KEYS} - new_components: dict[str, list[float]] = {key: [] for key in SET_STATE_COMPONENT_KEYS} - - for repeat_idx in range(warmup_repeats + measure_repeats): - record = repeat_idx >= warmup_repeats - - env._backend._refresh_link_pose_cache = types.MethodType( - _old_refresh_link_pose_cache, env._backend - ) - timing = _run_set_state_once(env, env_ids, qpos, qvel) - if record: - old_set_state.append(float(timing["dr_reset_set_state_wall_ms"])) - old_refresh.append(float(timing["dr_reset_set_state_refresh_cache_ms"])) - for key, values in old_components.items(): - values.append(float(timing.get(key, 0.0))) - - env._backend._refresh_link_pose_cache = original_refresh - timing = _run_set_state_once(env, env_ids, qpos, qvel) - if record: - new_set_state.append(float(timing["dr_reset_set_state_wall_ms"])) - new_refresh.append(float(timing["dr_reset_set_state_refresh_cache_ms"])) - for key, values in new_components.items(): - values.append(float(timing.get(key, 0.0))) - - env._backend._refresh_link_pose_cache = original_refresh - old_set_state_stats = _stats(old_set_state) - new_set_state_stats = _stats(new_set_state) - old_refresh_stats = _stats(old_refresh) - new_refresh_stats = _stats(new_refresh) - return SetStateAttributionResult( - backend="motrix", - num_envs=num_envs, - reset_count=reset_count, - repeats=measure_repeats, - old_set_state_ms=old_set_state_stats, - new_set_state_ms=new_set_state_stats, - old_refresh_cache_ms=old_refresh_stats, - new_refresh_cache_ms=new_refresh_stats, - old_component_ms={key: _stats(values) for key, values in old_components.items()}, - new_component_ms={key: _stats(values) for key, values in new_components.items()}, - set_state_speedup=_speedup(old_set_state_stats, new_set_state_stats), - refresh_cache_speedup=_speedup(old_refresh_stats, new_refresh_stats), - ) - finally: - env.close() - - -def benchmark_set_state_profile( - *, - backend: str, - num_envs: int, - reset_count: int, - warmup_repeats: int, - measure_repeats: int, - seed: int, -) -> SetStateProfileResult: - env = _compose_env("flashsac", "g1_walk_flat", backend, num_envs) - try: - env_ids = _sample_env_ids(num_envs, reset_count, seed) - plan = env._dr_manager._provider.build_reset_plan(env, env_ids) - qpos = np.asarray(plan.qpos) - qvel = np.asarray(plan.qvel) - randomization = plan.randomization - - set_state: list[float] = [] - components: dict[str, list[float]] = {key: [] for key in SET_STATE_COMPONENT_KEYS} - - for repeat_idx in range(warmup_repeats + measure_repeats): - record = repeat_idx >= warmup_repeats - t0 = time.perf_counter() - env._backend.set_state(env_ids, qpos, qvel, randomization=randomization) - wall_ms = (time.perf_counter() - t0) * 1000.0 - timing = env._backend.last_set_state_timing_ms - if record: - set_state.append(wall_ms) - for key, values in components.items(): - values.append(float(timing.get(key, 0.0))) - - return SetStateProfileResult( - backend=backend, - num_envs=num_envs, - reset_count=reset_count, - repeats=measure_repeats, - set_state_ms=_stats(set_state), - component_ms={key: _stats(values) for key, values in components.items()}, - ) - finally: - env.close() - - -def _print_stats(label: str, old: TimingStats, new: TimingStats, speedup: float) -> None: - print( - f" {label}: old={old.mean_ms:.6f} ms new={new.mean_ms:.6f} ms " - f"speedup={speedup:.2f}x" - ) - - -def _run_safely(label: str, fn: Callable[[], Any]) -> Any | None: - try: - return fn() - except Exception as exc: - print(f"{label}: ERROR {type(exc).__name__}: {exc}") - return None - - -def parse_args(argv: list[str] | None = None) -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--num-envs", type=int, default=8192) - parser.add_argument("--reset-count", type=int, default=256) - parser.add_argument("--warmup-repeats", type=int, default=5) - parser.add_argument("--measure-repeats", type=int, default=30) - parser.add_argument("--seed", type=int, default=673) - parser.add_argument("--body-pose-backends", default="motrix,mujoco") - parser.add_argument("--set-state-profile-backends", default="motrix,mujoco") - parser.add_argument("--skip-set-state", action="store_true") - parser.add_argument("--out-json", type=Path, default=None) - return parser.parse_args(argv) - - -def main(argv: list[str] | None = None) -> int: - args = parse_args(argv) - print(f"Device: {get_device_info_line()}") - - body_pose_results: list[BodyPoseAttributionResult] = [] - for backend in [part.strip() for part in args.body_pose_backends.split(",") if part.strip()]: - result = _run_safely( - f"body_pose/{backend}", - lambda backend=backend: benchmark_body_pose_attribution( - backend=backend, - num_envs=args.num_envs, - reset_count=args.reset_count, - warmup_repeats=args.warmup_repeats, - measure_repeats=args.measure_repeats, - seed=args.seed, - ), - ) - if result is None: - continue - body_pose_results.append(result) - print(f"body_pose/{backend}:") - _print_stats( - "build_observation_ms", - result.old_build_observation_ms, - result.new_build_observation_ms, - result.build_observation_speedup, - ) - _print_stats( - "body_pose_ms", - result.old_body_pose_ms, - result.new_body_pose_ms, - result.body_pose_speedup, - ) - - set_state_result = None - set_state_profiles: list[SetStateProfileResult] = [] - if not args.skip_set_state: - set_state_result = _run_safely( - "set_state/motrix", - lambda: benchmark_set_state_attribution( - num_envs=args.num_envs, - reset_count=args.reset_count, - warmup_repeats=args.warmup_repeats, - measure_repeats=args.measure_repeats, - seed=args.seed, - ), - ) - if set_state_result is not None: - print("set_state/motrix:") - _print_stats( - "set_state_ms", - set_state_result.old_set_state_ms, - set_state_result.new_set_state_ms, - set_state_result.set_state_speedup, - ) - _print_stats( - "refresh_cache_ms", - set_state_result.old_refresh_cache_ms, - set_state_result.new_refresh_cache_ms, - set_state_result.refresh_cache_speedup, - ) - for key in SET_STATE_COMPONENT_KEYS: - old = set_state_result.old_component_ms[key] - new = set_state_result.new_component_ms[key] - print(f" {key}: old={old.mean_ms:.6f} ms new={new.mean_ms:.6f} ms") - - for backend in [ - part.strip() for part in args.set_state_profile_backends.split(",") if part.strip() - ]: - profile = _run_safely( - f"set_state_profile/{backend}", - lambda backend=backend: benchmark_set_state_profile( - backend=backend, - num_envs=args.num_envs, - reset_count=args.reset_count, - warmup_repeats=args.warmup_repeats, - measure_repeats=args.measure_repeats, - seed=args.seed, - ), - ) - if profile is None: - continue - set_state_profiles.append(profile) - print(f"set_state_profile/{backend}:") - print(f" set_state_ms: current={profile.set_state_ms.mean_ms:.6f} ms") - for key in SET_STATE_COMPONENT_KEYS: - value = profile.component_ms[key] - if value.mean_ms > 0.0: - print(f" {key}: current={value.mean_ms:.6f} ms") - - if args.out_json is not None: - payload = { - "device": get_device_info_dict(), - "num_envs": args.num_envs, - "reset_count": args.reset_count, - "warmup_repeats": args.warmup_repeats, - "measure_repeats": args.measure_repeats, - "seed": args.seed, - "body_pose": [asdict(result) for result in body_pose_results], - "set_state": asdict(set_state_result) if set_state_result is not None else None, - "set_state_profiles": [asdict(result) for result in set_state_profiles], - } - args.out_json.parent.mkdir(parents=True, exist_ok=True) - args.out_json.write_text(json.dumps(payload, indent=2), encoding="utf-8") - print(f"Wrote JSON: {args.out_json}") - - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/src/unilab/base/backend/base.py b/src/unilab/base/backend/base.py index 2ef20e403..96cc56f53 100644 --- a/src/unilab/base/backend/base.py +++ b/src/unilab/base/backend/base.py @@ -281,11 +281,6 @@ def set_state( randomization: Optional backend randomization payload. """ - @property - def last_set_state_timing_ms(self) -> dict[str, float]: - """Return backend-internal timing from the most recent ``set_state`` call.""" - return {} - @abc.abstractmethod def get_dr_capabilities(self) -> DomainRandomizationCapabilities: """Return supported domain-randomization capabilities for this backend.""" diff --git a/src/unilab/base/backend/motrix/backend.py b/src/unilab/base/backend/motrix/backend.py index 130eac1c7..49aaa417f 100644 --- a/src/unilab/base/backend/motrix/backend.py +++ b/src/unilab/base/backend/motrix/backend.py @@ -294,7 +294,6 @@ def __init__( self._render_tracking_camera: MotrixTrackingCamera | None = None self.backend_type = "motrix" self._link_velocity_cache: np.ndarray | None = None - self._last_set_state_timing_ms: dict[str, float] = {} # Pre-cache link objects to avoid repeated get_link() lookups. self._link_cache: dict[int, "mtx.Link"] = {} @@ -334,10 +333,6 @@ def model(self): def data(self): return self._data - @property - def last_set_state_timing_ms(self) -> dict[str, float]: - return dict(self._last_set_state_timing_ms) - # ------------------------------------------------------------------ # # Model properties # # ------------------------------------------------------------------ # @@ -601,53 +596,21 @@ def set_state( qvel: np.ndarray, randomization: ResetRandomizationPayload | None = None, ) -> None: - total_t0 = time.perf_counter() - self._last_set_state_timing_ms = {} - - t0 = time.perf_counter() qpos_motrix = self._mujoco_qpos_to_motrix(qpos) - qpos_convert_ms = (time.perf_counter() - t0) * 1000.0 - t0 = time.perf_counter() # Create mask for batch operation mask = np.zeros(self._num_envs, dtype=bool) mask[env_indices] = True data_slice = self._data[mask] - slice_ms = (time.perf_counter() - t0) * 1000.0 - t0 = time.perf_counter() + # Batch set state data_slice.reset(self._model) - data_reset_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() self._clear_applied_body_forces(env_indices) - clear_forces_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() self._apply_init_geom_size_overrides(data_slice, env_indices) - geom_overrides_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() self._apply_reset_randomization(data_slice, env_indices, randomization) - randomization_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() data_slice.set_dof_vel(qvel) - set_dof_vel_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() data_slice.set_dof_pos(qpos_motrix, self._model) - set_dof_pos_ms = (time.perf_counter() - t0) * 1000.0 - data_write_ms = ( - data_reset_ms - + clear_forces_ms - + geom_overrides_ms - + randomization_ms - + set_dof_vel_ms - + set_dof_pos_ms - ) - t0 = time.perf_counter() if self._supports_position_actuator_gains and len(self._joint_dof_pos_indices) == int( self.num_actuators ): @@ -665,46 +628,10 @@ def set_state( else: ctrl = np.zeros((len(env_indices), self.num_actuators), dtype=self._np_dtype) data_slice.actuator_ctrls = np.ascontiguousarray(ctrl) - ctrl_ms = (time.perf_counter() - t0) * 1000.0 - t0 = time.perf_counter() self._model.forward_kinematic(data_slice) - forward_kinematic_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() self._refresh_link_pose_cache(env_indices, data_slice=data_slice) - refresh_cache_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() self._invalidate_link_velocity_cache() - invalidate_cache_ms = (time.perf_counter() - t0) * 1000.0 - - total_ms = (time.perf_counter() - total_t0) * 1000.0 - measured_ms = ( - qpos_convert_ms - + slice_ms - + data_write_ms - + ctrl_ms - + forward_kinematic_ms - + refresh_cache_ms - + invalidate_cache_ms - ) - self._last_set_state_timing_ms = { - "dr_reset_set_state_qpos_convert_ms": qpos_convert_ms, - "dr_reset_set_state_slice_ms": slice_ms, - "dr_reset_set_state_data_write_ms": data_write_ms, - "dr_reset_set_state_data_reset_ms": data_reset_ms, - "dr_reset_set_state_clear_forces_ms": clear_forces_ms, - "dr_reset_set_state_geom_overrides_ms": geom_overrides_ms, - "dr_reset_set_state_randomization_ms": randomization_ms, - "dr_reset_set_state_set_dof_vel_ms": set_dof_vel_ms, - "dr_reset_set_state_set_dof_pos_ms": set_dof_pos_ms, - "dr_reset_set_state_ctrl_ms": ctrl_ms, - "dr_reset_set_state_forward_kinematic_ms": forward_kinematic_ms, - "dr_reset_set_state_refresh_cache_ms": refresh_cache_ms, - "dr_reset_set_state_invalidate_cache_ms": invalidate_cache_ms, - "dr_reset_set_state_backend_internal_gap_ms": total_ms - measured_ms, - } def get_dr_capabilities(self) -> DomainRandomizationCapabilities: supported_reset_terms = { diff --git a/src/unilab/base/backend/mujoco/backend.py b/src/unilab/base/backend/mujoco/backend.py index 5ea6b76ee..d2c9e4013 100644 --- a/src/unilab/base/backend/mujoco/backend.py +++ b/src/unilab/base/backend/mujoco/backend.py @@ -319,7 +319,6 @@ def __init__( self._np_dtype = np_dtype if np_dtype is not None else get_global_dtype() self.backend_type = "mujoco" self._pending_xfrc_applied = np.zeros((num_envs, 6 * self._model.nbody), dtype=np.float64) - self._last_set_state_timing_ms: dict[str, float] = {} # Thread configuration. self._n_threads = min(num_envs, cpu_count() * 2) @@ -410,10 +409,6 @@ def _get_sensor_view(prefix, dim): self._tracked_linvel_b_all = _get_sensor_view("track_linvel_b", 3) self._tracked_angvel_b_all = _get_sensor_view("track_angvel_b", 3) - @property - def last_set_state_timing_ms(self) -> dict[str, float]: - return dict(self._last_set_state_timing_ms) - def _load_base_model(self) -> mujoco.MjModel: if isinstance(self._model_file, mujoco.MjModel): if self.add_body_sensors: @@ -883,71 +878,22 @@ def set_state( qvel: np.ndarray, randomization: ResetRandomizationPayload | None = None, ) -> None: - total_t0 = time.perf_counter() - self._last_set_state_timing_ms = {} if len(env_indices) == 0: return num_reset = len(env_indices) - - t0 = time.perf_counter() state_np = np.zeros((num_reset, self._physics_state.shape[1]), dtype=np.float64) - alloc_state_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() state_np[:, self._idx_qpos : self._idx_qpos + self.nq] = qpos - fill_qpos_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() state_np[:, self._idx_qvel : self._idx_qvel + self.nv] = qvel - fill_qvel_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - env_ids = np.asarray(env_indices, dtype=np.int32) - env_ids_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() - reset_randomization = self._translate_reset_randomization(randomization, num_reset) - randomization_ms = (time.perf_counter() - t0) * 1000.0 - t0 = time.perf_counter() state_out, sensor_np = self._pool.reset( # type: ignore[union-attr] - env_ids=env_ids, + env_ids=np.asarray(env_indices, dtype=np.int32), initial_state=state_np, - randomization=reset_randomization, + randomization=self._translate_reset_randomization(randomization, num_reset), ) - pool_reset_ms = (time.perf_counter() - t0) * 1000.0 - t0 = time.perf_counter() self._physics_state[env_indices] = state_out.astype(self._np_dtype) - physics_state_scatter_ms = (time.perf_counter() - t0) * 1000.0 - - t0 = time.perf_counter() self._sensor_data[env_indices] = sensor_np.astype(self._np_dtype) - sensor_scatter_ms = (time.perf_counter() - t0) * 1000.0 - - total_ms = (time.perf_counter() - total_t0) * 1000.0 - measured_ms = ( - alloc_state_ms - + fill_qpos_ms - + fill_qvel_ms - + env_ids_ms - + randomization_ms - + pool_reset_ms - + physics_state_scatter_ms - + sensor_scatter_ms - ) - self._last_set_state_timing_ms = { - "dr_reset_set_state_alloc_state_ms": alloc_state_ms, - "dr_reset_set_state_fill_qpos_ms": fill_qpos_ms, - "dr_reset_set_state_fill_qvel_ms": fill_qvel_ms, - "dr_reset_set_state_env_ids_ms": env_ids_ms, - "dr_reset_set_state_randomization_ms": randomization_ms, - "dr_reset_set_state_pool_reset_ms": pool_reset_ms, - "dr_reset_set_state_physics_state_scatter_ms": physics_state_scatter_ms, - "dr_reset_set_state_sensor_scatter_ms": sensor_scatter_ms, - "dr_reset_set_state_backend_internal_gap_ms": total_ms - measured_ms, - } def get_dr_capabilities(self) -> DomainRandomizationCapabilities: return DomainRandomizationCapabilities( diff --git a/src/unilab/base/np_env.py b/src/unilab/base/np_env.py index 018df89e8..65408d330 100644 --- a/src/unilab/base/np_env.py +++ b/src/unilab/base/np_env.py @@ -34,27 +34,6 @@ "dr_reset_plan_ms", "dr_reset_payload_filter_ms", "dr_reset_set_state_ms", - "dr_reset_set_state_alloc_state_ms", - "dr_reset_set_state_fill_qpos_ms", - "dr_reset_set_state_fill_qvel_ms", - "dr_reset_set_state_env_ids_ms", - "dr_reset_set_state_qpos_convert_ms", - "dr_reset_set_state_slice_ms", - "dr_reset_set_state_data_write_ms", - "dr_reset_set_state_data_reset_ms", - "dr_reset_set_state_clear_forces_ms", - "dr_reset_set_state_geom_overrides_ms", - "dr_reset_set_state_randomization_ms", - "dr_reset_set_state_set_dof_vel_ms", - "dr_reset_set_state_set_dof_pos_ms", - "dr_reset_set_state_ctrl_ms", - "dr_reset_set_state_forward_kinematic_ms", - "dr_reset_set_state_pool_reset_ms", - "dr_reset_set_state_physics_state_scatter_ms", - "dr_reset_set_state_sensor_scatter_ms", - "dr_reset_set_state_refresh_cache_ms", - "dr_reset_set_state_invalidate_cache_ms", - "dr_reset_set_state_backend_internal_gap_ms", "dr_reset_build_observation_ms", "dr_reset_internal_gap_ms", "dr_reset_observation_getters_ms", diff --git a/src/unilab/dr/manager.py b/src/unilab/dr/manager.py index f6469d6ca..6beddb281 100644 --- a/src/unilab/dr/manager.py +++ b/src/unilab/dr/manager.py @@ -71,7 +71,6 @@ def reset(self, env_ids: np.ndarray) -> tuple[dict[str, np.ndarray], dict]: "dr_reset_build_observation_ms": build_observation_ms, "dr_reset_internal_gap_ms": total_ms - measured_ms, } - timing.update(self._env._backend.last_set_state_timing_ms) provider_timing = getattr(self._provider, "last_reset_observation_timing_ms", {}) if isinstance(provider_timing, dict): timing.update(provider_timing) diff --git a/tests/benchmark/test_reset_done_attribution_benchmark.py b/tests/benchmark/test_reset_done_attribution_benchmark.py deleted file mode 100644 index 0aa776f8e..000000000 --- a/tests/benchmark/test_reset_done_attribution_benchmark.py +++ /dev/null @@ -1,41 +0,0 @@ -from __future__ import annotations - -import pytest -from benchmark import benchmark_reset_done_attribution as bench - - -def test_stats_reports_distribution() -> None: - stats = bench._stats([3.0, 1.0, 2.0]) - - assert stats.mean_ms == pytest.approx(2.0) - assert stats.median_ms == pytest.approx(2.0) - assert stats.min_ms == pytest.approx(1.0) - assert stats.max_ms == pytest.approx(3.0) - assert stats.samples_ms == [3.0, 1.0, 2.0] - - -def test_sample_env_ids_is_sorted_unique_and_seeded() -> None: - env_ids = bench._sample_env_ids(num_envs=16, reset_count=5, seed=673) - - assert env_ids.tolist() == sorted(env_ids.tolist()) - assert len(set(env_ids.tolist())) == 5 - assert env_ids.tolist() == bench._sample_env_ids(16, 5, 673).tolist() - - -def test_sample_env_ids_rejects_invalid_count() -> None: - with pytest.raises(ValueError, match="reset-count must be > 0"): - bench._sample_env_ids(num_envs=16, reset_count=0, seed=673) - - with pytest.raises(ValueError, match="reset-count must be <= num-envs"): - bench._sample_env_ids(num_envs=16, reset_count=17, seed=673) - - -def test_parse_args_defaults_to_issue_scale() -> None: - args = bench.parse_args([]) - - assert args.num_envs == 8192 - assert args.reset_count == 256 - assert args.measure_repeats == 30 - assert args.body_pose_backends == "motrix,mujoco" - assert args.set_state_profile_backends == "motrix,mujoco" - assert not args.skip_set_state diff --git a/tests/dr/test_manager.py b/tests/dr/test_manager.py index 1e3640fcd..c8150b30a 100644 --- a/tests/dr/test_manager.py +++ b/tests/dr/test_manager.py @@ -6,7 +6,6 @@ from typing import Any import numpy as np -import pytest from unilab.dr import ( DomainRandomizationCapabilities, @@ -123,10 +122,6 @@ class _FakeBackend: def __post_init__(self) -> None: self.last_randomization: ResetRandomizationPayload | None = None - @property - def last_set_state_timing_ms(self) -> dict[str, float]: - return {"dr_reset_set_state_forward_kinematic_ms": 1.25} - def get_dr_capabilities(self) -> DomainRandomizationCapabilities: return self.capabilities @@ -184,8 +179,6 @@ def test_manager_skips_unsupported_reset_terms_with_warning(caplog): "motrix backend does not support reset randomization terms: kp; skipping them." in caplog.text ) - timing = manager.last_reset_timing_ms - assert timing["dr_reset_set_state_forward_kinematic_ms"] == pytest.approx(1.25) def test_manager_keeps_supported_reset_terms_without_warning(caplog):