From 65546bbe45af515262939af19b463f3f2fdf2afc Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Mon, 24 Aug 2026 16:38:34 +0800 Subject: [PATCH] chore: remove temporary MBA profiling scaffolding (#1281) --- .../benchmark_offpolicy_collector_active.py | 128 +----------------- src/unilab/base/entity.py | 58 ++------ src/unilab/base/np_env.py | 88 ------------ src/unilab/base/reset_state.py | 18 +-- src/unilab/envs/manager_based_rl_env.py | 113 ---------------- src/unilab/managers/_types.py | 4 - src/unilab/managers/event_manager.py | 28 ---- src/unilab/managers/observation_manager.py | 54 -------- src/unilab/managers/reward_manager.py | 16 --- ...st_offpolicy_collector_active_benchmark.py | 101 +------------- .../locomotion/test_manager_gait_terms.py | 2 - tests/envs/mdp/test_observations.py | 1 - tests/envs/mdp/test_rewards.py | 1 - tests/envs/test_manager_based_rl_env.py | 80 +---------- tests/managers/conftest.py | 4 - .../test_event_command_metrics_recorder.py | 24 ---- 16 files changed, 20 insertions(+), 700 deletions(-) diff --git a/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py b/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py index b221b7467..80b351b4c 100644 --- a/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py +++ b/scripts/benchmark/rl/benchmark_offpolicy_collector_active.py @@ -101,19 +101,6 @@ "dr_reset_obs_get_body_pose_ms", "dr_reset_observation_compute_obs_ms", "dr_reset_observation_internal_gap_ms", - # Manager-Based (MBA) reset sub-steps (see RESET_DONE_DETAIL_TIMING_KEYS in - # src/unilab/base/np_env.py). Task-dependent per-term event timings with the - # mba_reset_event_term__ms prefix are discovered dynamically at sample - # time (JSON/console only; not part of the fixed CSV header). - "mba_reset_total_ms", - "mba_reset_curriculum_ms", - "mba_reset_event_apply_ms", - "mba_reset_command_reset_ms", - "mba_reset_set_state_ms", - "mba_reset_manager_reset_ms", - "mba_reset_command_compute_ms", - "mba_reset_obs_build_ms", - "mba_reset_internal_gap_ms", # Backend set_state internals (see BACKEND_SET_STATE_DETAIL_TIMING_KEYS in # src/unilab/base/np_env.py). All backends emit the same key set for column # stability; sub-keys that don't apply report 0.0. @@ -133,42 +120,9 @@ "set_state_pool_reset_ms", "set_state_state_scatter_ms", "set_state_internal_gap_ms", - # Manager-Based (MBA) update_state blocks, reported every step by - # ManagerBasedRlEnv.update_state() (see UPDATE_STATE_DETAIL_TIMING_KEYS in - # src/unilab/base/np_env.py). Task-dependent per-term and per-method keys - # (mba_obs_term_*, mba_reward_term_*, mba_getter__ms) are discovered - # dynamically at sample time (JSON/console only; not part of the fixed CSV - # header). - "mba_update_total_ms", - "mba_termination_ms", - "mba_reward_ms", - "mba_metrics_ms", - "mba_events_ms", - "mba_command_ms", - "mba_obs_compute_ms", - "mba_obs_map_ms", - "mba_termination_getter_ms", - "mba_reward_getter_ms", - "mba_metrics_getter_ms", - "mba_events_getter_ms", - "mba_command_getter_ms", - "mba_obs_compute_getter_ms", - "mba_obs_map_getter_ms", - "mba_update_internal_gap_ms", - "mba_getters_total_ms", - "mba_obs_noise_ms", - "mba_obs_clip_scale_ms", - "mba_obs_nan_check_ms", - "mba_obs_delay_ms", - "mba_obs_history_ms", - "mba_obs_concat_ms", ) NP_ENV_STEP_COUNT_KEYS = ("reset_done_count",) NP_ENV_STEP_SAMPLE_KEYS = (*NP_ENV_STEP_TIMING_KEYS, *NP_ENV_STEP_COUNT_KEYS) -# Task-dependent MBA timings emitted by ManagerBasedRlEnv (per-term reset event -# timings, per-term obs/reward timings, per-method getter timings); discovered -# dynamically from info["timing"] since term/method names vary per task. -MBA_DYNAMIC_KEY_PREFIX = "mba_" NP_RANDOM_PROFILE_FUNCTIONS = ( "uniform", "randint", @@ -206,15 +160,6 @@ ("dr_reset_obs_get_body_pose_ms", "dr_reset_obs_get_body_pose_ms"), ("dr_reset_observation_compute_obs_ms", "dr_reset_observation_compute_obs_ms"), ("dr_reset_observation_internal_gap_ms", "dr_reset_observation_internal_gap_ms"), - ("mba_reset_total_ms", "mba_reset_total_ms"), - ("mba_reset_curriculum_ms", "mba_reset_curriculum_ms"), - ("mba_reset_event_apply_ms", "mba_reset_event_apply_ms"), - ("mba_reset_command_reset_ms", "mba_reset_command_reset_ms"), - ("mba_reset_set_state_ms", "mba_reset_set_state_ms"), - ("mba_reset_manager_reset_ms", "mba_reset_manager_reset_ms"), - ("mba_reset_command_compute_ms", "mba_reset_command_compute_ms"), - ("mba_reset_obs_build_ms", "mba_reset_obs_build_ms"), - ("mba_reset_internal_gap_ms", "mba_reset_internal_gap_ms"), ("set_state_mask_ms", "set_state_mask_ms"), ("set_state_data_slice_ms", "set_state_data_slice_ms"), ("set_state_data_reset_ms", "set_state_data_reset_ms"), @@ -231,29 +176,6 @@ ("set_state_pool_reset_ms", "set_state_pool_reset_ms"), ("set_state_state_scatter_ms", "set_state_state_scatter_ms"), ("set_state_internal_gap_ms", "set_state_internal_gap_ms"), - ("mba_update_total_ms", "mba_update_total_ms"), - ("mba_termination_ms", "mba_termination_ms"), - ("mba_reward_ms", "mba_reward_ms"), - ("mba_metrics_ms", "mba_metrics_ms"), - ("mba_events_ms", "mba_events_ms"), - ("mba_command_ms", "mba_command_ms"), - ("mba_obs_compute_ms", "mba_obs_compute_ms"), - ("mba_obs_map_ms", "mba_obs_map_ms"), - ("mba_termination_getter_ms", "mba_termination_getter_ms"), - ("mba_reward_getter_ms", "mba_reward_getter_ms"), - ("mba_metrics_getter_ms", "mba_metrics_getter_ms"), - ("mba_events_getter_ms", "mba_events_getter_ms"), - ("mba_command_getter_ms", "mba_command_getter_ms"), - ("mba_obs_compute_getter_ms", "mba_obs_compute_getter_ms"), - ("mba_obs_map_getter_ms", "mba_obs_map_getter_ms"), - ("mba_update_internal_gap_ms", "mba_update_internal_gap_ms"), - ("mba_getters_total_ms", "mba_getters_total_ms"), - ("mba_obs_noise_ms", "mba_obs_noise_ms"), - ("mba_obs_clip_scale_ms", "mba_obs_clip_scale_ms"), - ("mba_obs_nan_check_ms", "mba_obs_nan_check_ms"), - ("mba_obs_delay_ms", "mba_obs_delay_ms"), - ("mba_obs_history_ms", "mba_obs_history_ms"), - ("mba_obs_concat_ms", "mba_obs_concat_ms"), ) @@ -685,13 +607,6 @@ def _run_active_window_case( ) else: env_step_timing_values["env_step_internal_gap_ms"] = None - # Task-dependent MBA dynamic keys (reset event terms, obs/reward - # terms, per-method getter timings); zero-filled by NpEnv on steps - # where they don't apply, once the key set has been discovered. - for key, value in _timing.items(): - if key.startswith(MBA_DYNAMIC_KEY_PREFIX) and key not in env_step_timing_values: - env_step_timing_values[key] = value - phase_start_ns = time.perf_counter_ns() next_obs_np, next_critic_np = split_obs_dict(state.obs) next_obs_np = np.asarray(next_obs_np, dtype=np.float32) @@ -1223,17 +1138,14 @@ def _format_np_env_value(result: CollectorResult, key: str, *, digits: int = 1) def _format_set_state_sub_ms(result: CollectorResult, key: str) -> str: """Format a backend set_state sub-timing as ``ms (%of set_state)``. - The percentage is relative to the outer set_state wall-clock measurement — - ``dr_reset_set_state_ms`` on the direct path, ``mba_reset_set_state_ms`` on - the MBA path — not to env_step_ms, so a reader can see which sub-step - dominates set_state. + The percentage is relative to ``dr_reset_set_state_ms`` (the outer + wall-clock measurement in DomainRandomizationManager), not to env_step_ms, + so a reader can see which sub-step dominates set_state. """ stat = result.env_step_timing_ms_per_vector_step.get(key) if stat is None: return "n/a" outer = result.env_step_timing_ms_per_vector_step.get("dr_reset_set_state_ms") - if outer is None or outer.mean_ms <= 0.0: - outer = result.env_step_timing_ms_per_vector_step.get("mba_reset_set_state_ms") if outer is None or outer.mean_ms <= 0.0: return f"{stat.mean_ms:.3f} (n/a)" pct = 100.0 * stat.mean_ms / outer.mean_ms @@ -1621,29 +1533,6 @@ def _print_result(result: CollectorResult) -> None: "dr_reset_observation_getters_ms", "dr_reset_obs_get_body_pose_ms", "dr_reset_observation_compute_obs_ms", - "mba_reset_total_ms", - "mba_reset_curriculum_ms", - "mba_reset_event_apply_ms", - "mba_reset_command_reset_ms", - "mba_reset_set_state_ms", - "mba_reset_manager_reset_ms", - "mba_reset_command_compute_ms", - "mba_reset_obs_build_ms", - "mba_reset_internal_gap_ms", - "mba_update_total_ms", - "mba_termination_ms", - "mba_reward_ms", - "mba_metrics_ms", - "mba_events_ms", - "mba_command_ms", - "mba_obs_compute_ms", - "mba_obs_map_ms", - "mba_update_internal_gap_ms", - "mba_getters_total_ms", - "mba_obs_noise_ms", - "mba_obs_clip_scale_ms", - "mba_obs_nan_check_ms", - "mba_obs_concat_ms", ): stat = result.env_step_timing_ms_per_vector_step.get(key) if stat is None: @@ -1653,17 +1542,6 @@ def _print_result(result: CollectorResult) -> None: f"pct_env={_env_step_child_env_pct(result, stat.mean_ms):5.1f}% " f"pct_active={_env_step_child_pct(result, stat.mean_ms):5.1f}%" ) - # Task-dependent MBA dynamic keys (reset event terms, obs/reward terms, - # per-method getter timings) not already printed above. - for key in sorted(result.env_step_timing_ms_per_vector_step): - if not key.startswith(MBA_DYNAMIC_KEY_PREFIX) or key in NP_ENV_STEP_SAMPLE_KEYS: - continue - stat = result.env_step_timing_ms_per_vector_step[key] - print( - f" {('np_env_' + key):<18} mean={stat.mean_ms:8.3f} ms " - f"pct_env={_env_step_child_env_pct(result, stat.mean_ms):5.1f}% " - f"pct_active={_env_step_child_pct(result, stat.mean_ms):5.1f}%" - ) def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 99f4cdfe1..34f36a6d2 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -8,7 +8,6 @@ from __future__ import annotations import re -import time from collections.abc import Iterator, Mapping, Sequence from contextlib import contextmanager from dataclasses import dataclass @@ -134,28 +133,6 @@ def _resolve_matching_names( return [item[1] for item in matches], [item[2] for item in matches] -class GetterTimingRecorder: - """Per-step wall-time accumulator for leaf backend state getters. - - One recorder per :class:`EntityData` instance; the owning env aggregates and - resets the recorders around its ``update_state`` pass (see - ``ManagerBasedRlEnv``). Timing is always on for actual backend reads: cache - hits add no synthetic getter time. - """ - - def __init__(self) -> None: - self.method_ms: dict[str, float] = {} - self.total_ms: float = 0.0 - - def record(self, method: str, elapsed_ms: float) -> None: - self.method_ms[method] = self.method_ms.get(method, 0.0) + elapsed_ms - self.total_ms += elapsed_ms - - def reset(self) -> None: - self.method_ms.clear() - self.total_ms = 0.0 - - _StateReadKey = tuple[str, tuple[int, ...] | None] @@ -246,11 +223,8 @@ def __init__( self._actuator_ctrl_range = actuator_ctrl_range self._control_buffer = control_buffer self._state_read_cache = state_read_cache - # Always-on leaf getter timing; aggregated/reset per update_state by the - # owning env (MBA update_state instrumentation, issue #1256). - self.getter_timing = GetterTimingRecorder() - def _timed_getter( + def _cached_getter( self, method: str, fn: Any, @@ -261,11 +235,7 @@ def _timed_getter( cached = self._state_read_cache.get(key) if cached is not None: return cached - t0 = time.perf_counter_ns() - try: - value = fn(*args) - finally: - self.getter_timing.record(method, (time.perf_counter_ns() - t0) / 1e6) + value = fn(*args) self._state_read_cache.put(key, value) return value @@ -280,7 +250,7 @@ def _require(self, value, capability: str): @property def root_link_pos_w(self) -> np.ndarray: ids = self._require(self._root_body_ids, "root body state") - return self._timed_getter( + return self._cached_getter( "body_pos_w", self._backend.get_body_pos_w, ids, @@ -290,7 +260,7 @@ def root_link_pos_w(self) -> np.ndarray: @property def root_link_quat_w(self) -> np.ndarray: ids = self._require(self._root_body_ids, "root body state") - return self._timed_getter( + return self._cached_getter( "body_quat_w", self._backend.get_body_quat_w, ids, @@ -300,7 +270,7 @@ def root_link_quat_w(self) -> np.ndarray: @property def root_link_lin_vel_w(self) -> np.ndarray: ids = self._require(self._root_body_ids, "root body state") - return self._timed_getter( + return self._cached_getter( "body_lin_vel_w", self._backend.get_body_lin_vel_w, ids, @@ -310,7 +280,7 @@ def root_link_lin_vel_w(self) -> np.ndarray: @property def root_link_ang_vel_w(self) -> np.ndarray: ids = self._require(self._root_body_ids, "root body state") - return self._timed_getter( + return self._cached_getter( "body_ang_vel_w", self._backend.get_body_ang_vel_w, ids, @@ -320,7 +290,7 @@ def root_link_ang_vel_w(self) -> np.ndarray: @property def root_link_lin_vel_b(self) -> np.ndarray: ids = self._require(self._root_body_ids, "root body state") - return self._timed_getter( + return self._cached_getter( "body_lin_vel_b", self._backend.get_body_lin_vel_b, ids, @@ -330,7 +300,7 @@ def root_link_lin_vel_b(self) -> np.ndarray: @property def root_link_ang_vel_b(self) -> np.ndarray: ids = self._require(self._root_body_ids, "root body state") - return self._timed_getter( + return self._cached_getter( "body_ang_vel_b", self._backend.get_body_ang_vel_b, ids, @@ -375,12 +345,12 @@ def default_root_state(self) -> np.ndarray: @property def joint_pos(self) -> np.ndarray: index = self._require(self._joint_pos_index, "joint position") - return self._timed_getter("dof_pos", self._backend.get_dof_pos)[:, index] + return self._cached_getter("dof_pos", self._backend.get_dof_pos)[:, index] @property def joint_vel(self) -> np.ndarray: index = self._require(self._joint_vel_index, "joint velocity") - return self._timed_getter("dof_vel", self._backend.get_dof_vel)[:, index] + return self._cached_getter("dof_vel", self._backend.get_dof_vel)[:, index] @property def joint_pos_biased(self) -> np.ndarray: @@ -410,7 +380,7 @@ def encoder_bias(self) -> np.ndarray: @property def body_link_pos_w(self) -> np.ndarray: ids = self._require(self._body_ids, "body state") - return self._timed_getter( + return self._cached_getter( "body_pos_w", self._backend.get_body_pos_w, ids, @@ -420,7 +390,7 @@ def body_link_pos_w(self) -> np.ndarray: @property def body_link_quat_w(self) -> np.ndarray: ids = self._require(self._body_ids, "body state") - return self._timed_getter( + return self._cached_getter( "body_quat_w", self._backend.get_body_quat_w, ids, @@ -430,7 +400,7 @@ def body_link_quat_w(self) -> np.ndarray: @property def body_link_lin_vel_w(self) -> np.ndarray: ids = self._require(self._body_ids, "body state") - return self._timed_getter( + return self._cached_getter( "body_lin_vel_w", self._backend.get_body_lin_vel_w, ids, @@ -440,7 +410,7 @@ def body_link_lin_vel_w(self) -> np.ndarray: @property def body_link_ang_vel_w(self) -> np.ndarray: ids = self._require(self._body_ids, "body state") - return self._timed_getter( + return self._cached_getter( "body_ang_vel_w", self._backend.get_body_ang_vel_w, ids, diff --git a/src/unilab/base/np_env.py b/src/unilab/base/np_env.py index 7c4f44035..9359f5a9d 100644 --- a/src/unilab/base/np_env.py +++ b/src/unilab/base/np_env.py @@ -45,20 +45,6 @@ "dr_reset_obs_get_body_pose_ms", "dr_reset_observation_compute_obs_ms", "dr_reset_observation_internal_gap_ms", - # Manager-Based (MBA) reset sub-steps, reported by ManagerBasedRlEnv.reset(). - # mba_reset_set_state_ms is the ResetStateTransaction commit wall time; its - # backend internals are merged into the set_state_* keys below. Per-term - # event timings (mba_reset_event_term__ms) are task-dependent and are - # zero-filled via the env's dynamic extra-key set instead of this tuple. - "mba_reset_total_ms", - "mba_reset_curriculum_ms", - "mba_reset_event_apply_ms", - "mba_reset_command_reset_ms", - "mba_reset_set_state_ms", - "mba_reset_manager_reset_ms", - "mba_reset_command_compute_ms", - "mba_reset_obs_build_ms", - "mba_reset_internal_gap_ms", # Backend-internal set_state sub-timings. All backends report the same key # set for column stability; sub-keys that don't apply report 0.0. "set_state_mask_ms", @@ -102,37 +88,6 @@ ) -# Fixed MBA update_state detail keys reported by ManagerBasedRlEnv.update_state() -# (issue #1256). All envs emit the same fixed key set for column stability; -# direct envs report 0.0. Task-dependent keys (mba_obs_term_*, mba_reward_term_*, -# mba_getter__ms) are zero-filled via the env's dynamic extra-key set. -UPDATE_STATE_DETAIL_TIMING_KEYS = ( - "mba_update_total_ms", - "mba_termination_ms", - "mba_reward_ms", - "mba_metrics_ms", - "mba_events_ms", - "mba_command_ms", - "mba_obs_compute_ms", - "mba_obs_map_ms", - "mba_termination_getter_ms", - "mba_reward_getter_ms", - "mba_metrics_getter_ms", - "mba_events_getter_ms", - "mba_command_getter_ms", - "mba_obs_compute_getter_ms", - "mba_obs_map_getter_ms", - "mba_update_internal_gap_ms", - "mba_getters_total_ms", - "mba_obs_noise_ms", - "mba_obs_clip_scale_ms", - "mba_obs_nan_check_ms", - "mba_obs_delay_ms", - "mba_obs_history_ms", - "mba_obs_concat_ms", -) - - @dataclass class NpEnvState: obs: dict[str, np.ndarray] @@ -159,12 +114,6 @@ def __init__(self, cfg: EnvCfg, backend: SimBackend, num_envs: int): self.step_counter = 0 self._dr_manager: DomainRandomizationManager | None = None self._init_randomization_applied = False - # Task-dependent reset detail keys (e.g. MBA per-event-term timings) - # merged into info["timing"]; zero-filled alongside the fixed keys. - self._reset_detail_extra_keys: set[str] = set() - # Task-dependent update_state detail keys (MBA obs/reward term and - # per-method getter timings); zero-filled every step. - self._update_detail_extra_keys: set[str] = set() self._nan_guard: NanGuard | None = None self._autoreset = True self._autoreset_reset_active = False @@ -281,18 +230,6 @@ def step(self, actions: np.ndarray) -> NpEnvState: timing["step_core_ms"] = step_core_time * 1000.0 timing["update_state_ms"] = update_state_time * 1000.0 timing["reset_done_ms"] = reset_done_time * 1000.0 - # update_state detail runs every step, so fixed keys are simply - # overwritten; dynamic keys discovered so far are zero-filled first. - for key in UPDATE_STATE_DETAIL_TIMING_KEYS: - timing[key] = 0.0 - for key in self._update_detail_extra_keys: - timing[key] = 0.0 - manager_update_timing = self._collect_manager_update_timing_ms() - if manager_update_timing: - self._update_detail_extra_keys.update( - key for key in manager_update_timing if key not in UPDATE_STATE_DETAIL_TIMING_KEYS - ) - timing.update(manager_update_timing) if backend_result is not None: backend_timing = backend_result.get("timing") if backend_timing: @@ -352,13 +289,6 @@ def _reset_done_envs(self) -> None: detail_timing["reset_done_reset_call_ms"] = (time.perf_counter() - t0) * 1000.0 if self._dr_manager is not None: detail_timing.update(self._dr_manager.last_reset_timing_ms) - manager_timing = self._collect_manager_reset_timing_ms() - if manager_timing: - self._reset_detail_extra_keys.update( - key for key in manager_timing if key not in RESET_DONE_DETAIL_TIMING_KEYS - ) - detail_timing.update(manager_timing) - t0 = time.perf_counter() for key in self._state.obs: self._state.obs[key][env_indices] = new_obs[key] @@ -393,24 +323,6 @@ def _reset_done_envs(self) -> None: def _clear_reset_done_detail_timing(self, timing: dict[str, Any]) -> None: for key in RESET_DONE_DETAIL_TIMING_KEYS: timing[key] = 0.0 - for key in self._reset_detail_extra_keys: - timing[key] = 0.0 - - def _collect_manager_reset_timing_ms(self) -> dict[str, float]: - """Detail timing from the most recent manager-driven reset. - - Direct envs report nothing (the DR manager merge covers them); - ManagerBasedRlEnv overrides this hook with MBA reset sub-step timings. - """ - return {} - - def _collect_manager_update_timing_ms(self) -> dict[str, float]: - """Detail timing from the most recent manager-driven update_state. - - Direct envs report nothing; ManagerBasedRlEnv overrides this hook with - MBA update_state block/term/getter timings. - """ - return {} def _resolve_nan_guard_model_file(self) -> str: scene = getattr(self._cfg, "scene", None) diff --git a/src/unilab/base/reset_state.py b/src/unilab/base/reset_state.py index b1b455d53..2bcdfefe5 100644 --- a/src/unilab/base/reset_state.py +++ b/src/unilab/base/reset_state.py @@ -54,8 +54,6 @@ def __init__( self._randomization_values: dict[str, np.ndarray] = {} self._randomization_dirty_masks: dict[str, np.ndarray] = {} self._requesting_terms: set[str] = set() - # Backend sub-timings from the most recent commit()'s set_state call. - self._last_set_state_timing_ms: dict[str, float] = {} self._last_commit_had_writes = False @property @@ -63,11 +61,6 @@ def active(self) -> bool: """Whether a reset lifecycle currently owns the transaction.""" return self._active - @property - def last_set_state_timing_ms(self) -> dict[str, float]: - """Backend-reported set_state sub-timings of the most recent commit.""" - return dict(self._last_set_state_timing_ms) - @property def last_commit_had_writes(self) -> bool: """Whether the most recent scoped commit submitted dirty rows to set_state.""" @@ -97,7 +90,6 @@ def begin(self, env_ids: np.ndarray) -> None: for mask in self._randomization_dirty_masks.values(): mask.fill(False) self._requesting_terms.clear() - self._last_set_state_timing_ms = {} self._last_commit_had_writes = False self._active = True @@ -549,20 +541,12 @@ def commit(self) -> dict | None: assert self._qvel is not None randomization = self._build_randomization_payload(dirty_ids) try: - result = self._backend.set_state( + return self._backend.set_state( dirty_ids, self._qpos[dirty_ids], self._qvel[dirty_ids], randomization=randomization, ) - backend_timing = result.get("timing") if isinstance(result, dict) else None - if isinstance(backend_timing, dict): - self._last_set_state_timing_ms = { - str(key): float(value) - for key, value in backend_timing.items() - if isinstance(value, (int, float)) - } - return result except (AttributeError, NotImplementedError) as exc: terms = ", ".join(sorted(self._requesting_terms)) raise NotImplementedError( diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 8ca8f1cae..fd9925dc2 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -8,7 +8,6 @@ import math import secrets -import time from dataclasses import dataclass, field from typing import Any @@ -254,12 +253,6 @@ def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int self._all_env_ids.setflags(write=False) self._has_transition = False self._uses_pre_step_control = False - # Sub-step wall time (ms) of the most recent partial reset; consumed by - # NpEnv._reset_done_envs via _collect_manager_reset_timing_ms(). - self._last_reset_timing_ms: dict[str, float] = {} - # Block/term/getter wall time (ms) of the most recent update_state; - # consumed by NpEnv.step via _collect_manager_update_timing_ms(). - self._last_update_timing_ms: dict[str, float] = {} self._load_managers() self._mapped_obs_dims = self._validate_observation_mapping() @@ -444,15 +437,6 @@ def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: self.action_manager.apply_action() return self._control - def _mba_getter_total_ms(self) -> float: - """Accumulated leaf backend getter time (ms) across all scene entities. - - EntityData records every leaf getter call into a per-entity - GetterTimingRecorder; update_state resets the recorders at block start - and reads deltas around blocks/terms for attribution (issue #1256). - """ - return sum(entity.data.getter_timing.total_ms for entity in self.scene.values()) - def update_state(self, state: NpEnvState) -> NpEnvState: # Physics stepping and reset/set_state lifecycles sit outside this private # scope. In-phase mutations explicitly invalidate it below. @@ -464,24 +448,10 @@ def _update_state_in_read_phase(self, state: NpEnvState) -> NpEnvState: state.info["log"] = log self.extras = state.info - # Reset leaf getter accumulators so this pass measures only update_state. - for entity in self.scene.values(): - entity.data.getter_timing.reset() - update_t0 = time.perf_counter() - update_timing: dict[str, float] = {} - np.add(state.info["steps"], 1, out=self.episode_length_buf) self.common_step_counter = self.step_counter + 1 self._sim_step_counter = self.common_step_counter * self._cfg.sim_substeps - def _mark(name: str, t0: float, getter0: float) -> tuple[float, float]: - now = time.perf_counter() - getter_now = self._mba_getter_total_ms() - update_timing[f"mba_{name}_ms"] = (now - t0) * 1000.0 - update_timing[f"mba_{name}_getter_ms"] = getter_now - getter0 - return now, getter_now - - t_prev, g_prev = time.perf_counter(), self._mba_getter_total_ms() self.termination_manager.compute() if self._cfg.is_finite_horizon: np.logical_or( @@ -494,17 +464,13 @@ def _mark(name: str, t0: float, getter0: float) -> tuple[float, float]: np.copyto(self.reset_terminated, self.termination_manager.terminated) np.copyto(self.reset_time_outs, self.termination_manager.time_outs) np.logical_or(self.reset_terminated, self.reset_time_outs, out=self.reset_buf) - t_prev, g_prev = _mark("termination", t_prev, g_prev) self.reward_buf = self.reward_manager.compute(dt=self.step_dt) log.update(self.reward_manager.step_reward_extras()) - t_prev, g_prev = _mark("reward", t_prev, g_prev) - update_timing.update(self.reward_manager.last_compute_timing_ms) if self._cfg.sim_substeps == 1: self.metrics_manager.compute_substep() self.metrics_manager.compute() - t_prev, g_prev = _mark("metrics", t_prev, g_prev) applied_runtime_event = False if "step" in self.event_manager.available_modes: @@ -518,7 +484,6 @@ def _mark(name: str, t0: float, getter0: float) -> tuple[float, float]: # interval/step capabilities; EventManager does not expose whether a # particular interval fired, so this boundary stays fail-closed. self.scene._invalidate_state_reads() - t_prev, g_prev = _mark("events", t_prev, g_prev) self._command_dt.fill(self.step_dt) self._command_dt[self.reset_buf] = 0.0 @@ -527,39 +492,12 @@ def _mark(name: str, t0: float, getter0: float) -> tuple[float, float]: if self._reset_state.last_commit_had_writes: self.scene._invalidate_state_reads() self.command_manager.post_compute() - t_prev, g_prev = _mark("command", t_prev, g_prev) manager_obs = self.observation_manager.compute(update_history=True) - t_prev, g_prev = _mark("obs_compute", t_prev, g_prev) - update_timing.update(self.observation_manager.last_compute_timing_ms) self.obs_buf = self._map_observations(manager_obs) - t_prev, g_prev = _mark("obs_map", t_prev, g_prev) self._has_transition = True - update_total_ms = (time.perf_counter() - update_t0) * 1000.0 - block_names = ( - "termination", - "reward", - "metrics", - "events", - "command", - "obs_compute", - "obs_map", - ) - update_timing["mba_update_total_ms"] = update_total_ms - update_timing["mba_update_internal_gap_ms"] = update_total_ms - sum( - update_timing[f"mba_{name}_ms"] for name in block_names - ) - getter_method_ms: dict[str, float] = {} - for entity in self.scene.values(): - for method, elapsed_ms in entity.data.getter_timing.method_ms.items(): - getter_method_ms[method] = getter_method_ms.get(method, 0.0) + elapsed_ms - update_timing["mba_getters_total_ms"] = g_prev - for method, elapsed_ms in getter_method_ms.items(): - update_timing[f"mba_getter_{method}_ms"] = elapsed_ms - self._last_update_timing_ms = update_timing - return state.replace( obs=self.obs_buf, reward=self.reward_buf, @@ -568,9 +506,6 @@ def _mark(name: str, t0: float, getter0: float) -> tuple[float, float]: info=state.info, ) - def _collect_manager_update_timing_ms(self) -> dict[str, float]: - return dict(self._last_update_timing_ms) - def _compute_truncated(self, state: NpEnvState) -> np.ndarray: del state self._no_truncation.fill(False) @@ -601,31 +536,17 @@ def reset( if self._has_transition and len(done_ids) > 0: self.recorder_manager.record_pre_reset(done_ids) - reset_t0 = time.perf_counter() log: dict[str, Any] = {} - t0 = time.perf_counter() self.curriculum_manager.compute(env_ids=ids) - curriculum_ms = (time.perf_counter() - t0) * 1000.0 - scoped_t0 = time.perf_counter() with self._reset_state.scoped(ids): - t0 = time.perf_counter() if "reset" in self.event_manager.available_modes: self.event_manager.apply( mode="reset", env_ids=ids, global_env_step_count=self.step_counter, ) - event_apply_ms = (time.perf_counter() - t0) * 1000.0 - t0 = time.perf_counter() log.update(self.command_manager.reset(ids)) - command_reset_ms = (time.perf_counter() - t0) * 1000.0 - # The transaction commits (backend set_state) on scoped-exit, so the - # remainder of the scoped block is the set_state wall time. - set_state_ms = ( - (time.perf_counter() - scoped_t0) * 1000.0 - event_apply_ms - command_reset_ms - ) - t0 = time.perf_counter() for manager in ( self.observation_manager, self.action_manager, @@ -636,7 +557,6 @@ def reset( self.termination_manager, ): log.update(manager.reset(ids)) - manager_reset_ms = (time.perf_counter() - t0) * 1000.0 self.episode_length_buf[ids] = 0 self._control[ids] = 0.0 @@ -644,17 +564,13 @@ def reset( if self._state is not None: self._state.info["steps"][ids] = 0 - t0 = time.perf_counter() self.command_manager.compute(dt=0.0, env_ids=ids) self.command_manager.post_compute() - command_compute_ms = (time.perf_counter() - t0) * 1000.0 - t0 = time.perf_counter() # Row-scoped reset rebuild (issue #1259 R2): the observation manager # returns only the reset rows, so no full-batch slice is needed here. manager_obs = self.observation_manager.compute(update_history=True, env_ids=ids) mapped_obs = self._map_observations(manager_obs, num_rows=len(ids)) reset_obs = {name: values.copy() for name, values in mapped_obs.items()} - obs_build_ms = (time.perf_counter() - t0) * 1000.0 if self._state is not None: for name, values in reset_obs.items(): @@ -669,37 +585,8 @@ def reset( self.obs_buf = self._state.obs if self._state is not None else mapped_obs self.extras = self._state.info if self._state is not None else {"log": log} self.recorder_manager.record_post_reset(ids) - - total_ms = (time.perf_counter() - reset_t0) * 1000.0 - measured_ms = ( - curriculum_ms - + event_apply_ms - + command_reset_ms - + set_state_ms - + manager_reset_ms - + command_compute_ms - + obs_build_ms - ) - reset_timing: dict[str, float] = { - "mba_reset_total_ms": total_ms, - "mba_reset_curriculum_ms": curriculum_ms, - "mba_reset_event_apply_ms": event_apply_ms, - "mba_reset_command_reset_ms": command_reset_ms, - "mba_reset_set_state_ms": set_state_ms, - "mba_reset_manager_reset_ms": manager_reset_ms, - "mba_reset_command_compute_ms": command_compute_ms, - "mba_reset_obs_build_ms": obs_build_ms, - "mba_reset_internal_gap_ms": total_ms - measured_ms, - } - for term_name, term_ms in self.event_manager.last_reset_term_timing_ms.items(): - reset_timing[f"mba_reset_event_term_{term_name}_ms"] = term_ms - reset_timing.update(self._reset_state.last_set_state_timing_ms) - self._last_reset_timing_ms = reset_timing return reset_obs, {"log": log} - def _collect_manager_reset_timing_ms(self) -> dict[str, float]: - return dict(self._last_reset_timing_ms) - def _normalize_reset_ids( self, env_indices: np.ndarray | None, diff --git a/src/unilab/managers/_types.py b/src/unilab/managers/_types.py index d24f01962..d59d21f5f 100644 --- a/src/unilab/managers/_types.py +++ b/src/unilab/managers/_types.py @@ -250,10 +250,6 @@ def max_episode_length(self) -> int: ... @property def max_episode_length_s(self) -> float: ... - def _mba_getter_total_ms(self) -> float: - """Accumulated leaf backend getter time (ms); 0.0 when uninstrumented.""" - ... - # Concrete task terms may still type their own richer env subclass. The # standalone manager core deliberately depends only on the properties above. diff --git a/src/unilab/managers/event_manager.py b/src/unilab/managers/event_manager.py index f169de872..fb98d7ea9 100644 --- a/src/unilab/managers/event_manager.py +++ b/src/unilab/managers/event_manager.py @@ -5,7 +5,6 @@ from __future__ import annotations -import time from copy import deepcopy from dataclasses import dataclass from typing import TYPE_CHECKING, Literal @@ -80,9 +79,6 @@ def __init__(self, cfg: dict[str, EventTermCfg | None], env: ManagerBasedRlEnv): self._mode_term_names: dict[EventMode, list[str]] = dict() self._mode_term_cfgs: dict[EventMode, list[EventTermCfg]] = dict() self._mode_class_term_cfgs: dict[EventMode, list[EventTermCfg]] = dict() - # Per-term wall time (ms) of the latest mode="reset" apply(); all reset - # term names are pre-registered at 0.0 in _prepare_terms. - self._reset_term_timing_ms: dict[str, float] = dict() super().__init__(env=env) @@ -117,15 +113,6 @@ def active_terms(self) -> dict[EventMode, list[str]]: def available_modes(self) -> list[EventMode]: return list(self._mode_term_names.keys()) - @property - def last_reset_term_timing_ms(self) -> dict[str, float]: - """Per-term wall time (ms) of the latest ``mode="reset"`` apply. - - All reset-mode term names are always present; terms gated out by - ``min_step_count_between_reset`` report 0.0 for that apply. - """ - return dict(self._reset_term_timing_ms) - # Methods. def get_term_cfg(self, term_name: str) -> EventTermCfg: @@ -182,10 +169,6 @@ def apply( if mode == "step" and dt is None: raise ValueError(f"Event mode '{mode}' requires the time-step of the environment.") - if mode == "reset": - for term_name in self._reset_term_timing_ms: - self._reset_term_timing_ms[term_name] = 0.0 - for index, term_cfg in enumerate(self._mode_term_cfgs[mode]): if mode == "interval": time_left = self._interval_term_time_left[index] @@ -220,12 +203,7 @@ def apply( if min_step_count == 0: self._reset_term_last_triggered_step_id[index][env_ids] = global_env_step_count self._reset_term_last_triggered_once[index][env_ids] = True - term_t0 = time.perf_counter() term_cfg.func(self._env, env_ids, **term_cfg.params) - term_name = self._mode_term_names[mode][index] - self._reset_term_timing_ms[term_name] += ( - time.perf_counter() - term_t0 - ) * 1000.0 else: last_triggered_step = self._reset_term_last_triggered_step_id[index][env_ids] triggered_at_least_once = self._reset_term_last_triggered_once[index][env_ids] @@ -241,12 +219,7 @@ def apply( self._reset_term_last_triggered_step_id[index][valid_env_ids] = ( global_env_step_count ) - term_t0 = time.perf_counter() term_cfg.func(self._env, valid_env_ids, **term_cfg.params) - term_name = self._mode_term_names[mode][index] - self._reset_term_timing_ms[term_name] += ( - time.perf_counter() - term_t0 - ) * 1000.0 else: term_cfg.func(self._env, env_ids, **term_cfg.params) @@ -294,7 +267,6 @@ def _prepare_terms(self) -> None: self._reset_term_last_triggered_step_id.append(step_count) no_trigger = np.zeros(self.num_envs, dtype=np.bool_) self._reset_term_last_triggered_once.append(no_trigger) - self._reset_term_timing_ms[term_name] = 0.0 func = term_cfg.func if hasattr(func, "model_fields"): diff --git a/src/unilab/managers/observation_manager.py b/src/unilab/managers/observation_manager.py index b52be349a..f47961c21 100644 --- a/src/unilab/managers/observation_manager.py +++ b/src/unilab/managers/observation_manager.py @@ -5,7 +5,6 @@ from __future__ import annotations -import time from copy import deepcopy from dataclasses import dataclass, field from typing import TYPE_CHECKING, Literal, Sequence @@ -158,11 +157,6 @@ def __init__(self, cfg: dict[str, ObservationGroupCfg | None], env: ManagerBased self._group_obs_dim[group_name] = group_term_dims self._obs_buffer: dict[str, np.ndarray | dict[str, np.ndarray]] | None = None - # Per-compute() timing breakdown consumed by the env's update_state - # instrumentation (issue #1256): per-term totals/getter shares and - # cross-term pipeline phases (noise/clip_scale/nan_check/delay/history/ - # concat), keyed with the mba_obs_* prefix. - self._last_compute_timing_ms: dict[str, float] = {} def __str__(self) -> str: msg = f" contains {len(self._group_obs_term_names)} groups.\n" @@ -241,11 +235,6 @@ def group_obs_term_dim(self) -> dict[str, list[tuple[int, ...]]]: def group_obs_concatenate(self) -> dict[str, bool]: return self._group_obs_concatenate - @property - def last_compute_timing_ms(self) -> dict[str, float]: - """Per-term and per-phase wall times (ms) of the latest compute() call.""" - return dict(self._last_compute_timing_ms) - # Methods. def get_term_cfg(self, group_name: str, term_name: str) -> ObservationTermCfg: @@ -374,7 +363,6 @@ def compute( return self._obs_buffer obs_buffer: dict[str, np.ndarray | dict[str, np.ndarray]] = dict() - self._last_compute_timing_ms = {} for group_name in self._group_obs_term_names: obs_buffer[group_name] = self.compute_group(group_name, update_history, env_ids) if env_ids is None: @@ -392,18 +380,6 @@ def compute_group( raise KeyError(f"Observation group '{group_name}' is disabled.") group_term_names = self._group_obs_term_names[group_name] group_obs: dict[str, np.ndarray] = {} - # Per-compute() instrumentation (issue #1256): per-term wall time with - # leaf-getter share, plus cross-term pipeline phase aggregation. Timing - # keys are merged into self._last_compute_timing_ms with mba_obs_* names. - phase_ms = { - "noise": 0.0, - "clip_scale": 0.0, - "nan_check": 0.0, - "delay": 0.0, - "history": 0.0, - "concat": 0.0, - } - term_timing: dict[str, float] = {} obs_terms = zip(group_term_names, self._group_obs_term_cfgs[group_name], strict=False) # Reset path (issue #1259 R2): when no term in this group uses delay or # history buffers, everything downstream of the term call is row @@ -413,8 +389,6 @@ def compute_group( # and per-row noise values identical to the full-batch path. row_scoped = env_ids is not None and not self._group_obs_temporal[group_name] for term_name, term_cfg in obs_terms: - term_t0 = time.perf_counter() - term_getter0 = self._env._mba_getter_total_ms() obs = term_cfg.func(self._env, **term_cfg.params) if not isinstance(obs, np.ndarray): raise TypeError( @@ -428,26 +402,21 @@ def compute_group( ) if not row_scoped: obs = obs.copy() - t0 = time.perf_counter() if isinstance(term_cfg.noise, noise_cfg.NoiseCfg): obs = term_cfg.noise.apply(obs, rng=self._env.rng) elif isinstance(term_cfg.noise, noise_cfg.NoiseModelCfg): obs = self._group_obs_class_instances[group_name][term_name](obs) - phase_ms["noise"] += time.perf_counter() - t0 if row_scoped: # Fresh row copy; safe for the in-place clip/scale below. obs = obs[env_ids] - t0 = time.perf_counter() if term_cfg.clip: np.clip(obs, term_cfg.clip[0], term_cfg.clip[1], out=obs) if term_cfg.scale is not None: scale = term_cfg.scale assert isinstance(scale, np.ndarray) np.multiply(obs, scale, out=obs) - phase_ms["clip_scale"] += time.perf_counter() - t0 # Check for NaN/Inf before delay/history buffers (per-term checking). - t0 = time.perf_counter() if group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled": obs = self._check_and_handle_nans( obs, @@ -455,9 +424,7 @@ def compute_group( policy=group_cfg.nan_policy, env_ids=env_ids if row_scoped else None, ) - phase_ms["nan_check"] += time.perf_counter() - t0 - t0 = time.perf_counter() if term_cfg.delay_max_lag > 0: delay_buffer = self._group_obs_term_delay_buffer[group_name][term_name] if env_ids is None or not delay_buffer.is_initialized: @@ -466,8 +433,6 @@ def compute_group( else: delay_buffer.backfill(obs, env_ids) obs = delay_buffer.peek() - phase_ms["delay"] += time.perf_counter() - t0 - t0 = time.perf_counter() if term_cfg.history_length > 0: circular_buffer = self._group_obs_term_history_buffer[group_name][term_name] if env_ids is None or not circular_buffer.is_initialized: @@ -482,16 +447,8 @@ def compute_group( group_obs[term_name] = circular_buffer.buffer else: group_obs[term_name] = obs - phase_ms["history"] += time.perf_counter() - t0 - - term_prefix = f"mba_obs_term_{group_name}_{term_name}" - term_timing[f"{term_prefix}_ms"] = (time.perf_counter() - term_t0) * 1000.0 - term_timing[f"{term_prefix}_getter_ms"] = ( - self._env._mba_getter_total_ms() - term_getter0 - ) # Final NaN check for non-per-term checking. - t0 = time.perf_counter() if not group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled": if self._group_obs_concatenate[group_name]: # Will check after concatenation below. @@ -504,16 +461,12 @@ def compute_group( policy=group_cfg.nan_policy, env_ids=env_ids if row_scoped else None, ) - phase_ms["nan_check"] += time.perf_counter() - t0 - t0 = time.perf_counter() if self._group_obs_concatenate[group_name]: result = np.concatenate( list(group_obs.values()), axis=self._group_obs_concatenate_dim[group_name] ) - phase_ms["concat"] += time.perf_counter() - t0 # Final check for concatenated result (non-per-term checking). - t0 = time.perf_counter() if not group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled": result = self._check_and_handle_nans( result, @@ -521,9 +474,7 @@ def compute_group( policy=group_cfg.nan_policy, env_ids=env_ids if row_scoped else None, ) - phase_ms["nan_check"] += time.perf_counter() - t0 else: - phase_ms["concat"] += time.perf_counter() - t0 result = group_obs if env_ids is not None and not row_scoped: @@ -535,11 +486,6 @@ def compute_group( else: result = result[env_ids] - timing = self._last_compute_timing_ms - for phase_name, elapsed_s in phase_ms.items(): - key = f"mba_obs_{phase_name}_ms" - timing[key] = timing.get(key, 0.0) + elapsed_s * 1000.0 - timing.update(term_timing) return result def _prepare_terms(self) -> None: diff --git a/src/unilab/managers/reward_manager.py b/src/unilab/managers/reward_manager.py index 261811909..262d64164 100644 --- a/src/unilab/managers/reward_manager.py +++ b/src/unilab/managers/reward_manager.py @@ -5,7 +5,6 @@ from __future__ import annotations -import time from copy import deepcopy from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -73,9 +72,6 @@ def __init__( self.cfg = deepcopy(cfg) super().__init__(env=env) self._episode_sums = dict() - # Per-compute() term timing consumed by the env's update_state - # instrumentation (issue #1256), keyed mba_reward_term__[getter_]ms. - self._last_compute_timing_ms: dict[str, float] = {} for term_name in self._term_names: self._episode_sums[term_name] = np.zeros(self.num_envs, dtype=np.float32) self._reward_buf = np.zeros(self.num_envs, dtype=np.float32) @@ -102,11 +98,6 @@ def __str__(self) -> str: def active_terms(self) -> list[str]: return self._term_names - @property - def last_compute_timing_ms(self) -> dict[str, float]: - """Per-term wall time and leaf-getter share (ms) of the latest compute().""" - return dict(self._last_compute_timing_ms) - # Methods. def reset(self, env_ids: np.ndarray | slice | None = None) -> dict[str, float]: @@ -126,26 +117,19 @@ def compute(self, dt: float) -> np.ndarray: raise ValueError(f"RewardManager received invalid dt {dt}.") self._reward_buf[:] = 0.0 scale = dt if self._scale_by_dt else 1.0 - term_timing: dict[str, float] = {} for term_idx, (name, term_cfg) in enumerate( zip(self._term_names, self._term_cfgs, strict=False) ): if term_cfg.weight == 0.0: self._step_reward[:, term_idx] = 0.0 continue - term_t0 = time.perf_counter() - term_getter0 = self._env._mba_getter_total_ms() value = term_cfg.func(self._env, **term_cfg.params) - term_getter_ms = self._env._mba_getter_total_ms() - term_getter0 self._check_term_shape(name, value) self._check_term_finite(name, value) value = value * term_cfg.weight * scale self._reward_buf += value self._episode_sums[name] += value self._step_reward[:, term_idx] = value / scale - term_timing[f"mba_reward_term_{name}_ms"] = (time.perf_counter() - term_t0) * 1000.0 - term_timing[f"mba_reward_term_{name}_getter_ms"] = term_getter_ms - self._last_compute_timing_ms = term_timing return self._reward_buf def step_reward_extras(self) -> dict[str, float]: diff --git a/tests/benchmark/test_offpolicy_collector_active_benchmark.py b/tests/benchmark/test_offpolicy_collector_active_benchmark.py index 206014e52..2fd6b31cb 100644 --- a/tests/benchmark/test_offpolicy_collector_active_benchmark.py +++ b/tests/benchmark/test_offpolicy_collector_active_benchmark.py @@ -256,6 +256,7 @@ def test_write_csv_includes_all_phase_columns(tmp_path) -> None: assert "env_step_overhead_ms" in header assert "physics_pct" in header assert "env_step_overhead_pct" in header + assert "mba_" not in header for key in bench.COLLECTOR_PHASES: assert key in header @@ -291,106 +292,6 @@ def test_write_csv_includes_backend_set_state_sub_timing_columns(tmp_path) -> No assert key in header, f"CSV header missing {key!r}" -def test_write_csv_includes_mba_reset_sub_step_columns(tmp_path) -> None: - """CSV headers must expose the fixed MBA reset sub-step keys.""" - result = _make_result(num_envs=2, throughput=2000.0, include_env_step_breakdown=True) - out_csv = tmp_path / "collector.csv" - - bench._write_csv(out_csv, [result]) - - header = out_csv.read_text(encoding="utf-8").splitlines()[0] - for key in ( - "mba_reset_total_ms", - "mba_reset_curriculum_ms", - "mba_reset_event_apply_ms", - "mba_reset_command_reset_ms", - "mba_reset_set_state_ms", - "mba_reset_manager_reset_ms", - "mba_reset_command_compute_ms", - "mba_reset_obs_build_ms", - "mba_reset_internal_gap_ms", - ): - assert key in header, f"CSV header missing {key!r}" - - -def test_write_csv_includes_mba_update_state_columns(tmp_path) -> None: - """CSV headers must expose the fixed MBA update_state block/phase keys.""" - result = _make_result(num_envs=2, throughput=2000.0, include_env_step_breakdown=True) - out_csv = tmp_path / "collector.csv" - - bench._write_csv(out_csv, [result]) - - header = out_csv.read_text(encoding="utf-8").splitlines()[0] - for key in ( - "mba_update_total_ms", - "mba_termination_ms", - "mba_reward_ms", - "mba_metrics_ms", - "mba_events_ms", - "mba_command_ms", - "mba_obs_compute_ms", - "mba_obs_map_ms", - "mba_termination_getter_ms", - "mba_reward_getter_ms", - "mba_metrics_getter_ms", - "mba_events_getter_ms", - "mba_command_getter_ms", - "mba_obs_compute_getter_ms", - "mba_obs_map_getter_ms", - "mba_update_internal_gap_ms", - "mba_getters_total_ms", - "mba_obs_noise_ms", - "mba_obs_clip_scale_ms", - "mba_obs_nan_check_ms", - "mba_obs_delay_ms", - "mba_obs_history_ms", - "mba_obs_concat_ms", - ): - assert key in header, f"CSV header missing {key!r}" - - -def test_print_result_includes_mba_update_and_dynamic_term_timings(capsys) -> None: - """Console breakdown shows MBA update blocks and dynamic per-term keys.""" - result = _make_result(num_envs=2, throughput=2000.0, include_env_step_breakdown=True) - result.env_step_timing_ms_per_vector_step["mba_update_total_ms"] = bench.TimingStats( - [5.0], 5.0, 5.0, 0.0, 5.0, 5.0 - ) - result.env_step_timing_ms_per_vector_step["mba_obs_term_actor_joint_pos_ms"] = ( - bench.TimingStats([2.0], 2.0, 2.0, 0.0, 2.0, 2.0) - ) - result.env_step_timing_ms_per_vector_step["mba_reward_term_pose_getter_ms"] = bench.TimingStats( - [1.0], 1.0, 1.0, 0.0, 1.0, 1.0 - ) - result.env_step_timing_ms_per_vector_step["mba_getter_dof_pos_ms"] = bench.TimingStats( - [0.5], 0.5, 0.5, 0.0, 0.5, 0.5 - ) - - bench._print_result(result) - - out = capsys.readouterr().out - assert "np_env_mba_update_total_ms" in out - assert "np_env_mba_obs_term_actor_joint_pos_ms" in out - assert "np_env_mba_reward_term_pose_getter_ms" in out - assert "np_env_mba_getter_dof_pos_ms" in out - - -def test_print_result_includes_mba_reset_and_per_term_event_timings(capsys) -> None: - """Console breakdown shows MBA reset sub-steps and dynamic per-term keys.""" - result = _make_result(num_envs=2, throughput=2000.0, include_env_step_breakdown=True) - result.env_step_timing_ms_per_vector_step["mba_reset_total_ms"] = bench.TimingStats( - [5.0], 5.0, 5.0, 0.0, 5.0, 5.0 - ) - result.env_step_timing_ms_per_vector_step["mba_reset_event_term_reset_root_ms"] = ( - bench.TimingStats([2.0], 2.0, 2.0, 0.0, 2.0, 2.0) - ) - - bench._print_result(result) - - out = capsys.readouterr().out - assert "np_env_mba_reset_total_ms" in out - assert "np_env_mba_reset_event_term_reset_root_ms" in out - - def test_format_set_state_detail_table_reports_sub_key_percentages() -> None: """The detail table shows each sub-key next to its share of ``dr_reset_set_state_ms`` (percentages must appear).""" diff --git a/tests/envs/locomotion/test_manager_gait_terms.py b/tests/envs/locomotion/test_manager_gait_terms.py index c52358eaa..d91bf2eb6 100644 --- a/tests/envs/locomotion/test_manager_gait_terms.py +++ b/tests/envs/locomotion/test_manager_gait_terms.py @@ -100,7 +100,6 @@ def _env(counter: int = 0, scene: _Scene | None = None) -> ManagerBasedRlEnv: step_dt=0.02, scene=scene or _Scene(), command_manager=_Commands(), - _mba_getter_total_ms=lambda: 0.0, ), ) @@ -137,7 +136,6 @@ def _parity_env() -> ManagerBasedRlEnv: scene=_ParityScene(), command_manager=_Commands(), max_episode_length_s=20.0, - _mba_getter_total_ms=lambda: 0.0, ), ) diff --git a/tests/envs/mdp/test_observations.py b/tests/envs/mdp/test_observations.py index 457f3cbad..e5dfd3b2a 100644 --- a/tests/envs/mdp/test_observations.py +++ b/tests/envs/mdp/test_observations.py @@ -154,7 +154,6 @@ def _env() -> tuple[ManagerBasedRlEnv, _Backend]: action_manager=_ActionManager(), command_manager=_CommandManager(), rng=np.random.default_rng(4), - _mba_getter_total_ms=lambda: 0.0, ), ) return env, backend diff --git a/tests/envs/mdp/test_rewards.py b/tests/envs/mdp/test_rewards.py index 9213fb675..b7fefd4b3 100644 --- a/tests/envs/mdp/test_rewards.py +++ b/tests/envs/mdp/test_rewards.py @@ -88,7 +88,6 @@ def _env() -> ManagerBasedRlEnv: terminated=np.asarray([False, True, False], dtype=np.bool_) ), max_episode_length_s=2.0, - _mba_getter_total_ms=lambda: 0.0, ), ) diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index fe1b83e8e..f4aa86973 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -1063,6 +1063,7 @@ def test_np_env_owns_substeps_autoreset_and_final_observation() -> None: np.testing.assert_array_equal(state.final_observation["obs"][:, 0], [2.0, 2.0]) assert state.info["_final_observation"].tolist() == [True, True] assert state.info["log"]["Episode_Termination/time_out"] == 2 + assert not any(key.startswith("mba_") for key in state.info["timing"]) pre_index = env.trace.index(("pre_reset", [0, 1])) post_reset_index = env.trace.index(("post_reset", [0, 1]), pre_index) @@ -1148,85 +1149,6 @@ def test_pure_reset_event_does_not_request_backend_state_capability() -> None: assert not hasattr(backend, "get_default_qpos") -def test_partial_reset_reports_mba_reset_timing_keys() -> None: - # No time_out term: only the explicit failure action triggers a reset. - cfg = _make_cfg() - cfg.terminations = {"failure": TerminationTermCfg(func=_failure)} - env, _ = _make_env(cfg) - env.init_state() - - state = env.step(np.array([[1.0], [0.0]], dtype=np.float32)) # env 0 autoresets - timing = state.info["timing"] - sub_step_keys = ( - "mba_reset_curriculum_ms", - "mba_reset_event_apply_ms", - "mba_reset_command_reset_ms", - "mba_reset_set_state_ms", - "mba_reset_manager_reset_ms", - "mba_reset_command_compute_ms", - "mba_reset_obs_build_ms", - "mba_reset_internal_gap_ms", - ) - for key in ("mba_reset_total_ms", *sub_step_keys, "mba_reset_event_term_reset_ms"): - assert key in timing - assert timing[key] >= 0.0 - assert sum(timing[key] for key in sub_step_keys) == pytest.approx(timing["mba_reset_total_ms"]) - - # Steps without a reset zero-fill both fixed and per-term MBA keys. - state = env.step(np.array([[0.0], [0.0]], dtype=np.float32)) - timing = state.info["timing"] - assert timing["mba_reset_total_ms"] == 0.0 - assert timing["mba_reset_event_term_reset_ms"] == 0.0 - - -def test_update_state_reports_mba_block_term_and_getter_timing_keys() -> None: - env, _ = _make_env() - env.init_state() - - state = env.step(np.array([[0.0], [0.0]], dtype=np.float32)) - timing = state.info["timing"] - block_names = ( - "termination", - "reward", - "metrics", - "events", - "command", - "obs_compute", - "obs_map", - ) - fixed_keys = ( - "mba_update_total_ms", - "mba_update_internal_gap_ms", - "mba_getters_total_ms", - "mba_obs_noise_ms", - "mba_obs_clip_scale_ms", - "mba_obs_nan_check_ms", - "mba_obs_delay_ms", - "mba_obs_history_ms", - "mba_obs_concat_ms", - *(f"mba_{name}_ms" for name in block_names), - *(f"mba_{name}_getter_ms" for name in block_names), - ) - # Task-dependent per-term keys discovered dynamically. - term_keys = ( - "mba_reward_term_track_ms", - "mba_reward_term_track_getter_ms", - "mba_obs_term_actor_policy_ms", - "mba_obs_term_actor_policy_getter_ms", - "mba_obs_term_value_critic_ms", - "mba_obs_term_value_critic_getter_ms", - ) - for key in (*fixed_keys, *term_keys): - assert key in timing, f"timing missing {key!r}" - assert timing[key] >= 0.0 - # Block decomposition closes against the update_state wall time. - assert sum(timing[f"mba_{name}_ms"] for name in block_names) + timing[ - "mba_update_internal_gap_ms" - ] == pytest.approx(timing["mba_update_total_ms"]) - # update_state detail is refreshed every step, including steps with no reset. - assert timing["mba_update_total_ms"] <= timing["update_state_ms"] + 1e-6 - - def test_update_state_reuses_backend_state_once_and_refreshes_after_physics() -> None: env, backend = _make_state_env() env.init_state() diff --git a/tests/managers/conftest.py b/tests/managers/conftest.py index 7271b567a..241c2b32c 100644 --- a/tests/managers/conftest.py +++ b/tests/managers/conftest.py @@ -52,10 +52,6 @@ def __init__(self, seed: int = 7, num_envs: int = 4) -> None: self.obs_buf: dict[str, np.ndarray] = {} self.reset_buf = np.zeros(num_envs, dtype=np.bool_) - def _mba_getter_total_ms(self) -> float: - """Stand-in for ManagerBasedRlEnv's leaf getter timing probe.""" - return 0.0 - @pytest.fixture def fake_env() -> FakeEnv: diff --git a/tests/managers/test_event_command_metrics_recorder.py b/tests/managers/test_event_command_metrics_recorder.py index e2bfb87d7..23bd67621 100644 --- a/tests/managers/test_event_command_metrics_recorder.py +++ b/tests/managers/test_event_command_metrics_recorder.py @@ -78,30 +78,6 @@ def model_mutation(env: FakeEnv, env_ids: np.ndarray | None) -> None: empty.apply("interval") -def test_event_reset_term_timing_is_tracked_per_term(fake_env: FakeEnv) -> None: - cfg = { - "always": EventTermCfg(func=_record, params={"label": "always"}, mode="reset"), - "gated": EventTermCfg( - func=_record, - params={"label": "gated"}, - mode="reset", - min_step_count_between_reset=100, - ), - } - manager = EventManager(cfg, fake_env) - - # First reset: the gated term fires once via the never-triggered exemption. - manager.apply("reset", env_ids=np.array([1, 3]), global_env_step_count=1) - # Second reset: the gated term is throttled out and must report 0.0. - manager.apply("reset", env_ids=np.array([1, 3]), global_env_step_count=2) - - timing = manager.last_reset_term_timing_ms - assert set(timing) == {"always", "gated"} - assert timing["always"] >= 0.0 - assert timing["gated"] == 0.0 - assert [label for label, _ in fake_env.calls] == ["always", "gated", "always"] - - def test_event_interval_rng_is_reproducible() -> None: cfg = { "interval": EventTermCfg(