Add control plugin API.

PiperOrigin-RevId: 493311715
Change-Id: Iccd99fb26004f39832818e378f01f9149a0e3e3b

doc/includes/references.h

@@ -1148,18 +1148,19 @@ struct mjModel_ {
   char*     names;                // names of all objects, 0-terminated       (nnames x 1)
 };
 typedef struct mjModel_ mjModel;
-typedef enum mjtPluginTypeBit_ {
+typedef enum mjtPluginCapabilityBit_ {
   mjPLUGIN_ACTUATOR = 1<<0,
   mjPLUGIN_SENSOR   = 1<<1,
   mjPLUGIN_PASSIVE  = 1<<2,
-} mjtPluginTypeBit;
+  mjPLUGIN_CONTROL  = 1<<3,
+} mjtPluginCapabilityBit;
 struct mjpPlugin_ {
   const char* name;    // globally unique name identifying the plugin
 
   int nattribute;                 // number of configuration attributes
   const char* const* attributes;  // name of configuration attributes
 
-  int type;            // bitfield of mjtPluginTypeBits specifying the plugin type
+  int capabilities;    // bitfield of mjtPluginCapabilityBit specifying optional plugin capabilities
   int needstage;       // an mjtStage enum value specifying the sensor computation stage
 
   // number of mjtNums needed to store the state of a plugin instance (required)
@@ -1181,10 +1182,13 @@ struct mjpPlugin_ {
   void (*reset)(const mjModel* m, mjData* d, int instance);
 
   // called when the plugin needs to update its outputs (required)
-  void (*compute)(const mjModel* m, mjData* d, int instance, int type);
+  void (*compute)(const mjModel* m, mjData* d, int instance, int capability_bit);
 
   // called when time integration occurs (optional)
   void (*advance)(const mjModel* m, mjData* d, int instance);
+
+  // called by mjv_updateScene (optional)
+  void (*visualize)(const mjModel*m, mjData* d, mjvScene* scn, int instance);
 };
 typedef struct mjpPlugin_ mjpPlugin;
 typedef enum mjtGridPos_ {        // grid position for overlay

include/mujoco/mjplugin.h

@@ -20,11 +20,12 @@
 #include <mujoco/mjvisualize.h>
 
 
-typedef enum mjtPluginTypeBit_ {
+typedef enum mjtPluginCapabilityBit_ {
   mjPLUGIN_ACTUATOR = 1<<0,
   mjPLUGIN_SENSOR   = 1<<1,
   mjPLUGIN_PASSIVE  = 1<<2,
-} mjtPluginTypeBit;
+  mjPLUGIN_CONTROL  = 1<<3,
+} mjtPluginCapabilityBit;
 
 struct mjpPlugin_ {
   const char* name;    // globally unique name identifying the plugin
@@ -32,7 +33,7 @@ struct mjpPlugin_ {
   int nattribute;                 // number of configuration attributes
   const char* const* attributes;  // name of configuration attributes
 
-  int type;            // bitfield of mjtPluginTypeBits specifying the plugin type
+  int capabilities;    // bitfield of mjtPluginCapabilityBit specifying optional plugin capabilities
   int needstage;       // an mjtStage enum value specifying the sensor computation stage
 
   // number of mjtNums needed to store the state of a plugin instance (required)
@@ -54,7 +55,7 @@ struct mjpPlugin_ {
   void (*reset)(const mjModel* m, mjData* d, int instance);
 
   // called when the plugin needs to update its outputs (required)
-  void (*compute)(const mjModel* m, mjData* d, int instance, int type);
+  void (*compute)(const mjModel* m, mjData* d, int instance, int capability_bit);
 
   // called when time integration occurs (optional)
   void (*advance)(const mjModel* m, mjData* d, int instance);

introspect/enums.py

@@ -550,14 +550,15 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjSTEREO_SIDEBYSIDE', 2),
          ]),
      )),
-    ('mjtPluginTypeBit',
+    ('mjtPluginCapabilityBit',
      EnumDecl(
-         name='mjtPluginTypeBit',
+         name='mjtPluginCapabilityBit',
-         declname='enum mjtPluginTypeBit_',
+         declname='enum mjtPluginCapabilityBit_',
          values=dict([
              ('mjPLUGIN_ACTUATOR', 1),
              ('mjPLUGIN_SENSOR', 2),
              ('mjPLUGIN_PASSIVE', 4),
+             ('mjPLUGIN_CONTROL', 8),
          ]),
      )),
     ('mjtGridPos',

plugin/elasticity/cable.cc

@@ -277,7 +277,7 @@ void Cable::RegisterPlugin() {
   mjp_defaultPlugin(&plugin);
 
   plugin.name = "mujoco.elasticity.cable";
-  plugin.type |= mjPLUGIN_PASSIVE;
+  plugin.capabilities |= mjPLUGIN_PASSIVE;
 
   const char* attributes[] = {"twist", "bend", "flat", "vmax"};
   plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
@@ -297,10 +297,11 @@ return 0;
     delete reinterpret_cast<Cable*>(d->plugin_data[instance]);
     d->plugin_data[instance] = 0;
   };
-  plugin.compute = +[](const mjModel* m, mjData* d, int instance, int type) {
+  plugin.compute =
-    auto* elasticity = reinterpret_cast<Cable*>(d->plugin_data[instance]);
+      +[](const mjModel* m, mjData* d, int instance, int capability_bit) {
-    elasticity->Compute(m, d, instance);
+        auto* elasticity = reinterpret_cast<Cable*>(d->plugin_data[instance]);
-  };
+        elasticity->Compute(m, d, instance);
+      };
   plugin.visualize =
       +[](const mjModel* m, mjData* d, mjvScene* scn, int instance) {
         auto* elasticity = reinterpret_cast<Cable*>(d->plugin_data[instance]);

plugin/elasticity/cable.h

@@ -30,8 +30,7 @@ class Cable {
  public:
   // Creates a new Cable instance (allocated with `new`) or
   // returns null on failure.
-  static std::optional<Cable> Create(const mjModel* m, mjData* d,
+  static std::optional<Cable> Create(const mjModel* m, mjData* d, int instance);
-                                          int instance);
   Cable(Cable&&) = default;
   ~Cable() = default;
 

plugin/elasticity/solid.cc

@@ -347,7 +347,7 @@ void Solid::RegisterPlugin() {
   mjp_defaultPlugin(&plugin);
 
   plugin.name = "mujoco.elasticity.solid";
-  plugin.type |= mjPLUGIN_PASSIVE;
+  plugin.capabilities |= mjPLUGIN_PASSIVE;
 
   const char* attributes[] = {"nx", "ny", "nz", "young", "poisson", "damping"};
   plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
@@ -367,10 +367,11 @@ void Solid::RegisterPlugin() {
     delete reinterpret_cast<Solid*>(d->plugin_data[instance]);
     d->plugin_data[instance] = 0;
   };
-  plugin.compute = +[](const mjModel* m, mjData* d, int instance, int type) {
+  plugin.compute =
-    auto* elasticity = reinterpret_cast<Solid*>(d->plugin_data[instance]);
+      +[](const mjModel* m, mjData* d, int instance, int capability_bit) {
-    elasticity->Compute(m, d, instance);
+        auto* elasticity = reinterpret_cast<Solid*>(d->plugin_data[instance]);
-  };
+        elasticity->Compute(m, d, instance);
+      };
 
   mjp_registerPlugin(&plugin);
 }

plugin/elasticity/solid.h

@@ -33,8 +33,7 @@ struct Stencil {
 class Solid {
  public:
   // Returns a new Solid instance or nullopt on failure.
-  static std::optional<Solid> Create(const mjModel* m, mjData* d,
+  static std::optional<Solid> Create(const mjModel* m, mjData* d, int instance);
-                                     int instance);
   Solid(Solid&&) = default;
 
   Solid& operator=(Solid&& other) = default;

src/engine/engine_core_smooth.c

@@ -1454,7 +1454,7 @@ void mj_passive(const mjModel* m, mjData* d) {
       if (!plugin) {
         mju_error_i("invalid plugin slot: %d", slot);
       }
-      if (plugin->type & mjPLUGIN_PASSIVE) {
+      if (plugin->capabilities & mjPLUGIN_PASSIVE) {
         if (!plugin->compute) {
           mju_error_i("`compute` is a null function pointer for plugin at slot %d", slot);
         }

src/engine/engine_forward.c

@@ -279,7 +279,7 @@ void mj_fwdActuation(const mjModel* m, mjData* d) {
       if (!plugin) {
         mju_error_i("invalid plugin slot: %d", slot);
       }
-      if (plugin->type & mjPLUGIN_ACTUATOR) {
+      if (plugin->capabilities & mjPLUGIN_ACTUATOR) {
         if (!plugin->compute) {
           mju_error_i("`compute` is a null function pointer for plugin at slot %d", slot);
         }
@@ -774,9 +774,31 @@ void mj_forwardSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor)
   }
 
   // acceleration-dependent
-  if (mjcb_control && !mjDISABLED(mjDSBL_ACTUATION)) {
+  if (!mjDISABLED(mjDSBL_ACTUATION)) {
-    mjcb_control(m, d);
+    // call legacy control callback if specified
-  }
+    if (mjcb_control) {
+      mjcb_control(m, d);
+    }
+
+    // handle control plugins
+    if (m->nplugin) {
+      const int nslot = mjp_pluginCount();
+      for (int i=0; i<m->nplugin; i++) {
+        const int slot = m->plugin[i];
+        const mjpPlugin* plugin = mjp_getPluginAtSlotUnsafe(slot, nslot);
+        if (!plugin) {
+          mju_error_i("invalid plugin slot: %d", slot);
+        }
+        if (plugin->capabilities & mjPLUGIN_CONTROL) {
+          if (!plugin->compute) {
+            mju_error_i("`compute` is a null function pointer for plugin at slot %d", slot);
+          }
+          plugin->compute(m, d, i, mjPLUGIN_CONTROL);
+        }
+      }
+    }
+}
+
   mj_fwdActuation(m, d);
   mj_fwdAcceleration(m, d);
   mj_fwdConstraint(m, d);

src/engine/engine_sensor.c

@@ -357,7 +357,7 @@ void mj_sensorPos(const mjModel* m, mjData* d) {
       if (!plugin) {
         mju_error_i("invalid plugin slot: %d", slot);
       }
-      if ((plugin->type & mjPLUGIN_SENSOR) &&
+      if ((plugin->capabilities & mjPLUGIN_SENSOR) &&
           (plugin->needstage==mjSTAGE_POS || plugin->needstage==mjSTAGE_NONE)) {
         if (!plugin->compute) {
           mju_error_i("`compute` is a null function pointer for plugin at slot %d", slot);
@@ -538,7 +538,7 @@ void mj_sensorVel(const mjModel* m, mjData* d) {
       if (!plugin) {
         mju_error_i("invalid plugin slot: %d", slot);
       }
-      if ((plugin->type & mjPLUGIN_SENSOR) && plugin->needstage==mjSTAGE_VEL) {
+      if ((plugin->capabilities & mjPLUGIN_SENSOR) && plugin->needstage==mjSTAGE_VEL) {
         if (!plugin->compute) {
           mju_error_i("`compute` is null for plugin at slot %d", slot);
         }
@@ -747,7 +747,7 @@ void mj_sensorAcc(const mjModel* m, mjData* d) {
       if (!plugin) {
         mju_error_i("invalid plugin slot: %d", slot);
       }
-      if ((plugin->type & mjPLUGIN_SENSOR) && plugin->needstage==mjSTAGE_ACC) {
+      if ((plugin->capabilities & mjPLUGIN_SENSOR) && plugin->needstage==mjSTAGE_ACC) {
         if (!plugin->compute) {
           mju_error_i("`compute` is null for plugin at slot %d", slot);
         }

src/user/user_model.cc

@@ -2633,7 +2633,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
       m->plugin_stateadr[i] = stateadr;
       m->plugin_statenum[i] = nstate;
       stateadr += nstate;
-      if (plugin->type & mjPLUGIN_SENSOR) {
+      if (plugin->capabilities & mjPLUGIN_SENSOR) {
         for (int sensor_id : plugin_to_sensors[i]) {
           if (!plugin->nsensordata) {
             mju_error_i("`reset` is null for plugin at slot %d", m->plugin[i]);

src/user/user_objects.cc

@@ -819,7 +819,7 @@ void mjCBody::Compile(void) {
 
     model->ResolvePlugin(this, plugin_name, plugin_instance_name, &plugin_instance);
     const mjpPlugin* plugin = mjp_getPluginAtSlot(plugin_instance->plugin_slot);
-    if (!(plugin->type & mjPLUGIN_PASSIVE)) {
+    if (!(plugin->capabilities & mjPLUGIN_PASSIVE)) {
       throw mjCError(this, "plugin '%s' does not support passive forces", plugin->name);
     }
   }
@@ -3618,7 +3618,7 @@ void mjCActuator::Compile(void) {
 
     model->ResolvePlugin(this, plugin_name, plugin_instance_name, &plugin_instance);
     const mjpPlugin* plugin = mjp_getPluginAtSlot(plugin_instance->plugin_slot);
-    if (!(plugin->type & mjPLUGIN_ACTUATOR)) {
+    if (!(plugin->capabilities & mjPLUGIN_ACTUATOR)) {
       throw mjCError(this, "plugin '%s' does not support actuators", plugin->name);
     }
   }
@@ -4006,7 +4006,7 @@ void mjCSensor::Compile(void) {
     {
       model->ResolvePlugin(this, plugin_name, plugin_instance_name, &plugin_instance);
       const mjpPlugin* plugin = mjp_getPluginAtSlot(plugin_instance->plugin_slot);
-      if (!(plugin->type & mjPLUGIN_SENSOR)) {
+      if (!(plugin->capabilities & mjPLUGIN_SENSOR)) {
         throw mjCError(this, "plugin '%s' does not support sensors", plugin->name);
       }
       needstage = static_cast<mjtStage>(plugin->needstage);

test/engine/engine_plugin_test.cc

@@ -206,7 +206,7 @@ int RegisterSensorPlugin() {
   plugin.nattribute = sizeof(attributes) / sizeof(*attributes);
   plugin.attributes = attributes;
 
-  plugin.type |= mjPLUGIN_SENSOR;
+  plugin.capabilities |= mjPLUGIN_SENSOR;
 
   plugin.nstate = +[](const mjModel* m, int instance) { return 3; };
   plugin.nsensordata =
@@ -250,7 +250,7 @@ int RegisterActuatorPlugin() {
   plugin.nattribute = sizeof(attributes) / sizeof(*attributes);
   plugin.attributes = attributes;
 
-  plugin.type |= mjPLUGIN_ACTUATOR;
+  plugin.capabilities |= mjPLUGIN_ACTUATOR;
 
   plugin.nstate = +[](const mjModel* m, int instance) { return 3; };
 
@@ -292,7 +292,7 @@ int RegisterPassivePlugin() {
   plugin.nattribute = sizeof(attributes) / sizeof(*attributes);
   plugin.attributes = attributes;
 
-  plugin.type |= mjPLUGIN_PASSIVE;
+  plugin.capabilities |= mjPLUGIN_PASSIVE;
 
   plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
 

unity/Runtime/Bindings/MujocoBindings.cs

@@ -361,10 +361,11 @@ public enum mjtLRMode : int{
   mjLRMODE_MUSCLEUSER = 2,
   mjLRMODE_ALL = 3,
 }
-public enum mjtPluginTypeBit : int{
+public enum mjtPluginCapabilityBit : int{
   mjPLUGIN_ACTUATOR = 1,
   mjPLUGIN_SENSOR = 2,
   mjPLUGIN_PASSIVE = 4,
+  mjPLUGIN_CONTROL = 8,
 }
 public enum mjtGridPos : int{
   mjGRID_TOPLEFT = 0,