Enhance installation instructions.

PiperOrigin-RevId: 719937540
Change-Id: Ib33cf8c0ca5da144ffc21221f6ea8e3299e38422

.gitignore

@@ -1,5 +1,7 @@
 .DS_Store
 __pycache__
 *.egg-info
+.venv
 MUJOCO_LOG.TXT
 mujoco_menagerie
+checkpoints/

README.md

@@ -28,11 +28,14 @@ pip install playground
 > [!IMPORTANT]
 > Requires Python 3.10 or later.
 
-1. `pip install -U "jax[cuda12]"`
-    * Verify GPU backend: python -c "import jax; print(jax.default_backend())" should print gpu
-2. `git clone [email protected]:google-deepmind/mujoco_playground.git`
-3. `cd mujoco_playground`
-4. `uv pip install -e ".[all]"`
+1. `git clone [email protected]: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`
+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"`
 
 #### Madrona-MJX (optional)
 

pyproject.toml

@@ -23,6 +23,7 @@ classifiers = [
     "Topic :: Scientific/Engineering",
 ]
 dependencies = [
+    "brax>=0.12.1",
     "etils",
     "flax",
     "jax",
@@ -31,20 +32,21 @@ dependencies = [
     "mujoco-mjx>=3.2.7",
     "mujoco>=3.2.7",
     "tqdm",
-    "brax>=0.12.1",
 ]
-keywords = ["mjx", "mujoco", "dm_control"]
+keywords = ["mjx", "mujoco", "sim2real", "reinforcement learning"]
 
 [project.optional-dependencies]
 test = [
     "absl-py",
+    "dm_control",
     "pytest",
     "pytest-xdist",
-    "dm_control",
 ]
 notebooks = [
     "matplotlib",
     "mediapy",
+    "jupyter",
+    "wandb",
 ]
 dev = [
     "playground[test]",