perf(reset): reset 路径 backend 边界 row-scoped 化与 set_state 计时补齐 (#1295) - #1300
Merged
TATP-233 merged 1 commit intoAug 25, 2026
Conversation
…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).
This was referenced Aug 25, 2026
This was referenced Aug 25, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #1295. Parent: #1292(杠杆 3)。Motrix Rust 边界仍属 #682,本 PR 不涉及。
改动
set_state 计时补齐
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 泄漏。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 占位保持列稳定;删除 legacyset_state_reset_ms/set_state_cache_refresh_ms(全仓库无引用)。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)。三后端同机对照(必测项,Ryzen 9 9950X3D + RTX 4090,8192 envs,sac/g1_motion_tracking)
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 子项(新增量化能力)