fix(envs): restore mjwarp collector Backend Step timing and surface CUDA MPS status

docs/sphinx/source/en/2-user_guide/1-training/3-logging.md

@@ -82,8 +82,10 @@ The terminal has explicit width tiers:
     uses `Iter Wall` as its denominator.
   - `Collector (own clock)` is measured in the collector subprocess and runs in
     parallel with the learner.
-  - `System` contains the buffer size, timeout rate, environment count, and
-    per-rank batch size.
+  - `System` contains the buffer size, timeout rate, environment count,
+    per-rank batch size, and the rank-local CUDA MPS status (`on` for a
+    validated request, `ambient` for a transparent environment-level daemon,
+    `off` otherwise).
 
 Metric rendering separates semantic groups: losses and policy state use stable short
 labels, training diagnostics are indented, and replay-ingress gauges render as
@@ -265,7 +267,11 @@ The terminal keeps the core cycle rows above visible. `Inference Request`,
 canonical TensorBoard / W&B fields, but appear in the terminal only when they
 reach at least 1% of the collector cycle. `Backend Step`, `Update State`, and
 `Reset Done` follow the same materiality rule because they are nested diagnostics
-rather than additional cycle slices.
+rather than additional cycle slices. For an event-backed asynchronous backend
+(mjwarp), the device physics wait surfaces at the first host read of the step
+outputs — the update-state termination drain, or the autoreset done-mask gate
+when update_state performs no host read — and is attributed to `Backend Step`,
+so the row reports real device work instead of only the kernel-launch overhead.
 
 `Learner Action Wait` is deliberately not named “Inference Wait”: it is not pure
 inference latency. If the collector finishes `Env Step + Replay Write` and submits its

docs/sphinx/source/zh_CN/2-user_guide/1-training/3-logging.md

@@ -76,7 +76,9 @@ key 的契约如下：
 - 120 列及以上时，timing 使用完整三列：
   - `Learner (Iter Wall)`：learner 主线程的一圈；所有行使用 `Iter Wall` 作分母。
   - `Collector (own clock)`：collector 子进程自己的采集时钟，与 learner 并行。
-  - `System`：buffer 大小、timeout rate、env 数和每 rank batch 大小。
+  - `System`：buffer 大小、timeout rate、env 数、每 rank batch 大小，以及 rank
+    本地 CUDA MPS 状态（验证通过的请求显示 `on`，环境级透明 daemon 显示
+    `ambient`，否则 `off`）。
 
 Metric 渲染按语义分组：loss 与 policy 状态使用稳定短标签，training 诊断缩进，
 replay-ingress gauge 显示为 `Ingress Backpressure`、`Ingress Depth` 等标签；它们
@@ -238,6 +240,10 @@ SAC / FlashSAC 每个 vectorized env tick 记录热路径互斥阶段，终端
 `Action+Backend` 明细属于安静诊断项：它们仍是 canonical TensorBoard / W&B 字段，
 但只有达到 collector cycle 的 1% 时才出现在终端。`Backend Step`、
 `Update State` 与 `Reset Done` 是嵌套诊断而非额外 cycle 切片，同样按该显著性规则显示。
+对于事件型异步后端（mjwarp），设备物理等待发生在 step 输出的第一次 host 读取
+——update_state 内 termination 阶段的 queue drain（若 update_state 没有 host 读取，
+则为 autoreset 的 done-mask 门控）——并归因到 `Backend Step`，因此该行反映真实设备
+工作而非仅 kernel 发射开销。
 
 `Learner Action Wait` 特意不叫 “Inference Wait”：它不是纯 inference latency。如果 collector
 在 learner update 期间先完成 `Env Step + Replay Write` 并提交下一 request，这一项会包含

src/unilab/base/torch_env.py

@@ -308,6 +308,14 @@ def step(self, actions: torch.Tensor) -> TorchEnvState:
         self._validate_state(self._state)
         update_state_ms = (time.perf_counter() - phase) * 1000.0
         update_state_cpu_ms = (_cpu_time() - phase_cpu) * 1000.0
+        if self.produces_completion_receipts:
+            # An event-backed asynchronous step drains its device work at the
+            # first update-state host read; re-attribute that measured window
+            # to the backend step phase so update_state reports manager work.
+            drain_ms = self._collect_step_core_drain_ms(self._state)
+            if drain_ms > 0.0:
+                step_core_ms += drain_ms
+                update_state_ms = max(update_state_ms - drain_ms, 0.0)
 
         self._state.info["steps"] += 1
         self.step_counter += 1
@@ -317,8 +325,17 @@ def step(self, actions: torch.Tensor) -> TorchEnvState:
         phase = time.perf_counter()
         phase_cpu = _cpu_time()
         # Preserve the semantic autoreset gate: an empty done mask must not
-        # consume reset RNG or publish a terminal carrier.
+        # consume reset RNG or publish a terminal carrier. When update_state
+        # made no host read of the step outputs, this gate is the first
+        # synchronization point after an event-backed asynchronous backend
+        # step, so its wait belongs to the backend step phase rather than to
+        # reset bookkeeping.
         did_reset = self._autoreset and bool((self._state.terminated | self._state.truncated).any())
+        step_core_ms += (time.perf_counter() - phase) * 1000.0
+        step_core_cpu_ms += (_cpu_time() - phase_cpu) * 1000.0
+
+        phase = time.perf_counter()
+        phase_cpu = _cpu_time()
         if did_reset:
             self._reset_done_envs()
         reset_done_ms = (time.perf_counter() - phase) * 1000.0
@@ -645,6 +662,18 @@ def _collect_reset_backend_timing_ms(self) -> dict[str, float]:
         """Return backend-owned sub-timings from the most recent tensor reset."""
         return {}
 
+    def _collect_step_core_drain_ms(self, state: TorchEnvState) -> float:
+        """Return the update-state window that drained an async backend step.
+
+        Only consulted when the backend publishes event-backed completion
+        receipts (``produces_completion_receipts``); synchronous backends drain
+        nothing, so the default is zero. An env whose update-state sub-phase
+        hosts the first read of asynchronous step outputs overrides this to
+        report that measured window.
+        """
+        del state
+        return 0.0
+
     def _clear_step_final_observation(self) -> None:
         assert self._state is not None
         self._state.final_observation = None

src/unilab/envs/manager_based_rl_env.py

@@ -1469,6 +1469,20 @@ def _collect_reset_backend_timing_ms(self) -> dict[str, float]:
         timing.update(self._last_reset_manager_timing_ms)
         return timing
 
+    def _collect_step_core_drain_ms(self, state: TorchEnvState) -> float:
+        # Termination is the first update-state reduction after the backend
+        # step; for an event-backed asynchronous step its window (measured as
+        # update_state_queue_drain_ms) drains the queued device work, which is
+        # backend step time rather than Manager work.
+        timing = state.info.get("timing")
+        if not isinstance(timing, dict):
+            return 0.0
+        value = timing.get("update_state_queue_drain_ms")
+        if isinstance(value, bool) or not isinstance(value, (int, float)):
+            return 0.0
+        value_float = float(value)
+        return value_float if math.isfinite(value_float) else 0.0
+
     def reset(
         self,
         env_ids: torch.Tensor | None = None,

src/unilab/rl/logging/offpolicy.py

@@ -872,6 +872,17 @@ def _build_reward_table(self, snapshot: _TerminalSnapshot | None = None) -> Tabl
             source_iteration=max(source_iterations) if source_iterations else None,
         )
 
+    def _cuda_mps_status_label(self) -> str | None:
+        """Render rank-local CUDA MPS state from runtime-manifest evidence."""
+        evidence = self._runtime_manifest.get("cuda_process_sharing")
+        if not isinstance(evidence, dict):
+            return None
+        if evidence.get("effective") == "mps":
+            return "[green]on[/]"
+        if evidence.get("ambient_server"):
+            return "[yellow]ambient[/]"
+        return "[dim]off[/]"
+
     def _build_timing_table(self, snapshot: _TerminalSnapshot | None = None) -> Table:
         snapshot = snapshot or self._terminal_snapshot
         terminal_width = self._console.width
@@ -1006,6 +1017,9 @@ def _show_collector_row(key: str) -> bool:
         system_items.append(("Envs", f"{self.num_envs:,}"))
         if batch_size_per_rank > 0:
             system_items.append(("Batch/Rank", f"{batch_size_per_rank:,}"))
+        cuda_mps_status = self._cuda_mps_status_label()
+        if cuda_mps_status is not None:
+            system_items.append(("CUDA MPS", cuda_mps_status))
         column_items: tuple[list[tuple[str, str]], ...]
         if terminal_width >= 120:
             column_items = (learner_items, collector_items, system_items)

tests/base/test_torch_env.py

@@ -3,6 +3,8 @@
 from __future__ import annotations
 
 import dataclasses
+import itertools
+import types
 from dataclasses import dataclass
 from typing import Any, cast
 from unittest.mock import MagicMock
@@ -19,6 +21,7 @@
     TensorRuntimeDiagnostic,
 )
 
+import unilab.base.torch_env as torch_env_module
 from unilab.base.base import EnvCfg, EnvPlayCapabilities
 from unilab.base.scene import SceneCfg
 from unilab.base.torch_env import TorchEnv, TorchEnvState
@@ -220,6 +223,66 @@ def test_step_without_done_preserves_eager_autoreset_gate() -> None:
     assert state.info["timing"]["reset_done_count"] == 0.0
 
 
+def test_step_done_gate_host_read_is_attributed_to_step_core(monkeypatch) -> None:
+    wall = itertools.count(start=0.0, step=0.001)
+    cpu = itertools.count(start=0.0, step=0.001)
+    monkeypatch.setattr(
+        torch_env_module,
+        "time",
+        types.SimpleNamespace(perf_counter=lambda: next(wall)),
+    )
+    monkeypatch.setattr(torch_env_module, "_cpu_time", lambda: next(cpu))
+
+    env = _StubTorchEnv(terminate=False)
+    env.init_state()
+    state = env.step(_actions(env))
+
+    timing = state.info["timing"]
+    # The fake clocks advance 1 ms per call, so each measured window spans
+    # exactly 1 ms. The backend call window and the done-gate host read both
+    # belong to step_core; reset bookkeeping keeps only its own empty window.
+    assert timing["step_core_ms"] == pytest.approx(2.0)
+    assert timing["step_core_cpu_ms"] == pytest.approx(2.0)
+    assert timing["reset_done_ms"] == pytest.approx(1.0)
+    assert timing["reset_done_cpu_ms"] == pytest.approx(1.0)
+    assert timing["env_step_other_cpu_ms"] >= 0.0
+
+
+def test_step_reattributes_update_state_drain_to_step_core(monkeypatch) -> None:
+    wall = itertools.count(start=0.0, step=0.001)
+    cpu = itertools.count(start=0.0, step=0.001)
+    monkeypatch.setattr(
+        torch_env_module,
+        "time",
+        types.SimpleNamespace(perf_counter=lambda: next(wall)),
+    )
+    monkeypatch.setattr(torch_env_module, "_cpu_time", lambda: next(cpu))
+
+    class _DrainEnv(_StubTorchEnv):
+        def _collect_step_core_drain_ms(self, state: TorchEnvState) -> float:
+            return 0.6
+
+    async_backend = _backend()
+    async_backend.backend_type = "mjwarp"
+    env = _DrainEnv(backend=async_backend, terminate=False)
+    env.init_state()
+    timing = env.step(_actions(env)).info["timing"]
+
+    # Fake clock: every window spans 1 ms. Backend launch (1 ms) + the drained
+    # update-state window (0.6 ms) + the done-gate host read (1 ms) all belong
+    # to step_core; update_state keeps only the remaining manager work.
+    assert timing["step_core_ms"] == pytest.approx(2.6)
+    assert timing["update_state_ms"] == pytest.approx(0.4)
+    assert timing["reset_done_ms"] == pytest.approx(1.0)
+
+    # Synchronous producers publish no completion receipt and drain nothing.
+    sync_env = _DrainEnv(terminate=False)
+    sync_env.init_state()
+    sync_timing = sync_env.step(_actions(sync_env)).info["timing"]
+    assert sync_timing["step_core_ms"] == pytest.approx(2.0)
+    assert sync_timing["update_state_ms"] == pytest.approx(1.0)
+
+
 def test_reset_done_count_remains_device_resident_before_timing_publication() -> None:
     env = _StubTorchEnv()
     env.init_state()

tests/envs/test_manager_based_rl_env.py

@@ -2335,6 +2335,9 @@ def test_update_state_publishes_drain_vs_host_termination_attribution() -> None:
     assert timing["update_state_state_replace_ms"] >= 0.0
     assert timing["update_state_child_sum_ms"] >= 0.0
     assert timing["update_state_command_preflight_ms"] >= 0.0
+    assert env._collect_step_core_drain_ms(state) == pytest.approx(
+        timing["update_state_queue_drain_ms"]
+    )
     observation_timing = {
         key: state.info["update_state_timing"][key]
         for key in (

tests/rl/logging/test_offpolicy_logger.py

@@ -324,6 +324,41 @@ def test_offpolicy_terminal_render_is_width_aware() -> None:
         assert "Wait Ms" not in output
 
 
+def test_offpolicy_logger_system_column_shows_cuda_mps_status() -> None:
+    logger = OffPolicyLogger(log_backend="none")
+    logger._console = Console(width=140, record=True, force_terminal=False)
+
+    def system_cells() -> tuple[list[str], list[str]]:
+        table = logger._build_timing_table()
+        labels = [cell for cell in table.columns[4].cells if cell]
+        values = [cell for cell in table.columns[5].cells if cell]
+        return labels, values
+
+    labels, _ = system_cells()
+    assert "CUDA MPS" not in labels
+
+    logger.update_runtime_manifest(
+        {"cuda_process_sharing": {"effective": "mps", "ambient_server": False}}
+    )
+    labels, values = system_cells()
+    assert labels[-1] == "CUDA MPS"
+    assert values[-1] == "[green]on[/]"
+
+    logger.update_runtime_manifest(
+        {"cuda_process_sharing": {"effective": None, "ambient_server": True}}
+    )
+    labels, values = system_cells()
+    assert labels[-1] == "CUDA MPS"
+    assert values[-1] == "[yellow]ambient[/]"
+
+    logger.update_runtime_manifest(
+        {"cuda_process_sharing": {"effective": None, "ambient_server": False}}
+    )
+    labels, values = system_cells()
+    assert labels[-1] == "CUDA MPS"
+    assert values[-1] == "[dim]off[/]"
+
+
 def test_offpolicy_logger_rejects_retired_collector_timing_names() -> None:
     logger = OffPolicyLogger(log_backend="none")