Merge branch 'add-quadruped' of github.com:yardenas/mujoco_playground into add-quadruped

Author: yardenas

CHANGELOG.md

@@ -2,6 +2,15 @@
 
 All notable changes to this project will be documented in this file.
 
+## Next release
+
+- Pass through the [MuJoCo Warp](https://github.com/google-deepmind/mujoco_warp)
+  (MjWarp) implementation to MJX, so that MuJoCo Playground environments can
+  train with MuJoCo Warp! DM Control Suite and Locomotion environments now
+  support MjWarp. You can pass through the implementation via the config
+  override
+  `registry.load('CartpoleBalance', config_overrides={'impl': 'warp'})`.
+
 ## [0.0.5] - 2025-06-23
 
 - Change `light_directional` to `light_type` following MuJoCo API change from version 3.3.2 to 3.3.3. Fixes https://github.com/google-deepmind/mujoco_playground/issues/142.

mujoco_playground/_src/collision.py

@@ -17,22 +17,29 @@
 from typing import Any, Tuple
 
 import jax
-import jax.numpy as jnp
+import jax.numpy as jp
 from mujoco import mjx
+from mujoco.mjx._src import types
 
 
 def get_collision_info(
     contact: Any, geom1: int, geom2: int
 ) -> Tuple[jax.Array, jax.Array]:
   """Get the distance and normal of the collision between two geoms."""
-  mask = (jnp.array([geom1, geom2]) == contact.geom).all(axis=1)
-  mask |= (jnp.array([geom2, geom1]) == contact.geom).all(axis=1)
-  idx = jnp.where(mask, contact.dist, 1e4).argmin()
+  mask = (jp.array([geom1, geom2]) == contact.geom).all(axis=1)
+  mask |= (jp.array([geom2, geom1]) == contact.geom).all(axis=1)
+  idx = jp.where(mask, contact.dist, 1e4).argmin()
   dist = contact.dist[idx] * mask[idx]
   normal = (dist < 0) * contact.frame[idx, 0, :3]
   return dist, normal
 
 
 def geoms_colliding(state: mjx.Data, geom1: int, geom2: int) -> jax.Array:
   """Return True if the two geoms are colliding."""
+  # if not isinstance(state._impl, types.DataJAX):
+  #   raise NotImplementedError(
+  #       "`geoms_colliding` only implemented for JAX MJX backend."
+  #   )
+  if not isinstance(state._impl, types.DataJAX):
+    return jp.array(False)
   return get_collision_info(state._impl.contact, geom1, geom2)[0] < 0  # pylint: disable=protected-access

mujoco_playground/_src/dm_control_suite/acrobot.py

@@ -36,6 +36,9 @@ def default_config() -> config_dict.ConfigDict:
       episode_length=1000,
       action_repeat=1,
       vision=False,
+      impl="jax",
+      nconmax=0,
+      njmax=0,
   )
 
 
@@ -62,7 +65,7 @@ def __init__(
         _XML_PATH.read_text(), self._model_assets
     )
     self._mj_model.opt.timestep = self.sim_dt
-    self._mjx_model = mjx.put_model(self._mj_model)
+    self._mjx_model = mjx.put_model(self._mj_model, impl=self._config.impl)
     self._post_init()
 
   def _post_init(self) -> None:
@@ -78,7 +81,14 @@ def reset(self, rng: jax.Array) -> mjx_env.State:
     qpos = jax.random.uniform(
         rng1, (self.mjx_model.nq,), minval=-jp.pi, maxval=jp.pi
     )
-    data = mjx_env.init(self.mjx_model, qpos=qpos)
+    data = mjx_env.make_data(
+        self.mj_model,
+        qpos=qpos,
+        impl=self.mjx_model.impl.value,
+        nconmax=self._config.nconmax,
+        njmax=self._config.njmax,
+    )
+    data = mjx.forward(self.mjx_model, data)
 
     metrics = {
         "distance": jp.zeros(()),

mujoco_playground/_src/dm_control_suite/ball_in_cup.py

@@ -35,6 +35,9 @@ def default_config() -> config_dict.ConfigDict:
       episode_length=1000,
       action_repeat=1,
       vision=False,
+      impl="jax",
+      nconmax=10_000,
+      njmax=25,
   )
 
 
@@ -58,7 +61,7 @@ def __init__(
         _XML_PATH.read_text(), self._model_assets
     )
     self._mj_model.opt.timestep = self.sim_dt
-    self._mjx_model = mjx.put_model(self._mj_model)
+    self._mjx_model = mjx.put_model(self._mj_model, impl=self._config.impl)
     self._post_init()
 
   def _post_init(self) -> None:
@@ -69,7 +72,13 @@ def _post_init(self) -> None:
     self._ball_size = self._mj_model.geom_size[geom_id, 0]
 
   def reset(self, rng: jax.Array) -> mjx_env.State:
-    data = mjx_env.init(self.mjx_model)
+    data = mjx_env.make_data(
+        self.mj_model,
+        impl=self.mjx_model.impl.value,
+        nconmax=self._config.nconmax,
+        njmax=self._config.njmax,
+    )
+    data = mjx.forward(self.mjx_model, data)
 
     metrics = {}
     info = {"rng": rng}

mujoco_playground/_src/dm_control_suite/cartpole.py

@@ -51,6 +51,9 @@ def default_config() -> config_dict.ConfigDict:
       action_repeat=1,
       vision=False,
       vision_config=default_vision_config(),
+      impl="jax",
+      nconmax=0,
+      njmax=2,
   )
 
 
@@ -95,7 +98,7 @@ def __init__(
         _XML_PATH.read_text(), self._model_assets
     )
     self._mj_model.opt.timestep = self.sim_dt
-    self._mjx_model = mjx.put_model(self._mj_model)
+    self._mjx_model = mjx.put_model(self._mj_model, impl=self._config.impl)
     self._post_init()
 
     if self._vision:
@@ -129,7 +132,7 @@ def _post_init(self) -> None:
     self._hinge_1_qposadr = self._mj_model.jnt_qposadr[hinge_1_jid]
 
   def _reset_swing_up(self, rng: jax.Array) -> jax.Array:
-    rng, rng1, rng2, rng3 = jax.random.split(rng, 4)
+    _, rng1, rng2, rng3 = jax.random.split(rng, 4)
 
     qpos = jp.zeros(self.mjx_model.nq)
     qpos = qpos.at[self._slider_qposadr].set(0.01 * jax.random.normal(rng1))
@@ -163,7 +166,15 @@ def reset(self, rng: jax.Array) -> mjx_env.State:
     rng, rng1 = jax.random.split(rng, 2)
     qvel = 0.01 * jax.random.normal(rng1, (self.mjx_model.nv,))
 
-    data = mjx_env.init(self.mjx_model, qpos=qpos, qvel=qvel)
+    data = mjx_env.make_data(
+        self.mj_model,
+        qpos=qpos,
+        qvel=qvel,
+        impl=self.mjx_model.impl.value,
+        nconmax=self._config.nconmax,
+        njmax=self._config.njmax,
+    )
+    data = mjx.forward(self.mjx_model, data)
 
     metrics = {
         "reward/upright": jp.zeros(()),

mujoco_playground/_src/dm_control_suite/cheetah.py

@@ -38,6 +38,9 @@ def default_config() -> config_dict.ConfigDict:
       episode_length=1000,
       action_repeat=1,
       vision=False,
+      impl="jax",
+      nconmax=100_000,
+      njmax=100,
   )
 
 
@@ -61,7 +64,7 @@ def __init__(
         _XML_PATH.read_text(), self._model_assets
     )
     self._mj_model.opt.timestep = self.sim_dt
-    self._mjx_model = mjx.put_model(self._mj_model)
+    self._mjx_model = mjx.put_model(self._mj_model, impl=self._config.impl)
     self._post_init()
 
   def _post_init(self) -> None:
@@ -81,7 +84,14 @@ def reset(self, rng: jax.Array) -> mjx_env.State:
         )
     )
 
-    data = mjx_env.init(self.mjx_model, qpos=qpos)
+    data = mjx_env.make_data(
+        self.mj_model,
+        qpos=qpos,
+        impl=self.mjx_model.impl.value,
+        nconmax=self._config.nconmax,
+        njmax=self._config.njmax,
+    )
+    data = mjx.forward(self.mjx_model, data)
 
     # Stabilize.
     data = mjx_env.step(self.mjx_model, data, jp.zeros(self.mjx_model.nu), 200)

mujoco_playground/_src/dm_control_suite/finger.py

@@ -48,6 +48,9 @@ def default_config() -> config_dict.ConfigDict:
       episode_length=1000,
       action_repeat=1,
       vision=False,
+      impl="jax",
+      nconmax=25_000,
+      njmax=25,
   )
 
 
@@ -93,7 +96,7 @@ def __init__(
     self._model_assets = common.get_assets()
     self._mj_model = _make_spin_model(_XML_PATH, self._model_assets)
     self._mj_model.opt.timestep = self.sim_dt
-    self._mjx_model = mjx.put_model(self._mj_model)
+    self._mjx_model = mjx.put_model(self._mj_model, impl=self._config.impl)
     self._post_init()
 
   def _post_init(self) -> None:
@@ -108,7 +111,13 @@ def reset(self, rng: jax.Array) -> mjx_env.State:
     )
     qpos = qpos.at[2].set(jax.random.uniform(rng1, minval=-jp.pi, maxval=jp.pi))
 
-    data = mjx_env.init(self.mjx_model, qpos)
+    data = mjx_env.make_data(
+        self.mj_model,
+        qpos=qpos,
+        impl=self.mjx_model.impl.value,
+        nconmax=self._config.nconmax,
+        njmax=self._config.njmax,
+    )
 
     metrics = {}
     info = {"rng": rng}

mujoco_playground/_src/dm_control_suite/fish.py

@@ -46,6 +46,9 @@ def default_config() -> config_dict.ConfigDict:
       episode_length=1000,
       action_repeat=1,
       vision=False,
+      impl="jax",
+      nconmax=0,
+      njmax=25,
   )
 
 
@@ -69,7 +72,7 @@ def __init__(
         _XML_PATH.read_text(), self._model_assets
     )
     self._mj_model.opt.timestep = self.sim_dt
-    self._mjx_model = mjx.put_model(self._mj_model)
+    self._mjx_model = mjx.put_model(self._mj_model, impl=self._config.impl)
     self._post_init()
 
   def _post_init(self) -> None:
@@ -104,7 +107,14 @@ def reset(self, rng: jax.Array) -> mjx_env.State:
         )
     )
 
-    data = mjx_env.init(self.mjx_model, qpos=qpos)
+    data = mjx_env.make_data(
+        self.mj_model,
+        qpos=qpos,
+        impl=self.mjx_model.impl.value,
+        nconmax=self._config.nconmax,
+        njmax=self._config.njmax,
+    )
+    data = mjx.forward(self.mjx_model, data)
 
     # Randomize target position.
     xyz = jax.random.uniform(

mujoco_playground/_src/dm_control_suite/hopper.py

@@ -41,6 +41,9 @@ def default_config() -> config_dict.ConfigDict:
       episode_length=1000,
       action_repeat=1,
       vision=False,
+      impl="jax",
+      nconmax=50_000,
+      njmax=50,
   )
 
 
@@ -78,7 +81,7 @@ def __init__(
         _XML_PATH.read_text(), self._model_assets
     )
     self._mj_model.opt.timestep = self.sim_dt
-    self._mjx_model = mjx.put_model(self._mj_model)
+    self._mjx_model = mjx.put_model(self._mj_model, impl=self._config.impl)
     self._post_init()
 
   def _post_init(self) -> None:
@@ -103,7 +106,14 @@ def reset(self, rng: jax.Array) -> mjx_env.State:
         )
     )
 
-    data = mjx_env.init(self.mjx_model, qpos=qpos)
+    data = mjx_env.make_data(
+        self.mj_model,
+        qpos=qpos,
+        impl=self.mjx_model.impl.value,
+        nconmax=self._config.nconmax,
+        njmax=self._config.njmax,
+    )
+    data = mjx.forward(self.mjx_model, data)
 
     metrics = {k: jp.zeros(()) for k in self._metric_keys}
     info = {"rng": rng}

mujoco_playground/_src/dm_control_suite/humanoid.py

@@ -42,6 +42,9 @@ def default_config() -> config_dict.ConfigDict:
       episode_length=1000,
       action_repeat=1,
       vision=False,
+      impl="jax",
+      nconmax=200_000,
+      njmax=250,
   )
 
 
@@ -72,7 +75,7 @@ def __init__(
         _XML_PATH.read_text(), self._model_assets
     )
     self._mj_model.opt.timestep = self.sim_dt
-    self._mjx_model = mjx.put_model(self._mj_model)
+    self._mjx_model = mjx.put_model(self._mj_model, impl=self._config.impl)
     self._post_init()
 
   def _post_init(self) -> None:
@@ -88,7 +91,13 @@ def _post_init(self) -> None:
   def reset(self, rng: jax.Array) -> mjx_env.State:
     # TODO(kevin): Add non-penetrating joint randomization.
 
-    data = mjx_env.init(self.mjx_model)
+    data = mjx_env.make_data(
+        self.mj_model,
+        impl=self.mjx_model.impl.value,
+        nconmax=self._config.nconmax,
+        njmax=self._config.njmax,
+    )
+    data = mjx.forward(self.mjx_model, data)
 
     metrics = {
         "reward/standing": jp.zeros(()),