From 225a32ef170fd9db6291f72c13e5e7300b1689c7 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Mon, 26 Oct 2026 00:38:34 +0800 Subject: [PATCH 01/18] bench(collector): freeze Genesis lifecycle attribution baseline --- .../benchmark_offpolicy_collector_active.py | 95 ++++++++++++++++++- src/unilab/base/backend_timing.py | 5 + src/unilab/base/entity.py | 12 ++- src/unilab/envs/manager_based_rl_env.py | 15 +++ tests/base/test_entity_tensor_read_plan.py | 4 +- ...st_offpolicy_collector_active_benchmark.py | 47 +++++++++ 6 files changed, 175 insertions(+), 3 deletions(-) diff --git a/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py b/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py index 06b964970..e6c1002c9 100644 --- a/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py +++ b/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py @@ -100,6 +100,8 @@ "update_state_ms", "update_state_cpu_ms", "update_state_read_boundary_ms", + "update_state_backend_state_read_ms", + "update_state_backend_sensor_projection_ms", "update_state_prelude_ms", "update_state_queue_drain_ms", "update_state_termination_ms", @@ -226,6 +228,17 @@ "set_state_tensor_mask_ms", "set_state_tensor_commit_forward_ms", "set_state_tensor_host_cache_refresh_ms", + # Genesis device-resident internals. Times are auxiliary sub-timings of + # step_core/reset; *_count fields are event counts, not milliseconds. + "backend_tensor_control_upload_ms", + "backend_tensor_control_h2d_count", + "backend_tensor_physics_ms", + "backend_tensor_state_publication_ms", + "backend_tensor_state_d2d_copy_count", + "backend_tensor_state_sync_count", + "backend_tensor_reset_validation_ms", + "backend_tensor_reset_apply_ms", + "backend_tensor_reset_state_publication_ms", ) ENV_STEP_COUNT_KEYS = ( "reset_done_count", @@ -251,6 +264,11 @@ ("apply_action_ms", "apply_action_ms"), ("step_core_ms", "step_core_ms"), ("update_state_ms", "update_state_ms"), + ("update_state_backend_state_read_ms", "update_state_backend_state_read_ms"), + ( + "update_state_backend_sensor_projection_ms", + "update_state_backend_sensor_projection_ms", + ), ("reset_done_ms", "reset_done_ms"), ("env_step_internal_gap_ms", "env_step_internal_gap_ms"), ("reset_done_terminal_obs_ms", "reset_done_terminal_obs_ms"), @@ -279,6 +297,12 @@ ("reset_done_motion_reset_publish_ms", "reset_done_motion_reset_publish_ms"), ("reset_done_reset_commit_ms", "reset_done_reset_commit_ms"), ("reset_done_reset_validation_ms", "reset_done_reset_validation_ms"), + ("backend_tensor_reset_validation_ms", "backend_tensor_reset_validation_ms"), + ("backend_tensor_reset_apply_ms", "backend_tensor_reset_apply_ms"), + ( + "backend_tensor_reset_state_publication_ms", + "backend_tensor_reset_state_publication_ms", + ), ("reset_done_command_refresh_ms", "reset_done_command_refresh_ms"), ("reset_done_command_post_compute_ms", "reset_done_command_post_compute_ms"), ("reset_done_motion_robot_refresh_ms", "reset_done_motion_robot_refresh_ms"), @@ -384,6 +408,13 @@ class CollectorResult: phase_ms_per_vector_step: dict[str, TimingStats] phase_pct: dict[str, float] notes: list[str] + # Lifecycle phases are intentionally outside total_active_ms. Construction + # includes backend materialization and the initial reset; setup includes the + # first transition/JIT drain; teardown closes the collector-owned backend. + construction_ms: float = 0.0 + collector_setup_ms: float = 0.0 + warmup_ms: float = 0.0 + teardown_ms: float = 0.0 # Aggregate rate over the final measured steps. This mirrors the runner # tail contract while remaining explicitly diagnostic, not acceptance. tail_step_count: int = 0 @@ -409,6 +440,7 @@ class CollectorResult: backend_runtime_diagnostics: dict[str, dict[str, bool | str | None]] = field( default_factory=dict ) + backend_tensor_profile: dict[str, Any] = field(default_factory=dict) # Public bounded-ingress ordering evidence collected after the active window. replay_ingress_diagnostics: dict[str, int | float] = field(default_factory=dict) # Context-only actor forward time; actions used by the measured collector @@ -571,6 +603,37 @@ def _backend_runtime_diagnostics(env: Any) -> dict[str, dict[str, bool | str | N } +def _backend_tensor_profile(env: Any) -> dict[str, Any]: + """Record and fail closed on the negotiated public tensor profile.""" + + backend = getattr(env, "_backend", None) + if backend is None: + return {} + execution = backend.get_tensor_execution() + capabilities = backend.get_tensor_capabilities() + profile = { + "backend_type": str(backend.backend_type), + "execution": execution.value, + "data_plane": capabilities.data_plane.value, + "process_topology": capabilities.process_topology.value, + "torch_devices": tuple(capabilities.torch_devices), + } + if profile["backend_type"] == "genesis" and ( + execution.value != "device_resident" + or capabilities.data_plane.value != "direct" + or capabilities.process_topology.value != "in_process" + or not any(str(device).startswith("cuda") for device in capabilities.torch_devices) + ): + raise RuntimeError( + "Genesis production collector benchmark requires the UniSim adapter's " + f"exact-CUDA DEVICE_RESIDENT/DIRECT profile; negotiated {execution.value}/" + f"{capabilities.data_plane.value}. Genesis 1.3.3 natively supports metal, " + "amdgpu, and cuda engines, but this benchmark must not silently use " + "portable Genesis, CPU, or another backend." + ) + return profile + + def _runtime_sim_backend(sim: str) -> str: return sim @@ -859,11 +922,13 @@ def _run_active_window_case( tail_steps: int = DEFAULT_TAIL_STEPS, policy_inference_context: TimingStats | None = None, policy_inference_context_evidence: dict[str, Any] | None = None, + construction_ms: float = 0.0, ) -> CollectorResult: from unilab.rl.contracts.observations import split_obs_dict from unilab.rl.runtime.ipc.replay_buffer import ReplayBuffer from unilab.rl.runtime.offpolicy.worker import timeout_bootstrap_mask + setup_started = time.perf_counter() replay_buffer = ReplayBuffer( capacity=case.replay_capacity_rows, obs_dim=case.obs_dim, @@ -876,8 +941,10 @@ def _run_active_window_case( env_device = env.device backend_runtime_diagnostics = _backend_runtime_diagnostics(env) + backend_tensor_profile = _backend_tensor_profile(env) actions = torch.zeros((case.num_envs, case.action_dim), dtype=torch.float32, device=env_device) state = env.step(actions) + collector_setup_ms = (time.perf_counter() - setup_started) * 1000.0 obs, critic = split_obs_dict(state.obs) current_ep_rewards = torch.zeros(case.num_envs, dtype=torch.float32, device=env_device) current_ep_lengths = torch.zeros(case.num_envs, dtype=torch.int64, device=env_device) @@ -898,6 +965,7 @@ def _run_active_window_case( numpy_random_call_samples: list[float] = [] random_profiler = _NumpyRandomProfiler() if profile_numpy_random else None total_active_ns = 0 + warmup_active_ns = 0 total_steps = int(warmup_steps) + int(measure_steps) if total_steps <= 0 or measure_steps <= 0: raise ValueError("measure_steps must be > 0 and warmup_steps must be >= 0") @@ -939,6 +1007,8 @@ def _run_active_window_case( state.info.get("update_state_timing", {}) if isinstance(state.info, dict) else {} ) physics_ms = _optional_timing_ms(_timing, "backend_physics_ms") + if physics_ms is None: + physics_ms = _optional_timing_ms(_timing, "backend_tensor_physics_ms") env_step_timing_values = { key: _optional_timing_ms(_update_state_timing, key) if key in UPDATE_STATE_DETAIL_TIMING_KEYS @@ -1036,6 +1106,8 @@ def _run_active_window_case( for key, value in env_step_timing_values.items(): if value is not None: env_step_timing_samples.setdefault(key, []).append(value) + else: + warmup_active_ns += int((env_step_ms + replay_ms + bookkeeping_ms) * 1e6) finally: if random_profiler is not None: random_profiler.uninstall() @@ -1080,6 +1152,9 @@ def _run_active_window_case( phase_ms_per_vector_step=phase_stats, phase_pct=phase_pct, notes=[], + construction_ms=construction_ms, + collector_setup_ms=collector_setup_ms, + warmup_ms=warmup_active_ns / 1.0e6, env_step_timing_ms_per_vector_step=env_step_timing_stats, physics_ms_per_vector_step=physics_stats, env_step_overhead_ms_per_vector_step=env_step_overhead_stats, @@ -1091,6 +1166,7 @@ def _run_active_window_case( _stats(numpy_random_call_samples) if numpy_random_call_samples else None ), backend_runtime_diagnostics=backend_runtime_diagnostics, + backend_tensor_profile=backend_tensor_profile, replay_ingress_diagnostics=replay_ingress_diagnostics, policy_inference_context=policy_inference_context, policy_inference_context_evidence=policy_inference_context_evidence or {}, @@ -1137,11 +1213,15 @@ def _build_and_run_case( ).build_task_env_cfg_override() env = None + result: CollectorResult | None = None + construction_ms = 0.0 try: + construction_started = time.perf_counter() env, obs_dim, critic_dim, action_dim = _make_env( cfg, env_cfg_override=env_cfg_override, ) + construction_ms = (time.perf_counter() - construction_started) * 1000.0 policy_context, policy_context_evidence = _measure_policy_inference_context( cfg, obs_dim=obs_dim, @@ -1163,7 +1243,7 @@ def _build_and_run_case( critic_dim=critic_dim, action_dim=action_dim, ) - return _run_active_window_case( + result = _run_active_window_case( case, env=env, warmup_steps=warmup_steps, @@ -1172,10 +1252,15 @@ def _build_and_run_case( tail_steps=tail_steps, policy_inference_context=policy_context, policy_inference_context_evidence=policy_context_evidence, + construction_ms=construction_ms, ) + return result finally: if env is not None: + teardown_started = time.perf_counter() env.close() + if result is not None: + result.teardown_ms = (time.perf_counter() - teardown_started) * 1000.0 _cleanup() @@ -1275,6 +1360,10 @@ def _write_csv(path: Path, results: list[CollectorResult]) -> None: "tail_step_count", "tail_active_steps_per_sec", "total_active_ms", + "construction_ms", + "collector_setup_ms", + "warmup_ms", + "teardown_ms", "env_step_ms", "replay_ms", "bookkeeping_ms", @@ -1315,6 +1404,10 @@ def _write_csv(path: Path, results: list[CollectorResult]) -> None: else "" ), "total_active_ms": result.total_active_ms, + "construction_ms": result.construction_ms, + "collector_setup_ms": result.collector_setup_ms, + "warmup_ms": result.warmup_ms, + "teardown_ms": result.teardown_ms, } for key in COLLECTOR_PHASES: row[key] = result.phase_ms_per_vector_step.get( diff --git a/src/unilab/base/backend_timing.py b/src/unilab/base/backend_timing.py index c821412cb..097d518ae 100644 --- a/src/unilab/base/backend_timing.py +++ b/src/unilab/base/backend_timing.py @@ -91,6 +91,9 @@ "set_state_tensor_mask_ms", "set_state_tensor_commit_forward_ms", "set_state_tensor_host_cache_refresh_ms", + "backend_tensor_reset_validation_ms", + "backend_tensor_reset_apply_ms", + "backend_tensor_reset_state_publication_ms", "tensor_reset_packed_d2h_ms", "tensor_reset_packed_d2h_bytes", "tensor_reset_packed_d2h_count", @@ -122,6 +125,8 @@ "update_state_motion_failure_stats_ms", "update_state_motion_step_sampler_ms", "update_state_motion_refresh_current_ms", + "update_state_backend_state_read_ms", + "update_state_backend_sensor_projection_ms", ) BACKEND_SET_STATE_DETAIL_TIMING_KEYS = ( diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 854ca79f9..a76ad3fda 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -8,6 +8,7 @@ from __future__ import annotations import re +import time from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass @@ -301,6 +302,7 @@ def __init__( self._packet: Mapping[str, torch.Tensor] = MappingProxyType({}) self._refreshed = False self._closed = False + self._direct_read_timing: dict[str, float] = {} @property def scene(self) -> "EntityScene": @@ -620,7 +622,7 @@ def transfer_stats(self) -> dict[str, int]: def last_timing(self) -> dict[str, float]: self._require_open() if self._host_plan_value is None: - return {} + return dict(self._direct_read_timing) return dict(self._host_plan_value.last_timing) def close(self) -> None: @@ -658,14 +660,22 @@ def _read_packet( ) from exc else: try: + state_started = time.perf_counter() # Every sensor must cross its public backend boundary. Genesis, # Newton, and Isaac workers may project derived values into # stable adapter-owned storage here; only a fresh call proves # that storage reflects the current post-step/post-reset state. raw = dict(backend.get_state_views(("qpos", "qvel"), device=self._device)) + state_ms = (time.perf_counter() - state_started) * 1000.0 + sensor_started = time.perf_counter() for name in self._packet_names: view = backend.get_sensor_view(name, device=self._device) raw[name] = view + sensor_ms = (time.perf_counter() - sensor_started) * 1000.0 + self._direct_read_timing = { + "state_read_ms": state_ms, + "sensor_projection_ms": sensor_ms, + } except (AttributeError, KeyError, TypeError, ValueError, NotImplementedError) as exc: raise type(exc)( "Manager scene device-resident tensor read on backend " diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 9bd1ae60b..4edb06410 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -1206,6 +1206,14 @@ def update_state(self, state: TorchEnvState) -> TorchEnvState: read_plan.refresh() boundary_ms = (time.perf_counter() - boundary_started) * 1000.0 timing["update_state_read_boundary_ms"] = boundary_ms + if read_plan is not None: + read_timing = read_plan.last_timing + timing["update_state_backend_state_read_ms"] = read_timing.get( + "state_read_ms", 0.0 + ) + timing["update_state_backend_sensor_projection_ms"] = read_timing.get( + "sensor_projection_ms", 0.0 + ) return self._update_state_in_read_phase(state) def _update_state_in_read_phase(self, state: TorchEnvState) -> TorchEnvState: @@ -1466,6 +1474,13 @@ def _compute_truncated(self, state: TorchEnvState) -> torch.Tensor: def _collect_reset_backend_timing_ms(self) -> dict[str, float]: timing = dict(super()._collect_reset_backend_timing_ms()) timing.update(self._reset_state.last_set_state_timing_ms) + for source, target in ( + ("tensor_reset_validation_ms", "backend_tensor_reset_validation_ms"), + ("tensor_reset_apply_ms", "backend_tensor_reset_apply_ms"), + ("tensor_state_publication_ms", "backend_tensor_reset_state_publication_ms"), + ): + if source in timing: + timing[target] = timing[source] timing.update(self._last_reset_manager_timing_ms) return timing diff --git a/tests/base/test_entity_tensor_read_plan.py b/tests/base/test_entity_tensor_read_plan.py index 550b87742..3375d5162 100644 --- a/tests/base/test_entity_tensor_read_plan.py +++ b/tests/base/test_entity_tensor_read_plan.py @@ -333,7 +333,9 @@ def test_device_resident_scene_plan_refreshes_every_sensor_boundary() -> None: # Genesis/Newton derived sensors are stable outputs, not native-live arrays. assert backend.sensor_view_calls == calls_after_first_refresh + len(plan.packet_names) assert plan.transfer_stats == {} - assert plan.last_timing == {} + assert set(plan.last_timing) == {"state_read_ms", "sensor_projection_ms"} + assert plan.last_timing["state_read_ms"] >= 0.0 + assert plan.last_timing["sensor_projection_ms"] >= 0.0 plan.close() diff --git a/tests/benchmark/test_offpolicy_collector_active_benchmark.py b/tests/benchmark/test_offpolicy_collector_active_benchmark.py index d6cd16a64..d11edd0aa 100644 --- a/tests/benchmark/test_offpolicy_collector_active_benchmark.py +++ b/tests/benchmark/test_offpolicy_collector_active_benchmark.py @@ -124,6 +124,48 @@ def test_cuda_tensor_backends_are_opt_in() -> None: assert bench._resolve_backend_selection(backend="mjwarp", all_backends=False) == ("mjwarp",) +def test_genesis_benchmark_fails_closed_on_non_exact_tensor_profile() -> None: + class UnsupportedBackend: + backend_type = "genesis" + + def get_tensor_execution(self): + return SimpleNamespace(value="unsupported") + + def get_tensor_capabilities(self): + return SimpleNamespace( + data_plane=SimpleNamespace(value="direct"), + process_topology=SimpleNamespace(value="in_process"), + torch_devices=("cpu",), + ) + + with pytest.raises(RuntimeError, match="exact-CUDA DEVICE_RESIDENT/DIRECT"): + bench._backend_tensor_profile(SimpleNamespace(_backend=UnsupportedBackend())) + + +def test_collector_result_records_lifecycle_and_tensor_profile() -> None: + result = _make_result() + result.construction_ms = 11.0 + result.collector_setup_ms = 12.0 + result.warmup_ms = 13.0 + result.teardown_ms = 14.0 + result.backend_tensor_profile = { + "backend_type": "genesis", + "execution": "device_resident", + "data_plane": "direct", + "process_topology": "in_process", + "torch_devices": ("cuda",), + } + + payload = bench._result_to_dict(result) + + assert payload["construction_ms"] == 11.0 + assert payload["collector_setup_ms"] == 12.0 + assert payload["warmup_ms"] == 13.0 + assert payload["teardown_ms"] == 14.0 + assert payload["backend_tensor_profile"]["execution"] == "device_resident" + assert payload["backend_tensor_profile"]["data_plane"] == "direct" + + def test_resolve_case_specs_deduplicates_explicit_specs() -> None: specs = bench._resolve_case_specs( "sac/g1_walk_flat/mujoco,sac/g1_walk_flat/mujoco,flashsac/g1_walk_flat/mujoco", @@ -313,6 +355,11 @@ def test_write_csv_includes_all_phase_columns(tmp_path) -> None: header = out_csv.read_text(encoding="utf-8").splitlines()[0] assert "bookkeeping_ms" in header + assert "construction_ms" in header + assert "collector_setup_ms" in header + assert "warmup_ms" in header + assert "teardown_ms" in header + assert "backend_tensor_reset_state_publication_ms" in header assert "env_step_overhead_ms" in header assert "physics_pct" in header assert "env_step_overhead_pct" in header From 2a829bb9db59e38f567d8128f2b8594029eef2be Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Mon, 26 Oct 2026 01:04:19 +0800 Subject: [PATCH 02/18] bench(genesis): freeze production collector schema --- .../benchmark_offpolicy_collector_active.py | 268 +++++++++++++++--- src/unilab/base/backend_factory.py | 2 + src/unilab/base/backend_timing.py | 3 + src/unilab/base/base.py | 8 + src/unilab/base/entity.py | 51 +++- src/unilab/base/torch_env.py | 14 +- .../task/g1_motion_tracking/base.yaml | 3 + src/unilab/envs/manager_based_rl_env.py | 19 +- src/unilab/rl/runtime/collector_benchmark.py | 121 ++++++++ tests/base/test_entity_tensor_read_plan.py | 24 +- tests/base/test_genesis_backend.py | 6 +- tests/base/test_torch_env.py | 6 + ...st_offpolicy_collector_active_benchmark.py | 128 ++++++++- 13 files changed, 584 insertions(+), 69 deletions(-) create mode 100644 src/unilab/rl/runtime/collector_benchmark.py diff --git a/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py b/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py index e6c1002c9..3c78cbceb 100644 --- a/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py +++ b/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py @@ -42,7 +42,7 @@ from collections import defaultdict, deque from dataclasses import asdict, dataclass, field from datetime import datetime, timezone -from importlib.metadata import PackageNotFoundError, version +from importlib.metadata import PackageNotFoundError, distribution, version from importlib.util import find_spec from pathlib import Path from statistics import mean, median, pstdev @@ -59,6 +59,7 @@ sys.path.append(str(ROOT_DIR)) from unilab.base.backend_timing import UPDATE_STATE_DETAIL_TIMING_KEYS +from unilab.rl.runtime.collector_benchmark import SCHEMA_ID, validate_envelope from unilab.utils.device import get_device_info_dict, get_device_info_line DEFAULT_OUTPUT_JSON = ( @@ -74,6 +75,11 @@ ) DEFAULT_ALGOS = ("sac", "flashsac") DEFAULT_BACKEND = "mujoco" +PRODUCTION_CASE = "flashsac/g1_motion_tracking/genesis" +PRODUCTION_SEED = 457 +PRODUCTION_NUM_ENVS = 2048 +PRODUCTION_SIM_DT = 0.006666666666666667 +PRODUCTION_CTRL_DT = 0.02 # CPU and CUDA tensor profiles have different dependency footprints. CPU host # bridge is default-safe; CUDA tensor backends are explicitly requested. BENCHMARK_BACKENDS = ("mujoco",) @@ -86,6 +92,7 @@ "torch", "mujoco-warp", "warp-lang", + "genesis-world", ) COLLECTOR_PHASES = ( "env_step_ms", @@ -101,7 +108,10 @@ "update_state_cpu_ms", "update_state_read_boundary_ms", "update_state_backend_state_read_ms", + "update_state_backend_tracked_sensor_projection_ms", + "update_state_backend_named_sensor_projection_ms", "update_state_backend_sensor_projection_ms", + "update_state_backend_sensor_projection_call_count", "update_state_prelude_ms", "update_state_queue_drain_ms", "update_state_termination_ms", @@ -231,10 +241,12 @@ # Genesis device-resident internals. Times are auxiliary sub-timings of # step_core/reset; *_count fields are event counts, not milliseconds. "backend_tensor_control_upload_ms", - "backend_tensor_control_h2d_count", + "backend_tensor_control_metadata_h2d_count", + "backend_tensor_control_metadata_h2d_bytes", "backend_tensor_physics_ms", "backend_tensor_state_publication_ms", - "backend_tensor_state_d2d_copy_count", + "backend_tensor_state_getter_copy_count", + "backend_tensor_state_mirror_copy_count", "backend_tensor_state_sync_count", "backend_tensor_reset_validation_ms", "backend_tensor_reset_apply_ms", @@ -265,6 +277,14 @@ ("step_core_ms", "step_core_ms"), ("update_state_ms", "update_state_ms"), ("update_state_backend_state_read_ms", "update_state_backend_state_read_ms"), + ( + "update_state_backend_tracked_sensor_projection_ms", + "update_state_backend_tracked_sensor_projection_ms", + ), + ( + "update_state_backend_named_sensor_projection_ms", + "update_state_backend_named_sensor_projection_ms", + ), ( "update_state_backend_sensor_projection_ms", "update_state_backend_sensor_projection_ms", @@ -373,6 +393,8 @@ class CollectorCase: critic_dim: int action_dim: int env_steps_per_sync: int + case_profile: str = "generic" + num_envs_source: str = "cli" # Exact environment shape and term-set fixation for Gate UP evidence. seed: int = 0 env_seed: int | None = None @@ -396,6 +418,8 @@ class TimingStats: std_ms: float min_ms: float max_ms: float + p50_ms: float = 0.0 + p95_ms: float = 0.0 @dataclass @@ -419,6 +443,9 @@ class CollectorResult: # tail contract while remaining explicitly diagnostic, not acceptance. tail_step_count: int = 0 tail_active_steps_per_sec: float | None = None + reset_rows_total: int = 0 + reset_vector_step_count: int = 0 + reset_rows_per_vector_step: TimingStats | None = None # Fine-grained timings reported by `TorchEnv.step()` inside env_step_ms. env_step_timing_ms_per_vector_step: dict[str, TimingStats] = field(default_factory=dict) # Backend-internal physics time per vector step (sub-part of env_step_ms). @@ -459,6 +486,8 @@ def _stats(samples_ms: list[float]) -> TimingStats: std_ms=pstdev(samples_ms) if len(samples_ms) > 1 else 0.0, min_ms=min(samples_ms), max_ms=max(samples_ms), + p50_ms=median(samples_ms), + p95_ms=float(np.percentile(samples_ms, 95)), ) @@ -589,48 +618,62 @@ def _configure_collector_cpu_threads(cap: int | None = None) -> int: def _backend_runtime_diagnostics(env: Any) -> dict[str, dict[str, bool | str | None]]: - backend = getattr(env, "_backend", None) - diagnostics = getattr(backend, "get_tensor_runtime_diagnostics", None) - if diagnostics is None: + get_diagnostics = getattr(env, "get_backend_runtime_diagnostics", None) + if get_diagnostics is None: + return {} + payload = get_diagnostics() + diagnostics = payload.get("diagnostics", {}) + if not isinstance(diagnostics, dict): return {} return { name: { - "requested": bool(diagnostic.requested), - "enabled": bool(diagnostic.enabled), - "disable_reason": diagnostic.disable_reason, + "requested": bool(diagnostic.get("requested")), + "enabled": bool(diagnostic.get("enabled")), + "disable_reason": diagnostic.get("disable_reason"), } - for name, diagnostic in diagnostics().items() + for name, diagnostic in diagnostics.items() + if isinstance(diagnostic, dict) } -def _backend_tensor_profile(env: Any) -> dict[str, Any]: +def _backend_tensor_profile( + env: Any, *, require_attribution: bool = False +) -> dict[str, Any]: """Record and fail closed on the negotiated public tensor profile.""" - backend = getattr(env, "_backend", None) - if backend is None: - return {} - execution = backend.get_tensor_execution() - capabilities = backend.get_tensor_capabilities() + get_diagnostics = getattr(env, "get_backend_runtime_diagnostics", None) + if get_diagnostics is None: + raise RuntimeError("collector benchmark environment has no public backend diagnostics") + payload = get_diagnostics() profile = { - "backend_type": str(backend.backend_type), - "execution": execution.value, - "data_plane": capabilities.data_plane.value, - "process_topology": capabilities.process_topology.value, - "torch_devices": tuple(capabilities.torch_devices), + "backend_type": str(payload.get("backend_type", "")), + **dict(payload.get("tensor_profile", {})), } + diagnostics = _backend_runtime_diagnostics(env) if profile["backend_type"] == "genesis" and ( - execution.value != "device_resident" - or capabilities.data_plane.value != "direct" - or capabilities.process_topology.value != "in_process" - or not any(str(device).startswith("cuda") for device in capabilities.torch_devices) + profile.get("execution") != "device_resident" + or profile.get("data_plane") != "direct" + or profile.get("process_topology") != "in_process" + or not any(str(device).startswith("cuda") for device in profile.get("torch_devices", ())) ): raise RuntimeError( "Genesis production collector benchmark requires the UniSim adapter's " - f"exact-CUDA DEVICE_RESIDENT/DIRECT profile; negotiated {execution.value}/" - f"{capabilities.data_plane.value}. Genesis 1.3.3 natively supports metal, " + "exact-CUDA DEVICE_RESIDENT/DIRECT profile; negotiated " + f"{profile.get('execution')}/{profile.get('data_plane')}. " + "Genesis 1.3.3 natively supports metal, " "amdgpu, and cuda engines, but this benchmark must not silently use " "portable Genesis, CPU, or another backend." ) + exact_cuda = diagnostics.get("exact_cuda_profile", {}) + if profile["backend_type"] == "genesis" and not bool(exact_cuda.get("enabled", False)): + raise RuntimeError( + "Genesis production collector benchmark did not activate exact-CUDA zero-copy; " + f"diagnostic: {exact_cuda.get('disable_reason') or 'missing diagnostic'}" + ) + if require_attribution and not bool( + diagnostics.get("attribution_diagnostics", {}).get("enabled", False) + ): + raise RuntimeError("Genesis production collector benchmark requires attribution diagnostics") return profile @@ -671,6 +714,31 @@ def _compose_offpolicy_cfg( return compose(config_name="config", overrides=overrides) +def _validate_production_cfg(cfg: DictConfig, *, num_envs_source: str) -> None: + expected = { + "algo.algo": "flashsac", + "training.sim_backend": "genesis", + "algo.num_envs": PRODUCTION_NUM_ENVS, + "env.sim_dt": PRODUCTION_SIM_DT, + "env.ctrl_dt": PRODUCTION_CTRL_DT, + "algo.seed": PRODUCTION_SEED, + "env.seed": PRODUCTION_SEED, + "env.genesis_attribution_diagnostics": True, + } + actual = {name: OmegaConf.select(cfg, name) for name in expected} + mismatches = [name for name, value in expected.items() if actual[name] != value] + if mismatches: + details = ", ".join( + f"{name}={actual[name]!r} (expected {expected[name]!r})" for name in mismatches + ) + raise ValueError(f"production Genesis collector profile mismatch: {details}") + if num_envs_source != "owner": + raise ValueError( + "production Genesis collector profile requires its 2048-env owner value; " + "do not pass --num-envs" + ) + + def _owner_config_path(algo: str, task: str, sim: str, profile: str | None = None) -> Path: return ( ROOT_DIR @@ -751,6 +819,8 @@ def _resolve_case_specs( algos_arg: str, backends: Sequence[str], ) -> list[str]: + if cases_arg == "production": + return [PRODUCTION_CASE] if cases_arg == "default": return _default_case_specs(backends) if cases_arg == "auto": @@ -781,6 +851,7 @@ def _build_case( obs_dim: int, critic_dim: int, action_dim: int, + production: bool = False, ) -> CollectorCase: num_envs = int(cfg.algo.num_envs) capacity_steps = max(1, int(replay_capacity_steps)) @@ -806,6 +877,8 @@ def _build_case( critic_dim=int(critic_dim), action_dim=int(action_dim), env_steps_per_sync=int(cfg.training.env_steps_per_sync), + case_profile="production" if production else "generic", + num_envs_source="owner" if production else "cli", seed=int(cfg.algo.seed), env_seed=( int(cfg.env.seed) @@ -923,6 +996,7 @@ def _run_active_window_case( policy_inference_context: TimingStats | None = None, policy_inference_context_evidence: dict[str, Any] | None = None, construction_ms: float = 0.0, + production: bool = False, ) -> CollectorResult: from unilab.rl.contracts.observations import split_obs_dict from unilab.rl.runtime.ipc.replay_buffer import ReplayBuffer @@ -941,7 +1015,7 @@ def _run_active_window_case( env_device = env.device backend_runtime_diagnostics = _backend_runtime_diagnostics(env) - backend_tensor_profile = _backend_tensor_profile(env) + backend_tensor_profile = _backend_tensor_profile(env, require_attribution=production) actions = torch.zeros((case.num_envs, case.action_dim), dtype=torch.float32, device=env_device) state = env.step(actions) collector_setup_ms = (time.perf_counter() - setup_started) * 1000.0 @@ -963,6 +1037,7 @@ def _run_active_window_case( env_step_timing_samples: dict[str, list[float]] = {key: [] for key in ENV_STEP_SAMPLE_KEYS} numpy_random_ms_samples: list[float] = [] numpy_random_call_samples: list[float] = [] + reset_rows_samples: list[float] = [] random_profiler = _NumpyRandomProfiler() if profile_numpy_random else None total_active_ns = 0 warmup_active_ns = 0 @@ -1072,6 +1147,7 @@ def _run_active_window_case( current_ep_rewards += rewards current_ep_lengths += 1 reset_indices = terminal_mask.nonzero(as_tuple=False).flatten() + reset_rows_samples.append(float(reset_indices.numel())) if reset_indices.numel() > 0: ep_rewards.extend(current_ep_rewards[reset_indices].detach().cpu().tolist()) ep_lengths.extend(current_ep_lengths[reset_indices].detach().cpu().tolist()) @@ -1131,9 +1207,11 @@ def _run_active_window_case( for key, stat in phase_stats.items() } + teardown_started = time.perf_counter() replay_ingress_diagnostics = replay_buffer.ingress_diagnostics() del replay_buffer _cleanup() + collector_teardown_ms = (time.perf_counter() - teardown_started) * 1000.0 physics_stats = _stats(aux_samples["physics_ms"]) if aux_samples["physics_ms"] else None env_step_overhead_stats = ( _stats(aux_samples["env_step_overhead_ms"]) if aux_samples["env_step_overhead_ms"] else None @@ -1149,12 +1227,18 @@ def _run_active_window_case( collector_active_steps_per_sec=collector_active_steps_per_sec, tail_step_count=tail_count, tail_active_steps_per_sec=tail_rate, + reset_rows_total=int(sum(reset_rows_samples)), + reset_vector_step_count=sum(value > 0.0 for value in reset_rows_samples), + reset_rows_per_vector_step=( + _stats(reset_rows_samples) if reset_rows_samples else None + ), phase_ms_per_vector_step=phase_stats, phase_pct=phase_pct, notes=[], construction_ms=construction_ms, collector_setup_ms=collector_setup_ms, warmup_ms=warmup_active_ns / 1.0e6, + teardown_ms=collector_teardown_ms, env_step_timing_ms_per_vector_step=env_step_timing_stats, physics_ms_per_vector_step=physics_stats, env_step_overhead_ms_per_vector_step=env_step_overhead_stats, @@ -1184,6 +1268,7 @@ def _build_and_run_case( extra_overrides: list[str], variant: str = "default", profile_numpy_random: bool = False, + production: bool = False, ) -> CollectorResult: from unilab.base.config_adapter import EnvConfigAdapter from unilab.training import ensure_registries @@ -1202,6 +1287,8 @@ def _build_and_run_case( num_envs=num_envs, extra_overrides=extra_overrides, ) + if production: + _validate_production_cfg(cfg, num_envs_source="owner") ensure_registries() # Mirror the real collector subprocess CPU budget for env/replay work. _configure_collector_cpu_threads() @@ -1242,6 +1329,7 @@ def _build_and_run_case( obs_dim=obs_dim, critic_dim=critic_dim, action_dim=action_dim, + production=production, ) result = _run_active_window_case( case, @@ -1253,15 +1341,18 @@ def _build_and_run_case( policy_inference_context=policy_context, policy_inference_context_evidence=policy_context_evidence, construction_ms=construction_ms, + production=production, ) return result finally: if env is not None: teardown_started = time.perf_counter() env.close() + _cleanup() if result is not None: - result.teardown_ms = (time.perf_counter() - teardown_started) * 1000.0 - _cleanup() + result.teardown_ms += (time.perf_counter() - teardown_started) * 1000.0 + else: + _cleanup() def _result_to_dict(result: CollectorResult) -> dict[str, Any]: @@ -1287,9 +1378,9 @@ def _package_versions() -> dict[str, str]: return values -def _git_fixation() -> dict[str, str | bool | None]: +def _git_fixation(repository: Path) -> dict[str, str | bool | None]: def git(*arguments: str) -> str: - return _run_text_command(["git", "-C", str(ROOT_DIR), *arguments]) + return _run_text_command(["git", "-C", str(repository), *arguments]) return { "branch": git("branch", "--show-current") or None, @@ -1298,19 +1389,58 @@ def git(*arguments: str) -> str: } +def _distribution_source(name: str) -> dict[str, Any]: + try: + raw = distribution(name).read_text("direct_url.json") + except PackageNotFoundError: + return {"installed": False} + if raw is None: + return {"installed": True, "direct_url": None} + try: + return {"installed": True, "direct_url": json.loads(raw)} + except json.JSONDecodeError: + return {"installed": True, "direct_url": {"invalid_json": raw}} + + def _runtime_fixation() -> dict[str, Any]: + unisim_repository = ROOT_DIR.parent / "unisim" fixation: dict[str, Any] = { "python_version": platform.python_version(), "platform": platform.platform(), "host": platform.node(), "kernel": platform.release(), "package_versions": _package_versions(), - "git": _git_fixation(), + "git": { + "unilab": _git_fixation(ROOT_DIR), + "unisim": ( + _git_fixation(unisim_repository) + if (unisim_repository / ".git").exists() + else None + ), + }, + "unisim_source": _distribution_source("unisim-core"), "torch_cuda": torch.version.cuda, "torch_cudnn": torch.backends.cudnn.version() if torch.backends.cudnn.is_available() else None, "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "environment": { + name: os.environ.get(name) + for name in ( + "CUDA_VISIBLE_DEVICES", + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "GS_PARA_LEVEL", + "GS_ENABLE_ZEROCOPY", + "GS_ENABLE_NDARRAY", + "GS_TORCH_FORCE_CPU_DEVICE", + "QD_ENABLE_CPU", + "QD_ENABLE_CUDA", + "QD_ENABLE_AMDGPU", + "QD_ENABLE_METAL", + COLLECTOR_CPU_THREADS_ENV, + ) + }, } if torch.cuda.is_available(): index = torch.cuda.current_device() @@ -1372,6 +1502,10 @@ def _write_csv(path: Path, results: list[CollectorResult]) -> None: "physics_ms", "env_step_overhead_ms", "reset_done_count", + "reset_rows_total", + "reset_vector_step_count", + "reset_rows_per_vector_step_p50", + "reset_rows_per_vector_step_p95", *(field_name for _, field_name in ENV_STEP_TIMING_CSV_FIELDS), "env_step_pct", "replay_pct", @@ -1437,6 +1571,15 @@ def _write_csv(path: Path, results: list[CollectorResult]) -> None: row["reset_done_count"] = ( reset_done_count.mean_ms if reset_done_count is not None else "" ) + row["reset_rows_total"] = result.reset_rows_total + row["reset_vector_step_count"] = result.reset_vector_step_count + reset_rows = result.reset_rows_per_vector_step + row["reset_rows_per_vector_step_p50"] = ( + reset_rows.p50_ms if reset_rows is not None else "" + ) + row["reset_rows_per_vector_step_p95"] = ( + reset_rows.p95_ms if reset_rows is not None else "" + ) for timing_key, field_name in ENV_STEP_TIMING_CSV_FIELDS: stat = result.env_step_timing_ms_per_vector_step.get(timing_key) row[field_name] = stat.mean_ms if stat is not None else "" @@ -2286,8 +2429,11 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.add_argument( "--num-envs", type=int, - default=DEFAULT_NUM_ENVS, - help=f"Override algo.num_envs. Default: {DEFAULT_NUM_ENVS}.", + default=None, + help=( + "Override algo.num_envs. Generic cases default to " + f"{DEFAULT_NUM_ENVS}; --cases production always uses its frozen owner value." + ), ) parser.add_argument( "--warmup-steps", @@ -2361,6 +2507,15 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: ), ) parser.add_argument("--continue-on-error", action="store_true") + parser.add_argument( + "--seed", + type=int, + default=None, + help=( + "Set both algo.seed and env.seed. Generic cases default to each owner; " + f"--cases production is frozen at {PRODUCTION_SEED}." + ), + ) args = parser.parse_args(argv) if args.repeat_variant_blocks < 1: raise ValueError("--repeat-variant-blocks must be at least 1") @@ -2386,6 +2541,24 @@ def main() -> int: if spec_sim in OPTIONAL_BACKENDS: _check_optional_backend_deps(spec_sim) + production = args.cases == "production" and specs == [PRODUCTION_CASE] + if production and args.num_envs is not None: + raise SystemExit( + "--cases production uses the frozen owner profile (2048 envs); do not pass --num-envs" + ) + seed = args.seed + if production and seed is None: + seed = PRODUCTION_SEED + if production and seed != PRODUCTION_SEED: + raise SystemExit(f"--cases production requires seed {PRODUCTION_SEED}") + effective_num_envs = DEFAULT_NUM_ENVS if args.num_envs is None else int(args.num_envs) + if production: + effective_num_envs = None + seed_overrides = [] if seed is None else [f"algo.seed={seed}", f"env.seed={seed}"] + production_overrides = ( + ["env.genesis_attribution_diagnostics=true"] if production else [] + ) + hardware_info = _get_benchmark_hardware_info() print(f"Device: {get_device_info_line()}") print(f"Cases: {', '.join(specs)}") @@ -2411,10 +2584,16 @@ def main() -> int: measure_steps=int(args.measure_steps), replay_capacity_steps=int(args.replay_capacity_steps), tail_steps=int(args.tail_steps), - num_envs=args.num_envs, - extra_overrides=[*args.override, *variant_extra_overrides], + num_envs=effective_num_envs, + extra_overrides=[ + *seed_overrides, + *production_overrides, + *args.override, + *variant_extra_overrides, + ], variant=variant, profile_numpy_random=bool(args.profile_numpy_random), + production=production, ) ] results.extend(case_results) @@ -2428,6 +2607,7 @@ def main() -> int: raise payload = { + "schema": SCHEMA_ID, "created_at": datetime.now(timezone.utc).isoformat(), "benchmark": "offpolicy_collector_active", "definition": ( @@ -2436,6 +2616,17 @@ def main() -> int: "terminal-contract handling, replay writes, and collector-side bookkeeping." ), "device": hardware_info, + "phase_relationships": { + "additive_active": list(COLLECTOR_PHASES), + "lifecycle_excluded_from_active": [ + "construction_ms", + "collector_setup_ms", + "warmup_ms", + "teardown_ms", + ], + "tail_subset_of_active": True, + "attribution_children_non_additive": True, + }, "command": ("uv run " + shlex.join([str(Path(__file__).resolve()), *sys.argv[1:]])), "args": { "cases": args.cases, @@ -2443,6 +2634,8 @@ def main() -> int: "backend": args.backend, "all_backends": args.all_backends, "num_envs": args.num_envs, + "effective_num_envs": effective_num_envs, + "seed": seed, "warmup_steps": args.warmup_steps, "measure_steps": args.measure_steps, "tail_steps": args.tail_steps, @@ -2457,6 +2650,7 @@ def main() -> int: "results": [_result_to_dict(result) for result in results], "errors": errors, } + validate_envelope(payload) _write_json(args.out_json, payload) print(f"Wrote JSON: {args.out_json}") if args.out_csv is not None: diff --git a/src/unilab/base/backend_factory.py b/src/unilab/base/backend_factory.py index f7655c367..ab5ff80c6 100644 --- a/src/unilab/base/backend_factory.py +++ b/src/unilab/base/backend_factory.py @@ -262,6 +262,8 @@ def env_backend_kwargs(cfg: "EnvCfg") -> dict[str, Any]: # its own device. if cfg.genesis_device_id is not None: result["genesis_device_id"] = cfg.genesis_device_id + if cfg.genesis_attribution_diagnostics: + result["attribution_diagnostics"] = True result["newton_use_cuda_graph"] = cfg.newton_use_cuda_graph return result diff --git a/src/unilab/base/backend_timing.py b/src/unilab/base/backend_timing.py index 097d518ae..858a2a0f8 100644 --- a/src/unilab/base/backend_timing.py +++ b/src/unilab/base/backend_timing.py @@ -126,7 +126,10 @@ "update_state_motion_step_sampler_ms", "update_state_motion_refresh_current_ms", "update_state_backend_state_read_ms", + "update_state_backend_tracked_sensor_projection_ms", + "update_state_backend_named_sensor_projection_ms", "update_state_backend_sensor_projection_ms", + "update_state_backend_sensor_projection_call_count", ) BACKEND_SET_STATE_DETAIL_TIMING_KEYS = ( diff --git a/src/unilab/base/base.py b/src/unilab/base/base.py index bf9f779ce..1a253fd4f 100644 --- a/src/unilab/base/base.py +++ b/src/unilab/base/base.py @@ -87,6 +87,9 @@ class EnvCfg: # must be selected before ``gs.init`` so each data-parallel rank gets its # own simulator device; ``None`` keeps Genesis' current/default device. genesis_device_id: Optional[int] = None + # Opt in to Genesis' adapter-owned collector attribution counters. The + # default path reports only the existing step/reset timing keys. + genesis_attribution_diagnostics: bool = False # ``genesis`` drops the MJCF global