Skip to content

perf: Motrix reset 热路径 Rust/motrixsim 边界优化(#679 延伸) #682

Description

@TATP-233

目标

#679 已完成 Python 侧 (owner-layer) 优化 + 完整 sub-timing 报告的基础上,继续攻 Motrix `set_state` 里剩余 82%+ 的 Rust 边界开销,把 `g1_walk_flat` 上 `dr_reset_set_state_ms` 从当前的 9.46 ms 再压 30-50%。

背景

#679 已经把 `SimBackend.set_state` contract 扩展为返回 `{"timing": {...}}`,Motrix backend 16 个 sub-key 在远端 Xeon 8568Y+ 上完整落地(详见 #679 comment 4905826906)。该 comment 里 `g1_walk_flat numba_accelerated` 的实测 breakdown 为:

sub-key ms 占 set_state 归属
FK (`forward_kinematic`) 4.561 48.2% Rust 边界
Data reset (`data_slice.reset(model)`) 1.691 17.9% Rust 边界
Data slice (`self._data[mask]`) 0.909 9.6% Rust 边界
Actuator ctrl 0.814 8.6% Python + Rust
Refresh pose (`get_link_poses`) 0.746 7.9% Rust 边界
Set dof vel/pos 0.662 7.0% Rust 边界
QPos convert + Mask + 其它 Python 侧 ≤ 0.05 ≤ 0.5% Python (已优化到极限)
Total `dr_reset_set_state_ms` 9.461 100%

关键观察:Python 侧攻击面在 #679 已经压到 <0.5%(Mask 0.1% + QPos convert 0.3%),剩下 82%+ 全在 Rust/motrixsim 边界。要继续压 reset_done,必须走 motrixsim 侧改动 + backend 接口重构。

范围

按 sub-key 占比排序的 3 条主线:

主线 A(P0):Rust 侧 fused set_state

把当前 `set_state` 里对 `data_slice` 的多次独立 Rust 调用合并成一个 motrixsim-side batched 调用。当前是:
```
data_slice.reset(model) # 1.691 ms
data_slice.set_dof_vel(qvel) # 0.349 ms
data_slice.set_dof_pos(qpos, model) # 0.313 ms
data_slice.actuator_ctrls = ctrl # 0.814 ms
model.forward_kinematic(data_slice) # 4.561 ms
```
共 7.7 ms 分 5 次穿越 Python-Rust 边界。fused API 期望:`model.reset_and_step_kinematics(data, env_ids, qpos, qvel, ctrl)` 一次穿越完成。

预期收益:合并 4-5 次边界穿越 + 一次拿到所有 buffer 一起 batched 处理,g1_walk_flat 上从 7.7 ms 降到 3-4 ms(-40 到 -50%)。

依赖:motrixsim 侧新增 fused 接口;本仓库 `SimBackend.set_state` contract 保持不变,只在 Motrix 后端切实现。

主线 B(P1):index-based `forward_kinematic`

FK 单项 4.56 ms 是 `g1_walk_flat set_state` 里最大项。当前是对 `data_slice = self._data[mask]` 执行 FK;motrixsim 内部仍走完整 batch 计算。期望:`model.forward_kinematic_partial(data, env_ids)` 只对 done 的 env 子集算,其余 env 跳过。

预期收益:当前每次 reset 只涉及 ~500-2000 env(reset_ratio 5-25%),full-batch FK 计算浪费 75-95% 的 env。若 motrixsim 侧支持 index-based FK,FK 时间线性缩放到 reset_ratio × current,g1_walk_flat `4.56 → ~0.5-2 ms`。

依赖:motrixsim 侧 FK API 支持 env-index subset;如现有 `data[mask]` 索引后 FK 就已经只对子集算,则本项 no-op —— 需要先用 motrixsim profiler 确认 FK 内部是否已经跳过 mask=False 的 env。

主线 C(P2):lazy / batched `get_link_poses` refresh

当前 `set_state` 尾部无条件调 `self._refresh_link_pose_cache(env_indices, data_slice=data_slice)` (0.746 ms),而下一个 step 的 `update_state` 才真正读 link_poses。可以:

  • 把 `_link_poses` 改成 lazy:`set_state` 只标记 dirty,`update_state` 或 `get_link_pos_w` 第一次读时才 refresh。
  • 或者把 refresh 合并进主线 A 的 fused kernel(一次 Rust 调用同时算 FK + link_poses)。

预期收益:0.5-0.7 ms,取决于是否能与下一 step 的 physics/update_state 隔行重叠。

主线 D(可选,P2):`data_slice.reset()` 精简

`data_slice.reset(model)` 1.691 ms 里,motrixsim 目前会重置 data 的所有字段(应用 gravity / joint range / actuator gain 等 override)。对 g1_walk_flat 这种无 DR 或极简 DR 的 task,可以支持"skip `data.reset()`, only overwrite qpos/qvel/actuator_ctrls"的快路径。需要 motrixsim 侧新增 `data_slice.reset_light(model)` 或类似 API。

预期收益:0.5-1.0 ms(依 DR 配置)。

非目标

完成标准

  • 主线 A 或 B 至少一条落地,且 `g1_walk_flat numba_accelerated` 远端 Xeon 8568Y+ 8192 env A/B 上 `dr_reset_set_state_ms` 相对 perf: 深度优化 reset_done/backend set_state 热路径 #679 后的 baseline (9.461 ms) 下降 ≥ 30%(目标 ≤ 6.5 ms)。
  • `reset_done_ms` 同口径下降 ≥ 25%(目标 ≤ 10 ms)。
  • 新接口在 `SimBackend` 抽象层可选(`get_dr_capabilities` 类似的能力探测),MuJoCo 后端不受影响。
  • #679 引入的 set_state 16 个 sub-key breakdown 出对比报告,证明收益点确实落在 FK / Data reset / Data slice 上。
  • `make test-all` 通过;`tests/base/test_sim_backend_set_state_timing.py` 的 schema 契约不变。

前置工作

  • 与 motrixsim 团队对齐上述 A/B/C/D 4 条 API 需求,评估哪一条能在本 issue 周期内到位。
  • 用 motrixsim 侧 profiler(如果有)先确认 FK 是否真的跑完整 batch —— 若已经按 mask 跳过,则主线 B 直接降级为 no-op,主线 A 成为主要攻击面。
  • 保持 `benchmark/benchmark_motrix_set_state_ab.py` 作为 microbenchmark 入口;本 issue 落地时新增 fused-kernel 变体到该脚本。

Pointers

  • Motrix `set_state` 现状:`src/unilab/base/backend/motrix/backend.py` L592-`set_state`(instrumented + Python-side optimized)
  • SimBackend contract:`src/unilab/base/backend/base.py` L267-`set_state` 已定义为 `dict | None` 返回
  • DR 合并 backend 计时:`src/unilab/dr/manager.py` L51-L67
  • Benchmark 表格:`benchmark/benchmark_offpolicy_collector_active.py`
  • 微基准:`benchmark/benchmark_motrix_set_state_ab.py`

Refs: #679, #679 comment 4905826906, #665, PR #681

Activity

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

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions