diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst index 8cf40b72..a8ebff10 100644 --- a/doc/XMLreference.rst +++ b/doc/XMLreference.rst @@ -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` other than ``RK4``. Recall from the - :ref:`numerical integration` 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`. + + Enable discrete-time inverse dynamics + This flag **enables** discrete-time inverse dynamics with :ref:`mj_inverse` for all + :ref:`integrators` other than ``RK4``. Recall from the :ref:`numerical + integration` 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` 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: diff --git a/doc/changelog.rst b/doc/changelog.rst index e8dbb0ed..d664df1c 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -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`. 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` for standalone free bodies in ``implicit`` and ``implicitfast`` + :ref:`integrators`. 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` 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`. diff --git a/doc/computation/index.rst b/doc/computation/index.rst index 67b1623d..74421918 100644 --- a/doc/computation/index.rst +++ b/doc/computation/index.rst @@ -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 + `__: + + .. 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 + `__ 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 `__ 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` 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` 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` 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 diff --git a/mjx/mujoco/mjx/_src/forward_test.py b/mjx/mujoco/mjx/_src/forward_test.py index ba200b72..751aeb70 100644 --- a/mjx/mujoco/mjx/_src/forward_test.py +++ b/mjx/mujoco/mjx/_src/forward_test.py @@ -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') diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index e8e08f00..b0d15b16 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -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); diff --git a/src/engine/engine_forward.h b/src/engine/engine_forward.h index 1735ca6c..0a3fd5e0 100644 --- a/src/engine/engine_forward.h +++ b/src/engine/engine_forward.h @@ -18,6 +18,7 @@ #include #include #include +#include #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 ----------------------------------------------- diff --git a/src/engine/engine_util_solve.c b/src/engine/engine_util_solve.c index 426f214e..259f7c59 100644 --- a/src/engine/engine_util_solve.c +++ b/src/engine/engine_util_solve.c @@ -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 diff --git a/src/engine/engine_util_solve.h b/src/engine/engine_util_solve.h index 523bc04a..a50ab4ee 100644 --- a/src/engine/engine_util_solve.h +++ b/src/engine/engine_util_solve.h @@ -17,6 +17,7 @@ #include #include +#include #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]); diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index 18ae65b3..e5978c9b 100644 --- a/test/engine/engine_derivative_test.cc +++ b/test/engine/engine_derivative_test.cc @@ -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(duration / dt_large); + int substeps = static_cast(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 diff --git a/test/engine/engine_forward_test.cc b/test/engine/engine_forward_test.cc index fce6cf68..4b0caf9c 100644 --- a/test/engine/engine_forward_test.cc +++ b/test/engine/engine_forward_test.cc @@ -17,6 +17,7 @@ #include "src/engine/engine_forward.h" #include "src/engine/engine_derivative.h" +#include #include #include #include @@ -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"( + + + + + + + + + + )"; + + // auto-aligned: CoM at joint origin + static constexpr char xml2[] = R"( + + + + + + + + + + )"; + + // non-aligned: CoM offset from joint origin + static constexpr char xml3[] = R"( + + + + + + + + + + )"; + 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"( + + + + + + + + + + )"; + + // non-aligned: CoM offset from joint origin + static constexpr char xml2[] = R"( + + + + + + + + + + )"; + + 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"( diff --git a/test/engine/engine_inverse_test.cc b/test/engine/engine_inverse_test.cc index 2bfc2d6d..988e73e6 100644 --- a/test/engine/engine_inverse_test.cc +++ b/test/engine/engine_inverse_test.cc @@ -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); diff --git a/test/engine/engine_sleep_test.cc b/test/engine/engine_sleep_test.cc index 88067eff..1245dc83 100644 --- a/test/engine/engine_sleep_test.cc +++ b/test/engine/engine_sleep_test.cc @@ -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"( + + + + + + + + + + )"; + + 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";