mujoco

// 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 "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;

  // 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);

  // warmstart solver
  warmstart(m, d);
  mju_zeroInt(d->solver_niter, mjNISLAND);

  // check if islands are supported
  int islands_supported = !mjDISABLED(mjDSBL_ISLAND)    &&
                          nisland > 0                   &&
                          m->opt.noslip_iterations == 0 &&
                          (m->opt.solver == mjSOL_CG || m->opt.solver == mjSOL_NEWTON);

  // run solver over constraint islands
  if (islands_supported) {
    int nidof = d->nidof;

    // copy inputs to islands (vel+acc deps, pos-dependent already copied in mj_island)
    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);
  }

  // run solver over all constraints
  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;

    default:
      mjERROR("unknown solver type %d", m->opt.solver);
    }
  }

  // run noslip solver if enabled
  if (m->opt.noslip_iterations > 0) {
    mj_solNoSlip(m, d, m->opt.noslip_iterations);
  }

  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);
}


// context for flex interp reduced banded factorization/solve
typedef struct {
  mjtNum* H;              // banded Cholesky-factored matrix (ndof x nband)
  int* dof_indices;       // global DOF index for each local flex DOF
  int ndof;               // number of flex DOFs
  int nband;              // half-bandwidth + 1 (number of band columns)
  int ncoupling;          // number of off-diagonal coupling terms
  mjtNum* coupling_val;   // coupling coefficient values
  int* coupling_row;      // local flex row index for each coupling term
  int* coupling_col;      // global DOF column index for each coupling term
} FlexInterpContext;


// collect flex DOFs for one flex, marking seen_dof and incrementing count
static void flexInterp_collect(const mjModel* m, int f,
                               int* chain_dofs, int* seen_dof, int* count) {
  int nodenum = m->flex_nodenum[f];
  int nodeadr = m->flex_nodeadr[f];
  for (int n=0; n < nodenum; n++) {
    int b = m->flex_nodebodyid[nodeadr+n];
    int chain_nnz;
    if (m->body_dofnum[b] == 0) {
      // pinned node: use bodyChain to get parent DOFs
      chain_nnz = mj_bodyChain(m, b, chain_dofs);
    } else {
      // regular flex node: use body's own DOFs only
      chain_nnz = m->body_dofnum[b];
      for (int j=0; j < chain_nnz; j++) {
        chain_dofs[j] = m->body_dofadr[b] + j;
      }
    }
    for (int i=0; i < chain_nnz; i++) {
      int dof = chain_dofs[i];
      if (!seen_dof[dof]) {
        seen_dof[dof] = 1;
        (*count)++;
      }
    }
  }
}


// build and factor the reduced banded matrix for flex interp DOFs
//   mark/free stack handled by caller
static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv) {
  FlexInterpContext ctx = {0};

  int* chain_dofs = mjSTACKALLOC(d, nv, int);
  int* seen_dof = mjSTACKALLOC(d, nv, int);
  mju_fillInt(seen_dof, 0, nv);

  // count flex DOFs
  int ndof = 0;
  for (int f=0; f < m->nflex; f++) {
    if (m->flex_interp[f]) {
      flexInterp_collect(m, f, chain_dofs, seen_dof, &ndof);
    }
  }
  if (ndof == 0) {
    return ctx;
  }

  // allocate and build global-to-local mapping
  int* dof_indices = mjSTACKALLOC(d, ndof, int);
  int* global2local = mjSTACKALLOC(d, nv, int);
  mju_fillInt(global2local, -1, nv);

  // collect unique DOFs in order
  int cnt = 0;
  mju_fillInt(seen_dof, 0, nv);
  for (int f=0; f < m->nflex; f++) {
    if (m->flex_interp[f]) {
      int nodenum = m->flex_nodenum[f];
      int nodeadr = m->flex_nodeadr[f];
      for (int n=0; n < nodenum; n++) {
        int b = m->flex_nodebodyid[nodeadr+n];
        int chain_nnz;
        if (m->body_dofnum[b] == 0) {
          // pinned node: use bodyChain to get parent DOFs
          chain_nnz = mj_bodyChain(m, b, chain_dofs);
        } else {
          // regular flex node: use body's own DOFs only
          chain_nnz = m->body_dofnum[b];
          for (int j=0; j < chain_nnz; j++) {
            chain_dofs[j] = m->body_dofadr[b] + j;
          }
        }
        for (int i=0; i < chain_nnz; i++) {
          int dof = chain_dofs[i];
          if (!seen_dof[dof]) {
            seen_dof[dof] = 1;
            dof_indices[cnt] = dof;
            global2local[dof] = cnt;
            cnt++;
          }
        }
      }
    }
  }

  // select sparse matrix format based on integrator
  int implicit = (m->opt.integrator == mjINT_IMPLICIT);
  const int* rownnz = implicit ? m->D_rownnz : m->M_rownnz;
  const int* rowadr = implicit ? m->D_rowadr : m->M_rowadr;
  const int* colind = implicit ? m->D_colind : m->M_colind;
  const mjtNum* source = implicit ? d->qLU : d->qH;

  // get precomputed bandwidth
  int bandwidth = 0;
  for (int f=0; f < m->nflex; f++) {
    if (m->flex_interp[f]) {
      if (m->flex_bandwidth[f] > bandwidth) {
        bandwidth = m->flex_bandwidth[f];
      }
    }
  }

  // compute ncoupling from sparse matrix entries
  int ncoupling = 0;
  for (int i=0; i < ndof; i++) {
    int row = dof_indices[i];
    int start = rowadr[row];
    int end = start + rownnz[row];
    for (int k=start; k < end; k++) {
      int local_j = global2local[colind[k]];
      if (local_j < 0) {
        ncoupling++;
      }
    }
  }

  // nband = bandwidth + 1 (includes diagonal)
  int nband = bandwidth + 1;

  // cap nband at ndof (dense fallback for small systems)
  if (nband > ndof) nband = ndof;

  // allocate coupling storage
  mjtNum* coupling_val = NULL;
  int* coupling_row = NULL;
  int* coupling_col = NULL;
  if (ncoupling > 0) {
    coupling_val = mjSTACKALLOC(d, ncoupling, mjtNum);
    coupling_row = mjSTACKALLOC(d, ncoupling, int);
    coupling_col = mjSTACKALLOC(d, ncoupling, int);
  }

  // build H_flex (banded) from qLU (implicit) or qH (implicitfast)
  mjtNum* H = mjSTACKALLOC(d, ndof*nband, mjtNum);
  mju_zero(H, ndof*nband);

  int coup_cnt = 0;
  for (int i=0; i < ndof; i++) {
    int row = dof_indices[i];
    int start = rowadr[row];
    int end = start + rownnz[row];
    for (int k=start; k < end; k++) {
      int col = colind[k];
      int local_j = global2local[col];
      if (local_j >= 0) {
        // store lower triangle only: row i, col local_j, where i >= local_j
        if (i >= local_j) {
          H[i*nband + nband-1-(i-local_j)] = source[k];
        } else {
          // upper triangle entry: store symmetrically in lower triangle
          H[local_j*nband + nband-1-(local_j-i)] = source[k];
        }
      } else if (coup_cnt < ncoupling) {
        coupling_val[coup_cnt] = source[k];
        coupling_row[coup_cnt] = i;
        coupling_col[coup_cnt] = col;
        coup_cnt++;
      }
    }
  }

  // add flex stiffness in banded format and factorize
  mjd_flexInterp_addH(m, d, H, dof_indices, ndof, nband, m->opt.timestep);
  mju_cholFactorBand(H, ndof, nband, 0, 0, 0);

  // store results in context
  ctx.H = H;
  ctx.dof_indices = dof_indices;
  ctx.ndof = ndof;
  ctx.nband = nband;
  ctx.ncoupling = ncoupling;
  ctx.coupling_val = coupling_val;
  ctx.coupling_row = coupling_row;
  ctx.coupling_col = coupling_col;
  return ctx;
}


// solve the reduced banded system for flex interp DOFs, overwrite qacc
static void flexInterp_solve(const mjModel* m, mjData* d, const FlexInterpContext* ctx,
                             mjtNum* qacc, const mjtNum* qfrc, int nv) {
  int ndof = ctx->ndof;
  mjtNum* qfrc_flex = mjSTACKALLOC(d, ndof, mjtNum);
  mjtNum* res = mjSTACKALLOC(d, nv, mjtNum);

  mjtNum h = m->opt.timestep;
  mjtNum damp = (m->nflex > 0 && m->flex_damping) ? m->flex_damping[0] : 0;
  mjtNum scl = h*h + h*damp;
  mjtNum factor = (scl > mjMINVAL) ? (h/scl) : 0;

  // velocity correction: -h * K * v
  mju_zero(res, nv);
  mjd_flexInterp_mulKD(m, d, res, d->qvel, h);

  for (int i=0; i < ndof; i++) {
    int global_dof = ctx->dof_indices[i];
    qfrc_flex[i] = qfrc[global_dof] + res[global_dof] * factor;
  }

  // coupling correction: qfrc_flex -= H_coupling * qacc_parent
  for (int k=0; k < ctx->ncoupling; k++) {
    qfrc_flex[ctx->coupling_row[k]] -= ctx->coupling_val[k] * qacc[ctx->coupling_col[k]];
  }

  // solve with banded Cholesky and scatter back
  mju_cholSolveBand(qfrc_flex, ctx->H, qfrc_flex, ndof, ctx->nband, 0);
  mju_scatter(qacc, qfrc_flex, ctx->dof_indices, ndof);
}


// return 1 if free joint is eligible for midpoint quaternion integration:
//   standalone 6-DOF tree with no children
static int midpoint_eligible(const mjModel* m, int jnt) {
  if (m->jnt_type[jnt] == mjJNT_FREE) {
    int body = m->jnt_bodyid[jnt];
    int treeid = m->dof_treeid[m->jnt_dofadr[jnt]];
    return m->tree_dofnum[treeid] == 6 &&
           m->body_subtreemass[body] == m->body_mass[body];
  }

  return 0;
}


// return 1 if the body's CoM is at the joint origin (no translational-rotational coupling)
static int midpoint_aligned(const mjModel* m, int jnt) {
  int body = m->jnt_bodyid[jnt];
  return m->body_ipos[3*body+0] == 0 &&
         m->body_ipos[3*body+1] == 0 &&
         m->body_ipos[3*body+2] == 0;
}


// implicit midpoint integration for 3D rotation of a single body
//
// solves the Euler rigid body equation in the inertial frame:
//   I * (w_new - w) / h = tau - w_mid x (I*w_mid)
// where w_mid = (w + w_new) / 2 is solved via Newton iteration.
//
// inputs:
//   inertia:   principal moments of inertia (3)
//   w:         initial angular velocity in principal axes frame (3)
//   tau:       external torque in principal axes frame (3)
//   h:         timestep
// outputs:
//   w_mid:     midpoint angular velocity in principal axes frame (3)
// returns:     number of Newton iterations
static int midpointNewton(const mjtNum inertia[3], const mjtNum w[3],
                          const mjtNum tau[3], mjtNum h, mjtNum w_mid[3]) {
  // precompute constants
  mjtNum i2h = 2.0 / h;
  mjtNum dI[3] = {inertia[2]-inertia[1], inertia[0]-inertia[2], inertia[1]-inertia[0]};
  mjtNum i2h_I[3] = {i2h*inertia[0], i2h*inertia[1], i2h*inertia[2]};

  // initialize solution to previous angular velocity
  mji_copy3(w_mid, w);

  // Newton iteration
  int niter;
  for (niter=0; niter < 100; niter++) {
    // compute Coriolis term
    mjtNum Iw[3] = {inertia[0]*w_mid[0], inertia[1]*w_mid[1], inertia[2]*w_mid[2]};
    mjtNum coriolis[3];
    mji_cross(coriolis, w_mid, Iw);

    // residual: f = i2h*I*(w_mid - w) + w_mid x (I*w_mid) - tau
    mjtNum f[3];
    for (int k=0; k < 3; k++) {
      f[k] = i2h_I[k]*(w_mid[k] - w[k]) + coriolis[k] - tau[k];
    }

    // check convergence
    mjtNum fnorm = mju_norm3(f);
#ifndef mjUSESINGLE
    mjtNum tol = 1e-13;
#else
    mjtNum tol = 1e-6f;
#endif
    if (fnorm < tol*(1 + i2h*mju_norm3(Iw))) break;

    // Jacobian: J = i2h*diag(I) + d(w x Iw)/dw
    mjtNum J[9];
    J[0] = i2h_I[0];        J[1] = w_mid[2]*dI[0];  J[2] = w_mid[1]*dI[0];
    J[3] = w_mid[2]*dI[1];  J[4] = i2h_I[1];        J[5] = w_mid[0]*dI[1];
    J[6] = w_mid[1]*dI[2];  J[7] = w_mid[0]*dI[2];  J[8] = i2h_I[2];

    // solve J*delta = -f for search direction delta
    mjtNum neg_f[3] = {-f[0], -f[1], -f[2]};
    mjtNum delta[3];
    mju_solve3(delta, J, neg_f);

    // backtracking line search
    mjtNum step = 1.0;
    for (int ls=0; ls < 20; ls++) {
      // candidate step
      mjtNum w_try[3], Iw_try[3];
      for (int k=0; k < 3; k++) {
        w_try[k] = w_mid[k] + step*delta[k];
        Iw_try[k] = inertia[k]*w_try[k];
      }
      mjtNum coriolis_try[3];
      mji_cross(coriolis_try, w_try, Iw_try);

      // residual at candidate step
      mjtNum f_try[3];
      for (int k=0; k < 3; k++) {
        f_try[k] = i2h_I[k]*(w_try[k] - w[k]) + coriolis_try[k] - tau[k];
      }

      // accept step if residual decreased, otherwise backtrack
      if (mju_norm3(f_try) < fnorm) {
        mji_copy3(w_mid, w_try);
        break;
      }
      step *= 0.5;
    }
  }

  return niter;
}


// implicit midpoint integration for one free body
//
// solves the Euler rigid body equation in the inertial frame:
//   I * dw/dt = tau - w x (I*w)
// using the implicit midpoint rule:
//   I * (w_new - w_old) / h = tau_mid - w_mid x (I*w_mid)
// where w_mid = (w_old + w_new) / 2 is solved via Newton iteration.
//
// inputs:
//   mass:      body mass
//   inertia:   principal moments of inertia
//   ipos:      CoM offset from joint origin, in body frame
//   iquat:     inertial quaternion (body_iquat)
//   xquat:     body orientation in world frame
//   qvel_old:  current velocity (lin in world : rot in body)
//   qfrc:      external force (lin in world : rot in body)
//   gravity:   gravitational acceleration in world frame (NULL: no gravity)
//   h:         timestep
// outputs:
//   qvel_new:  next velocity (lin in world : rot in body)
int mj_midpoint(mjtNum mass, const mjtNum inertia[3], const mjtNum ipos[3],
                const mjtNum iquat[4], const mjtNum xquat[4], const mjtNum qvel_old[6],
                const mjtNum qfrc[6], const mjtNum gravity[3], mjtNum h,
                mjtNum qvel_new[6]) {
  // transform angular velocity and torque to inertial frame
  mjtNum iquat_neg[4], w[3], tau[3];
  mji_negQuat(iquat_neg, iquat);
  mji_rotVecQuat(w, qvel_old+3, iquat_neg);  // qvel+3 (angular) is in body frame
  mji_rotVecQuat(tau, qfrc+3, iquat_neg);    // qfrc+3 (angular) is in body frame

  // check for translational-rotational coupling
  int aligned = (ipos[0] == 0 && ipos[1] == 0 && ipos[2] == 0);

  mjtNum r_com[3];    // joint-to-CoM vector in inertial frame
  mjtNum tau_com[3];  // torque at CoM in inertial frame
  mjtNum rot_x2i[4];  // quaternion rotation from world to inertial frame
  mjtNum force[3];    // external force in inertial frame

  // compute torque at CoM in inertial frame
  if (aligned) {
    mji_copy3(tau_com, tau);
  } else {
    // rotation from world to inertial frame
    mjtNum xquat_neg[4];
    mji_negQuat(xquat_neg, xquat);
    mji_mulQuat(rot_x2i, iquat_neg, xquat_neg);

    // force and CoM offset in inertial frame
    mji_rotVecQuat(force, qfrc, rot_x2i);
    mji_rotVecQuat(r_com, ipos, iquat_neg);

    // torque at CoM in inertial frame
    mjtNum rxf[3];
    mji_cross(rxf, r_com, force);
    mji_sub3(tau_com, tau, rxf);
  }

  // solve for midpoint angular velocity
  mjtNum w_mid[3];
  int niter = midpointNewton(inertia, w, tau_com, h, w_mid);

  // next and mid angular velocities in inertial frame, rotate both to body frame
  mjtNum w_new[3], w_new_body[3], w_mid_body[3];
  for (int k=0; k < 3; k++) {
    w_new[k] = 2.0*w_mid[k] - w[k];
  }
  mji_rotVecQuat(w_new_body, w_new, iquat);
  mji_rotVecQuat(w_mid_body, w_mid, iquat);
  mji_copy3(qvel_new+3, w_new_body);

  // === aligned: return
  if (aligned) {
    return niter;
  }

  // === non-aligned: solve for translational velocity

  // rotate linear velocity to inertial frame
  mjtNum v[3];
  mji_rotVecQuat(v, qvel_old, rot_x2i);

  // current CoM velocities (rot, lin) in inertial frame
  mjtNum wxr[3];
  mji_cross(wxr, w, r_com);
  mjtNum vcom[3];
  mji_add3(vcom, v, wxr);

  // right-hand side for midpoint CoM velocity
  mjtNum i2h = 2.0 / h;
  mjtNum b[3];
  for (int k=0; k < 3; k++) {
    b[k] = force[k]/mass + i2h*vcom[k];
  }

  // add gravity, if any
  if (gravity) {
    mjtNum g_inertial[3];
    mji_rotVecQuat(g_inertial, gravity, rot_x2i);
    mji_addTo3(b, g_inertial);
  }

  // analytic solution for (i2h*Id + [w_mid]x) * vcom_mid = b
  mjtNum wnorm2 = mju_dot3(w_mid, w_mid);
  mjtNum denom = i2h*i2h + wnorm2;
  mjtNum w_dot_b = mju_dot3(w_mid, b);
  mjtNum w_cross_b[3];
  mji_cross(w_cross_b, w_mid, b);
  mjtNum vcom_mid[3];
  for (int k=0; k < 3; k++) {
    vcom_mid[k] = (i2h*b[k] + (w_dot_b/i2h)*w_mid[k] - w_cross_b[k]) / denom;
  }

  // recover midpoint and new joint velocity in inertial frame
  mjtNum wxr_mid[3];
  mji_cross(wxr_mid, w_mid, r_com);
  mjtNum v_mid[3], v_new[3];
  for (int k=0; k < 3; k++) {
    v_mid[k] = vcom_mid[k] - wxr_mid[k];
    v_new[k] = 2.0*v_mid[k] - v[k];
  }

  // estimate new orientation
  mjtNum axis[3];
  mji_copy3(axis, w_mid_body);
  mjtNum wnorm = mju_normalize3(axis);
  mjtNum qrot_new[4];
  mji_axisAngle2Quat(qrot_new, axis, h*wnorm);
  mjtNum xquat_new[4];
  mji_mulQuat(xquat_new, xquat, qrot_new);

  // v_new (linear): inertial → body → world using new orientation
  mjtNum v_body[3];
  mji_rotVecQuat(v_body, v_new, iquat);
  mji_rotVecQuat(qvel_new, v_body, xquat_new);

  return niter;
}


// compute next velocities via midpoint integration for eligible free bodies
//   qfrc:       total force (qfrc_smooth + qfrc_constraint)
//   free_jntid: list of eligible free joint IDs
//   nfree:      number of eligible free joints
//   qvel_old:   output array for old velocities (6 per joint)
//   qvel_new:   output array for new velocities (6 per joint)
//   dofadr:     output array for DOF addresses (1 per joint)
static void midpoint(const mjModel* m, const mjData* d, const mjtNum* qfrc,
                     const int* free_jntid, int nfree,
                     mjtNum* qvel_old, mjtNum* qvel_new, int* dofadr) {
  for (int i=0; i < nfree; i++) {
    int j = free_jntid[i];
    int body = m->jnt_bodyid[j];

    // save DOF address
    int adr = m->jnt_dofadr[j];
    dofadr[i] = adr;

    // save old (current) velocity, needed after mj_advance (which overwrites qvel)
    mju_copy(qvel_old+6*i, d->qvel+adr, 6);

    // compute external force = qfrc + qfrc_bias (undo bias subtraction)
    mjtNum qfrc_total[6];
    mju_add(qfrc_total, qfrc+adr, d->qfrc_bias+adr, 6);

    // gravity handled inside mj_midpoint (accelerating frame of reference)
    const mjtNum* gravity = mjDISABLED(mjDSBL_GRAVITY) ? NULL : m->opt.gravity;

    // midpoint solver for free joint j
    mj_midpoint(m->body_mass[body], m->body_inertia+3*body, m->body_ipos+3*body,
                m->body_iquat+4*body, d->xquat+4*body,
                d->qvel+adr, qfrc_total, gravity, m->opt.timestep, qvel_new+6*i);
  }
}


// fully implicit in velocity, possibly skipping factorization
void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
  TM_START;
  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
  int has_flex_interp = 0;
  for (int f=0; f < m->nflex; f++) {
    if (m->flex_interp[f]) {
      has_flex_interp = 1;
      break;
    }
  }

  // flex interp context (populated during factorization)
  FlexInterpContext flex = {0};

  // 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");
    }

    // flex: reduced dense factorization
    if (has_flex_interp && !sleep_filter) {
      flex = flexInterp_factor(m, d, nv);
    }

    // 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: reduced dense solve
  if (flex.H) {
    flexInterp_solve(m, d, &flex, qacc, qfrc, nv);
  }

  // count and list joints of free bodies eligible for midpoint integration
  int nfree = 0;
  int* free_jntid = NULL;
  if (!mjENABLED(mjENBL_INVDISCRETE)) {
    free_jntid = mjSTACKALLOC(d, m->njnt, int);
    for (int j=0; j < m->njnt; j++) {
      // add to list if eligible and awake
      if (midpoint_eligible(m, j) && d->tree_awake[m->dof_treeid[m->jnt_dofadr[j]]]) {
        free_jntid[nfree++] = j;
      }
    }
  }

  // compute midpoint velocities (used to update positions)
  int* dofadr = NULL;
  mjtNum* qvel_old = NULL;
  mjtNum* qvel_new = NULL;
  mjtNum* qvel_mid = NULL;
  if (nfree) {
    // allocate arrays, call midpoint solver for all eligible free joints
    dofadr = mjSTACKALLOC(d, nfree, int);
    qvel_new = mjSTACKALLOC(d, 6*nfree, mjtNum);
    qvel_old = mjSTACKALLOC(d, 6*nfree, mjtNum);
    midpoint(m, d, qfrc, free_jntid, nfree, qvel_old, qvel_new, dofadr);

    // build qvel_mid = d->qvel + h*qacc for all DOFs, then overwrite midpoint DOFs
    qvel_mid = mjSTACKALLOC(d, m->nv, mjtNum);
    mju_addScl(qvel_mid, d->qvel, qacc, m->opt.timestep, m->nv);
    for (int i=0; i < nfree; i++) {
      int adr = dofadr[i];
      int start = midpoint_aligned(m, free_jntid[i]) ? 3 : 0;
      for (int k=start; k < 6; k++) {
        qvel_mid[adr+k] = 0.5*(qvel_new[6*i+k] + qvel_old[6*i+k]);
      }
    }
  }

  // advance state and time (use qvel_mid if allocated, NULL otherwise)
  mj_advance(m, d, d->act_dot, qacc, qvel_mid);

  // overwrite midpoint DOFs with true next velocity and acceleration
  if (nfree) {
    mjtNum h_inv = 1.0 / m->opt.timestep;
    for (int i=0; i < nfree; i++) {
      // skip sleeping tree (may have been put to sleep during mj_advance)
      int adr = dofadr[i];
      if (!d->tree_awake[m->dof_treeid[adr]]) {
        continue;
      }

      // overwrite 3 or 6 midpoint DOFs with true next velocity and acceleration
      int start = midpoint_aligned(m, free_jntid[i]) ? 3 : 0;
      for (int k=start; k < 6; k++) {
        d->qvel[adr+k] = qvel_new[6*i+k];
        d->qacc[adr+k] = (qvel_new[6*i+k] - qvel_old[6*i+k]) * h_inv;
      }
    }
  }

  mj_freeStack(d);

  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);
}