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

PiperOrigin-RevId: 740378879
Change-Id: If14b326942529494f172e7aedcae30195798b458

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 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  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 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  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 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  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 @@ class SpecsTest(absltest.TestCase):
         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()));