Skip to content

perf(reset): reset 路径 backend 边界 row-scoped 化与 set_state 计时补齐 (#1295) - #1300

Merged
TATP-233 merged 1 commit into
dev/issue-1292-motion-tracking-manager-perffrom
perf/issue-1295-reset-row-scoped
Aug 25, 2026
Merged

perf(reset): reset 路径 backend 边界 row-scoped 化与 set_state 计时补齐 (#1295)#1300
TATP-233 merged 1 commit into
dev/issue-1292-motion-tracking-manager-perffrom
perf/issue-1295-reset-row-scoped

Conversation

@TATP-233

Copy link
Copy Markdown
Collaborator

Fixes #1295. Parent: #1292(杠杆 3)。Motrix Rust 边界仍属 #682,本 PR 不涉及。

改动

set_state 计时补齐

  • 关键发现:manager-based 路径下 ResetStateTransaction.commit() 的 backend timing 返回值一直被丢弃,benchmark 的 set_state_* key 因此全为 0。现在在 commit 内捕获外层 wall-clock(dr_reset_set_state_ms)+ backend 子项,ManagerBasedRlEnv 经新 hook _collect_reset_backend_timing_ms 上报进 info["timing"];收集时按 RESET_DONE_DETAIL_TIMING_KEYS 过滤,顺带修掉旧 key 的 stale 泄漏。
  • mjwarp set_state 从 2 个 collapsed key 改为共享 schema:新增 set_state_reset_upload_ms / set_state_reset_forward_ms / set_state_host_cache_refresh_ms 实测值 + set_state_internal_gap_ms;mujoco/motrix 以 0.0 占位保持列稳定;删除 legacy set_state_reset_ms / set_state_cache_refresh_ms(全仓库无引用)。
  • benchmark keyset / CSV 列同步扩展,新增 mjwarp set_state 明细表。

reset 路径 row-scoped

  • SimBackend 新增 get_body_lin_vel_w_rows / get_body_ang_vel_w_rows(默认全量+切片),mujoco/motrix 覆盖为行级 gather;EntityData 增加 4 个 *_rows 薄封装。
  • MotionCommand._refresh_robot_state 行分支改用 row getter,消除每次 reset 4 次全 batch fancy-index 拷贝。
  • ManagerBasedRlEnv.reset() 的 command/obs compute 段包进 scene._scoped_state_reads,跨 term 重复 getter 去重(scope 在 set_state commit 之后进入,值为 post-reset 新鲜读取)。

明确不做:obs term 契约行级化(该 task 热 obs term 全部配置 UniformNoiseCfg,noise 需 full-batch draw 保持共享 RNG 流 bitwise parity,行级化只能省 term 内算术,性价比不成立);Motrix Rust 边界(#682)。

Parity

row getter 逐行 gather 与"全量取再切片"bitwise 一致;noise draw 路径未动。既有 parity 锚点全绿:test_motion_command_partial_reset.py(行刷新 vs 全量刷新)、test_observation_partial_reset.py(bitwise)、test_sim_backend_smoke.py(新增 vel rows parity)。另修复 HEAD 既有 bug:test_mjwarp_backend.py::test_real_cuda_init_reset_step 重复 monkeypatch 块导致 CUDA graph 下无限递归。

Validation

  • make test-all 本地已通过(含 docs 契约;support matrix 已 regenerate)。
  • slow lane 全量 273 passed(含 mjwarp CUDA 新增 set_state schema 测试)。
  • 按用户要求跳过远程 CI 等待。

三后端同机对照(必测项,Ryzen 9 9950X3D + RTX 4090,8192 envs,sac/g1_motion_tracking)

backend reset_done 基线 reset_done 本 PR Δ throughput 基线 → 本 PR
mujoco 9.38 ms (8.2% env) 6.92 ms (6.1%) −26% 70.8k → 71.5k env/s
motrix 20.63 ms (13.8%) 14.20 / 13.34 ms(两次复测) −31~35% 54.0k → 60.4k / 64.0k env/s
mjwarp 14.11 ms (21.9%) 14.79 / 14.63 ms(两次复测) ≈ 持平 123.5k → 129.2k / 127.2k env/s

mjwarp 持平原因已被新 keyset 量化:reset_done 14.6 ms 中 set_state 占 6.0 ms,其中 reset_forward 5.2 ms(86.8%)为 device 侧 forward,不在 host row-scoping 范围;host 侧收益被 device 边界主导掩盖。注意:基线为单次采样,未改动的 update_state 亦存在 run 间波动,但 reset_done 下降幅度大且 motrix 两次复测方向一致。

结果文件(未入库):scripts/benchmark/outputs/offpolicy_collector_active/results_issue1295_{baseline,after}.json/csv + 两个 rerun JSON。

mjwarp set_state 子项(新增量化能力)

sub-key ms share
reset_upload 0.114 1.9%
reset_forward 5.215 86.8%
host_cache_refresh 0.622 10.3%
internal_gap 0.010 0.2%

…t_state timing (#1295)

- Surface ResetStateTransaction commit timings (outer wall-clock + backend
  set_state sub-keys) into info["timing"] for manager-based envs; previously
  the backend timing dict was discarded and all set_state_* keys sampled as 0.
- mjwarp set_state now reports the shared keyset with granular
  set_state_reset_upload_ms / set_state_reset_forward_ms /
  set_state_host_cache_refresh_ms (mujoco/motrix emit 0.0 placeholders for
  column stability); legacy collapsed keys removed, stale-key leak fixed by
  filtering collected keys against RESET_DONE_DETAIL_TIMING_KEYS.
- Add SimBackend.get_body_lin_vel_w_rows / get_body_ang_vel_w_rows with
  row-gather overrides on mujoco/motrix; MotionCommand._refresh_robot_state
  uses them on the partial-reset path instead of full-batch reads + slice.
- Wrap reset command/obs compute in scene._scoped_state_reads so repeated
  getter calls across terms are deduplicated.
- Benchmark: extend NP_ENV_STEP_TIMING_KEYS / CSV fields and add the mjwarp
  set_state detail table.

Same-host sac/g1_motion_tracking (8192 envs) reset_done: mujoco 9.38 -> 6.92
ms (-26%), motrix 20.63 -> 13.3-14.2 ms (-31~35%), mjwarp ~flat (set_state
device forward 5.2 ms of 14.6 ms dominates; now quantified via new keyset).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant