审计：G1 motion tracking/WBT collector 的 update state 与 reset 为什么比 MotrixLab 慢，以及可借鉴的优化方向

[TATP-233]

审计：G1 motion tracking/WBT collector 的 update state 与 reset 为什么比 MotrixLab 慢？

只读架构审计。本文不修改代码，不声称已经完成性能验证；所有数字均来自当前仓库 checked-in benchmark 与本地 TensorBoard 日志。

TL;DR

MotrixLab 的 g1-wbt-* 快，不是因为它的 MotrixSim Torch backend 有一个“魔法 reset”，而是因为它走的是另一条执行架构：

• CPU/NumPy ArrayEnv;
• 原生批量 read/write program;
• 生成并缓存后的 Numba fused kernel;
• 全程 CPU-resident manager 状态;
• collector 本地推理与异步 ring。

UniLab 的 sim_backend=motrix 使用的是 UniLab Manager-Based Torch 生命周期 + unisim Motrix host bridge，并不等同于 MotrixLab 自己的 WBT runtime。

同时有两个重要校正：

1. UniLab mjwarp 的长时间运行并不是数量级慢。本地 2048 env FlashSAC/G1MotionTracking/mjwarp 日志中，collector_env_step_ms 平均约 18.94 ms/vector step，即约 108k env-steps/s；MotrixLab checked-in g1-wbt-backflip benchmark 是 110,975 env-steps/s，折算约 18.4 ms/vector step。但硬件、任务、算法阶段和语义并不完全相同，只能说明总吞吐差距没有 phase 级数字看起来那么夸张。
2. mjwarp 的 update_state_ms 不是纯 manager 计算时间。mjwarp backend step 经常只是提交 GPU 工作，step_core_ms 很小；到 update_state 的第一个 reduction/bool/sync 时才排空 GPU 队列，因此 physics 时间会被归因到 update_state_ms。这和 MotrixLab CPU 同步 physics 后再执行 manager kernel 的计时模型不是同一种口径。

---

1. 实测数据

1.1 UniLab 本机 TensorBoard

run	env 数	env.step 平均	update_state	reset_done	备注
logs/flash_sac/G1MotionTracking/2026-10-20_11-47-57_mjwarp	2048	18.94 ms	13.90 ms	4.14 ms	24,998 个样本，较稳定
logs/flash_sac/G1MotionTracking/2026-10-21_18-18-01_mujoco	2048	27.45 ms	3.98 ms	4.86 ms	483 个样本，短 run，可能含瞬态
logs/flash_sac/G1MotionTracking/2026-10-21_18-20-11_motrix	2048	31.59 ms	7.09 ms	7.02 ms	487 个样本，短 run

配置均为：

• task: G1MotionTracking;
• 2048 env;
• sim_dt=0.006666...;
• ctrl_dt=0.02;
• motion: motions/g1/dance1_subject2_part.npz;
• sampling mode: adaptive;
• truncate_on_clip_end: true.

1.2 UniLab checked-in collector attribution

来自 scripts/benchmark/outputs/issue2307/phase_b_reprofile/attribution_primary_null_current.csv，FlashSAC/G1MotionTracking/mjwarp/2048：

• collector active throughput: 145,607 env-steps/s;
• env_step_ms: 13.72 ms;
• physics: 5.77 ms;
• update_state_ms: 4.41 ms;
• reset_done_ms: 3.06 ms;
• 平均 reset 行数: 67 / 2048;
• mjwarp 实际 selected state commit:
  • set_state_tensor_mask_ms: 0.013 ms;
  • set_state_tensor_commit_forward_ms: 0.084 ms.

这说明 mjwarp 的 backend selected reset commit 本身很便宜；慢的是周边 manager/motion/read/observation 生命周期。

1.3 MotrixLab checked-in benchmark

g1-wbt-backflip provenance：

• mean_env_steps_per_s: 110,975;
• num_envs: 2048;
• EPYC 9J14 96-Core / 192 threads + RTX 4090;
• motrix-env-core 0.4.0b0, motrixsim 0.10.1, Torch 2.7.

文件：third-party/MotrixLab/docs/source/_static/data/performance/g1-wbt-backflip.json。

g1-wbt-dance 曲线点粗略折算约 39.7k env-steps/s，但该文件缺少同等 provenance，不建议直接引用为强结论。

---

2. MotrixLab 为什么快

2.1 它的 G1 WBT 不是 Torch Manager 环境

MotrixLab 的 g1-wbt-dance 注册的是 CPU/NumPy ManagerEnv：

• motrix_envs/src/motrix_envs/locomotion/wbt/g1/dance.py:13-30

FastSAC collector 构建：

• motrix_rl_builtin/src/motrix_rl_builtin/fastsac/async_impl/worker.py:79-92

因为它是 ArrayEnv，走 NumPy wrapper：

• motrix_rl_builtin/src/motrix_rl_builtin/fastsac/wrap_np.py:35-51

因此 WBT reset/update 的 manager 状态本身不经历 Torch CUDA carrier、Torch manager term dispatch、GPU bool 同步。GPU 主要用于本地 actor 推理和 learner 更新。

2.2 每步只有一次完整 simulator read

ArrayEnv.step 的生命周期：

apply_action -> physics -> compute_transition -> episode bookkeeping -> final obs -> reset -> compute_observation

证据：motrix_env_core/src/motrix_env_core/array/env.py:324-357。

注释还明确：

• compute_transition 是唯一 full simulator read;
• reset 后只做 selected-row read;
• compute_observation 只从已刷新数据构造，不再查 simulator。

证据：同文件 :287-322。

ManagerEnv.compute_transition 实际执行：

1. 一次 batched simulator read;
2. 一个 fused evaluate kernel;
3. 少量 host-side command on_transition()。

证据：motrix_env_core/src/motrix_env_core/numba/manager/env.py:778-786。

read plan 的输入 bundle 是持久 flat tuple，simulator-backed array 原地重写，不每步重建：

• motrix_env_core/src/motrix_env_core/numba/manager/env.py:859-877.

2.3 manager 项被编译成少量 fused kernel

MotrixLab 将 command/reward/termination/observation/reset 编译成三类 kernel：

• evaluate: codegen.py:46-85;
• observe: codegen.py:87-101;
• reset: codegen.py:103-130.

路径前缀为 motrix_env_core/src/motrix_env_core/numba/manager/compiler/。

编译选项：

numba.njit(cache=True, nogil=True, parallel=True)

证据：compiler.py:987-997。

效果：

• 一个 kernel 处理所有 reward terms;
• 一个 kernel 处理所有 termination terms;
• command update/advance 也在同一 evaluate kernel;
• observation 是另一个 fused kernel;
• reset kernel 只遍历 reset rows;
• 没有 UniLab 这种每个 manager term 一次 Python function call / validation / timing boundary。

2.4 Numba 调度器针对短 kernel 优化

MotrixLab 明确优先加载 TBB，因为默认 OpenMP 在短 kernel 上有 thread-wakeup storm：

• motrix_env_core/src/motrix_env_core/numba/manager/env.py:344-353;
• motrix_env_core/src/motrix_env_core/numba/threading.py:4-17.

对 G1 WBT 这种每个 lane 计算量不大、但每步都要跑的 kernel，调度开销很关键。

2.5 reset 是真正的 selected-row

MotrixLab 区分：

1. episode reset;
2. sim-only rematerialization，例如 motion clip wrap。

证据：motrix_env_core/src/motrix_env_core/numba/manager/env.py:700-752。

行为：

• episode rows 执行 command/action lifecycle reset;
• sim-reset-only rows 不执行 episode lifecycle reset;
• 两类 rows 取 union;
• 一次 reset kernel;
• 一次 selected-row simulator reread.

reset kernel 的循环是：

for row in numba.prange(env_ids.shape[0])

而不是 num_envs：

• codegen.py:110-130.

reset outputs 使用 compact prefix，scatter 到 selected rows，再执行一次 native write program：

• motrix_env_core/src/motrix_env_core/numba/manager/sim_reset.py:96-109.

reset 后只做 selected-row reread：

• motrix_env_core/src/motrix_env_core/numba/manager/env.py:680-698.

结论：MotrixLab reset 的工作量与 reset 行数相关，而不是与 num_envs 相关。

2.6 motion 数据和输出 buffer 都是持久连续布局

Motion corpus 是 frame-major contiguous arrays，所有 lane 共享：

• motrix_envs/src/motrix_envs/motion/tracked.py:52-96.

WBT command 的 per-env 状态只有 frame index、clip end、target buffers：

• motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py:95-120.

reset 时 joint/root/body state 直接按 frame index gather：

• motrix_envs/src/motrix_envs/locomotion/wbt/mdp/reset.py:33-180.

observations 使用两个预分配 slot，通过 swap 保证 collector 保留 obs[t]，下一 kernel 写 obs[t+1]：

• motrix_env_core/src/motrix_env_core/numba/manager/env.py:304-338;
• 初始化 :490-515;
• step 前 swap :530-535.

2.7 collector 是本地推理 + 异步 ring

FastSAC collector 拥有本地 read-only actor/normalizer：

• motrix_rl_builtin/.../async_impl/collector.py:59-76.

推理路径：

• torch.compile(..., mode="reduce-overhead");
• pinned obs/action staging;
• optional AMP;
• 一次 H2D、一次推理、一次 D2H、一次必要同步。

证据：collector.py:135-160, :201-221.

collector 每步只做 action sample、env.step、ring push、vectorized episode bookkeeping、current obs 替换：

• collector.py:247-335.

它不会像 UniLab 当前 collector 一样每 tick 都等 learner 返回 action。不过 UniLab 实测 learner_action_wait_ms 平均约 1.1 ms，这不是最大瓶颈；本地 actor 属于架构增强，不是第一优先修复点。

---

3. UniLab 为什么慢

以下路径均在 src/unilab/ 下。

3.1 mjwarp 的 update_state_ms 包含 GPU queue drain

TorchEnv.step 中：

• backend step: base/torch_env.py:263-285;
• update state: base/torch_env.py:290-306.

代码自己也说明 termination 前的 queue drain 会排空 pre-update GPU work：

• envs/manager_based_rl_env.py:1240-1249.

因此 mjwarp 常见现象是：

• backend_ms 很小；
• update_state_ms 很大;
• 实际包含 pending physics kernel completion、sensor projection、manager reduction、首次 host sync.

不能简单解释成 “update state manager 代码本身慢”。

3.2 manager 生命周期是通用、分层、多 Python boundary

ManagerBasedRlEnv.update_state 依次执行：

• read boundary refresh;
• termination;
• reset flags;
• reward;
• metrics;
• events;
• command preflight;
• command compute;
• command read refresh;
• command post compute;
• observations;
• mapping/publication.

证据：envs/manager_based_rl_env.py:1213-1345.

每一层还有 timing、validation、manager dispatch。对 2048 env、大量小 tensor 操作来说，Python/Torch dispatch 和 kernel launch overhead 明显。

3.3 HOST_BRIDGE backend 默认可能把 manager carrier 放到 CUDA

training.sim_carrier_device=auto 时，如果 learner 是 CUDA，HOST_BRIDGE backend 会选择 CUDA Torch carrier：

• base/process_device.py:169-215;
• training/tensor_runtime.py:30-43.

这对 MuJoCo/Motrix 意味着：

• CPU physics;
• CPU ↔ GPU packed transfer;
• Torch manager compute;
• 某些 bool/all/any 检查触发 host sync.

MotrixLab 的对应 WBT path 全程 CPU-resident，只在 inference 处使用 GPU。

注意：简单设置 training.sim_carrier_device=cpu 未必自动变快，因为 UniLab CPU Torch manager graph 也可能有很多小 kernel；但它是一个必要的对照实验。

3.4 motion command 存在重复 host 同步

MotionCommand.command 每次访问都会：

torch.all(torch.isfinite(self._command), out=self._command_finite_valid)
if not bool(self._command_finite_valid):
    self._refresh_motion_torch()

证据：tasks/motion_tracking/common/manager_terms.py:851-856.

command 会在多处访问：

• command manager compute/validation: managers/command_manager.py:383-388;
• observation: envs/mdp/observations.py:401-408;
• actor/critic 都有 command observation: conf/flashsac/task/g1_motion_tracking/base.yaml:65-125.

因此一个 step 可能多次执行 isfinite(...).all() -> bool()，每次都可能排空 CUDA 队列。

3.5 reset 行数少，但 motion reset 计算很多是 full-width

实际 benchmark 中平均只有 67/2048 行 reset，但 _resample_command 会：

• reset_mask.any() host sync: manager_terms.py:957-963;
• full-width frame sampling;
• full-width motion packet gather;
• full-width pose/velocity/joint noise RNG;
• full-width reset state construction;
• fixed-width mask write;
• full-width motion packet merge.

证据：manager_terms.py:957-1023，full-width merge :1211-1275.

sampler 本身也是 full-width：

• tasks/motion_tracking/common/tensor_sampler.py:218-298;
• adaptive path 使用 torch.multinomial(probabilities, self.num_envs, ...).

UniLab 已经有 selected-row scatter kernel：

• manager_terms.py:269-318;
• row-oriented ingest :1181-1209.

说明 row-scoped reset 并非没有基础设施，而是当前 hot path 没有完全用上。

3.6 reset 后 full read，尽管已有 selected read API

UniLab reset 后总是：

read_plan.refresh()

证据：envs/manager_based_rl_env.py:1602-1607.

但 SceneTensorReadPlan 明确提供：

• refresh(): full batch;
• refresh_selected(): selected post-reset rows.

证据：base/entity.py:339-360.

MuJoCo/Motrix host bridge 已实现 selected packed read：

• third-party/unisim/src/unisim/backend/mujoco/tensor_io.py:430-452;
• third-party/unisim/src/unisim/backend/motrix/tensor.py:1058-1093.

生产路径目前没有调用 refresh_selected()。这是低风险、直接可借鉴 MotrixLab 的点。

3.7 autoreset 和 reset commit 有多次 mask/count/host sync

TorchEnv._reset_done_envs()：

• done.any() 转 bool: base/torch_env.py:503-508;
• full-width final observation copy: :513-516;
• torch.nonzero 生成 row IDs: :521-527.

Manager reset 又把 row IDs scatter 回 full-width mask：

• envs/manager_based_rl_env.py:1498-1522.

backend packed reset 最终仍消费 fixed-width mask。

链路中存在：

• mask -> bool;
• mask -> rows;
• rows -> mask;
• count -> .item();
• reset finite/row validation;
• selected payload D2H.

可以保留 public env_ids 契约，但内部 carry 一个 reset mask/count verdict，减少重复同步和转换。

3.8 reset observation full-width 且合并两次

reset 后：

1. 计算 full-width reset observations;
2. merge 到 _state.obs;
3. 又 merge 到 obs_buf.

证据：envs/manager_based_rl_env.py:1609-1642.

代码注释说明 publication intentionally stays full-width。正确性优先，但性能上可以改成一次 authoritative masked write，另一个 view alias 它，或让 observation manager write-through。

---

4. 共通的可借鉴原则

虽然架构不同，但以下原则是共通的。

4.1 一条边界，一次数据交换

MotrixLab：

• one full simulator read per step;
• one selected read per reset;
• one native write per reset.

UniLab 已经有 packed read plan 和 selected read API，应更严格执行同样的边界纪律。

4.2 selected reset 就应该 selected compute

MotrixLab reset kernel 只遍历 reset rows.

UniLab backend selected reset 已经比较快，但 motion sampling/RNG/gather/merge/read/observation 仍是 full-width。这是最明显的借鉴点。

4.3 fused owner graph

MotrixLab 把 manager terms 编译成 evaluate/observe/reset 三个 kernel.

UniLab 不一定要照搬 Numba，但可以为目标 owner 引入 compiled/fused fast path：

• Torch CUDA graph / compiled module;
• CPU host-bridge owner graph 编译;
• 至少把 command/reward/termination 的 per-term Python dispatch 合并成 owner-owned fused kernel.

4.4 持久 buffer 和稳定 ABI

MotrixLab：

• persistent flat inputs;
• persistent kernel outputs;
• double-buffered observations;
• contiguous motion table.

UniLab 已经大量使用 scratch buffers，但仍有 reset/full obs merge、selector conversion、timing dict 等额外对象和拷贝。

4.5 collector 与 learner 解耦

MotrixLab 使用 collector local actor + async SPSC ring.

UniLab 当前依赖 learner inference tick，但实测 wait 约 1 ms，优先级低于 reset/read/motion 优化。

---

5. 建议优化路线

P0：先补齐 attribution，再动手

现有 benchmark CSV 已有 reset detail keys，但生产 TensorBoard 只导出了部分。建议先用：

scripts/benchmark/rl/benchmark_offpolicy_collector_active.py

对同一任务跑 mjwarp/mujoco/motrix，并记录：

• reset_done_motion_sampler_ms;
• reset_done_motion_packet_ms;
• reset_done_motion_reset_rng_ms;
• reset_done_motion_reset_values_ms;
• reset_done_motion_reset_write_ms;
• reset_done_command_refresh_ms;
• reset_done_observation_ms;
• update_state_command_*;
• update_state_observation_*;
• transfer counts / synchronization counts.

当前 _resample_command 只记录总的：

• reset_done_motion_reset_values_ms.

见 manager_terms.py:1021-1023.

而 schema 已定义很多细分 key：

• base/backend_timing.py:10-97.

先把 reset 的 3 ms 拆开，避免猜错热点。

P0：motion command finite validity 只验证一次

建议：

• 在 _refresh_motion_torch() 发布 carrier 时计算/发布 validity;
• 一个 phase 内 getter 只读 cache;
• observation getter 不再重复 isfinite(...).all();
• debug 模式保留 per-access strict check.

预期收益：

• 减少 update_state 中的多次 GPU reduction/host sync;
• 让 mjwarp 的 phase attribution 更真实.

风险：

• 必须保持 invalid carrier 的 recovery 语义.

P0：HOST_BRIDGE reset 后使用 refresh_selected()

这是最直接的 MotrixLab 对齐点：

• MuJoCo/Motrix 已有 selected packed read;
• SceneTensorReadPlan.refresh_selected() 已存在;
• 测试已有局部覆盖;
• device-resident backend 可以继续 normal refresh.

预期收益：

• 对 mujoco/motrix 明显;
• 对 mjwarp 无直接收益.

风险：

• 需要确认所有 backend plan 都正确记录最近 reset rows;
• backend 必须公开能力，不能探测 private 方法.

P0：row-scoped motion reset

保留 backend fixed-width mask contract，但内部：

• compact reset rows;
• sampler 只 sample k = reset_count 帧;
• RNG shape 改为 selected rows;
• motion packet gather 只 gather selected rows;
• reset values 只构造 selected rows;
• 使用已有 selected-row ingest kernel scatter 到 full-width carrier.

这是和 MotrixLab reset kernel 最等价的方向。

风险：

• RNG stream 语义变化，需要 motion parity/reset determinism 测试;
• 相关测试包括：
  • tests/tasks/test_motion_term_parity.py;
  • tests/tasks/test_tensor_motion_sampler.py;
  • tests/tasks/locomotion/g1/test_g1_reset_owner.py.

P1：carry reset mask/count 一次通过生命周期

统一内部 reset verdict：

mask + count + rows(optional) + backend-needed representation

public reset 继续使用 env_ids，backend 继续使用 fixed-width mask，但转换只做一次。

P1：final observation 和 reset observation 改成 selected-row publication

MotrixLab 只为需要 terminal observation 的 rows 捕获 compact snapshot.

UniLab 当前只要 any done，就 copy 全部 obs groups：

• base/torch_env.py:510-516.

建议：

• replay 继续接受 full-width terminal carrier，或改为 compact contract;
• env 内部只 copy selected rows;
• reset obs 一次 masked write，避免 _state.obs 和 obs_buf 双重 merge.

如果改成 compact terminal observation，需要同时更新 replay buffer/contract 和测试。

P1：device-resident sensor refresh 增加 batched public backend capability

当前 device-resident read plan 会逐 sensor 调 get_sensor_view：

• base/entity.py:648-668.

mjwarp 每次 tracked-body refresh 又可能跑多个 kernel。建议通过公开 SimBackend interface 增加批量 sensor-view refresh capability，一次提交所有需要的 sensor projections.

注意架构约束：必须先扩展 public backend interface，env/training 不能调用 backend private 方法。

P1：HOST_BRIDGE owner graph CPU/compiled fast path

对 MuJoCo/Motrix：

• CPU physics;
• selected reset;
• CPU-resident manager state;
• fused CPU manager kernel;
• 控制和 reset payload 批量 crossing.

可以借鉴 MotrixLab 的 Numba compiled owner graph，但不一定要复制实现。关键是消除每个 term 的 Python/Torch dispatch.

P2：collector local actor / ring 改造

MotrixLab 的本地 actor 架构值得借鉴，但不是第一优先：

• UniLab 当前 learner wait 约 1 ms;
• 本地 actor 会引入 policy staleness/version consistency 语义;
• 需要明确 weight snapshot 一致性策略.

可选机制：

• flat parameter buffer;
• versioned/seqlock weight publication;
• fixed-shape compiled actor;
• pinned staging;
• CUDA IPC transition ring.

P2：successor-slot replay，避免重复 next_obs

MotrixLab ring 不单独存 next_obs，由下一 slot 推导，拷贝量近似减半。

UniLab replay 当前每步传：

• obs;
• next_obs;
• critic;
• next_critic.

见 rl/runtime/offpolicy/worker.py:843-856.

这是 replay contract 改造，需谨慎处理 terminal rows 和 bootstrap masks.

---

6. 不建议直接照搬的地方

6.1 “CPU carrier 就一定快”不成立

MotrixLab 快是因为 CPU + native compiled manager graph，不是单纯因为 CPU。UniLab CPU Torch manager 如果仍是大量小 op 和 Python dispatch，切换 sim_carrier_device=cpu 可能不变甚至变慢。

建议只作为对照实验，不作为最终方案。

6.2 clip-end semantics 不同

UniLab G1 motion SAC/FlashSAC 当前：

truncate_on_clip_end: true

并配置为 timeout termination：

• conf/flashsac/task/g1_motion_tracking/base.yaml:162-169.

MotrixLab 支持 motion clip wrap 触发 sim-only rematerialization，不是 episode end：

• motrix_envs/src/motrix_envs/locomotion/wbt/mdp/command.py:235-290;
• motrix_env_core/src/motrix_env_core/numba/manager/env.py:700-752.

这会改变：

• episode boundary;
• replay bootstrap;
• action history reset;
• episode metrics;
• reward stream.

不能当纯性能优化直接搬，需要任务 owner/ADR 决策。

6.3 硬件和任务不完全可比

MotrixLab benchmark 在 192-thread EPYC 上跑，UniLab 本机是 9950X3D2 16 core。任务 dt、solver、DR、motion corpus、算法阶段也不同。

phase 时间可以直接启发架构，但 throughput 数字不能机械对比。

---

7. 推荐落地顺序

1. 补齐 reset/update timing attribution：低风险，先确认 3 ms reset、4 ms update 的真实构成。
2. motion command finite validity cache：低/中风险，直接减少 update_state host sync。
3. HOST_BRIDGE reset 后调用 refresh_selected()：低风险，现有能力/测试已经存在。
4. row-scoped motion reset sampling/gather/ingest：中风险，预期对 reset 收益最大。
5. reset mask/count verdict 复用：中风险，减少重复 nonzero/any/item。
6. selected final obs + single masked reset obs publication：需要协调 replay contract。
7. backend public batched sensor refresh capability：主要针对 mjwarp device-resident path。
8. HOST_BRIDGE compiled owner graph：投入最大，但最接近 MotrixLab 的长期解法。
9. collector local actor / successor-slot ring：提升吞吐，但优先级低于 env reset/read/motion。

---

审计方法与限制

• 只读审计，未修改文件。
• 使用 rg/nl -ba 检索代码路径，jq 解析 run config 与 benchmark JSON，uv run python 聚合 TensorBoard scalar，git log/show 追溯 reset timing 与 env_ids selector 演进。
• 未运行完整 benchmark；所有性能数字来自 checked-in artifact 或已有日志。
• MotrixSim native query/write program 的内部实现不可直接从仓库审计，只能审计其 Python 调用与数据布局。
• 本机 mujoco/motrix 短 run 样本较少，可能受启动/策略阶段影响，只作为方向性参考。