perf: scope partial-reset observation rebuild to reset rows - #1270
Merged
TATP-233 merged 1 commit intoAug 23, 2026
Merged
Conversation
Issue #1269 (roadmap #1259 R2). On the partial-reset path the ObservationManager computed the full batch through the whole noise/clip/scale/NaN/delay/history/concat pipeline and the env then sliced out the reset rows, discarding the rest. compute(env_ids=...) now returns only the reset rows: groups without delay/history terms run the row-independent pipeline stages on the reset rows only, while term calls and noise draws stay full-batch so the (num_envs, ...) term contract and the shared RNG stream are bit-identical to the full-batch path. Groups with temporal buffers keep the full pipeline and only slice the final output.
7 tasks
This was referenced Aug 24, 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.
Summary
Fixes #1269. Parent: #1259(Wave 2 R2).
MBA G1 motion tracking 的 partial reset 只重置约 251/8192 个 env,但 reset 路径上 ObservationManager 的 noise→clip→scale→NaN→delay/history→concat 管线全程全 batch 计算,env 再切出 reset 行,其余行的结果被丢弃。#1261 合并后
mba_reset_obs_build_ms成为 reset 路径最大剩余热点(MuJoCo 36%、Motrix 58%)。本 PR 把 reset 路径的行独立管线限定到 reset 行:ObservationManager.compute(update_history=True, env_ids=ids)现在只返回 reset 行(按 ids 顺序),不再由 env 侧全量切片;正常 step 管线(env_ids=None)逐行代码不变。(num_envs, ...)公共契约不变;full-shape noise 抽取保证共享 RNG 流与逐行 noise 值和全量实现逐位一致),noise 之后切片,clip/scale/NaN/concat 只处理 reset 行。_obs_buffer缓存(ObservationManager.reset本就将其失效),下一个 per-step compute 正常重建。_map_observations增加num_rows校验参数;NaN 诊断在行级路径报告真实 env id。Validation
11bb0ac3,git status --short --branch工作树干净(仅有本仓库既有的未跟踪benchmark/目录,与本 PR 无关)。make test-all本地通过(最终提交上运行;含 mypy/pyright/ruff/format 与全部 pytest)。tests/managers/test_observation_partial_reset.py(6 个 manager 级测试:reset 行与全量 compute 切片逐位一致、RNG 流一致、时序组 fallback、缓存语义、NaN 诊断 env id)与tests/envs/test_observation_partial_reset.py(env 级:reset_obs 行形状/写回/untouched 行逐位不变/下一步恢复全量)。sac/g1_motion_tracking× mujoco/motrix、sac/g1_wbt_obs(history fallback)、flashsac/g1_walk_flat):obs / reward / terminated / truncated / reset_obs 全部 bitwise 一致(4 case × 156 arrays)。sac/g1_motion_tracking(8192 envs,同机同 commit 前后):mba_reset_obs_build_ms:MuJoCo 5.67 → 2.94 ms(-48%),Motrix 21.54 → 14.41 ms(-33%)mba_reset_total_ms:MuJoCo 15.67 → 13.01 ms(-17%),Motrix 37.33 → 30.04 ms(-20%)备注
mba_reset_*_ms/mba_obs_*_ms持久计时键属于优化期的临时脚手架,正式训练不需要;待 Roadmap: 优化必选的 Manager-Based runtime collector 吞吐并完成 #1042 性能验收 #1259 下所有优化完成后统一删除(见 perf: 限定 MotionCommand partial-reset 的全量重计算范围 #1261 评论)。make test-all为准。observation_manager.pyowner/file)。