Updates to Playground for warp.

PiperOrigin-RevId: 797037846
Change-Id: Ibbfdc590c78b1a3dd32e6e09e5575256174eec61

Author: btaba

mujoco_playground/_src/dm_control_suite/finger.py

@@ -50,7 +50,7 @@ def default_config() -> config_dict.ConfigDict:
       vision=False,
       impl="jax",
       nconmax=25_000,
-      njmax=25,
+      njmax=50,
   )
 
 
@@ -114,10 +114,11 @@ def reset(self, rng: jax.Array) -> mjx_env.State:
     data = mjx_env.make_data(
         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,
     )
+    data = mjx.forward(self.mjx_model, data)
 
     metrics = {}
     info = {"rng": rng}
@@ -249,7 +250,14 @@ 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, qpos=qpos, impl=self._config.impl)
+    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)
 
     target_angle = jax.random.uniform(rng2, minval=-jp.pi, maxval=jp.pi)
     hinge_x = data.xanchor[self._hinge_joint_id, 0]

mujoco_playground/_src/locomotion/t1/joystick.py

@@ -640,12 +640,12 @@ def _reward_alive(self) -> jax.Array:
     return jp.array(1.0)
 
   def _cost_collision(self, data: mjx.Data) -> jax.Array:
-    return jp.array([
+    return jp.array(
         data.sensordata[
             self._mj_model.sensor_adr[self._left_foot_right_foot_found_sensor]
         ]
         > 0
-    ])
+    )
 
   # Pose-related rewards.
 

mujoco_playground/_src/locomotion/t1/xmls/sensor.xml

@@ -1,13 +1,13 @@
 <mujoco>
   <sensor>
-    <contact name="left_foot_1_floor_found" geom1="left_foot_1" geom2="floor" num="1" data="found"/>
-    <contact name="left_foot_2_floor_found" geom1="left_foot_2" geom2="floor" num="1" data="found"/>
-    <contact name="left_foot_3_floor_found" geom1="left_foot_3" geom2="floor" num="1" data="found"/>
-    <contact name="left_foot_4_floor_found" geom1="left_foot_4" geom2="floor" num="1" data="found"/>
-    <contact name="right_foot_1_floor_found" geom1="right_foot_1" geom2="floor" num="1" data="found"/>
-    <contact name="right_foot_2_floor_found" geom1="right_foot_2" geom2="floor" num="1" data="found"/>
-    <contact name="right_foot_3_floor_found" geom1="right_foot_3" geom2="floor" num="1" data="found"/>
-    <contact name="right_foot_4_floor_found" geom1="right_foot_4" geom2="floor" num="1" data="found"/>
-    <contact name="left_foot_right_foot_found" geom1="left_foot" geom2="right_foot" num="1" data="found"/>
+    <contact name="left_foot_1_floor_found" geom1="left_foot_1" geom2="floor" num="1" reduce="mindist" data="found"/>
+    <contact name="left_foot_2_floor_found" geom1="left_foot_2" geom2="floor" num="1" reduce="mindist" data="found"/>
+    <contact name="left_foot_3_floor_found" geom1="left_foot_3" geom2="floor" num="1" reduce="mindist" data="found"/>
+    <contact name="left_foot_4_floor_found" geom1="left_foot_4" geom2="floor" num="1" reduce="mindist" data="found"/>
+    <contact name="right_foot_1_floor_found" geom1="right_foot_1" geom2="floor" num="1" reduce="mindist" data="found"/>
+    <contact name="right_foot_2_floor_found" geom1="right_foot_2" geom2="floor" num="1" reduce="mindist" data="found"/>
+    <contact name="right_foot_3_floor_found" geom1="right_foot_3" geom2="floor" num="1" reduce="mindist" data="found"/>
+    <contact name="right_foot_4_floor_found" geom1="right_foot_4" geom2="floor" num="1" reduce="mindist" data="found"/>
+    <contact name="left_foot_right_foot_found" geom1="left_foot" geom2="right_foot" num="1" reduce="mindist" data="found"/>
   </sensor>
 </mujoco>