Skip to content

feat: assemble motion tracking Numba plan - #928

Merged
TATP-233 merged 1 commit into
refactor/issue-902-mjwarp-squashed-deliveryfrom
feat/issue-923-motion-auto-numba
Aug 5, 2026
Merged

feat: assemble motion tracking Numba plan#928
TATP-233 merged 1 commit into
refactor/issue-902-mjwarp-squashed-deliveryfrom
feat/issue-923-motion-auto-numba

Conversation

@TATP-233

@TATP-233 TATP-233 commented Aug 5, 2026

Copy link
Copy Markdown
Collaborator

Fixes #923

Summary

  • register motion tracking preamble, reward, observation, and termination Numba item functions on the existing resolved term plan
  • extend the shared fused assembler narrowly for bool/int scalar parameters, typed tuple parameters, parameter-schema cache keys, and cold-path workspace view binding
  • replace the motion tracking hand-written full aggregator and fixed term order with the automatic fused plan while preserving task-owned noise and lifecycle orchestration
  • reuse preallocated actor/critic, reward, termination, and shared workspace buffers; remove the legacy accelerator path
  • update the focused benchmark from reward/termination wording to full fused plan-state measurements

Owner boundaries

  • no backend, runner, NpEnv lifecycle, training ABI, checkpoint, sim2sim, asset, or regular CI changes
  • motion reference sampling, adaptive failure accounting, frame advance, clip-end behavior, teleport, and final observation remain task-owned
  • this PR completes only the G1 motion-tracking pilot; it does not promote other task families

Validation

  • make test-all
    • 1667 passed, 37 skipped, 280 deselected, 1 xfailed
    • Ruff, mypy, Pyright, and benchmark import smoke passed
  • uv run pytest -q tests/envs/test_g1_motion_tracking_numba.py -m slow
    • 2 passed
  • uv run pytest -q tests/envs/test_env_configs.py -k 'g1_motion_tracking_clip_end'
    • 5 passed
  • shared assembler execution now explicitly covers bool/int scalar and typed tuple parameters
  • actor owner-noise / clean critic boundary, output buffer reuse, and first-JIT reporting are covered

Performance

Command:

uv run python -m benchmark.env.benchmark_g1_motion_tracking_numba --profiles sac_default --num-envs 2048 8192 --threads 8 --iters 80 --warmup 8 --no-e2e

8-thread full fused plan-state results versus the pre-migration hand-written _compute_update_state_kernel baseline:

envs current old baseline retained
2048 10.43M env/s 10.72M env/s 97.3%
8192 12.24M env/s 11.60M env/s 105.5%

Both exceed the 90% acceptance gate. Parity maxima were about 1.0e-7 for reward and 7.2e-7 for critic observation, with zero termination mismatches.

Cold/cache measurements at N=2048, 8 threads:

  • cache miss assembly: 0.813 ms; first JIT execute: 11876.861 ms
  • cache-hit assembly: 0.001 ms; first warm execute: 0.198 ms

Scope size

  • 9 files
  • +640 / -1045 lines

@TATP-233
TATP-233 requested a review from caozx1110 as a code owner August 5, 2026 19:21
@TATP-233
TATP-233 merged commit 71108ca into refactor/issue-902-mjwarp-squashed-delivery Aug 5, 2026
@TATP-233
TATP-233 deleted the feat/issue-923-motion-auto-numba branch August 6, 2026 11:46
TATP-233 added a commit that referenced this pull request Aug 6, 2026
…auto-numba"

This reverts commit 71108ca, reversing
changes made to b03da66.
lemon-star608 pushed a commit that referenced this pull request Aug 21, 2026
lemon-star608 pushed a commit that referenced this pull request Aug 21, 2026
…auto-numba"

This reverts commit 2ccec24, reversing
changes made to b70440a.
TATP-233 added a commit that referenced this pull request Sep 4, 2026
TATP-233 added a commit that referenced this pull request Sep 4, 2026
…auto-numba"

This reverts commit 83edbf2, reversing
changes made to 90daa2f.
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