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:
committed by
Copybara-Service
parent
2ea01bf2f6
commit
228264c92b
@@ -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)
|
||||
|
||||
@@ -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 ) \
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
)),
|
||||
])
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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
@@ -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;
|
||||
|
||||
@@ -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 -----------------------------------------------
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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>
|
||||
Reference in New Issue
Block a user