diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 5118a90a..d84d5f85 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -286,6 +286,10 @@ struct mjData_ { mjtNum* subtree_linvel; // linear velocity of subtree com (nbody x 3) mjtNum* subtree_angmom; // angular momentum about subtree com (nbody x 3) + // computed by mj_Euler + mjtNum* qH; // L'*D*L factorization of modified M (nM x 1) + mjtNum* qHDiagInv; // 1/diag(D) of modified M (nv x 1) + // computed by mj_implicit int* D_rownnz; // non-zeros in each row (nv x 1) int* D_rowadr; // address of each row in D_colind (nv x 1) diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 617b79e8..c4091869 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -509,6 +509,8 @@ X( mjtNum, efc_aref, njmax, 1 ) \ X( mjtNum, subtree_linvel, nbody, 3 ) \ X( mjtNum, subtree_angmom, nbody, 3 ) \ + X( mjtNum, qH, nM, 1 ) \ + X( mjtNum, qHDiagInv, nv, 1 ) \ X( int, D_rownnz, nv, 1 ) \ X( int, D_rowadr, nv, 1 ) \ X( int, D_colind, nD, 1 ) \ diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 2672f90b..23c39642 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -1000,7 +1000,7 @@ MJAPI int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus); MJAPI int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum mat[9]); -//---------------------- Miscellaneous -------------------------------------------------- +//---------------------- Miscellaneous ------------------------------------------------------------- // Muscle active force, prm = (range[2], force, scale, lmin, lmax, vmax, fpmax, fvmax). MJAPI mjtNum mju_muscleGain(mjtNum len, mjtNum vel, const mjtNum lengthrange[2], @@ -1080,6 +1080,13 @@ MJAPI char* mju_strncpy(char *dst, const char *src, int n); MJAPI mjtNum mju_sigmoid(mjtNum x); +//---------------------- Derivatives --------------------------------------------------------------- + +// finite differenced state-transition and control-transition matrices dy = A*dx + B*du +MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, + mjtNum* A, mjtNum* B); + + #if defined(__cplusplus) } #endif diff --git a/introspect/functions.py b/introspect/functions.py index f4994335..3291ddf6 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -6873,4 +6873,44 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Sigmoid function over 0<=x<=1 constructed from half-quadratics.', )), + ('mjd_transitionFD', + FunctionDecl( + name='mjd_transitionFD', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='m', + type=PointerType( + inner_type=ValueType(name='mjModel', is_const=True), + ), + ), + FunctionParameterDecl( + name='d', + type=PointerType( + inner_type=ValueType(name='mjData'), + ), + ), + FunctionParameterDecl( + name='eps', + type=ValueType(name='mjtNum'), + ), + FunctionParameterDecl( + name='centered', + type=ValueType(name='mjtByte'), + ), + FunctionParameterDecl( + name='A', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + FunctionParameterDecl( + name='B', + type=PointerType( + inner_type=ValueType(name='mjtNum'), + ), + ), + ), + doc='finite differenced state-transition and control-transition matrices dy = A*dx + B*du', # pylint: disable=line-too-long + )), ]) diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index e52bd545..f280bf03 100644 --- a/src/engine/engine_core_smooth.c +++ b/src/engine/engine_core_smooth.c @@ -855,20 +855,20 @@ void mj_crb(const mjModel* m, mjData* d) { -// sparse L'*D*L factorizaton of the inertia matrix M, assumed spd -void mj_factorM(const mjModel* m, mjData* d) { +// sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd +void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv, + mjtNum* qLDiagSqrtInv) { int cnt; int Madr_kk, Madr_ki; mjtNum tmp; // local copies of key variables - mjtNum* qLD = d->qLD; int* dof_Madr = m->dof_Madr; int* dof_parentid = m->dof_parentid; int nv = m->nv; // copy M into LD - mju_copy(d->qLD, d->qM, m->nM); + mju_copy(qLD, M, m->nM); // dense backward loop over dofs (regular only, simple diagonal already copied) for (int k=nv-1; k>=0; k--) { @@ -912,21 +912,31 @@ void mj_factorM(const mjModel* m, mjData* d) { // compute 1/diag(D), 1/sqrt(diag(D)) for (int i=0; iqLDiagInv[i] = 1.0/qLD[dof_Madr[i]]; - d->qLDiagSqrtInv[i] = 1.0/mju_sqrt(qLD[dof_Madr[i]]); + mjtNum qLDi = qLD[dof_Madr[i]]; + qLDiagInv[i] = 1.0/qLDi; + if (qLDiagSqrtInv) { + qLDiagSqrtInv[i] = 1.0/mju_sqrt(qLDi); + } } } +// sparse L'*D*L factorizaton of the inertia matrix M, assumed spd +void mj_factorM(const mjModel* m, mjData* d) { + mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv, d->qLDiagSqrtInv); +} + + + // sparse backsubstitution: x = inv(L'*D*L)*y // L is in lower triangle of qLD; D is on diagonal of qLD // handle n vectors at once -void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) { +void mj_solveLD(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n, + const mjtNum* qLD, const mjtNum* qLDiagInv) { mjtNum tmp; // local copies of key variables - mjtNum *qLD = d->qLD, *qLDiagInv = d->qLDiagInv; int* dof_Madr = m->dof_Madr; int* dof_parentid = m->dof_parentid; int nv = m->nv; @@ -1039,6 +1049,14 @@ void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) { +// sparse backsubstitution: x = inv(L'*D*L)*y +// use factorization in d +void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) { + mj_solveLD(m, d, x, y, n, d->qLD, d->qLDiagInv); +} + + + // half of sparse backsubstitution: x = sqrt(inv(D))*inv(L')*y void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n) { // local copies of key variables diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h index 403d0579..edc34d83 100644 --- a/src/engine/engine_core_smooth.h +++ b/src/engine/engine_core_smooth.h @@ -48,10 +48,18 @@ void mj_crbSkip(const mjModel* m, mjData* d, int skipsimple); // composite rigid body inertia algorithm MJAPI void mj_crb(const mjModel* m, mjData* d); +// sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd +MJAPI void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv, + mjtNum* qLDiagSqrtInv); + // sparse L'*D*L factorizaton of the inertia matrix M, assumed spd MJAPI void mj_factorM(const mjModel* m, mjData* d); // sparse backsubstitution: x = inv(L'*D*L)*y +MJAPI void mj_solveLD(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n, + const mjtNum* qLD, const mjtNum* qLDiagInv); + +// sparse backsubstitution: x = inv(L'*D*L)*y, use factorization in d MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n); // half of sparse backsubstitution: x = sqrt(inv(D))*inv(L')*y diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index b40e4558..65948af3 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -19,17 +19,15 @@ #include #include -#include "engine/engine_core_smooth.h" #include "engine/engine_forward.h" -#include "engine/engine_callback.h" #include "engine/engine_core_constraint.h" #include "engine/engine_io.h" +#include "engine/engine_inverse.h" #include "engine/engine_macro.h" #include "engine/engine_support.h" #include "engine/engine_util_blas.h" #include "engine/engine_util_errmem.h" #include "engine/engine_util_misc.h" -#include "engine/engine_util_sparse.h" #include "engine/engine_util_spatial.h" @@ -194,6 +192,151 @@ static void mjd_mulInertVec_vel(mjtNum D[36], const mjtNum i[10]) +//--------------------------- utility functions for mjd_stepFD ------------------------------------- + +// get state=[qpos; qvel; act] and optionally sensordata +static void getState(const mjModel* m, const mjData* d, mjtNum* state, mjtNum* sensordata) { + int nq = m->nq, nv = m->nv, na = m->na; + + mju_copy(state, d->qpos, nq); + mju_copy(state+nq, d->qvel, nv); + mju_copy(state+nq+nv, d->act, na); + if (sensordata) { + mju_copy(sensordata, d->sensordata, m->nsensordata); + } +} + + + +// set state=[qpos; qvel; act] and optionally warmstart accelerations +static void setState(const mjModel* m, mjData* d, const mjtNum* state, const mjtNum* ctrl, + const mjtNum* warmstart) { + int nq = m->nq, nv = m->nv, na = m->na; + + mju_copy(d->qpos, state, nq); + mju_copy(d->qvel, state+nq, nv); + mju_copy(d->act, state+nq+nv, na); + if (ctrl) { + mju_copy(d->ctrl, ctrl, m->nu); + } + if (warmstart) { + mju_copy(d->qacc_warmstart, warmstart, nv); + } +} + + + +// dx = (x2 - x1) / h +static void diff(mjtNum* restrict dx, const mjtNum* x1, const mjtNum* x2, mjtNum h, int n) { + mjtNum inv_h = 1/h; + for (int i=0; inq, nv = m->nv, na = m->na; + + if (nq == nv) { + diff(ds, s1, s2, h, nq+nv+na); + } else { + mj_differentiatePos(m, ds, h, s1, s2); + diff(ds+nv, s1+nq, s2+nq, h, nv+na); + } +} + + + +// finite-difference two vectors, forward, backward or centered +static void clampedDiff(mjtNum* dx, const mjtNum* x, const mjtNum* x_plus, const mjtNum* x_minus, + mjtNum h, int nx) { + if (x_plus && !x_minus) { + // forward differencing + diff(dx, x, x_plus, h, nx); + } else if (!x_plus && x_minus) { + // backward differencing + diff(dx, x_minus, x, h, nx); + } else if (x_plus && x_minus) { + // centered differencing + diff(dx, x_plus, x_minus, 2*h, nx); + } else { + // differencing failed, write zeros + mju_zero(dx, nx); + } +} + + + +// finite-difference two state vectors, forward, backward or centered +static void clampedStateDiff(const mjModel* m, mjtNum* ds, const mjtNum* s, const mjtNum* s_plus, + const mjtNum* s_minus, mjtNum h) { + if (s_plus && !s_minus) { + // forward differencing + stateDiff(m, ds, s, s_plus, h); + } else if (!s_plus && s_minus) { + // backward differencing + stateDiff(m, ds, s_minus, s, h); + } else if (s_plus && s_minus) { + // centered differencing + stateDiff(m, ds, s_minus, s_plus, 2*h); + } else { + // differencing failed, write zeros + mju_zero(ds, m->nq + m->nv + m->na); + } +} + + + +// check if two numbers are inside a given range +static int inRange(const mjtNum x1, const mjtNum x2, const mjtNum* range) { + return x1 >= range[0] && x1 <= range[1] && + x2 >= range[0] && x2 <= range[1]; +} + + + +// advance simulation using control callback, skipstage is mjtStage +void mj_stepSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor) { + TM_START; + + // common to all integrators + mj_checkPos(m, d); + mj_checkVel(m, d); + mj_forwardSkip(m, d, skipstage, skipsensor); + mj_checkAcc(m, d); + + // compare forward and inverse solutions if enabled + if (mjENABLED(mjENBL_FWDINV)) { + mj_compareFwdInv(m, d); + } + + // use selected integrator + switch(m->opt.integrator) { + case mjINT_EULER: + mj_EulerSkip(m, d, skipstage >= mjSTAGE_POS); + break; + + case mjINT_RK4: + // ignore skipstage + mj_RungeKutta(m, d, 4); + break; + + case mjINT_IMPLICIT: + mj_implicitSkip(m, d, skipstage >= mjSTAGE_VEL); + break; + + default: + mju_error("Invalid integrator"); + } + + TM_END(mjTIMER_STEP); +} + + + //------------------------- derivatives of component functions ------------------------------------- // derivative of cvel, cdof_dot w.r.t qvel @@ -776,39 +919,7 @@ static void mjd_actuator_vel(const mjModel* m, mjData* d, mjtNum* DfDv) { -//------------------------- main entry points ------------------------------------------------------ - -// Analytical derivative: -// d->qDeriv = d (qfrc_actuator + qfrc_passive - qfrc_bias) / d qvel. -void mjd_smooth_vel(const mjModel *m, mjData *d) { - int nv = m->nv; - - // allocate space - mjMARKSTACK; - mjtNum *DfDv = mj_stackAlloc(d, nv*nv); - - // clear DfDv - mju_zero(DfDv, nv*nv); - - // DfDv = d (qfrc_actuator + qfrc_passive - qfrc_bias) / d qvel - mjd_actuator_vel(m, d, DfDv); - mjd_passive_vel(m, d, DfDv); - mjd_rne_vel(m, d, DfDv); - - // copy dense DfDv to sparse qDeriv - for (int i=0; iD_rownnz[i]; j++) { - int adr = d->D_rowadr[i] + j; - d->qDeriv[adr] = DfDv[i*nv + d->D_colind[adr]]; - } - } - - mjFREESTACK; -} - - - -// Centered finite difference approximation to mj_derivative. +// centered finite difference approximation to mjd_smooth_vel void mjd_smooth_velFD(const mjModel* m, mjData* d, mjtNum eps) { int nv = m->nv; @@ -870,3 +981,315 @@ void mjd_smooth_velFD(const mjModel* m, mjData* d, mjtNum eps) { mjFREESTACK; } + + +//------------------------- main entry points ------------------------------------------------------ + +// analytical derivative of smooth forces w.r.t velocities: +// d->qDeriv = d (qfrc_actuator + qfrc_passive - qfrc_bias) / d qvel +void mjd_smooth_vel(const mjModel *m, mjData *d) { + int nv = m->nv; + + // allocate space + mjMARKSTACK; + mjtNum *DfDv = mj_stackAlloc(d, nv*nv); + + // clear DfDv + mju_zero(DfDv, nv*nv); + + // DfDv = d (qfrc_actuator + qfrc_passive - qfrc_bias) / d qvel + mjd_actuator_vel(m, d, DfDv); + mjd_passive_vel(m, d, DfDv); + mjd_rne_vel(m, d, DfDv); + + // copy dense DfDv to sparse qDeriv + for (int i=0; iD_rownnz[i]; j++) { + int adr = d->D_rowadr[i] + j; + d->qDeriv[adr] = DfDv[i*nv + d->D_colind[adr]]; + } + } + + mjFREESTACK; +} + + +// finite differenced Jacobian of (next_state, sensors) = mj_step(state, control) +// all outputs are optional +// output dimensions (transposed w.r.t common convention): +// DyDq: (nv x 2*nv+na) +// DyDv: (nv x 2*nv+na) +// DyDa: (na x 2*nv+na) +// DyDu: (nu x 2*nv+na) +// DsDq: (nv x nsensordata) +// DsDv: (nv x nsensordata) +// DsDa: (na x nsensordata) +// DsDu: (nu x nsensordata) +// single-letter shortcuts: +// inputs: q=qpos, v=qvel, a=act, u=ctrl +// outputs: y=next_state (concatenated next qpos, qvel, act), s=sensordata +void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, + mjtNum* DyDq, mjtNum* DyDv, mjtNum* DyDa, mjtNum* DyDu, + mjtNum* DsDq, mjtNum* DsDv, mjtNum* DsDa, mjtNum* DsDu) { + int nq = m->nq, nv = m->nv, na = m->na, nu = m->nu, ns = m->nsensordata; + int ndx = 2*nv+na; // row length of Dy Jacobians + mjMARKSTACK; + + // states + mjtNum *state = mj_stackAlloc(d, nq+nv+na); // current state + mjtNum *next = mj_stackAlloc(d, nq+nv+na); // next state + mjtNum *next_plus = mj_stackAlloc(d, nq+nv+na); // forward-nudged next state + mjtNum *next_minus = mj_stackAlloc(d, nq+nv+na); // backward-nudged next state + + // warmstart accelerations + mjtNum *warmstart = mjDISABLED(mjDSBL_WARMSTART) ? NULL : mj_stackAlloc(d, nv); + + // sensors + int skipsensor = !DsDq && !DsDv && !DsDa && !DsDu; + mjtNum *sensor = skipsensor ? NULL : mj_stackAlloc(d, ns); // sensor values + mjtNum *sensor_plus = skipsensor ? NULL : mj_stackAlloc(d, ns); // forward-nudged sensors + mjtNum *sensor_minus = skipsensor ? NULL : mj_stackAlloc(d, ns); // backward-nudged sensors + + // controls + mjtNum *ctrl = mj_stackAlloc(d, nu); + + // save current inputs + mju_copy(ctrl, d->ctrl, nu); + getState(m, d, state, NULL); + if (warmstart) { + mju_copy(warmstart, d->qacc_warmstart, nv); + } + + // step input + mj_step(m, d); + + // save output + getState(m, d, next, sensor); + + // restore input + setState(m, d, state, ctrl, warmstart); + + // finite-difference controls: skip=mjSTAGE_VEL, handle ctrl at range limits + if (DyDu || DsDu) { + for (int i=0; iactuator_ctrllimited[i]; + // nudge forward, if possible given ctrlrange + int nudge_fwd = !limited || inRange(ctrl[i], ctrl[i]+eps, m->actuator_ctrlrange+2*i); + if (nudge_fwd) { + // nudge forward + d->ctrl[i] += eps; + + // step, get nudged output + mj_stepSkip(m, d, mjSTAGE_VEL, skipsensor); + getState(m, d, next_plus, sensor_plus); + + // reset + setState(m, d, state, ctrl, warmstart); + } + + // nudge backward, if possible given ctrlrange + int nudge_back = (centered || !nudge_fwd) && + (!limited || inRange(ctrl[i]-eps, ctrl[i], m->actuator_ctrlrange+2*i)); + if (nudge_back) { + // nudge backward + d->ctrl[i] -= eps; + + // step, get nudged output + mj_stepSkip(m, d, mjSTAGE_VEL, skipsensor); + getState(m, d, next_minus, sensor_minus); + + // reset + setState(m, d, state, ctrl, warmstart); + } + + // difference states + if (DyDu) { + clampedStateDiff(m, DyDu+i*ndx, next, nudge_fwd ? next_plus : NULL, + nudge_back ? next_minus : NULL, eps); + } + + // difference sensors + if (DsDu) { + clampedDiff(DsDu+i*ns, sensor, nudge_fwd ? sensor_plus : NULL, + nudge_back ? sensor_minus : NULL, eps, ns); + } + } + } + + // finite-difference activations: skip=mjSTAGE_VEL + if (DyDa || DsDa) { + for (int i=0; iact[i] += eps; + + // step, get nudged output + mj_stepSkip(m, d, mjSTAGE_VEL, skipsensor); + getState(m, d, next_plus, sensor_plus); + + // reset + setState(m, d, state, NULL, warmstart); + + // nudge backward + if (centered) { + // nudge backward + d->act[i] -= eps; + + // step, get nudged output + mj_stepSkip(m, d, mjSTAGE_VEL, skipsensor); + getState(m, d, next_minus, sensor_minus); + + // reset + setState(m, d, state, NULL, warmstart); + } + + // difference states + if (DyDa) { + if (!centered) { + stateDiff(m, DyDa+i*ndx, next, next_plus, eps); + } else { + stateDiff(m, DyDa+i*ndx, next_minus, next_plus, 2*eps); + } + } + + // difference sensors + if (DsDa) { + if (!centered) { + diff(DsDa+i*ns, sensor, sensor_plus, eps, ns); + } else { + diff(DsDa+i*ns, sensor_minus, sensor_plus, 2*eps, ns); + } + } + } + } + + + // finite-difference velocities: skip=mjSTAGE_POS + if (DyDv || DsDv) { + for (int i=0; iqvel[i] += eps; + + // step, get nudged output + mj_stepSkip(m, d, mjSTAGE_POS, skipsensor); + getState(m, d, next_plus, sensor_plus); + + // reset + setState(m, d, state, NULL, warmstart); + + // nudge backward + if (centered) { + // nudge + d->qvel[i] -= eps; + + // step, get nudged output + mj_stepSkip(m, d, mjSTAGE_POS, skipsensor); + getState(m, d, next_minus, sensor_minus); + + // reset + setState(m, d, state, NULL, warmstart); + } + + // difference states + if (DyDv) { + if (!centered) { + stateDiff(m, DyDv+i*ndx, next, next_plus, eps); + } else { + stateDiff(m, DyDv+i*ndx, next_minus, next_plus, 2*eps); + } + } + + // difference sensors + if (DsDv) { + if (!centered) { + diff(DsDv+i*ns, sensor, sensor_plus, eps, ns); + } else { + diff(DsDv+i*ns, sensor_minus, sensor_plus, 2*eps, ns); + } + } + } + } + + // finite-difference positions: skip=mjSTAGE_NONE + if (DyDq || DsDq) { + mjtNum *dpos = mj_stackAlloc(d, nv); // allocate position perturbation + for (int i=0; iqpos, dpos, eps); + + // step, get nudged output + mj_stepSkip(m, d, mjSTAGE_NONE, skipsensor); + getState(m, d, next_plus, sensor_plus); + + // reset + setState(m, d, state, NULL, warmstart); + + // nudge backward + if (centered) { + // nudge backward + mju_zero(dpos, nv); + dpos[i] = 1; + mj_integratePos(m, d->qpos, dpos, -eps); + + // step, get nudged output + mj_stepSkip(m, d, mjSTAGE_NONE, skipsensor); + getState(m, d, next_minus, sensor_minus); + + // reset + setState(m, d, state, NULL, warmstart); + } + + // difference states + if (DyDq) { + if (!centered) { + stateDiff(m, DyDq+i*ndx, next, next_plus, eps); + } else { + stateDiff(m, DyDq+i*ndx, next_minus, next_plus, 2*eps); + } + } + + // difference sensors + if (DsDq) { + if (!centered) { + diff(DsDq+i*ns, sensor, sensor_plus, eps, ns); + } else { + diff(DsDq+i*ns, sensor_minus, sensor_plus, 2*eps, ns); + } + } + } + } + + mjFREESTACK; +} + + + +// finite differenced state-transition and control-transition matrices dy = A*dx + B*du +// required output matrix dimensions: +// A: (2*nv+na x 2*nv+na) +// B: (2*nv+na x nu) +void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, + mjtNum* A, mjtNum* B) { + int nv = m->nv, na = m->na, nu = m->nu; + int ndx = 2*nv+na; // row length of Jacobians + mjMARKSTACK; + + // allocate transposed matrices + mjtNum *AT = mj_stackAlloc(d, ndx*ndx); // state-transition matrix (transposed) + mjtNum *BT = B ? mj_stackAlloc(d, nu*ndx) : NULL; // control-transition matrix (transposed) + + // get Jacobians + if (A) { + mjd_stepFD(m, d, eps, centered, AT, AT+ndx*nv, AT+ndx*2*nv, BT, NULL, NULL, NULL, NULL); + } else { + mjd_stepFD(m, d, eps, centered, NULL, NULL, NULL, BT, NULL, NULL, NULL, NULL); + } + + // transpose + if (A) mju_transpose(A, AT, ndx, ndx); + if (B) mju_transpose(B, BT, nu, ndx); + + mjFREESTACK; +} diff --git a/src/engine/engine_derivative.h b/src/engine/engine_derivative.h index b7c6cd01..ee989571 100644 --- a/src/engine/engine_derivative.h +++ b/src/engine/engine_derivative.h @@ -24,7 +24,7 @@ extern "C" { #endif // analytical derivative of smooth forces w.r.t velocities: -// d->qDeriv = d (qfrc_actuator + qfrc_passive - qfrc_bias) / d qvel. +// d->qDeriv = d (qfrc_actuator + qfrc_passive - qfrc_bias) / d qvel MJAPI void mjd_smooth_vel(const mjModel* m, mjData* d); // centered finite difference approximation to mjd_smooth_vel @@ -36,6 +36,12 @@ MJAPI void mjd_passive_vel(const mjModel* m, mjData* d, mjtNum* DfDv); // add forward finite difference approximation of (d qfrc_passive / d qvel) to DfDv MJAPI void mjd_passive_velFD(const mjModel* m, mjData* d, mjtNum eps, mjtNum* DfDv); +// advance simulation using control callback, skipstage is mjtStage +MJAPI void mj_stepSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor); + +// finite differenced state-transition and control-transition matrices dy = A*dx + B*du +MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered, + mjtNum* A, mjtNum* B); #ifdef __cplusplus } diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 401610d1..1bfbaf2a 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -485,16 +485,11 @@ static void mj_advance(const mjModel* m, mjData* d, d->time += m->opt.timestep; } +// Euler integrator, semi-implicit in velocity, possibly skipping factorisation +void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) { - -// Euler integrator, semi-implicit in velocity -void mj_Euler(const mjModel* m, mjData* d) { int i, nv = m->nv, nM = m->nM; mjMARKSTACK; - mjtNum* saveM = mj_stackAlloc(d, nM); - mjtNum* saveLD = mj_stackAlloc(d, nM); - mjtNum* saveLDiagInv = mj_stackAlloc(d, nv); - mjtNum* saveLDiagSqrtInv = mj_stackAlloc(d, nv); mjtNum* qfrc = mj_stackAlloc(d, nv); mjtNum* qacc = mj_stackAlloc(d, nv); @@ -512,29 +507,22 @@ void mj_Euler(const mjModel* m, mjData* d) { // damping: integrate implicitly else { - // save M and factorization - mju_copy(saveM, d->qM, nM); - mju_copy(saveLD, d->qLD, nM); - mju_copy(saveLDiagInv, d->qLDiagInv, nv); - mju_copy(saveLDiagSqrtInv, d->qLDiagSqrtInv, nv); + if (!skipfactor) { + mjtNum* MhB = mj_stackAlloc(d, nM); - // add hB to diagonal of M - for (i=0; iqM[m->dof_Madr[i]] += m->opt.timestep * m->dof_damping[i]; + // MhB = M + h*diag(B) + mju_copy(MhB, d->qM, m->nM); + for (i=0; idof_Madr[i]] += m->opt.timestep * m->dof_damping[i]; + } + + // factor + mj_factorI(m, d, MhB, d->qH, d->qHDiagInv, 0); } - // factor - mj_factorM(m, d); - // solve mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv); - mj_solveM(m, d, qacc, qfrc, 1); - - // restore M and factorization - mju_copy(d->qM, saveM, nM); - mju_copy(d->qLD, saveLD, nM); - mju_copy(d->qLDiagInv, saveLDiagInv, nv); - mju_copy(d->qLDiagSqrtInv, saveLDiagSqrtInv, nv); + mj_solveLD(m, d, qacc, qfrc, 1, d->qH, d->qHDiagInv); } // advance state and time @@ -545,6 +533,13 @@ void mj_Euler(const mjModel* m, mjData* d) { +// 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, @@ -653,26 +648,28 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) { //-------------------------- top-level API --------------------------------------------------------- -// fully implicit in velocity -void mj_implicit(const mjModel *m, mjData *d) { +// fully implicit in velocity, possibly skipping factorization +void mj_implicitSkip(const mjModel *m, mjData *d, int skipfactor) { int nv = m->nv; mjMARKSTACK; mjtNum *qfrc = mj_stackAlloc(d, nv); mjtNum *qacc = mj_stackAlloc(d, nv); - // construct sparse structure in d->D_xxx - mj_makeMSparse(m, d, d->D_rownnz, d->D_rowadr, d->D_colind); + if (!skipfactor) { + // construct sparse structure in d->D_xxx + mj_makeMSparse(m, d, d->D_rownnz, d->D_rowadr, d->D_colind); - // compute analytical derivative qDeriv - mjd_smooth_vel(m, d); + // compute analytical derivative qDeriv + mjd_smooth_vel(m, d); - // set qLU = qM - dt*qDeriv - mj_setMSparse(m, d, d->qLU, d->D_rownnz, d->D_rowadr, d->D_colind); - mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD); + // set qLU = qM - dt*qDeriv + mj_setMSparse(m, d, d->qLU, d->D_rownnz, d->D_rowadr, d->D_colind); + mju_addToScl(d->qLU, d->qDeriv, -m->opt.timestep, m->nD); - // factorize qLU, use qacc as scratch space - mju_factorLUSparse(d->qLU, nv, (int*)qacc, d->D_rownnz, d->D_rowadr, d->D_colind); + // factorize qLU, use qacc as scratch space + mju_factorLUSparse(d->qLU, nv, (int*)qacc, d->D_rownnz, d->D_rowadr, d->D_colind); + } // set qfrc = qfrc_smooth + qfrc_constraint mju_add(qfrc, d->qfrc_smooth, d->qfrc_constraint, nv); @@ -688,6 +685,13 @@ void mj_implicit(const mjModel *m, mjData *d) { +// fully implicit in velocity +void mj_implicit(const mjModel *m, mjData *d) { + mj_implicitSkip(m, d, 0); +} + + + // forward dynamics with skip; skipstage is mjtStage void mj_forwardSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor) { TM_START; diff --git a/src/engine/engine_forward.h b/src/engine/engine_forward.h index d75e64e4..4b2f7e6b 100644 --- a/src/engine/engine_forward.h +++ b/src/engine/engine_forward.h @@ -43,21 +43,27 @@ MJAPI void mj_step2(const mjModel* m, mjData* d); MJAPI void mj_forward(const mjModel* m, mjData* d); // forward dynamics with skip; skipstage is mjtStage -MJAPI void mj_forwardSkip(const mjModel* m, mjData* d, - int skipstage, int skipsensor); +MJAPI void mj_forwardSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor); + //-------------------------------- integrators ----------------------------------------------------- -// Euler integrator, semi-implicit in velocity -MJAPI void mj_Euler(const mjModel* m, mjData* d); - // Runge Kutta explicit order-N integrator MJAPI void mj_RungeKutta(const mjModel* m, mjData* d, int N); +// Euler integrator, semi-implicit in velocity +MJAPI void mj_Euler(const mjModel* m, mjData* d); + +// Euler integrator, semi-implicit in velocity, possibly skipping factorisation +MJAPI void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor); + // fully implicit in velocity MJAPI void mj_implicit(const mjModel *m, mjData *d); +// fully implicit in velocity, possibly skipping factorization +MJAPI void mj_implicitSkip(const mjModel *m, mjData *d, int skipfactor); + //-------------------------------- solver components ----------------------------------------------- diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc index 933de44a..cd94504e 100644 --- a/test/engine/engine_derivative_test.cc +++ b/test/engine/engine_derivative_test.cc @@ -21,6 +21,7 @@ #include #include #include +#include "src/engine/engine_core_smooth.h" #include "src/engine/engine_derivative.h" #include "src/engine/engine_io.h" #include "src/engine/engine_support.h" @@ -33,10 +34,12 @@ namespace { using ::testing::Pointwise; using ::testing::DoubleNear; +using ::testing::Eq; +using ::testing::Each; using DerivativeTest = MujocoTest; // errors smaller than this are ignored -static const mjtNum absolute_tolerance = 1e-7; +static const mjtNum absolute_tolerance = 1e-9; // corrected relative error static mjtNum RelativeError(mjtNum a, mjtNum b) { @@ -46,10 +49,12 @@ static mjtNum RelativeError(mjtNum a, mjtNum b) { } // expect two 2D arrays to have elementwise relative error smaller than eps -static void CompareMatrices(mjtNum* Actual, mjtNum* Expected, - int nrow, int ncol, mjtNum eps) { - for (int i=0; i AsVector(const mjtNum* array, int n) { return std::vector(array, array + n); } @@ -73,7 +93,10 @@ static const char* const kDampedActuatorsPath = "engine/testdata/derivative/damped_actuators.xml"; static const char* const kDamperActuatorsPath = "engine/testdata/damper.xml"; - +static const char* const kDampedPendulumPath = + "engine/testdata/derivative/damped_pendulum.xml"; +static const char* const kLinearPath = + "engine/testdata/derivative/linear.xml"; // compare analytic and finite-difference d_smooth/d_qvel TEST_F(DerivativeTest, SmoothDvel) { // run test on all models @@ -92,9 +115,9 @@ TEST_F(DerivativeTest, SmoothDvel) { // take 100 steps so we have some velocities, then call forward mj_resetData(model, data); if (model->nu) { - data->ctrl[0]=0.1; + data->ctrl[0] = 0.1; } - for (int i=0; i<100; i++) { + for (int i=0; i < 100; i++) { mj_step(model, data); } mj_forward(model, data); @@ -141,7 +164,7 @@ TEST_F(DerivativeTest, PassiveDvel) { // take 100 steps so we have some velocities, then call forward mj_resetData(model, data); - for (int i=0; i<100; i++) { + for (int i=0; i < 100; i++) { mj_step(model, data); } mj_forward(model, data); @@ -163,7 +186,249 @@ TEST_F(DerivativeTest, PassiveDvel) { mju_free(DfDv_analytic); mj_deleteData(data); mj_deleteModel(model); +} +// ----------------------- derivatives of mj_step() ---------------------------- + +// mj_stepSkip computes the same next state as mj_step +TEST_F(DerivativeTest, StepSkip) { + const std::string xml_path = GetTestDataFilePath(kDampedPendulumPath); + mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + mjData* data = mj_makeData(model); + int nq = model->nq; + int nv = model->nv; + + // disable warmstarts so we don't need to save qacc_warmstart + model->opt.disableflags |= mjDSBL_WARMSTART; + + for (const mjtIntegrator integrator : {mjINT_EULER, mjINT_IMPLICIT}) { + model->opt.integrator = integrator; + + // reset, take 20 steps, save initial state + mj_resetData(model, data); + for (int i=0; i < 20; i++) { + mj_step(model, data); + } + std::vector qpos = AsVector(data->qpos, nq); + std::vector qvel = AsVector(data->qvel, nv); + + // take one more step, save next state + mj_step(model, data); + std::vector qpos_next = AsVector(data->qpos, nq); + std::vector qvel_next = AsVector(data->qvel, nv); + + // reset state, take step again, compare (assert mj_step is deterministic) + mju_copy(data->qpos, qpos.data(), nq); + mju_copy(data->qvel, qvel.data(), nv); + mj_step(model, data); + EXPECT_THAT(AsVector(data->qpos, nq), Pointwise(Eq(), qpos_next)); + EXPECT_THAT(AsVector(data->qvel, nv), Pointwise(Eq(), qvel_next)); + + // reset state, change ctrl, call mj_stepSkip, save next state + mju_copy(data->qpos, qpos.data(), nq); + mju_copy(data->qvel, qvel.data(), nv); + data->ctrl[0] = 1; + mj_stepSkip(model, data, mjSTAGE_VEL, 0); // skipping both POS and VEL + std::vector qpos_next_dctrl = AsVector(data->qpos, nq); + std::vector qvel_next_dctrl = AsVector(data->qvel, nv); + + // reset state (ctrl remains unchanged), call full mj_step, compare + mju_copy(data->qpos, qpos.data(), nq); + mju_copy(data->qvel, qvel.data(), nv); + mj_step(model, data); + EXPECT_THAT(AsVector(data->qpos, nq), Pointwise(Eq(), qpos_next_dctrl)); + EXPECT_THAT(AsVector(data->qvel, nv), Pointwise(Eq(), qvel_next_dctrl)); + + // reset state, change velocity, call mj_stepSkip, save next state + mju_copy(data->qpos, qpos.data(), nq); + mju_copy(data->qvel, qvel.data(), nv); + data->qvel[0] += 1; + mj_stepSkip(model, data, mjSTAGE_POS, 0); // skipping POS + std::vector qpos_next_dvel = AsVector(data->qpos, nq); + std::vector qvel_next_dvel = AsVector(data->qvel, nv); + + // reset state, change velocity, call full mj_step, compare + mju_copy(data->qpos, qpos.data(), nq); + mju_copy(data->qvel, qvel.data(), nv); + data->qvel[0] += 1; + mj_step(model, data); + EXPECT_THAT(AsVector(data->qpos, nq), Pointwise(Eq(), qpos_next_dvel)); + EXPECT_THAT(AsVector(data->qvel, nv), Pointwise(Eq(), qvel_next_dvel)); + } + + mj_deleteData(data); + mj_deleteModel(model); +} + + +// Analytic transition matrices for linear dynamical system xn = A*x + B*u +// given modified mass matrix H (`data->qH`) and +// Ac = H^-1 [diag(-stiffness) diag(-damping)] +// we have +// A = eye(2*nv) + dt [dt*Ac + [zeros(3) eye(3)]; Ac] +// given the moment arm matrix K (`data->actuator_moment`) and Bc = H^-1 K +// B = dt*[Bc*dt; Bc] +static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) { + int nv = m->nv, nu = m->nu; + mjtNum dt = m->opt.timestep; + mjMARKSTACK; + + // === state-transition matrix A + if (A) { + mjtNum *Ac = mj_stackAlloc(d, 2*nv*nv); + // Ac = H^-1 [diag(-stiffness) diag(-damping)] + mju_zero(Ac, 2*nv*nv); + for (int i=0; i < nv; i++) { + Ac[i*nv + i] = -m->jnt_stiffness[i]; + Ac[nv*nv + i*nv + i] = -m->dof_damping[i]; + } + mj_solveLD(m, d, Ac, Ac, 2*nv, d->qH, d->qHDiagInv); + + // A = [dt*Ac; Ac] + mju_transpose(A, Ac, 2*nv, nv); + mju_scl(A, A, dt, nv*2*nv); + mju_transpose(A+2*nv*nv, Ac, 2*nv, nv); + + // Add eye(nv) to top right quadrant of A + for (int i=0; i < nv; i++) { + A[i*2*nv + nv + i] += 1; + } + + // A *= dt + mju_scl(A, A, dt, 2*nv*2*nv); + + // A += eye(2*nv) + for (int i=0; i < 2*nv; i++) { + A[i*2*nv + i] += 1; + } + } + + // === control-transition matrix B + if (B) { + mjtNum *Bc = mj_stackAlloc(d, nu*nv); + mjtNum *BcT = mj_stackAlloc(d, nv*nu); + mj_solveLD(m, d, Bc, d->actuator_moment, nu, d->qH, d->qHDiagInv); + mju_transpose(BcT, Bc, nu, nv); + mju_scl(B, BcT, dt*dt, nu*nv); + mju_scl(B+nu*nv, BcT, dt, nu*nv); + } + + mjFREESTACK; +} + +// compare FD derivatives to analytic derivatives of linear dynamical system +TEST_F(DerivativeTest, LinearSystem) { + const std::string xml_path = GetTestDataFilePath(kLinearPath); + mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + mjData* data = mj_makeData(model); + int nv = model->nv, nu = model->nu; + + // set ctrl, integrate for 20 steps + data->ctrl[0] = .1; + data->ctrl[1] = -.1; + for (int i=0; i < 20; i++) { + mj_step(model, data); + } + + // analytic A and B + mjtNum* A = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*2*nv); + mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*nu); + + LinearSystem(model, data, A, B); + + PrintMatrix(A, 2*nv, 2*nv); + PrintMatrix(B, 2*nv, nu); + + // forward differenced A and B + mjtNum eps = 1e-6; + mjtNum* AFD = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*2*nv); + mjtNum* BFD = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*nu); + + mjd_transitionFD(model, data, eps, /*centered=*/0, AFD, BFD); + + PrintMatrix(AFD, 2*nv, 2*nv); + PrintMatrix(BFD, 2*nv, nu); + + // expect FD and analytic derivatives to be similar to eps precision + CompareMatrices(A, AFD, 2*nv, 2*nv, eps); + CompareMatrices(B, BFD, 2*nv, nu, eps); + + // central differenced A and B + mjtNum* AFDc = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*2*nv); + mjtNum* BFDc = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*nu); + mjd_transitionFD(model, data, eps, /*centered=*/1, AFDc, BFDc); + + // expect central derivatives to be closer to analytic solution + EXPECT_LT(CompareMatrices(A, AFDc, 2*nv, 2*nv, eps), + CompareMatrices(A, AFD, 2*nv, 2*nv, eps)); + EXPECT_LT(CompareMatrices(B, BFDc, 2*nv, nu, eps), + CompareMatrices(B, BFD, 2*nv, nu, eps)); + + mju_free(BFDc); + mju_free(AFDc); + mju_free(BFD); + mju_free(AFD); + mju_free(B); + mju_free(A); + mj_deleteData(data); + mj_deleteModel(model); +} + +// check ctrl derivatives at the range limit +TEST_F(DerivativeTest, ClampedCtrlDerivatives) { + const std::string xml_path = GetTestDataFilePath(kLinearPath); + mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, nullptr, 0); + mjData* data = mj_makeData(model); + int nv = model->nv, nu = model->nu; + + // set ctrl, integrate for 20 steps + data->ctrl[0] = .1; + data->ctrl[1] = -.1; + for (int i=0; i < 20; i++) { + mj_step(model, data); + } + + // analytic B + mjtNum* B = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*nu); + + LinearSystem(model, data, NULL, B); + + // forward differenced A and B + mjtNum eps = 1e-6; + mjtNum* BFD = (mjtNum*) mju_malloc(sizeof(mjtNum)*2*nv*nu); + + // set ctrl to the limits, request forward differences + data->ctrl[0] = 1; + data->ctrl[1] = -1; + mjd_transitionFD(model, data, eps, /*centered=*/0, NULL, BFD); + // expect FD and analytic derivatives to be similar to eps precision + CompareMatrices(B, BFD, 2*nv, nu, eps); + + // ctrl remains at limits, request central differences + mjd_transitionFD(model, data, eps, /*centered=*/1, NULL, BFD); + // expect FD and analytic derivatives to be similar to eps precision + CompareMatrices(B, BFD, 2*nv, nu, eps); + + // set ctrl beyond limits, request forward differences + data->ctrl[0] = 2; + data->ctrl[1] = -2; + mjd_transitionFD(model, data, eps, /*centered=*/0, NULL, BFD); + // expect derivatives to be 0 + EXPECT_THAT(AsVector(BFD, 2*nv*nu), Each(Eq(0.0))); + + // expect ctrl to remain unchanged (despite intenal clamping) + EXPECT_EQ(data->ctrl[0], 2.0); + EXPECT_EQ(data->ctrl[1], -2.0); + + // ctrl remains beyond limits, request centered differences + mjd_transitionFD(model, data, eps, /*centered=*/1, NULL, BFD); + // expect derivatives to be 0 + EXPECT_THAT(AsVector(BFD, 2*nv*nu), Each(Eq(0.0))); + + mju_free(BFD); + mju_free(B); + mj_deleteData(data); + mj_deleteModel(model); } } // namespace diff --git a/test/engine/testdata/derivative/damped_pendulum.xml b/test/engine/testdata/derivative/damped_pendulum.xml new file mode 100644 index 00000000..663f224a --- /dev/null +++ b/test/engine/testdata/derivative/damped_pendulum.xml @@ -0,0 +1,31 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/derivative/linear.xml b/test/engine/testdata/derivative/linear.xml new file mode 100644 index 00000000..d136c9d2 --- /dev/null +++ b/test/engine/testdata/derivative/linear.xml @@ -0,0 +1,28 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + +