Dynamic model editing.

PiperOrigin-RevId: 640838852
Change-Id: I559b1d60f2a3a2d866788d4006be31757454c619

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;
 }