Skip to content

perf(managers): continue reducing update_state cost (#1316) - #1323

Merged
TATP-233 merged 3 commits into
dev/issue-1304-motion-numba-body-statefrom
dev/issue-1316-update-state-numba
Aug 26, 2026
Merged

perf(managers): continue reducing update_state cost (#1316)#1323
TATP-233 merged 3 commits into
dev/issue-1304-motion-numba-body-statefrom
dev/issue-1316-update-state-numba

Conversation

@TATP-233

Copy link
Copy Markdown
Collaborator

Summary

This integration PR completes roadmap #1316 through three bounded child PRs:

The two arithmetic replacements are Numba-only in production; NumPy remains only in independent test oracles. The observation stage contains RNG, ownership copies, finite scanning, and concatenation rather than a standalone arithmetic loop, so it was optimized directly instead of maintaining a nominal Numba path. Motrix, public contracts, runner/lifecycle, and config ownership are unchanged.

Closes #1316
Milestone: M2 - Manager update_state throughput

Branch contract

  • Declared base: dev/issue-1304-motion-numba-body-state@9cca010d2c97910a6704a3da80b0d7f347103591
  • Integration head: dev/issue-1316-update-state-numba@d09d05aa39123ec1001210500e61b6539cb22caf
  • Child heads were created from the latest integration branch and merged in order.
  • This PR does not target, checkout, push, merge, or rewrite main or dev/issue-1042-manager-based-api.

Direct cumulative A/B

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

Backend Metric Roadmap base (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 interleaved A/B, using its immediate base 1d18060e, measured observation_manager.compute at -7.24% MuJoCo and -7.93% MJWarp. The direct roadmap-wide observation deltas above are smaller because all three stages are initialized and sampled together; both measurements agree on direction without claiming sub-noise end-to-end effects for #1319 alone.

MJWarp PCIe accounting

Metric Roadmap base (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 happen inside backend.step(). They are therefore already included in step_core_ms and env_step_total_ms, while remaining outside update_state; the table reports them separately and does not double-count them.

Remaining bottlenecks

At the final head, the largest measured update_state partition is still 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. For the full environment step, backend physics remains dominant; on MJWarp, the stable 0.66 ms D2H barrier is part of the real collector cost.

Hardware: AMD Ryzen 9 9950X3D2 (16 cores / 32 threads), NVIDIA GeForce RTX 4090 48 GiB, Linux 7.0.0-30-generic, NumPy 2.4.4, Numba 0.67.0, 32 Numba threads.

Validation

  • every child final head passed local make test-all before its PR was created
  • exact integration head d09d05aa passed make test-all: 2309 passed, 28 skipped, 281 deselected, 1 xfailed; mypy, pyright, Ruff, and benchmark import smoke passed
  • parity covers full batch, partial-reset rows, untouched rows, dtype/shape, quaternion semantics, output ownership, fixed-seed noise, delay/history, and NaN policies
  • working tree is clean

Per #1313 governance this is a non-main integration PR into the declared non-protected base. Remote CI is intentionally not scheduled or awaited; the complete gate is the recorded local make test-all.

@TATP-233
TATP-233 requested a review from caozx1110 as a code owner August 26, 2026 13:15
@TATP-233
TATP-233 merged commit 4fa6631 into dev/issue-1304-motion-numba-body-state Aug 26, 2026
@TATP-233
TATP-233 deleted the dev/issue-1316-update-state-numba branch August 26, 2026 15:11
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