diff --git a/doc/changelog.rst b/doc/changelog.rst index 43e5e34e..04cb26a5 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -18,7 +18,8 @@ MJX ^^^ - Added ``mocap_pos`` and ``mocap_quat`` in kinematics. - Added support for :ref:`spatial tendons ` with pulleys and external sphere and cylinder wrapping. -- Add sphere-cylinder and sphere-ellipsoid collision functions (:github:issue:`2126`). +- Added sphere-cylinder and sphere-ellipsoid collision functions (:github:issue:`2126`). +- Fixed a bug with frictionloss constraints. Bug fixes ^^^^^^^^^ diff --git a/mjx/mujoco/mjx/_src/solver.py b/mjx/mujoco/mjx/_src/solver.py index a288c833..4c1bcbf9 100644 --- a/mjx/mujoco/mjx/_src/solver.py +++ b/mjx/mujoco/mjx/_src/solver.py @@ -161,7 +161,7 @@ class _LSPoint(PyTreeNode): linear_pos = (x >= rf)[:, None] qf = linear_neg * jp.array([f * (-0.5 * rf - ctx.Jaref), -f * jv, z]).T qf += linear_pos * jp.array([f * (-0.5 * rf + ctx.Jaref), f * jv, z]).T - quad = jp.where(f[:, None] > 0, qf, quad) + quad = jp.where((linear_neg | linear_pos) & (f[:, None] > 0), qf, quad) if m.opt.cone == ConeType.ELLIPTIC: mu, u0 = ctx.fri[:, 0], ctx.u[:, 0] diff --git a/mjx/mujoco/mjx/_src/solver_test.py b/mjx/mujoco/mjx/_src/solver_test.py index f29bd68e..989e4015 100644 --- a/mjx/mujoco/mjx/_src/solver_test.py +++ b/mjx/mujoco/mjx/_src/solver_test.py @@ -132,5 +132,20 @@ class SolverTest(parameterized.TestCase): nnz = dx.efc_J.any(axis=1) _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force') + def test_quad_frictionloss(self): + """Test a case with quadratic frictionloss constraints.""" + m = test_util.load_test_file('quadratic_frictionloss.xml') + d = mujoco.MjData(m) + mujoco.mj_resetDataKeyframe(m, d, 0) + mujoco.mj_forward(m, d) + + dx = jax.jit(mjx.solve)(mjx.put_model(m), mjx.put_data(m, d)) + + _assert_attr_eq(d, dx, 'qacc') + _assert_attr_eq(d, dx, 'qfrc_constraint') + nnz = dx.efc_J.any(axis=1) + _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force') + + if __name__ == '__main__': absltest.main() diff --git a/mjx/mujoco/mjx/test_data/quadratic_frictionloss.xml b/mjx/mujoco/mjx/test_data/quadratic_frictionloss.xml new file mode 100644 index 00000000..e62f81ae --- /dev/null +++ b/mjx/mujoco/mjx/test_data/quadratic_frictionloss.xml @@ -0,0 +1,21 @@ + + +