From ea0f80f5958779a0fa1aaf48ddd2c45f37006eb6 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Wed, 21 Oct 2026 20:41:22 +0800 Subject: [PATCH] perf(env): collapse duplicate action/control scalar syncs on collector hot path Issue #2482 partial groundwork. CUDA-carrier runs without MPS expose every blocking scalar read on the apply_action path to dual-CUDA-context switching overhead. The collector step previously performed two hard bool() reads of 0-d finite verdicts per step: - TorchEnv._validate_control ran torch.isfinite(control).all() and read the scalar every step. - ActionManager.process_action(reuse_recorded_verdict=True) re-read the verdict that step() had already recorded. The reuse path now records without reading, and JointPositionAction keeps one device-resident verdict for its affine target. TorchEnv._validate_control publishes the manager verdict at the control-validation boundary, so the contract stays fail closed with one small verdict read instead of a full control reduction read. The verdict is cleared per step and never leaks. Benchmark parity was confirmed locally (RTX 4090, no-MPS, mujoco SAC g1_walk_flat, 2048 envs, 3x200 iters, tail-100; A/B by stashing only actions.py): baseline 64.9k steps/s / 5.65 ms apply_action vs patched 65.2k / 5.57 ms. The remaining no-MPS apply_action penalty is attributed outside this change (learner action publication non_blocking=False and host-bridge write_control hard sync) and is tracked for follow-up work. --- src/unilab/base/torch_env.py | 18 +++++++++++++++ src/unilab/envs/mdp/actions/actions.py | 7 ++++++ src/unilab/managers/action_manager.py | 14 +++++------ tests/base/test_torch_env.py | 32 ++++++++++++++++++++++++++ 4 files changed, 64 insertions(+), 7 deletions(-) diff --git a/src/unilab/base/torch_env.py b/src/unilab/base/torch_env.py index 7428f8438..c8e2e70e7 100644 --- a/src/unilab/base/torch_env.py +++ b/src/unilab/base/torch_env.py @@ -106,6 +106,7 @@ def __init__( self._final_observation_scratch: dict[str, torch.Tensor] | None = None self._action_finite_valid = torch.ones((), dtype=torch.bool, device=self._device) self._state_finite_valid = torch.ones((), dtype=torch.bool, device=self._device) + self._pending_control_finite_check: bool | None = None self._reset_obs_finite_valid = torch.ones((), dtype=torch.bool, device=self._device) self._reset_done_count = torch.zeros((), dtype=torch.int64, device=self._device) self._tensor_runtime_bound = False @@ -249,6 +250,9 @@ def step(self, actions: torch.Tensor) -> TorchEnvState: phase = time.perf_counter() phase_cpu = _cpu_time() + # Manager action owners publish one device-resident control verdict + # during this phase; it must not leak into the next step. + self._pending_control_finite_check = None ctrl = self.apply_action(actions, self._state) if self._nan_guard is not None: bad_ctrl_ids = self._nan_guard.check_ctrl(ctrl, step=self.step_counter) @@ -388,6 +392,14 @@ def _validate_control(self, ctrl: torch.Tensor) -> None: raise ValueError("TorchEnv control must be contiguous") if ctrl.device != self.device: raise ValueError(f"TorchEnv control device must be {self.device}, got {ctrl.device}") + verdict = self._publish_control_finite_verdict() + if verdict is not None: + # The action owner records one device-resident control verdict per + # step. Publishing it here keeps the fail-closed contract while + # removing the per-step scalar host read from the CUDA hot path. + if not verdict: + raise ValueError("TorchEnv control contains NaN or Inf") + return valid = torch.isfinite(ctrl).all() if not bool(valid): raise ValueError("TorchEnv control contains NaN or Inf") @@ -442,6 +454,12 @@ def _record_action_finite(self, actions: torch.Tensor) -> None: """Record the public action finite verdict on its persistent carrier.""" self._action_finite_valid.copy_(torch.isfinite(actions).all()) + def _publish_control_finite_verdict(self) -> bool | None: + """Publish and consume the manager control verdict, if recorded.""" + verdict = self._pending_control_finite_check + self._pending_control_finite_check = None + return verdict + def _require_floating_float(self, value: torch.Tensor) -> None: if not value.is_floating_point(): raise TypeError("TorchEnvState floating outputs must be floating-point") diff --git a/src/unilab/envs/mdp/actions/actions.py b/src/unilab/envs/mdp/actions/actions.py index 98a505635..8e01895d4 100644 --- a/src/unilab/envs/mdp/actions/actions.py +++ b/src/unilab/envs/mdp/actions/actions.py @@ -14,6 +14,7 @@ import torch +from unilab.base.torch_env import TorchEnv from unilab.managers.action_manager import ActionTerm, ActionTermCfg if TYPE_CHECKING: @@ -92,6 +93,7 @@ def __init__(self, cfg: BaseActionCfg, env: ManagerBasedRlEnv): ) self._processed_actions = torch.zeros_like(self._raw_actions) self._affine_shape = (self.num_envs, self.action_dim) + self._processed_actions_finite_valid = torch.ones((), dtype=torch.bool, device=self._device) self._scale = self._cold_affine(self._resolve_affine(cfg.scale, default=1.0, label="scale")) self._offset = self._cold_affine( self._resolve_affine(cfg.offset, default=0.0, label="offset") @@ -214,6 +216,7 @@ def _apply_affine(self, source: torch.Tensor, destination: torch.Tensor) -> None min=self._clip[..., 0], max=self._clip[..., 1], ) + self._processed_actions_finite_valid.copy_(torch.isfinite(destination).all()) def reset(self, env_ids: torch.Tensor | None = None) -> None: """Zero the raw actions of the selected envs.""" @@ -278,6 +281,10 @@ def apply_actions(self) -> None: self._entity.set_joint_position_target( self._tensor_target, joint_ids=self._target_ids ) + if isinstance(self._env, TorchEnv): + # Publish the affine-target device verdict for the owner + # boundary; the control-validation consumer owns the read. + self._env._pending_control_finite_check = bool(self._processed_actions_finite_valid) return processed = self._entity_values(self._processed_actions) encoder_bias = self._entity.data.encoder_bias_tensor.to( diff --git a/src/unilab/managers/action_manager.py b/src/unilab/managers/action_manager.py index 261a700dd..1e7f11ee9 100644 --- a/src/unilab/managers/action_manager.py +++ b/src/unilab/managers/action_manager.py @@ -235,13 +235,13 @@ def process_action(self, action: torch.Tensor, *, reuse_recorded_verdict: bool = raise ValueError( f"ActionManager action device must be {self._device}, received {action.device}" ) - if reuse_recorded_verdict: - valid = self._action_finite_valid - else: - valid = torch.isfinite(action).all() - self._action_finite_valid.copy_(valid) - if not bool(valid): - raise ValueError("ActionManager received an action containing NaN or Inf.") + if not reuse_recorded_verdict: + self._action_finite_valid.copy_(torch.isfinite(action).all()) + if not bool(self._action_finite_valid): + raise ValueError("ActionManager received an action containing NaN or Inf.") + # ``step`` records the device verdict once and publishes it at the + # control-validation boundary, so the reuse path stays free of + # blocking scalar host reads. # Shift history: prev_prev ← prev ← current ← new. self._prev_prev_action[:] = self._prev_action self._prev_action[:] = self._action diff --git a/tests/base/test_torch_env.py b/tests/base/test_torch_env.py index 50f8693c6..1e8659191 100644 --- a/tests/base/test_torch_env.py +++ b/tests/base/test_torch_env.py @@ -358,6 +358,38 @@ def apply_action(self, actions: torch.Tensor, state: TorchEnvState) -> torch.Ten env.step(_actions(env)) +def test_nonfinite_manager_control_verdict_fails_closed_without_host_read() -> None: + class _ManagerControlEnv(_StubTorchEnv): + def apply_action(self, actions: torch.Tensor, state: TorchEnvState) -> torch.Tensor: + control = actions * torch.nan + self._pending_control_finite_check = False + return control + + env = _ManagerControlEnv() + env.init_state() + with pytest.raises(ValueError, match="control contains NaN or Inf"): + env.step(_actions(env)) + + +def test_manager_control_verdict_does_not_leak_between_steps() -> None: + class _SpuriousVerdictEnv(_StubTorchEnv): + first = True + + def apply_action(self, actions: torch.Tensor, state: TorchEnvState) -> torch.Tensor: + if self.first: + self.first = False + self._pending_control_finite_check = False + return actions * torch.nan + return actions + + env = _SpuriousVerdictEnv() + env.init_state() + with pytest.raises(ValueError, match="control contains NaN or Inf"): + env.step(_actions(env)) + state = env.step(_actions(env)) + assert state.obs["obs"].shape == (env.num_envs, 3) + + def test_tensor_nan_guard_detects_and_dumps_at_explicit_host_boundary(tmp_path) -> None: class _NonfiniteObsEnv(_StubTorchEnv): def update_state(self, state: TorchEnvState) -> TorchEnvState: