Skip to content

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

Merged
TATP-233 merged 4 commits into
mainfrom
fix/issue-2495-mjwarp-backend-step-timing
Oct 22, 2026
Merged

TATP-233 merged 4 commits into
mainfrom
fix/issue-2495-mjwarp-backend-step-timing

Conversation

@TATP-233

@TATP-233 TATP-233 commented Oct 21, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

  • TorchEnv.step:修复事件型异步后端(mjwarp)的 Env Step 计时归因,解决 collector Env Step 分解缺 Backend Step 行的 bug。mjwarp 的 step_tensor 只发射 device work;物理等待发生在 step 输出之后第一个 host 读取处:
    1. ManagerBasedRlEnv.update_state 内 termination 阶段的 queue drain(仓库已测为 update_state_queue_drain_ms,注释本就声明其"drain pre-update GPU work",但此前 step_core_ms 从未包含它)——经新增保护协议 TorchEnv._collect_step_core_drain_ms 从 update_state_ms 重归入 step_core_ms;
    2. 若 update_state 无 host 读取,则 autoreset done-mask 门控是第一个同步点——该窗口从 reset_done_ms 拆出归入 step_core_ms。
      仅 produces_completion_receipts(mjwarp)的后端做重归因;同步后端(mujoco/motrix/genesis)drain 为零,归因不变;各 phase 合计与 env_step_total_ms 不变。
  • OffPolicyLogger:System 列新增 CUDA MPS 行,消费既有 runtime_manifest["cuda_process_sharing"] 证据——验证通过的显式请求显示 on,环境级透明 daemon 显示 ambient,否则 off。补上 Work: HOST_BRIDGE CUDA 载体吞吐归因收口——环境级 MPS 观测缺口、旋钮合并与 action 路径根本修复 #2479 观测缺口的终端侧。
  • 文档:en/zh_CN logging 页同步 System 列与 Backend Step 归因说明。

行为影响:仅计时归因与终端显示变化,不改变 step/reset 语义与执行路径。

Linked Work

Validation

  • make test-all passed on the final local head before this PR was created or updated
  • Additional task-specific validation listed below

Commands actually run:

uv run pytest tests/base/test_torch_env.py tests/envs/test_manager_based_rl_env.py tests/rl/logging/test_offpolicy_logger.py tests/rl/runtime/offpolicy/worker.py -q  # 193+ passed
make check                                                                                                                          # ruff/mypy/pyright pass
make test                                                                                                                           # 2764 passed, 38 skipped
uv run pytest tests/scripts/test_check_docs.py -q                                                                                   # 20 passed
make test-all                                                                                                                       # check + test-cov + benchmark smoke 34/34, 35/35 pass(最终 head GATE_EXIT=0)
# 2048-env 探针(g1_motion_tracking/mjwarp 原地 env):step_core_ms 0.05 -> 5.52,update_state_ms 8.5 -> 3.8,与 env_step_total 平衡
uv run train --algo flashsac --task g1_motion_tracking --sim mjwarp algo.max_iterations=8 algo.num_envs=2048
# 140 列终端实跑确认:Env Step 14.7ms = Apply Action 0.4 + Backend Step 5.0 + Update State 4.4 + Reset Done 4.6;System 列显示 "CUDA MPS off"(本机无 daemon,正确)

Remote CI route:

Impact

  • Backend impact: none(仅计时归因;mjwarp 显示受益,mujoco/motrix 数值不变)
  • Platform impact: Linux(MPS 指示仅消费 manifest 证据;无证据时不渲染该行)
  • Training effect expected: no

Artifacts

  • W&B: n/a
  • benchmark result: n/a(非性能行为变更,未跑 benchmark;探针数字见上)
  • video / screenshot: n/a
  • ONNX / checkpoint: n/a

Checklist

…iming

For an event-backed asynchronous backend step (mjwarp), step_tensor only
launches device work; the physics wait surfaces at the first host read of
the step outputs, which is the autoreset done-mask gate. That window was
measured inside reset_done_ms, so the collector Env Step breakdown showed
real device work as Reset Done while Backend Step collapsed to launch
overhead and disappeared under the terminal materiality filter. Measure
the gate separately and add it to step_core_ms/step_core_cpu_ms.

Issue #2495.
Render the rank-local cuda_process_sharing runtime-manifest evidence as a
CUDA MPS row: on for a validated request, ambient for a transparent
environment-level daemon, off otherwise. This closes the terminal-side
observation gap from #2479 where environment-level MPS silently changed
benchmark behavior without any run-level signal.

Issue #2495.
The first host read of an event-backed asynchronous backend step (mjwarp)
lands inside ManagerBasedRlEnv.update_state: the termination phase window,
already measured as update_state_queue_drain_ms, is where the host drains
the queued device work. Move that window from update_state_ms into
step_core_ms via the new TorchEnv._collect_step_core_drain_ms protocol so
the collector Env Step breakdown shows the physics wait as Backend Step
(2048-env probe: step_core 0.05 -> 5.52 ms, update_state 8.5 -> 3.8 ms).
Only receipt-publishing backends re-attribute; synchronous producers drain
nothing and keep their existing attribution.

Issue #2495.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant