Set actuator_force in MJX. Fixes #2068.

PiperOrigin-RevId: 676867005
Change-Id: I80514a050ff40cb6d43a95b74c4a36c4991db588

doc/changelog.rst

@@ -13,6 +13,11 @@ General
   <https://github.com/google-deepmind/mujoco/blob/main/src/model/trampoline.xml>`__ for examples of flex objects
   that previously required these plugins.
 
+Bug fixes
+^^^^^^^^^
+- Fixed a bug where ``actuator_force`` was not set in MJX (:github:issue:`2068`).
+
+
 Version 3.2.3 (Sep 16, 2024)
 ----------------------------
 

mjx/mujoco/mjx/_src/forward.py

@@ -192,7 +192,9 @@ def fwd_actuation(m: Model, d: Data) -> Data:
   actfrcrange = actfrcrange[m.dof_jntid]
   qfrc_actuator = jp.clip(qfrc_actuator, actfrcrange[:, 0], actfrcrange[:, 1])
 
-  d = d.replace(act_dot=act_dot, qfrc_actuator=qfrc_actuator)
+  d = d.replace(
+      act_dot=act_dot, qfrc_actuator=qfrc_actuator, actuator_force=force
+  )
   return d
 
 

mjx/mujoco/mjx/_src/forward_test.py

@@ -52,9 +52,15 @@ class ForwardTest(absltest.TestCase):
     mx = mjx.put_model(m)
 
     # fwd_actuation
-    dx = jax.jit(mjx.fwd_actuation)(mx, mjx.put_data(m, d))
+    dx = mjx.put_data(m, d).replace(
+        act_dot=np.zeros_like(d.act_dot),
+        qfrc_actuator=np.zeros_like(d.qfrc_actuator),
+        actuator_force=np.zeros_like(d.actuator_force),
+    )
+    dx = jax.jit(mjx.fwd_actuation)(mx, dx)
     _assert_attr_eq(d, dx, 'act_dot')
     _assert_attr_eq(d, dx, 'qfrc_actuator')
+    _assert_attr_eq(d, dx, 'actuator_force')
 
     # fwd_accleration (fwd_position and fwd_velocity already tested elsewhere)
     dx = jax.jit(mjx.fwd_acceleration)(mx, mjx.put_data(m, d))