Fix #212.

PiperOrigin-RevId: 814756605
Change-Id: Ic17f2ad28529a674a0f4ed10a9f42cc65e4b750e

Author: btaba

pyproject.toml

@@ -34,7 +34,6 @@ dependencies = [
     "mujoco>=3.3.6.dev",
     "orbax-checkpoint>=0.11.22",
     "tqdm",
-    "warp-lang>=1.9.0.dev",
 ]
 keywords = ["mjx", "mujoco", "sim2real", "reinforcement learning"]
 
@@ -54,6 +53,9 @@ warp-lang = { index = "nvidia" }
 mujoco-mjx = { git = "https://github.com/google-deepmind/mujoco", rev = "main", subdirectory = "mjx" }
 
 [project.optional-dependencies]
+warp = [
+    "warp-lang>=1.9.0.dev",
+]
 test = [
     "absl-py",
     "dm_control",
@@ -80,6 +82,7 @@ learning = [
 ]
 all = [
     "playground[dev]",
+    "playground[warp]",
     "playground[notebooks]",
     "playground[learning]",
 ]