Fix bug in joint range for MJX.

PiperOrigin-RevId: 575023489 Change-Id: I717333064ed4d6971a841bec60e8afd3019120a7
2023-10-19 15:39:15 -07:00
parent dc8bac2f4d
commit dd841c9283
4 changed files with 43 additions and 3 deletions
@@ -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)
 --------------------------------
@@ -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,
@@ -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')

@@ -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"/>