Add gravcomp.

PiperOrigin-RevId: 627791303
Change-Id: I87c9254967eb492d0c38153ff41d7458ffb69136

mjx/mujoco/mjx/_src/forward.py

@@ -174,16 +174,17 @@ def fwd_actuation(m: Model, d: Data) -> Data:
 
   qfrc_actuator = d.actuator_moment.T @ force
 
+  if m.ngravcomp:
+    # actuator-level gravity compensation, skip if added as passive force
+    qfrc_actuator += d.qfrc_gravcomp * m.jnt_actgravcomp[m.dof_jntid]
+
   # clamp qfrc_actuator
   actfrcrange = jp.where(
       m.jnt_actfrclimited[:, None],
       m.jnt_actfrcrange,
       jp.array([-jp.inf, jp.inf]),
   )
-  ids = sum(
-      ([i] * JointType(j).dof_width() for i, j in enumerate(m.jnt_type)), []
-  )
-  actfrcrange = jp.take(actfrcrange, jp.array(ids), axis=0)
+  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)