Files
Mujoco_WASM/src/engine/engine_forward.c
T
Kyle Bayes b935d4153c Add new mju_threadpool API function, and delete old threading API.
PiperOrigin-RevId: 922838541
Change-Id: Id9f7e0fb298ffde61fcc49a802dc78971858ce51
2026-05-28 10:11:44 -07:00

2071 lines
63 KiB
C

// Copyright 2021 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "engine/engine_forward.h"
#include <stddef.h>
#include <stdio.h>
#include <mujoco/mjdata.h>
#include <mujoco/mjmacro.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjsan.h> // IWYU pragma: keep
#include <mujoco/mjplugin.h>
#include "engine/engine_callback.h"
#include "engine/engine_collision_driver.h"
#include "engine/engine_core_constraint.h"
#include "engine/engine_core_smooth.h"
#include "engine/engine_core_util.h"
#include "engine/engine_derivative.h"
#include "engine/engine_inverse.h"
#include "engine/engine_island.h"
#include "engine/engine_macro.h"
#include "engine/engine_memory.h"
#include "engine/engine_passive.h"
#include "engine/engine_plugin.h"
#include "engine/engine_sensor.h"
#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"
#include "engine/engine_util_solve.h"
#include "engine/engine_util_sparse.h"
#include "engine/engine_thread.h"
//--------------------------- check values ---------------------------------------------------------
// check positions, reset if bad
void mj_checkPos(const mjModel* m, mjData* d) {
int nq = m->nq;
const mjtNum* qpos = d->qpos;
for (int i=0; i < nq; i++) {
if (mju_isBad(qpos[i])) {
mj_warning(d, mjWARN_BADQPOS, i);
if (!mjDISABLED(mjDSBL_AUTORESET)) {
mj_resetData(m, d);
}
d->warning[mjWARN_BADQPOS].number++;
d->warning[mjWARN_BADQPOS].lastinfo = i;
return;
}
}
}
// check velocities, reset if bad
void mj_checkVel(const mjModel* m, mjData* d) {
int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < m->nv;
int nv = sleep_filter ? d->nv_awake : m->nv;
for (int j=0; j < nv; j++) {
int i = sleep_filter ? d->dof_awake_ind[j] : j;
if (mju_isBad(d->qvel[i])) {
mj_warning(d, mjWARN_BADQVEL, i);
if (!mjDISABLED(mjDSBL_AUTORESET)) {
mj_resetData(m, d);
}
d->warning[mjWARN_BADQVEL].number++;
d->warning[mjWARN_BADQVEL].lastinfo = i;
return;
}
}
}
// check accelerations, reset if bad
void mj_checkAcc(const mjModel* m, mjData* d) {
int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < m->nv;
int nv = sleep_filter ? d->nv_awake : m->nv;
for (int j=0; j < nv; j++) {
int i = sleep_filter ? d->dof_awake_ind[j] : j;
if (mju_isBad(d->qacc[i])) {
mj_warning(d, mjWARN_BADQACC, i);
if (!mjDISABLED(mjDSBL_AUTORESET)) {
mj_resetData(m, d);
}
d->warning[mjWARN_BADQACC].number++;
d->warning[mjWARN_BADQACC].lastinfo = i;
if (!mjDISABLED(mjDSBL_AUTORESET)) {
mj_forward(m, d);
}
return;
}
}
}
//-------------------------- solver components -----------------------------------------------------
// kinematics-related computations
void mj_fwdKinematics(const mjModel* m, mjData* d) {
mj_kinematics(m, d);
mj_comPos(m, d);
mj_camlight(m, d);
mj_flex(m, d);
mj_tendon(m, d);
if (mj_wakeTendon(m, d)) {
mj_updateSleep(m, d);
}
}
// position-dependent computations
void mj_fwdPosition(const mjModel* m, mjData* d) {
TM_START1;
// clear position-dependent flags for lazy evaluation
d->flg_energypos = 0;
TM_START;
mj_fwdKinematics(m, d);
TM_END(mjTIMER_POS_KINEMATICS);
// inertia, timed internally (POS_INERTIA)
mj_makeM(m, d);
mj_factorM(m, d);
// collision, timed internally (POS_COLLISION)
mj_collision(m, d);
if (mj_wakeCollision(m, d)) {
mj_updateSleep(m, d);
mj_collision(m, d);
}
if (mj_wakeEquality(m, d)) {
mj_updateSleep(m, d);
}
TM_RESTART;
mj_makeConstraint(m, d);
mj_island(m, d);
TM_END(mjTIMER_POS_MAKE);
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);
}
// velocity-dependent computations
void mj_fwdVelocity(const mjModel* m, mjData* d) {
TM_START;
// clear velocity-dependent flags for lazy evaluation
d->flg_subtreevel = 0;
d->flg_energyvel = 0;
// flexedge velocity: skip interp and rigid flexes (edge Jacobians are zero)
mju_zero(d->flexedge_velocity, m->nflexedge);
for (int f = 0; f < m->nflex; f++) {
if (m->flex_rigid[f] || m->flex_interp[f]) continue;
int adr = m->flex_edgeadr[f];
int num = m->flex_edgenum[f];
mju_mulMatVecSparse(d->flexedge_velocity + adr, d->flexedge_J, d->qvel, num,
m->flexedge_J_rownnz + adr, m->flexedge_J_rowadr + adr,
m->flexedge_J_colind, NULL);
}
// tendon velocity: always sparse
mju_mulMatVecSparse(d->ten_velocity, d->ten_J, d->qvel, m->ntendon,
m->ten_J_rownnz, m->ten_J_rowadr, m->ten_J_colind, NULL);
// actuator velocity: always sparse
if (!mjDISABLED(mjDSBL_ACTUATION)) {
mju_mulMatVecSparse(d->actuator_velocity, d->actuator_moment, d->qvel, m->nu,
d->moment_rownnz, d->moment_rowadr, d->moment_colind, NULL);
} else {
mju_zero(d->actuator_velocity, m->nu);
}
// com-based velocities, passive forces, constraint references
mj_comVel(m, d);
mj_passive(m, d);
mj_referenceConstraint(m, d);
// compute qfrc_bias with abbreviated RNE (without acceleration)
mj_rne(m, d, 0, d->qfrc_bias);
// add bias force due to tendon armature
mj_tendonBias(m, d, d->qfrc_bias);
TM_END(mjTIMER_VELOCITY);
}
// helper for DC motor: computes control voltage from PID state
static mjtNum dcmotorVoltage(mjtNum ctrl, mjtNum length, mjtNum velocity,
mjtNum x_I, const mjtNum* gainprm) {
int input_mode = (int)gainprm[8];
mjtNum Vmax = gainprm[7];
mjtNum voltage;
// get voltage
if (input_mode > 0) {
mjtNum kp = gainprm[4]; // proportional gain
mjtNum ki = gainprm[5]; // integral gain
mjtNum kd = gainprm[6]; // derivative gain
if (input_mode == 1) {
// position mode
voltage = kp * (ctrl - length) + ki * x_I - kd * velocity;
} else {
// velocity mode
voltage = kp * (ctrl - velocity) + ki * (x_I - length);
}
} else {
voltage = ctrl;
}
// clip voltage
if (Vmax > 0) voltage = mju_clip(voltage, -Vmax, Vmax);
return voltage;
}
// clamp vector to range
static void clampVec(mjtNum* vec, const mjtNum* range, const mjtBool* limited, int n,
const int* index) {
for (int i=0; i < n; i++) {
int j = index ? index[i] : i;
if (limited[i]) {
vec[j] = mju_clip(vec[j], range[2*i], range[2*i + 1]);
}
}
}
// (qpos, qvel, ctrl, act) => (qfrc_actuator, actuator_force, act_dot)
void mj_fwdActuation(const mjModel* m, mjData* d) {
TM_START;
int nv = m->nv, nu = m->nu, ntendon = m->ntendon;
mjtNum gain, bias, tau;
mjtNum *force = d->actuator_force;
// clear actuator_force
mju_zero(force, nu);
int sleep_filter = mjENABLED(mjENBL_SLEEP);
// disabled or no actuation: return
if (nu == 0 || mjDISABLED(mjDSBL_ACTUATION)) {
mju_zero(d->qfrc_actuator, nv);
return;
}
// any tendon transmission targets with force limits
int tendon_frclimited = 0;
// local copy of ctrl
mj_markStack(d);
mjtNum *ctrl = mjSTACKALLOC(d, nu, mjtNum);
// read from ctrl or history buffer for delayed actuators
for (int i = 0; i < nu; i++) {
int interp = m->actuator_history[2*i+1];
ctrl[i] = m->actuator_delay[i] ? mj_readCtrl(m, d, i, d->time, interp) : d->ctrl[i];
}
// clamp local copy
if (!mjDISABLED(mjDSBL_CLAMPCTRL)) {
clampVec(ctrl, m->actuator_ctrlrange, m->actuator_ctrllimited, nu, NULL);
}
// check controls, set all to 0 if any are bad
for (int i=0; i < nu; i++) {
if (mju_isBad(ctrl[i])) {
mj_warning(d, mjWARN_BADCTRL, i);
mju_zero(ctrl, nu);
break;
}
}
// act_dot for stateful actuators
for (int i=0; i < nu; i++) {
if (sleep_filter && mj_sleepState(m, d, mjOBJ_ACTUATOR, i) == mjS_ASLEEP) {
continue;
}
int act_first = m->actuator_actadr[i];
if (act_first < 0) {
continue;
}
// zero act_dot for actuator plugins
int actnum = m->actuator_actnum[i];
if (actnum) {
mju_zero(d->act_dot + act_first, actnum);
}
// extract info
const mjtNum* dynprm = m->actuator_dynprm + i*mjNDYN;
mjtDyn dyntype = m->actuator_dyntype[i];
// index into the last element in act. For most actuators it's also the
// first element, but actuator plugins might store their own state in act
int act_last = act_first + actnum - 1;
// compute act_dot according to dynamics type
switch (dyntype) {
case mjDYN_INTEGRATOR: // simple integrator
d->act_dot[act_last] = ctrl[i];
break;
case mjDYN_FILTER: // linear filter: dynprm = tau
case mjDYN_FILTEREXACT:
tau = mju_max(mjMINVAL, dynprm[0]);
d->act_dot[act_last] = (ctrl[i] - d->act[act_last]) / tau;
break;
case mjDYN_MUSCLE: // muscle model: dynprm = (tau_act, tau_deact)
d->act_dot[act_last] = mju_muscleDynamics(ctrl[i], d->act[act_last], dynprm);
break;
case mjDYN_DCMOTOR: { // DC motor: up to 5 optional states
const mjtNum* gainprm = m->actuator_gainprm + mjNGAIN*i;
// verify allocated state size matches parameters; SHOULD NOT OCCUR
if (mj_dcmotorSlots(dynprm, gainprm).num_slots != actnum) {
mjERROR("inconsistent state array dimension in DC motor (actuator %d)", i);
}
int adr = act_first;
mjtNum velocity = d->actuator_velocity[i];
mjtNum R = gainprm[0]; // resistance
mjtNum K = gainprm[1]; // motor constant
mjtNum ki = gainprm[5]; // integral gain
mjtNum te = dynprm[0]; // electrical time constant
// slot order: slew, integral, temperature, bristle, current
// controller state: slew rate limiting
mjtNum slew_s = dynprm[7]; // slew rate limit
if (slew_s > 0) {
mjtNum u_prev = d->act[adr];
mjtNum slew = slew_s * m->opt.timestep;
mjtNum u_eff = mju_clip(ctrl[i], u_prev - slew, u_prev + slew);
d->act_dot[adr] = (u_eff - u_prev) / m->opt.timestep;
ctrl[i] = u_eff;
adr++;
}
// controller state: integral state
mjtNum x_I = 0;
if (ki > 0) {
x_I = d->act[adr];
int input_mode = (int)gainprm[8];
mjtNum Imax = dynprm[8]; // integral clamp
mjtNum act_dot = ctrl[i]; // default raw accumulator for voltage and velocity modes
// position mode
if (input_mode == 1) {
act_dot = ctrl[i] - d->actuator_length[i];
}
// clamp act_dot based on integral state
if (Imax > 0) {
if (x_I >= Imax) {
act_dot = mju_min(act_dot, 0);
} else if (x_I <= -Imax) {
act_dot = mju_max(act_dot, 0);
}
}
d->act_dot[adr] = act_dot;
adr++;
}
// compute physical voltage to feed into current and temperature equations
mjtNum V = dcmotorVoltage(ctrl[i], d->actuator_length[i], velocity, x_I, gainprm);
// temperature: dT/dt = (R*i^2 - T/RT) / C, where T = delta above ambient
mjtNum RT = dynprm[2]; // thermal resistance
if (RT > 0) {
mjtNum C = dynprm[3]; // thermal capacitance
mjtNum Ta = dynprm[4]; // ambient temperature
mjtNum alpha = gainprm[2]; // temperature coefficient
mjtNum T0 = gainprm[3]; // reference temperature
mjtNum T = d->act[adr]; // temperature rise above ambient
R *= 1 + alpha * (T + Ta - T0);
// get current: from act_last if stateful, from (V - K*omega)/R if stateless
mjtNum current = (te > 0) ? d->act[act_last] : (V - K * velocity) / R;
d->act_dot[adr] = (R*current*current - T / RT) / C;
adr++;
}
// LuGre bristle state: dz/dt = v - sigma0 * |v| / g(v) * z
mjtNum sigma0 = dynprm[5]; // bristle stiffness
if (sigma0 > 0) {
const mjtNum* biasprm = m->actuator_biasprm + mjNBIAS*i;
mjtNum F_C = biasprm[3]; // Coulomb friction
mjtNum F_S = biasprm[4]; // static friction
mjtNum v_S = biasprm[5]; // Stribeck velocity
mjtNum z = d->act[adr]; // bristle state
mjtNum g = mj_lugreStribeck(velocity, F_C, F_S, v_S);
mjtNum a = -sigma0 * mju_abs(velocity) / mju_max(mjMINVAL, g);
d->act_dot[adr] = a * z + velocity;
adr++;
}
// current state: di/dt = (V/R - K/R*omega - i) / te
if (te > 0) {
mjtNum dimax = dynprm[1]; // current rate limit (di/dt)_max
mjtNum i_dot = (V/R - K/R*velocity - d->act[act_last]) / te;
if (dimax > 0) {
i_dot = mju_clip(i_dot, -dimax, dimax);
}
d->act_dot[act_last] = i_dot;
}
break;
}
default: // user dynamics
if (mjcb_act_dyn) {
if (actnum == 1) {
// scalar activation dynamics, get act_dot
d->act_dot[act_last] = mjcb_act_dyn(m, d, i);
} else {
// higher-order dynamics, mjcb_act_dyn writes into act_dot directly
mjcb_act_dyn(m, d, i);
}
}
}
}
// get act_dot from actuator plugins
if (m->nplugin) {
const int nslot = mjp_pluginCount();
for (int i=0; i < m->nplugin; i++) {
const int slot = m->plugin[i];
const mjpPlugin* plugin = mjp_getPluginAtSlotUnsafe(slot, nslot);
if (!plugin) {
mjERROR("invalid plugin slot: %d", slot);
}
if (plugin->capabilityflags & mjPLUGIN_ACTUATOR) {
if (plugin->actuator_act_dot) {
plugin->actuator_act_dot(m, d, i);
}
}
}
}
// force = gain .* [ctrl/act] + bias
for (int i=0; i < nu; i++) {
// skip if sleeping
if (sleep_filter && mj_sleepState(m, d, mjOBJ_ACTUATOR, i) == mjS_ASLEEP) {
continue;
}
// skip if disabled
if (mj_actuatorDisabled(m, i)) {
continue;
}
// skip actuator plugins -- these are handled after builtin actuator types
if (m->actuator_plugin[i] >= 0) {
continue;
}
// check for tendon transmission with force limits
if (ntendon && !tendon_frclimited && m->actuator_trntype[i] == mjTRN_TENDON) {
tendon_frclimited = m->tendon_actfrclimited[m->actuator_trnid[2*i]];
}
// extract info
const mjtNum* dynprm = m->actuator_dynprm + mjNDYN*i;
const mjtNum* gainprm = m->actuator_gainprm + mjNGAIN*i;
mjtGain gaintype = m->actuator_gaintype[i];
int actnum = m->actuator_actnum[i];
// handle according to gain type
switch (gaintype) {
case mjGAIN_FIXED: // fixed gain: prm = gain
gain = gainprm[0];
break;
case mjGAIN_AFFINE: // affine: prm = [const, kp, kv]
gain = gainprm[0] + gainprm[1]*d->actuator_length[i] + gainprm[2]*d->actuator_velocity[i];
break;
case mjGAIN_MUSCLE: // muscle gain
gain = mju_muscleGain(d->actuator_length[i],
d->actuator_velocity[i],
m->actuator_lengthrange+2*i,
m->actuator_acc0[i],
gainprm);
break;
case mjGAIN_DCMOTOR: { // DC motor: gain = K or K/R
mjtNum R = gainprm[0]; // resistance
mjtNum K = gainprm[1]; // motor constant
mjDCMotorSlots slots = mj_dcmotorSlots(dynprm, gainprm);
// verify allocated state size matches parameters; SHOULD NOT OCCUR
if (slots.num_slots != actnum) {
mjERROR("inconsistent state array dimension in DC motor (actuator %d)", i);
}
int adr = m->actuator_actadr[i];
// adjust R for temperature if enabled
if (slots.temperature >= 0) {
mjtNum T = d->act[adr + slots.temperature];
mjtNum alpha = gainprm[2]; // temperature coefficient
mjtNum T0 = gainprm[3]; // reference temperature
mjtNum Ta = dynprm[4]; // ambient temperature
R *= 1 + alpha * (T + Ta - T0);
}
// stateful current: gain = K, force = K * act[last] (generic path)
// stateless: gain = K/R, force = K/R * ctrl (condition below)
gain = (dynprm[0] > 0) ? K : K / mju_max(mjMINVAL, R);
// controller: compute voltage, override ctrl[i] for force computation
if ((int)gainprm[8] > 0) {
mjtNum x_I = (slots.integral >= 0) ? d->act[adr + slots.integral] : 0;
ctrl[i] = dcmotorVoltage(ctrl[i], d->actuator_length[i],
d->actuator_velocity[i], x_I, gainprm);
}
break;
}
default: // user gain
if (mjcb_act_gain) {
gain = mjcb_act_gain(m, d, i);
} else {
gain = 1;
}
}
// set force = gain .* [ctrl/act]
// DC motor without current state: use ctrl even if other activations exist
int dcmotor_no_current = (gaintype == mjGAIN_DCMOTOR && dynprm[0] <= 0);
if (actnum == 0 || dcmotor_no_current) {
force[i] = gain * ctrl[i];
} else {
// use last activation variable associated with actuator i
int act_adr = m->actuator_actadr[i] + actnum - 1;
mjtNum act;
if (m->actuator_actearly[i]) {
act = mj_nextActivation(m, d, i, act_adr, d->act_dot[act_adr]);
} else {
act = d->act[act_adr];
}
force[i] = gain * act;
}
// extract bias info
const mjtNum* biasprm = m->actuator_biasprm + mjNBIAS*i;
mjtBias biastype = m->actuator_biastype[i];
// handle according to bias type
switch (biastype) {
case mjBIAS_NONE: // none
bias = 0.0;
break;
case mjBIAS_AFFINE: // affine: biasprm = [const, kp, kv]
bias = biasprm[0] + biasprm[1]*d->actuator_length[i] + biasprm[2]*d->actuator_velocity[i];
break;
case mjBIAS_MUSCLE: // muscle passive force
bias = mju_muscleBias(d->actuator_length[i],
m->actuator_lengthrange+2*i,
m->actuator_acc0[i],
biasprm);
break;
case mjBIAS_DCMOTOR: { // DC motor: back-EMF only (current-limited)
bias = 0;
// back-EMF (stateless only; for stateful current it's in the ODE)
mjtNum te = m->actuator_dynprm[mjNDYN*i]; // electrical time constant
if (te <= 0) {
mjtNum K = gainprm[1]; // motor constant
bias -= gain * K * d->actuator_velocity[i];
}
break;
}
default: // user bias
if (mjcb_act_bias) {
bias = mjcb_act_bias(m, d, i);
} else {
bias = 0;
}
}
// add bias
force[i] += bias;
}
// handle actuator plugins
if (m->nplugin) {
const int nslot = mjp_pluginCount();
for (int i=0; i < m->nplugin; i++) {
const int slot = m->plugin[i];
const mjpPlugin* plugin = mjp_getPluginAtSlotUnsafe(slot, nslot);
if (!plugin) {
mjERROR("invalid plugin slot: %d", slot);
}
if (plugin->capabilityflags & mjPLUGIN_ACTUATOR) {
if (!plugin->compute) {
mjERROR("`compute` is a null function pointer for plugin at slot %d", slot);
}
plugin->compute(m, d, i, mjPLUGIN_ACTUATOR);
}
}
}
// clamp tendon total actuator force
if (tendon_frclimited) {
// compute total force for each tendon
mjtNum* tendon_total_force = mjSTACKALLOC(d, ntendon, mjtNum);
mju_zero(tendon_total_force, ntendon);
for (int i=0; i < nu; i++) {
if (m->actuator_trntype[i] == mjTRN_TENDON) {
int tendon_id = m->actuator_trnid[2*i];
if (m->tendon_actfrclimited[tendon_id]) {
tendon_total_force[tendon_id] += force[i];
}
}
}
// scale tendon actuator forces if limited and outside range
for (int i=0; i < nu; i++) {
if (m->actuator_trntype[i] != mjTRN_TENDON) {
continue;
}
int tendon_id = m->actuator_trnid[2*i];
mjtNum tendon_force = tendon_total_force[tendon_id];
if (m->tendon_actfrclimited[tendon_id] && tendon_force) {
const mjtNum* range = m->tendon_actfrcrange + 2 * tendon_id;
if (tendon_force < range[0]) {
force[i] *= range[0] / tendon_force;
} else if (tendon_force > range[1]) {
force[i] *= range[1] / tendon_force;
}
}
}
}
// clamp actuator_force
clampVec(force, m->actuator_forcerange, m->actuator_forcelimited, nu, NULL);
// add DC motor mechanical forces (not subject to current limits)
for (int i=0; i < nu; i++) {
if (m->actuator_biastype[i] != mjBIAS_DCMOTOR) {
continue;
}
if (sleep_filter && mj_sleepState(m, d, mjOBJ_ACTUATOR, i) == mjS_ASLEEP) {
continue;
}
if (mj_actuatorDisabled(m, i) || m->actuator_plugin[i] >= 0) {
continue;
}
const mjtNum* biasprm = m->actuator_biasprm + mjNBIAS*i;
const mjtNum* dynprm = m->actuator_dynprm + mjNDYN*i;
// cogging torque
mjtNum A = biasprm[0];
if (A != 0) {
mjtNum Np = biasprm[1];
mjtNum phi = biasprm[2];
force[i] += A * mju_sin(Np*d->actuator_length[i] + phi);
}
// LuGre friction
mjtNum sigma0 = dynprm[5];
if (sigma0 > 0) {
mjtNum sigma1 = dynprm[6];
mjDCMotorSlots slots = mj_dcmotorSlots(dynprm, m->actuator_gainprm + mjNGAIN*i);
int adr = m->actuator_actadr[i] + slots.bristle;
mjtNum z = d->act[adr];
mjtNum z_dot = d->act_dot[adr];
force[i] -= sigma0 * z + sigma1 * z_dot;
}
}
// qfrc_actuator = moment' * force
mju_mulMatTVecSparse(d->qfrc_actuator, d->actuator_moment, force, nu, nv,
d->moment_rownnz, d->moment_rowadr, d->moment_colind);
// actuator-level gravity compensation
if (m->ngravcomp && !mjDISABLED(mjDSBL_GRAVITY) && mju_norm3(m->opt.gravity)) {
// number of dofs for each joint type: {mjJNT_FREE, mjJNT_BALL, mjJNT_SLIDE, mjJNT_HINGE}
static const int jnt_dofnum[4] = {6, 3, 1, 1};
int njnt = m->njnt;
for (int i=0; i < njnt; i++) {
// skip if gravcomp added as passive force
if (!m->jnt_actgravcomp[i]) {
continue;
}
// add gravcomp force
int dofnum = jnt_dofnum[m->jnt_type[i]];
int dofadr = m->jnt_dofadr[i];
mju_addTo(d->qfrc_actuator + dofadr, d->qfrc_gravcomp + dofadr, dofnum);
}
}
// clamp qfrc_actuator to joint-level actuator force limits
clampVec(d->qfrc_actuator, m->jnt_actfrcrange, m->jnt_actfrclimited, m->njnt, m->jnt_dofadr);
mj_freeStack(d);
TM_END(mjTIMER_ACTUATION);
}
// add up all non-constraint forces, compute qacc_smooth
void mj_fwdAcceleration(const mjModel* m, mjData* d) {
int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < m->nv;
int nv;
const int* index;
// qfrc_smooth = qfrc_passive - qfrc_bias + qfrc_applied + qfrc_actuator
if (!sleep_filter) {
nv = m->nv;
index = NULL;
mju_sub(d->qfrc_smooth, d->qfrc_passive, d->qfrc_bias, nv);
mju_addTo(d->qfrc_smooth, d->qfrc_applied, nv);
mju_addTo(d->qfrc_smooth, d->qfrc_actuator, nv);
} else {
nv = d->nv_awake;
index = d->dof_awake_ind;
mju_subInd(d->qfrc_smooth, d->qfrc_passive, d->qfrc_bias, index, nv);
mju_addToInd(d->qfrc_smooth, d->qfrc_applied, index, nv);
mju_addToInd(d->qfrc_smooth, d->qfrc_actuator, index, nv);
}
// qfrc_smooth += project(xfrc_applied)
mj_xfrcAccumulate(m, d, d->qfrc_smooth);
// copy for in-place solve: qacc_smooth = qfrc_smooth
if (!sleep_filter) {
mju_copy(d->qacc_smooth, d->qfrc_smooth, nv);
} else {
mju_copyInd(d->qacc_smooth, d->qfrc_smooth, index, nv);
}
// qacc_smooth = M \ qfrc_smooth
mj_solveLD(d->qacc_smooth, d->qLD, d->qLDiagInv, nv, 1,
m->M_rownnz, m->M_rowadr, m->M_colind, index);
}
// warmstart/init solver
static void warmstart(const mjModel* m, mjData* d) {
int nv = m->nv, nefc = d->nefc;
// warmstart with best of (qacc_warmstart, qacc_smooth)
if (!mjDISABLED(mjDSBL_WARMSTART)) {
mj_markStack(d);
mjtNum* jar = mjSTACKALLOC(d, nefc, mjtNum);
// start with qacc = qacc_warmstart
mju_copy(d->qacc, d->qacc_warmstart, nv);
// compute jar(qacc_warmstart)
mj_mulJacVec(m, d, jar, d->qacc_warmstart);
mju_subFrom(jar, d->efc_aref, nefc);
// update constraints, save cost(qacc_warmstart)
mjtNum cost_warmstart;
mj_constraintUpdate(m, d, jar, &cost_warmstart, 0);
// PGS
if (m->opt.solver == mjSOL_PGS) {
// cost(force_warmstart)
mjtNum PGS_warmstart = mju_dot(d->efc_force, d->efc_b, nefc);
mjtNum* ARf = mjSTACKALLOC(d, nefc, mjtNum);
if (mj_isSparse(m))
mju_mulMatVecSparse(ARf, d->efc_AR, d->efc_force, nefc,
d->efc_AR_rownnz, d->efc_AR_rowadr,
d->efc_AR_colind, NULL);
else {
mju_mulMatVec(ARf, d->efc_AR, d->efc_force, nefc, nefc);
}
PGS_warmstart += 0.5*mju_dot(d->efc_force, ARf, nefc);
// use zero if better
if (PGS_warmstart > 0) {
mju_zero(d->efc_force, nefc);
mju_zero(d->qfrc_constraint, nv);
}
}
// non-PGS
else {
// add Gauss to cost(qacc_warmstart)
mjtNum* Ma = mjSTACKALLOC(d, nv, mjtNum);
mj_mulM(m, d, Ma, d->qacc_warmstart);
for (int i=0; i < nv; i++) {
cost_warmstart += 0.5*(Ma[i]-d->qfrc_smooth[i])*(d->qacc_warmstart[i]-d->qacc_smooth[i]);
}
// cost(qacc_smooth)
mjtNum cost_smooth;
mj_constraintUpdate(m, d, d->efc_b, &cost_smooth, 0);
// use qacc_smooth if better
if (cost_warmstart > cost_smooth) {
mju_copy(d->qacc, d->qacc_smooth, nv);
}
}
// have island structure: unconstrained qacc = qacc_smooth
if (d->nisland > 0) {
// loop over unconstrained dofs in map_idof2dof[nidof, nv)
for (int i=d->nidof; i < nv; i++) {
int dof = d->map_idof2dof[i];
d->qacc[dof] = d->qacc_smooth[dof];
}
}
mj_freeStack(d);
}
// coldstart with qacc = qacc_smooth, efc_force = 0
else {
mju_copy(d->qacc, d->qacc_smooth, nv);
mju_zero(d->efc_force, nefc);
}
}
// mju_dispatch callback: solve one island
static void solveIslandTask(const mjModel* m, mjData* d, void* arg, int thread_id, int island) {
if (m->opt.solver == mjSOL_NEWTON) {
mj_solNewton_island(m, d, island, m->opt.iterations);
} else {
mj_solCG_island(m, d, island, m->opt.iterations);
}
}
// compute efc_b, efc_force, qfrc_constraint; update qacc
void mj_fwdConstraint(const mjModel* m, mjData* d) {
TM_START;
int nv = m->nv, nefc = d->nefc, nisland = d->nisland, nidof;
// always clear qfrc_constraint
mju_zero(d->qfrc_constraint, nv);
// no constraints: copy unconstrained acc, clear forces, return
if (!nefc) {
mju_copy(d->qacc, d->qacc_smooth, nv);
mju_zeroInt(d->solver_niter, mjNISLAND);
TM_END(mjTIMER_CONSTRAINT);
return;
}
// compute efc_b = J*qacc_smooth - aref
mj_mulJacVec(m, d, d->efc_b, d->qacc_smooth);
mju_subFrom(d->efc_b, d->efc_aref, nefc);
// check for invalid solver type
if (m->opt.solver != mjSOL_PGS && m->opt.solver != mjSOL_CG && m->opt.solver != mjSOL_NEWTON) {
mjERROR("unknown solver type %d", m->opt.solver);
}
// warmstart solver
warmstart(m, d);
mju_zeroInt(d->solver_niter, mjNISLAND);
// check if islands are supported
int islands_supported = !mjDISABLED(mjDSBL_ISLAND) && nisland > 0;
// run solver over constraint islands
if (islands_supported) {
switch ((mjtSolver) m->opt.solver) {
case mjSOL_PGS:
for (int island=0; island < nisland; island++) {
mj_solPGS_island(m, d, island, m->opt.iterations);
}
break;
case mjSOL_CG:
case mjSOL_NEWTON:
// copy inputs to islands (vel+acc deps, pos-dependent already copied in mj_island)
nidof = d->nidof;
mju_gather(d->ifrc_smooth, d->qfrc_smooth, d->map_idof2dof, nidof);
mju_gather(d->ifrc_constraint, d->qfrc_constraint, d->map_idof2dof, nidof);
mju_gather(d->iacc_smooth, d->qacc_smooth, d->map_idof2dof, nidof);
mju_gather(d->iacc, d->qacc, d->map_idof2dof, nidof);
mju_gather(d->iefc_force, d->efc_force, d->map_iefc2efc, nefc);
mju_gather(d->iefc_aref, d->efc_aref, d->map_iefc2efc, nefc);
mju_dispatch(m, d, solveIslandTask, NULL, nisland);
// copy back solver outputs (scatter dofs since ni <= nv)
mju_scatter(d->qacc, d->iacc, d->map_idof2dof, nidof);
mju_scatter(d->qfrc_constraint, d->ifrc_constraint, d->map_idof2dof, nidof);
mju_gather(d->efc_force, d->iefc_force, d->map_efc2iefc, nefc);
break;
}
// run noslip solver per island if enabled
if (m->opt.noslip_iterations > 0) {
for (int island=0; island < nisland; island++) {
mj_solNoSlip_island(m, d, island, m->opt.noslip_iterations);
}
}
}
// run solver over all constraints (monolithic)
else {
switch ((mjtSolver) m->opt.solver) {
case mjSOL_PGS: // PGS
mj_solPGS(m, d, m->opt.iterations);
break;
case mjSOL_CG: // CG
mj_solCG(m, d, m->opt.iterations);
break;
case mjSOL_NEWTON: // Newton
mj_solNewton(m, d, m->opt.iterations);
break;
}
// run noslip solver if enabled
if (m->opt.noslip_iterations > 0) {
mj_solNoSlip(m, d, m->opt.noslip_iterations);
}
}
// dual solvers: map efc_force to joint space (always monolithic)
if (m->opt.solver == mjSOL_PGS || m->opt.noslip_iterations > 0) {
mj_dualFinish(m, d);
}
TM_END(mjTIMER_CONSTRAINT);
}
//-------------------------- state advancement and integration ------------------------------------
// 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;
// advance history buffers
if (m->nhistory > 0) {
// advance ctrl history buffers
for (int i = 0; i < nu; i++) {
int nsample = m->actuator_history[2*i];
if (nsample == 0) continue;
// get history buffer pointer and insert ctrl at current time
mjtNum* buf = d->history + m->actuator_historyadr[i];
*mju_historyInsert(buf, nsample, /*dim=*/1, d->time) = d->ctrl[i];
}
// advance sensor history buffers
for (int i = 0; i < nsensor; i++) {
int nsample = m->sensor_history[2*i];
if (nsample == 0) continue;
// get history buffer parameters
int dim = m->sensor_dim[i];
mjtNum* buf = d->history + m->sensor_historyadr[i];
mjtNum delay = m->sensor_delay[i];
mjtNum interval = m->sensor_interval[2*i];
if (interval > 0) {
// interval mode: if condition is satisfied, compute; otherwise copy
mjtNum time_prev = buf[0]; // first slot stores previous sensor tick
if (time_prev + interval <= d->time) {
buf[0] += interval; // advance by exact interval (continuous time)
mjtNum* slot = mju_historyInsert(buf, nsample, dim, d->time);
if (delay > 0) {
// have delay, compute sensor
mj_computeSensor(m, d, i, slot);
} else {
// no delay, copy from sensordata (already computed)
mju_copy(slot, d->sensordata + m->sensor_adr[i], dim);
}
}
} else if (delay > 0) {
// delay-only mode: always compute and insert
mjtNum* slot = mju_historyInsert(buf, nsample, dim, d->time);
mj_computeSensor(m, d, i, slot);
} else {
// history-only mode: copy from sensordata (already computed)
mjtNum* slot = mju_historyInsert(buf, nsample, dim, d->time);
mju_copy(slot, d->sensordata + m->sensor_adr[i], dim);
}
}
}
// advance activations
if (m->na && !mjDISABLED(mjDSBL_ACTUATION)) {
for (int i=0; i < nu; i++) {
int actadr = m->actuator_actadr[i];
int actadr_end = actadr + m->actuator_actnum[i];
for (int j=actadr; j < actadr_end; j++) {
// if disabled, set act_dot to 0
d->act[j] = mj_nextActivation(m, d, i, j, mj_actuatorDisabled(m, i) ? 0 : act_dot[j]);
}
}
}
// put islands to sleep according to velocity tolerance
if (mj_sleep(m, d)) {
// if any trees put to sleep (qvel set to 0), recompute all velocity-dependent quantities
mj_forwardSkip(m, d, mjSTAGE_POS, 0);
// update sleep indices
mj_updateSleep(m, d);
}
// advance velocities
int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->ntree_awake < m->ntree;
if (sleep_filter) {
mju_addToSclInd(d->qvel, qacc, d->dof_awake_ind, m->opt.timestep, d->nv_awake);
} else {
mju_addToScl(d->qvel, qacc, m->opt.timestep, m->nv);
}
// advance positions with qvel if given, d->qvel otherwise (semi-implicit)
const int* index = sleep_filter ? d->body_awake_ind : NULL;
int nbody = sleep_filter ? d->nbody_awake : m->nbody;
mj_integratePosInd(m, d->qpos, qvel ? qvel : d->qvel, m->opt.timestep, index, nbody);
// advance time
d->time += m->opt.timestep;
// advance plugin states
if (m->nplugin) {
const int nslot = mjp_pluginCount();
for (int i = 0; i < m->nplugin; ++i) {
const int slot = m->plugin[i];
const mjpPlugin* plugin = mjp_getPluginAtSlotUnsafe(slot, nslot);
if (!plugin) {
mjERROR("invalid plugin slot: %d", slot);
}
if (plugin->advance) {
plugin->advance(m, d, i);
}
}
}
// save qacc for next step warmstart
mju_copy(d->qacc_warmstart, d->qacc, m->nv);
}
// Euler integrator, semi-implicit in velocity, possibly skipping factorisation
void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
TM_START;
mj_markStack(d);
mjtNum* qfrc = mjSTACKALLOC(d, m->nv, mjtNum);
mjtNum* qacc = mjSTACKALLOC(d, m->nv, mjtNum);
// sleep filtering
int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < m->nv;
int nv = sleep_filter ? d->nv_awake : m->nv;
const int* dof_awake_ind = sleep_filter ? d->dof_awake_ind : NULL;
// check for dof damping if disable flag is not set
int dof_damping = 0;
if (!mjDISABLED(mjDSBL_EULERDAMP) && !mjDISABLED(mjDSBL_DAMPER)) {
for (int v=0; v < nv; v++) {
int i = sleep_filter ? dof_awake_ind[v] : v;
if (m->dof_damping[i] > 0 ||
!mju_isZero(m->dof_dampingpoly + mjNPOLY*i, mjNPOLY) ||
m->jnt_actuatorid[m->dof_jntid[i]] != -1) {
dof_damping = 1;
break;
}
}
}
// no damping or disabled: explicit velocity integration
if (!dof_damping) {
if (sleep_filter) {
mju_copyInd(qacc, d->qacc, dof_awake_ind, nv);
} else {
mju_copy(qacc, d->qacc, nv);
}
}
// damping: integrate implicitly
else {
if (!skipfactor) {
// qH = M
if (sleep_filter) {
mju_copySparse(d->qH, d->M, m->M_rownnz, m->M_rowadr, dof_awake_ind, d->nv_awake);
} else {
mju_copy(d->qH, d->M, m->nC);
}
// qH += h*diag(B)
for (int v=0; v < nv; v++) {
int i = sleep_filter ? dof_awake_ind[v] : v;
mjtNum qv = d->qvel[i];
mjtNum poly[mjNPOLY];
mju_copy(poly, m->dof_dampingpoly + mjNPOLY*i, mjNPOLY);
mjtNum damping = m->dof_damping[i]
+ mj_actuatorDamping(m, mjOBJ_JOINT, m->dof_jntid[i], poly);
mjtNum damp_deriv = mjd_xPolyForce(damping, poly, qv, mjNPOLY, 1);
d->qH[m->M_rowadr[i] + m->M_rownnz[i] - 1] += m->opt.timestep * damp_deriv;
}
// factorize in-place
mj_factorI(d->qH, d->qHDiagInv, nv, m->M_rownnz, m->M_rowadr, m->M_colind, dof_awake_ind);
}
// solve
if (sleep_filter) {
mju_addInd(qfrc, d->qfrc_smooth, d->qfrc_constraint, dof_awake_ind, nv);
mju_copyInd(qacc, qfrc, dof_awake_ind, nv);
} else {
mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
mju_copy(qacc, qfrc, nv);
}
mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1,
m->M_rownnz, m->M_rowadr, m->M_colind, dof_awake_ind);
}
// advance state and time
mj_advance(m, d, d->act_dot, qacc, NULL);
mj_freeStack(d);
TM_END(mjTIMER_ADVANCE);
}
// Euler integrator, semi-implicit in velocity
void mj_Euler(const mjModel* m, mjData* d) {
mj_EulerSkip(m, d, 0);
}
// RK4 tableau
const mjtNum RK4_A[9] = {
0.5, 0, 0,
0, 0.5, 0,
0, 0, 1
};
const mjtNum RK4_B[4] = {
1.0/6.0, 1.0/3.0, 1.0/3.0, 1.0/6.0
};
// Runge Kutta explicit order-N integrator
// (A,B) is the tableau, C is set to row_sum(A)
void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
int nv = m->nv, nq = m->nq, na = m->na;
mjtNum h = m->opt.timestep, time = d->time;
mjtNum C[9], T[9], *X[10], *F[10], *dX;
const mjtNum* A = (N == 4 ? RK4_A : 0);
const mjtNum* B = (N == 4 ? RK4_B : 0);
// check order
if (!A) {
mjERROR("supported RK orders: N=4");
}
// allocate space for intermediate solutions
mj_markStack(d);
dX = mjSTACKALLOC(d, 2*nv+na, mjtNum);
for (int i=0; i < N; i++) {
X[i] = mjSTACKALLOC(d, nq+nv+na, mjtNum);
F[i] = mjSTACKALLOC(d, nv+na, mjtNum);
}
// precompute C and T; C,T,A have size (N-1)
for (int i=1; i < N; i++) {
// C(i) = sum_j A(i,j)
C[i-1] = 0;
for (int j=0; j < i; j++) {
C[i-1] += A[(i-1)*(N-1)+j];
}
// compute T
T[i-1] = d->time + C[i-1]*h;
}
// init X[0], F[0]; mj_forward() was already called
mju_copy(X[0], d->qpos, nq);
mju_copy(X[0]+nq, d->qvel, nv);
mju_copy(F[0], d->qacc, nv);
if (na) {
mju_copy(X[0]+nq+nv, d->act, na);
mju_copy(F[0]+nv, d->act_dot, na);
}
// compute the remaining X[i], F[i]
for (int i=1; i < N; i++) {
// compute dX
mju_zero(dX, 2*nv+na);
for (int j=0; j < i; j++) {
mju_addToScl(dX, X[j]+nq, A[(i-1)*(N-1)+j], nv);
mju_addToScl(dX+nv, F[j], A[(i-1)*(N-1)+j], nv+na);
}
// compute X[i] = X[0] '+' dX
mju_copy(X[i], X[0], nq+nv+na);
mj_integratePos(m, X[i], dX, h);
mju_addToScl(X[i]+nq, dX+nv, h, nv+na);
// set X[i], T[i-1] in mjData
mju_copy(d->qpos, X[i], nq);
mju_copy(d->qvel, X[i]+nq, nv);
if (na) {
mju_copy(d->act, X[i]+nq+nv, na);
}
d->time = T[i-1];
// evaluate F[i]
mj_forwardSkip(m, d, mjSTAGE_NONE, 1); // 1: do not recompute sensors and energy
mju_copy(F[i], d->qacc, nv);
if (na) {
mju_copy(F[i]+nv, d->act_dot, na);
}
}
// compute dX for final update (using B instead of A)
mju_zero(dX, 2*nv+na);
for (int j=0; j < N; j++) {
mju_addToScl(dX, X[j]+nq, B[j], nv);
mju_addToScl(dX+nv, F[j], B[j], nv+na);
}
// reset state and time
d->time = time;
mju_copy(d->qpos, X[0], nq);
mju_copy(d->qvel, X[0]+nq, nv);
mju_copy(d->act, X[0]+nq+nv, na);
// advance state and time
mj_advance(m, d, dX+2*nv, dX+nv, dX);
mj_freeStack(d);
}
// return 1 if any flex needs implicit stiffness treatment (interp or bending)
static int flex_has_implicit_stiffness(const mjModel* m) {
for (int f=0; f < m->nflex; f++) {
if (m->flex_rigid[f]) {
continue;
}
// interpolated flex with stiffness
if (m->flex_interp[f] &&
m->flex_edgeequality[f] != 3 &&
m->flex_stiffness[m->flex_stiffnessadr[f]] != 0) {
return 1;
}
// standard flex with bending
if (!m->flex_interp[f] && m->flex_dim[f] == 2 &&
m->flex_bendingadr[f] >= 0) {
return 1;
}
}
return 0;
}
// preconditioned CG solve for implicit flex interp
// solves (M - h*qDeriv - (h^2+h*d)*K) * qacc = qfrc - h*K*qvel
// where K is the flex stiffness, using the already-factored standard system
// (M - h*qDeriv) as a preconditioner
static void flexInterp_cgsolve(const mjModel* m, mjData* d,
mjtNum* qacc, const mjtNum* qfrc, int nv) {
mjtNum h = m->opt.timestep;
int implicit = (m->opt.integrator == mjINT_IMPLICIT);
mj_markStack(d);
// allocate CG work vectors
mjtNum* rhs = mjSTACKALLOC(d, nv, mjtNum);
mjtNum* r = mjSTACKALLOC(d, nv, mjtNum);
mjtNum* z = mjSTACKALLOC(d, nv, mjtNum);
mjtNum* p = mjSTACKALLOC(d, nv, mjtNum);
mjtNum* Ap = mjSTACKALLOC(d, nv, mjtNum);
mjtNum* temp = mjSTACKALLOC(d, nv, mjtNum);
// precompute K_rot cache: same layout as m->flex_stiffness
int krot_size = m->nflexstiffness;
mjtNum* K_rot_cache = mjSTACKALLOC(d, krot_size, mjtNum);
mju_zero(K_rot_cache, krot_size);
mjd_flexInterp_cacheKrot(m, d, K_rot_cache);
// build RHS: rhs = qfrc
mju_copy(rhs, qfrc, nv);
// flex_interp velocity correction: rhs += h*K_interp*qvel (K_interp is NSD)
mjd_flexInterp_mul(m, d, rhs, d->qvel, h, 0, K_rot_cache);
// standard flex bending velocity correction: rhs -= h*K_bend*qvel
mjd_flexBend_mul(m, d, rhs, d->qvel, -h, 0); // rhs -= h*K_bend*v
// --- helper: compute Ap = A*x ---
// A*x = (M - h*qDeriv)*x - (h^2+h*d)*K_interp*x + (h^2+h*d)*K_bend*x
#define FLEX_CG_MATVEC(Ap_out, x_in) \
mju_mulMatVecSparse(Ap_out, d->qDeriv, x_in, nv, m->D_rownnz, m->D_rowadr, \
m->D_colind, NULL); \
mju_mulSymVecSparse(temp, d->M, x_in, nv, m->M_rownnz, m->M_rowadr, \
m->M_colind); \
mju_addScl(Ap_out, temp, Ap_out, -h, nv); \
mjd_flexInterp_mul(m, d, Ap_out, x_in, -(h*h), -h, K_rot_cache); \
mjd_flexBend_mul(m, d, Ap_out, x_in, h*h, h)
// --- helper: preconditioner solve z = (M - h*qDeriv)^{-1} * r ---
#define FLEX_CG_PRECOND(z_out, r_in) \
if (implicit) { \
mju_solveLUSparse(z_out, d->qLU, r_in, nv, m->D_rownnz, m->D_rowadr, \
m->D_diag, m->D_colind, NULL); \
} else { \
mju_copy(z_out, r_in, nv); \
mj_solveLD(z_out, d->qH, d->qHDiagInv, nv, 1, m->M_rownnz, m->M_rowadr, \
m->M_colind, NULL); \
}
// initial residual: r = rhs - A*qacc
FLEX_CG_MATVEC(Ap, qacc);
mju_sub(r, rhs, Ap, nv);
// check if already converged
mjtNum rnorm = mju_dot(r, r, nv);
mjtNum tol = 1e-10 * mju_dot(rhs, rhs, nv);
if (rnorm < tol || rnorm < mjMINVAL) {
mj_freeStack(d);
return;
}
// z = precond(r), p = z
FLEX_CG_PRECOND(z, r);
mju_copy(p, z, nv);
mjtNum rz = mju_dot(r, z, nv);
// CG iterations
int maxiter = 50;
for (int iter=0; iter < maxiter; iter++) {
FLEX_CG_MATVEC(Ap, p);
// alpha = rz / dot(p, Ap)
mjtNum pAp = mju_dot(p, Ap, nv);
if (mju_abs(pAp) < mjMINVAL) break;
mjtNum alpha = rz / pAp;
// qacc += alpha * p
mju_addToScl(qacc, p, alpha, nv);
// r -= alpha * Ap
mju_addToScl(r, Ap, -alpha, nv);
// check convergence
rnorm = mju_dot(r, r, nv);
if (rnorm < tol || rnorm < mjMINVAL) break;
// z = precond(r)
FLEX_CG_PRECOND(z, r);
// beta = rz_new / rz
mjtNum rz_new = mju_dot(r, z, nv);
mjtNum beta = rz_new / mju_max(mjMINVAL, rz);
// p = z + beta * p
mju_addScl(p, z, p, beta, nv);
rz = rz_new;
}
#undef FLEX_CG_MATVEC
#undef FLEX_CG_PRECOND
mj_freeStack(d);
}
// return 1 if free joint is eligible for midpoint quaternion integration:
// standalone 6-DOF tree with no children, awake, and unconstrained
static int midpoint_eligible(const mjModel* m, const mjData* d, int jnt) {
if (m->jnt_type[jnt] != mjJNT_FREE) {
return 0;
}
int body = m->jnt_bodyid[jnt];
int adr = m->jnt_dofadr[jnt];
int tree = m->dof_treeid[adr];
// must be standalone 6-DOF tree with no children
if (m->tree_dofnum[tree] != 6 ||
m->body_subtreemass[body] != m->body_mass[body]) {
return 0;
}
// must be awake
if (!d->tree_awake[tree]) {
return 0;
}
// must be unconstrained
if (d->nefc) {
// islands enabled: O(1) lookup
if (!mjDISABLED(mjDSBL_ISLAND)) {
if (d->dof_island[adr] >= 0) {
return 0;
}
}
// islands disabled: check if any constraint involves this tree
else {
for (int c=0; c < d->nefc; c++) {
int type = d->efc_type[c];
int id = d->efc_id[c];
// contact: check if either geom belongs to this body
if (type == mjCNSTR_CONTACT_FRICTIONLESS ||
type == mjCNSTR_CONTACT_PYRAMIDAL ||
type == mjCNSTR_CONTACT_ELLIPTIC) {
int g1 = d->contact[id].geom[0];
int g2 = d->contact[id].geom[1];
if (g1 >= 0 && m->geom_bodyid[g1] == body) return 0;
if (g2 >= 0 && m->geom_bodyid[g2] == body) return 0;
}
// connect or weld: check if either body is this body
else if (type == mjCNSTR_EQUALITY &&
(m->eq_type[id] == mjEQ_CONNECT || m->eq_type[id] == mjEQ_WELD)) {
int b1 = m->eq_obj1id[id];
int b2 = m->eq_obj2id[id];
if (m->eq_objtype[id] == mjOBJ_SITE) {
b1 = m->site_bodyid[b1];
b2 = m->site_bodyid[b2];
}
if (b1 == body || b2 == body) return 0;
}
// tendon limit or friction: check first two trees
else if (type == mjCNSTR_LIMIT_TENDON || type == mjCNSTR_FRICTION_TENDON) {
if (m->tendon_treeid[2*id] == tree ||
m->tendon_treeid[2*id+1] == tree) return 0;
}
}
}
}
// otherwise eligible
return 1;
}
// 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;
int nD = m->nD, nC = m->nC;
mj_markStack(d);
mjtNum* qfrc = mjSTACKALLOC(d, m->nv, mjtNum);
mjtNum* qacc = mjSTACKALLOC(d, m->nv, mjtNum);
// sleep filtering
int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nv_awake < m->nv;
int nv = sleep_filter ? d->nv_awake : m->nv;
const int* dof_awake_ind = sleep_filter ? d->dof_awake_ind : NULL;
// set qfrc = qfrc_smooth + qfrc_constraint
if (sleep_filter) {
mju_addInd(qfrc, d->qfrc_smooth, d->qfrc_constraint, dof_awake_ind, nv);
} else {
mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv);
}
// check for flex_interp that needs implicit treatment
int has_flex_stiffness = !sleep_filter && flex_has_implicit_stiffness(m);
// factorization
if (!skipfactor) {
// implicit
if (m->opt.integrator == mjINT_IMPLICIT) {
// compute analytical derivative qDeriv
mjd_smooth_vel(m, d, /* flg_bias = */ 1);
// gather qLU <- M (lower to full)
mju_gatherMasked(d->qLU, d->M, m->mapM2D, nD);
// set qLU = M - dt*qDeriv
mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, nD);
}
// implicitfast
else if (m->opt.integrator == mjINT_IMPLICITFAST) {
// compute analytical derivative qDeriv; skip rne derivative
mjd_smooth_vel(m, d, /* flg_bias = */ 0);
// modified mass matrix: gather qH <- qDeriv (full to lower)
mju_gather(d->qH, d->qDeriv, m->mapD2M, nC);
// set qH = M - dt*qDeriv
mju_addScl(d->qH, d->M, d->qH, -m->opt.timestep, nC);
} else {
mjERROR("integrator must be implicit or implicitfast");
}
// standard factorization (implicit / implicitfast)
if (m->opt.integrator == mjINT_IMPLICIT) {
int* scratch = mjSTACKALLOC(d, nv, int);
mju_factorLUSparse(d->qLU, nv, scratch, m->D_rownnz, m->D_rowadr, m->D_colind, dof_awake_ind);
} else {
mj_factorI(d->qH, d->qHDiagInv, nv, m->M_rownnz, m->M_rowadr, m->M_colind, dof_awake_ind);
}
}
// standard sparse solve
if (m->opt.integrator == mjINT_IMPLICIT) {
mju_solveLUSparse(qacc, d->qLU, qfrc, nv, m->D_rownnz, m->D_rowadr, m->D_diag, m->D_colind,
dof_awake_ind);
} else {
// implicitfast
if (sleep_filter) {
mju_copyInd(qacc, qfrc, dof_awake_ind, nv);
} else {
mju_copy(qacc, qfrc, nv);
}
mj_solveLD(qacc, d->qH, d->qHDiagInv, nv, 1, m->M_rownnz, m->M_rowadr, m->M_colind, dof_awake_ind);
}
// flex: CG correction for implicit flex stiffness
if (has_flex_stiffness) {
flexInterp_cgsolve(m, d, qacc, qfrc, m->nv);
}
// count and list joints of free bodies eligible for midpoint integration
int nfree = 0;
int* free_jntid = NULL;
if (!mjENABLED(mjENBL_INVDISCRETE) &&
m->opt.integrator == mjINT_IMPLICITFAST &&
m->opt.density == 0 && m->opt.viscosity == 0) {
free_jntid = mjSTACKALLOC(d, m->njnt, int);
for (int j=0; j < m->njnt; j++) {
if (midpoint_eligible(m, d, 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);
TM_END(mjTIMER_ADVANCE);
}
// fully implicit in velocity
void mj_implicit(const mjModel* m, mjData* d) {
mj_implicitSkip(m, d, 0);
}
//-------------------------- top-level API ---------------------------------------------------------
// forward dynamics with skip; skipstage is mjtStage
void mj_forwardSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor) {
TM_START;
// position-dependent
if (skipstage < mjSTAGE_POS) {
mj_fwdPosition(m, d);
if (!skipsensor) {
mj_sensorPos(m, d);
}
if (!d->flg_energypos) {
if (mjENABLED(mjENBL_ENERGY)) {
mj_energyPos(m, d);
} else {
d->energy[0] = d->energy[1] = 0;
}
}
}
// velocity-dependent
if (skipstage < mjSTAGE_VEL) {
mj_fwdVelocity(m, d);
if (!skipsensor) {
mj_sensorVel(m, d);
}
if (mjENABLED(mjENBL_ENERGY) && !d->flg_energyvel) {
mj_energyVel(m, d);
}
}
// acceleration-dependent
if (mjcb_control && !mjDISABLED(mjDSBL_ACTUATION)) {
mjcb_control(m, d);
}
mj_fwdActuation(m, d);
mj_fwdAcceleration(m, d);
mj_fwdConstraint(m, d);
if (!skipsensor) {
d->flg_rnepost = 0; // clear flag for lazy evaluation
mj_sensorAcc(m, d);
}
TM_END(mjTIMER_FORWARD);
}
// forward dynamics
void mj_forward(const mjModel* m, mjData* d) {
mj_forwardSkip(m, d, mjSTAGE_NONE, 0);
}
// advance simulation using control callback
void mj_step(const mjModel* m, mjData* d) {
TM_START;
// common to all integrators
mj_checkPos(m, d);
mj_checkVel(m, d);
mj_forward(m, d);
mj_checkAcc(m, d);
// compare forward and inverse solutions if enabled
if (mjENABLED(mjENBL_FWDINV)) {
mj_compareFwdInv(m, d);
}
// use selected integrator
switch ((mjtIntegrator) m->opt.integrator) {
case mjINT_EULER:
mj_Euler(m, d);
break;
case mjINT_RK4:
mj_RungeKutta(m, d, 4);
break;
case mjINT_IMPLICIT:
case mjINT_IMPLICITFAST:
mj_implicit(m, d);
break;
default:
mjERROR("invalid integrator");
}
TM_END(mjTIMER_STEP);
}
// advance simulation in two phases: before input is set by user
void mj_step1(const mjModel* m, mjData* d) {
TM_START;
mj_checkPos(m, d);
mj_checkVel(m, d);
mj_fwdPosition(m, d);
mj_sensorPos(m, d);
if (!d->flg_energypos) {
if (mjENABLED(mjENBL_ENERGY)) {
mj_energyPos(m, d);
} else {
d->energy[0] = d->energy[1] = 0;
}
}
mj_fwdVelocity(m, d);
mj_sensorVel(m, d);
if (mjENABLED(mjENBL_ENERGY) && !d->flg_energyvel) {
mj_energyVel(m, d);
}
if (mjcb_control) {
mjcb_control(m, d);
}
TM_END(mjTIMER_STEP);
}
// >>>> user can modify ctrl and q/xfrc_applied between step1 and step2 <<<<
// advance simulation in two phases: after input is set by user
void mj_step2(const mjModel* m, mjData* d) {
TM_START;
mj_fwdActuation(m, d);
mj_fwdAcceleration(m, d);
mj_fwdConstraint(m, d);
d->flg_rnepost = 0; // clear flag for lazy evaluation
mj_sensorAcc(m, d);
mj_checkAcc(m, d);
// compare forward and inverse solutions if enabled
if (mjENABLED(mjENBL_FWDINV)) {
mj_compareFwdInv(m, d);
}
// integrate with Euler or implicit; RK4 defaults to Euler
if (m->opt.integrator == mjINT_IMPLICIT || m->opt.integrator == mjINT_IMPLICITFAST) {
mj_implicit(m, d);
} else {
mj_Euler(m, d);
}
d->timer[mjTIMER_STEP].number--;
TM_END(mjTIMER_STEP);
}