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

Move Newton Hessian memory allocation from stack to arena.

Browse Source 
PiperOrigin-RevId: 684054446
Change-Id: I4be0b216810ebbd9d7fa119c6a84eb08d3c16b83
This commit is contained in:
Yuval Tassa
2024-10-09 08:53:40 -07:00
committed by Copybara-Service
parent 356d172f56
commit 9a0dc20821
4 changed files with 92 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/bindings_test.py
+10 -2
View File
@@ -1360,12 +1360,20 @@ Euler integrator, semi-implicit in velocity.
# (e.g. because the internal output buffer is too small)
self.assertIn('mujoco', mujoco.mj_printSchema(flg_html, flg_pad))
def test_pickle_mjdata_before_step(self):
data2 = pickle.loads(pickle.dumps(self.data))
attr_to_compare = (
'time', 'qpos', 'qvel', 'qacc', 'xpos', 'mocap_pos',
'warning', 'energy', 'contact', 'efc_J', 'L'
)
self._assert_attributes_equal(data2, self.data, attr_to_compare)
def test_pickle_mjdata(self):
mujoco.mj_step(self.model, self.data)
data2 = pickle.loads(pickle.dumps(self.data))
attr_to_compare = (
'time', 'qpos', 'qvel', 'qacc', 'xpos', 'mocap_pos',
'warning', 'energy', 'contact', 'efc_J'
'warning', 'energy', 'contact', 'efc_J', 'L'
)
self._assert_attributes_equal(data2, self.data, attr_to_compare)
for _ in range(10):
@@ -1379,7 +1387,7 @@ Euler integrator, semi-implicit in velocity.
data2 = pickle.loads(pickle.dumps(self.data))
attr_to_compare = (
'time', 'qpos', 'qvel', 'qacc', 'xpos', 'mocap_pos',
'warning', 'energy', 'contact', 'efc_J'
'warning', 'energy', 'contact', 'efc_J', 'L'
)
self._assert_attributes_equal(data2, self.data, attr_to_compare)
for _ in range(10):
python/mujoco/serialization.h
+4 -1
View File
@@ -15,6 +15,7 @@
#ifndef MUJOCO_PYTHON_SERIALIZATION_H_
#define MUJOCO_PYTHON_SERIALIZATION_H_
#include <cstddef>
#include <iostream>
#include <mujoco/mjtnum.h>
@@ -56,7 +57,9 @@ inline void WriteBytes(std::ostream& output, const void* src,
// Start by writing nbytes itself, so it can be validated at the time of
// reading.
WriteInt(output, nbytes);
output.write(reinterpret_cast<const char*>(src), nbytes);
if (src) {
output.write(reinterpret_cast<const char*>(src), nbytes);
}
}
inline void ReadBytes(std::istream& input, void* dest, std::size_t nbytes) {
python/mujoco/structs.cc
+32 -12
View File
@@ -755,7 +755,7 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
X(energy);
#undef X
// Write buffer contents
// Write buffer and arena contents
{
MJDATA_POINTERS_PREAMBLE((this->model_->get()))
@@ -764,18 +764,24 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
MJDATA_POINTERS
#undef X
bool is_sparse_newton = this->model_->get()->opt.solver == mjSOL_NEWTON &&
mj_isSparse(this->model_->get());
#undef MJ_M
#define MJ_M(x) this->model_->get()->x
#undef MJ_D
#define MJ_D(x) this->ptr_->x
#define X(type, name, nr, nc) \
if ((nr) * (nc)) { \
WriteBytes(output, ptr_->name, sizeof(type) * (nr) * (nc)); \
#define X(type, name, nr, nc) \
if ((nr) * (nc)) { \
WriteBytes(output, ptr_->name, \
ptr_->name ? sizeof(type) * (nr) * (nc) : 0); \
}
MJDATA_ARENA_POINTERS_CONTACT
MJDATA_ARENA_POINTERS_SOLVER
// MJDATA_ARENA_POINTERS_NEWTON // TODO: tassa - Add after allocation exists
if (is_sparse_newton) {
MJDATA_ARENA_POINTERS_NEWTON
}
if (mj_isDual(this->model_->get())) {
MJDATA_ARENA_POINTERS_DUAL
}
@@ -805,6 +811,8 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
bool is_dual = mj_isDual(&m);
bool is_sparse_newton = m.opt.solver == mjSOL_NEWTON && mj_isSparse(&m);
raw::MjData* d = mj_makeData(&m);
if (!d) {
throw py::value_error("Failed to create mjData.");
@@ -833,7 +841,7 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
X(energy);
#undef X
// Read buffer contents
// Read buffer and arena contents
{
MJDATA_POINTERS_PREAMBLE((&m))
@@ -846,16 +854,28 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
#define MJ_M(x) m.x
#undef MJ_D
#define MJ_D(x) d->x
#define X(type, name, nr, nc) \
if ((nr) * (nc)) { \
d->name = static_cast<decltype(d->name)>( \
mj_arenaAllocByte(d, sizeof(type) * (nr) * (nc), alignof(type))); \
ReadBytes(input, d->name, sizeof(type) * (nr) * (nc)); \
// arena pointers might be null, so we need to check the size before allocating.
#define X(type, name, nr, nc) \
if ((nr) * (nc)) { \
std::size_t actual_nbytes = ReadInt(input); \
if (actual_nbytes) { \
if (actual_nbytes != sizeof(type) * (nr) * (nc)) { \
input.setstate(input.rdstate() | std::ios_base::failbit); \
} else { \
d->name = static_cast<decltype(d->name)>( \
mj_arenaAllocByte(d, sizeof(type) * (nr) * (nc), alignof(type))); \
input.read(reinterpret_cast<char*>(d->name), actual_nbytes); \
} \
} else { \
d->name = nullptr; \
} \
}
MJDATA_ARENA_POINTERS_CONTACT
MJDATA_ARENA_POINTERS_SOLVER
// MJDATA_ARENA_POINTERS_NEWTON // TODO: tassa - Add after allocation exists
if (is_sparse_newton) {
MJDATA_ARENA_POINTERS_NEWTON
}
if (is_dual) {
MJDATA_ARENA_POINTERS_DUAL
}
src/engine/engine_solver.c
+46 -42
View File
@@ -765,6 +765,8 @@ void mj_solNoSlip(const mjModel* m, mjData* d, int maxiter) {
// CG context
struct _mjCGContext {
int flg_Newton; // 1: Newton, 0: CG
// island-related
int island; // current island index, -1 if monolithic
int nv; // number of dofs
@@ -782,15 +784,6 @@ struct _mjCGContext {
mjtNum* search; // linesearch vector (nv x 1)
mjtNum* quad; // quadratic polynomials for constraint costs (nefc x 3)
// Hessian (Newton only)
int flg_Newton; // 1: Newton, 0: CG (const)
int nnz; // total number of non-zeros
mjtNum* H; // Cholesky factorization of Hessian (nv x nv)
mjtNum* Hcone; // with cone contributions if present (nv x nv)
int* rownnz; // non-zeros in row (nv X 1)
int* rowadr; // row address (nv x 1)
int* colind; // column indices (nv x nv)
// globals
mjtNum cost; // constraint + Gauss cost
mjtNum quadGauss[3]; // quadratic polynomial for Gauss cost
@@ -837,12 +830,11 @@ static void CGallocate(const mjModel* m, mjData* d, mjCGContext* ctx,
// Hessian (Newton only)
ctx->flg_Newton = flg_Newton;
if (flg_Newton) {
ctx->H = mj_stackAllocNum(d, nv*nv);
ctx->Hcone = mj_stackAllocNum(d, nv*nv);
ctx->rownnz = mj_stackAllocInt(d, nv);
ctx->rowadr = mj_stackAllocInt(d, nv);
ctx->colind = mj_stackAllocInt(d, nv*nv);
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");
}
}
@@ -894,10 +886,10 @@ static void CGupdateGradient(const mjModel* m, const mjData* d, mjCGContext* ctx
// TODO: b/295296178 - add island support to Newton solver
if (ctx->flg_Newton) {
if (mj_isSparse(m)) {
mju_cholSolveSparse(ctx->Mgrad, (ctx->ncone ? ctx->Hcone : ctx->H),
ctx->grad, nv, ctx->rownnz, ctx->rowadr, ctx->colind);
mju_cholSolveSparse(ctx->Mgrad, (ctx->ncone ? d->Lcone : d->L),
ctx->grad, nv, d->L_rownnz, d->L_rowadr, d->L_colind);
} else {
mju_cholSolve(ctx->Mgrad, (ctx->ncone ? ctx->Hcone : ctx->H), ctx->grad, nv);
mju_cholSolve(ctx->Mgrad, (ctx->ncone ? d->Lcone : d->L), ctx->grad, nv);
}
}
@@ -1378,7 +1370,7 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) {
int* LTJ_ind = mj_stackAllocInt(d, nv);
// start with Hcone = H
mju_copy(ctx->Hcone, ctx->H, ctx->nnz);
mju_copy(d->Lcone, d->L, d->nnzL);
// add contributions
for (int i=0; i < nefc; i++) {
@@ -1410,8 +1402,8 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) {
mju_copyInt(LTJ_ind, d->efc_J_colind+d->efc_J_rowadr[i+r], nnz);
// update
mju_cholUpdateSparse(ctx->Hcone, LTJ_row, nv, 1,
ctx->rownnz, ctx->rowadr, ctx->colind, nnz, LTJ_ind, d);
mju_cholUpdateSparse(d->Lcone, LTJ_row, nv, 1,
d->L_rownnz, d->L_rowadr, d->L_colind, nnz, LTJ_ind, d);
}
}
@@ -1427,7 +1419,7 @@ static void HessianCone(const mjModel* m, mjData* d, mjCGContext* ctx) {
// update
for (int r=0; r < dim; r++) {
mju_cholUpdate(ctx->Hcone, LTJ+r*nv, nv, 1);
mju_cholUpdate(d->Lcone, LTJ+r*nv, nv, 1);
}
}
@@ -1448,6 +1440,21 @@ 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
@@ -1472,21 +1479,21 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
// compute H = J'*D*J
// TODO(b/266802572): remove uncompressed layout
mju_sqrMatTDUncompressedInit(ctx->rowadr, nv);
mju_sqrMatTDSparse(ctx->H, d->efc_J, d->efc_JT, D, nefc, nv,
ctx->rownnz, ctx->rowadr, ctx->colind,
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, ctx->H, ctx->rownnz, ctx->rowadr, ctx->colind,
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(ctx->H, nv, mjMINVAL,
ctx->rownnz, ctx->rowadr, ctx->colind, d);
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) {
@@ -1494,14 +1501,11 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
}
// compress layout of H
mju_compressSparse(ctx->H, nv, nv, ctx->rownnz, ctx->rowadr, ctx->colind);
mju_compressSparse(d->L, nv, nv, d->L_rownnz, d->L_rowadr, d->L_colind);
// count nnz
ctx->nnz = 0;
for (int i=0; i < nv; i++) {
ctx->nnz += ctx->rownnz[i];
}
if (ctx->nnz > nv*nv) { // SHOULD NOT OCCUR
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");
}
}
@@ -1509,14 +1513,14 @@ static void HessianDirect(const mjModel* m, mjData* d, mjCGContext* ctx) {
// dense
else {
// compute H = M + J'*D*J
mju_sqrMatTD(ctx->H, d->efc_J, D, nefc, nv);
mj_addMDense(m, d, ctx->H);
mju_sqrMatTD(d->L, d->efc_J, D, nefc, nv);
mj_addMDense(m, d, d->L);
// factorize H
mju_cholFactor(ctx->H, nv, mjMINVAL);
mju_cholFactor(d->L, nv, mjMINVAL);
// set nnz
ctx->nnz = nv*nv;
d->nnzL = nv*nv;
}
mj_freeStack(d);
@@ -1572,12 +1576,12 @@ static void HessianIncremental(const mjModel* m, mjData* d,
mju_copyInt(vec_ind, d->efc_J_colind+adr, nnz);
// sparse update
rank = mju_cholUpdateSparse(ctx->H, vec, nv, flag_update,
ctx->rownnz, ctx->rowadr, ctx->colind, nnz, vec_ind,
rank = mju_cholUpdateSparse(d->L, vec, nv, flag_update,
d->L_rownnz, d->L_rowadr, d->L_colind, nnz, vec_ind,
d);
} else {
mju_scl(vec, d->efc_J+i*nv, mju_sqrt(d->efc_D[i]), nv);
rank = mju_cholUpdate(ctx->H, vec, nv, flag_update);
rank = mju_cholUpdate(d->L, vec, nv, flag_update);
}
ctx->nupdate++;
@@ -1758,7 +1762,7 @@ static void mj_solCGNewton(const mjModel* m, mjData* d, int island, int maxiter,
// set solver_nnz
if (flg_Newton) {
if (mj_isSparse(m)) {
d->solver_nnz[island_stat] = 2*ctx.nnz - nv;
d->solver_nnz[island_stat] = 2*d->nnzL - nv;
} else {
d->solver_nnz[island_stat] = nv*nv;
}