// 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 #include #include #include #include #include // IWYU pragma: keep #include #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 "thread/thread_pool.h" #include "thread/thread_task.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 ----------------------------------------------------- // args for internal functions in mj_fwdPosition struct mjFwdPositionArgs_ { const mjModel* m; mjData* d; }; typedef struct mjFwdPositionArgs_ mjFwdPositionArgs; // wrapper for mj_crb and mj_factorM void* mj_inertialThreaded(void* args) { mjFwdPositionArgs* forward_args = (mjFwdPositionArgs*) args; mj_makeM(forward_args->m, forward_args->d); mj_factorM(forward_args->m, forward_args->d); return NULL; } // wrapper for mj_collision void* mj_collisionThreaded(void* args) { mjFwdPositionArgs* forward_args = (mjFwdPositionArgs*) args; mj_collision(forward_args->m, forward_args->d); return NULL; } // 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); // no threadpool: inertia and collision on main thread if (!d->threadpool) { // inertia, timed internally (POS_INERTIA) mj_makeM(m, d); mj_factorM(m, d); // collision, timed internally (POS_COLLISION) mj_collision(m, d); } // have threadpool: inertia and collision on separate threads else { mjTask tasks[2]; mjFwdPositionArgs forward_args; forward_args.m = m; forward_args.d = d; mju_defaultTask(&tasks[0]); tasks[0].func = mj_inertialThreaded; tasks[0].args = &forward_args; mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[0]); mju_defaultTask(&tasks[1]); tasks[1].func = mj_collisionThreaded; tasks[1].args = &forward_args; mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[1]); mju_taskJoin(&tasks[0]); mju_taskJoin(&tasks[1]); } 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_transmission(m, d); TM_ADD(mjTIMER_POS_KINEMATICS); TM_RESTART; mj_projectConstraint(m, d); TM_END(mjTIMER_POS_PROJECT); 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: always sparse mju_mulMatVecSparse(d->flexedge_velocity, d->flexedge_J, d->qvel, m->nflexedge, m->flexedge_J_rownnz, m->flexedge_J_rowadr, 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 mjtByte* 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); } } // struct encapsulating arguments to thread task struct mjSolIslandArgs_ { const mjModel* m; mjData* d; int island; }; typedef struct mjSolIslandArgs_ mjSolIslandArgs; // extract arguments, pass to CG solver static void* CG_wrapper(void* args) { mjSolIslandArgs* solargs = (mjSolIslandArgs*) args; mj_solCG_island(solargs->m, solargs->d, solargs->island, solargs->m->opt.iterations); return NULL; } // extract arguments, pass to Newton solver static void* Newton_wrapper(void* args) { mjSolIslandArgs* solargs = (mjSolIslandArgs*) args; mj_solNewton_island(solargs->m, solargs->d, solargs->island, solargs->m->opt.iterations); return NULL; } // CG solver, multi-threaded over islands static void solve_threaded(const mjModel* m, mjData* d, int flg_Newton) { mj_markStack(d); // allocate array of arguments to be passed to threads mjSolIslandArgs* sol_island_args = mjSTACKALLOC(d, d->nisland, mjSolIslandArgs); mjTask* tasks = mjSTACKALLOC(d, d->nisland, mjTask); for (int island = 0; island < d->nisland; ++island) { sol_island_args[island].m = m; sol_island_args[island].d = d; sol_island_args[island].island = island; mju_defaultTask(&tasks[island]); tasks[island].func = flg_Newton ? Newton_wrapper : CG_wrapper; tasks[island].args = &sol_island_args[island]; mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[island]); } for (int island = 0; island < d->nisland; ++island) { mju_taskJoin(&tasks[island]); } mj_freeStack(d); } // 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); // solve per island, with or without threads if (!d->threadpool) { // no threadpool, loop over islands for (int island=0; island < nisland; 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); } } } else { // have threadpool, solve using threads solve_threaded(m, d, m->opt.solver == mjSOL_NEWTON); } // 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); }