Add signature to mjSpec and mjModel and use it to perform safe bind to mjModel and mjData.

PiperOrigin-RevId: 740378879
Change-Id: If14b326942529494f172e7aedcae30195798b458

Author: quagla

doc/includes/references.h

@@ -414,6 +414,9 @@ struct mjData_ {
 
   // thread pool pointer
   uintptr_t threadpool;
+
+  // compilation signature
+  uint64_t  signature;       // also held by the mjSpec that compiled the model
 };
 typedef struct mjData_ mjData;
 typedef enum mjtDisableBit_ {     // disable default feature bitflags
@@ -1451,6 +1454,9 @@ struct mjModel_ {
 
   // paths
   char*     paths;                // paths to assets, 0-terminated            (npaths x 1)
+
+  // compilation signature
+  uint64_t  signature;            // also held by the mjSpec that compiled this model
 };
 typedef struct mjModel_ mjModel;
 struct mjResource_ {
@@ -1713,6 +1719,7 @@ typedef enum mjtOrientation_ {     // type of orientation specifier
 } mjtOrientation;
 typedef struct mjsElement_ {       // element type, do not modify
   mjtObj elemtype;                 // element type
+  uint64_t signature;              // compilation signature
 } mjsElement;
 typedef struct mjsCompiler_ {      // compiler options
   mjtByte autolimits;              // infer "limited" attribute based on range

include/mujoco/mjdata.h

@@ -442,6 +442,9 @@ struct mjData_ {
 
   // thread pool pointer
   uintptr_t threadpool;
+
+  // compilation signature
+  uint64_t  signature;       // also held by the mjSpec that compiled the model
 };
 typedef struct mjData_ mjData;
 

include/mujoco/mjmodel.h

@@ -1155,6 +1155,9 @@ struct mjModel_ {
 
   // paths
   char*     paths;                // paths to assets, 0-terminated            (npaths x 1)
+
+  // compilation signature
+  uint64_t  signature;            // also held by the mjSpec that compiled this model
 };
 typedef struct mjModel_ mjModel;
 

include/mujoco/mjspec.h

@@ -23,6 +23,7 @@
 // this is a C-API
 #ifdef __cplusplus
 #include <cstddef>
+#include <cstdint>
 #include <string>
 #include <vector>
 
@@ -119,6 +120,7 @@ typedef enum mjtOrientation_ {     // type of orientation specifier
 
 typedef struct mjsElement_ {       // element type, do not modify
   mjtObj elemtype;                 // element type
+  uint64_t signature;              // compilation signature
 } mjsElement;
 
 

python/mujoco/codegen/generate_spec_bindings.py

@@ -613,12 +613,27 @@ def generate_find() -> None:
     print(code)
 
 
+def generate_signature() -> None:
+  """Generate signature functions."""
+  for key, _, _, _, _ in SPECS:
+    elem = key.removeprefix('mjs')
+    titlecase = 'Mjs' + elem
+    code = f"""\n
+      {key}.def_property_readonly("signature",
+      [](raw::{titlecase}& self) -> uint64_t {{
+        return mjs_getSpec(self.element)->element->signature;
+      }});
+    """
+    print(code)
+
+
 def main(argv: Sequence[str]) -> None:
   if len(argv) > 1:
     raise app.UsageError('Too many command-line arguments.')
   generate()
   generate_add()
   generate_find()
+  generate_signature()
 
 
 if __name__ == '__main__':

python/mujoco/introspect/structs.py

@@ -4454,6 +4454,11 @@
                  doc='paths to assets, 0-terminated',
                  array_extent=('npaths',),
              ),
+             StructFieldDecl(
+                 name='signature',
+                 type=ValueType(name='uint64_t'),
+                 doc='also held by the mjSpec that compiled this model',
+             ),
          ),
      )),
     ('mjThreadPool',
@@ -6037,6 +6042,11 @@
                  type=ValueType(name='uintptr_t'),
                  doc='thread pool pointer',
              ),
+             StructFieldDecl(
+                 name='signature',
+                 type=ValueType(name='uint64_t'),
+                 doc='also held by the mjSpec that compiled the model',
+             ),
          ),
      )),
     ('mjvPerturb',
@@ -9044,6 +9054,11 @@
                  type=ValueType(name='mjtObj'),
                  doc='element type',
              ),
+             StructFieldDecl(
+                 name='signature',
+                 type=ValueType(name='uint64_t'),
+                 doc='compilation signature',
+             ),
          ),
      )),
     ('mjsCompiler',

python/mujoco/specs_test.py

@@ -1153,6 +1153,13 @@ def test_bind(self):
         AttributeError, "object has no attribute 'invalid'"
     ):
       print(mj_model.bind(joints).invalid)
+    invalid_spec = mujoco.MjSpec()
+    invalid_spec.worldbody.add_body(name='main')
+    with self.assertRaisesRegex(
+        ValueError,
+        'The mjSpec does not match mjModel. Please recompile the mjSpec.',
+    ):
+      print(mj_model.bind(invalid_spec.body('main')))
 
   def test_incorrect_hfield_size(self):
     nrow = 300

python/mujoco/structs.cc

@@ -1710,6 +1710,10 @@ This is useful for example when the MJB is not available as a file on disk.)"));
         // Return the full bytes array of concatenated paths
         return m.paths_bytes;
       });
+  mjModel.def_property_readonly(
+      "signature", [](const MjModelWrapper& m) -> const uint64_t& {
+        return m.get()->signature;
+      });
 
 #define XGROUP(MjModelGroupedViews, field, nfield, FIELD_XMACROS)             \
   mjModel.def(                                                                \
@@ -1730,7 +1734,13 @@ This is useful for example when the MJB is not available as a file on disk.)"));
   mjModel.def(                                                        \
       "bind_scalar",                                                  \
       [](MjModelWrapper& m, spectype& spec) -> auto& {                \
-        return m.indexer().field##_by_name(mjs_getString(spec.name)); \
+        if (mjs_getSpec(spec.element)->element->signature !=          \
+            m.get()->signature) {                                     \
+          throw py::value_error(                                      \
+              "The mjSpec does not match mjModel. Please recompile "  \
+              "the mjSpec.");                                         \
+        }                                                             \
+        return m.indexer().field(mjs_getId(spec.element)); \
       },                                                              \
       py::return_value_policy::reference_internal,                    \
       py::arg_v("spec", py::none()));
@@ -2018,6 +2028,10 @@ This is useful for example when the MJB is not available as a file on disk.)"));
         std::istringstream input(b, std::ios::in | std::ios::binary);
         return MjDataWrapper::Deserialize(input);
       }));
+  mjData.def_property_readonly(
+      "signature", [](const MjDataWrapper& d) -> uint64_t {
+        return d.get()->signature;
+      });
 
 #define X(type, var)                                             \
   mjData.def_property(                                           \
@@ -2076,7 +2090,13 @@ This is useful for example when the MJB is not available as a file on disk.)"));
   mjData.def(                                                         \
       "bind_scalar",                                                  \
       [](MjDataWrapper& d, spectype& spec) -> auto& {                 \
-        return d.indexer().field##_by_name(mjs_getString(spec.name)); \
+        if (mjs_getSpec(spec.element)->element->signature !=          \
+            d.get()->signature) {                                     \
+          throw py::value_error(                                      \
+                 "The mjSpec does not match mjData. Please recompile "\
+                 "the mjSpec.");                                      \
+        }                                                             \
+        return d.indexer().field(mjs_getId(spec.element));            \
       },                                                              \
       py::return_value_policy::reference_internal,                    \
       py::arg_v("spec", py::none()));

src/engine/engine_io.c

@@ -2014,6 +2014,9 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
       }
     }
   }
+
+  // copy signature from model
+  d->signature = m->signature;
 }
 
 

src/user/user_flexcomp.cc

@@ -403,12 +403,14 @@ 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,6 +201,7 @@ 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";
@@ -213,6 +214,9 @@ mjCModel::mjCModel() {
 
   // the source spec is the model itself, overwritten in the copy constructor
   source_spec_ = &spec;
+
+  // set the signature
+  spec.element->signature = 0;
 }
 
 
@@ -289,6 +293,7 @@ void mjCModel::CopyList(std::vector<T*>& dest,
     // copy the element from the other model to this model
     if (deepcopy_) {
       source[i]->ForgetKeyframes();
+      candidate->uid = GetUid();
     } else {
       candidate->AddRef();
     }
@@ -499,6 +504,9 @@ mjCModel& mjCModel::operator+=(const mjCModel& other) {
     nq = nv = na = nu = nmocap = 0;
   }
 
+  // update signature before we reset the tree lists
+  spec.element->signature = Signature();
+
   PointToLocal();
   return *this;
 }
@@ -633,6 +641,9 @@ mjCModel& mjCModel::operator-=(const mjCBody& subtree) {
   RemoveFromList(sensors_, oldmodel);
   RemovePlugins();
 
+  // update signature before we reset the tree lists
+  spec.element->signature = Signature();
+
   return *this;
 }
 
@@ -803,6 +814,9 @@ void mjCModel::DeleteElement(mjsElement* el) {
       break;
   }
 
+  // update signature before we reset the tree lists
+  spec.element->signature = Signature();
+
   ResetTreeLists();  // in case of a nested delete
   MakeTreeLists();
   ProcessLists(/*checkrepeat=*/false);
@@ -1019,7 +1033,9 @@ 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;
 }
 
@@ -1030,7 +1046,9 @@ 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;
 }
 
@@ -4572,6 +4590,28 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
     mju::strcpy_arr(errInfo.message, warningtext);
     errInfo.warning = true;
   }
+
+  // save signature
+  m->signature = Signature();
+}
+
+
+
+uint64_t mjCModel::Signature() {
+  std::string uid_str;
+  for (int i = 0; i < mjNOBJECT; ++i) {
+    if (i == mjOBJ_XBODY || i == mjOBJ_UNKNOWN || i == mjOBJ_DOF) {
+      continue;
+    }
+    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);
 }
 
 

src/user/user_model.h

@@ -16,6 +16,7 @@
 #define MUJOCO_SRC_USER_USER_MODEL_H_
 
 #include <array>
+#include <cstdint>
 #include <functional>
 #include <map>
 #include <string>
@@ -324,6 +325,9 @@ 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_;
@@ -440,11 +444,14 @@ class mjCModel : public mjCModel_, private mjSpec {
   void MarkPluginInstance(std::unordered_map<std::string, bool>& instances,
                           const std::vector<T*>& list);
 
+  // generate a signature for the model
+  uint64_t Signature();
 
   mjListKeyMap ids;   // map from object names to ids
   mjCError errInfo;   // last error info
   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
 };
 #endif  // MUJOCO_SRC_USER_USER_MODEL_H_