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

PiperOrigin-RevId: 660891329
Change-Id: Ie6c282a397de2f52843775ec747495e9be46d307

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)