Skip to content

perf: scope partial-reset observation rebuild to reset rows - #1270

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

perf: scope partial-reset observation rebuild to reset rows#1270
TATP-233 merged 1 commit into
dev/issue-1042-manager-based-apifrom
perf/issue-1269-obs-partial-reset

Conversation

@TATP-233

Copy link
Copy Markdown
Collaborator

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)逐行代码不变。
  • 无 delay/history 的组:term func 与 noise 保持全 batch(term 的 (num_envs, ...) 公共契约不变;full-shape noise 抽取保证共享 RNG 流与逐行 noise 值和全量实现逐位一致),noise 之后切片,clip/scale/NaN/concat 只处理 reset 行。
  • 有 delay/history 的组:保持全量管线(buffer backfill 本来就是 per-row 写入),仅切片最终输出;时序语义不变。
  • reset 路径不再填充 _obs_buffer 缓存(ObservationManager.reset 本就将其失效),下一个 per-step compute 正常重建。
  • _map_observations 增加 num_rows 校验参数;NaN 诊断在行级路径报告真实 env id。

Validation

  • 最终提交 11bb0ac3git 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 行逐位不变/下一步恢复全量)。
  • 跨 commit A/B(固定种子 30 步 + 3 次 partial reset,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%)
    • collector 吞吐:MuJoCo 67.2k → 72.8k env/s(+8%),Motrix 46.5k → 54.0k env/s(+16%)
    • 剩余 obs_build 主要是 term func 内部的全量 backend getter(U1 边界,不在本 PR)。

备注

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.
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