Skip to content

perf: scope MotionCommand partial-reset recompute to reset rows - #1265

Merged
TATP-233 merged 1 commit into
dev/issue-1042-manager-based-apifrom
perf/issue-1261-motion-command-partial-reset
Aug 23, 2026
Merged

perf: scope MotionCommand partial-reset recompute to reset rows#1265
TATP-233 merged 1 commit into
dev/issue-1042-manager-based-apifrom
perf/issue-1261-motion-command-partial-reset

Conversation

@TATP-233

Copy link
Copy Markdown
Collaborator

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)。
  • 跨 commit A/B(固定种子 30 步 + 多次 partial reset,flip/box × mujoco/motrix):obs / reward / command / time_steps / terminated / motion 输入全部 bitwise 一致,reset 行 metrics bitwise 一致。唯一差异:reset 时刻 untouched 行的 error metrics —— baseline 的全量重算会把 set_state FK 重算的浮点抖动写入这些行(baseline 自身此时值就与它上一步 per-step 值差 1e-4~0.25 量级),scoped 版保留 per-step 值;这些值下一步全量 compute 即被覆盖,不被日志/reward/obs 消费。
  • 基准 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%)
    • collector 吞吐:MuJoCo 64.5k → 68.9k env/s,Motrix 32.5k → 45.1k env/s
    • flashsac/g1_walk_flat 等其他 case 不经过 MotionCommand 路径,仅签名对齐无行为变化。

备注

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
@TATP-233
TATP-233 requested a review from caozx1110 as a code owner August 23, 2026 13:17
@TATP-233

Copy link
Copy Markdown
Collaborator Author

本地复测报告(第三方机器重复验证)

环境: Linux, Python 3.12.3, NVIDIA RTX 4090 (48GB), commit 6c8ab882(baseline: 2a3bcea7

1. 新增专项测试

uv run pytest tests/envs/test_motion_command_partial_reset.py -v

  • test_motion_command_partial_reset_row_parity[ppo-g1_flip_tracking-G1FlipTracking] PASSED
  • test_motion_command_partial_reset_row_parity[ppo-g1_box_tracking-G1BoxTracking] PASSED

2. make test-all

✅ 通过:2190 passed, 28 skipped, 275 deselected, 1 xfailed(与 PR 描述一致),benchmark smoke 32/33 + 33/34(仅 mlx 平台可选项跳过)。

3. 基准 A/B(同机同配置,sac/g1_motion_tracking,8192 envs,10 warmup + 100 measure steps)

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。

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