chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Remove mjData.qLDiagSqrtInv, add corresponding argument to mj_solveM2.

Browse Source 
- `qLDiagSqrtInv` is only required for the dual solvers. It is now computed as-needed rather than unconditionally.
- `mj_solveM2` now requires a new input array `sqrtInvD` which contains the square root of the inverse diagonal D (formerly saved in `qLDiagSqrtInv`).

PiperOrigin-RevId: 710805133
Change-Id: I0622d6a8da3882916824e9c10bad9223c122c321
This commit is contained in:
Yuval Tassa
2024-12-30 15:23:41 -08:00
committed by Copybara-Service
parent ec322641b7
commit 69c9ac074a
19 changed files with 53 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+9
View File
@@ -13,6 +13,15 @@ Python bindings
from multiple threads. To run multiple threaded rollouts simultaneously, use the new class ``Rollout`` which
encapsulates the thread pool. Contribution by :github:user:`aftersomemath`.
General
^^^^^^^
.. admonition:: Breaking API changes (minor)
:class: attention
- The field ``mjData.qLDiagSqrtInv`` has been removed. This field is only required for the dual solvers. It is now
computed as-needed rather than unconditionally. Relatedly, added the corresponding argument to :ref:`mj_solveM2`.
Bug fixes
^^^^^^^^^
- Fixed a bug in the box-sphere collider, depth was incorrect for deep penetrations (:github:issue:`2206`).
doc/includes/references.h
+2 -2
View File
@@ -273,7 +273,6 @@ struct mjData_ {
// computed by mj_fwdPosition/mj_factorM
mjtNum* qLD; // L'*D*L factorization of M (sparse) (nM x 1)
mjtNum* qLDiagInv; // 1/diag(D) (nv x 1)
mjtNum* qLDiagSqrtInv; // 1/sqrt(diag(D)) (nv x 1)
// computed by mj_collisionTree
mjtNum* bvh_aabb_dyn; // global bounding box (center, size) (nbvhdynamic x 6)
@@ -3221,7 +3220,8 @@ void mj_transmission(const mjModel* m, mjData* d);
void mj_crb(const mjModel* m, mjData* d);
void mj_factorM(const mjModel* m, mjData* d);
void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
const mjtNum* sqrtInvD, int n);
void mj_comVel(const mjModel* m, mjData* d);
void mj_passive(const mjModel* m, mjData* d);
void mj_subtreeVel(const mjModel* m, mjData* d);
include/mujoco/mjdata.h
-1
View File
@@ -301,7 +301,6 @@ struct mjData_ {
// computed by mj_fwdPosition/mj_factorM
mjtNum* qLD; // L'*D*L factorization of M (sparse) (nM x 1)
mjtNum* qLDiagInv; // 1/diag(D) (nv x 1)
mjtNum* qLDiagSqrtInv; // 1/sqrt(diag(D)) (nv x 1)
// computed by mj_collisionTree
mjtNum* bvh_aabb_dyn; // global bounding box (center, size) (nbvhdynamic x 6)
include/mujoco/mjxmacro.h
-1
View File
@@ -632,7 +632,6 @@
X ( mjtNum, qM, nM, 1 ) \
X ( mjtNum, qLD, nM, 1 ) \
X ( mjtNum, qLDiagInv, nv, 1 ) \
X ( mjtNum, qLDiagSqrtInv, nv, 1 ) \
XMJV( mjtNum, bvh_aabb_dyn, nbvhdynamic, 6 ) \
XMJV( mjtByte, bvh_active, nbvh, 1 ) \
X ( mjtNum, flexedge_velocity, nflexedge, 1 ) \
include/mujoco/mujoco.h
+2 -1
View File
@@ -365,7 +365,8 @@ MJAPI void mj_factorM(const mjModel* m, mjData* d);
MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
// Half of linear solve: x = sqrt(inv(D))*inv(L')*y
MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
const mjtNum* sqrtInvD, int n);
// Compute cvel, cdof_dot.
MJAPI void mj_comVel(const mjModel* m, mjData* d);
introspect/functions.py
+6
View File
@@ -1868,6 +1868,12 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='sqrtInvD',
type=PointerType(
inner_type=ValueType(name='mjtNum', is_const=True),
),
),
FunctionParameterDecl(
name='n',
type=ValueType(name='int'),
introspect/structs.py
-8
View File
@@ -5213,14 +5213,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
doc='1/diag(D)',
array_extent=('nv',),
),
StructFieldDecl(
name='qLDiagSqrtInv',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='1/sqrt(diag(D))',
array_extent=('nv',),
),
StructFieldDecl(
name='bvh_aabb_dyn',
type=PointerType(
mjx/mujoco/mjx/_src/io.py
-1
View File
@@ -310,7 +310,6 @@ def make_data(
'qM': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float),
'qLD': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float),
'qLDiagInv': (m.nv, float) if support.is_sparse(m) else (0, float),
'qLDiagSqrtInv': (m.nv, float),
'bvh_aabb_dyn': (m.nbvhdynamic, 6, float),
'bvh_active': (m.nbvh, jp.uint8),
'flexedge_velocity': (m.nflexedge, float),
mjx/mujoco/mjx/_src/types.py
-2
View File
@@ -1274,7 +1274,6 @@ class Data(PyTreeNode):
if dense: (nv, nv)
qLDiagInv: 1/diag(D) if sparse: (nv,)
if dense: (0,)
qLDiagSqrtInv: 1/sqrt(diag(D)) (nv,)
bvh_aabb_dyn: global bounding box (center, size) (nbvhdynamic, 6)
bvh_active: volume has been added to collisions (nbvh,)
flexedge_velocity: flex edge velocities (nflexedge,)
@@ -1404,7 +1403,6 @@ class Data(PyTreeNode):
qM: jax.Array # pylint:disable=invalid-name
qLD: jax.Array # pylint:disable=invalid-name
qLDiagInv: jax.Array # pylint:disable=invalid-name
qLDiagSqrtInv: jax.Array # pylint:disable=invalid-name
bvh_aabb_dyn: jax.Array = _restricted_to('mujoco')
bvh_active: jax.Array = _restricted_to('mujoco')
# position, velocity dependent:
python/mujoco/functions.cc
+6 -2
View File
@@ -234,7 +234,7 @@ PYBIND11_MODULE(_functions, pymodule) {
DEF_WITH_OMITTED_PY_ARGS(traits::mj_solveM2, "n")(
pymodule,
[](const raw::MjModel* m, raw::MjData* d, Eigen::Ref<EigenArrayXX> x,
Eigen::Ref<const EigenArrayXX> y) {
Eigen::Ref<const EigenArrayXX> y, Eigen::Ref<const EigenArrayXX> sqrtInvD) {
if (x.rows() != y.rows()) {
throw py::type_error(
"the first dimension of x and y should be of the same size");
@@ -247,8 +247,12 @@ PYBIND11_MODULE(_functions, pymodule) {
throw py::type_error(
"the last dimension of y should be of size nv");
}
if (sqrtInvD.size() != m->nv) {
throw py::type_error(
"the size of sqrtInvD should be nv");
}
return InterceptMjErrors(::mj_solveM2)(
m, d, x.data(), y.data(), y.rows());
m, d, x.data(), y.data(), sqrtInvD.data(), y.rows());
});
Def<traits::mj_comVel>(pymodule);
Def<traits::mj_passive>(pymodule);
python/mujoco/indexer_xmacro.h
-1
View File
@@ -379,7 +379,6 @@
X( mjtNum, , xaxis, njnt, 3 ) \
X( mjtNum, , cdof, nv, 6 ) \
X( mjtNum, , qLDiagInv, nv, 1 ) \
X( mjtNum, , qLDiagSqrtInv, nv, 1 ) \
X( mjtNum, , cdof_dot, nv, 6 ) \
X( mjtNum, , qfrc_bias, nv, 1 ) \
X( mjtNum, , qfrc_passive, nv, 1 ) \
src/engine/engine_core_constraint.c
+8 -2
View File
@@ -2067,6 +2067,12 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
mj_markStack(d);
// inverse square root of D from inertia LDL decomposition
mjtNum* sqrtInvD = mjSTACKALLOC(d, nv, mjtNum);
for (int i=0; i < nv; i++) {
sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]);
}
// space for backsubM2(J')' and its traspose
mjtNum* JM2 = mjSTACKALLOC(d, nefc*nv, mjtNum);
mjtNum* JM2T = mjSTACKALLOC(d, nv*nefc, mjtNum);
@@ -2140,7 +2146,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
// process if not zero
if (xi) {
// x(i) /= sqrt(L(i,i))
JM2[adr+i] *= d->qLDiagSqrtInv[colind[adr+i]];
JM2[adr+i] *= sqrtInvD[colind[adr+i]];
// x(j) -= L(i,j) * x(i)
int Madr_ij = m->dof_Madr[colind[adr+i]]+1;
@@ -2191,7 +2197,7 @@ void mj_projectConstraint(const mjModel* m, mjData* d) {
// dense
else {
// JM2 = backsubM2(J')'
mj_solveM2(m, d, JM2, d->efc_J, nefc);
mj_solveM2(m, d, JM2, d->efc_J, sqrtInvD, nefc);
// construct JM2T
mju_transpose(JM2T, JM2, nefc, nv);
src/engine/engine_core_smooth.c
+7 -11
View File
@@ -774,7 +774,7 @@ void mj_tendon(const mjModel* m, mjData* d) {
L[i] += (mju_dist3(wpnt, wpnt+3) + wlen + mju_dist3(wpnt+6, wpnt+9))/divisor;
}
// accumulate moments if consequtive points are in different bodies
// accumulate moments if consecutive points are in different bodies
for (int k=0; k < (wlen < 0 ? 1 : 3); k++) {
if (wbody[k] != wbody[k+1]) {
// get 3D position difference, normalize
@@ -1387,8 +1387,7 @@ void mj_crb(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) {
void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv) {
int cnt;
int Madr_kk, Madr_ki;
mjtNum tmp;
@@ -1445,9 +1444,6 @@ void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNu
for (int i=0; i < nv; i++) {
mjtNum qLDi = qLD[dof_Madr[i]];
qLDiagInv[i] = 1.0/qLDi;
if (qLDiagSqrtInv) {
qLDiagSqrtInv[i] = 1.0/mju_sqrt(qLDi);
}
}
}
@@ -1456,7 +1452,7 @@ void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNu
// sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
void mj_factorM(const mjModel* m, mjData* d) {
TM_START;
mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv, d->qLDiagSqrtInv);
mj_factorI(m, d, d->qM, d->qLD, d->qLDiagInv);
TM_ADD(mjTIMER_POS_INERTIA);
}
@@ -1685,10 +1681,10 @@ void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* restrict x, int
// 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) {
void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
const mjtNum* sqrtInvD, int n) {
// local copies of key variables
mjtNum* qLD = d->qLD;
mjtNum* qLDiagSqrtInv = d->qLDiagSqrtInv;
int* dof_Madr = m->dof_Madr;
int* dof_parentid = m->dof_parentid;
int nv = m->nv;
@@ -1720,7 +1716,7 @@ void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n)
// x <- sqrt(inv(D)) * x
for (int i=0; i < nv; i++) {
x[i+offset] *= qLDiagSqrtInv[i]; // x(i) /= sqrt(L(i,i))
x[i+offset] *= sqrtInvD[i]; // x(i) /= sqrt(L(i,i))
}
}
}
@@ -1781,7 +1777,7 @@ void mj_comVel(const mjModel* m, mjData* d) {
default:
// in principle we should use the new velocity to compute cdofdot,
// but it makes no difference becase crossMotion(cdof, cdof) = 0,
// but it makes no difference because crossMotion(cdof, cdof) = 0,
// and using the old velocity may be more accurate numerically
mju_crossMotion(cdofdot+6*j, cvel, d->cdof+6*(bda+j));
src/engine/engine_core_smooth.h
+3 -3
View File
@@ -49,8 +49,7 @@ MJAPI void mj_transmission(const mjModel* m, mjData* d);
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);
MJAPI void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv);
// sparse L'*D*L factorizaton of the inertia matrix M, assumed spd
MJAPI void mj_factorM(const mjModel* m, mjData* d);
@@ -71,7 +70,8 @@ MJAPI void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, in
MJAPI void mj_solveM_island(const mjModel* m, const mjData* d, mjtNum* x, int island);
// half of sparse backsubstitution: x = sqrt(inv(D))*inv(L')*y
MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
const mjtNum* sqrtInvD, int n);
//-------------------------- velocity --------------------------------------------------------------
src/engine/engine_forward.c
+2 -2
View File
@@ -803,7 +803,7 @@ void mj_EulerSkip(const mjModel* m, mjData* d, int skipfactor) {
}
// factor
mj_factorI(m, d, MhB, d->qH, d->qHDiagInv, 0);
mj_factorI(m, d, MhB, d->qH, d->qHDiagInv);
}
// solve
@@ -986,7 +986,7 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
mju_addScl(MhB, d->qM, MhB, -m->opt.timestep, nM);
// factorize
mj_factorI(m, d, MhB, d->qH, d->qHDiagInv, NULL);
mj_factorI(m, d, MhB, d->qH, d->qHDiagInv);
}
// solve for qacc: (qM - dt*qDeriv) * qacc = qfrc
src/engine/engine_print.c
-1
View File
@@ -1107,7 +1107,6 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
}
printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format);
printArray("QLDIAGSQRTINV", m->nv, 1, d->qLDiagSqrtInv, fp, float_format);
// B sparse structure
printSparsity("B: body-dof matrix", m->nbody, m->nv, d->B_rowadr, NULL, d->B_rownnz, NULL,
src/engine/engine_support.c
+2 -3
View File
@@ -1087,7 +1087,6 @@ void mj_mulM_island(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum
void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec) {
int nv = m->nv;
const mjtNum* qLD = d->qLD;
const mjtNum* qLDiagSqrtInv = d->qLDiagSqrtInv;
const int* dofMadr = m->dof_Madr;
mju_zero(res, nv);
@@ -1117,9 +1116,9 @@ void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec)
}
}
// res = sqrt(D) * res
// res *= sqrt(D)
for (int i=0; i < nv; i++) {
res[i] /= qLDiagSqrtInv[i];
res[i] *= mju_sqrt(qLD[dofMadr[i]]);
}
}
src/engine/engine_vis_interact.c
+5 -1
View File
@@ -541,6 +541,7 @@ void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPertur
mjtNum* jac = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* jacM2 = mjSTACKALLOC(d, 3*nv, mjtNum);
mjtNum* sqrtInvD = mjSTACKALLOC(d, nv, mjtNum);
// invalid selected body: return
if (sel <= 0 || sel >= m->nbody) {
@@ -554,8 +555,11 @@ void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPertur
mju_addTo3(selpos, d->xpos+3*sel);
// compute average spatial inertia at selection point
for (int i=0; i < nv; i++) {
sqrtInvD[i] = 1 / mju_sqrt(d->qLD[m->dof_Madr[i]]);
}
mj_jac(m, d, jac, NULL, selpos, sel);
mj_solveM2(m, d, jacM2, jac, 3);
mj_solveM2(m, d, jacM2, jac, sqrtInvD, 3);
mjtNum invmass = mju_dot(jacM2+0*nv, jacM2+0*nv, nv) +
mju_dot(jacM2+1*nv, jacM2+1*nv, nv) +
mju_dot(jacM2+2*nv, jacM2+2*nv, nv);
unity/Runtime/Bindings/MjBindings.cs
+1 -2
View File
@@ -4917,7 +4917,6 @@ public unsafe struct mjData_ {
public double* qM;
public double* qLD;
public double* qLDiagInv;
public double* qLDiagSqrtInv;
public double* bvh_aabb_dyn;
public byte* bvh_active;
public double* flexedge_velocity;
@@ -6679,7 +6678,7 @@ public static unsafe extern void mj_factorM(mjModel_* m, mjData_* d);
public static unsafe extern void mj_solveM(mjModel_* m, mjData_* d, double* x, double* y, int n);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mj_solveM2(mjModel_* m, mjData_* d, double* x, double* y, int n);
public static unsafe extern void mj_solveM2(mjModel_* m, mjData_* d, double* x, double* y, double* sqrtInvD, int n);
[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
public static unsafe extern void mj_comVel(mjModel_* m, mjData_* d);