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

Stop over-allocating memory in Newton solver.

Browse Source 
PiperOrigin-RevId: 684797464
Change-Id: I3389416cb69a8564b5119f6e3944b579893ea511
This commit is contained in:
Yuval Tassa
2024-10-11 05:08:27 -07:00
committed by Copybara-Service
parent 9aca4bccac
commit 2dd518734f
8 changed files with 220 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+5 -1
View File
@@ -8,9 +8,13 @@ Upcoming version (not yet released)
General
^^^^^^^
- The Newton solver no longer requires ``nv*nv`` memory allocation, allowing for much larger models. See e.g.,
`100_humanoids.xml <https://github.com/google-deepmind/mujoco/blob/main/model/humanoid/100_humanoids.xml>`__.
Two quadratic-memory allocations still remain to be fully sparsified: ``mjData.actuator_moment`` and the matrices used
by the PGS solver.
- Removed the :at:`solid` and :at:`membrane` plugins and moved the associated computations into the engine. See `3D
example model <https://github.com/google-deepmind/mujoco/blob/main/model/flex/floppy.xml>`__ and `2D example model
<https://github.com/google-deepmind/mujoco/blob/main/src/model/trampoline.xml>`__ for examples of flex objects
<https://github.com/google-deepmind/mujoco/blob/main/model/flex/trampoline.xml>`__ for examples of flex objects
that previously required these plugins.
- Replaced the function ``mjs_setActivePlugins`` with :ref:`mjs_activatePlugin`.
model/humanoid/100_humanoids.xml
+50
View File
@@ -0,0 +1,50 @@
<!-- Copyright 2021 DeepMind Technologies Limited
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
-->
<mujoco model="100 Humanoids">
<option timestep="0.005"/>
<size memory="100M"/>
<asset>
<texture type="skybox" builtin="gradient" rgb1=".3 .5 .7" rgb2="0 0 0" width="512" height="512"/>
<texture name="floor" type="2d" builtin="checker" width="512" height="512" rgb1=".1 .2 .3" rgb2=".2 .3 .4"/>
<material name="floor" texture="floor" texrepeat="1 1" texuniform="true" reflectance=".2"/>
<model name="humanoid" file="humanoid.xml"/>
</asset>
<visual>
<map force="0.1" zfar="30"/>
<rgba haze="0.15 0.25 0.35 1"/>
<quality shadowsize="4096" numslices="16" numstacks="8"/>
<global offwidth="800" offheight="800"/>
</visual>
<worldbody>
<geom name="floor" size="10 10 .05" type="plane" material="floor" condim="3"/>
<light directional="true" diffuse=".4 .4 .4" specular="0.1 0.1 0.1" pos="0 0 5" dir="0 0 -1" castshadow="false"/>
<light name="spotlight" mode="targetbodycom" target="world" diffuse="1 1 1" specular="0.3 0.3 0.3" pos="-6 -6 4" cutoff="60"/>
<replicate count="10" euler="0 0 36" sep="-">
<frame pos="1 0 0">
<replicate count="10" euler="0 0 15" sep="-" offset="0.5 0 0">
<attach model="humanoid" body="torso" prefix="_"/>
</replicate>
</frame>
</replicate>
<replicate count="5" euler="0 0 72" sep="-">
<light pos="0 -4 4"/>
</replicate>
</worldbody>
</mujoco>
src/engine/engine_print.c
+22 -16
View File
@@ -756,7 +756,7 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
mjERROR("attempting to print mjData when stack is in use");
}
mjtNum *M;
mjtNum *M = NULL;
mj_markStack(d);
// check format string
@@ -780,8 +780,10 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
return;
}
// allocate full inertia
M = mj_stackAllocNum(d, m->nv*m->nv);
// allocate full inertia if it's small
if (m->nv <= 200) {
M = mj_stackAllocNum(d, m->nv*m->nv);
}
#ifdef MEMORY_SANITIZER
// If memory sanitizer is active, d->buffer will be marked as poisoned, even
@@ -969,13 +971,15 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
printArray("ACTUATOR_MOMENT", m->nu, m->nv, d->actuator_moment, fp, float_format);
printArray("CRB", m->nbody, 10, d->crb, fp, float_format);
// construct and print full M matrix
mj_fullM(m, M, d->qM);
printArray("QM", m->nv, m->nv, M, fp, float_format);
if (M) {
// construct and print full M matrix
mj_fullM(m, M, d->qM);
printArray("QM", m->nv, m->nv, M, fp, float_format);
// construct and print full LD matrix
mj_fullM(m, M, d->qLD);
printArray("QLD", m->nv, m->nv, M, fp, float_format);
// construct and print full LD matrix
mj_fullM(m, M, d->qLD);
printArray("QLD", m->nv, m->nv, M, fp, float_format);
}
printArray("QLDIAGINV", m->nv, 1, d->qLDiagInv, fp, float_format);
printArray("QLDIAGSQRTINV", m->nv, 1, d->qLDiagSqrtInv, fp, float_format);
@@ -1064,14 +1068,16 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
}
fprintf(fp, "\n\n");
// print qDeriv
mju_sparse2dense(M, d->qDeriv, m->nv, m->nv, d->D_rownnz, d->D_rowadr, d->D_colind);
printArray("QDERIV", m->nv, m->nv, M, fp, float_format);
if (M) {
// print qDeriv
mju_sparse2dense(M, d->qDeriv, m->nv, m->nv, d->D_rownnz, d->D_rowadr, d->D_colind);
printArray("QDERIV", m->nv, m->nv, M, fp, float_format);
// print qLU
mju_sparse2dense(M, d->qLU, m->nv, m->nv, d->D_rownnz, d->D_rowadr,
d->D_colind);
printArray("QLU", m->nv, m->nv, M, fp, float_format);
// print qLU
mju_sparse2dense(M, d->qLU, m->nv, m->nv, d->D_rownnz, d->D_rowadr,
d->D_colind);
printArray("QLU", m->nv, m->nv, M, fp, float_format);
}
// contact
fprintf(fp, "CONTACT\n");
src/engine/engine_solver.c
+130 -59
View File
@@ -830,11 +830,29 @@ static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx,
// Hessian (Newton only)
ctx->flg_Newton = flg_Newton;
if (flg_Newton && mj_isSparse(m)) {
d->L_rowadr = mj_arenaAllocByte(d, sizeof(int) * nv, _Alignof(int));
if (!d->L_rowadr) mjERROR("failed to allocate L_rowadr");
d->L_rownnz = mj_arenaAllocByte(d, sizeof(int) * nv, _Alignof(int));
if (!d->L_rownnz) mjERROR("failed to allocate L_rownnz");
if (flg_Newton) {
// sparse: allocate L_rowadr, L_rownnz
if (mj_isSparse(m)) {
d->L_rowadr = mj_arenaAllocByte(d, sizeof(int) * nv, _Alignof(int));
if (!d->L_rowadr) mjERROR("failed to allocate L_rowadr");
d->L_rownnz = mj_arenaAllocByte(d, sizeof(int) * nv, _Alignof(int));
if (!d->L_rownnz) mjERROR("failed to allocate L_rownnz");
// zero nnzL, clear pointers (compute and allocate later in HessianDirect)
d->nnzL = 0;
d->L_colind = NULL;
d->L = NULL;
d->Lcone = NULL;
}
// dense: allocate L
else if (d->nnzL != nv*nv) {
d->L = mj_arenaAllocByte(d, sizeof(mjtNum) * nv*nv, _Alignof(mjtNum));
if (!d->L) mjERROR("failed to allocate L");
// set dense nnzL
d->nnzL = nv*nv;
}
}
}
@@ -1362,6 +1380,16 @@ static mjtNum CGsearch(const mjModel* m, const mjData* d, mjCGContext* ctx) {
static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) {
int nv = m->nv, nefc = d->nefc;
mjtNum local[36];
// allocate Lcone if required
if (!d->Lcone) {
d->Lcone = mj_arenaAllocByte(d, sizeof(mjtNum) * d->nnzL, _Alignof(mjtNum));
}
if (!d->Lcone) mjERROR("failed to allocate Lcone");
// start with Hcone = H
mju_copy(d->Lcone, d->L, d->nnzL);
mj_markStack(d);
// storage for L'*J
@@ -1369,9 +1397,6 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) {
mjtNum* LTJ_row = mj_stackAllocNum(d, nv);
int* LTJ_ind = mj_stackAllocInt(d, nv);
// start with Hcone = H
mju_copy(d->Lcone, d->L, d->nnzL);
// add contributions
for (int i=0; i < nefc; i++) {
if (d->efc_state[i] == mjCNSTRSTATE_CONE) {
@@ -1440,21 +1465,6 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) {
// TODO: b/295296178 - add island support to Newton solver
static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
int nv = m->nv, nefc = d->nefc;
// allocate Hessian on arena if not already allocated
if (!d->nnzL) {
int nnz = nv*nv;
if (mj_isSparse(m)) {
d->L_colind = mj_arenaAllocByte(d, sizeof(int) * nnz, _Alignof(int));
if (!d->L_colind) mjERROR("failed to allocate L_colind");
}
d->L = mj_arenaAllocByte(d, sizeof(mjtNum) * nnz, _Alignof(mjtNum));
if (!d->L) mjERROR("failed to allocate L");
d->Lcone = mj_arenaAllocByte(d, sizeof(mjtNum) * nnz, _Alignof(mjtNum));
if (!d->Lcone) mjERROR("failed to allocate Lcone");
d->nnzL = nnz;
}
mj_markStack(d);
// compute D corresponding to quad states
@@ -1469,44 +1479,109 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
// sparse
if (mj_isSparse(m)) {
// fill-in reduced sparse inertia matrix C (no off-diagonals for simple dofs)
int nC = m->nC;
mjtNum* C = mj_stackAllocNum(d, nC);
for (int i=0; i < nC; i++) {
// copy values of reduced sparse inertia matrix C, get nnz
int nnz_C = m->nC;
mjtNum* C = mj_stackAllocNum(d, nnz_C);
for (int i=0; i < nnz_C; i++) {
C[i] = d->qM[d->mapM2C[i]];
}
// compute H = J'*D*J
// allocate and initialize Hessian rowadr, rownnz; get nnz for J'*J
int* H_rowadr = mj_stackAllocInt(d, nv);
int* H_rownnz = mj_stackAllocInt(d, nv);
mju_sqrMatTDSparseInit(H_rownnz, H_rowadr, nv,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
d);
int nnz_JTJ = H_rowadr[nv-1] + H_rownnz[nv-1];
// TODO(b/266802572): remove uncompressed layout
mju_sqrMatTDUncompressedInit(d->L_rowadr, nv);
mju_sqrMatTDSparse(d->L, d->efc_J, d->efc_JT, D, nefc, nv,
d->L_rownnz, d->L_rowadr, d->L_colind,
d->efc_J_rownnz, d->efc_J_rowadr,
d->efc_J_colind, NULL,
d->efc_JT_rownnz, d->efc_JT_rowadr,
d->efc_JT_colind, d->efc_JT_rowsuper, d);
// compute H = M + J'*D*J
mj_addMSparse(m, d, d->L, d->L_rownnz, d->L_rowadr, d->L_colind,
C, d->C_rownnz, d->C_rowadr, d->C_colind);
// factorize H, uncompressed layout
int rank = mju_cholFactorSparse(d->L, nv, mjMINVAL,
d->L_rownnz, d->L_rowadr, d->L_colind, d);
// rank-defficient, SHOULD NOT OCCUR
if (rank != nv) {
mjERROR("rank-defficient Hessian");
// shift H rowadr to make room for C
int shift = 0;
for (int r = 0; r < nv - 1; r++) {
shift += d->C_rownnz[r];
H_rowadr[r + 1] += shift;
}
// compress layout of H
mju_compressSparse(d->L, nv, nv, d->L_rownnz, d->L_rowadr, d->L_colind);
// allocate Hessian H, colind
int nnz_H = nnz_C + nnz_JTJ;
mjtNum* H = mj_stackAllocNum(d, nnz_H);
int* H_colind = mj_stackAllocInt(d, nnz_H);
// count nnz
d->nnzL = d->L_rowadr[nv-1] + d->L_rownnz[nv-1];
if (d->nnzL > nv*nv) { // SHOULD NOT OCCUR
mjERROR("more nonzero values than elements in sparse direct-solver Hessian");
// compute H = J'*D*J
mju_sqrMatTDSparse(H, d->efc_J, d->efc_JT, D, nefc, nv,
H_rownnz, H_rowadr, H_colind,
d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind, NULL,
d->efc_JT_rownnz, d->efc_JT_rowadr, d->efc_JT_colind, d->efc_JT_rowsuper,
d);
// add mass matrix; H = J'*D*J + C
mj_addMSparse(m, d, H, H_rownnz, H_rowadr, H_colind,
C, d->C_rownnz, d->C_rowadr, d->C_colind);
// count row and total non-zeros of reverse-Cholesky factor L
int* parent = mj_stackAllocInt(d, nv);
int* flag = mj_stackAllocInt(d, nv);
int nnz_L = mju_cholFactorNNZ(d->L_rownnz, parent, flag, H_rownnz, H_rowadr, H_colind, nv);
// allocate L_colind, L on arena if required
if (!d->nnzL) {
// nnzL is 0 but pointers are allocated; SHOULD NOT OCCUR
if (d->L_colind || d->L) {
mjERROR("nnzL is 0 but L_colind or L or Lcone are allocated");
}
// allocate on arena
d->L_colind = mj_arenaAllocByte(d, sizeof(int) * nnz_L, _Alignof(int));
if (!d->L_colind) mjERROR("failed to allocate L_colind");
d->L = mj_arenaAllocByte(d, sizeof(mjtNum) * nnz_L, _Alignof(mjtNum));
if (!d->L) mjERROR("failed to allocate L");
// set nnzL
d->nnzL = nnz_L;
} else if (d->nnzL != nnz_L) {
// nnzL is nonzero but not equal to computed value; SHOULD NOT OCCUR
mjERROR("nnzL is nonzero but not equal to computed value");
}
// compute L row adresses: L_rowadr = cumsum(L_rownnz)
d->L_rowadr[0] = 0;
for (int r=1; r < nv; r++) {
d->L_rowadr[r] = d->L_rowadr[r-1] + d->L_rownnz[r-1];
}
// copy H lower-triangle into L
for (int r = 0; r < nv; r++) {
// count H non-zeros up to diagonal (inclusive) for row r
const int* colind = H_colind + H_rowadr[r];
int rownnz = 1;
while (rownnz < nv && colind[rownnz - 1] < r) {
rownnz++;
}
// last row element is not the diagonal; SHOULD NOT OCCUR
if (colind[rownnz - 1] != r) {
mjERROR("Newton solver Hessian has zero diagonal on row %d", r);
}
// copy values and column indices
mju_copy(d->L + d->L_rowadr[r], H + H_rowadr[r], rownnz);
mju_copyInt(d->L_colind + d->L_rowadr[r], H_colind + H_rowadr[r], rownnz);
// set L_rownnz
d->L_rownnz[r] = rownnz;
}
// in-place sparse factorization L = chol(H)
int rank = mju_cholFactorSparse(d->L, nv, mjMINVAL, d->L_rownnz, d->L_rowadr, d->L_colind, d);
// rank-deficient; SHOULD NOT OCCUR
if (rank != nv) {
mjERROR("rank-deficient Hessian");
}
// pre-counted nnzL does not match post-factorization nnzL; SHOULD NOT OCCUR
if (d->nnzL != d->L_rowadr[nv-1] + d->L_rownnz[nv-1]) {
mjERROR("mismatch between pre-counted and post-factorization L nonzeros");
}
}
@@ -1518,9 +1593,6 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
// factorize H
mju_cholFactor(d->L, nv, mjMINVAL);
// set nnz
d->nnzL = nv*nv;
}
mj_freeStack(d);
@@ -1538,8 +1610,7 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
// incremental update to Hessian
// TODO: b/295296178 - add island support to Newton solver
static void HessianIncremental(const mjModel* m, mjData* d,
mjCGContext* ctx, const int* oldstate) {
static void HessianIncremental(const mjModel* m, mjData* d, mjCGContext* ctx, const int* oldstate) {
int rank, nv = m->nv, nefc = d->nefc;
mj_markStack(d);
src/engine/engine_util_solve.c
+5 -18
View File
@@ -141,28 +141,15 @@ int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus) {
//---------------------------- sparse Cholesky -----------------------------------------------------
// sparse reverse-order Cholesky decomposition: mat = L'*L; return 'rank'
// mat must have uncompressed layout; rownnz is modified to end at diagonal
// mat must be lower-triangular, have preallocated space for fill-in
int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
int* rownnz, int* rowadr, int* colind,
int* rownnz, const int* rowadr, int* colind,
mjData* d) {
int rank = n;
mj_markStack(d);
mjtNum* buf = mj_stackAllocNum(d, n);
int* buf_ind = mj_stackAllocInt(d, n);
mjtNum* sparse_buf = mj_stackAllocNum(d, n);
// shrink rows so that rownnz ends at diagonal
for (int r=0; r < n; r++) {
// shrink
while (rownnz[r] > 0 && colind[rowadr[r]+rownnz[r]-1] > r) {
rownnz[r]--;
}
// check
if (rownnz[r] == 0 || colind[rowadr[r]+rownnz[r]-1] != r) {
mjERROR("matrix must have non-zero diagonal");
}
}
// backpass over rows
for (int r=n-1; r >= 0; r--) {
@@ -191,7 +178,7 @@ int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
// mat(c,0:c) = mat(c,0:c) - mat(r,c) * mat(r,0:c)
int nnz_c = mju_combineSparse(mat + rowadr[c], mat+rowadr[r], 1, -mat[adr+i],
rownnz[c], i+1, colind+rowadr[c], colind+rowadr[r],
sparse_buf, buf_ind);
buf, buf_ind);
// assign new nnz to row c
rownnz[c] = nnz_c;
@@ -251,7 +238,7 @@ void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
// sparse reverse-order Cholesky rank-one update: L'*L +/- x*x'; return rank
// x is sparse, change in sparsity pattern of mat is not allowed
int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
int* rownnz, int* rowadr, int* colind, int x_nnz, int* x_ind,
const int* rownnz, const int* rowadr, int* colind, int x_nnz, int* x_ind,
mjData* d) {
mj_markStack(d);
int* buf_ind = mj_stackAllocInt(d, n);
src/engine/engine_util_solve.h
+4 -3
View File
@@ -33,8 +33,9 @@ MJAPI void mju_cholSolve(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
MJAPI int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus);
// sparse reverse-order Cholesky decomposition: mat = L'*L; return 'rank'
// mat must have uncompressed layout; rownnz is modified to end at diagonal
int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag, int* rownnz, int* rowadr, int* colind,
// mat must be lower-triangular, have preallocated space for fill-in
int mju_cholFactorSparse(mjtNum* mat, int n, mjtNum mindiag,
int* rownnz, const int* rowadr, int* colind,
mjData* d);
// sparse reverse-order Cholesky solve
@@ -44,7 +45,7 @@ void mju_cholSolveSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int
// sparse reverse-order Cholesky rank-one update: L'*L +/i x*x'; return rank
// x is sparse, change in sparsity pattern of mat is not allowed
int mju_cholUpdateSparse(mjtNum* mat, mjtNum* x, int n, int flg_plus,
int* rownnz, int* rowadr, int* colind, int x_nnz, int* x_ind,
const int* rownnz, const int* rowadr, int* colind, int x_nnz, int* x_ind,
mjData* d);
// band-dense Cholesky decomposition
test/sample/testspeed_test.sh
+2 -1
View File
@@ -28,7 +28,8 @@ test_model() {
local iterations=10
# for particularly slow models, only run 2 steps under ASAN, or skip.
if [[ ${TESTSPEED_ASAN:-0} != 0 ]]; then
if [[ "$model" == */composite/particle.xml ||
if [[ "$model" == */humanoid/100_humanoids.xml ||
"$model" == */composite/particle.xml ||
"$model" == */replicate/bunnies.xml ||
"$model" == */replicate/leaves.xml ||
"$model" == */replicate/particle.xml ||
test/xml/xml_native_writer_test.cc
+2 -1
View File
@@ -1312,7 +1312,8 @@ TEST_F(XMLWriterTest, WriteReadCompare) {
std::string xml = p.path().string();
// if file is meant to fail, skip it
if (absl::StrContains(p.path().string(), "malformed_") ||
if (absl::StrContains(p.path().string(), "100_humanoids") ||
absl::StrContains(p.path().string(), "malformed_") ||
absl::StrContains(p.path().string(), "touch_grid") ||
absl::StrContains(p.path().string(), "gmsh_") ||
absl::StrContains(p.path().string(), "shark_") ||