Add sensors to MJX. Implements jointpos and actuatorpos as examples.

PiperOrigin-RevId: 660891329
Change-Id: Ie6c282a397de2f52843775ec747495e9be46d307

mjx/mujoco/mjx/__init__.py

@@ -32,6 +32,9 @@ from mujoco.mjx._src.io import put_data
 from mujoco.mjx._src.io import put_model
 from mujoco.mjx._src.passive import passive
 from mujoco.mjx._src.ray import ray
+from mujoco.mjx._src.sensor import sensor_pos
+from mujoco.mjx._src.sensor import sensor_vel
+from mujoco.mjx._src.sensor import sensor_acc
 from mujoco.mjx._src.smooth import camlight
 from mujoco.mjx._src.smooth import com_pos
 from mujoco.mjx._src.smooth import com_vel

mjx/mujoco/mjx/_src/forward.py

@@ -25,6 +25,7 @@ from mujoco.mjx._src import constraint
 from mujoco.mjx._src import math
 from mujoco.mjx._src import passive
 from mujoco.mjx._src import scan
+from mujoco.mjx._src import sensor
 from mujoco.mjx._src import smooth
 from mujoco.mjx._src import solver
 from mujoco.mjx._src import support
@@ -351,9 +352,12 @@ def rungekutta4(m: Model, d: Data) -> Data:
 def forward(m: Model, d: Data) -> Data:
   """Forward dynamics."""
   d = fwd_position(m, d)
+  d = sensor.sensor_pos(m, d)
   d = fwd_velocity(m, d)
+  d = sensor.sensor_vel(m, d)
   d = fwd_actuation(m, d)
   d = fwd_acceleration(m, d)
+  d = sensor.sensor_acc(m, d)
 
   if d.efc_J.size == 0:
     d = d.replace(qacc=d.qacc_smooth)

mjx/mujoco/mjx/_src/io.py

@@ -101,6 +101,8 @@ def put_model(
       (m.actuator_gaintype, types.GainType, mujoco.mjtGain),
       (m.actuator_trntype, types.TrnType, mujoco.mjtTrn),
       (m.eq_type, types.EqType, mujoco.mjtEq),
+      # TODO(taylorhowell): causes Menagerie test to fail
+      # (m.sensor_type, types.SensorType, mujoco.mjtSensor),
       (m.wrap_type, types.WrapType, mujoco.mjtWrap),
   ):
     missing = set(enum_field) - set(enum_type)

mjx/mujoco/mjx/_src/sensor.py

@@ -0,0 +1,54 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Sensor functions."""
+
+import jax
+# pylint: disable=g-importing-member
+from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import SensorType
+from typing import Tuple
+# pylint: enable=g-importing-member
+import numpy as np
+
+
+def sensor_pos(m: Model, d: Data) -> Data:
+  """Compute position-dependent sensors values."""
+
+  sensordata = d.sensordata
+  if np.isin(SensorType.JOINTPOS, m.sensor_type):
+    # jointpos
+    i = m.sensor_type == SensorType.JOINTPOS
+    objid = m.sensor_objid[i]
+    adr = m.sensor_adr[i]
+    sensordata = sensordata.at[adr].set(d.qpos[m.jnt_qposadr[objid]])
+  if np.isin(SensorType.ACTUATORPOS, m.sensor_type):
+    # actuatorpos
+    i = m.sensor_type == SensorType.ACTUATORPOS
+    objid = m.sensor_objid[i]
+    adr = m.sensor_adr[i]
+    sensordata = sensordata.at[adr].set(d.actuator_length[objid])
+
+  return d.replace(sensordata=sensordata)
+
+
+def sensor_vel(m: Model, d: Data) -> Data:
+  """Compute velocity-dependent sensors values."""
+  return d
+
+
+def sensor_acc(m: Model, d: Data) -> Data:
+  """Compute acceleration/force-dependent sensors values."""
+  return d

mjx/mujoco/mjx/_src/sensor_test.py

@@ -0,0 +1,67 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Tests for sensor functions."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+import mujoco
+from mujoco import mjx
+from mujoco.mjx._src import test_util
+import numpy as np
+
+# tolerance for difference between MuJoCo and MJX smooth calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+def _assert_attr_eq(a, b, attr):
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
+
+
+class SensorTest(parameterized.TestCase):
+
+  @parameterized.parameters('no_sensor.xml', 'sensor.xml')
+  def test_sensor(self, filename):
+    """Tests MJX sensor functions match MuJoCo sensor functions."""
+    m = test_util.load_test_file(filename)
+    d = mujoco.MjData(m)
+    # give the system a little kick to ensure we have non-identity rotations
+    d.qvel = np.random.random(m.nv)
+    # apply control for activation dynamics
+    d.ctrl = np.clip(
+        np.random.random(m.nu),
+        m.actuator_ctrlrange[:, 0],
+        m.actuator_ctrlrange[:, 1],
+    )
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+
+    mujoco.mj_forward(m, d)
+    dx = jax.jit(mjx.forward)(mx, dx)
+
+    # sensor values
+    _assert_eq(d.sensordata, dx.sensordata, 'sensordata')
+
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/_src/types.py

@@ -100,7 +100,6 @@ class GeomType(enum.IntEnum):
     MESH: mesh
     SDF: signed distance field
   """
-
   PLANE = mujoco.mjtGeom.mjGEOM_PLANE
   HFIELD = mujoco.mjtGeom.mjGEOM_HFIELD
   SPHERE = mujoco.mjtGeom.mjGEOM_SPHERE
@@ -122,7 +121,6 @@ class ConvexMesh(PyTreeNode):
     edge: edge indexes for all edges in the convex mesh
     edge_face_normal: indexes for face normals adjacent to edges in `edge`
   """
-
   vert: jax.Array
   face: jax.Array
   face_normal: jax.Array
@@ -273,7 +271,6 @@ class CamLightType(enum.IntEnum):
     TARGETBODY: pos fixed in body, rot tracks target body
     TARGETBODYCOM: pos fixed in body, rot tracks target subtree com
   """
-
   FIXED = mujoco.mjtCamLight.mjCAMLIGHT_FIXED
   TRACK = mujoco.mjtCamLight.mjCAMLIGHT_TRACK
   TRACKCOM = mujoco.mjtCamLight.mjCAMLIGHT_TRACKCOM
@@ -281,6 +278,17 @@ class CamLightType(enum.IntEnum):
   TARGETBODYCOM = mujoco.mjtCamLight.mjCAMLIGHT_TARGETBODYCOM
 
 
+class SensorType(enum.IntEnum):
+  """Type of sensor.
+
+  Members:
+    JOINTPOS: joint position
+    ACTUATORPOS: actuator position
+  """
+  JOINTPOS = mujoco.mjtSensor.mjSENS_JOINTPOS
+  ACTUATORPOS = mujoco.mjtSensor.mjSENS_ACTUATORPOS
+
+
 class Option(PyTreeNode):
   """Physics options.
 

mjx/mujoco/mjx/test_data/no_sensor.xml

@@ -0,0 +1,9 @@
+<!-- For validating model with no sensor -->
+<mujoco model="no_sensor">
+  <worldbody>
+    <body>
+      <joint type="hinge"/>
+      <geom size="1"/>
+    </body>
+  </worldbody>
+</mujoco>

mjx/mujoco/mjx/test_data/sensor.xml

@@ -0,0 +1,29 @@
+<!-- For validating sensors:
+
+* position-dependent:
+-jointpos
+-actuatorpos
+* velocity-dependent:
+* acceleration/force-dependent:
+-->
+<mujoco model="sensor">
+  <worldbody>
+    <!-- body 0 -->
+    <body name="body0" pos="1 2 3">
+      <joint name="hinge0" type="hinge" axis="1 0 0"/>
+      <geom size="1"/>
+    </body>
+  </worldbody>
+
+  <actuator>
+    <motor name="motor0" joint="hinge0" ctrlrange="-1 1" ctrllimited="true"/>
+  </actuator>
+
+  <sensor>
+    <!-- position-dependent sensors -->
+    <jointpos name="jointpos0" joint="hinge0"/>
+    <actuatorpos name="actuatorpos0" actuator="motor0"/>
+    <!-- velocity-dependent sensors -->
+    <!-- acceleration/force-dependent sensors -->
+  </sensor>
+</mujoco>