Skip to content

roadmap: macOS tensor runtime for MuJoCo/mjbatch off-policy training #2055

Closed
Closed
roadmap: macOS tensor runtime for MuJoCo/mjbatch off-policy training#2055

Description

@TATP-233
TATP-233
opened on Oct 4, 2026

Summary

  • Goal: run the completed npenv→torchenv runtime switch on macOS (Apple Silicon) for the MuJoCo/mjbatch backend, and deliver working off-policy training on macOS: SAC + g1_walk_flat and FlashSAC + g1_motion_tracking. Final acceptance is one full SAC g1_walk_flat training run on macOS.
  • Recommended path: one integration branch per repo (dev/issue-<N>-macos-tensor-runtime) based on develop/tensor-runtime (siblings based on the tensor_runtime_workspace.json pins); materialize the workspace; fix macOS breakages in the owner layers behind device/platform gates; iterate with algo.max_iterations=100 smoke trainings; finish with the repo gate and the full acceptance run.
  • Delivery boundary: no npenv-runtime revert, no np-rng revert, no contract crossings, no main or develop/tensor-runtime mutation in any repo, no PyPI publication, no CUDA performance/runnability regression, no training-semantics change. Multi-GPU DP stays CUDA-only. No macOS CI job is added (follow-up candidate).
  • Estimated review scale: small-to-medium; expected mostly config/platform-gate fixes across UniLab, unisim, unilab_rl; mjbatch_uni expected zero or minimal changes.
  • Durable responsibility: UniLab owns training/runtime glue and owner YAML; uni_rl owns learners/collectors/replay/IPC; unisim owns the SimBackend tensor contract and the mujoco adapter; mjbatch_uni owns the CPU batch executor.
  • Decisions required: (1) accept MPS as the macOS off-policy learner/replay device (uni_rl rejects a pure-CPU learner by design: require_offpolicy_replay_device accepts cuda|mps only); a pure-CPU off-policy learner is out of scope; (2) sibling bases are the workspace manifest pins, not branch tips; (3) macOS CI coverage deferred.

Repository evidence

  • develop/tensor-runtime (90068b8) completed the npenv→torchenv switch; the old NpEnv runtime and the tensor-runtime switch are removed (ae851eb4, 759153cd).
  • MuJoCo/mjbatch is a HOST_BRIDGE backend: the Manager env is CPU-tensor authoritative by design on every platform (src/unilab/envs/manager_based_rl_env.py:372-377); unisim declares macos_tensor_profile = "CPU-authoritative host bridge only" for mujoco (unisim/support.py, tensor-runtime branch).
  • macOS learner path exists: get_default_device() resolves cuda → xpu → mps → cpu (src/unilab/utils/device.py:17-25); uni_rl implements an MPS replay pipeline and MPS fused optimizers, tested on CPU-only torch builds.
  • mjbatch-uni 0.2.4 already builds and steps on macOS arm64 (verified: editable arm64 .so, live 4-sim CPU Batch.step()); cpu_ids affinity is Linux-only there.
  • Current blocker on a macOS workstation: sibling checkouts on main lack the tensor-runtime API (HostBridgeTransferPlan et al.), so unilab.base.torch_env fails to import; the workspace manifest pins develop/tensor-runtime commits (unisim 09d5caed, unilab-rl 75680c23, mjbatch-uni b3a3b480).
  • No macOS job exists in CI; the heaviest integration coverage for the two tasks lives in the slow lane.

Scope

Included:

  1. Workspace materialization on macOS: integration branches from the declared bases, uv sync, import gates green.
  2. macOS adaptations in owner layers only, behind device/platform gates, with CUDA behavior preserved: env/learner device resolution, replay device gates, CPU-affinity degradation, viewer/playback guards as needed.
  3. Training enablement and iteration-efficiency smoke runs (algo.max_iterations=100):
    • uv run train --algo sac --task g1_walk_flat --sim mujoco
    • uv run train --algo flashsac --task g1_motion_tracking --sim mujoco
  4. Repo gates on the integration head: make check, make test, focused slow contract tests for the two tasks; make test-all before the final PR.
  5. Acceptance: one full SAC g1_walk_flat --sim mujoco training run on macOS (owner YAML default algo.max_iterations), normal shutdown, metrics/logs inspected.

Explicitly excluded: pure-CPU off-policy learner device, multi-GPU on macOS, new public contracts, new ADRs unless a structural decision emerges, macOS CI jobs, PyPI publication, main/develop/tensor-runtime mutation.

Acceptance criteria

  • import unilab.base.torch_env and backend construction succeed on macOS with the workspace materialized.
  • SAC + g1_walk_flat --sim mujoco trains on macOS (CPU physics, MPS learner) for algo.max_iterations=100 with normal shutdown and expected metric schema.
  • FlashSAC + g1_motion_tracking --sim mujoco trains on macOS for algo.max_iterations=100 with normal shutdown.
  • make check, make test, and focused slow tests for the two tasks pass on the integration head.
  • Full SAC g1_walk_flat --sim mujoco training completes on macOS (default iteration budget, normal shutdown, collector exit 0, no residual processes).
  • No diff touches npenv-runtime, np-rng, CUDA paths' behavior, main, or develop/tensor-runtime.

Governance

Declared base: develop/tensor-runtime (UniLab) and the tensor_runtime_workspace.json sibling pins. Integration branch per repo: dev/issue-<N>-macos-tensor-runtime. No PyPI release. Related: #1706 (NpEnv removal), #2052 (Newton canonical tensor runtime), tensor-runtime sibling dependency transition policy in 5-contributing_workflow.md.

Activity

  1. TATP-233 commented on Oct 4, 2026

    @TATP-233
    TATP-233
    on Oct 4, 2026
    CollaboratorAuthor

    Progress update (macOS bring-up complete, acceptance run in flight)

    Root cause of the first macOS failure

    flashsac + g1_motion_tracking --sim mujoco failed at the first autoreset with
    RuntimeError: ManagerBased tensor reset-state mutation requires an active reset
    (TensorMotionCommand._resample_command → write_motion_state_tensor_to_sim).

    Diagnosis: the env-side reset dispatch in ManagerBasedRlEnv.reset only opens a
    tensor transaction through the packed host bridge when
    ResetStateTransaction.can_commit_packed() is true. That probe duck-types
    SimBackend.get_public_state_widths(), which did not exist at the sibling
    commits pinned by tensor_runtime_workspace.json (unisim 09d5caed). The
    contract exists at the develop/tensor-runtime branch tips (unisim
    5e835923, PublicStateWidths + mujoco/genesis/mjwarp implementations; the
    mujoco host plan's apply_reset already validates (rows, nq=36)/(rows, nv=35)
    parity with the Manager layout). The stale pins also predated unilab-rl
    single-device/inference-transport fixes (1d83704, e5b3488).

    Base realignment (recorded decision)

    Sibling integration branches are based on the develop/tensor-runtime branch
    tips, not the stale manifest pins; the manifest was updated on the UniLab
    integration branch (commit d2950202). This matches UniLab HEAD's own test
    expectations (tests/envs/test_env_configs.py calls
    backend.get_public_state_widths()), so the pins were stale relative to the
    declared base branch. Bases now:

    • unisim 5e835923 (= tip), unilab-rl 036fcba6 (= tip), mjbatch-uni b3a3b480 (= tip == main, unchanged)
    • UniLab dev/issue-2055-macos-tensor-runtime from develop/tensor-runtime (90068b82)

    With this, can_commit_packed() == True for the mujoco G1 owners and the reset
    dispatch takes scoped_device_event_tensor_with_host_commit — no UniLab
    runtime code change was needed for the training path itself.

    Training evidence (macOS arm64, torch 2.14.0, MPS learner, CPU physics)

    Run Result
    uv run train --algo sac --task g1_walk_flat --sim mujoco algo.max_iterations=100 completed 100/100, 225,280 steps, ~18.2k steps/s, normal_completion, collector exit 0, cleanup errors []
    uv run train --algo flashsac --task g1_motion_tracking --sim mujoco algo.max_iterations=100 completed 100/100, 405,504 steps, normal_completion, collector exit 0

    run_summary.json confirms env_public_device=cpu, learner_device=mps,
    cuda_context_initialized=false, tensor_native_env=true,
    cuda_graph.backend=eager (correctly inert off CUDA), metric_schema_version=1.
    No residual collector processes after either run. SAC re-run after the base
    change: same green result.

    Gates on the integration head

    • make check: pass (ruff, mypy 144 files, pyright 0 errors).
    • make test (non-slow): 1801 passed, 59 skipped — was 1724 passed /
      62 failed / 18 errors before the test-layer fixes (commit f2fb4519):
      • the fake-runtime/mock-worker CUDA-only backend suites
        (genesis/isaacgym/isaacsim/fixed-variants) now bypass the factory
        platform gate (which has dedicated coverage in
        test_cuda_backend_platform_preflight.py), restoring 58 failed + 18
        errored tests on all CUDA-less hosts;
      • targeted skips: CUDA observation naming test, mjwarp-extra soak command
        test, Linux-/proc soak resource test; dropped the never-committed
        outputs/velb_compare/micro_velb.py benchmark-scope entry.
    • Focused slow lane: test_env_configs.py -k mujoco 5 passed (now includes a
      new selected-row reset regression case through the packed host-bridge tensor
      command commit, commit b6c293fe), G1 owner contract 17 passed / 4 skipped
      (newton cases now skip off-CUDA), flashsac owner contract 10 passed.

    Acceptance

    Full-budget sac + g1_walk_flat --sim mujoco (owner default
    algo.max_iterations=5000) is running unattended with
    training.play_render_mode=record (avoids the interactive mjpython viewer;
    playback + ONNX export still exercised). Result follows.

  2. TATP-233 commented on Oct 4, 2026

    @TATP-233
    TATP-233
    on Oct 4, 2026
    CollaboratorAuthor

    Acceptance complete — roadmap delivered on the integration head

    Full-budget acceptance run on macOS arm64 (torch 2.14.0, CPU physics + MPS learner):

    uv run train --algo sac --task g1_walk_flat --sim mujoco training.play_render_mode=record
    • 5000/5000 iterations, 10,260,480 env steps in 11.8 min (~14,305 steps/s,
      457,760 replay rows/s)       — normal_completion, collector exitcode 0,
      cleanup errors [], no residual processes.
    • Learned gait: final mean reward 278.98, best 303.60, mean
      episode length 946.85 (from ~49 at init).
    • Post-train play on the final checkpoint: ONNX export verified
      (max_diff 7.84e-07), 800-frame playback video rendered
      (play_render_mode=record to avoid the interactive mjpython viewer in an
      unattended run; playback + export paths fully exercised).
    • run_summary.json: env_public_device=cpu, learner_device=mps,
      cuda_context_initialized=false, cuda_graph.backend=eager,
      metric_schema_version=1.

    Pre-PR gate on the integration head: make test-all pass (check + coverage
    lane + benchmark smoke 35/35, 36/36).

    PR: #2061 (dev/issue-2055-macos-tensor-runtime → develop/tensor-runtime).
    Sibling repos (unisim, unilab_rl, mjbatch_uni) required zero code changes —
    their integration branches track the develop/tensor-runtime tips recorded in
    the updated tensor_runtime_workspace.json. All acceptance criteria checked.
    Remote CI will run when the PR to main is opened later (per workflow, local
    results + review are the gate for a non-main base).

  3. TATP-233
    closed this as completedon Oct 8, 2026
  4. EtherealTide
    added 3 commits that reference this issue on Oct 8, 2026

    chore(workspace): realign sibling pins with develop/tensor-runtime heads

    d295020

    test(platform): make the non-slow suite pass on CUDA-less hosts

    f2fb451

    test(g1): cover mujoco motion selected reset and skip newton cases of…

    b6c293f
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions