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: