调研：FlashSAC/MJWarp collector update_state 的 CUDA Graph 可行性与分阶段方案

[TATP-233]

结论

以 uv run --extra mjwarp train --algo flashsac --task g1_motion_tracking --sim mjwarp 为例，当前 collector tick 里的 ManagerBasedRlEnv.update_state() 不能直接整体捕获成 CUDA Graph。原因不是缺少设备常驻数据，而是 Manager 热路径仍有多个 host 同步、动态行选择、Python 分支和 Torch/Warp 跨 stream 依赖。

同时，MJWarp 的物理 step/reset/forward 已经在后端内部使用固定地址 Warp CUDA Graph；本机 RTX 4090 上诊断为 cuda_graph.enabled=true。因此这个议题讨论的是 Manager update-state 阶段的额外图化，不是给物理 step 补 CUDA Graph。

当前调用链

FlashSAC collector
  → env.step()
  → TorchEnv.step()
    → backend.step_tensor()          # MJWarp physics graph replay
    → update_state()                 # termination/reward/metrics/events/command/observation
    → _validate_state()
    → autoreset row selection/reset
  → transition extraction
  → CUDA-resident replay ingress

关键入口：

• Collector 调用环境：.venv/lib/python3.13/site-packages/uni_rl/offpolicy/worker.py#L767
• TorchEnv.step() 的 physics / update / validate / reset 顺序：src/unilab/base/torch_env.py#L283-L315
• update_state() 入口与读边界：src/unilab/envs/manager_based_rl_env.py#L1199-L1213
• MJWarp 物理 graph 初始化与 replay：.venv/lib/python3.13/site-packages/unisim/backend/mjwarp/backend.py#L1552-L1638

观测到的基线

在 RTX 4090、MJWarp 物理 graph 已启用的条件下：

• 1024 env、collector active 基准：update_state mean 约 3.275 ms，约占 active 窗口 25.9%；其中 update_state_queue_drain 约 1.07 ms。
• 8192 env：update_state mean 约 4.457 ms，约占 active 窗口 16.3%。
• 独立 1024-env owner 采样中，update_state median 约 3.4 ms，physics step_core 约 5.2 ms，reset-done 约 4.5 ms。
• 直接在一个已 step 的 2-env owner 上尝试 torch.cuda.CUDAGraph() 捕获 update_state() 失败，出现 Warp CUDA 905（dependency on uncaptured work in another stream），随后 PyTorch 报 cudaErrorStreamCaptureInvalidated。

queue_drain 是等待既有设备工作完成的时间，不是 kernel launch overhead；把 Manager kernel 图化不能直接消除这一段。

直接整体捕获的阻塞点

1. Host-side finite gates

Reward 每步执行：

finite = bool(torch.isfinite(self._reward_buf).all())

位置：src/unilab/managers/reward_manager.py#L190-L206。这会同步 CUDA stream。

Observation 拼接后也检查：

if not finite.all():

位置：src/unilab/managers/observation_manager.py#L707-L734。

TorchEnv._validate_state() 在 update_state() 之后又把 obs + reward concat 后做一次 finite reduction 和 host bool 同步：src/unilab/base/torch_env.py#L406-L426。

这些检查保证了错误暴露位置，但如果放进 captured region，每步仍然同步，CUDA Graph 的 launch-side 收益会被吞掉；直接捕获也会被同步操作破坏。

2. Motion sampler 的动态行输出

TensorMotionSampler.step_full() 返回：

done.nonzero(as_tuple=False).flatten()

位置：src/unilab/tasks/motion_tracking/common/tensor_sampler.py#L197-L216。

TensorMotionCommand._update_command() 随后读取 wrap_rows.numel() 并按结果分支：src/unilab/tasks/motion_tracking/common/manager_terms.py#L1393-L1400。

当前 G1 motion tracking 配置 truncate_on_clip_end: true，该分支实际不会走 resample，但仍为 .numel() 支付每步 host 同步。

3. Command finite / timer gates

Command 的 per-env dt 校验包含 bool(torch.isfinite(dt_tensor).all())：src/unilab/managers/command_manager.py#L216-L225。

timer resampling 使用 torch.nonzero(...) 并读取 numel()：src/unilab/managers/command_manager.py#L234-L250。全量 command 校验也有 host bool：src/unilab/managers/command_manager.py#L540-L548。

4. ManagerBasedRlEnv.step() 的 legacy reward log

每步 super().step() 返回后，环境执行：

means = self._reward_log_means.detach().cpu().tolist()

位置：src/unilab/envs/manager_based_rl_env.py#L1379-L1387，调用点在 src/unilab/envs/manager_based_rl_env.py#L1169-L1183。这保证即使 update-state 内核完成，collector 仍会很快遇到 D2H 同步。

5. Autoreset / replay / IPC 边界

• autoreset 依赖 bool((terminated | truncated).any())：src/unilab/base/torch_env.py#L310-L315。
• selected reset 是动态行集合，天然不适合固定 shape whole graph。
• collector 与 learner inference ring、replay ingress 存在跨进程 IPC 和 stream ordering 边界。
• replay ingress 发布和 terminal patch 也有保守同步/host gate。

因此 collector whole tick graph 不现实；可讨论的最多是 update-state 内部的固定形状子图。

6. Torch/Warp graph 与 stream ownership

MJWarp 的 read refresh 和 sensor view 依赖后端 stream 语义。如果 Manager 直接用 Torch CUDA Graph 捕获这些调用，容易出现本次实验中的 “dependency created on uncaptured work in another stream”。任何跨 backend graph/refresh 能力都应作为 SimBackend 公共能力声明并由 owner 协商，不能从 env/training 侧探测或调用 backend private graph 对象。

建议的分阶段路线

Phase 0：先建立可复现实验面

使用现有 timing 和 collector active benchmark，固定 1024 与 8192 两个规模，成对交替 eager/candidate，而不是只报告单次均值。

uv run --extra mjwarp scripts/benchmark/rl/benchmark_offpolicy_collector_active.py \
  --cases flashsac/g1_motion_tracking/mjwarp \
  --num-envs 1024 --warmup-steps 20 --measure-steps 100 \
  --out-json /tmp/flashsac_mjwarp_1024.json

Phase 1：先做 async hot path，不急着 graph

优先消除每步同步：

1. truncate_on_clip_end=true 时让 motion wrap 状态保留为 device mask，不生成/读取动态 rows；truncate 语义不变。
2. 合并 reward/observation/state finite 检查，或改为 device-side flag + 明确的批量/异常边界。
3. Reward per-term means 不再每步 .cpu().tolist()；改为 collector metrics interval 或 device-resident logging。
4. Command timer resample 对当前任务可以先保持全批 masked update，动态 rows 只在确实需要 resample 时进入图外路径。
5. 固定 Manager 输出和 scratch buffer，减少每步 clone/cat/rand 分配。

这一阶段即使不 graph 化，也可能降低 update_state_nonattributed_ms 与同步等待。

Phase 2：局部捕获固定形状 Manager kernel

候选区域按风险从低到高：

1. 固定全批 termination flag aggregation。
2. 固定 reward buffer zeroing、term aggregation、weighted accumulation。
3. Motion frame advance / current feature gather 的固定全批部分。
4. Observation noise/transform/concat，使用稳定输出 carrier。
5. MJWarp tracked-sensor refresh：固定 shape 且稳定地址，但它属于 backend owner，应通过公共 SimBackend capability 协商，而不是 Manager 私有捕获。

autoreset row selection、terminal observation patch、异常诊断和 finite diagnostics 应留在 graph 外。

Phase 3：显式配置与 fail-closed 诊断

建议新增类似：

{
  "update_state_cuda_graph": {
    "requested": true,
    "enabled": true,
    "disable_reason": null
  }
}

不可捕获时回退 eager 并说明原因。诊断应区分：

• Manager term 不满足 static shape；
• backend 不支持 update-state graph/stream ownership；
• capture 失败；
• 用户配置非法。

配置 owner 应放在任务/training owner YAML，而不是让 training.sim_backend 变成独立行为开关。

验收建议

1. 正确性
   • FlashSAC/G1/MJWarp rollout 与 selected-reset 测试通过。
   • Seeded eager vs graph 输出 obs/reward/terminated/truncated 在既定 tolerance 内一致。
   • timeout/truncation bootstrap、terminal observation、reset randomness 语义不变。
2. 捕获面
   • captured region 内无 .item()、bool(tensor)、.numel()、nonzero()、CPU transfer。
   • graph 外的 finite/autoreset/异常诊断仍有明确错误归属。
3. 性能
   • paired benchmark 至少 3 个 eager/candidate 交替 block。
   • 1024 env 的 update_state median 改善建议 ≥15%，collector active throughput 不回退（建议 ≥3%）。
   • 同时报告 8192 env，避免只为小 batch launch overhead 优化。
4. 架构
   • 所有 backend-specific graph/stream 能力通过 SimBackend 接口声明。
   • 无 env/training 侧 backend-private probe。
   • 不可用能力 fail closed，并输出诊断。

开放问题

1. finite 校验是否可以接受 device flag + 周期性/边界检查，还是必须每 transition 在 host 立即失败？
2. reward per-term logging 的消费方是否要求每步精确 host 值，还是可以跟随 collector metrics interval？
3. tracked-sensor refresh 的 graph ownership 应完全放在 MJWarp backend，还是由 Manager 发起但通过公共 capability 协商？
4. 1024 env 的 launch overhead 收益是否足以抵消图内存和 capture 复杂度；8192 env 下是否更适合只做 async/stable-buffer 优化？

我倾向先做 Phase 1 的 async hot path：它不引入 CUDA Graph 生命周期和 Torch/Warp stream 语义风险，且能直接处理当前最明显的每步同步。等 captured region 真正无同步后，再评估局部 graph 的实际收益。