{ "cells": [ { "cell_type": "markdown", "id": "fe88c0e3", "metadata": {}, "source": [ "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tayalmanan28/MuJoCo-Tutorial/blob/main/tutorial/07_mjx.ipynb)\n", "\n", "⚡ **This notebook requires a GPU runtime.** In Colab: Runtime → Change runtime type → T4 GPU\n", "\n", "# Tutorial 7: MJX — GPU-Accelerated Simulation\n", "\n", "MJX is MuJoCo's JAX backend for massively parallel simulation on GPU.\n", "Perfect for reinforcement learning where you need thousands of environments." ] }, { "cell_type": "code", "execution_count": null, "id": "e5be1573", "metadata": {}, "outputs": [], "source": [ "!pip install -q mujoco mujoco-mjx jax[cuda12] mediapy" ] }, { "cell_type": "code", "execution_count": null, "id": "cde1b063", "metadata": {}, "outputs": [], "source": [ "import mujoco\n", "from mujoco import mjx\n", "import jax\n", "import jax.numpy as jnp\n", "import numpy as np\n", "import mediapy as media\n", "import time\n", "\n", "print(f\"JAX devices: {jax.devices()}\")\n", "print(f\"MuJoCo version: {mujoco.__version__}\")" ] }, { "cell_type": "markdown", "id": "ee43d2fe", "metadata": {}, "source": [ "## MJX Basics\n", "\n", "MJX converts MuJoCo models to JAX-compatible pytrees that run on GPU:\n", "\n", "1. Create a standard `MjModel`\n", "2. Convert to MJX with `mjx.put_model()`\n", "3. Use `mjx.step()` instead of `mujoco.mj_step()`" ] }, { "cell_type": "code", "execution_count": null, "id": "032a343b", "metadata": {}, "outputs": [], "source": [ "xml = \"\"\"\n", "\n", " \n", "\"\"\"\n", "\n", "# Standard MuJoCo model\n", "mj_model = mujoco.MjModel.from_xml_string(xml)\n", "mj_data = mujoco.MjData(mj_model)\n", "\n", "# Convert to MJX\n", "mjx_model = mjx.put_model(mj_model)\n", "mjx_data = mjx.put_data(mj_model, mj_data)\n", "\n", "print(f\"MJX model type: {type(mjx_model)}\")\n", "print(f\"MJX data qpos: {mjx_data.qpos}\")" ] }, { "cell_type": "markdown", "id": "5da4e3e8", "metadata": {}, "source": [ "## JIT-Compiled Stepping" ] }, { "cell_type": "code", "execution_count": null, "id": "e3ceabf6", "metadata": {}, "outputs": [], "source": [ "# JIT compile the step function\n", "jit_step = jax.jit(mjx.step)\n", "\n", "# First call triggers compilation\n", "mjx_data = jit_step(mjx_model, mjx_data)\n", "\n", "# Benchmark: 1000 steps\n", "t0 = time.time()\n", "for _ in range(1000):\n", " mjx_data = jit_step(mjx_model, mjx_data)\n", "mjx_data.qpos.block_until_ready() # wait for GPU\n", "elapsed = time.time() - t0\n", "\n", "print(f\"1000 steps in {elapsed*1000:.1f}ms ({1000/elapsed:.0f} steps/sec)\")\n", "print(f\"Final height: {float(mjx_data.qpos[2]):.4f}\")" ] }, { "cell_type": "markdown", "id": "a4ec3c40", "metadata": {}, "source": [ "## Batched (Parallel) Simulation\n", "\n", "The real power of MJX: simulate thousands of environments simultaneously using `jax.vmap`." ] }, { "cell_type": "code", "execution_count": null, "id": "60ac78f9", "metadata": {}, "outputs": [], "source": [ "batch_size = 4096\n", "\n", "# Create batch of initial states with random velocities\n", "rng = jax.random.PRNGKey(0)\n", "random_vels = jax.random.normal(rng, (batch_size, 6)) * jnp.array([2, 2, 5, 0.5, 0.5, 0.5])\n", "\n", "# Reset base data\n", "mj_data_reset = mujoco.MjData(mj_model)\n", "mjx_data_base = mjx.put_data(mj_model, mj_data_reset)\n", "\n", "# Batch: each environment gets a different initial velocity\n", "batched_data = jax.vmap(lambda v: mjx_data_base.replace(qvel=v))(random_vels)\n", "\n", "print(f\"Batched qpos shape: {batched_data.qpos.shape}\") # (4096, 7)\n", "print(f\"Simulating {batch_size} environments in parallel...\")" ] }, { "cell_type": "code", "execution_count": null, "id": "42782767", "metadata": {}, "outputs": [], "source": [ "# Batched step function\n", "@jax.jit\n", "def batched_step(data):\n", " return jax.vmap(mjx.step, in_axes=(None, 0))(mjx_model, data)\n", "\n", "# Simulate 500 steps for all environments\n", "n_steps = 500\n", "\n", "t0 = time.time()\n", "data = batched_data\n", "for _ in range(n_steps):\n", " data = batched_step(data)\n", "data.qpos.block_until_ready()\n", "elapsed = time.time() - t0\n", "\n", "total_steps = batch_size * n_steps\n", "print(f\"{total_steps:,} total steps in {elapsed:.2f}s\")\n", "print(f\"Throughput: {total_steps/elapsed:,.0f} steps/sec\")\n", "print(f\"\\nFinal heights — mean: {float(data.qpos[:, 2].mean()):.3f}, std: {float(data.qpos[:, 2].std()):.3f}\")" ] }, { "cell_type": "markdown", "id": "d0aa98e8", "metadata": {}, "source": [ "## Rendering MJX Results\n", "\n", "MJX runs on GPU but rendering uses CPU MuJoCo. Transfer data back:" ] }, { "cell_type": "code", "execution_count": null, "id": "4ef5e821", "metadata": {}, "outputs": [], "source": [ "# Simulate one trajectory and render\n", "mj_data_viz = mujoco.MjData(mj_model)\n", "mjx_data_viz = mjx.put_data(mj_model, mj_data_viz)\n", "mjx_data_viz = mjx_data_viz.replace(qvel=jnp.array([3.0, 0, 8.0, 0, 0, 0]))\n", "\n", "renderer = mujoco.Renderer(mj_model, height=300, width=400)\n", "frames = []\n", "\n", "for i in range(300):\n", " mjx_data_viz = jit_step(mjx_model, mjx_data_viz)\n", " \n", " if i % 5 == 0: # render every 5th step (for 60fps at dt=0.002)\n", " # Transfer back to CPU for rendering\n", " mjx.get_data_into(mj_data_viz, mj_model, mjx_data_viz)\n", " renderer.update_scene(mj_data_viz)\n", " frames.append(renderer.render().copy())\n", "\n", "renderer.close()\n", "media.show_video(frames, fps=30)" ] }, { "cell_type": "markdown", "id": "2079396c", "metadata": {}, "source": [ "## When to Use MJX vs CPU MuJoCo\n", "\n", "| Use Case | Recommendation |\n", "|----------|---------------|\n", "| Single environment, interactive | CPU MuJoCo + `mujoco.viewer` |\n", "| RL training (PPO, SAC) | MJX — batch 1000s of envs |\n", "| System identification | MJX — differentiable physics |\n", "| Trajectory optimization | MJX — autodiff through dynamics |\n", "| Complex contact-rich tasks | CPU MuJoCo (MJX has some contact limitations) |\n", "\n", "### Key Differences\n", "- MJX supports most but not all MuJoCo features (check [docs](https://mujoco.readthedocs.io/en/stable/mjx.html#feature-parity))\n", "- MJX is **differentiable** — you can compute gradients through the physics\n", "- Ideal batch size: 1024–8192 for best GPU utilization" ] } ], "metadata": { "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }