Add Booster T1 robot.

This PR adds code to train a joystick policy for the T1 humanoid as well as the ONNX file that was deployed on real hardware, see https://www.youtube.com/shorts/GY81u3aqGEw

PiperOrigin-RevId: 726173305
Change-Id: I1acad2a87ebf80f1b9ec62e8db341c3afa709301

mujoco_playground/_src/locomotion/__init__.py

@@ -37,6 +37,8 @@
 from mujoco_playground._src.locomotion.spot import getup as spot_getup
 from mujoco_playground._src.locomotion.spot import joystick as spot_joystick
 from mujoco_playground._src.locomotion.spot import joystick_gait_tracking as spot_joystick_gait_tracking
+from mujoco_playground._src.locomotion.t1 import joystick as t1_joystick
+from mujoco_playground._src.locomotion.t1 import randomize as t1_randomize
 
 _envs = {
     "BarkourJoystick": barkour_joystick.Joystick,
@@ -71,6 +73,12 @@
     "SpotJoystickGaitTracking": (
         spot_joystick_gait_tracking.JoystickGaitTracking
     ),
+    "T1JoystickFlatTerrain": functools.partial(
+        t1_joystick.Joystick, task="flat_terrain"
+    ),
+    "T1JoystickRoughTerrain": functools.partial(
+        t1_joystick.Joystick, task="rough_terrain"
+    ),
 }
 
 _cfgs = {
@@ -94,6 +102,8 @@
     "SpotFlatTerrainJoystick": spot_joystick.default_config,
     "SpotGetup": spot_getup.default_config,
     "SpotJoystickGaitTracking": spot_joystick_gait_tracking.default_config,
+    "T1JoystickFlatTerrain": t1_joystick.default_config,
+    "T1JoystickRoughTerrain": t1_joystick.default_config,
 }
 
 _randomizer = {
@@ -110,6 +120,8 @@
     "Go1Getup": go1_randomize.domain_randomize,
     "Go1Handstand": go1_randomize.domain_randomize,
     "Go1Footstand": go1_randomize.domain_randomize,
+    "T1JoystickFlatTerrain": t1_randomize.domain_randomize,
+    "T1JoystickRoughTerrain": t1_randomize.domain_randomize,
 }
 
 

mujoco_playground/_src/locomotion/t1/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2025 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

mujoco_playground/_src/locomotion/t1/base.py

@@ -0,0 +1,113 @@
+# Copyright 2025 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Base classes for T1."""
+
+from typing import Any, Dict, Optional, Union
+
+from etils import epath
+import jax
+from ml_collections import config_dict
+import mujoco
+from mujoco import mjx
+
+from mujoco_playground._src import mjx_env
+from mujoco_playground._src.locomotion.t1 import t1_constants as consts
+
+
+def get_assets() -> Dict[str, bytes]:
+  assets = {}
+  mjx_env.update_assets(assets, consts.ROOT_PATH / "xmls", "*.xml")
+  mjx_env.update_assets(assets, consts.ROOT_PATH / "xmls" / "assets")
+  path = mjx_env.MENAGERIE_PATH / "booster_t1"
+  mjx_env.update_assets(assets, path, "*.xml")
+  mjx_env.update_assets(assets, path / "assets")
+  return assets
+
+
+class T1Env(mjx_env.MjxEnv):
+  """Base class for T1 environments."""
+
+  def __init__(
+      self,
+      xml_path: str,
+      config: config_dict.ConfigDict,
+      config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None,
+  ) -> None:
+    super().__init__(config, config_overrides)
+
+    self._mj_model = mujoco.MjModel.from_xml_string(
+        epath.Path(xml_path).read_text(), assets=get_assets()
+    )
+    self._mj_model.opt.timestep = self.sim_dt
+
+    self._mj_model.vis.global_.offwidth = 3840
+    self._mj_model.vis.global_.offheight = 2160
+
+    self._mjx_model = mjx.put_model(self._mj_model)
+    self._xml_path = xml_path
+
+  # Sensor readings.
+
+  def get_gravity(self, data: mjx.Data) -> jax.Array:
+    """Return the gravity vector in the world frame."""
+    return mjx_env.get_sensor_data(
+        self.mj_model, data, f"{consts.GRAVITY_SENSOR}"
+    )
+
+  def get_global_linvel(self, data: mjx.Data) -> jax.Array:
+    """Return the linear velocity of the robot in the world frame."""
+    return mjx_env.get_sensor_data(
+        self.mj_model, data, f"{consts.GLOBAL_LINVEL_SENSOR}"
+    )
+
+  def get_global_angvel(self, data: mjx.Data) -> jax.Array:
+    """Return the angular velocity of the robot in the world frame."""
+    return mjx_env.get_sensor_data(
+        self.mj_model, data, f"{consts.GLOBAL_ANGVEL_SENSOR}"
+    )
+
+  def get_local_linvel(self, data: mjx.Data) -> jax.Array:
+    """Return the linear velocity of the robot in the local frame."""
+    return mjx_env.get_sensor_data(
+        self.mj_model, data, f"{consts.LOCAL_LINVEL_SENSOR}"
+    )
+
+  def get_accelerometer(self, data: mjx.Data) -> jax.Array:
+    """Return the accelerometer readings in the local frame."""
+    return mjx_env.get_sensor_data(
+        self.mj_model, data, f"{consts.ACCELEROMETER_SENSOR}"
+    )
+
+  def get_gyro(self, data: mjx.Data) -> jax.Array:
+    """Return the gyroscope readings in the local frame."""
+    return mjx_env.get_sensor_data(self.mj_model, data, f"{consts.GYRO_SENSOR}")
+
+  # Accessors.
+
+  @property
+  def xml_path(self) -> str:
+    return self._xml_path
+
+  @property
+  def action_size(self) -> int:
+    return self._mjx_model.nu
+
+  @property
+  def mj_model(self) -> mujoco.MjModel:
+    return self._mj_model
+
+  @property
+  def mjx_model(self) -> mjx.Model:
+    return self._mjx_model