Move Newton Hessian memory allocation from stack to arena.

PiperOrigin-RevId: 684054446
Change-Id: I4be0b216810ebbd9d7fa119c6a84eb08d3c16b83

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):

python/mujoco/serialization.h

@@ -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

@@ -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)                                   \
+#define X(type, name, nr, nc)                     \
-  if ((nr) * (nc)) {                                            \
+  if ((nr) * (nc)) {                              \
-    WriteBytes(output, ptr_->name, sizeof(type) * (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)                                              \
+// arena pointers might be null, so we need to check the size before allocating.
-  if ((nr) * (nc)) {                                                       \
+#define X(type, name, nr, nc)                                                 \
-    d->name = static_cast<decltype(d->name)>(                              \
+  if ((nr) * (nc)) {                                                          \
-        mj_arenaAllocByte(d, sizeof(type) * (nr) * (nc), alignof(type)));  \
+    std::size_t actual_nbytes = ReadInt(input);                               \
-    ReadBytes(input, d->name, sizeof(type) * (nr) * (nc));                 \
+    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

@@ -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) {
+  if (flg_Newton && mj_isSparse(m)) {
-    ctx->H      = mj_stackAllocNum(d, nv*nv);
+    d->L_rowadr = mj_arenaAllocByte(d, sizeof(int) * nv, _Alignof(int));
-    ctx->Hcone  = mj_stackAllocNum(d, nv*nv);
+    if (!d->L_rowadr) mjERROR("failed to allocate L_rowadr");
-    ctx->rownnz = mj_stackAllocInt(d, nv);
+    d->L_rownnz = mj_arenaAllocByte(d, sizeof(int) * nv, _Alignof(int));
-    ctx->rowadr = mj_stackAllocInt(d, nv);
+    if (!d->L_rownnz) mjERROR("failed to allocate L_rownnz");
-    ctx->colind = mj_stackAllocInt(d, nv*nv);
   }
 }
 
@@ -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),
+      mju_cholSolveSparse(ctx->Mgrad, (ctx->ncone ? d->Lcone : d->L),
-                          ctx->grad, nv, ctx->rownnz, ctx->rowadr, ctx->colind);
+                          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,
+          mju_cholUpdateSparse(d->Lcone, LTJ_row, nv, 1,
-                               ctx->rownnz, ctx->rowadr, ctx->colind, nnz, LTJ_ind, d);
+                               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_sqrMatTDUncompressedInit(d->L_rowadr, nv);
-    mju_sqrMatTDSparse(ctx->H, d->efc_J, d->efc_JT, D, nefc, nv,
+    mju_sqrMatTDSparse(d->L, d->efc_J, d->efc_JT, D, nefc, nv,
-                       ctx->rownnz, ctx->rowadr, ctx->colind,
+                       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,
+    int rank = mju_cholFactorSparse(d->L, nv, mjMINVAL,
-                                    ctx->rownnz, ctx->rowadr, ctx->colind, d);
+                                    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;
+    d->nnzL = d->L_rowadr[nv-1] + d->L_rownnz[nv-1];
-    for (int i=0; i < nv; i++) {
+    if (d->nnzL > nv*nv) {  // SHOULD NOT OCCUR
-      ctx->nnz += ctx->rownnz[i];
-    }
-    if (ctx->nnz > 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);
+    mju_sqrMatTD(d->L, d->efc_J, D, nefc, nv);
-    mj_addMDense(m, d, ctx->H);
+    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,
+        rank = mju_cholUpdateSparse(d->L, vec, nv, flag_update,
-                                    ctx->rownnz, ctx->rowadr, ctx->colind, nnz, vec_ind,
+                                    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;
       }