chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add mjENBL_DIAGEXACT for exact constraint diagonal. Fixes #2472

Browse Source 
PiperOrigin-RevId: 916932908
Change-Id: Id23ac39b5cd996afc52990719a4e07c0cc7de600
This commit is contained in:
Yuval Tassa
2026-05-17 16:47:20 -07:00
committed by Copybara-Service
parent 66d764a116
commit 71d1014e70
21 changed files with 234 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/XMLreference.rst
+13
View File
@@ -697,6 +697,19 @@ from its default.
order for the :ref:`sleep-init<body-sleep>` policy to take effect. Second, it must be set in order for static
quantities to be computed. See :ref:`implementation notes<siSleep>` for more details.
.. _option-flag-diagexact:
:at:`diagexact`: :at-val:`[disable, enable], "disable"`
This flag enables computation of the exact diagonal of the constraint-space inertia matrix :math:`A = J M^{-1} J^T`,
replacing the body-based approximation normally used. The exact diagonal is computed from the whitened Jacobian
:math:`Y = J M^{-1/2}` as :math:`A_{ii} = \|Y_i\|^2`. This provides a more accurate
:ref:`impedance<soParameters>` computation, which can improve solver quality for models with complex kinematic
coupling. See :ref:`Diagonal approximation <soExactDiag>` for details on the approximation errors that this flag
eliminates. The cost is one back-substitution with the Cholesky factor of the mass matrix per active constraint row;
if dual solvers are used (:ref:`PGS<option-solver>` or :ref:`NoSlip<option-noslip_iterations>`), the cost is
negligible since :math:`Y` is computed anyway. Consider enabling this flag when observing divergence or poor
constraint quality, particularly in models with highly anisotropic body inertias or bodies operating far from the
initial configuration ``qpos0``.
.. _compiler:
doc/XMLschema.rst
+3
View File
@@ -175,6 +175,9 @@
.. grid-item::
:ref:`sleep<option-flag-sleep>`
.. grid-item::
:ref:`diagexact<option-flag-diagexact>`
.. dropdown:: :ref:`compiler<compiler>` |*|
doc/changelog.rst
+6
View File
@@ -7,6 +7,12 @@ Upcoming version (not yet released)
General
^^^^^^^
- Added ``mjData.efc_Y``, the whitened constraint Jacobian :math:`Y = J M^{-1/2}`, allocated in the arena when
dual solvers (PGS or NoSlip) are used or when :ref:`diagexact<option-flag-diagexact>` is enabled.
- Added the :ref:`diagexact<option-flag-diagexact>` enable flag, which computes the exact diagonal of the
constraint-space inertia matrix at the current configuration, replacing the default compile-time approximation.
This improves solver quality for models with anisotropic inertias or complex kinematic coupling. See
:ref:`Exact diagonal <soExactDiag>` for details.
- The pseudo-random constraint visitation order in the :ref:`PGS solver<soAlgorithms>`, introduced in the previous
release, now uses a fixed seed. The previous implementation seeded with ``mjData.time``, which introduced subtle yet
undesirable time dependence.
doc/computation/index.rst
+14
View File
@@ -1514,6 +1514,20 @@ constraint would satisfy
and so we would achieve the desired interpolation effect. This of course does not hold exactly in general, but the goal
here is to construct a sensible and intuitive parameterization of the constraint model and get the scaling right.
.. _soExactDiag:
**Diagonal approximation:** The approximation has three sources of error: (i) it is frozen at ``qpos0`` rather than
evaluated at the current configuration; (ii) it averages the directional inverse inertia into a scalar, assuming
isotropy; and (iii) it treats the contributions of different bodies as independent, ignoring kinematic coupling through
shared DOFs. These errors are usually modest, but can become significant for models with highly anisotropic inertias or
long kinematic chains that operate far from ``qpos0``. In severe cases — particularly when the averaged inertia becomes
near-zero despite finite directional inertia — the regularizer :math:`R` becomes near-zero, making constraints
infinitely hard and causing divergence. The :ref:`diagexact<option-flag-diagexact>` flag replaces the approximation with
the exact diagonal :math:`A_{ii} = \|Y_i\|^2`, where :math:`Y = J M^{-1/2}` is the whitened Jacobian, computed at the
current configuration. This eliminates all three sources of error at a modest runtime cost: computing :math:`Y` requires
a back-substitution with the Cholesky factor of the mass matrix for each active constraint row; if
:ref:`dual solvers<soAlgorithms>` are used (PGS or NoSlip), the cost is negligible since :math:`Y` is computed anyway.
Next we explain how the reference acceleration is computed. As already mentioned, we use a spring-damper model
parameterized by *damping* and *stiffness* coefficients element-wise:
doc/includes/references.h
+2 -2
View File
@@ -508,10 +508,10 @@ typedef enum mjtEnableBit_ { // enable optional feature bitflags
mjENBL_ENERGY = 1<<1, // energy computation
mjENBL_FWDINV = 1<<2, // record solver statistics
mjENBL_INVDISCRETE = 1<<3, // discrete-time inverse dynamics
// experimental features:
mjENBL_SLEEP = 1<<4, // sleeping
mjENBL_DIAGEXACT = 1<<5, // exact diagonal of constraint inertia
mjNENABLE = 5 // number of enable flags
mjNENABLE = 6 // number of enable flags
} mjtEnableBit;
typedef enum mjtJoint_ { // type of degree of freedom
mjJNT_FREE = 0, // global position and orientation (quat) (7)
doc/modeling.rst
+19
View File
@@ -1855,6 +1855,25 @@ in a visible way, and the energy fluctuates around the initial value instead of
</worldbody>
.. _CConstraintImpedance:
Constraint accuracy
~~~~~~~~~~~~~~~~~~~
MuJoCo's :ref:`constraint impedance<soParameters>` computation relies on an approximate diagonal of the constraint-space
inertia matrix, computed once at compile time from the initial configuration ``qpos0``.
In the vast majority of models this approximation is entirely adequate. However, in certain situations—such as models
with highly anisotropic inertias, complex kinematic chains, or bodies operating far from ``qpos0``—the approximation
may become inaccurate. This can occasionally manifest as unexplained solver divergence (``badqacc`` warnings),
excessive penetration, unrealistic slip, or poor solver convergence. A useful diagnostic is the
:ref:`fwdinv<option-flag-fwdinv>` flag: if the forward-inverse discrepancy is large, inaccurate constraint scaling may
be a contributing factor.
If you suspect that the compile-time approximation is insufficient for your model, you can enable the
:ref:`diagexact<option-flag-diagexact>` flag to compute the exact diagonal at runtime. See :ref:`Diagonal approximation
<soExactDiag>` for details on the underlying mechanics and performance implications.
.. |image3| image:: images/modeling/tendonwraps.png
:width: 500px
.. |image4| image:: images/modeling/particle.png
include/mujoco/mjmodel.h
+2 -2
View File
@@ -81,10 +81,10 @@ typedef enum mjtEnableBit_ { // enable optional feature bitflags
mjENBL_ENERGY = 1<<1, // energy computation
mjENBL_FWDINV = 1<<2, // record solver statistics
mjENBL_INVDISCRETE = 1<<3, // discrete-time inverse dynamics
// experimental features:
mjENBL_SLEEP = 1<<4, // sleeping
mjENBL_DIAGEXACT = 1<<5, // exact diagonal of constraint inertia
mjNENABLE = 5 // number of enable flags
mjNENABLE = 6 // number of enable flags
} mjtEnableBit;
python/mujoco/bindings_test.py
+1 -1
View File
@@ -966,7 +966,7 @@ Euler integrator, semi-implicit in velocity.
self.assertEqual(mujoco.mjtEnableBit.mjENBL_OVERRIDE, 1 << 0)
self.assertEqual(mujoco.mjtEnableBit.mjENBL_ENERGY, 1 << 1)
self.assertEqual(mujoco.mjtEnableBit.mjENBL_FWDINV, 1 << 2)
self.assertEqual(mujoco.mjtEnableBit.mjNENABLE, 5)
self.assertEqual(mujoco.mjtEnableBit.mjNENABLE, 6)
self.assertEqual(mujoco.mjtGeom.mjGEOM_PLANE, 0)
self.assertEqual(mujoco.mjtGeom.mjGEOM_HFIELD, 1)
self.assertEqual(mujoco.mjtGeom.mjGEOM_SPHERE, 2)
python/mujoco/introspect/enums.py
+2 -1
View File
@@ -60,7 +60,8 @@ ENUMS: Mapping[str, EnumDecl] = dict([
('mjENBL_FWDINV', 4),
('mjENBL_INVDISCRETE', 8),
('mjENBL_SLEEP', 16),
('mjNENABLE', 5),
('mjENBL_DIAGEXACT', 32),
('mjNENABLE', 6),
]),
)),
('mjtJoint',
python/mujoco/introspect/enums_test.py
+2 -1
View File
@@ -43,7 +43,8 @@ class EnumsTest(absltest.TestCase):
('mjENBL_FWDINV', 1<<2),
('mjENBL_INVDISCRETE', 1<<3),
('mjENBL_SLEEP', 1<<4),
('mjNENABLE', 5)))
('mjENBL_DIAGEXACT', 1<<5),
('mjNENABLE', 6)))
# values mostly increment by one with occasional overrides
def test_mjtGeom(self): # pylint: disable=invalid-name
src/engine/engine_core_constraint.c
+84 -25
View File
@@ -2812,15 +2812,10 @@ void mj_makeConstraint(const mjModel* m, mjData* d) {
}
// compute efc_AR
void mj_projectConstraint(const mjModel* m, mjData* d) {
// compute Y = J*M^{-1/2}; if flg_diagexact, overwrite efc_diagApprox with ||Y_i||^2
static void mj_makeY(const mjModel* m, mjData* d, int flg_diagexact) {
int nefc = d->nefc, nv = m->nv;
// nothing to do
if (nefc == 0 || !mj_isDual(m)) {
return;
}
mj_markStack(d);
// inverse square root of D from inertia LDL decomposition
@@ -2843,10 +2838,8 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
return;
}
// markers for merged dofs, initialized to -1
int* marker = mjSTACKALLOC(d, nv, int);
// pre-count Y_rownnz, Y_rowadr, nY (total nonzeros)
int* marker = mjSTACKALLOC(d, nv, int);
d->nY = computeY_precount(d->efc_Y_rownnz, d->efc_Y_rowadr, nefc, nv,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
m->M_rownnz, m->M_rowadr, m->M_colind, marker);
@@ -2867,12 +2860,57 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
d->efc_J, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
m->dof_parentid);
// in-place sparse back-substitution: Y <- Y * M^-1/2
computeY_backsub(d->efc_Y, d->efc_Y_rownnz, d->efc_Y_rowadr,
d->efc_Y_colind, nefc,
d->qLD, m->M_rownnz, m->M_rowadr, m->M_colind, sqrtInvD);
// overwrite diagApprox with exact diagonal: diagApprox[i] = ||Y_i||^2
if (flg_diagexact) {
for (int i=0; i < nefc; i++) {
int adr = d->efc_Y_rowadr[i];
int nnz = d->efc_Y_rownnz[i];
d->efc_diagApprox[i] = mju_dot(d->efc_Y+adr, d->efc_Y+adr, nnz);
}
}
}
// dense Y = backsubM2(J')' and its transpose
else {
// arena-allocate efc_Y
d->nY = nefc * nv;
d->efc_Y = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nY, _Alignof(mjtNum));
if (!d->efc_Y) {
mj_warning(d, mjWARN_CNSTRFULL, d->narena);
mj_clearEfc(d);
d->parena = d->ncon * sizeof(mjContact);
mj_freeStack(d);
return;
}
// Y = backsubM2(J')'
mj_solveM2(m, d, d->efc_Y, d->efc_J, sqrtInvD, nefc);
// overwrite diagApprox with exact diagonal: diagApprox[i] = ||Y_i||^2
if (flg_diagexact) {
for (int i=0; i < nefc; i++) {
d->efc_diagApprox[i] = mju_dot(d->efc_Y+i*nv, d->efc_Y+i*nv, nv);
}
}
}
mj_freeStack(d);
}
// assemble AR = Y*Y' + diag(R) for dual solver
static void mj_makeAR(const mjModel* m, mjData* d) {
int nefc = d->nefc, nv = m->nv;
mj_markStack(d);
// sparse
if (mj_isSparse(m)) {
// Y supernodes are identical to J supernodes
const int* Y_rowsuper = d->efc_J_rowsuper;
@@ -2934,20 +2972,6 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
// dense Y = backsubM2(J')' and its transpose
else {
// arena-allocate efc_Y
d->nY = nefc * nv;
d->efc_Y = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nY, _Alignof(mjtNum));
if (!d->efc_Y) {
mj_warning(d, mjWARN_CNSTRFULL, d->narena);
mj_clearEfc(d);
d->parena = d->ncon * sizeof(mjContact);
mj_freeStack(d);
return;
}
// Y = backsubM2(J')'
mj_solveM2(m, d, d->efc_Y, d->efc_J, sqrtInvD, nefc);
// arena-allocate efc_AR
d->nA = nefc * nefc;
d->efc_AR = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nA, _Alignof(mjtNum));
@@ -2976,6 +3000,41 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
}
// compute efc_Y, optionally efc_diagApprox, optionally efc_AR
void mj_projectConstraint(const mjModel* m, mjData* d) {
int nefc = d->nefc;
// nothing to do
if (!nefc) {
return;
}
int isDual = mj_isDual(m);
int diagexact = mjENABLED(mjENBL_DIAGEXACT);
// compute Y = J*M^{-1/2}; overwrite diagApprox if diagexact
if (isDual || diagexact) {
mj_makeY(m, d, diagexact);
}
// recompute impedance from exact diagonal
if (diagexact && d->nefc) {
mj_makeImpedance(m, d);
// re-gather island D/R
if (d->nisland) {
mju_gather(d->iefc_D, d->efc_D, d->map_iefc2efc, d->nefc);
mju_gather(d->iefc_R, d->efc_R, d->map_iefc2efc, d->nefc);
}
}
// assemble AR for dual solver
if (isDual && d->nefc) {
mj_makeAR(m, d);
}
}
// compute efc_vel, efc_aref
void mj_referenceConstraint(const mjModel* m, mjData* d) {
int nefc = d->nefc;
src/engine/engine_forward.c
+4 -4
View File
@@ -207,14 +207,14 @@ void mj_fwdPosition(const mjModel* m, mjData* d) {
mj_island(m, d);
TM_END(mjTIMER_POS_MAKE);
TM_RESTART;
mj_transmission(m, d);
TM_ADD(mjTIMER_POS_KINEMATICS);
TM_RESTART;
mj_projectConstraint(m, d);
TM_END(mjTIMER_POS_PROJECT);
TM_RESTART;
mj_transmission(m, d);
TM_ADD(mjTIMER_POS_KINEMATICS);
TM_END1(mjTIMER_POSITION);
}
src/engine/engine_inverse.c
+7
View File
@@ -59,6 +59,13 @@ void mj_invPosition(const mjModel* m, mjData* d) {
mj_makeConstraint(m, d);
TM_END(mjTIMER_POS_MAKE);
// compute exact diagonal if enabled
if (mjENABLED(mjENBL_DIAGEXACT)) {
TM_RESTART;
mj_projectConstraint(m, d);
TM_END(mjTIMER_POS_PROJECT);
}
TM_RESTART;
mj_transmission(m, d);
TM_ADD(mjTIMER_POS_KINEMATICS);
src/engine/engine_memory.h
+1
View File
@@ -76,6 +76,7 @@ static inline void mj_clearEfc(mjData* d) {
#undef X
d->nefc = 0;
d->nisland = 0;
d->nJ = d->nY = d->nA = 0;
d->contact = (mjContact*) d->arena;
// if any contacts are allocated, clear their efc_address
src/engine/engine_support.c
+2 -1
View File
@@ -77,7 +77,8 @@ const char* mjENABLESTRING[mjNENABLE] = {
"Energy",
"Fwdinv",
"InvDiscrete",
"Sleep"
"Sleep",
"DiagExact"
};
src/xml/xml_native_reader.cc
+3 -1
View File
@@ -118,7 +118,8 @@ std::vector<const char*> MJCF[nMJCF] = {
{"flag", "?", "constraint", "equality", "frictionloss", "limit", "contact",
"spring", "damper", "gravity", "clampctrl", "warmstart", "filterparent", "actuation",
"refsafe", "sensor", "midphase", "eulerdamp", "autoreset", "nativeccd", "island",
"override", "energy", "fwdinv", "invdiscrete", "multiccd", "sleep"},
"override", "energy", "fwdinv", "invdiscrete", "multiccd", "sleep",
"diagexact"},
{">"},
{"size", "*", "memory", "njmax", "nconmax", "nstack", "nuserdata", "nkey",
@@ -1297,6 +1298,7 @@ void mjXReader::Option(XMLElement* section, mjOption* opt) {
READENBL("fwdinv", mjENBL_FWDINV)
READENBL("invdiscrete", mjENBL_INVDISCRETE)
READENBL("sleep", mjENBL_SLEEP)
READENBL("diagexact", mjENBL_DIAGEXACT)
#undef READENBL
}
}
src/xml/xml_native_writer.cc
+1
View File
@@ -1126,6 +1126,7 @@ void mjXWriter::Option(XMLElement* root) {
WRITEENBL("fwdinv", mjENBL_FWDINV)
WRITEENBL("invdiscrete", mjENBL_INVDISCRETE)
WRITEENBL("sleep", mjENBL_SLEEP)
WRITEENBL("diagexact", mjENBL_DIAGEXACT)
#undef WRITEENBL
}
test/engine/engine_forward_test.cc
+37
View File
@@ -917,6 +917,43 @@ TEST_F(ImplicitIntegratorTest, MidpointEligibility) {
mj_deleteModel(m);
}
// model with degenerate translational inertia
TEST_F(ForwardTest, DegenerateInertia) {
static constexpr char xml[] = R"(
<mujoco>
<option integrator="implicitfast" cone="elliptic">
<flag gravity="disable" diagexact="enable"/>
</option>
<worldbody>
<body name="1" pos="0.05 0.3 0">
<joint name="1" axis="0 1 0"/>
<geom type="capsule" size="0.1 0.5"/>
</body>
<body name="2">
<joint name="2" axis="1 0 0" stiffness="1" springref="90"/>
<geom type="capsule" size="0.1 0.5"/>
</body>
</worldbody>
</mujoco>
)";
char error[1024];
mjModel* model = LoadModelFromString(xml, error, sizeof(error));
ASSERT_THAT(model, NotNull()) << error;
mjData* data = mj_makeData(model);
for (int i = 0; i < 1000; i++) {
mj_step(model, data);
EXPECT_EQ(data->warning[mjWARN_BADQACC].number, 0)
<< "divergence at timestep " << i;
if (data->warning[mjWARN_BADQACC].number != 0) {
break;
}
}
mj_deleteData(data);
mj_deleteModel(model);
}
TEST_F(ForwardTest, ControlClamping) {
static constexpr char xml[] = R"(
<mujoco>
test/engine/engine_inverse_test.cc
+28 -18
View File
@@ -48,28 +48,38 @@ TEST_F(InverseTest, ForwardInverseMatch) {
// solver names for diagnostics
const char* solver_name[] = {"PGS", "CG", "Newton"};
for (mjtSolver solver : {mjSOL_PGS, mjSOL_CG, mjSOL_NEWTON}) {
model->opt.solver = solver;
mj_resetData(model, data);
// simulate, call mj_forward
for (int i = 0; i < kSteps; ++i) {
mj_step(model, data);
for (int diagexact = 0; diagexact < 2; diagexact++) {
if (diagexact) {
model->opt.enableflags |= mjENBL_DIAGEXACT;
} else {
model->opt.enableflags &= ~mjENBL_DIAGEXACT;
}
mj_forward(model, data);
// call built-in testing function
mj_compareFwdInv(model, data);
for (mjtSolver solver : {mjSOL_PGS, mjSOL_CG, mjSOL_NEWTON}) {
model->opt.solver = solver;
mj_resetData(model, data);
// per-solver tolerances
mjtNum epsilon;
switch (solver) {
case mjSOL_PGS: epsilon = MjTol(1e-6, 1e-2); break;
case mjSOL_CG: epsilon = MjTol(1e-3, 1e-1); break;
case mjSOL_NEWTON: epsilon = MjTol(1e-10, 5e-3); break;
// simulate, call mj_forward
for (int i = 0; i < kSteps; ++i) {
mj_step(model, data);
}
mj_forward(model, data);
// call built-in testing function
mj_compareFwdInv(model, data);
// per-solver tolerances
mjtNum epsilon;
switch (solver) {
case mjSOL_PGS: epsilon = MjTol(1e-6, 1e-2); break;
case mjSOL_CG: epsilon = MjTol(1e-3, 1e0); break;
case mjSOL_NEWTON: epsilon = MjTol(1e-10, 1e-2); break;
}
EXPECT_LT(data->solver_fwdinv[0], epsilon)
<< solver_name[solver] << " diagexact=" << diagexact;
EXPECT_LT(data->solver_fwdinv[1], epsilon)
<< solver_name[solver] << " diagexact=" << diagexact;
}
EXPECT_LT(data->solver_fwdinv[0], epsilon) << solver_name[solver];
EXPECT_LT(data->solver_fwdinv[1], epsilon) << solver_name[solver];
}
mj_deleteData(data);
unity/Runtime/Bindings/MjBindings.cs
+2 -1
View File
@@ -146,7 +146,8 @@ public enum mjtEnableBit : int{
mjENBL_FWDINV = 4,
mjENBL_INVDISCRETE = 8,
mjENBL_SLEEP = 16,
mjNENABLE = 5,
mjENBL_DIAGEXACT = 32,
mjNENABLE = 6,
}
public enum mjtJoint : int{
mjJNT_FREE = 0,
wasm/codegen/generated/bindings.cc
+1
View File
@@ -11009,6 +11009,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
.value("mjENBL_FWDINV", mjENBL_FWDINV)
.value("mjENBL_INVDISCRETE", mjENBL_INVDISCRETE)
.value("mjENBL_SLEEP", mjENBL_SLEEP)
.value("mjENBL_DIAGEXACT", mjENBL_DIAGEXACT)
.value("mjNENABLE", mjNENABLE);
enum_<mjtEq>("mjtEq")
.value("mjEQ_CONNECT", mjEQ_CONNECT)