Update environments to utilize contact sensor.

PiperOrigin-RevId: 795391597
Change-Id: I877b54f441b00712f5280dad5a31b38a87fdfff8

Author: DeepMind

CHANGELOG.md

@@ -10,6 +10,7 @@ All notable changes to this project will be documented in this file.
   support MjWarp. You can pass through the implementation via the config
   override
   `registry.load('CartpoleBalance', config_overrides={'impl': 'warp'})`.
+- Update environments to utilize contact sensors and remove `collision.py`.
 
 ## [0.0.5] - 2025-06-23
 

mujoco_playground/_src/collision.py

@@ -112,9 +112,9 @@ 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.make_data(
-        self.mj_model,
+        self._mj_model,
         qpos=qpos,
-        impl=self.mjx_model.impl.value,
+        impl=self._mjx_model.impl.value,
         nconmax=self._config.nconmax,
         njmax=self._config.njmax,
     )
@@ -229,7 +229,7 @@ def __init__(
         _XML_PATH, target_radius, 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:
@@ -249,7 +249,7 @@ 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._config.impl)
 
     target_angle = jax.random.uniform(rng2, minval=-jp.pi, maxval=jp.pi)
     hinge_x = data.xanchor[self._hinge_joint_id, 0]

mujoco_playground/_src/dm_control_suite/finger.py

@@ -23,7 +23,6 @@
 import mujoco
 from mujoco import mjx
 from mujoco_playground._src import mjx_env
-from mujoco_playground._src.collision import geoms_colliding
 from mujoco_playground._src.locomotion.apollo import constants as consts
 import numpy as np
 
@@ -105,6 +104,31 @@ def __init__(
         [self._mj_model.site(name).id for name in consts.FEET_SITES]
     )
 
+    # Contact sensor IDs.
+    self._left_feet_floor_found_sensor = [
+        self._mj_model.sensor(foot_geom + "_floor_found").id
+        for foot_geom in consts.LEFT_FEET_GEOMS
+    ]
+    self._right_feet_floor_found_sensor = [
+        self._mj_model.sensor(foot_geom + "_floor_found").id
+        for foot_geom in consts.RIGHT_FEET_GEOMS
+    ]
+    self._left_hand_left_thigh_found_sensor = self._mj_model.sensor(
+        "collision_l_hand_plate_collision_capsule_body_l_thigh_found"
+    ).id
+    self._right_hand_right_thigh_found_sensor = self._mj_model.sensor(
+        "collision_r_hand_plate_collision_capsule_body_r_thigh_found"
+    ).id
+    self._left_foot_right_foot_found_sensor = self._mj_model.sensor(
+        "collision_l_sole_collision_r_sole_found"
+    ).id
+    self._left_shin_right_shin_found_sensor = self._mj_model.sensor(
+        "collision_capsule_body_l_shin_collision_capsule_body_r_shin_found"
+    ).id
+    self._left_thigh_right_thigh_found_sensor = self._mj_model.sensor(
+        "collision_capsule_body_l_thigh_collision_capsule_body_r_thigh_found"
+    ).id
+
   # Sensor readings.
 
   def get_gravity(self, data: mjx.Data) -> jax.Array:
@@ -144,12 +168,18 @@ def get_gyro(self, data: mjx.Data) -> jax.Array:
   def get_feet_ground_contacts(self, data: mjx.Data) -> jax.Array:
     """Return an array indicating whether each foot is in contact with the ground."""
     left_feet_contact = jp.array([
-        geoms_colliding(data, geom_id, self._floor_geom_id)
-        for geom_id in self._left_feet_geom_id
+        data.sensordata[
+            self._mj_model.sensor_adr[self._mj_model.sensor_adr[sensorid]]
+        ]
+        > 0
+        for sensorid in self._left_feet_floor_found_sensor
     ])
     right_feet_contact = jp.array([
-        geoms_colliding(data, geom_id, self._floor_geom_id)
-        for geom_id in self._right_feet_geom_id
+        data.sensordata[
+            self._mj_model.sensor_adr[self._mj_model.sensor_adr[sensorid]]
+        ]
+        > 0
+        for sensorid in self._right_feet_floor_found_sensor
     ])
     return jp.hstack([jp.any(left_feet_contact), jp.any(right_feet_contact)])
 

mujoco_playground/_src/locomotion/apollo/base.py

@@ -22,7 +22,6 @@
 from mujoco import mjx
 from mujoco.mjx._src import math
 from mujoco_playground._src import gait, mjx_env
-from mujoco_playground._src.collision import geoms_colliding
 from mujoco_playground._src.locomotion.apollo import base
 from mujoco_playground._src.locomotion.apollo import constants as consts
 
@@ -377,27 +376,17 @@ def _cost_action_rate(self, act: jax.Array, last_act: jax.Array) -> jax.Array:
     return jp.sum(jp.square(act - last_act))
 
   def _cost_collision(self, data: mjx.Data) -> jax.Array:
+    adr = self._mj_model.sensor_adr
     # Hand - thigh.
-    c = geoms_colliding(data, self._left_hand_geom_id, self._left_thigh_geom_id)
-    c |= geoms_colliding(
-        data, self._right_hand_geom_id, self._right_thigh_geom_id
-    )
+    c = data.sensordata[adr[self._left_hand_left_thigh_found_sensor]] > 0
+    c |= data.sensordata[adr[self._right_hand_right_thigh_found_sensor]] > 0
     # Foot - foot.
-    c |= geoms_colliding(
-        data, self._left_foot_geom_id, self._right_foot_geom_id
-    )
+    c |= data.sensordata[adr[self._left_foot_right_foot_found_sensor]] > 0
     # Shin - shin.
-    c |= geoms_colliding(
-        data,
-        self._left_shin_geom_id,
-        self._right_shin_geom_id,
-    )
+    c |= data.sensordata[adr[self._left_shin_right_shin_found_sensor]] > 0
     # Thigh - thigh.
-    c |= geoms_colliding(
-        data,
-        self._left_thigh_geom_id,
-        self._right_thigh_geom_id,
-    )
+    c |= data.sensordata[adr[self._left_thigh_right_thigh_found_sensor]] > 0
+
     return jp.any(c)
 
   def _cost_pose(self, qpos: jax.Array, commands: jax.Array) -> jax.Array:

mujoco_playground/_src/locomotion/apollo/joystick.py

@@ -23,6 +23,16 @@
     <geom name="floor" size="0 0 0.01" type="plane" material="groundplane" priority="1" friction="0.6" condim="3"/>
   </worldbody>
 
+  <sensor>
+    <contact name="collision_l_sole_floor_found" geom1="collision_l_sole" geom2="floor" reduce="mindist" num="1" data="found"/>
+    <contact name="collision_r_sole_floor_found" geom1="collision_r_sole" geom2="floor" reduce="mindist" num="1" data="found"/>
+    <contact name="collision_l_hand_plate_collision_capsule_body_l_thigh_found" geom1="collision_l_hand_plate" geom2="collision_capsule_body_l_thigh" reduce="mindist" num="1" data="found"/>
+    <contact name="collision_r_hand_plate_collision_capsule_body_r_thigh_found" geom1="collision_r_hand_plate" geom2="collision_capsule_body_r_thigh" reduce="mindist" num="1" data="found"/>
+    <contact name="collision_l_sole_collision_r_sole_found" geom1="collision_l_sole" geom2="collision_r_sole" reduce="mindist" num="1" data="found"/>
+    <contact name="collision_capsule_body_l_shin_collision_capsule_body_r_shin_found" geom1="collision_capsule_body_l_shin" geom2="collision_capsule_body_r_shin" reduce="mindist" num="1" data="found"/>
+    <contact name="collision_capsule_body_l_thigh_collision_capsule_body_r_thigh_found" geom1="collision_capsule_body_l_thigh" geom2="collision_capsule_body_r_thigh" reduce="mindist" num="1" data="found"/>
+  </sensor>
+
   <keyframe>
     <key name="stand"
       qpos="

mujoco_playground/_src/locomotion/apollo/xmls/scene_mjx_feetonly_flat_terrain.xml

@@ -24,7 +24,6 @@
 from mujoco import mjx
 import numpy as np
 
-from mujoco_playground._src import collision
 from mujoco_playground._src import mjx_env
 
 _FEET_SITES = [
@@ -119,9 +118,30 @@ def __init__(
     super().__init__(config, config_overrides)
     xml_path = mjx_env.MENAGERIE_PATH / "google_barkour_vb" / "scene_mjx.xml"
     self._xml_path = xml_path.as_posix()
-    xml = epath.Path(xml_path).read_text()
     self._model_assets = get_assets()
-    mj_model = mujoco.MjModel.from_xml_string(xml, assets=self._model_assets)
+    mj_spec = mujoco.MjSpec.from_file(
+        xml_path.as_posix(), assets=self._model_assets
+    )
+    # add contact sensors
+    feet_floor_found_sensor = []
+    for geom in _FEET_GEOMS:
+      name = f"{geom}_floor_found"
+      mj_spec.add_sensor(
+          name=name,
+          type=mujoco.mjtSensor.mjSENS_CONTACT,
+          objtype=mujoco.mjtObj.mjOBJ_GEOM,
+          objname=geom,
+          reftype=mujoco.mjtObj.mjOBJ_GEOM,
+          refname="floor",
+          intprm=[1, 1, 1],  # data=found, reduce=mindist
+          datatype=mujoco.mjtDataType.mjDATATYPE_REAL,
+          needstage=mujoco.mjtStage.mjSTAGE_ACC,
+          dim=1,
+      )
+      feet_floor_found_sensor.append(name)
+    self._feet_floor_found_sensor = feet_floor_found_sensor
+    # compile spec
+    mj_model = mj_spec.compile()
     mj_model.vis.global_.offwidth = 3840
     mj_model.vis.global_.offheight = 2160
     mj_model.dof_damping[6:] = 0.5239
@@ -239,8 +259,11 @@ def step(self, state: mjx_env.State, action: jax.Array) -> mjx_env.State:
     torso_z = data.xpos[self._torso_body_id, -1]
 
     contact = jp.array([
-        collision.geoms_colliding(data, geom_id, self._floor_geom_id)
-        for geom_id in self._feet_geom_id
+        data.sensordata[
+            self._mj_model.sensor_adr[self._mj_model.sensor(sensor).id]
+        ]
+        > 0
+        for sensor in self._feet_floor_found_sensor
     ])
     contact_filt = contact | state.info["last_contact"]
     first_contact = (state.info["feet_air_time"] > 0.0) * contact_filt

mujoco_playground/_src/locomotion/barkour/joystick.py

@@ -25,7 +25,6 @@
 
 from mujoco_playground._src import gait
 from mujoco_playground._src import mjx_env
-from mujoco_playground._src.collision import geoms_colliding
 from mujoco_playground._src.locomotion.berkeley_humanoid import base as berkeley_humanoid_base
 from mujoco_playground._src.locomotion.berkeley_humanoid import berkeley_humanoid_constants as consts
 
@@ -184,6 +183,12 @@ def _post_init(self) -> None:
     qpos_noise_scale[faa_ids] = self._config.noise_config.scales.faa_pos
     self._qpos_noise_scale = jp.array(qpos_noise_scale)
 
+    # Contact sensor IDs.
+    self._feet_floor_found_sensor = [
+        self._mj_model.sensor(f"{geom}_floor_found").id
+        for geom in consts.FEET_GEOMS
+    ]
+
   def reset(self, rng: jax.Array) -> mjx_env.State:
     qpos = self._init_q
     qvel = jp.zeros(self.mjx_model.nv)
@@ -264,9 +269,10 @@ def reset(self, rng: jax.Array) -> mjx_env.State:
     metrics["swing_peak"] = jp.zeros(())
 
     contact = jp.array([
-        geoms_colliding(data, geom_id, self._floor_geom_id)
-        for geom_id in self._feet_geom_id
+        data.sensordata[self._mj_model.sensor_adr[sensor_id]] > 0
+        for sensor_id in self._feet_floor_found_sensor
     ])
+
     obs = self._get_obs(data, info, contact)
     reward, done = jp.zeros(2)
     return mjx_env.State(data, obs, reward, done, metrics, info)
@@ -299,8 +305,8 @@ def step(self, state: mjx_env.State, action: jax.Array) -> mjx_env.State:
     state.info["motor_targets"] = motor_targets
 
     contact = jp.array([
-        geoms_colliding(data, geom_id, self._floor_geom_id)
-        for geom_id in self._feet_geom_id
+        data.sensordata[self._mj_model.sensor_adr[sensor_id]] > 0
+        for sensor_id in self._feet_floor_found_sensor
     ])
     contact_filt = contact | state.info["last_contact"]
     first_contact = (state.info["feet_air_time"] > 0.0) * contact_filt

mujoco_playground/_src/locomotion/berkeley_humanoid/joystick.py

@@ -24,6 +24,8 @@
       friction="0.6" condim="3"/>
   </worldbody>
 
+  <include file="sensor.xml"/>
+
   <keyframe>
     <key name="home"
       qpos="

mujoco_playground/_src/locomotion/berkeley_humanoid/xmls/scene_mjx_feetonly_flat_terrain.xml

@@ -23,6 +23,8 @@
       friction="1.0"/>
   </worldbody>
 
+  <include file="sensor.xml"/>
+
   <keyframe>
     <key name="home"
       qpos="