perf(env): collapse duplicate action/control scalar syncs on collector hot path

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")

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(

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

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: