Fix bug in joint range for MJX.

PiperOrigin-RevId: 575023489
Change-Id: I717333064ed4d6971a841bec60e8afd3019120a7

doc/changelog.rst

@@ -14,7 +14,7 @@ MJX Bug fixes
 ^^^^^^^^^
 
 1. Fix typos and supported fields in the docs (fixes :github:issue:`1105` and :github:issue:`1106`).
-
+2. Fix bug where mixed `jnt_limited` joints are not being constrained correctly.
 
 Version 3.0.0 (October 18, 2023)
 --------------------------------

mjx/mujoco/mjx/_src/constraint.py

@@ -228,13 +228,19 @@ def _instantiate_limit(m: Model, d: Data) -> _Efc:
 
     return jp.stack(js), jp.stack(rs), jp.stack(arefs)
 
+  jnt_range = jp.where(
+      m.jnt_limited[:, None],
+      m.jnt_range,
+      jp.array([-jp.inf, jp.inf]),
+  )
+
   j, r, aref = scan.flat(
       m,
       fn,
       'jjjjjqvv',
       'jjj',
       m.jnt_type,
-      m.jnt_range,
+      jnt_range,
       m.jnt_solref,
       m.jnt_solimp,
       m.jnt_margin,

mjx/mujoco/mjx/_src/constraint_test.py

@@ -24,6 +24,7 @@ from mujoco.mjx._src import constraint
 from mujoco.mjx._src import test_util
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import DisableBit
+from mujoco.mjx._src.types import SolverType
 # pylint: enable=g-importing-member
 import numpy as np
 
@@ -79,6 +80,39 @@ class ConstraintTest(parameterized.TestCase):
           fname,
       )
 
+  _JNT_RANGE = """
+    <mujoco>
+      <worldbody>
+        <body pos="0 0 1">
+          <joint type="slide" axis="1 0 0" range="-1.8 1.8" solreflimit=".08 1"
+           damping="5e-4"/>
+          <geom type="box" size="0.2 0.15 0.1" mass="1"/>
+          <body>
+            <joint axis="0 1 0" damping="2e-6"/>
+            <geom type="capsule" fromto="0 0 0 0 0 1" size="0.045" mass=".1"/>
+          </body>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  def test_jnt_range(self):
+    """Tests that mixed joint ranges are respected."""
+    m = mujoco.MjModel.from_xml_string(self._JNT_RANGE)
+    m.opt.solver = SolverType.CG.value
+    d = mujoco.MjData(m)
+    d.qpos = np.array([2.0, 15.0])
+
+    mx = mjx.device_put(m)
+    dx = mjx.device_put(d)
+    efc = jax.jit(constraint._instantiate_limit)(mx, dx)
+
+    # first joint is outside the joint range
+    np.testing.assert_array_almost_equal(efc.J[0, 0], -1.0)
+
+    # second joint does not hit joint range
+    np.testing.assert_array_almost_equal(efc.aref[1], 0.0)
+
   def test_disable_refsafe(self):
     m = test_util.load_test_file('ant.xml')
 

mjx/mujoco/mjx/test_data/cherry_pendulum.xml

@@ -5,7 +5,7 @@
   </default>
   <worldbody>
     <body>
-      <joint axis="1 0 0" type="hinge"/>
+      <joint axis="1 0 0" type="hinge" range="-45 45"/>
       <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
       <body pos="0 0.5 0">
         <joint axis="1 0 0" type="hinge"/>