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

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Fields

No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions