Add actuator and sensor delays. Fixes #1004

PiperOrigin-RevId: 866478839
Change-Id: Id21a6da0f98454c8fa39ea5af8a5e213d6eae497

mjx/mujoco/mjx/_src/io.py

@@ -571,6 +571,7 @@ def _make_data_public_fields(m: types.Model) -> Dict[str, Any]:
       'time': (float_,),
       'qvel': (m.nv, float_),
       'act': (m.na, float_),
+      'history': (m.nhistory, float_),
       'plugin_state': (m.npluginstate, float_),
       'qacc_warmstart': (m.nv, float_),
       'ctrl': (m.nu, float_),
@@ -875,7 +876,7 @@ def _make_data_warp(
 
   fields = _make_data_public_fields(m)
   for k in fields:
-    if k in {'userdata', 'plugin_state'}:
+    if k in {'userdata', 'plugin_state', 'history'}:
       continue
     if not hasattr(dw, k):
       raise ValueError(f'Public data field {k} not found in Warp data.')
@@ -1729,6 +1730,7 @@ _STATE_MAP = {
     mujoco.mjtState.mjSTATE_QPOS: 'qpos',
     mujoco.mjtState.mjSTATE_QVEL: 'qvel',
     mujoco.mjtState.mjSTATE_ACT: 'act',
+    mujoco.mjtState.mjSTATE_HISTORY: 'history',
     mujoco.mjtState.mjSTATE_WARMSTART: 'qacc_warmstart',
     mujoco.mjtState.mjSTATE_CTRL: 'ctrl',
     mujoco.mjtState.mjSTATE_QFRC_APPLIED: 'qfrc_applied',
@@ -1752,6 +1754,7 @@ def _state_elem_size(m: types.Model, state_enum: mujoco.mjtState) -> int:
       'qpos',
       'qvel',
       'act',
+      'history',
       'qacc_warmstart',
       'ctrl',
       'qfrc_applied',
@@ -1767,6 +1770,7 @@ def _state_elem_size(m: types.Model, state_enum: mujoco.mjtState) -> int:
             'qpos': 'nq',
             'qvel': 'nv',
             'act': 'na',
+            'history': 'nhistory',
             'qacc_warmstart': 'nv',
             'ctrl': 'nu',
             'qfrc_applied': 'nv',

mjx/mujoco/mjx/_src/types.py

@@ -641,7 +641,14 @@ class ModelC(PyTreeNode):
   tendon_treenum: jax.Array
   tendon_treeid: jax.Array
   actuator_plugin: jax.Array
+  actuator_history: jax.Array
+  actuator_historyadr: jax.Array
+  actuator_delay: jax.Array
   sensor_plugin: jax.Array
+  sensor_history: jax.Array
+  sensor_historyadr: jax.Array
+  sensor_delay: jax.Array
+  sensor_interval: jax.Array
   plugin: jax.Array
   plugin_stateadr: jax.Array
   B_rownnz: jax.Array  # pylint:disable=invalid-name
@@ -719,6 +726,7 @@ class Model(PyTreeNode):
     nuserdata: number of elements in userdata
     nsensordata: number of elements in sensor data vector
     npluginstate: number of plugin state values
+    nhistory: number of history buffer elements
     opt: physics options
     stat: model statistics
     qpos0: qpos values at default pose
@@ -768,6 +776,7 @@ class Model(PyTreeNode):
   nuserdata: int
   nsensordata: int
   npluginstate: int
+  nhistory: int
   opt: Option
   stat: Union[Statistic, StatisticWarp]
   qpos0: jax.Array
@@ -1206,6 +1215,7 @@ class Data(PyTreeNode):
     qpos: position
     qvel: velocity
     act: actuator activation
+    history: actuator history buffer
     qacc_warmstart: warm start for solver
     plugin_state: plugin state values
     ctrl: control input
@@ -1255,6 +1265,7 @@ class Data(PyTreeNode):
   qpos: jax.Array
   qvel: jax.Array
   act: jax.Array
+  history: jax.Array
   qacc_warmstart: jax.Array
   plugin_state: jax.Array
   # control: