Adds support for mjOption.impratio to MJX.

PiperOrigin-RevId: 607456034
Change-Id: I7f6549a4332fd2bc0331d76a8906b971ce28dbcf

doc/changelog.rst

@@ -18,6 +18,7 @@ MJX
    ``mjx.get_data_into`` function call.
 3. Fixed a bug in ``mjx.euler`` that applied incorrect damping when using dense mass matrices.
 4. Fixed a bug in ``mjx.solve`` that was causing slow convergence when using ``mjSOL_NEWTON`` in :ref:`mjtSolver`.
+5. Added support for :ref:`mjOption.impratio<mjOption>` to ``mjx.Model``.
 
 Version 3.1.2 (February 05, 2024)
 -----------------------------------

mjx/mujoco/mjx/_src/constraint.py

@@ -297,7 +297,7 @@ def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]:
     for diff_tan, friction in zip(diff_con[1:], c.friction[:2]):
       for f in (friction, -friction):
         js.append(diff_con[0] + diff_tan * f)
-        invweights.append((t + f * f * t) * 2 * f * f)
+        invweights.append((t + f * f * t) * 2 * f * f / m.opt.impratio)
 
     active = dist < 0
     j, invweight = jp.stack(js) * active, jp.stack(invweights)

mjx/mujoco/mjx/_src/types.py

@@ -218,6 +218,7 @@ class Option(PyTreeNode):
 
   Attributes:
     timestep:         timestep
+    impratio:         ratio of friction-to-normal contact impedance
     tolerance:        main solver tolerance
     ls_tolerance:     CG/Newton linesearch tolerance
     gravity:          gravitational acceleration                 (3,)
@@ -238,9 +239,11 @@ class Option(PyTreeNode):
     disableflags:     bit flags for disabling standard features
   """
   timestep: jax.Array
+  # unsupported: apirate
+  impratio: jax.Array
   tolerance: jax.Array
   ls_tolerance: jax.Array
-  # unsupported: apirate, impratio, noslip_tolerance, mpr_tolerance
+  # unsupported: noslip_tolerance, mpr_tolerance
   gravity: jax.Array
   wind: jax.Array
   density: jax.Array

mjx/mujoco/mjx/test_data/constraints.xml

@@ -5,7 +5,7 @@
 * solref, solimp
 -->
 <mujoco>
-  <option timestep="0.015"/>
+  <option timestep="0.015" impratio="1.5"/>
 
   <default>
     <default class="box">