perf(managers): continue reducing update_state cost (#1316) - #1323
Merged
TATP-233 merged 3 commits intoAug 26, 2026
Merged
Conversation
This was referenced Aug 26, 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
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
dev/issue-1304-motion-numba-body-state@9cca010d2c97910a6704a3da80b0d7f347103591dev/issue-1316-update-state-numba@d09d05aa39123ec1001210500e61b6539cb22cafmainordev/issue-1042-manager-based-api.Direct cumulative A/B
Fresh interleaved comparison of exact roadmap base
9cca010dagainst 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.command_manager.computecommand_manager.post_computeobservation_manager.computeupdate_statestep_coreenv_step_totalcommand_manager.computecommand_manager.post_computeobservation_manager.computeupdate_statestep_coreenv_step_totalThe isolated #1319 interleaved A/B, using its immediate base
1d18060e, measuredobservation_manager.computeat -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
backend_physics_msbackend_host_cache_refresh_msbackend_control_upload_msThe approximately 0.66 ms GPU-to-host copy and synchronization happen inside
backend.step(). They are therefore already included instep_core_msandenv_step_total_ms, while remaining outsideupdate_state; the table reports them separately and does not double-count them.Remaining bottlenecks
At the final head, the largest measured
update_statepartition is stillobservation_manager.compute(MuJoCo/MJWarp 4.20/3.58 ms), followed by complete reward+termination term calls (2.94/2.62 ms) andcommand_manager.compute(2.11/1.75 ms).command_manager.post_computeis 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
make test-allbefore its PR was createdd09d05aapassedmake test-all: 2309 passed, 28 skipped, 281 deselected, 1 xfailed; mypy, pyright, Ruff, and benchmark import smoke passedPer #1313 governance this is a non-
mainintegration PR into the declared non-protected base. Remote CI is intentionally not scheduled or awaited; the complete gate is the recorded localmake test-all.