目标
在 #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
目标
在 #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 为:
关键观察: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。可以:
预期收益: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 配置)。
非目标
完成标准
前置工作
Pointers
Refs: #679, #679 comment 4905826906, #665, PR #681