Fix a bug in mj_recompile.

The qposadr_ and dofadr_ variables of the joints were not being reset to -1 when a joint is copied. This could lead to errors when calling mj_recompile after attaching a compiled spec.

PiperOrigin-RevId: 686152004
Change-Id: I9e058089ce87ea0d8e51a7f43475bc48c6527a26

src/user/user_api.cc

@@ -96,21 +96,26 @@ mjModel* mj_compile(mjSpec* s, const mjVFS* vfs) {
   mjCModel* modelC = static_cast<mjCModel*>(s->element);
   std::string state_name = "state";
   mjtNum time = 0;
-  if (d) {
-    time = d->time;
-    modelC->SaveState(state_name, d->qpos, d->qvel, d->act, d->ctrl, d->mocap_pos, d->mocap_quat);
-  }
-  if (!modelC->Compile(vfs, &m)) {
+  try {
     if (d) {
-      mj_deleteData(d);
+      time = d->time;
+      modelC->SaveState(state_name, d->qpos, d->qvel, d->act, d->ctrl, d->mocap_pos, d->mocap_quat);
     }
+    if (!modelC->Compile(vfs, &m)) {
+      if (d) {
+        mj_deleteData(d);
+      }
+      return -1;
+    };
+    if (d) {
+      modelC->MakeData(m, &d);
+      modelC->RestoreState(state_name, m->qpos0, m->body_pos, m->body_quat, d->qpos, d->qvel,
+                          d->act, d->ctrl, d->mocap_pos, d->mocap_quat);
+      d->time = time;
+    }
+  } catch (mjCError& e) {
+    modelC->SetError(e);
     return -1;
-  };
-  if (d) {
-    modelC->MakeData(m, &d);
-    modelC->RestoreState(state_name, m->qpos0, m->body_pos, m->body_quat, d->qpos, d->qvel,
-                         d->act, d->ctrl, d->mocap_pos, d->mocap_quat);
-    d->time = time;
   }
   return 0;
 }

src/user/user_model.cc

@@ -353,9 +353,7 @@ mjCModel& mjCModel::operator+=(const mjCModel& other) {
 
   // create global lists
   mjCBody *world = bodies_[0];
-  if (compiled) {
-    ResetTreeLists();
-  }
+  ResetTreeLists();
   MakeLists(world);
   ProcessLists(/*checkrepeat=*/false);
 
@@ -3034,11 +3032,15 @@ template <class T>
 void mjCModel::SaveState(const std::string& state_name, const T* qpos, const T* qvel, const T* act,
                          const T* ctrl, const T* mpos, const T* mquat) {
   for (auto joint : joints_) {
-    if (joint->qposadr_ == -1 || joint->dofadr_ == -1) {
-      throw mjCError(nullptr, "SaveState: joint %s has no address", joint->name.c_str());
+    if (joint->qposadr_ < -1 || joint->dofadr_ < -1) {
+      throw mjCError(nullptr, "SaveState: joint %s has invalid address", joint->name.c_str());
+    }
+    if (qpos && joint->qposadr_ != -1) {
+      mjuu_copyvec(joint->qpos(state_name), qpos + joint->qposadr_, joint->nq());
+    }
+    if (qvel && joint->dofadr_ != -1) {
+      mjuu_copyvec(joint->qvel(state_name), qvel + joint->dofadr_, joint->nv());
     }
-    if (qpos) mjuu_copyvec(joint->qpos(state_name), qpos + joint->qposadr_, joint->nq());
-    if (qvel) mjuu_copyvec(joint->qvel(state_name), qvel + joint->dofadr_, joint->nv());
   }
 
   for (unsigned int i=0; i<actuators_.size(); i++) {

src/user/user_objects.cc

@@ -1951,6 +1951,8 @@ mjCJoint& mjCJoint::operator=(const mjCJoint& other) {
     this->spec = other.spec;
     *static_cast<mjCJoint_*>(this) = static_cast<const mjCJoint_&>(other);
     *static_cast<mjsJoint*>(this) = static_cast<const mjsJoint&>(other);
+    qposadr_ = -1;
+    dofadr_ = -1;
   }
   PointToLocal();
   return *this;

src/user/user_objects.h

@@ -425,8 +425,6 @@ class mjCJoint_ : public mjCBase {
   mjCBody* body;                   // joint's body
 
   // variable used for temporarily storing the state of the joint
-  int qposadr_;                                        // address of dof in data->qpos
-  int dofadr_;                                         // address of dof in data->qvel
   std::map<std::string, std::array<mjtNum, 7>> qpos_;  // qpos at the previous step
   std::map<std::string, std::array<mjtNum, 6>> qvel_;  // qvel at the previous step
 
@@ -473,6 +471,10 @@ class mjCJoint : public mjCJoint_, private mjsJoint {
  private:
   int Compile(void);               // compiler; return dofnum
   void PointToLocal(void);
+
+  // variables that should not be copied during copy assignment
+  int qposadr_;                                        // address of dof in data->qpos
+  int dofadr_;                                         // address of dof in data->qvel
 };
 
 

test/user/user_api_test.cc

@@ -1368,6 +1368,54 @@ TEST_F(MujocoTest, PreserveState) {
   mj_deleteModel(m_expected);
 }
 
+TEST_F(MujocoTest, RecompileAttach) {
+  std::array<char, 1000> er;
+  std::string field = "";
+
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <body name="body">
+        <joint type="slide" axis="0 0 1"/>
+        <geom size=".2"/>
+      </body>
+    </worldbody>
+  </mujoco>)";
+
+  mjSpec* parent = mj_makeSpec();
+  EXPECT_THAT(parent, NotNull());
+
+  mjSpec* child = mj_parseXMLString(xml, 0, er.data(), er.size());
+  EXPECT_THAT(child, NotNull());
+
+  mjsFrame* frame1 = mjs_addFrame(mjs_findBody(parent, "world"), 0);
+  mjs_attachBody(frame1, mjs_findBody(child, "body"), "child-", "-1");
+
+  mjModel* model = mj_compile(parent, 0);
+  EXPECT_THAT(model, NotNull());
+  EXPECT_THAT(model->nq, 1);
+  mjData* data = mj_makeData(model);
+  EXPECT_THAT(data, NotNull());
+
+  for (int i = 0; i < 100; i++) {
+    mj_step(model, data);
+  }
+
+  mjsFrame* frame2 = mjs_addFrame(mjs_findBody(parent, "world"), 0);
+  mjs_attachBody(frame2, mjs_findBody(child, "body"), "child-", "-2");
+
+  EXPECT_EQ(mj_recompile(parent, 0, model, data), 0);
+  EXPECT_THAT(model, NotNull());
+  EXPECT_THAT(model->nq, 2);
+  EXPECT_NE(data->qpos[0], data->qpos[1]);
+  EXPECT_EQ(data->qpos[1], 0);
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+  mj_deleteSpec(child);
+  mj_deleteSpec(parent);
+}
+
 TEST_F(MujocoTest, AttachMocap) {
   std::array<char, 1000> er;
   mjtNum tol = 0;