Change mjSpec and mjModel signature mechanism.

The signature now contains the necessary information to safely perform `bind`. The private UIDs are now removed.

Fixes an issue of changing signature when compiling a copy of an mjSpec.

PiperOrigin-RevId: 744670626
Change-Id: Id3c66419cf2afbe78e91bc4b37d2299f5ec00ab1

mjx/mujoco/mjx/_src/support_test.py

@@ -349,7 +349,7 @@ class SupportTest(parameterized.TestCase):
     self.assertEqual(
         str(e.exception),
         'mjSpec signature does not match mjx.Model signature:'
-        ' 5495345807332648606 != 270010677651259353',
+        ' 4300287342280373816 != 9887180086914550999',
     )
 
   _CONTACTS = """

src/user/user_flexcomp.cc

@@ -403,14 +403,12 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz) {
   mjCFlex* flex = model->AddFlex();
   mjsFlex* pf = &flex->spec;
   int id = flex->id;
-  int uid = flex->uid;
 
   *flex = def.Flex();
   flex->PointToLocal();
 
   flex->model = model;
   flex->id = id;
-  flex->uid = uid;
   mjs_setString(pf->name, name.c_str());
   mjs_setInt(pf->elem, element.data(), element.size());
   mjs_setFloat(pf->texcoord, texcoord.data(), texcoord.size());

src/user/user_model.cc

@@ -201,7 +201,6 @@ mjCModel::mjCModel() {
   world->mass = 0;
   mjuu_zerovec(world->inertia, 3);
   world->id = 0;
-  world->uid = GetUid();
   world->parent = nullptr;
   world->weldid = 0;
   world->name = "world";
@@ -297,7 +296,6 @@ void mjCModel::CopyList(std::vector<T*>& dest,
     // copy the element from the other model to this model
     if (deepcopy_) {
       source[i]->ForgetKeyframes();
-      candidate->uid = source[i]->uid;
     } else {
       candidate->AddRef();
     }
@@ -1038,7 +1036,6 @@ template <class T>
 T* mjCModel::AddObject(vector<T*>& list, string type) {
   T* obj = new T(this);
   obj->id = (int)list.size();
-  obj->uid = GetUid();
   list.push_back(obj);
   spec.element->signature = Signature();
   return obj;
@@ -1051,7 +1048,6 @@ T* mjCModel::AddObjectDefault(vector<T*>& list, string type, mjCDef* def) {
   T* obj = new T(this, def ? def : defaults_[0]);
   obj->id = (int)list.size();
   obj->classname = def ? def->name : "main";
-  obj->uid = GetUid();
   list.push_back(obj);
   spec.element->signature = Signature();
   return obj;
@@ -3615,7 +3611,7 @@ void mjCModel::SaveState(const std::string& state_name, const T* qpos, const T*
   }
 
   for (auto body : bodies_) {
-    if (!body->spec.mocap) {
+    if (!body->spec.mocap || body->mocapid == -1) {
       continue;
     }
     if (mpos) {
@@ -4388,6 +4384,9 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
     bodies_[i]->subtreedofs = 0;
   }
 
+  // initialize spec signature (needed if the user changed sensor or joint types)
+  spec.element->signature = Signature();
+
   // fill missing names and check that they are all filled
   for (const auto& asset : meshes_) asset->CopyFromSpec();
   for (const auto& asset : skins_) asset->CopyFromSpec();
@@ -4651,7 +4650,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
 
   // special cases that are not caused by user edits
   if (compiler.fusestatic || compiler.discardvisual ||
-      !spec.element->signature || !pairs_.empty() || !excludes_.empty()) {
+      !pairs_.empty() || !excludes_.empty()) {
     spec.element->signature = m->signature;
   }
 
@@ -4663,21 +4662,78 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
 
 
 
-uint64_t mjCModel::Signature() {
+std::string mjCModel::PrintTree(const mjCBody* body, std::string indent) {
-  std::string uid_str;
+  std::string tree;
-  for (int i = 0; i < mjNOBJECT; ++i) {
+  tree += indent + "<body>\n";
-    if (i == mjOBJ_XBODY || i == mjOBJ_UNKNOWN || i == mjOBJ_DOF) {
+  indent += "  ";
-      continue;
+  for (const auto& joint : body->joints) {
-    }
+    tree += indent + "<joint>" + std::to_string(joint->nq()) + "</joint>\n";
-    if (object_lists_[i] == nullptr) {
-      throw mjCError(0, "object list %s is null", std::to_string(i).c_str());
-    }
-    uid_str += '|';
-    for (mjCBase* object : *object_lists_[i]) {
-      uid_str += std::to_string(object->uid) + " ";
-    }
   }
-  return mj_hashString(uid_str.c_str(), UINT64_MAX);
+  for (uint64_t i = 0; i < body->geoms.size(); ++i) {
+    tree += indent + "<geom/>\n";
+  }
+  for (uint64_t i = 0; i < body->sites.size(); ++i) {
+    tree += indent + "<site/>\n";
+  }
+  for (uint64_t i = 0; i < body->cameras.size(); ++i) {
+    tree += indent + "<camera/>\n";
+  }
+  for (uint64_t i = 0; i < body->lights.size(); ++i) {
+    tree += indent + "<light/>\n";
+  }
+  for (uint64_t i = 0; i < body->bodies.size(); ++i) {
+    tree += PrintTree(body->bodies[i], indent);
+  }
+  indent.pop_back();
+  indent.pop_back();
+  tree += indent + "</body>\n";
+  return tree;
+}
+
+
+
+uint64_t mjCModel::Signature() {
+  std::string tree = "\n" + PrintTree(bodies_[0]);
+  for (unsigned int i = 0; i < flexes_.size(); ++i) {
+    tree += "<flex/>\n";
+  }
+  for (unsigned int i = 0; i < meshes_.size(); ++i) {
+    tree += "<mesh/>\n";
+  }
+  for (unsigned int i = 0; i < skins_.size(); ++i) {
+    tree += "<skin/>\n";
+  }
+  for (unsigned int i = 0; i < hfields_.size(); ++i) {
+    tree += "<heightfield/>\n";
+  }
+  for (unsigned int i = 0; i < textures_.size(); ++i) {
+    tree += "<texture/>\n";
+  }
+  for (unsigned int i = 0; i < materials_.size(); ++i) {
+    tree += "<material/>\n";
+  }
+  for (unsigned int i = 0; i < pairs_.size(); ++i) {
+    tree += "<pair/>\n";
+  }
+  for (unsigned int i = 0; i < excludes_.size(); ++i) {
+    tree += "<exclude/>\n";
+  }
+  for (unsigned int i = 1; i < equalities_.size(); ++i) {
+    tree += "<equality/>\n";
+  }
+  for (unsigned int i = 0; i < tendons_.size(); ++i) {
+    tree += "<tendon/>\n";
+  }
+  for (unsigned int i = 0; i < actuators_.size(); ++i) {
+    tree += "<actuator/>\n";
+  }
+  for (unsigned int i = 0; i < sensors_.size(); ++i) {
+    tree += "<sensor>" + std::to_string(sensors_[i]->spec.type) + "<sensor/>\n";
+  }
+  for (unsigned int i = 0; i < keys_.size(); ++i) {
+    tree += "<key/>\n";
+  }
+  return mj_hashString(tree.c_str(), UINT64_MAX);
 }
 
 

src/user/user_model.h

@@ -324,9 +324,6 @@ class mjCModel : public mjCModel_, private mjSpec {
   // set attached flag
   void SetAttached(bool deepcopy) { attached_ |= !deepcopy; }
 
-  // get new uid
-  int GetUid() { return uid_count_++; }
-
  private:
   // settings for each defaults class
   std::vector<mjCDef*> defaults_;
@@ -441,6 +438,9 @@ class mjCModel : public mjCModel_, private mjSpec {
   void MarkPluginInstance(std::unordered_map<std::string, bool>& instances,
                           const std::vector<T*>& list);
 
+  // print the tree of a body
+  std::string PrintTree(const mjCBody* body, std::string indent = "");
+
   // generate a signature for the model
   uint64_t Signature();
 
@@ -449,7 +449,6 @@ class mjCModel : public mjCModel_, private mjSpec {
   std::vector<mjKeyInfo> key_pending_;  // attached keyframes
   bool deepcopy_;     // copy objects when attaching
   bool attached_ = false;  // true if model is attached to a parent model
-  int uid_count_ = 0;  // unique id count for all objects
   std::unordered_map<const mjsCompiler*, mjSpec*> compiler2spec_;  // map from compiler to spec
 };
 #endif  // MUJOCO_SRC_USER_USER_MODEL_H_

src/user/user_objects.cc

@@ -818,6 +818,7 @@ mjCBody::mjCBody(mjCModel* _model) {
   mjuu_zerovec(xpos0, 3);
   mjuu_setvec(xquat0, 1, 0, 0, 0);
   last_attached = nullptr;
+  mocapid = -1;
 
   // clear object lists
   bodies.clear();
@@ -840,7 +841,6 @@ mjCBody::mjCBody(mjCModel* _model) {
 
 mjCBody::mjCBody(const mjCBody& other, mjCModel* _model) {
   model = _model;
-  uid = other.uid;
   mjSpec* origin = model->FindSpec(other.compiler);
   compiler = origin ? &origin->compiler : &model->spec.compiler;
   *this = other;
@@ -944,7 +944,6 @@ mjCBody& mjCBody::operator+=(const mjCFrame& other) {
   frames.back()->frame = other.frame;
   if (model->deepcopy_) {
     frames.back()->NameSpace(other_model);
-    frames.back()->uid = other.uid;
   } else {
     frames.back()->AddRef();
   }
@@ -1037,8 +1036,6 @@ void mjCBody::CopyList(std::vector<T*>& dst, const std::vector<T*>& src,
     // increment refcount if shallow copy is made
     if (!model->deepcopy_) {
       dst.back()->AddRef();
-    } else {
-      dst.back()->uid = src[i]->uid;
     }
 
     // set namespace
@@ -1256,7 +1253,6 @@ mjCBody* mjCBody::AddBody(mjCDef* _def) {
   obj->parent = this;
 
   // update signature
-  obj->uid = model->GetUid();
   model->spec.element->signature = model->Signature();
   return obj;
 }
@@ -1271,7 +1267,6 @@ mjCFrame* mjCBody::AddFrame(mjCFrame* _frame) {
   model->MakeTreeLists();
 
   // update signature
-  obj->uid = model->GetUid();
   model->spec.element->signature = model->Signature();
   return obj;
 }
@@ -1295,7 +1290,6 @@ mjCJoint* mjCBody::AddFreeJoint() {
 
 
   // update signature
-  obj->uid = model->GetUid();
   model->spec.element->signature = model->Signature();
   return obj;
 }
@@ -1318,7 +1312,6 @@ mjCJoint* mjCBody::AddJoint(mjCDef* _def) {
 
 
   // update signature
-  obj->uid = model->GetUid();
   model->spec.element->signature = model->Signature();
   return obj;
 }
@@ -1341,7 +1334,6 @@ mjCGeom* mjCBody::AddGeom(mjCDef* _def) {
 
 
   // update signature
-  obj->uid = model->GetUid();
   model->spec.element->signature = model->Signature();
   return obj;
 }
@@ -1364,7 +1356,6 @@ mjCSite* mjCBody::AddSite(mjCDef* _def) {
 
 
   // update signature
-  obj->uid = model->GetUid();
   model->spec.element->signature = model->Signature();
   return obj;
 }
@@ -1387,7 +1378,6 @@ mjCCamera* mjCBody::AddCamera(mjCDef* _def) {
 
 
   // update signature
-  obj->uid = model->GetUid();
   model->spec.element->signature = model->Signature();
   return obj;
 }
@@ -1410,7 +1400,6 @@ mjCLight* mjCBody::AddLight(mjCDef* _def) {
 
 
   // update signature
-  obj->uid = model->GetUid();
   model->spec.element->signature = model->Signature();
   return obj;
 }

src/user/user_objects.h

@@ -219,7 +219,6 @@ class mjCBoundingVolumeHierarchy : public mjCBoundingVolumeHierarchy_ {
 class mjCBase_ : public mjsElement {
  public:
   int id;                 // object id
-  int uid;                // unique identifier
   std::string name;       // object name
   std::string classname;  // defaults class name
   std::string info;       // error message info set by the user

test/user/user_api_test.cc

@@ -550,6 +550,10 @@ TEST_F(PluginTest, RecompileCompare) {
         mjModel* m_new = mj_compile(s, nullptr);
         mjModel* m_copy = mj_compile(s_copy, nullptr);
 
+        // compare signature
+        EXPECT_EQ(m_old->signature, m_new->signature) << xml;
+        EXPECT_EQ(m_old->signature, m_copy->signature) << xml;
+
         ASSERT_THAT(m_new, NotNull())
             << "Failed to recompile " << xml << ": " << mjs_getError(s);
         ASSERT_THAT(m_copy, NotNull())