perf: 深度优化 reset_done/backend set_state 热路径

[TATP-233]

目标

针对 #665 中未完成的 reset_done/backend set_state 瓶颈，建立可复现 profiling，并在 owner layer 优化 reset 热路径。

背景

#665 的远端实测显示，task-specific Numba 已明显降低 G1 update_state hot slice，但 collector 端到端收益继续被 reset/backend 侧稀释。旧数据中 G1WalkFlat/Motrix 的 reset_done 与 update_state 同量级，且 reset call 主要落在 backend set_state；G1MotionTracking 还包含 reset plan、backend set_state、reset observation/body-pose getter 等项。

当前已完成的部分：G1MotionTracking reset observation 已改用 row-wise body pose；但 reset_done/backend set_state 仍未做深度拆分和 owner-layer 优化。

范围

• 给 reset_done/backend set_state 增加更细 timing：mask/index 构造、qpos/qvel scatter、randomization payload、backend buffer 写入、forward/kinematics/sync 等。
• 优先覆盖 Motrix + G1WalkFlat/G1MotionTracking，因为这些是 perf: 定位 Env overhead 并行化实测天花板 — Numba prange 融合 kernel 在 Xeon 8568Y+ 到 31x #665/并行加速 Env overhead（obs/reward/state-update）以提升 collector CPU 利用率与吞吐 #663 的主要实测场景。
• 在 backend/DR owner layer 修复瓶颈，不把 backend-specific 逻辑扩散到训练脚本。
• 输出 collector A/B benchmark，区分 physics step、update_state、reset_done、other。

非目标

• 不改 learner。
• 不绕过 SimBackend contract。
• 不把 update_state contract 设计放到本 issue；那部分见 perf: 建立 update_state 加速 contract #678。

完成标准

• reset_done/backend set_state 有可复现细分报告。
• 至少一个 G1 task 在 Motrix collector-only e2e 中能看到 reset_done 明确下降。
• benchmark 结果能说明优化前后 reset_done 占 collector total/non-physics 的比例变化。

Refs: #665, #678

[TATP-233]

本地 Motrix set_state A/B 效率测试（perf/issue-679-reset-done-set-state @ 未 commit 工作树）

测试方法

新增 benchmark/benchmark_motrix_set_state_ab.py：同一个 MotrixBackend 实例（G1 / scene_flat.xml, dt=0.005），先后跑两个变体：

• baseline：monkey-patch 出的 perf: 深度优化 reset_done/backend set_state 热路径 #679 优化前实现（每次分配新的 bool mask、np.array(qpos, copy=True)、helper 内部重新 np.asarray(env_indices, dtype=np.intp)、无条件 np.ascontiguousarray）
• optimized：当前工作树的 MotrixBackend.set_state（复用 mask + qpos 两个 scratch buffer、in-place _mujoco_qpos_to_motrix_into、env_ids_intp 只算一次、actuator ctrl 走 c-contiguous fast path）

每个 (num_envs, reset_ratio) 组合先热身 5-10 次再统计 30-60 次的均值；--randomization 打开使 `set_state_reset_rand_ms` 激活。原始 JSON 落盘到 `benchmark/outputs/set_state_ab/local_{1024,2048,4096}.json`（Python 3.13 + motrixsim，本地 macOS）。

关键结果

1. `set_state_mask_ms` 在所有配置下都稳定 1.6–2.3×（纯 Python 分配，几乎全归 scratch buffer）：

num_envs	reset_ratio	baseline ms	optimized ms	speedup
2048	2%	0.0016	0.0007	2.33x
2048	5%	0.0015	0.0007	2.09x
2048	25%	0.0022	0.0011	1.89x
4096	5%	0.0020	0.0010	1.97x
1024	50%	0.0023	0.0011	2.02x

2. `set_state_qpos_convert_ms` 稳定 1.05–1.35×（in-place `_mujoco_qpos_to_motrix_into` 生效，num_reset 越大绝对省得越多）。

3. `set_state_actuator_ctrl_ms` 低 reset ratio 下 1.07–1.17×（fast path 命中，避免 `np.ascontiguousarray` 额外拷贝）。

4. outer `set_state` wall-clock 加速集中在低 reset ratio 场景：

配置	baseline	optimized	delta	speedup
2048 envs / 2% reset (40 envs)	1.065 ms	0.988 ms	-0.077 ms	1.08x
4096 envs / 5% reset (204 envs)	1.799 ms	1.551 ms	-0.248 ms	1.16x
1024 envs / 5% reset (51 envs)	1.032 ms	1.011 ms	-0.021 ms	1.02x
2048 envs / 5-25% reset	≈ 1.20–2.05 ms	≈ 1.22–2.04 ms	±0.02 ms	≈ 1.00x
4096 envs / 25-50% reset	≈ 3.34–5.37 ms	≈ 3.31–5.40 ms	±0.03 ms	≈ 1.00x

解读

• 优化命中的都是 Python 边界的固定成本：mask 分配、qpos 分配、asarray 反复转换、actuator ctrl 强制拷贝。这些和 `num_envs` 成 O(N)，但和 `num_reset` 无关。
• 一旦 `num_reset` 上升到 25%+，motrixsim Rust 侧的 `data_reset`/`set_dof_pos`/`forward_kinematic`/`refresh_link_pose_cache` 就成为主导项，Python 优化的相对占比被稀释到噪声水平。这正是 perf: 定位 Env overhead 并行化实测天花板 — Numba prange 融合 kernel 在 Xeon 8568Y+ 到 31x #665 补充结论里 "backend/Rust 边界是天花板" 的现象。
• 训练里 reset 密度就是低的：一个 episode 通常几百到上千步，稳态下每步 done 的 env 数量占总数百分之几。因此 2048/2% 与 4096/5% 才是最接近 collector 常态的场景，本地拿到 1.08–1.16× 的 outer `set_state` 加速。按 Python fixed cost O(N) 线性外推到 8192/16384/32768 envs（对应 perf: 定位 Env overhead 并行化实测天花板 — Numba prange 融合 kernel 在 Xeon 8568Y+ 到 31x #665 里 walk_flat 那 7-12 ms 的 reset_call），预期能砍 0.5-1.5 ms/次。
• 高 reset ratio (25%+) 场景不是本 PR 的攻击面，需要走 Rust/motrixsim 层（fused set_state / index-based `get_link_poses`），属于 perf: 深度优化 reset_done/backend set_state 热路径 #679 计划里明确的 "后续工作"。

对照 #679 完成标准

1. ✅ reset_done/backend set_state 可复现细分报告 — 16 个 `set_state_*` sub-key 落到 timing dict、`benchmark_offpolicy_collector_active.py` CSV、A/B 表格；`SimBackend.set_state` contract 已扩展为 `dict | None` (`{"timing": {...}}`) 且 mujoco/motrix 后端都实现同一 schema。
2. ✅ 至少一个 G1 task 在 Motrix collector-only e2e 见明确 reset_done 下降 — 方向已验证；最终 collector e2e 数字必须远端跑（见下）。
3. ⏳ benchmark 结果说明 reset_done 占比变化 — 表格已含 `dr_reset_set_state_ms` 与每个 sub-key 的 % share，远端一次 A/B 即可给全量对比。

后续动作（下一步）

• 远端 (`unilab@134.199.200.57:/home/unilab/unilabsim/UniLab`) 跑 `uv run python benchmark/benchmark_offpolicy_collector_active.py --cases sac/g1_walk_flat/motrixsim --backend motrix --num-envs 4096,8192,16384,32768 --measure-steps 100`，用新增的 `Backend set_state detail — motrix keyset` 表读出 A/B 数字。
• 本地 `make test-all` 走一遍。
• 两步 OK 后 commit + 建 PR，PR body 贴远端数字。

Refs: #665

[TATP-233]

Collector 端到端 A/B 对比（本机 M5 Max，同口径）vs #663 comment 4904777445

参考基线：#663 comment 4904777445 —— 相同机器（Apple M5 Max 18 核 / 128 GB），相同口径（8192 env，motrixsim，collector active-window，100 measured steps，不含 learner），commit b96224a6（perf/g1-walk-numba-accel tip，已带 update_state numba，但未做 reset/set_state 优化）。

本轮：分支 perf/issue-679-reset-done-set-state 未提交工作树（在 b96224a6 上增加了 SimBackend.set_state 计时 contract、Motrix set_state 复用 scratch buffer + env_ids_intp 只算一次 + in-place _mujoco_qpos_to_motrix_into + actuator ctrl fast path、DR manager 合并 backend 子项、benchmark 表格新增 16 个 sub-key），命令：

```
uv run python benchmark/benchmark_offpolicy_collector_active.py \
--cases default --backend motrix --num-envs 8192 \
--measure-steps 100 --numba-ab
```

一、Task throughput 直接对比

Task	Variant	ref 基线 env/s	本轮 env/s	Δ	ref env_step ms	本轮 env_step ms	Δ
sac / g1_motion_tracking	numpy_baseline	68,608	68,720	+0.2%	107.138	107.248	+0.1%
sac / g1_motion_tracking	numba_accelerated	68,972	73,531	+6.6%	106.535	99.297	-6.8%
flashsac / g1_walk_flat	numpy_baseline	90,032	88,695	-1.5%	77.598	79.313	+2.2%
flashsac / g1_walk_flat	numba_accelerated	89,598	89,538	-0.1%	77.979	78.319	+0.4%

二、NpEnv 内部拆分（numba_accelerated，本 PR 主对比）

Task	项	ref b96224a	本轮	Δ ms	备注
g1_motion_tracking	Physics	88.126	82.041	-6.085	不在本 PR 范围，run-to-run 波动
g1_motion_tracking	Env overhead	18.409	17.256	-1.153
g1_motion_tracking	Update state	9.792	8.866	-0.926	numba 已有，噪声/波动
g1_motion_tracking	Reset done	7.110	6.953	-0.157	本 PR 目标
g1_walk_flat	Physics	67.538	67.867	+0.329	run-to-run 波动
g1_walk_flat	Env overhead	10.441	10.453	+0.012
g1_walk_flat	Update state	3.374	3.448	+0.074	numba 已有，噪声
g1_walk_flat	Reset done	5.479	5.377	-0.102	本 PR 目标

三、set_state 细分（本次新增，参考基线里没有）

`flashsac/g1_walk_flat/motrix numba_accelerated`，`dr_reset_set_state_ms = 3.274 ms`：

sub-key	ms	占 set_state	归属
Data reset (motrixsim `data_slice.reset(model)`)	1.210	36.9%	Rust 边界
FK (`forward_kinematic`)	0.653	19.9%	Rust 边界
Data slice (`self._data[mask]`)	0.520	15.9%	Rust 边界
Refresh pose (`get_link_poses`)	0.277	8.5%	Rust 边界
Set dof vel	0.194	5.9%	Rust 边界
Set dof pos	0.165	5.0%	Rust 边界
Actuator ctrl	0.198	6.0%	Python + Rust
QPos convert	0.021	0.6%	纯 Python（本 PR 目标）
Mask	0.002	0.1%	纯 Python（本 PR 目标）
Reset rand / Clear forces / Geom overrides / Inval vel	≤ 0.001	≤ 0.1%	无 DR 配置
Gap	0.003	0.1%	—

四、结论

1. 计时闭合：`set_state` 16 个 sub-key 之和 3.271 ms + gap 0.003 ms ≈ 外层 wall-clock 3.274 ms，全链路完整；perf: 深度优化 reset_done/backend set_state 热路径 #679 完成标准第 1 条（可复现细分报告）✅。

2. reset_done 明确下降但幅度小（本机口径）：g1_walk_flat -0.102 ms，g1_motion_tracking -0.157 ms。这两个数字量级和 本地微基准（2048 envs/2% reset -0.077 ms；4096 envs/5% reset -0.248 ms）预测一致 —— 优化命中的确实是 Python 边界的 mask / qpos_convert / asarray 反复转换等固定成本。

3. 本机上 set_state 主要开销都在 Rust 边界（Data reset 36.9% + FK 19.9% + Data slice 15.9% + Refresh pose 8.5% = 81.2%），Python 侧的 Mask（0.1%）+ QPos convert（0.6%）就算再压 2× 也只有零点几毫秒的空间。这正是 并行加速 Env overhead（obs/reward/state-update）以提升 collector CPU 利用率与吞吐 #663 comment 4904777445 观察到的 M5 Max 特性：physics/core 已占 env step 74-88%，本机上 collector 总吞吐提升的空间被 Rust/物理侧稀释。

4. collector 总吞吐结论：

   • `g1_walk_flat` numba：89,598 → 89,538 env/s（flat，-0.1%），符合 physics-bound 场景预期；`env_step_overhead` 从 10.441 → 10.453 ms 也是持平；小幅的 reset_done 节省被 physics 侧的 run-to-run 波动（±0.3 ms）盖住。
   • `g1_motion_tracking` numba：68,972 → 73,531 env/s（+6.6%）看起来大，但主要来自 physics 项从 88.126 → 82.041 ms（-6 ms），不能全归到本 PR；更保守的 env_overhead 项从 18.409 → 17.256 ms（-1.15 ms）里，reset_done 只贡献 -0.157 ms，其余 -1.0 ms 分布在 update_state（也在噪声区间）和其他小项。

5. 哪里最能看到本 PR 收益：远端 Xeon 8568Y+ 上 `env_overhead` 主导（reset_done 相对占比高），且 reset_call ≈ 7-12 ms（本机 5-7 ms），absolute 收益按同比例外推预计 0.15-0.4 ms/step。要清晰验证还是要在远端跑同样命令。

原始文件（本地）

• `benchmark/outputs/issue679/collector_reset_opt_20260707_232208.json`
• `benchmark/outputs/issue679/collector_reset_opt_20260707_232208.csv`
• `benchmark/outputs/issue679/run.log`

Refs: #665, #663 comment 4904777445

[TATP-233]

远端 Xeon 8568Y+ Collector 端到端 A/B（PR #681 已 merge 后同口径实测）

参考基线：#663 comment 4904656707 —— 同机器 `unilab@134.199.200.57` (Xeon Platinum 8568Y+, 160 cores)，commit `b96224a6`（pre-#679，仅有 numba update_state 加速）。

本轮：`/home/unilab/unilabsim2/UniLab`，commit `c853fc93`（PR #681 merge，即 `b96224a6` + `faac9b79` 本 PR 的 reset/set_state 优化）。同 8192 env、同 motrixsim、同 100 measured steps、同 `--numba-ab` 命令。

一、Task throughput

Task	Variant	ref env/s	本轮 env/s	Δ	ref env_step ms	本轮 env_step ms	Δ ms
sac / g1_motion_tracking	numpy_baseline	86,194	83,371	-3.3%	81.959	84.275	+2.316
sac / g1_motion_tracking	numba_accelerated	104,804	107,501	+2.6%	63.415	63.833	+0.418
flashsac / g1_walk_flat	numpy_baseline	116,215	112,475	-3.2%	50.221	52.220	+1.999
flashsac / g1_walk_flat	numba_accelerated	127,777	125,592	-1.7%	43.784	43.792	+0.008

`numpy_baseline` 两 task 都掉 3% 左右，说明机器基线有 ±3% 级别的环境噪声（Xeon 8568Y+ 上 physics 波动大，同 #663 comment 4904777445 已观察到）。numba_accelerated 上 motion_tracking +2.6% / walk_flat -1.7% 都在这条噪声带内 —— throughput 不是本 PR 的直接读数。

二、Reset done（本 PR 主目标）

Task	Variant	ref reset_done	本轮 reset_done	Δ ms	ref reset_call	本轮 reset_call	Δ ms
g1_motion_tracking	numba	12.349	13.146	+0.797	—	12.486	—
g1_walk_flat	numba	14.688	13.354	-1.334 (-9.1%)	13.943	12.576	-1.367 (-9.8%)
g1_walk_flat	numpy	13.907	13.077	-0.830 (-6.0%)	—	12.393	—

• `g1_walk_flat` reset_done 在 numpy 和 numba 两个 variant 上都稳定下降 ~1.0-1.3 ms（不含 physics，机器噪声小）。这是本 PR 的直接收益。
• `g1_motion_tracking` reset_done +0.797 ms：分解后主要在 DR plan（motion sampler，Python-heavy，非本 PR 范围）+0.638 ms 和 DR build_observation +0.156 ms 上；本 PR 触及的 `dr_reset_set_state` 本 task 上是 4.962 ms 且 ref 未细分无法直接对齐。属 motion_tracking task 特有的 Python 路径波动，不是本 PR 引入。

三、DR set_state 拆分（本 PR 新增，ref 没有此细分）

`g1_walk_flat numba_accelerated`：

sub-key	ms	占 set_state	归属
FK (`forward_kinematic`)	4.561	48.2%	Rust 边界（未优化）
Data reset	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	0.746	7.9%	Rust 边界
Set dof vel/pos	0.662	7.0%	Rust 边界
QPos convert	0.031	0.3%	纯 Python（本 PR 目标）
Mask	0.007	0.1%	纯 Python（本 PR 目标）
Reset rand / Clear forces / Geom overrides / Inval vel	≤ 0.004	≤ 0.1%	无 DR
Gap	0.013	0.1%	—
Total dr_reset_set_state_ms	9.461	100%

对比 ref：`g1_walk_flat numba_accelerated dr_reset_set_state_ms = 10.916 ms`（ref）→ 9.461 ms（本轮），-1.455 ms，1.15×。

`g1_motion_tracking numba_accelerated` set_state 分布更均衡：Data reset 27.7% + Data slice 16.6% + FK 15.7% + Actuator ctrl 13.1% + Refresh pose 11.0%，总 4.962 ms。

四、对齐检查

Update state（本 PR 不改）：g1_motion_tracking numba 21.031 → 19.062 ms；g1_walk_flat numba 4.183 → 4.536 ms（同量级，numba path 正常）。

Physics（本 PR 不改）：g1_motion_tracking numba 27.570 → 28.844 ms（+1.274）；g1_walk_flat numba 21.979 → 22.986 ms（+1.007），跟 CPU% 从 55/44 降到 52.5/46.2 一致 —— 机器整体负载比 ref 那天略高，与 numpy_baseline 的 -3% 观察一致。

五、对 #679 完成标准的验收

1. ✅ "reset_done/backend set_state 有可复现细分报告"：16 个 set_state sub-key 全部在远端 Xeon 上落地，gap ≤ 0.3%，计时完全闭合。JSON + CSV + log 存 `/home/unilab/unilabsim2/UniLab/benchmark/outputs/issue679/collector_remote_20260707_155738.*`。

2. ✅ "至少一个 G1 task 在 Motrix collector-only e2e 中能看到 reset_done 明确下降"：`g1_walk_flat` 上 reset_done `-1.334 ms (1.10×)`、reset_call `-1.367 ms (1.11×)`、dr_reset_set_state `-1.455 ms (1.15×)`；numpy 和 numba 两个 variant 同向验证。

3. ✅ "benchmark 结果能说明优化前后 reset_done 占 collector total/non-physics 的比例变化"：

   • `g1_walk_flat numba` reset_done 占 env_step：14.688/43.784 = 33.5% → 13.354/43.792 = 30.5%（-3.0 pp）
   • 占 non-physics env_overhead：14.688/21.805 = 67.4% → 13.354/20.806 = 64.2%（-3.2 pp）

六、下一阶段最大空间

set_state 里 FK 单项 4.561 ms 占 48.2%，Data reset + Data slice + Refresh pose 共 25.4%，合计 Rust 侧 82.5%。本 PR 攻击面（Python-side Mask + QPos convert + env_ids_intp hoist + actuator ctrl fast path）已达上限。剩下必须走 motrixsim 侧改动：

• fused set_state：把 `data_slice.reset` + `set_dof_vel` + `set_dof_pos` + `forward_kinematic` 合成一个 Rust-side batched 调用，减少多次 Python-Rust 边界穿越。
• index-based `forward_kinematic`：目前 FK 是对 `data_slice` 执行的，可以让 motrixsim 支持只对 done 的 env 子集做 FK，跳过其他 env。这条能直接砍掉 walk_flat 4.56 ms 的一大块。
• batched `get_link_poses`：Refresh pose 0.746 ms 目前也是每次 set_state 都 refresh，可以延迟到实际读之前 lazy 化。

这些都需要 motrixsim 侧改动 + 新 backend 接口，属于 #679 计划里明确的"后续工作"，建议在 #679 关闭前另开子 issue 跟进。

七、`g1_motion_tracking` reset_done +0.8 ms 的说明

看似回归，但分解到 sub-key 后落在：

• DR plan（motion sampler）：3.751 → 4.389 ms，+0.638 ms
• DR build_observation：2.950 → 3.106 ms，+0.156 ms

都在 task-specific 观测组装（motion sampler + relative transform）路径上，不在 backend set_state。本 PR 触及的 `dr_reset_set_state` 该 task 上是 4.962 ms（ref 未细分），Data reset 27.7% + Data slice 16.6% + FK 15.7% 分布均匀，没有单点异常。这条属于 motion_tracking task Python-heavy 路径在 Xeon 上的 run-to-run 波动，若需要该 task 也降 reset_done，得走"motion sampler numba 化 / row-wise body-pose getter"深耕，超出本 PR 边界。

八、原始文件（远端）

• `/home/unilab/unilabsim2/UniLab/benchmark/outputs/issue679/collector_remote_20260707_155738.json`
• `/home/unilab/unilabsim2/UniLab/benchmark/outputs/issue679/collector_remote_20260707_155738.csv`
• `/home/unilab/unilabsim2/UniLab/benchmark/outputs/issue679/remote_run_20260707_155738.log`

Refs: #665, #663 comment 4904656707, PR #681

[TATP-233]

收尾 & 关闭

三条完成标准全部达标：

1. ✅ 可复现细分报告：16 个 set_state sub-key 在 M5 Max + Xeon 8568Y+ 两台机器落地，gap ≤ 0.3% 计时闭合。
2. ✅ G1 task reset_done 明确下降：g1_walk_flat numba 上 reset_done -1.334 ms (1.10×)、dr_reset_set_state -1.455 ms (1.15×)；numpy variant 同向验证。
3. ✅ 占比变化：reset_done 占 env_step 33.5% → 30.5%，占 non-physics 67.4% → 64.2%。

Delivery：PR #681 已 merge 到 perf/g1-walk-numba-accel；本地微基准 / 本地 e2e / 远端 Xeon e2e 三条实测 comment 已归档。

后续攻击面 → #682：Motrix reset 热路径 Rust/motrixsim 边界优化（fused set_state / index-based FK / lazy get_link_poses / light data_slice.reset）。远端 Xeon 实测暴露出 FK 单项占 set_state 48.2%，是下一阶段最大空间，但必须走 motrixsim 侧改动。

未覆盖但已记录：

• g1_motion_tracking reset_done 上升 +0.797 ms，落在 motion sampler (DR plan) 而非本 PR 触及的 backend set_state 路径；若需要该 task 也降 reset_done，属于 motion sampler numba 化的独立工作项，超出本 issue 边界。

Closing。Refs → #682, PR #681, #665。