diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py
index e0ad5cbe..2b442d77 100644
--- a/mjx/mujoco/mjx/__init__.py
+++ b/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
diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py
index b161593c..a90ef854 100644
--- a/mjx/mujoco/mjx/_src/forward.py
+++ b/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)
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index ff6cd29e..202cfd9d 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/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)
diff --git a/mjx/mujoco/mjx/_src/sensor.py b/mjx/mujoco/mjx/_src/sensor.py
new file mode 100644
index 00000000..a59e9c58
--- /dev/null
+++ b/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
diff --git a/mjx/mujoco/mjx/_src/sensor_test.py b/mjx/mujoco/mjx/_src/sensor_test.py
new file mode 100644
index 00000000..d30e928e
--- /dev/null
+++ b/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()
diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index 0500eefe..d5c3bdcc 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/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.
diff --git a/mjx/mujoco/mjx/test_data/no_sensor.xml b/mjx/mujoco/mjx/test_data/no_sensor.xml
new file mode 100644
index 00000000..5ccc660b
--- /dev/null
+++ b/mjx/mujoco/mjx/test_data/no_sensor.xml
@@ -0,0 +1,9 @@
+
+
+
+
+
+
+
+
+
diff --git a/mjx/mujoco/mjx/test_data/sensor.xml b/mjx/mujoco/mjx/test_data/sensor.xml
new file mode 100644
index 00000000..810abd1d
--- /dev/null
+++ b/mjx/mujoco/mjx/test_data/sensor.xml
@@ -0,0 +1,29 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+