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

Implement midpoint integrator for free bodies.

Browse Source 
PiperOrigin-RevId: 899043541
Change-Id: I0bb38f6ad94e189b45ab16777a04ad6fefc6adf7
This commit is contained in:
Yuval Tassa
2026-04-13 09:37:51 -07:00
committed by Copybara-Service
parent d9b5d8babb
commit 0c337799bd
12 changed files with 913 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/XMLreference.rst
+16 -7
View File
@@ -657,13 +657,22 @@ from its default.
.. _option-flag-invdiscrete:
:at:`invdiscrete`: :at-val:`[disable, enable], "disable"`
This flag enables discrete-time inverse dynamics with :ref:`mj_inverse` for all
:ref:`integrators<option-integrator>` other than ``RK4``. Recall from the
:ref:`numerical integration<geIntegration>` section that the one-step integrators (``Euler``, ``implicit`` and
``implicitfast``), modify the mass matrix :math:`M \rightarrow M-hD`. This implies that finite-differenced
accelerations :math:`(v_{t+h} - v_t)/h` will not correspond to the continuous-time acceleration ``mjData.qacc``.
When this flag is enabled, :ref:`mj_inverse` will interpret ``qacc`` as having been computed from the difference of
two sequential velocities, and undo the above modification.
This dual-purpose flag enables discrete-time inverse dynamics and disables :ref:`midpoint integration<geMidpoint>`.
Enable discrete-time inverse dynamics
This flag **enables** discrete-time inverse dynamics with :ref:`mj_inverse` for all
:ref:`integrators<option-integrator>` other than ``RK4``. Recall from the :ref:`numerical
integration<geIntegration>` section that the one-step integrators (``Euler``, ``implicit`` and ``implicitfast``),
modify the mass matrix :math:`M \rightarrow M-hD`. This implies that finite-differenced accelerations
:math:`(v_{t+h} - v_t)/h` will not correspond to the continuous-time acceleration ``mjData.qacc``. When this flag
is enabled, :ref:`mj_inverse` will interpret ``qacc`` as having been computed from the difference of two sequential
velocities, and undo the above modification.
Disable midpoint integration
Additionally and relatedly, this flag **disables** :ref:`midpoint integration<geMidpoint>` for free bodies, which
would otherwise break the linear relationship between finite-differenced velocities and forces assumed by discrete
inverse dynamics. Note that disabling midpoint integration might be useful for debugging or for other reasons,
regardless or whether inverse dynamics are used.
.. _option-flag-multiccd:
doc/changelog.rst
+7 -2
View File
@@ -27,8 +27,13 @@ General
(``jnt_stiffness``, ``dof_damping``, etc.) continue to hold the linear coefficient and are unchanged.
The polynomial order is defined by the new constant :ref:`mjNPOLY<glNumericSizes>`. A future breaking C-API change
may unify the linear and higher-order coefficients into a single array.
- Introduced :ref:`mjpEncoder`, the counterpart to :ref:`mjpDecoder` for encoding of :ref:`mjSpec` and :ref:`mjModel` into :ref:`mjResource`.
- Added :ref:`midpoint integration<geMidpoint>` for standalone free bodies in ``implicit`` and ``implicitfast``
:ref:`integrators<geIntegrators>`. This applies the implicit midpoint rule to the rotational dynamics of free bodies
with no children, exactly conserving kinetic energy and angular momentum in the absence of external torques. The
:ref:`invdiscrete<option-flag-invdiscrete>` flag now also disables midpoint integration, providing an opt-out
mechanism.
- Introduced :ref:`mjpEncoder`, the counterpart to :ref:`mjpDecoder` for encoding of :ref:`mjSpec` and :ref:`mjModel`
into :ref:`mjResource`.
- Added :ref:`mj_encode`, :ref:`mjp_registerEncoder`, :ref:`mjp_defaultEncoder`, and :ref:`mjp_findEncoder`.
doc/computation/index.rst
+51 -5
View File
@@ -573,6 +573,49 @@ Solving for :math:`v_{t+h}`, we obtain the implicit-in-velocity update
\widehat{M} &\equiv M-h D
\end{aligned}
.. _geMidpoint:
Midpoint integration for free bodies
The implicit-in-velocity update :eq:`eq_implicit_update` treats the acceleration as a function of velocity and
linearizes. While effective for damping-like forces, it is sub-optimal for rotational dynamics, where
Coriolis and gyroscopic forces are *quadratic* in angular velocity. For this case, a better approach is to directly
discretize the rotational equations of motion using the *midpoint method*.
Consider a rigid body rotating in its principal-axis frame with angular velocity
:math:`\omega \in \mathbb{R}^3` and diagonal inertia tensor :math:`I = \text{diag}(I_1, I_2, I_3)`. The rotational
dynamics are given by `Euler's rotation equation
<https://en.wikipedia.org/wiki/Euler%27s_equations_(rigid_body_dynamics)>`__:
.. math::
I \dot{\omega} + \omega \times I\omega = \tau
where :math:`\tau` is the external torque in the principal-axis frame.
Evaluating the velocities at the midpoint, :math:`\omega_\text{mid} = (\omega_t + \omega_{t+h})/2`, gives:
.. math::
\frac{2}{h} I (\omega_\text{mid} - \omega_t) + \omega_\text{mid} \times I \omega_\text{mid} = \tau
This is a system of 3 nonlinear equations in 3 unknowns :math:`\omega_\text{mid}`, solved at each timestep using
Newton's method with a backtracking line search. After solving, the new velocity is recovered as
:math:`\omega_{t+h} = 2\omega_\text{mid} - \omega_t`.
**Properties.** The midpoint method preserves all `quadratic first integrals
<https://doi.org/10.1007/3-540-30666-8>`__ of the ODE. For Euler's equations, these are the
kinetic energy :math:`H = \frac{1}{2}\omega^T I\omega` and the squared angular momentum
:math:`C = \frac{1}{2}|I\omega|^2`, both conserved exactly in the absence of external torque. Since :math:`C` is the
Casimir function of the `Lie-Poisson <https://en.wikipedia.org/wiki/Poisson_bracket>`__ structure, the midpoint
method is a symmetric (time-reversible) and second-order accurate *Poisson integrator*.
**Eligibility.** Midpoint integration is only applied to free bodies with no child bodies.
**Performance.** While the midpoint method carries computational overhead, we've found it to be
negligible compared to the rest of the pipeline, on the order of 1% in the worst case.
**Disabling.** Because midpoint integration solves a nonlinear equation for the next velocity, it breaks the linear
relationship between finite-differenced velocities and forces assumed by discrete inverse dynamics. Therefore,
setting the :ref:`invdiscrete<option-flag-invdiscrete>` flag disables midpoint integration, and also provides a
general opt-out mechanism for this integrator.
.. _geIntegrators:
Integrators
@@ -610,6 +653,9 @@ Fast implicit-in-velocity (``implicitfast``)
derivatives are also the main source of asymmetry of :math:`D`, by dropping them and symmetrizing, we can use the
faster :math:`L^TL` rather than :math:`LU` decomposition.
Both ``implicit`` and ``implicitfast`` apply :ref:`midpoint integration<geMidpoint>` to eligible free bodies,
providing exact energy conservation for spinning objects at negligible additional cost.
4th-order Runge-Kutta (``RK4``)
One advantage of our continuous-time formulation is that we can use higher order integrators such as Runge-Kutta or
multistep methods. MuJoCo implements the fixed-step `4th-order Runge-Kutta method
@@ -641,11 +687,11 @@ Fast implicit-in-velocity (``implicitfast``)
The ``implicitfast`` integrator has similar computational cost to ``Euler``, yet provides
increased stability, and is therefore a strict improvement. It is the recommended integrator for most models.
**implicit**:
The benefit over ``implicitfast`` is the implicit integration of Coriolis and centripetal forces, including
gyroscopic forces. The most common case where integrating such forces implicitly leads to noticeable improvement is
when free objects with asymmetric inertia are spinning quickly. `gyroscopic.xml <../_static/gyroscopic.xml>`__
shows an ellipsoid rolling on an inclined plane which quickly diverges with ``implicitfast`` but is stable with
``implicit``.
The benefit over ``implicitfast`` is the implicit integration of Coriolis and centripetal forces for *coupled*
rotational systems such as multi-link pendula. Both ``implicitfast`` and ``implicit`` apply :ref:`midpoint
integration<geMidpoint>` to eligible free bodies with no children, for example
`gyroscopic.xml <../_static/gyroscopic.xml>`__ shows an ellipsoid rolling on an
inclined plane; both ``implicitfast`` and ``implicit`` handle this case well, while ``Euler`` quickly diverges.
**RK4**:
This integrator is best for systems which are energy conserving, or almost energy-conserving. `pendulum.xml
<../_static/pendulum.xml>`__ shows a complicated pendulum mechanism which diverges quickly using ``Euler`` or
mjx/mujoco/mjx/_src/forward_test.py
+2
View File
@@ -78,6 +78,8 @@ class ForwardTest(absltest.TestCase):
# implicitfast
m.opt.integrator = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
# TODO(team): remove this override when the mjx feature matches mujoco
m.opt.enableflags |= mujoco.mjtEnableBit.mjENBL_INVDISCRETE
dx = jax.jit(mjx.implicit)(mx, mjx.put_data(m, d))
mujoco.mj_implicit(m, d)
_assert_attr_eq(d, dx, 'qpos')
src/engine/engine_forward.c
+350 -3
View File
@@ -38,6 +38,7 @@
#include "engine/engine_sleep.h"
#include "engine/engine_solver.h"
#include "engine/engine_support.h"
#include "engine/engine_inline.h"
#include "engine/engine_util_blas.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
@@ -1043,7 +1044,10 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
//-------------------------- state advancement and integration ------------------------------------
// advance state and time given activation derivatives, acceleration, and optional velocity
// advance state and time
// act_dot: activation derivatives
// qacc: acceleration used to update d->qvel (d->qvel += h*qacc)
// qvel: optional velocity used for position integration; if NULL, use d->qvel
static void mj_advance(const mjModel* m, mjData* d,
const mjtNum* act_dot, const mjtNum* qacc, const mjtNum* qvel) {
int nu = m->nu, nsensor = m->nsensor;
@@ -1546,6 +1550,293 @@ static void flexInterp_solve(const mjModel* m, mjData* d, const FlexInterpContex
}
// return 1 if free joint is eligible for midpoint quaternion integration:
// standalone 6-DOF tree with no children
static int midpoint_eligible(const mjModel* m, int jnt) {
if (m->jnt_type[jnt] == mjJNT_FREE) {
int body = m->jnt_bodyid[jnt];
int treeid = m->dof_treeid[m->jnt_dofadr[jnt]];
return m->tree_dofnum[treeid] == 6 &&
m->body_subtreemass[body] == m->body_mass[body];
}
return 0;
}
// return 1 if the body's CoM is at the joint origin (no translational-rotational coupling)
static int midpoint_aligned(const mjModel* m, int jnt) {
int body = m->jnt_bodyid[jnt];
return m->body_ipos[3*body+0] == 0 &&
m->body_ipos[3*body+1] == 0 &&
m->body_ipos[3*body+2] == 0;
}
// implicit midpoint integration for 3D rotation of a single body
//
// solves the Euler rigid body equation in the inertial frame:
// I * (w_new - w) / h = tau - w_mid x (I*w_mid)
// where w_mid = (w + w_new) / 2 is solved via Newton iteration.
//
// inputs:
// inertia: principal moments of inertia (3)
// w: initial angular velocity in principal axes frame (3)
// tau: external torque in principal axes frame (3)
// h: timestep
// outputs:
// w_mid: midpoint angular velocity in principal axes frame (3)
// returns: number of Newton iterations
static int midpointNewton(const mjtNum inertia[3], const mjtNum w[3],
const mjtNum tau[3], mjtNum h, mjtNum w_mid[3]) {
// precompute constants
mjtNum i2h = 2.0 / h;
mjtNum dI[3] = {inertia[2]-inertia[1], inertia[0]-inertia[2], inertia[1]-inertia[0]};
mjtNum i2h_I[3] = {i2h*inertia[0], i2h*inertia[1], i2h*inertia[2]};
// initialize solution to previous angular velocity
mji_copy3(w_mid, w);
// Newton iteration
int niter;
for (niter=0; niter < 100; niter++) {
// compute Coriolis term
mjtNum Iw[3] = {inertia[0]*w_mid[0], inertia[1]*w_mid[1], inertia[2]*w_mid[2]};
mjtNum coriolis[3];
mji_cross(coriolis, w_mid, Iw);
// residual: f = i2h*I*(w_mid - w) + w_mid x (I*w_mid) - tau
mjtNum f[3];
for (int k=0; k < 3; k++) {
f[k] = i2h_I[k]*(w_mid[k] - w[k]) + coriolis[k] - tau[k];
}
// check convergence
mjtNum fnorm = mju_norm3(f);
#ifndef mjUSESINGLE
mjtNum tol = 1e-13;
#else
mjtNum tol = 1e-6f;
#endif
if (fnorm < tol*(1 + i2h*mju_norm3(Iw))) break;
// Jacobian: J = i2h*diag(I) + d(w x Iw)/dw
mjtNum J[9];
J[0] = i2h_I[0]; J[1] = w_mid[2]*dI[0]; J[2] = w_mid[1]*dI[0];
J[3] = w_mid[2]*dI[1]; J[4] = i2h_I[1]; J[5] = w_mid[0]*dI[1];
J[6] = w_mid[1]*dI[2]; J[7] = w_mid[0]*dI[2]; J[8] = i2h_I[2];
// solve J*delta = -f for search direction delta
mjtNum neg_f[3] = {-f[0], -f[1], -f[2]};
mjtNum delta[3];
mju_solve3(delta, J, neg_f);
// backtracking line search
mjtNum step = 1.0;
for (int ls=0; ls < 20; ls++) {
// candidate step
mjtNum w_try[3], Iw_try[3];
for (int k=0; k < 3; k++) {
w_try[k] = w_mid[k] + step*delta[k];
Iw_try[k] = inertia[k]*w_try[k];
}
mjtNum coriolis_try[3];
mji_cross(coriolis_try, w_try, Iw_try);
// residual at candidate step
mjtNum f_try[3];
for (int k=0; k < 3; k++) {
f_try[k] = i2h_I[k]*(w_try[k] - w[k]) + coriolis_try[k] - tau[k];
}
// accept step if residual decreased, otherwise backtrack
if (mju_norm3(f_try) < fnorm) {
mji_copy3(w_mid, w_try);
break;
}
step *= 0.5;
}
}
return niter;
}
// implicit midpoint integration for one free body
//
// solves the Euler rigid body equation in the inertial frame:
// I * dw/dt = tau - w x (I*w)
// using the implicit midpoint rule:
// I * (w_new - w_old) / h = tau_mid - w_mid x (I*w_mid)
// where w_mid = (w_old + w_new) / 2 is solved via Newton iteration.
//
// inputs:
// mass: body mass
// inertia: principal moments of inertia
// ipos: CoM offset from joint origin, in body frame
// iquat: inertial quaternion (body_iquat)
// xquat: body orientation in world frame
// qvel_old: current velocity (lin in world : rot in body)
// qfrc: external force (lin in world : rot in body)
// gravity: gravitational acceleration in world frame (NULL: no gravity)
// h: timestep
// outputs:
// qvel_new: next velocity (lin in world : rot in body)
int mj_midpoint(mjtNum mass, const mjtNum inertia[3], const mjtNum ipos[3],
const mjtNum iquat[4], const mjtNum xquat[4], const mjtNum qvel_old[6],
const mjtNum qfrc[6], const mjtNum gravity[3], mjtNum h,
mjtNum qvel_new[6]) {
// transform angular velocity and torque to inertial frame
mjtNum iquat_neg[4], w[3], tau[3];
mji_negQuat(iquat_neg, iquat);
mji_rotVecQuat(w, qvel_old+3, iquat_neg); // qvel+3 (angular) is in body frame
mji_rotVecQuat(tau, qfrc+3, iquat_neg); // qfrc+3 (angular) is in body frame
// check for translational-rotational coupling
int aligned = (ipos[0] == 0 && ipos[1] == 0 && ipos[2] == 0);
mjtNum r_com[3]; // joint-to-CoM vector in inertial frame
mjtNum tau_com[3]; // torque at CoM in inertial frame
mjtNum rot_x2i[4]; // quaternion rotation from world to inertial frame
mjtNum force[3]; // external force in inertial frame
// compute torque at CoM in inertial frame
if (aligned) {
mji_copy3(tau_com, tau);
} else {
// rotation from world to inertial frame
mjtNum xquat_neg[4];
mji_negQuat(xquat_neg, xquat);
mji_mulQuat(rot_x2i, iquat_neg, xquat_neg);
// force and CoM offset in inertial frame
mji_rotVecQuat(force, qfrc, rot_x2i);
mji_rotVecQuat(r_com, ipos, iquat_neg);
// torque at CoM in inertial frame
mjtNum rxf[3];
mji_cross(rxf, r_com, force);
mji_sub3(tau_com, tau, rxf);
}
// solve for midpoint angular velocity
mjtNum w_mid[3];
int niter = midpointNewton(inertia, w, tau_com, h, w_mid);
// next and mid angular velocities in inertial frame, rotate both to body frame
mjtNum w_new[3], w_new_body[3], w_mid_body[3];
for (int k=0; k < 3; k++) {
w_new[k] = 2.0*w_mid[k] - w[k];
}
mji_rotVecQuat(w_new_body, w_new, iquat);
mji_rotVecQuat(w_mid_body, w_mid, iquat);
mji_copy3(qvel_new+3, w_new_body);
// === aligned: return
if (aligned) {
return niter;
}
// === non-aligned: solve for translational velocity
// rotate linear velocity to inertial frame
mjtNum v[3];
mji_rotVecQuat(v, qvel_old, rot_x2i);
// current CoM velocities (rot, lin) in inertial frame
mjtNum wxr[3];
mji_cross(wxr, w, r_com);
mjtNum vcom[3];
mji_add3(vcom, v, wxr);
// right-hand side for midpoint CoM velocity
mjtNum i2h = 2.0 / h;
mjtNum b[3];
for (int k=0; k < 3; k++) {
b[k] = force[k]/mass + i2h*vcom[k];
}
// add gravity, if any
if (gravity) {
mjtNum g_inertial[3];
mji_rotVecQuat(g_inertial, gravity, rot_x2i);
mji_addTo3(b, g_inertial);
}
// analytic solution for (i2h*Id + [w_mid]x) * vcom_mid = b
mjtNum wnorm2 = mju_dot3(w_mid, w_mid);
mjtNum denom = i2h*i2h + wnorm2;
mjtNum w_dot_b = mju_dot3(w_mid, b);
mjtNum w_cross_b[3];
mji_cross(w_cross_b, w_mid, b);
mjtNum vcom_mid[3];
for (int k=0; k < 3; k++) {
vcom_mid[k] = (i2h*b[k] + (w_dot_b/i2h)*w_mid[k] - w_cross_b[k]) / denom;
}
// recover midpoint and new joint velocity in inertial frame
mjtNum wxr_mid[3];
mji_cross(wxr_mid, w_mid, r_com);
mjtNum v_mid[3], v_new[3];
for (int k=0; k < 3; k++) {
v_mid[k] = vcom_mid[k] - wxr_mid[k];
v_new[k] = 2.0*v_mid[k] - v[k];
}
// estimate new orientation
mjtNum axis[3];
mji_copy3(axis, w_mid_body);
mjtNum wnorm = mju_normalize3(axis);
mjtNum qrot_new[4];
mji_axisAngle2Quat(qrot_new, axis, h*wnorm);
mjtNum xquat_new[4];
mji_mulQuat(xquat_new, xquat, qrot_new);
// v_new (linear): inertial → body → world using new orientation
mjtNum v_body[3];
mji_rotVecQuat(v_body, v_new, iquat);
mji_rotVecQuat(qvel_new, v_body, xquat_new);
return niter;
}
// compute next velocities via midpoint integration for eligible free bodies
// qfrc: total force (qfrc_smooth + qfrc_constraint)
// free_jntid: list of eligible free joint IDs
// nfree: number of eligible free joints
// qvel_old: output array for old velocities (6 per joint)
// qvel_new: output array for new velocities (6 per joint)
// dofadr: output array for DOF addresses (1 per joint)
static void midpoint(const mjModel* m, const mjData* d, const mjtNum* qfrc,
const int* free_jntid, int nfree,
mjtNum* qvel_old, mjtNum* qvel_new, int* dofadr) {
for (int i=0; i < nfree; i++) {
int j = free_jntid[i];
int body = m->jnt_bodyid[j];
// save DOF address
int adr = m->jnt_dofadr[j];
dofadr[i] = adr;
// save old (current) velocity, needed after mj_advance (which overwrites qvel)
mju_copy(qvel_old+6*i, d->qvel+adr, 6);
// compute external force = qfrc + qfrc_bias (undo bias subtraction)
mjtNum qfrc_total[6];
mju_add(qfrc_total, qfrc+adr, d->qfrc_bias+adr, 6);
// gravity handled inside mj_midpoint (accelerating frame of reference)
const mjtNum* gravity = mjDISABLED(mjDSBL_GRAVITY) ? NULL : m->opt.gravity;
// midpoint solver for free joint j
mj_midpoint(m->body_mass[body], m->body_inertia+3*body, m->body_ipos+3*body,
m->body_iquat+4*body, d->xquat+4*body,
d->qvel+adr, qfrc_total, gravity, m->opt.timestep, qvel_new+6*i);
}
}
// fully implicit in velocity, possibly skipping factorization
void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
TM_START;
@@ -1640,8 +1931,64 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
flexInterp_solve(m, d, &flex, qacc, qfrc, nv);
}
// advance state and time
mj_advance(m, d, d->act_dot, qacc, NULL);
// count and list joints of free bodies eligible for midpoint integration
int nfree = 0;
int* free_jntid = NULL;
if (!mjENABLED(mjENBL_INVDISCRETE)) {
free_jntid = mjSTACKALLOC(d, m->njnt, int);
for (int j=0; j < m->njnt; j++) {
// add to list if eligible and awake
if (midpoint_eligible(m, j) && d->tree_awake[m->dof_treeid[m->jnt_dofadr[j]]]) {
free_jntid[nfree++] = j;
}
}
}
// compute midpoint velocities (used to update positions)
int* dofadr = NULL;
mjtNum* qvel_old = NULL;
mjtNum* qvel_new = NULL;
mjtNum* qvel_mid = NULL;
if (nfree) {
// allocate arrays, call midpoint solver for all eligible free joints
dofadr = mjSTACKALLOC(d, nfree, int);
qvel_new = mjSTACKALLOC(d, 6*nfree, mjtNum);
qvel_old = mjSTACKALLOC(d, 6*nfree, mjtNum);
midpoint(m, d, qfrc, free_jntid, nfree, qvel_old, qvel_new, dofadr);
// build qvel_mid = d->qvel + h*qacc for all DOFs, then overwrite midpoint DOFs
qvel_mid = mjSTACKALLOC(d, m->nv, mjtNum);
mju_addScl(qvel_mid, d->qvel, qacc, m->opt.timestep, m->nv);
for (int i=0; i < nfree; i++) {
int adr = dofadr[i];
int start = midpoint_aligned(m, free_jntid[i]) ? 3 : 0;
for (int k=start; k < 6; k++) {
qvel_mid[adr+k] = 0.5*(qvel_new[6*i+k] + qvel_old[6*i+k]);
}
}
}
// advance state and time (use qvel_mid if allocated, NULL otherwise)
mj_advance(m, d, d->act_dot, qacc, qvel_mid);
// overwrite midpoint DOFs with true next velocity and acceleration
if (nfree) {
mjtNum h_inv = 1.0 / m->opt.timestep;
for (int i=0; i < nfree; i++) {
// skip sleeping tree (may have been put to sleep during mj_advance)
int adr = dofadr[i];
if (!d->tree_awake[m->dof_treeid[adr]]) {
continue;
}
// overwrite 3 or 6 midpoint DOFs with true next velocity and acceleration
int start = midpoint_aligned(m, free_jntid[i]) ? 3 : 0;
for (int k=start; k < 6; k++) {
d->qvel[adr+k] = qvel_new[6*i+k];
d->qacc[adr+k] = (qvel_new[6*i+k] - qvel_old[6*i+k]) * h_inv;
}
}
}
mj_freeStack(d);
src/engine/engine_forward.h
+8 -1
View File
@@ -18,6 +18,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjexport.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjtnum.h>
#ifdef __cplusplus
extern "C" {
@@ -46,7 +47,6 @@ MJAPI void mj_forward(const mjModel* m, mjData* d);
MJAPI void mj_forwardSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor);
//-------------------------------- integrators -----------------------------------------------------
// Runge Kutta explicit order-N integrator
@@ -64,6 +64,13 @@ MJAPI void mj_implicit(const mjModel *m, mjData *d);
// fully implicit in velocity, possibly skipping factorization
MJAPI void mj_implicitSkip(const mjModel *m, mjData *d, int skipfactor);
// implicit midpoint integration for 6 DOFs (translation + rotation) of a single body
// returns number of Newton iterations
MJAPI int mj_midpoint(mjtNum mass, const mjtNum inertia[3], const mjtNum ipos[3],
const mjtNum iquat[4], const mjtNum xquat[4], const mjtNum qvel[6],
const mjtNum qfrc[6], const mjtNum gravity[3], mjtNum h,
mjtNum qvel_new[6]);
//-------------------------------- solver components -----------------------------------------------
src/engine/engine_util_solve.c
+31
View File
@@ -883,6 +883,37 @@ void mju_solveLUSparse(mjtNum* res, const mjtNum* LU, const mjtNum* vec, int n,
}
//--------------------------- 3x3 linear solve -----------------------------------------------------
// solve 3x3 linear system A*x = b using Gaussian elimination
void mju_solve3(mjtNum x[3], const mjtNum A[9], const mjtNum b[3]) {
mjtNum M[3][4] = {
{A[0], A[1], A[2], b[0]},
{A[3], A[4], A[5], b[1]},
{A[6], A[7], A[8], b[2]}
};
for (int i=0; i<3; i++) {
mjtNum pivot = M[i][i];
for (int j=i; j<4; j++) {
M[i][j] /= pivot;
}
for (int k=0; k<3; k++) {
if (k != i) {
mjtNum factor = M[k][i];
for (int j=i; j<4; j++) {
M[k][j] -= factor * M[i][j];
}
}
}
}
x[0] = M[0][3];
x[1] = M[1][3];
x[2] = M[2][3];
}
//--------------------------- eigen decomposition --------------------------------------------------
// eigenvalue decomposition of symmetric 3x3 matrix
src/engine/engine_util_solve.h
+4
View File
@@ -17,6 +17,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjexport.h>
#include <mujoco/mjtnum.h>
#ifdef __cplusplus
extern "C" {
@@ -103,6 +104,9 @@ void mju_solveLUSparse(mjtNum *res, const mjtNum *LU, const mjtNum* vec, int n,
const int *rownnz, const int *rowadr, const int* diag, const int *colind,
const int *index);
// solve 3x3 linear system A*x = b using Gaussian elimination
void mju_solve3(mjtNum x[3], const mjtNum A[9], const mjtNum b[3]);
// eigenvalue decomposition of symmetric 3x3 matrix
MJAPI int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum mat[9]);
test/engine/engine_derivative_test.cc
+78 -2
View File
@@ -320,17 +320,19 @@ TEST_F(DerivativeTest, PassiveDvel) {
mj_forward(model, data);
// get analytic derivatives
mju_zero(data->qDeriv, model->nD);
mjd_passive_vel(model, data);
mju_copy(qDerivAnalytic, data->qDeriv, nD);
// clear qDeriv, get finite-difference derivatives
mju_zero(data->qDeriv, nD);
mju_zero(qDerivFD, nD);
mjtNum eps = MjTol(1e-6, 1e-3);
mjtNum eps = MjTol(1e-6, 1e-4);
mjd_passive_velFD(model, data, eps);
// expect FD and analytic derivatives to be similar to tol precision
EXPECT_THAT(AsVector(data->qDeriv, nD),
Pointwise(MjNear(1e-4, 1e-3), AsVector(qDerivAnalytic, nD)));
Pointwise(MjNear(1e-6, 1e-4), AsVector(qDerivAnalytic, nD)));
}
mju_free(qDerivFD);
@@ -1733,5 +1735,79 @@ TEST_F(DerivativeTest, FlexInterpDerivativesDeformed) {
mj_deleteModel(model);
}
TEST_F(DerivativeTest, MidpointFluidAccuracy) {
const std::string xml_path =
GetTestDataFilePath(kTumblingThinObjectEllipsoidPath);
char error[1024];
mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
ASSERT_THAT(m, NotNull()) << error;
mjtNum dt_small = 1e-4;
mjtNum dt_large = m->opt.timestep; // 2e-3, the default
mjtNum duration = 0.5;
mjData* d_ref = mj_makeData(m);
mjData* d_midpoint = mj_makeData(m);
mjData* d_nomidpoint = mj_makeData(m);
// give initial angular velocity for tumbling
mj_resetData(m, d_ref);
mj_resetData(m, d_midpoint);
mj_resetData(m, d_nomidpoint);
d_ref->qvel[3] = 5;
d_ref->qvel[4] = 3;
d_ref->qvel[5] = 1;
d_midpoint->qvel[3] = 5;
d_midpoint->qvel[4] = 3;
d_midpoint->qvel[5] = 1;
d_nomidpoint->qvel[3] = 5;
d_nomidpoint->qvel[4] = 3;
d_nomidpoint->qvel[5] = 1;
int nsteps_large = static_cast<int>(duration / dt_large);
int substeps = static_cast<int>(dt_large / dt_small);
mjtNum error_midpoint = 0;
mjtNum error_nomidpoint = 0;
for (int i = 0; i < nsteps_large; i++) {
// reference: RK4 at small timestep
m->opt.integrator = mjINT_RK4;
m->opt.timestep = dt_small;
m->opt.enableflags &= ~mjENBL_INVDISCRETE;
for (int j = 0; j < substeps; j++) {
mj_step(m, d_ref);
}
// implicit with midpoint (default)
m->opt.integrator = mjINT_IMPLICIT;
m->opt.timestep = dt_large;
m->opt.enableflags &= ~mjENBL_INVDISCRETE;
mj_step(m, d_midpoint);
// implicit without midpoint
m->opt.enableflags |= mjENBL_INVDISCRETE;
mj_step(m, d_nomidpoint);
// accumulate position errors
for (int k = 0; k < 7; k++) {
mjtNum diff_mid = d_ref->qpos[k] - d_midpoint->qpos[k];
mjtNum diff_nomid = d_ref->qpos[k] - d_nomidpoint->qpos[k];
error_midpoint += diff_mid * diff_mid;
error_nomidpoint += diff_nomid * diff_nomid;
}
}
// expect midpoint to be more accurate
EXPECT_LT(error_midpoint, error_nomidpoint)
<< "implicit midpoint should be more accurate than implicit without "
<< "midpoint for a free body with fluid forces";
mj_deleteData(d_nomidpoint);
mj_deleteData(d_midpoint);
mj_deleteData(d_ref);
mj_deleteModel(m);
}
} // namespace
} // namespace mujoco
test/engine/engine_forward_test.cc
+305
View File
@@ -17,6 +17,7 @@
#include "src/engine/engine_forward.h"
#include "src/engine/engine_derivative.h"
#include <algorithm>
#include <array>
#include <cmath>
#include <cstdlib>
@@ -478,6 +479,310 @@ TEST_F(ImplicitIntegratorTest, EnergyConservation) {
mj_deleteModel(model);
}
// Energy and angmom conservation for free body with implicitfast (IMR)
TEST_F(ImplicitIntegratorTest, ConservationMidpoint) {
// aligned: CoM at joint origin
static constexpr char xml1[] = R"(
<mujoco>
<option integrator="implicitfast" timestep="0.01">
<flag energy="enable" gravity="disable"/>
</option>
<worldbody>
<body>
<freejoint/>
<geom type="box" size=".1 .2 .3" mass="1" euler="10 20 30"/>
</body>
</worldbody>
</mujoco>
)";
// auto-aligned: CoM at joint origin
static constexpr char xml2[] = R"(
<mujoco>
<option integrator="implicitfast" timestep="0.01">
<flag energy="enable" gravity="disable"/>
</option>
<worldbody>
<body>
<freejoint align="true"/>
<geom type="box" size=".1 .2 .3" mass="1" euler="10 20 30" pos=".03 .02 .01"/>
</body>
</worldbody>
</mujoco>
)";
// non-aligned: CoM offset from joint origin
static constexpr char xml3[] = R"(
<mujoco>
<option integrator="implicitfast" timestep="0.01">
<flag energy="enable" gravity="disable"/>
</option>
<worldbody>
<body>
<freejoint/>
<geom type="box" size=".1 .2 .3" mass="1" euler="10 20 30" pos=".03 .02 .01"/>
</body>
</worldbody>
</mujoco>
)";
int xml_idx = 1;
for (auto xml : {xml1, xml2, xml3}) {
SCOPED_TRACE(testing::Message() << "XML case " << xml_idx++);
char error[1024];
mjModel* model = LoadModelFromString(xml, error, sizeof(error));
ASSERT_THAT(model, NotNull()) << error;
mjData* data = mj_makeData(model);
const int nstep = 500;
mjtNum energy_drift[2], angmom_drift[2]; // [0]=midpoint, [1]=rk4
for (int integrator : {mjINT_IMPLICITFAST, mjINT_RK4}) {
int idx = (integrator == mjINT_IMPLICITFAST) ? 0 : 1;
model->opt.integrator = integrator;
// reset
mj_resetData(model, data);
data->qvel[3] = 1.0;
data->qvel[4] = 2.0;
data->qvel[5] = 3.0;
mj_forward(model, data);
mjtNum initial_energy = data->energy[1];
mjtNum initial_angmom[3];
mj_subtreeVel(model, data);
mju_copy3(initial_angmom, data->subtree_angmom);
for (int i=0; i < nstep; i++) {
mj_step(model, data);
}
energy_drift[idx] = fabs(data->energy[1] - initial_energy);
mj_subtreeVel(model, data);
mjtNum angmom_err[3];
mju_sub3(angmom_err, data->subtree_angmom, initial_angmom);
angmom_drift[idx] = mju_norm3(angmom_err);
}
// midpoint should conserve energy better than RK4 (double only)
#ifndef mjUSESINGLE
EXPECT_LT(energy_drift[0], energy_drift[1]);
#endif
// both should conserve angular momentum well
EXPECT_LT(angmom_drift[0], MjTol(1e-3, 1e-2));
EXPECT_LT(angmom_drift[1], MjTol(1e-3, 1e-2));
mj_deleteData(data);
mj_deleteModel(model);
}
}
// verify second-order convergence of midpoint integration
TEST_F(ImplicitIntegratorTest, MidpointConvergenceOrder) {
// aligned: CoM at joint origin
static constexpr char xml1[] = R"(
<mujoco>
<option integrator="implicitfast">
<flag gravity="disable"/>
</option>
<worldbody>
<body>
<freejoint/>
<geom type="box" size=".1 .2 .3" mass="1" euler="10 20 30"/>
</body>
</worldbody>
</mujoco>
)";
// non-aligned: CoM offset from joint origin
static constexpr char xml2[] = R"(
<mujoco>
<option integrator="implicitfast">
<flag gravity="disable"/>
</option>
<worldbody>
<body>
<freejoint/>
<geom type="box" size=".1 .2 .3" mass="1" euler="10 20 30"
pos=".05 .03 .02"/>
</body>
</worldbody>
</mujoco>
)";
int xml_idx = 1;
for (auto xml : {xml1, xml2}) {
SCOPED_TRACE(testing::Message() << "XML case " << xml_idx++);
char error[1024];
mjModel* model = LoadModelFromString(xml, error, sizeof(error));
ASSERT_THAT(model, NotNull()) << error;
mjtNum T = 1.0;
mjtNum h_coarse = 0.02;
mjtNum quat_coarse[4], quat_fine[4], quat_ref[4];
auto run = [&](mjtNum h, mjtNum quat_out[4]) {
model->opt.timestep = h;
mjData* data = mj_makeData(model);
data->qvel[3] = 1.0;
data->qvel[4] = 2.0;
data->qvel[5] = 3.0;
int nstep = (int)(T / h + 0.5);
for (int i = 0; i < nstep; i++) {
mj_step(model, data);
}
mju_copy4(quat_out, data->qpos + 3);
mj_deleteData(data);
};
run(h_coarse, quat_coarse);
run(h_coarse / 2, quat_fine);
run(h_coarse / 16, quat_ref);
// quaternion distance: ||quat - quat_ref|| (handles sign ambiguity)
auto quat_dist = [](const mjtNum a[4], const mjtNum b[4]) -> mjtNum {
mjtNum pos = 0, neg = 0;
for (int i = 0; i < 4; i++) {
pos += (a[i] - b[i]) * (a[i] - b[i]);
neg += (a[i] + b[i]) * (a[i] + b[i]);
}
return mju_sqrt(mju_min(pos, neg));
};
mjtNum err_coarse = quat_dist(quat_coarse, quat_ref);
mjtNum err_fine = quat_dist(quat_fine, quat_ref);
// second-order: error ratio should be ~4 when halving timestep
mjtNum ratio = err_coarse / err_fine;
EXPECT_GT(ratio, 3.5);
EXPECT_LT(ratio, 4.5);
mj_deleteModel(model);
}
}
// verify that Newton iteration in mj_midpoint converges quickly (aligned case)
TEST_F(ImplicitIntegratorTest, MidpointNewtonConvergence) {
// inertia ratios: symmetric, mildly asymmetric, extremely asymmetric
mjtNum inertias[][3] = {
{1.0, 1.0, 1.0},
{1.0, 2.0, 3.0},
{0.01, 1.0, 100.0},
{1.0, 1.0, 1000.0},
};
mjtNum timesteps[] = {0.001, 0.01, 0.1};
mjtNum velocities[][3] = {
{1.0, 2.0, 3.0},
{100.0, 0.0, 0.0},
{10.0, 10.0, 10.0},
{0.01, 0.01, 100.0},
};
mjtNum q_identity[4] = {1, 0, 0, 0};
mjtNum torques[][3] = {
{0, 0, 0},
{10.0, 20.0, 30.0},
{100.0, 0.0, 0.0},
{0.0, 0.0, 100.0},
};
int max_iter = 0;
int total_iter = 0;
int ncases = 0;
for (auto& I : inertias) {
for (mjtNum h : timesteps) {
for (auto& w : velocities) {
for (auto& tau : torques) {
mjtNum vel[6] = {0, 0, 0, w[0], w[1], w[2]};
mjtNum tau_ext[6] = {0, 0, 0, tau[0], tau[1], tau[2]};
mjtNum v_new[6];
mjtNum ipos[3] = {0, 0, 0};
int niter = mj_midpoint(1.0, I, ipos, q_identity, q_identity, vel,
tau_ext, NULL, h, v_new);
EXPECT_LT(niter, 10)
<< "Failed for I=(" << I[0] << "," << I[1] << "," << I[2] << ")"
<< " h=" << h
<< " w=(" << w[0] << "," << w[1] << "," << w[2] << ")"
<< " tau=(" << tau[0] << "," << tau[1] << "," << tau[2] << ")";
max_iter = std::max(max_iter, niter);
total_iter += niter;
ncases++;
}
}
}
}
EXPECT_LE(max_iter, 4);
EXPECT_LT((mjtNum)total_iter / ncases, 2.0);
}
// verify that Newton iteration in mj_midpoint converges quickly (non-aligned)
TEST_F(ImplicitIntegratorTest, MidpointFullNewtonConvergence) {
mjtNum masses[] = {0.1, 1.0, 10.0};
mjtNum inertias[][3] = {
{1.0, 1.0, 1.0},
{1.0, 2.0, 3.0},
{0.01, 1.0, 100.0},
};
mjtNum offsets[][3] = {
{0.1, 0.0, 0.0},
{0.05, 0.03, 0.02},
{0.0, 0.0, 0.5},
};
mjtNum timesteps[] = {0.001, 0.01, 0.1};
mjtNum velocities[][6] = {
{1.0, 0.0, 0.0, 1.0, 2.0, 3.0},
{0.0, 0.0, 0.0, 10.0, 10.0, 10.0},
{5.0, 5.0, 5.0, 0.01, 0.01, 100.0},
};
mjtNum q_identity[4] = {1, 0, 0, 0};
mjtNum forces[][6] = {
{0, 0, 0, 0, 0, 0},
{10.0, 20.0, 30.0, 1.0, 2.0, 3.0},
};
int max_iter = 0;
int total_iter = 0;
int ncases = 0;
for (mjtNum mass : masses) {
for (auto& I : inertias) {
for (auto& r : offsets) {
for (mjtNum h : timesteps) {
for (auto& vel : velocities) {
for (auto& frc : forces) {
mjtNum v_new[6];
int niter = mj_midpoint(mass, I, r, q_identity, q_identity,
vel, frc, NULL, h, v_new);
EXPECT_LT(niter, 10)
<< "Failed for mass=" << mass
<< " I=(" << I[0] << "," << I[1] << "," << I[2] << ")"
<< " r=(" << r[0] << "," << r[1] << "," << r[2] << ")"
<< " h=" << h;
max_iter = std::max(max_iter, niter);
total_iter += niter;
ncases++;
}
}
}
}
}
}
EXPECT_LE(max_iter, 6);
EXPECT_LT((mjtNum)total_iter / ncases, 3.0);
}
TEST_F(ForwardTest, ControlClamping) {
static constexpr char xml[] = R"(
<mujoco>
test/engine/engine_inverse_test.cc
+8 -9
View File
@@ -73,9 +73,16 @@ TEST_F(InverseTest, DiscreteInverseMatch) {
mjtNum* qvel_next = (mjtNum*)mju_malloc(nv * sizeof(mjtNum));
mjtNum* qacc_fd = (mjtNum*)mju_malloc(nv * sizeof(mjtNum));
for (auto integrator : {mjINT_EULER, mjINT_IMPLICIT, mjINT_IMPLICITFAST}) {
for (auto integrator : {mjINT_EULER, mjINT_IMPLICIT}) {
model->opt.integrator = integrator;
for (bool invdiscrete : {false, true}) {
// set/unset mjENBL_INVDISCRETE flag (affects both forward and inverse)
if (invdiscrete) {
model->opt.enableflags |= mjENBL_INVDISCRETE;
} else {
model->opt.enableflags &= ~mjENBL_INVDISCRETE;
}
// simulate
mj_resetData(model, data);
for (int i = 0; i < kSteps; ++i) {
@@ -98,17 +105,9 @@ TEST_F(InverseTest, DiscreteInverseMatch) {
mj_forward(model, data);
mju_copy(data->qacc, qacc_fd, nv);
// set/unset mjENBL_INVDISCRETE flag
if (invdiscrete) {
model->opt.enableflags |= mjENBL_INVDISCRETE;
} else {
model->opt.enableflags &= ~mjENBL_INVDISCRETE;
}
// call built-in testing function
mj_compareFwdInv(model, data);
// depending on mjENBL_INVDISCRETE flag, expect mismatch to be small/large
if (invdiscrete) {
mjtNum epsilon = MjTol(1e-9, 0.05);
EXPECT_LT(data->solver_fwdinv[0], epsilon);
test/engine/engine_sleep_test.cc
+53
View File
@@ -533,6 +533,59 @@ TEST_F(SleepTest, Equality) {
mj_deleteModel(m);
}
// Test that the midpoint integrator doesn't break the sleep qvel=0 invariant.
// A standalone free body (eligible for midpoint) with high viscosity should
// eventually go to sleep, and after sleeping, qvel/qacc must be exactly zero.
TEST_F(SleepTest, MidpointSleepZeroVelocity) {
static constexpr char xml[] = R"(
<mujoco>
<option integrator="implicitfast" viscosity="10"
sleep_tolerance="0.01">
<flag sleep="enable" gravity="disable" constraint="disable"
contact="disable"/>
</option>
<worldbody>
<body>
<freejoint/>
<geom type="box" size=".1 .2 .3" mass="1" euler="10 20 30"
pos=".03 .02 .01"/>
</body>
</worldbody>
</mujoco>
)";
char error[1024];
mjModel* m = LoadModelFromString(xml, error, sizeof(error));
ASSERT_THAT(m, NotNull()) << error;
mjData* d = mj_makeData(m);
// give initial velocity (both translational and angular)
d->qvel[0] = 0.5;
d->qvel[1] = 0.5;
d->qvel[2] = 0.5;
d->qvel[3] = 1.0;
d->qvel[4] = 2.0;
d->qvel[5] = 3.0;
// step until body goes to sleep
for (int step = 0; step < 1000; step++) {
mj_step(m, d);
if (d->ntree_awake == 0) break;
}
// body should have gone to sleep
ASSERT_EQ(d->ntree_awake, 0) << "body did not go to sleep";
// qvel and qacc must be exactly zero for sleeping body
for (int i = 0; i < 6; i++) {
EXPECT_EQ(d->qvel[i], 0.0) << "qvel[" << i << "] not zero after sleep";
EXPECT_EQ(d->qacc[i], 0.0) << "qacc[" << i << "] not zero after sleep";
}
mj_deleteData(d);
mj_deleteModel(m);
}
static const char* const kInitIslandFailModel =
"engine/testdata/sleep/init_island_fail.xml";