Dynamic model editing.

PiperOrigin-RevId: 640838852
Change-Id: I559b1d60f2a3a2d866788d4006be31757454c619

src/engine/engine_io.c

@@ -1093,7 +1093,7 @@ static void mj_setPtrData(const mjModel* m, mjData* d) {
 
 
 // initialize plugins, copy into d (required for deletion)
-static void _initPlugin(const mjModel* m, mjData* d) {
+void mj_initPlugin(const mjModel* m, mjData* d) {
   d->nplugin = m->nplugin;
   for (int i = 0; i < m->nplugin; ++i) {
     d->plugin[i] = m->plugin[i];
@@ -1203,7 +1203,7 @@ mjData* mj_makeData(const mjModel* m) {
   mjData* d = NULL;
   mj_makeRawData(&d, m);
   if (d) {
-    _initPlugin(m, d);
+    mj_initPlugin(m, d);
     mj_resetData(m, d);
   }
   return d;
@@ -1219,7 +1219,7 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) {
   // allocate new data if needed
   if (!dest) {
     mj_makeRawData(&dest, m);
-    _initPlugin(m, dest);
+    mj_initPlugin(m, dest);
   }
 
   // check sizes

src/engine/engine_io.h

@@ -110,6 +110,9 @@ MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
 // mjData arena allocate
 MJAPI void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment);
 
+// init plugins
+MJAPI void mj_initPlugin(const mjModel* m, mjData* d);
+
 #ifndef ADDRESS_SANITIZER
 
 // mjData mark stack frame

src/user/user_api.cc

@@ -83,6 +83,16 @@ mjModel* mjs_compile(mjSpec* s, const mjVFS* vfs) {
 
 
 
+// recompile spec into existing model and data while preserving the state
+void mjs_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d) {
+  mjCModel* modelC = static_cast<mjCModel*>(s->element);
+  modelC->SaveState(d);
+  modelC->Compile(vfs, &m);
+  modelC->RestoreState(m, &d);
+}
+
+
+
 // attach body to a frame of the parent
 int mjs_attachBody(mjsFrame* parent, const mjsBody* child,
                    const char* prefix, const char* suffix) {

src/user/user_api.h

@@ -17,6 +17,7 @@
 
 #include <math.h>
 #include <stddef.h>
+#include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjtnum.h>
@@ -754,6 +755,9 @@ MJAPI mjSpec* mjs_createSpec(void);
 // Compile spec to model.
 MJAPI mjModel* mjs_compile(mjSpec* s, const mjVFS* vfs);
 
+// Recompile spec to model preserving the current state.
+MJAPI void mjs_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d);
+
 // Copy spec.
 MJAPI mjSpec* mjs_copySpec(const mjSpec* s);
 

src/user/user_model.cc

@@ -1843,6 +1843,8 @@ void mjCModel::CopyTree(mjModel* m) {
       int jid = pj->id;
 
       // set joint fields
+      pj->qposadr_ = qposadr;
+      pj->dofadr_ = dofadr;
       m->jnt_type[jid] = pj->type;
       m->jnt_group[jid] = pj->group;
       m->jnt_limited[jid] = (mjtByte)pj->is_limited();
@@ -2584,6 +2586,8 @@ void mjCModel::CopyObjects(mjModel* m) {
     m->actuator_trnid[2*i+1] = pac->trnid[1];
     m->actuator_actnum[i] = pac->actdim + pac->plugin_actdim;
     m->actuator_actadr[i] = m->actuator_actnum[i] ? adr : -1;
+    pac->actadr_ = m->actuator_actadr[i];
+    pac->actnum_ = m->actuator_actnum[i];
     adr += m->actuator_actnum[i];
     m->actuator_group[i] = pac->group;
     m->actuator_ctrllimited[i] = (mjtByte)pac->is_ctrllimited();
@@ -2716,6 +2720,75 @@ void mjCModel::CopyObjects(mjModel* m) {
 
 
 
+// save the current state
+void mjCModel::SaveState(const mjData* d) {
+  for (auto joint : joints_) {
+    switch (joint->type) {
+      case mjJNT_FREE:
+        mju_copy(joint->qpos, d->qpos + joint->qposadr_, 7);
+        mju_copy(joint->qvel, d->qvel + joint->dofadr_, 6);
+        break;
+      case mjJNT_BALL:
+        mju_copy(joint->qpos, d->qpos + joint->qposadr_, 4);
+        mju_copy(joint->qvel, d->qvel + joint->dofadr_, 3);
+        break;
+      case mjJNT_HINGE:
+      case mjJNT_SLIDE:
+        mju_copy(joint->qpos, d->qpos + joint->qposadr_, 1);
+        mju_copy(joint->qvel, d->qvel + joint->dofadr_, 1);
+        break;
+    }
+  }
+
+  for (auto actuator : actuators_) {
+    if (actuator->actadr_ != -1) {
+      actuator->act.assign(actuator->actnum_, 0);
+      mju_copy(actuator->act.data(), d->act + actuator->actadr_, actuator->actnum_);
+    }
+  }
+}
+
+
+
+// restore the previous state
+void mjCModel::RestoreState(const mjModel* m, mjData** dest) {
+  mj_makeRawData(dest, m);
+  mjData* d = *dest;
+  if (d) {
+    mj_initPlugin(m, d);
+    mj_resetData(m, d);
+  }
+
+  for (auto joint : joints_) {
+    if (!mjuu_defined(joint->qpos[0]) || !mjuu_defined(joint->qvel[0])) {
+      continue;
+    }
+    switch (joint->type) {
+      case mjJNT_FREE:
+        mju_copy(d->qpos + joint->qposadr_, joint->qpos, 7);
+        mju_copy(d->qvel + joint->dofadr_, joint->qvel, 6);
+        break;
+      case mjJNT_BALL:
+        mju_copy(d->qpos + joint->qposadr_, joint->qpos, 4);
+        mju_copy(d->qvel + joint->dofadr_, joint->qvel, 3);
+        break;
+      case mjJNT_HINGE:
+      case mjJNT_SLIDE:
+        mju_copy(d->qpos + joint->qposadr_, joint->qpos, 1);
+        mju_copy(d->qvel + joint->dofadr_, joint->qvel, 1);
+        break;
+    }
+  }
+
+  for (auto actuator : actuators_) {
+    if (mjuu_defined(actuator->act[0])) {
+      mju_copy(d->act + actuator->actadr_, actuator->act.data(), actuator->actnum_);
+    }
+  }
+}
+
+
+
 //------------------------------- FUSE STATIC ------------------------------------------------------
 
 template <class T>
@@ -3075,7 +3148,7 @@ static void warninghandler(const char* msg) {
 
 
 // compiler
-mjModel* mjCModel::Compile(const mjVFS* vfs) {
+mjModel* mjCModel::Compile(const mjVFS* vfs, mjModel** m) {
   if (compiled) {
     // clear kinematic tree
     for (int i=0; i<bodies_.size(); i++) {
@@ -3096,7 +3169,7 @@ mjModel* mjCModel::Compile(const mjVFS* vfs) {
   // setjmp returns, and therefore to pass nullptr directly to the
   // mj_deleteModel and mj_deleteData calls in the subsequent catch block,
   // without ever reading the actual pointer values.
-  mjModel* volatile m = nullptr;
+  mjModel* volatile model = (m && *m) ? *m : nullptr;
   mjData* volatile data = nullptr;
 
   // save error and warning handlers
@@ -3118,10 +3191,10 @@ mjModel* mjCModel::Compile(const mjVFS* vfs) {
       // TryCompile resulted in an mju_error which was converted to a longjmp.
       throw mjCError(0, "engine error: %s", errortext);
     }
-    TryCompile(*const_cast<mjModel**>(&m), *const_cast<mjData**>(&data), vfs);
+    TryCompile(*const_cast<mjModel**>(&model), *const_cast<mjData**>(&data), vfs);
   } catch (mjCError err) {
     // deallocate everything allocated in Compile
-    mj_deleteModel(m);
+    mj_deleteModel(model);
     mj_deleteData(data);
     mjCBody* world = bodies_[0];
     Clear();
@@ -3140,7 +3213,7 @@ mjModel* mjCModel::Compile(const mjVFS* vfs) {
   _mjPRIVATE__set_tls_error_fn(save_error);
   _mjPRIVATE__set_tls_warning_fn(save_warning);
   compiled = true;
-  return m;
+  return model;
 }
 
 

src/user/user_model.h

@@ -168,7 +168,7 @@ class mjCModel : public mjCModel_, private mjSpec {
 
   mjSpec spec;
 
-  mjModel* Compile(const mjVFS* vfs = nullptr);  // construct mjModel
+  mjModel* Compile(const mjVFS* vfs = nullptr, mjModel** m = nullptr);  // construct mjModel
   bool CopyBack(const mjModel*);                 // DECOMPILER: copy numeric back
   void FuseStatic();                             // fuse static bodies with parent
   void FuseReindex(mjCBody* body);               // reindex elements during fuse
@@ -254,9 +254,6 @@ class mjCModel : public mjCModel_, private mjSpec {
                      const std::string& plugin_instance_name,
                      mjCPlugin** plugin_instance);
 
-  void TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs);
-  mjModel* _Compile(const mjVFS* vfs);
-
   // clear objects allocated by Compile
   void Clear();
 
@@ -267,6 +264,10 @@ class mjCModel : public mjCModel_, private mjSpec {
   template <class T> void DeleteMaterial(std::vector<T*>& list,
                                          std::string_view name = "");
 
+  // save/restore the current state
+  void SaveState(const mjData* d);
+  void RestoreState(const mjModel* m, mjData** dest);
+
  private:
   // settings for each defaults class
   std::vector<mjCDef*> defaults_;
@@ -275,6 +276,7 @@ class mjCModel : public mjCModel_, private mjSpec {
   std::vector<std::pair<const mjpPlugin*, int>> active_plugins_;
 
   // compile phases
+  void TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs);
   void MakeLists(mjCBody* body);        // make lists of bodies, geoms, joints, sites
   void SetNuser();                      // set nuser fields
   void IndexAssets(bool discard);       // convert asset names into indices

src/user/user_objects.cc

@@ -1733,6 +1733,10 @@ mjCJoint::mjCJoint(mjCModel* _model, mjCDef* _def) {
 
   // in case this joint is not compiled
   CopyFromSpec();
+
+  // no previous state when a joint is created
+  qpos[0] = mjNAN;
+  qvel[0] = mjNAN;
 }
 
 
@@ -5003,6 +5007,9 @@ mjCActuator::mjCActuator(mjCModel* _model, mjCDef* _def) {
 
   // point to local
   PointToLocal();
+
+  // no previous state when an actuator is created
+  act.push_back(mjNAN);
 }
 
 

src/user/user_objects.h

@@ -423,10 +423,14 @@ class mjCJoint : public mjCJoint_, private mjsJoint {
   bool is_limited() const;
   bool is_actfrclimited() const;
 
-
  private:
   int Compile(void);               // compiler; return dofnum
   void PointToLocal(void);
+
+  int qposadr_;                    // address of dof in data->qpos
+  int dofadr_;                     // address of dof in data->qvel
+  mjtNum qpos[7];                  // qpos at the previous step
+  mjtNum qvel[6];                  // qvel at the previous step
 };
 
 
@@ -1412,6 +1416,10 @@ class mjCActuator : public mjCActuator_, private mjsActuator {
   void NameSpace(const mjCModel* m);
 
   mjCBase* ptarget;  // transmission target
+
+  int actadr_;              // address of dof in data->act
+  int actnum_;              // number of dofs in data->act
+  std::vector<mjtNum> act;  // act at the previous step
 };
 
 

test/user/user_api_test.cc

@@ -747,5 +747,124 @@ TEST_F(MujocoTest, DetachBody) {
   TestDetachBody(/*compile=*/true);
 }
 
+TEST_F(MujocoTest, PreserveState) {
+  std::array<char, 1000> er;
+  std::string field = "";
+
+  static constexpr char xml_full[] = R"(
+  <mujoco>
+    <worldbody>
+      <body name="detachable" pos="1 0 0">
+        <joint type="hinge" axis="0 0 1" name="hinge"/>
+        <geom type="sphere" size=".1"/>
+      </body>
+      <body name="persistent">
+        <joint type="slide" axis="0 0 1" name="slide"/>
+        <geom type="sphere" size=".2"/>
+      </body>
+    </worldbody>
+    <actuator>
+      <position name="hinge" joint="hinge" timeconst=".01"/>
+      <position name="slide" joint="slide" timeconst=".01"/>
+    </actuator>
+  </mujoco>)";
+
+  static constexpr char xml_expected[] = R"(
+  <mujoco>
+    <worldbody>
+      <body name="persistent">
+        <joint type="slide" axis="0 0 1" name="slide"/>
+        <geom type="sphere" size=".2"/>
+      </body>
+      <body name="newbody" pos="2 0 0">
+        <joint type="slide" axis="0 0 1"/>
+        <geom type="sphere" size=".3"/>
+      </body>
+    </worldbody>
+    <actuator>
+      <position name="slide" joint="slide" timeconst=".01"/>
+    </actuator>
+  </mujoco>)";
+
+  // load spec
+  mjSpec* spec = mj_parseXMLString(xml_full, 0, er.data(), er.size());
+  EXPECT_THAT(spec, NotNull()) << er.data();
+
+  // compile models
+  mjModel* model = mjs_compile(spec, 0);
+  EXPECT_THAT(model, NotNull());
+  mjModel* m_expected = LoadModelFromString(xml_expected, er.data(), er.size());
+  EXPECT_THAT(m_expected, NotNull());
+
+  // create data
+  mjData* data = mj_makeData(model);
+  EXPECT_THAT(data, NotNull());
+  mjData* d_expected = mj_makeData(m_expected);
+  EXPECT_THAT(d_expected, NotNull());
+
+  // set ctrl
+  data->ctrl[0] = 1;
+  data->ctrl[1] = 2;
+  d_expected->ctrl[0] = 2;
+
+  // step models
+  mj_step(model, data);
+  mj_step(m_expected, d_expected);
+
+  // detach subtree
+  mjsBody* body = mjs_findBody(spec, "detachable");
+  EXPECT_THAT(body, NotNull());
+  EXPECT_THAT(mjs_detachBody(spec, body), 0);
+
+  // add body
+  mjsBody* newbody = mjs_addBody(mjs_findBody(spec, "world"), 0);
+  EXPECT_THAT(newbody, NotNull());
+
+  // add geom and joint
+  mjsGeom* geom = mjs_addGeom(newbody, 0);
+  mjsJoint* joint = mjs_addJoint(newbody, 0);
+
+  // set properties
+  newbody->pos[0] = 2;
+  geom->size[0] = .3;
+  joint->type = mjJNT_SLIDE;
+  joint->axis[0] = 0;
+  joint->axis[1] = 0;
+  joint->axis[2] = 1;
+  joint->ref = d_expected->qpos[m_expected->nq-1];
+
+  // compile new model
+  mjs_recompile(spec, 0, model, data);
+  EXPECT_THAT(model, NotNull());
+
+  // compare qpos
+  EXPECT_EQ(model->nq, m_expected->nq);
+  for (int i = 0; i < model->nq; ++i) {
+    EXPECT_EQ(data->qpos[i], d_expected->qpos[i]) << i;
+  }
+
+  // compare qvel
+  EXPECT_EQ(model->nv, m_expected->nv);
+  for (int i = 0; i < model->nv-1; ++i) {
+    EXPECT_EQ(data->qvel[i], d_expected->qvel[i]) << i;
+  }
+
+  // second body was added after stepping so qvel should be zero
+  EXPECT_EQ(data->qvel[model->nv-1], 0);
+
+  // compare act
+  EXPECT_EQ(model->na, m_expected->na);
+  for (int i = 0; i < model->na; ++i) {
+    EXPECT_EQ(data->act[i], d_expected->act[i]) << i;
+  }
+
+  // destroy everything
+  mj_deleteData(data);
+  mj_deleteData(d_expected);
+  mjs_deleteSpec(spec);
+  mj_deleteModel(model);
+  mj_deleteModel(m_expected);
+}
+
 }  // namespace
 }  // namespace mujoco