perf: scope MotionCommand partial-reset recompute to reset rows - #1265
Merged
TATP-233 merged 1 commit intoAug 23, 2026
Merged
Conversation
On the MBA partial-reset path the command manager recomputed motion reference, robot state, relative transforms and metrics over the full batch even though only the reset rows changed. Restrict the reset-path work to the reset rows: - CommandTerm.compute forwards the existing env_ids to the private _update_metrics hook (public compute/post_compute signatures and the manager lifecycle are unchanged). - MotionCommand refreshes motion data, robot state, relative transforms and error metrics row-scoped on the reset path; _update_command hands the reset ids to post_compute through a private field. Sampler-stat metrics stay full-batch since they track global sampler scalars. - BoxMotionCommand scopes its object-state refresh the same way. mba_reset_command_compute_ms (sac/g1_motion_tracking, 8192 envs): mujoco 21.57 -> 5.77 ms (-73%), motrix 53.12 -> 10.85 ms (-80%). Seeded flip/box rollouts on both backends are bitwise identical to the baseline in obs/reward/command/time_steps/terminated and reset-row metrics; reset-time untouched-row error metrics keep the per-step values instead of the baseline's post-set_state FK-recompute jitter. Fixes #1261
Collaborator
Author
本地复测报告(第三方机器重复验证)环境: Linux, Python 3.12.3, NVIDIA RTX 4090 (48GB), commit 1. 新增专项测试
2.
|
| Backend | mba_reset_command_compute_ms baseline → PR |
变化 | Collector 吞吐 baseline → PR |
|---|---|---|---|
| MuJoCo | 14.68 → 3.39 ms(median 14.16 → 3.32) | -77% | 53.5k → 55.9k env/s(+4.3%) |
| Motrix | 24.66 → 6.32 ms(median 23.11 → 6.14) | -74% | 46.6k → 44.6k env/s(-4.4%,噪声范围内) |
绝对值与 PR 描述不同(机器不同),但相对降幅与 PR 结论一致(-73% / -80%),均远超 ≥50% 目标。Motrix 侧吞吐未复现 PR 中的提升(本机 baseline 吞吐已较高),但目标指标 mba_reset_command_compute_ms 的优化效果确认有效。
结论
PR 的测试与性能结论在本机可复现,LGTM。
7 tasks
This was referenced Aug 23, 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 #1261. Parent: #1259.
MBA G1 motion tracking 的 partial reset 只重置约 251/8192 个 env,但 reset 路径上 MotionCommand 对全 batch 重做 command / relative-state / metrics 计算。本 PR 把 reset 路径的工作限定到 reset 行:
CommandTerm.compute把已有的env_ids传入私有钩子_update_metrics(env_ids);公共compute/post_compute签名与 Manager-Based lifecycle 不变。MotionCommand:_refresh_motion/_refresh_robot_state/_refresh_relative_state/_update_metrics支持行级 scoped;_update_command通过私有字段把 reset ids 传给post_compute。sampler 统计是全局标量,保持全量 fill。motion sampling、RNG 消耗、termination、observation 语义不变。BoxMotionCommand的 object 状态刷新同样 scoped。Validation
make test-all本地通过(2190 passed;附带清理了工作区中 refactor(tasks): remove legacy env task packages #1190 遗留的纯.pyc空包目录导致的 1 个与本改动无关的预存失败)。tests/envs/test_motion_command_partial_reset.py:untouched 行逐位不变、reset 行与全量重算参考一致、reset obs 校验(plain MotionCommand + BoxMotionCommand)。sac/g1_motion_tracking(8192 envs,同机同 commit 前后):mba_reset_command_compute_ms:MuJoCo 21.57 → 5.77 ms(-73%),Motrix 53.12 → 10.85 ms(-80%)(目标 ≥50%)flashsac/g1_walk_flat等其他 case 不经过 MotionCommand 路径,仅签名对齐无行为变化。备注
mba_reset_*_ms持久计时键属于优化期的临时脚手架,正式训练不需要;待 Roadmap: 优化必选的 Manager-Based runtime collector 吞吐并完成 #1042 性能验收 #1259 下所有优化完成后统一删除(见 issue 评论)。make test-all为准。