Merge branch 'main' into warp-rsl

Author: btaba

CHANGELOG.md

@@ -2,7 +2,7 @@
 
 All notable changes to this project will be documented in this file.
 
-## Next release
+## [0.1.0] - 2026-01-07
 
 - Pass through the [MuJoCo Warp](https://github.com/google-deepmind/mujoco_warp)
   (MjWarp) implementation to MJX, so that MuJoCo Playground environments can
@@ -16,6 +16,7 @@ All notable changes to this project will be documented in this file.
 - Update AutoResetWrapper to allow full resets on done. Fixes #179. Also
   provides a means for doing curriculum learning via
   `state.info['AutoResetWrapper_done_count']`, see #140.
+- Update dependencies to use `mujoco>=3.4` and `warp-lang>=1.11`.
 
 ## [0.0.5] - 2025-06-23
 

README.md

@@ -16,7 +16,7 @@ Features include:
 For more details, check out the project [website](https://playground.mujoco.org/).
 
 > [!NOTE]
-> We now support training with both the MuJoCo MJX JAX implementation, as well as the [MuJoCo Warp](https://github.com/google-deepmind/mujoco_warp) implementation at HEAD. See MuJoCo 3.3.5 [release notes](https://mujoco.readthedocs.io/en/stable/changelog.html#version-3-3-5-august-8-2025) under `MJX` for more details.
+> We now support training with both the MuJoCo MJX JAX implementation, as well as the [MuJoCo Warp](https://github.com/google-deepmind/mujoco_warp) implementation at HEAD. See this [discussion post](https://github.com/google-deepmind/mujoco_playground/discussions/197) for more details.
 
 ## Installation
 
@@ -26,34 +26,50 @@ You can install MuJoCo Playground directly from PyPI:
 pip install playground
 ```
 
-> [!WARNING]
-> The `playground` release may depend on pre-release versions of `mujoco` and
-> `warp-lang`, in which case you can try `pip install playground
-> --extra-index-url=https://py.mujoco.org
-> --extra-index-url=https://pypi.nvidia.com/warp-lang/`.
-> If there are still version mismatches, please open a github issue, and install
-> from source.
+> [!IMPORTANT]
+> We recommend users to install [from source](#from-source) to get the latest features and bug fixes from MuJoCo.
 
-### From Source
+### <a id="from-source">From Source</a>
 
 > [!IMPORTANT]
 > Requires Python 3.10 or later.
 
 1. `git clone git@github.com:google-deepmind/mujoco_playground.git && cd mujoco_playground`
 2. [Install uv](https://docs.astral.sh/uv/getting-started/installation/), a faster alternative to `pip`
-3. Create a virtual environment: `uv venv --python 3.11`
+3. Create a virtual environment: `uv venv --python 3.12`
 4. Activate it: `source .venv/bin/activate`
-5. Install CUDA 12 jax: `uv pip install -U "jax[cuda12]"`
-    * Verify GPU backend: `python -c "import jax; print(jax.default_backend())"` should print gpu
-6. Install playground: `uv pip install -e ".[all]"`
-7. Verify installation (and download Menagerie): `python -c "import mujoco_playground"`
+5. Install CUDA 12 jax: `uv pip install -U "jax[cuda12]" --index-url https://pypi.org/simple`
+    * Verify GPU backend: `python -c "import jax; print(jax.default_backend())"` should print gpu. `unset LD_LIBRARY_PATH` may need to be run before running this command.
+6. Install playground from source: `uv --no-config sync --all-extras`
+7. Verify installation: `uv --no-config run python -c "import mujoco_playground; print('Success')"`
+    * **Note**: Menagerie assets will be downloaded automatically the first time you load a locomotion or manipulation environment. You can trigger this with: `uv --no-config run python -c "from mujoco_playground import locomotion; locomotion.load('G1JoystickFlatTerrain')"`
 
 #### Madrona-MJX (optional)
 
 For vision-based environments, please refer to the installation instructions in the [Madrona-MJX](https://github.com/shacklettbp/madrona_mjx?tab=readme-ov-file#installation) repository.
 
 ## Getting started
 
+### Running from CLI
+For basic usage, navigate to the repo's directory, install [from source](#from-source) with `jax[cuda12]`, and run:
+
+```bash
+train-jax-ppo --env_name CartpoleBalance
+```
+
+To train with [MuJoCo Warp](https://github.com/google-deepmind/mujoco_warp):
+
+```bash
+train-jax-ppo --env_name CartpoleBalance --impl warp
+```
+
+Or with `uv`:
+
+```bash
+uv --no-config run train-jax-ppo --env_name CartpoleBalance --impl warp
+uv --no-config run train-rsl-ppo --env_name CartpoleBalance --impl warp
+```
+
 ### Basic Tutorials
 | Colab | Description |
 |-------|-------------|
@@ -74,15 +90,6 @@ For vision-based environments, please refer to the installation instructions in
 | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/training_vision_1.ipynb) | Training CartPole from Vision |
 | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/training_vision_2.ipynb) | Robotic Manipulation from Vision |
 
-## Running from CLI
-> [!IMPORTANT]
-> Assumes installation from source.
-
-For basic usage, navigate to the repo's directory and run:
-```bash
-python learning/train_jax_ppo.py --env_name CartpoleBalance
-```
-
 ### Training Visualization
 
 To interactively view trajectories throughout training with [rscope](https://github.com/Andrew-Luo1/rscope/tree/main), install it (`pip install rscope`) and run:

learning/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2026 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Learning scripts for MuJoCo Playground."""

learning/notebooks/locomotion.ipynb

@@ -29,9 +29,9 @@
         "id": "_UbO9uhtBSX5"
       },
       "source": [
-        "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eCopyright 2025 DeepMind Technologies Limited.\u003c/small\u003e\u003c/p\u003e\n",
-        "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eLicensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at \u003ca href=\"http://www.apache.org/licenses/LICENSE-2.0\"\u003ehttp://www.apache.org/licenses/LICENSE-2.0\u003c/a\u003e.\u003c/small\u003e\u003c/small\u003e\u003c/p\u003e\n",
-        "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eUnless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.\u003c/small\u003e\u003c/small\u003e\u003c/p\u003e"
+        "> <p><small><small>Copyright 2025 DeepMind Technologies Limited.</small></p>\n",
+        "> <p><small><small>Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at <a href=\"http://www.apache.org/licenses/LICENSE-2.0\">http://www.apache.org/licenses/LICENSE-2.0</a>.</small></small></p>\n",
+        "> <p><small><small>Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.</small></small></p>"
       ]
     },
     {
@@ -40,11 +40,11 @@
         "id": "dNIJkb_FM2Ux"
       },
       "source": [
-        "# Locomotion in The Playground! \u003ca href=\"https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/locomotion.ipynb\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/\u003e\u003c/a\u003e\n",
+        "# Locomotion in The Playground! <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/locomotion.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/></a>\n",
         "\n",
         "In this notebook, we'll walk through a few locomotion environments available in MuJoCo Playground.\n",
         "\n",
-        "**A Colab runtime with GPU acceleration is required.** If you're using a CPU-only runtime, you can switch using the menu \"Runtime \u003e Change runtime type\".\n"
+        "**A Colab runtime with GPU acceleration is required.** If you're using a CPU-only runtime, you can switch using the menu \"Runtime > Change runtime type\".\n"
       ]
     },
     {
@@ -107,7 +107,7 @@
         "  print('Checking that the installation succeeded:')\n",
         "  import mujoco\n",
         "\n",
-        "  mujoco.MjModel.from_xml_string('\u003cmujoco/\u003e')\n",
+        "  mujoco.MjModel.from_xml_string('<mujoco/>')\n",
         "except Exception as e:\n",
         "  raise e from RuntimeError(\n",
         "      'Something went wrong during installation. Check the shell output above '\n",
@@ -142,7 +142,7 @@
         "\n",
         "# Graphics and plotting.\n",
         "print(\"Installing mediapy:\")\n",
-        "!command -v ffmpeg \u003e/dev/null || (apt update \u0026\u0026 apt install -y ffmpeg)\n",
+        "!command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)\n",
         "!pip install -q mediapy\n",
         "import mediapy as media\n",
         "import matplotlib.pyplot as plt\n",
@@ -476,11 +476,11 @@
         "command = jp.array([x_vel, y_vel, yaw_vel])\n",
         "\n",
         "state = jit_reset(rng)\n",
-        "if state.info[\"steps_since_last_pert\"] \u003c state.info[\"steps_until_next_pert\"]:\n",
+        "if state.info[\"steps_since_last_pert\"] < state.info[\"steps_until_next_pert\"]:\n",
         "  rng = sample_pert(rng)\n",
         "state.info[\"command\"] = command\n",
         "for i in range(env_cfg.episode_length):\n",
-        "  if state.info[\"steps_since_last_pert\"] \u003c state.info[\"steps_until_next_pert\"]:\n",
+        "  if state.info[\"steps_since_last_pert\"] < state.info[\"steps_until_next_pert\"]:\n",
         "    rng = sample_pert(rng)\n",
         "  act_rng, rng = jax.random.split(rng)\n",
         "  ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
@@ -651,7 +651,7 @@
         "    print(f\"Setting x to {x}\")\n",
         "    command = jp.array([x, 0, 0])\n",
         "  state.info[\"command\"] = command\n",
-        "  if state.info[\"steps_since_last_pert\"] \u003c state.info[\"steps_until_next_pert\"]:\n",
+        "  if state.info[\"steps_since_last_pert\"] < state.info[\"steps_until_next_pert\"]:\n",
         "    rng = sample_pert(rng)\n",
         "  act_rng, rng = jax.random.split(rng)\n",
         "  ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",

learning/notebooks/manipulation.ipynb

@@ -29,9 +29,9 @@
         "id": "_UbO9uhtBSX5"
       },
       "source": [
-        "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eCopyright 2025 DeepMind Technologies Limited.\u003c/small\u003e\u003c/p\u003e\n",
-        "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eLicensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at \u003ca href=\"http://www.apache.org/licenses/LICENSE-2.0\"\u003ehttp://www.apache.org/licenses/LICENSE-2.0\u003c/a\u003e.\u003c/small\u003e\u003c/small\u003e\u003c/p\u003e\n",
-        "\u003e \u003cp\u003e\u003csmall\u003e\u003csmall\u003eUnless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.\u003c/small\u003e\u003c/small\u003e\u003c/p\u003e"
+        "> <p><small><small>Copyright 2025 DeepMind Technologies Limited.</small></p>\n",
+        "> <p><small><small>Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at <a href=\"http://www.apache.org/licenses/LICENSE-2.0\">http://www.apache.org/licenses/LICENSE-2.0</a>.</small></small></p>\n",
+        "> <p><small><small>Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.</small></small></p>"
       ]
     },
     {
@@ -40,11 +40,11 @@
         "id": "dNIJkb_FM2Ux"
       },
       "source": [
-        "# Manipulation in The Playground! \u003ca href=\"https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/manipulation.ipynb\"\u003e\u003cimg src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/\u003e\u003c/a\u003e\n",
+        "# Manipulation in The Playground! <a href=\"https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/manipulation.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" width=\"140\" align=\"center\"/></a>\n",
         "\n",
         "In this notebook, we'll walk through a couple manipulation environments available in MuJoCo Playground.\n",
         "\n",
-        "**A Colab runtime with GPU acceleration is required.** If you're using a CPU-only runtime, you can switch using the menu \"Runtime \u003e Change runtime type\".\n"
+        "**A Colab runtime with GPU acceleration is required.** If you're using a CPU-only runtime, you can switch using the menu \"Runtime > Change runtime type\".\n"
       ]
     },
     {
@@ -107,7 +107,7 @@
         "  print('Checking that the installation succeeded:')\n",
         "  import mujoco\n",
         "\n",
-        "  mujoco.MjModel.from_xml_string('\u003cmujoco/\u003e')\n",
+        "  mujoco.MjModel.from_xml_string('<mujoco/>')\n",
         "except Exception as e:\n",
         "  raise e from RuntimeError(\n",
         "      'Something went wrong during installation. Check the shell output above '\n",
@@ -142,7 +142,7 @@
         "\n",
         "# Graphics and plotting.\n",
         "print(\"Installing mediapy:\")\n",
-        "!command -v ffmpeg \u003e/dev/null || (apt update \u0026\u0026 apt install -y ffmpeg)\n",
+        "!command -v ffmpeg >/dev/null || (apt update && apt install -y ffmpeg)\n",
         "!pip install -q mediapy\n",
         "import mediapy as media\n",
         "import matplotlib.pyplot as plt\n",

learning/train_jax_ppo.py

@@ -40,7 +40,11 @@
 from mujoco_playground.config import locomotion_params
 from mujoco_playground.config import manipulation_params
 import tensorboardX
-import wandb
+
+try:
+  import wandb
+except ImportError:
+  wandb = None
 
 
 xla_flags = os.environ.get("XLA_FLAGS", "")
@@ -298,6 +302,11 @@ def main(argv):
 
   # Initialize Weights & Biases if required
   if _USE_WANDB.value and not _PLAY_ONLY.value:
+    if wandb is None:
+      raise ImportError(
+          "wandb is required for --use_wandb. "
+          "Install via: pip install wandb"
+      )
     wandb.init(project="mjxrl", name=exp_name)
     wandb.config.update(env_cfg.to_dict())
     wandb.config.update({"env_name": _ENV_NAME.value})
@@ -530,5 +539,10 @@ def step(carry, _):
     print(f"Rollout video saved as 'rollout{i}.mp4'.")
 
 
-if __name__ == "__main__":
+def run():
+  """Entry point for uv/pip script."""
   app.run(main)
+
+
+if __name__ == "__main__":
+  run()