MuJoCo MJX — legged-robot rollouts (GPU)

super run . --script run.sh

"""Batched GPU rollouts of a legged robot (Ant) with MuJoCo MJX physics via
brax — a robotics trajectory dataset plus a smoke that the MJX / JAX-CUDA stack
runs on supercomputer.

brax's `backend="mjx"` runs the MuJoCo MJX rigid-body engine under JAX/XLA, so we
step 1024 Ant robots in parallel on the GPU under a random policy and record the
per-step reward (the raw material for offline RL / system-id on a legged robot).
Writes a summary + reward trace to /workspace/out/mujoco_mjx.json.
"""

import json
import os
import time
from pathlib import Path

import jax
import jax.numpy as jp
from brax import envs

OUT = Path("/workspace/out/mujoco_mjx.json")
OUT.parent.mkdir(parents=True, exist_ok=True)

ENV_NAME = "ant"      # a legged (quadruped) robot
NUM_ENVS = 1024
STEPS = 200
# Soak knob (stress tier): keep rolling out STEPS-sized chunks until this many
# seconds have elapsed; the artifact is written only after the loop. 0 = 1 chunk.
STRESS_DURATION_S = int(os.environ.get("SUPERCOMPUTER_STRESS_DURATION_S", "0") or "0")


def main() -> None:
    dev = jax.devices()[0]
    assert dev.platform == "gpu", f"JAX not on GPU: {dev}"

    env = envs.get_environment(env_name=ENV_NAME, backend="mjx")  # MuJoCo MJX physics
    reset = jax.jit(jax.vmap(env.reset))
    step = jax.jit(jax.vmap(env.step))

    key = jax.random.PRNGKey(0)
    state = reset(jax.random.split(key, NUM_ENVS))

    # Compile (jit) + warm up before timing so the throughput number is real.
    key, sub = jax.random.split(key)
    act = jax.random.uniform(sub, (NUM_ENVS, env.action_size), minval=-1.0, maxval=1.0)
    state = step(state, act)
    jax.block_until_ready(state.reward)

    rewards = []
    deadline = time.monotonic() + STRESS_DURATION_S
    chunks = 0
    t0 = time.time()
    while True:
        for _ in range(STEPS):
            key, sub = jax.random.split(key)
            act = jax.random.uniform(sub, (NUM_ENVS, env.action_size), minval=-1.0, maxval=1.0)
            state = step(state, act)
            rewards.append(float(jp.mean(state.reward)))
        chunks += 1
        if time.monotonic() >= deadline:
            break
    jax.block_until_ready(state.reward)
    dt = time.time() - t0

    OUT.write_text(json.dumps({
        "env": ENV_NAME,
        "robot": "ant_quadruped",
        "physics": "mujoco_mjx",
        "num_envs": NUM_ENVS,
        "steps": STEPS * chunks,
        "rollout_chunks": chunks,
        "total_env_steps": NUM_ENVS * STEPS * chunks,
        "sim_steps_per_sec": round(NUM_ENVS * STEPS * chunks / dt, 1),
        "wall_seconds": round(dt, 2),
        "action_dim": int(env.action_size),
        "mean_reward": round(sum(rewards) / len(rewards), 4),
        "reward_trace_head": [round(r, 4) for r in rewards[:10]],
        "device": str(dev),
    }, indent=2) + "\n")
    print("wrote", OUT)


if __name__ == "__main__":
    main()

#!/usr/bin/env bash
# Quadruped (Unitree Go1) locomotion rollouts batched on GPU with MuJoCo MJX —
# produces a robotics trajectory dataset and proves the MJX / JAX-CUDA physics
# stack works end-to-end on supercomputer. See README.md.
set -euo pipefail

# Fail loudly if JAX can't see the GPU — MJX physics is the whole point.
python -c "import jax; d=jax.devices(); print('jax devices:', d); assert d[0].platform=='gpu', 'JAX is on CPU'"

exec python "$(dirname "$0")/rollout.py"

{
    "sku": "gpu_1x_l4",
    "image": {
        "base": "pytorch/pytorch:2.8.0-cuda12.8-cudnn9-devel",
        "apt": ["build-essential", "git", "libgl1", "libegl1", "libglfw3", "ca-certificates"],
        "pip": ["jax[cuda12]", "brax"],
        "env": {
            "PYTHONUNBUFFERED": "1",
            "MUJOCO_GL": "egl",
            "XLA_PYTHON_CLIENT_PREALLOCATE": "false",
            "SSL_CERT_FILE": "/etc/ssl/certs/ca-certificates.crt",
            "PIP_CERT": "/etc/ssl/certs/ca-certificates.crt",
            "REQUESTS_CA_BUNDLE": "/etc/ssl/certs/ca-certificates.crt"
        }
    }
}