Fixes #1668

PiperOrigin-RevId: 634483090 Change-Id: I30b55bbf419e57a7b2966dfca4da6bb57d5879ad
google-deepmind · May 16, 2024 · 86e6fc6 · 86e6fc6
1 parent e885a3d
commit 86e6fc6
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 6 deletions.
diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py
@@ -73,8 +73,8 @@ def _kbi(
   k = 1 / (dmax * dmax * timeconst * timeconst * dampratio * dampratio)
   b = 2 / (dmax * timeconst)
   # TODO(robotics-simulation): check various solparam settings in model gen test
-  k = jp.where(dampratio <= 0, -dampratio / (dmax * dmax), k)
-  b = jp.where(timeconst <= 0, -timeconst / dmax, b)
+  k = jp.where(dampratio <= 0, -solref[0] / (dmax * dmax), k)
+  b = jp.where(timeconst <= 0, -solref[1] / dmax, b)
 
   imp_x = jp.abs(pos) / width
   imp_a = (1.0 / jp.power(mid, power - 1)) * jp.power(imp_x, power)

diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py
@@ -104,21 +104,21 @@ def test_disable_equality(self):
     ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m))
     self.assertEqual(ne, 0)
     self.assertEqual(nf, 0)
-    self.assertEqual(nl, 2)
+    self.assertEqual(nl, 3)
     self.assertEqual(nc, 148)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 150)  # only joint range, contact
+    self.assertEqual(dx.efc_J.shape[0], 151)  # only joint range, contact
 
   def test_disable_contact(self):
     m = test_util.load_test_file('constraints.xml')
     m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONTACT
     ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m))
     self.assertEqual(ne, 10)
     self.assertEqual(nf, 0)
-    self.assertEqual(nl, 2)
+    self.assertEqual(nl, 3)
     self.assertEqual(nc, 0)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 12)  # only joint range, limit
+    self.assertEqual(dx.efc_J.shape[0], 13)  # only joint range, limit
 
 
 if __name__ == '__main__':

diff --git a/mjx/mujoco/mjx/test_data/constraints.xml b/mjx/mujoco/mjx/test_data/constraints.xml
@@ -38,6 +38,11 @@
       <geom class="box"/>
     </body>
 
+    <body name="beam5" pos="-2 0 0" euler="45 45 0">
+      <joint name="joint5" type="ball" range="0 45" solreflimit="-100 -1"/> <!-- test negative solref -->
+      <geom class="box" contype="0" conaffinity="0"/>
+    </body>
+
     <body name="box_condim1" pos="4 0 0">
       <freejoint/>
       <geom class="box" condim="1"/>