Added functionality for user payload data in mjSpec.

PiperOrigin-RevId: 738343134
Change-Id: I6217362ced9993afcd5a2914c78718ffe3bf753a

doc/APIreference/functions.rst

@@ -4458,6 +4458,33 @@ Resolve alternative orientations to quat, return error if any.
 
 Transform body into a frame.
 
+.. _mjs_setUserValue:
+
+`mjs_setUserValue <#mjs_setUserValue>`__
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mjs_setUserValue
+
+Set user payload, overriding the existing value for the specified key if present.
+
+.. _mjs_getUserValue:
+
+`mjs_getUserValue <#mjs_getUserValue>`__
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mjs_getUserValue
+
+Return user payload or NULL if none found.
+
+.. _mjs_deleteUserValue:
+
+`mjs_deleteUserValue <#mjs_deleteUserValue>`__
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mjs_deleteUserValue
+
+Delete user payload.
+
 .. _ElementInitialization:
 
 Element initialization

doc/includes/references.h

@@ -3677,6 +3677,9 @@ void mjs_setFrame(mjsElement* dest, mjsFrame* frame);
 const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* sequence,
                                    const mjsOrientation* orientation);
 mjsFrame* mjs_bodyToFrame(mjsBody** body);
+void mjs_setUserValue(mjsElement* element, const char* key, const void* data);
+const void* mjs_getUserValue(mjsElement* element, const char* key);
+void mjs_deleteUserValue(mjsElement* element, const char* key);
 void mjs_defaultSpec(mjSpec* spec);
 void mjs_defaultOrientation(mjsOrientation* orient);
 void mjs_defaultBody(mjsBody* body);

include/mujoco/mujoco.h

@@ -1648,6 +1648,15 @@ MJAPI const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const c
 // Transform body into a frame.
 MJAPI mjsFrame* mjs_bodyToFrame(mjsBody** body);
 
+// Set user payload, overriding the existing value for the specified key if present.
+MJAPI void mjs_setUserValue(mjsElement* element, const char* key, const void* data);
+
+// Return user payload or NULL if none found.
+MJAPI const void* mjs_getUserValue(mjsElement* element, const char* key);
+
+// Delete user payload.
+MJAPI void mjs_deleteUserValue(mjsElement* element, const char* key);
+
 //---------------------------------- Element initialization  ---------------------------------------
 
 // Default spec attributes.

python/mujoco/introspect/functions.py

@@ -10471,6 +10471,74 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Transform body into a frame.',
      )),
+    ('mjs_setUserValue',
+     FunctionDecl(
+         name='mjs_setUserValue',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='element',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjsElement'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='key',
+                 type=PointerType(
+                     inner_type=ValueType(name='char', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='data',
+                 type=PointerType(
+                     inner_type=ValueType(name='void', is_const=True),
+                 ),
+             ),
+         ),
+         doc='Set user payload, overriding the existing value for the specified key if present.',  # pylint: disable=line-too-long
+     )),
+    ('mjs_getUserValue',
+     FunctionDecl(
+         name='mjs_getUserValue',
+         return_type=PointerType(
+             inner_type=ValueType(name='void', is_const=True),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='element',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjsElement'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='key',
+                 type=PointerType(
+                     inner_type=ValueType(name='char', is_const=True),
+                 ),
+             ),
+         ),
+         doc='Return user payload or NULL if none found.',
+     )),
+    ('mjs_deleteUserValue',
+     FunctionDecl(
+         name='mjs_deleteUserValue',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='element',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjsElement'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='key',
+                 type=PointerType(
+                     inner_type=ValueType(name='char', is_const=True),
+                 ),
+             ),
+         ),
+         doc='Delete user payload.',
+     )),
     ('mjs_defaultSpec',
      FunctionDecl(
          name='mjs_defaultSpec',

src/user/user_api.cc

@@ -769,6 +769,30 @@ mjsFrame* mjs_bodyToFrame(mjsBody** body) {
 
 
 
+// set user payload
+void mjs_setUserValue(mjsElement* element, const char* key, const void* data) {
+  mjCBase* baseC = static_cast<mjCBase*>(element);
+  baseC->SetUserValue(key, data);
+}
+
+
+
+// return user payload or NULL if none found
+const void* mjs_getUserValue(mjsElement* element, const char* key) {
+  mjCBase* baseC = static_cast<mjCBase*>(element);
+  return baseC->GetUserValue(key);
+}
+
+
+
+// delete user payload
+void mjs_deleteUserValue(mjsElement* element, const char* key) {
+  mjCBase* baseC = static_cast<mjCBase*>(element);
+  baseC->DeleteUserValue(key);
+}
+
+
+
 // get id
 int mjs_getId(mjsElement* element) {
   if (!element) {

src/user/user_api.h

@@ -378,6 +378,15 @@ MJAPI const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const c
 // Transform body into a frame.
 MJAPI mjsFrame* mjs_bodyToFrame(mjsBody** body);
 
+// Set user payload.
+MJAPI void mjs_setUserValue(mjsElement* element, const char* key, const void* data);
+
+// Return user payload or NULL if none found.
+MJAPI const void* mjs_getUserValue(mjsElement* element, const char* key);
+
+// Delete user payload.
+MJAPI void mjs_deleteUserValue(mjsElement* element, const char* key);
+
 
 //---------------------------------- Initialization  -----------------------------------------------
 

src/user/user_objects.cc

@@ -766,6 +766,22 @@ void mjCBase::SetFrame(mjCFrame* _frame) {
 }
 
 
+void mjCBase::SetUserValue(std::string_view key, const void* data) {
+  user_payload_[std::string(key)] = data;
+}
+
+
+const void* mjCBase::GetUserValue(std::string_view key) {
+  auto found = user_payload_.find(std::string(key));
+  return found != user_payload_.end() ? found->second : nullptr;
+}
+
+
+void mjCBase::DeleteUserValue(std::string_view key) {
+  user_payload_.erase(std::string(key));
+}
+
+
 //------------------ class mjCBody implementation --------------------------------------------------
 
 // constructor

src/user/user_objects.h

@@ -23,11 +23,14 @@
 #include <map>
 #include <string>
 #include <string_view>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjplugin.h>
 #include <mujoco/mjspec.h>
-#include <mujoco/mujoco.h>
+#include <mujoco/mjtnum.h>
 #include "user/user_cache.h"
 #include "user/user_util.h"
 #include <tiny_obj_loader.h>
@@ -274,12 +277,20 @@ class mjCBase : public mjCBase_ {
     }
   }
 
+  // Set and get user payload
+  void SetUserValue(std::string_view key, const void* data);
+  const void* GetUserValue(std::string_view key);
+  void DeleteUserValue(std::string_view key);
+
  protected:
   mjCBase();                                 // constructor
   mjCBase(const mjCBase& other);             // copy constructor
 
   // reference count for allowing deleting an attached object
   int refcount = 1;
+
+  // user payload
+  std::unordered_map<std::string, const void*> user_payload_;
 };
 
 
@@ -781,7 +792,7 @@ class mjCLight : public mjCLight_, private mjsLight {
 
 class mjCFlex_ : public mjCBase {
  protected:
-  int nvert;                              // number of verices
+  int nvert;                              // number of vertices
   int nnode;                              // number of nodes
   int nedge;                              // number of edges
   int nelem;                              // number of elements

test/user/user_api_test.cc

@@ -2585,5 +2585,20 @@ TEST_F(MujocoTest, SetFrameReverseOrder) {
   mj_deleteSpec(copy);
 }
 
+TEST_F(MujocoTest, UserValue) {
+  mjSpec* spec = mj_makeSpec();
+  EXPECT_THAT(spec, NotNull());
+  mjsBody* body = mjs_addBody(mjs_findBody(spec, "world"), nullptr);
+  EXPECT_THAT(body, NotNull());
+  std::string data = "data";
+  mjs_setUserValue(body->element, "key", data.data());
+  EXPECT_THAT(mjs_getUserValue(body->element, "invalid_key"), IsNull());
+  const void* payload = mjs_getUserValue(body->element, "key");
+  EXPECT_STREQ(static_cast<const char*>(payload), data.c_str());
+  mjs_deleteUserValue(body->element, "key");
+  EXPECT_THAT(mjs_getUserValue(body->element, "key"), IsNull());
+  mj_deleteSpec(spec);
+}
+
 }  // namespace
 }  // namespace mujoco