diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index 6036366a..bbbd7022 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -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): diff --git a/python/mujoco/serialization.h b/python/mujoco/serialization.h index 9fb6a1b3..d38abefd 100644 --- a/python/mujoco/serialization.h +++ b/python/mujoco/serialization.h @@ -15,6 +15,7 @@ #ifndef MUJOCO_PYTHON_SERIALIZATION_H_ #define MUJOCO_PYTHON_SERIALIZATION_H_ +#include #include #include @@ -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(src), nbytes); + if (src) { + output.write(reinterpret_cast(src), nbytes); + } } inline void ReadBytes(std::istream& input, void* dest, std::size_t nbytes) { diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc index 3fe757a8..ff4a9f4d 100644 --- a/python/mujoco/structs.cc +++ b/python/mujoco/structs.cc @@ -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_castname)>( \ - 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_castname)>( \ + mj_arenaAllocByte(d, sizeof(type) * (nr) * (nc), alignof(type))); \ + input.read(reinterpret_cast(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 } diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index e23b0d8a..99c02368 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -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; }