Add efficient finite-difference Jacobians of mj_step.

- Add `qH` and `qHDiagInv` to `mjData` to save factorized modified inertia.
- Add `mj_EulerSkip`, `mj_implicitSkip`, to `engine_forward.c`.
- Using the above functions, implement `mj_stepSkip` in `engine_derivative.c`.
- Add `mjd_stepFD` and `mjd_transitionFD` to `engine_derivative.c` to compute `mj_step` Jacobians.
  - Exploit "Skip" functionality for speed.
  - Correctly handle quaternion derivatives.
  - Handle warmstarts and control limits.

PiperOrigin-RevId: 456584811
Change-Id: Iee8541f11e7b66feb8f431cb102d9bbe65461f79
This commit is contained in:
Yuval Tassa
2022-06-22 12:48:04 -07:00
committed by Copybara-Service
parent 2ea01bf2f6
commit 228264c92b
13 changed files with 938 additions and 96 deletions
+4
View File
@@ -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)
+2
View File
@@ -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 ) \
+8 -1
View File
@@ -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
+40
View File
@@ -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
)),
])
+26 -8
View File
@@ -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; i<nv; i++) {
d->qLDiagInv[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
+8
View File
@@ -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
+459 -36
View File
@@ -19,17 +19,15 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#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; i<n; i++) {
dx[i] = inv_h * (x2[i] - x1[i]);
}
}
// finite-difference two state vectors ds = (s2 - s1) / h
static void stateDiff(const mjModel* m, mjtNum* ds, const mjtNum* s1, const mjtNum* s2, mjtNum h) {
int nq = m->nq, 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; i<nv; i++) {
for (int j=0; j<d->D_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; i<nv; i++) {
for (int j=0; j<d->D_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; i<nu; i++) {
int limited = m->actuator_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; i<na; i++) {
// nudge forward
d->act[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; i<nv; i++) {
// nudge forward
d->qvel[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; i<nv; i++) {
// nudge forward
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_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;
}
+7 -1
View File
@@ -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
}
+40 -36
View File
@@ -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; i<nv; i++) {
d->qM[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; i<nv; i++) {
MhB[m->dof_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;
+11 -5
View File
@@ -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 -----------------------------------------------
+274 -9
View File
@@ -21,6 +21,7 @@
#include <gtest/gtest.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mujoco.h>
#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<nrow; i++) {
for (int j=0; j<ncol; j++) {
// return maximum absolute error
static mjtNum CompareMatrices(mjtNum* Actual, mjtNum* Expected,
int nrow, int ncol, mjtNum eps) {
mjtNum max_error = 0;
for (int i=0; i < nrow; i++) {
for (int j=0; j < ncol; j++) {
mjtNum actual = Actual[i*ncol+j];
mjtNum expected = Expected[i*ncol+j];
EXPECT_LT(RelativeError(actual, expected), eps)
@@ -57,10 +62,25 @@ static void CompareMatrices(mjtNum* Actual, mjtNum* Expected,
<< "\nexpected = " << expected
<< "\nactual = " << actual
<< "\ndiff = " << expected-actual;
max_error = mjMAX(mju_abs(actual-expected), max_error);
}
}
return max_error;
}
// utility function for matrix printing
static void PrintMatrix(mjtNum* mat, int nrow, int ncol) {
std::cerr.precision(5);
std::cerr << "\n";
for (int r=0; r < nrow; r++) {
for (int c=0; c < ncol; c++) {
std::cerr << std::fixed << std::setw(9) << mat[c + r*ncol] << " ";
}
std::cerr << "\n";
}
}
std::vector<mjtNum> AsVector(const mjtNum* array, int n) {
return std::vector<mjtNum>(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<mjtNum> qpos = AsVector(data->qpos, nq);
std::vector<mjtNum> qvel = AsVector(data->qvel, nv);
// take one more step, save next state
mj_step(model, data);
std::vector<mjtNum> qpos_next = AsVector(data->qpos, nq);
std::vector<mjtNum> 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<mjtNum> qpos_next_dctrl = AsVector(data->qpos, nq);
std::vector<mjtNum> 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<mjtNum> qpos_next_dvel = AsVector(data->qpos, nq);
std::vector<mjtNum> 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
+31
View File
@@ -0,0 +1,31 @@
<mujoco>
<default>
<joint damping=".01" frictionloss=".1"/>
</default>
<worldbody>
<light pos="0 0 1"/>
<geom type="plane" size="1 1 .01" pos="0 0 -1"/>
<body pos="0.15 0 0">
<joint name="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
<body pos="0.1 0 0">
<joint type="slide" axis="1 0 0" stiffness="200"/>
<geom type="capsule" size="0.015" fromto="-.1 0 0 .1 0 0"/>
<body pos=".1 0 0">
<joint type="ball"/>
<geom type="box" size=".02" fromto="0 0 0 0 .1 0"/>
<body pos="0 .1 0">
<joint axis="1 0 0"/>
<geom type="capsule" size="0.02" fromto="0 0 0 0 0 .1"/>
</body>
</body>
</body>
</body>
</worldbody>
<actuator>
<motor joint="hinge"/>
</actuator>
</mujoco>
+28
View File
@@ -0,0 +1,28 @@
<mujoco>
<!-- linear dynamical system with fixed, anaylticaly computable derivatives -->
<default>
<joint axis="1 0 0" damping=".01" stiffness=".1"/>
<geom type="box" size=".05 .05 .05"/>
</default>
<worldbody>
<light pos="0 0 1"/>
<body>
<joint name="joint0"/>
<geom/>
<body pos=".15 0 0">
<joint name="joint1"/>
<geom/>
<body pos=".15 0 0">
<joint/>
<geom/>
</body>
</body>
</body>
</worldbody>
<actuator>
<motor joint="joint0" ctrllimited="true" ctrlrange="-1 1"/>
<motor joint="joint1" ctrllimited="true" ctrlrange="-1 1"/>
</actuator>
</mujoco>