Skip to content

perf(managers): Numba 热点、body-state 与 update_state 加速 (#1304, #1316) - #1314

Merged
TATP-233 merged 21 commits into
dev/issue-1042-manager-based-apifrom
dev/issue-1304-motion-numba-body-state
Aug 27, 2026
Merged

perf(managers): Numba 热点、body-state 与 update_state 加速 (#1304, #1316)#1314
TATP-233 merged 21 commits into
dev/issue-1042-manager-based-apifrom
dev/issue-1304-motion-numba-body-state

Conversation

@TATP-233

@TATP-233 TATP-233 commented Aug 26, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Merge policy: do not auto-merge. This integration PR intentionally waits for maintainer approval and manual merge into dev/issue-1042-manager-based-api.

Summary

本 PR 现在集成两个已完成 roadmap 的有界工作：

  • roadmap: motion tracking Numba 热点与统一 body-state copy 收敛 #1304：将 termination/anchor_pos 与四个固定 reward/motion_body_* 热点改为预热、并行 Numba kernels；通过正式 SimBackend.copy_body_state_w() contract 消除四组 full-batch getter/gather；MuJoCo 与 MJWarp 共用 fused host-copy kernel。
  • roadmap: 扩大 Numba 优化并继续降低 Manager-Based update_state 开销 #1316：继续优化 command_manager.compute、command_manager.post_compute、observation_manager.compute 三个最大 manager 分区。
  • command metrics 和 relative/body-frame transforms 的 production call graph 只保留 Numba 数学；对应旧 runtime NumPy modules/helpers 已移除，NumPy 只作为独立测试 oracle。
  • observation pipeline 的大头是 RNG、ownership copy、finite scan 和 concat，没有值得单独 Numba 调度的数学循环；因此直接消除临时数组与重复扫描，没有维护第二条 nominal Numba 路径。
  • Motrix 按已确认范围不参与加速或 benchmark：没有 Motrix 专用 kernel，依赖仍为 0.8.2；文档定义的 Manager-facing backend/env/config/runner 公共边界保持兼容。内部 term 实现、buffer reuse、Numba runtime 和下文列出的首次 metrics 初始化行为属于明确适配，不再笼统声明所有 lifecycle/bitwise 行为无变化。

Linked work

Manager-Based API contract boundary

本 PR 保留文档定义的 Manager-facing 公共边界：

  • Hydra owner YAML 仍是 reward/termination 的配置 source of truth；dotted callable、params、weight、body subset 与 timeout 配置继续由 manager 冷路径物化。
  • 四个 motion_body_* reward 仍是四个独立 MBA term 和四个独立 Numba kernel，并未融合成一个不可配置的 reward；RewardManager 仍逐项应用 weight 与 dt。
  • termination/anchor_pos 仍由 TerminationManager 聚合，并保留 bool shape/dtype 校验和 termination/timeout split。
  • body-state 加速通过正式 Entity/SimBackend.copy_body_state_w() contract 接入；task/env 不访问 MuJoCo 或 MJWarp 私有对象。
  • ManagerBasedRlEnv、NpEnvState、reset/step 核心顺序和 runner/IPC 路径没有修改。

以下是在本 PR 中明确披露的内部或 source-level 适配，不属于“所有实现细节完全不变”的声明：

  • bad_anchor_pos_z_only 从 module function 改为 ManagerTermBase class。Hydra/Manager 的 dotted callable 与 (cfg, env) 冷路径实例化兼容；旧式直接 Python 函数调用不再 source-compatible。
  • 四个 body reward 和 anchor termination 返回各自 term-owned reusable buffer。Manager 在应用 weight/dt 或聚合 done 时安全消费；外部调用者不能把返回引用当成跨调用不变的快照。
  • Numba 标量累加和 quaternion/math 运算按声明容差与 NumPy oracle 对齐，不声明 bitwise identity。
  • MotionCommand 构造冷路径会执行一次 _update_metrics() 以完成 kernel 预热，因此首次 reset 前的 metrics buffer 可以是已计算值，而不是历史实现的全零值。reset/step 核心调用顺序未改，但本 PR 不再把这一初始化细节描述为 lifecycle 完全无变化。
  • Numba 是直接 runtime dependency；未显式设置环境变量时，motion kernels 默认选择 workqueue threading layer，并把 Numba worker mask 限制到最多 8 线程。用户设置的 NUMBA_THREADING_LAYER / NUMBA_NUM_THREADS 优先。

因此，本 PR 的 compatibility claim 限定为 documented Manager-facing public boundary；内部 term 结构、buffer ownership、JIT/thread runtime 和数值执行属于明确 adaptation，数值、shape/dtype、body subset/order 与 output reuse 等近风险边界已有专项测试。

Final head and validation

  • Final PR head: 4fa6631a3c9c465db86cd671275319168ce36fe0
  • Working tree: clean
  • Exact-head make test-all: passed
    • Ruff format/check: passed
    • mypy: passed
    • pyright: passed
    • pytest: 2309 passed, 28 skipped, 281 deselected, 1 xfailed
    • benchmark module/script import smoke: passed; only platform-optional MLX skipped
  • Near-risk parity covers full batch, partial-reset rows, untouched rows, quaternion semantics, body subset/order, dtype/shape, input immutability, output ownership/reuse, fixed-seed noise, delay/history, and NaN policies.

Remote CI override: per the maintainer/user instruction for these non-main roadmap PRs, validation is local-only; remote CI is intentionally skipped and not awaited.

Performance: #1304 phase

Same host: AMD Ryzen 9 9950X3D2, NVIDIA RTX 4090 48 GiB, 8192 envs, warmup 10 / measure 100. Values below preserve the original #1304 collector benchmark convention (median of three independent run means).

Backend Initial update_state #1304 final Delta Initial env/s #1304 final Delta
MuJoCo 29.766 ms 19.986 ms -32.9% (1.49×) 70,297 76,895 +9.4%
MJWarp 28.704 ms 18.894 ms -34.2% (1.52×) 122,524 149,283 +21.8%

Motrix intentionally has no performance result or accelerated support claim.

Performance: #1316 direct cumulative A/B

Fresh interleaved comparison of exact #1316 base 9cca010d against the final tree: g1_motion_tracking, 8192 envs, warmup 10, measure 100, fixed action/env seeds, three repetitions. Phase values are the median of three per-run medians; instrumented manager values are the median of three per-run means.

Backend Metric Before (ms) Final (ms) Delta
MuJoCo command_manager.compute 7.364839 2.113478 -71.30%
MuJoCo command_manager.post_compute 4.332488 0.710599 -83.60%
MuJoCo observation_manager.compute 4.431236 4.199805 -5.22%
MuJoCo update_state 19.417052 10.322279 -46.84%
MuJoCo step_core 52.925145 52.876745 -0.09%
MuJoCo env_step_total 84.086522 76.055615 -9.55%
MJWarp command_manager.compute 6.747341 1.753317 -74.01%
MJWarp command_manager.post_compute 4.262291 0.586890 -86.23%
MJWarp observation_manager.compute 3.811295 3.584991 -5.94%
MJWarp update_state 17.762506 8.971272 -49.49%
MJWarp step_core 14.396119 14.419870 +0.16%
MJWarp env_step_total 49.853370 40.943064 -17.87%

The isolated #1319 immediate-base A/B measured observation at -7.24% MuJoCo / -7.93% MJWarp. The direct roadmap-wide values are deliberately conservative; both methods agree on direction.

MJWarp PCIe accounting

Metric Before (ms) Final (ms) Delta
backend_physics_ms 13.604566 13.669142 +0.47%
backend_host_cache_refresh_ms 0.658934 0.659782 +0.13%
backend_control_upload_ms 0.103609 0.097939 -5.47%

The approximately 0.66 ms GPU-to-host copy and synchronization occur inside backend.step(). They are already included in step_core_ms and env_step_total_ms, remain outside update_state, and are not double-counted.

Remaining bottlenecks

At the final head, the largest measured update_state partition is observation_manager.compute (MuJoCo/MJWarp 4.20/3.58 ms), followed by complete reward+termination term calls (2.94/2.62 ms) and command_manager.compute (2.11/1.75 ms). command_manager.post_compute is now 0.71/0.59 ms. At whole-step scope, backend physics dominates; MJWarp's stable 0.66 ms D2H barrier is part of real collector cost.

Impact

  • Backend: MuJoCo and MJWarp accelerated; Motrix remains correct but intentionally non-accelerated
  • Runtime math: one production expression for every Numba-replaced formula; no NumPy fallback path
  • Dependency: Numba remains a direct runtime dependency; Motrix pin/lock unchanged
  • Manager-Based boundary: documented manager/config/env/backend/runner public contract remains compatible; internal adaptations are documented above
  • Expected training semantics: numerically equivalent within the tested tolerances, not claimed as bitwise identity; collector update_state cost is lower

Checklist

  • all child issues stayed within their declared owner/scope
  • tests added at reward, command, observation, Entity, and backend ownership boundaries
  • no second reward/transform runtime expression
  • no Motrix-specific Numba kernel
  • exact final head passed local make test-all
  • final PR left open for maintainer approval and manual merge

@TATP-233
TATP-233 requested a review from caozx1110 as a code owner August 26, 2026 09:23
This was referenced Aug 26, 2026
Closed
Closed
Closed
Closed
Closed
Merged
@TATP-233 TATP-233 changed the title perf(motion): Numba 热点与 MuJoCo/MJWarp body-state 加速 (#1304) perf(managers): Numba 热点、body-state 与 update_state 加速 (#1304, #1316) Aug 26, 2026
TATP-233 and others added 7 commits August 27, 2026 01:03
@TATP-233
perf(env): confine DP collector host compute to the per-rank CPU block
512fc82 
Multi-rank off-policy collectors pin their MuJoCo BatchEnvPool workers to
a per-rank CPU block via EnvCfg.cpu_ids, but the collector's host-side
compute did not follow: Numba's parallel kernels sized their pool from the
host CPU count and drifted across rank boundaries, and the OpenBLAS pool
spawned at import kept the host-wide mask.

NpEnv.__init__ now applies apply_env_cpu_runtime(cfg.cpu_ids) on the cold
path: the process is confined to the block (existing threads pinned
individually via /proc/self/task, later threads — including Numba's
lazily-launched pool — inherit the mask) and Numba's pool is sized to
len(cpu_ids) unless NUMBA_NUM_THREADS is set explicitly. cpu_ids=None keeps
the single-rank path bit-identical. Backend-agnostic: any env declaring
cpu_ids (e.g. motrix once it grows affinity support) gets the same
confinement.
@TATP-233
perf(mjwarp): reduce motion tracking reset latency
a482a1f
@TATP-233
Merge pull request #1324 from unilabsim/perf/dp-collector-numba-cpu-b…
c2b110d 
…lock

perf(env): confine DP collector host compute to the per-rank CPU block
@TATP-233
Merge pull request #1327 from unilabsim/perf/issue-1325-mjwarp-reset
f734cab 
perf(mjwarp): reduce g1 motion tracking reset latency
@TATP-233
benchmark(env): add MuJoCo pool thread-scaling and env-step phase-CPU…
bc586c4 
… probes (#1328)

Diagnostic probes for the SAC/MuJoCo single-GPU collector CPU
under-utilization report: pool thread-count scaling on the G1 scene, and
per-phase wall/CPU attribution of a full task env step. New files only;
no behavior change.
@TATP-233
Merge pull request #1329 from unilabsim/perf/issue-1328-mujoco-pool-t…
b23e632 
…hread-sizing

benchmark(env): add MuJoCo pool thread-scaling and env-step phase-CPU probes (#1328)
@TATP-233
fix(base): type-check and test cpu_runtime on non-Linux hosts
56032a6 
os.sched_setaffinity/sched_getaffinity are Linux-only, so mypy on darwin
rejected the direct attribute access (attr-defined) and the unit tests'
monkeypatch.setattr/delattr failed because the attributes do not exist.

Resolve the affinity symbols via getattr at call time (identical runtime
semantics, still monkeypatchable) and pass raising=False to the test
monkeypatch seams so they work whether or not the host exposes them.
@TATP-233 TATP-233 mentioned this pull request Aug 27, 2026
Merged
TATP-233 and others added 4 commits August 27, 2026 14:29
@TATP-233
Merge pull request #1330 from unilabsim/fix/cpu-runtime-darwin-mypy
07e6fe6 
fix(base): cpu_runtime 在非 Linux 主机上的类型检查与测试修复
@TATP-233 @claude
fix(logging): make collector reward reporting timely
411a6ee 
Reward displays (tensorboard reward/mean and the terminal logger) lagged
badly on off-policy and APPO runs:

- collectors sent metrics only every num_envs * 10 env steps, so the
  reported reward changed just once per ~10 learner iterations;
- runners then averaged the last 100 (off-policy) or 50 (APPO) reports,
  each already a rolling 100-episode mean, delaying the visible curve by
  ~1000 iterations.

Report metrics every collector cycle, keep the runner-side window at the
last 10 reports, and bound the per-worker episode reward/length buffers
with deque(maxlen=100) instead of lists that grew for the whole run.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@TATP-233 @claude
fix(env): keep per-step reward log entries through autoreset
22fc439 
ManagerBasedRlEnv.reset() replaced state.info["log"] with the reset-only
extras (Episode_Reward/*), wiping the fresh per-step reward/* entries that
_update_state_in_read_phase() had just computed for the current transition.
On any step where at least one env resets — with thousands of envs, nearly
every step — collectors therefore saw no reward/* keys at all, so the
per-term reward components in tensorboard and the terminal logger stayed
frozen at one stale value for thousands of iterations (observed as long
flat staircases on reward/motion_* etc.).

Merge instead of replace on the autoreset path: the pre-reset per-step
entries stay, reset extras layer on top. Standalone (non-autoreset) resets
are unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@TATP-233
Merge pull request #1331 from unilabsim/fix/offpolicy-reward-log-cadence
457b1a9 
fix(logging): make collector reward reporting timely
@TATP-233
TATP-233 merged commit c07cfa7 into dev/issue-1042-manager-based-api Aug 27, 2026
@TATP-233
TATP-233 deleted the dev/issue-1304-motion-numba-body-state branch August 27, 2026 08:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@caozx1110 caozx1110 Awaiting requested review from caozx1110

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@TATP-233