Add plugin_state support to MJX API.

Move fr3 controllers states to plugin_state.

Reset plugin_state during device.reset().

PiperOrigin-RevId: 786433231
Change-Id: I7c95cc884acab23fb8c59419b573ee0a6b4cf803

mjx/mujoco/mjx/_src/io.py

@@ -393,6 +393,7 @@ def _make_data_public_fields(m: types.Model) -> Dict[str, Any]:
       'time': (float_,),
       'qvel': (m.nv, float_),
       'act': (m.na, float_),
+      'plugin_state': (m.npluginstate, float_),
       'qacc_warmstart': (m.nv, float_),
       'ctrl': (m.nu, float_),
       'qfrc_applied': (m.nv, float_),

mjx/mujoco/mjx/_src/types.py

@@ -579,6 +579,7 @@ class ModelC(PyTreeNode):
   sensor_plugin: jax.Array
   sensor_intprm: jax.Array
   plugin: jax.Array
+  plugin_stateadr: jax.Array
 
 
 class ModelJAX(PyTreeNode):
@@ -636,6 +637,7 @@ class Model(PyTreeNode):
   ngravcomp: int
   nuserdata: int
   nsensordata: int
+  npluginstate: int
   opt: Option
   stat: Statistic
   qpos0: jax.Array
@@ -1084,6 +1086,7 @@ class Data(PyTreeNode):
   qvel: jax.Array
   act: jax.Array
   qacc_warmstart: jax.Array
+  plugin_state: jax.Array
   # control:
   ctrl: jax.Array
   qfrc_applied: jax.Array