Add updateScene callback for plugins.

PiperOrigin-RevId: 492497689
Change-Id: I4a15cc5fa39dc7c9b85d132727f82033d4aec2c8

doc/changelog.rst

@@ -67,6 +67,8 @@ General
   ``dlopen`` on a POSIX system, and to ``LoadLibraryA`` on Windows. The second function scans a specified directory for
   all dynamic libraries file and loads each library found. Dynamic libraries opened by these functions are assumed to
   register one or more MuJoCo plugins on load.
+- Add an optional ``visualize`` callback to plugins, which is called during ``mjv_updateScene``. This callback allows
+  custom plugin visualizations. Enable stree visualization for the Cable plugin as an example.
 - Sensors of type :ref:`user<sensor-user>` no longer require :at:`objtype` and :at:`objname`. If unspecified, the
   objtype will be :ref:`mjOBJ_UNKNOWN<mjtObj>`. ``user`` sensors :at:`datatype` default is now :at-val:`"real"`.
 - Add support for capsules in URDF import.

include/mujoco/mjplugin.h

@@ -17,6 +17,8 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
+#include <mujoco/mjvisualize.h>
+
 
 typedef enum mjtPluginTypeBit_ {
   mjPLUGIN_ACTUATOR = 1<<0,
@@ -56,6 +58,9 @@ struct mjpPlugin_ {
 
   // 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;
 

introspect/enums.py

@@ -408,107 +408,6 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjNTIMER', 13),
          ]),
      )),
-    ('mjtPluginTypeBit',
-     EnumDecl(
-         name='mjtPluginTypeBit',
-         declname='enum mjtPluginTypeBit_',
-         values=dict([
-             ('mjPLUGIN_ACTUATOR', 1),
-             ('mjPLUGIN_SENSOR', 2),
-             ('mjPLUGIN_PASSIVE', 4),
-         ]),
-     )),
-    ('mjtGridPos',
-     EnumDecl(
-         name='mjtGridPos',
-         declname='enum mjtGridPos_',
-         values=dict([
-             ('mjGRID_TOPLEFT', 0),
-             ('mjGRID_TOPRIGHT', 1),
-             ('mjGRID_BOTTOMLEFT', 2),
-             ('mjGRID_BOTTOMRIGHT', 3),
-         ]),
-     )),
-    ('mjtFramebuffer',
-     EnumDecl(
-         name='mjtFramebuffer',
-         declname='enum mjtFramebuffer_',
-         values=dict([
-             ('mjFB_WINDOW', 0),
-             ('mjFB_OFFSCREEN', 1),
-         ]),
-     )),
-    ('mjtFontScale',
-     EnumDecl(
-         name='mjtFontScale',
-         declname='enum mjtFontScale_',
-         values=dict([
-             ('mjFONTSCALE_50', 50),
-             ('mjFONTSCALE_100', 100),
-             ('mjFONTSCALE_150', 150),
-             ('mjFONTSCALE_200', 200),
-             ('mjFONTSCALE_250', 250),
-             ('mjFONTSCALE_300', 300),
-         ]),
-     )),
-    ('mjtFont',
-     EnumDecl(
-         name='mjtFont',
-         declname='enum mjtFont_',
-         values=dict([
-             ('mjFONT_NORMAL', 0),
-             ('mjFONT_SHADOW', 1),
-             ('mjFONT_BIG', 2),
-         ]),
-     )),
-    ('mjtButton',
-     EnumDecl(
-         name='mjtButton',
-         declname='enum mjtButton_',
-         values=dict([
-             ('mjBUTTON_NONE', 0),
-             ('mjBUTTON_LEFT', 1),
-             ('mjBUTTON_RIGHT', 2),
-             ('mjBUTTON_MIDDLE', 3),
-         ]),
-     )),
-    ('mjtEvent',
-     EnumDecl(
-         name='mjtEvent',
-         declname='enum mjtEvent_',
-         values=dict([
-             ('mjEVENT_NONE', 0),
-             ('mjEVENT_MOVE', 1),
-             ('mjEVENT_PRESS', 2),
-             ('mjEVENT_RELEASE', 3),
-             ('mjEVENT_SCROLL', 4),
-             ('mjEVENT_KEY', 5),
-             ('mjEVENT_RESIZE', 6),
-         ]),
-     )),
-    ('mjtItem',
-     EnumDecl(
-         name='mjtItem',
-         declname='enum mjtItem_',
-         values=dict([
-             ('mjITEM_END', -2),
-             ('mjITEM_SECTION', -1),
-             ('mjITEM_SEPARATOR', 0),
-             ('mjITEM_STATIC', 1),
-             ('mjITEM_BUTTON', 2),
-             ('mjITEM_CHECKINT', 3),
-             ('mjITEM_CHECKBYTE', 4),
-             ('mjITEM_RADIO', 5),
-             ('mjITEM_RADIOLINE', 6),
-             ('mjITEM_SELECT', 7),
-             ('mjITEM_SLIDERINT', 8),
-             ('mjITEM_SLIDERNUM', 9),
-             ('mjITEM_EDITINT', 10),
-             ('mjITEM_EDITNUM', 11),
-             ('mjITEM_EDITTXT', 12),
-             ('mjNITEM', 13),
-         ]),
-     )),
     ('mjtCatBit',
      EnumDecl(
          name='mjtCatBit',
@@ -651,4 +550,105 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjSTEREO_SIDEBYSIDE', 2),
          ]),
      )),
+    ('mjtPluginTypeBit',
+     EnumDecl(
+         name='mjtPluginTypeBit',
+         declname='enum mjtPluginTypeBit_',
+         values=dict([
+             ('mjPLUGIN_ACTUATOR', 1),
+             ('mjPLUGIN_SENSOR', 2),
+             ('mjPLUGIN_PASSIVE', 4),
+         ]),
+     )),
+    ('mjtGridPos',
+     EnumDecl(
+         name='mjtGridPos',
+         declname='enum mjtGridPos_',
+         values=dict([
+             ('mjGRID_TOPLEFT', 0),
+             ('mjGRID_TOPRIGHT', 1),
+             ('mjGRID_BOTTOMLEFT', 2),
+             ('mjGRID_BOTTOMRIGHT', 3),
+         ]),
+     )),
+    ('mjtFramebuffer',
+     EnumDecl(
+         name='mjtFramebuffer',
+         declname='enum mjtFramebuffer_',
+         values=dict([
+             ('mjFB_WINDOW', 0),
+             ('mjFB_OFFSCREEN', 1),
+         ]),
+     )),
+    ('mjtFontScale',
+     EnumDecl(
+         name='mjtFontScale',
+         declname='enum mjtFontScale_',
+         values=dict([
+             ('mjFONTSCALE_50', 50),
+             ('mjFONTSCALE_100', 100),
+             ('mjFONTSCALE_150', 150),
+             ('mjFONTSCALE_200', 200),
+             ('mjFONTSCALE_250', 250),
+             ('mjFONTSCALE_300', 300),
+         ]),
+     )),
+    ('mjtFont',
+     EnumDecl(
+         name='mjtFont',
+         declname='enum mjtFont_',
+         values=dict([
+             ('mjFONT_NORMAL', 0),
+             ('mjFONT_SHADOW', 1),
+             ('mjFONT_BIG', 2),
+         ]),
+     )),
+    ('mjtButton',
+     EnumDecl(
+         name='mjtButton',
+         declname='enum mjtButton_',
+         values=dict([
+             ('mjBUTTON_NONE', 0),
+             ('mjBUTTON_LEFT', 1),
+             ('mjBUTTON_RIGHT', 2),
+             ('mjBUTTON_MIDDLE', 3),
+         ]),
+     )),
+    ('mjtEvent',
+     EnumDecl(
+         name='mjtEvent',
+         declname='enum mjtEvent_',
+         values=dict([
+             ('mjEVENT_NONE', 0),
+             ('mjEVENT_MOVE', 1),
+             ('mjEVENT_PRESS', 2),
+             ('mjEVENT_RELEASE', 3),
+             ('mjEVENT_SCROLL', 4),
+             ('mjEVENT_KEY', 5),
+             ('mjEVENT_RESIZE', 6),
+         ]),
+     )),
+    ('mjtItem',
+     EnumDecl(
+         name='mjtItem',
+         declname='enum mjtItem_',
+         values=dict([
+             ('mjITEM_END', -2),
+             ('mjITEM_SECTION', -1),
+             ('mjITEM_SEPARATOR', 0),
+             ('mjITEM_STATIC', 1),
+             ('mjITEM_BUTTON', 2),
+             ('mjITEM_CHECKINT', 3),
+             ('mjITEM_CHECKBYTE', 4),
+             ('mjITEM_RADIO', 5),
+             ('mjITEM_RADIOLINE', 6),
+             ('mjITEM_SELECT', 7),
+             ('mjITEM_SLIDERINT', 8),
+             ('mjITEM_SLIDERNUM', 9),
+             ('mjITEM_EDITINT', 10),
+             ('mjITEM_EDITNUM', 11),
+             ('mjITEM_EDITTXT', 12),
+             ('mjNITEM', 13),
+         ]),
+     )),
 ])

model/plugin/belt.xml

@@ -59,6 +59,7 @@
         <!--Units are in Pa (SI)-->
         <config key="twist" value="1e8"/>
         <config key="bend" value="1e7"/>
+        <config key="vmax" value="0.2"/>
       </plugin>
       <joint kind="main" damping="0.1"/>
       <geom type="box" size="0.03535533906 0.02 0.0015" rgba=".8 .2 .1 1" group="3"/>

model/plugin/cable.xml

@@ -35,6 +35,7 @@
         <!--Units are in Pa (SI)-->
         <config key="twist" value="1e7"/>
         <config key="bend" value="4e6"/>
+        <config key="vmax" value="0.05"/>
       </plugin>
       <joint kind="main" damping=".015"/>
       <geom type="capsule" size=".005" rgba=".8 .2 .1 1" condim="1"/>

model/plugin/coil.xml

@@ -70,6 +70,7 @@
         <!--Units are in Pa (SI)-->
         <config key="twist" value="5e8"/>
         <config key="bend" value="15e8"/>
+        <config key="vmax" value="2"/>
       </plugin>
       <joint kind="main" damping="0.12" armature="0.01"/>
       <geom type="capsule" size=".005" rgba=".8 .2 .1 1"/>
@@ -80,6 +81,7 @@
         <!--Units are in Pa (SI)-->
         <config key="twist" value="5e8"/>
         <config key="bend" value="15e8"/>
+        <config key="vmax" value="2"/>
       </plugin>
       <joint kind="main" damping="0.1" armature="0.01"/>
       <geom type="capsule" size=".005" rgba=".8 .2 .1 1"/>

plugin/elasticity/cable.cc

@@ -26,9 +26,36 @@
 namespace mujoco::plugin::elasticity {
 namespace {
 
+// Jet color palette
+void scalar2rgba(float rgba[4], mjtNum stress[3], mjtNum vmin, mjtNum vmax) {
+  // L2 norm of the stress
+  mjtNum v = mju_norm3(stress);
+  v = v < vmin ? vmin : v;
+  v = v > vmax ? vmax : v;
+  mjtNum dv = vmax - vmin;
+
+  if (v < (vmin + 0.25 * dv)) {
+    rgba[0] = 0;
+    rgba[1] = 4 * (v - vmin) / dv;
+    rgba[2] = 1;
+  } else if (v < (vmin + 0.5 * dv)) {
+    rgba[0] = 0;
+    rgba[1] = 1;
+    rgba[2] = 1 + 4 * (vmin + 0.25 * dv - v) / dv;
+  } else if (v < (vmin + 0.75 * dv)) {
+    rgba[0] = 4 * (v - vmin - 0.5 * dv) / dv;
+    rgba[1] = 1;
+    rgba[2] = 0;
+  } else {
+    rgba[0] = 1;
+    rgba[1] = 1 + 4 * (vmin + 0.75 * dv - v) / dv;
+    rgba[2] = 0;
+  }
+}
+
 // compute quaternion difference between two frames in joint coordinates
 void QuatDiff(mjtNum* quat, const mjtNum body_quat[4],
-              const mjtNum joint_quat[4], bool pullback) {
+              const mjtNum joint_quat[4], bool pullback = false) {
   if (pullback == 0) {
     // contribution in local coordinates
     mju_mulQuat(quat, body_quat, joint_quat);
@@ -40,7 +67,7 @@ void QuatDiff(mjtNum* quat, const mjtNum body_quat[4],
   }
 }
 
-// compute local force given material properties, orientation,
+// compute local stress given material properties, orientation,
 // and reference curvature
 //   inputs:
 //     stiffness   - material parameters
@@ -49,14 +76,15 @@ void QuatDiff(mjtNum* quat, const mjtNum body_quat[4],
 //     xquat       - cartesian orientation of the body (optional)
 //     scl         - scaling of the force
 //   outputs:
-//     qfrc        - local torque contribution
+//     stress      - local stress contribution
-void LocalForce(mjtNum qfrc[3], const mjtNum stiffness[4],
+void LocalStress(mjtNum stress[3],
+                const mjtNum stiffness[4],
                 const mjtNum quat[4], const mjtNum omega0[3],
-                const mjtNum xquat[4], mjtNum scl) {
+                bool pullback = false) {
-  mjtNum omega[3], lfrc[3];
+  mjtNum omega[3];
 
   // compute curvature
-  mju_quat2Vel(omega, quat, scl);
+  mju_quat2Vel(omega, quat, 1.0);
 
   // subtract omega0 in reference configuration
   mjtNum tmp[] = {
@@ -65,15 +93,15 @@ void LocalForce(mjtNum qfrc[3], const mjtNum stiffness[4],
       - stiffness[2]*(omega[2] - omega0[2]) / stiffness[3],
   };
 
-  // rotate into global frame
-  if (xquat) {
-    mju_rotVecQuat(lfrc, tmp, xquat);
-  } else {
-    mju_copy3(lfrc, tmp);
-  }
 
-  // add to total qfrc
+  // pull-back into the other body frame
-  mju_addToScl3(qfrc, lfrc, scl);
+  if (pullback) {
+    mjtNum invquat[4];
+    mju_negQuat(invquat, quat);
+    mju_rotVecQuat(stress, tmp, invquat);
+  } else {
+    mju_copy3(stress, tmp);
+  }
 }
 
 // reads numeric attributes
@@ -105,7 +133,7 @@ Cable::Cable(const mjModel* m, mjData* d, int instance) {
   std::string flat = mj_getPluginConfig(m, instance, "flat");
   mjtNum G = strtod(mj_getPluginConfig(m, instance, "twist"), nullptr);
   mjtNum E = strtod(mj_getPluginConfig(m, instance, "bend"), nullptr);
-
+  vmax = strtod(mj_getPluginConfig(m, instance, "vmax"), nullptr);
   // count plugin bodies
   n = 0;
   for (int i = 1; i < m->nbody; i++) {
@@ -120,6 +148,7 @@ Cable::Cable(const mjModel* m, mjData* d, int instance) {
   prev.assign(n, 0);         // index of previous body
   next.assign(n, 0);         // index of next body
   omega0.assign(3*n, 0);     // reference curvature
+  stress.assign(3*n, 0);     // mechanical stress
   stiffness.assign(4*n, 0);  // material parameters
 
   // run forward kinematics to populate xquat (mjData not yet initialized)
@@ -189,16 +218,17 @@ void Cable::Compute(const mjModel* m, mjData* d, int instance) {
 
     // elastic forces
     mjtNum quat[4] = {0};
-    mjtNum xfrc[3] = {0};
+    mjtNum lfrc[3] = {0};
 
     // local orientation
     if (prev[b]) {
       int qadr = m->jnt_qposadr[m->body_jntadr[i]] + m->body_dofnum[i]-3;
-      QuatDiff(quat, m->body_quat+4*i, d->qpos+qadr, 0);
+      QuatDiff(quat, m->body_quat+4*i, d->qpos+qadr);
 
       // contribution of orientation i-1 to xfrc i
-      LocalForce(xfrc, stiffness.data()+4*b, quat, omega0.data()+3*b,
+      LocalStress(stress.data() + 3 * b, stiffness.data() + 4 * b, quat,
-                 d->xquat+4*(i+prev[b]), 1);
+                  omega0.data() + 3 * b, true);
+      mju_addToScl3(lfrc, stress.data() + 3 * b, 1.0);
     }
 
     if (next[b]) {
@@ -207,19 +237,40 @@ void Cable::Compute(const mjModel* m, mjData* d, int instance) {
 
       // local orientation
       int qadr = m->jnt_qposadr[m->body_jntadr[in]] + m->body_dofnum[in]-3;
-      QuatDiff(quat, m->body_quat+4*in, d->qpos+qadr, 1);
+      QuatDiff(quat, m->body_quat+4*in, d->qpos+qadr);
 
       // contribution of orientation i+1 to xfrc i
-      LocalForce(xfrc, stiffness.data()+4*bn, quat, omega0.data()+3*bn,
+      LocalStress(stress.data() + 3 * bn, stiffness.data() + 4 * bn, quat,
-                 d->xquat+4*i, -1);
+                  omega0.data() + 3 * bn);
+      mju_addToScl3(lfrc, stress.data() + 3 * bn, -1.0);
     }
 
     // convert from global coordinates and apply torque to com
+    mjtNum xfrc[3] = {0};
+    mju_rotVecQuat(xfrc, lfrc, d->xquat+4*i);
     mj_applyFT(m, d, 0, xfrc, d->xpos+3*i, i, d->qfrc_passive);
   }
 }
 
+void Cable::Visualize(const mjModel* m, mjData* d, mjvScene* scn,
+                      int instance) {
+if(!vmax) {
+  return;
+}
 
+for (int b = 0; b < n; b++)  {
+    int i = i0 + b;
+    int bn = b + next[b];
+
+    // set geometry color based on stress norm
+    mjtNum stress_m[3] = {0};
+    mjtNum *stress_l = prev[b] ? stress.data()+3*b : stress.data()+3*bn;
+    mjtNum *stress_r = next[b] ? stress.data()+3*bn : stress.data()+3*b;
+    mju_add3(stress_m, stress_l, stress_r);
+    mju_scl3(stress_m, stress_m, 0.5);
+    scalar2rgba(m->geom_rgba + 4*m->body_geomadr[i], stress_m, 0, vmax);
+  }
+}
 
 void Cable::RegisterPlugin() {
   mjpPlugin plugin;
@@ -228,7 +279,7 @@ void Cable::RegisterPlugin() {
   plugin.name = "mujoco.elasticity.cable";
   plugin.type |= mjPLUGIN_PASSIVE;
 
-  const char* attributes[] = {"twist", "bend", "flat"};
+  const char* attributes[] = {"twist", "bend", "flat", "vmax"};
   plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
   plugin.attributes = attributes;
   plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
@@ -250,8 +301,13 @@ return 0;
     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]);
+        elasticity->Visualize(m, d, scn, instance);
+      };
 
   mjp_registerPlugin(&plugin);
 }
 
-}  // namespace mujoco::plugin::elasticity
+}  // namespace mujoco::plugin::elasticity

plugin/elasticity/cable.h

@@ -21,6 +21,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjtnum.h>
+#include <mujoco/mjvisualize.h>
 
 
 namespace mujoco::plugin::elasticity {
@@ -35,6 +36,7 @@ class Cable {
   ~Cable() = default;
 
   void Compute(const mjModel* m, mjData* d, int instance);
+  void Visualize(const mjModel* m, mjData* d, mjvScene* scn, int instance);
 
   static void RegisterPlugin();
 
@@ -44,6 +46,8 @@ class Cable {
   std::vector<int> next;          // indices of next bodies       (n x 1)
   std::vector<mjtNum> stiffness;  // stiffness parameters         (n x 4)
   std::vector<mjtNum> omega0;     // reference curvature          (n x 3)
+  std::vector<mjtNum> stress;     // mechanical stress            (n x 3)
+  mjtNum vmax;                    // max value in colormap
 
  private:
   Cable(const mjModel* m, mjData* d, int instance);

src/engine/engine_vis_visualize.c

@@ -23,6 +23,7 @@
 #include <mujoco/mjvisualize.h>
 #include "engine/engine_array_safety.h"
 #include "engine/engine_macro.h"
+#include "engine/engine_plugin.h"
 #include "engine/engine_support.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
@@ -2020,6 +2021,22 @@ void mjv_updateScene(const mjModel* m, mjData* d, const mjvOption* opt,
   if (opt->flags[mjVIS_SKIN]) {
     mjv_updateActiveSkin(m, d, scn, opt);
   }
+
+  // update plugin
+  if (m->nplugin) {
+    const int nslot = mjp_pluginCount();
+    // iterate over plugins, call visualize if defined
+    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->visualize) {
+        plugin->visualize(m, d, scn, i);
+      }
+    }
+  }
 }