diff --git a/doc/changelog.rst b/doc/changelog.rst index 5f976430..1fe612a6 100644 --- a/doc/changelog.rst +++ b/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` no longer require :at:`objtype` and :at:`objname`. If unspecified, the objtype will be :ref:`mjOBJ_UNKNOWN`. ``user`` sensors :at:`datatype` default is now :at-val:`"real"`. - Add support for capsules in URDF import. diff --git a/include/mujoco/mjplugin.h b/include/mujoco/mjplugin.h index 025e2733..466f230e 100644 --- a/include/mujoco/mjplugin.h +++ b/include/mujoco/mjplugin.h @@ -17,6 +17,8 @@ #include #include +#include + 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; diff --git a/introspect/enums.py b/introspect/enums.py index f9152bf4..89bb479b 100755 --- a/introspect/enums.py +++ b/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), + ]), + )), ]) diff --git a/model/plugin/belt.xml b/model/plugin/belt.xml index 050785ca..cdb4d181 100644 --- a/model/plugin/belt.xml +++ b/model/plugin/belt.xml @@ -59,6 +59,7 @@ + diff --git a/model/plugin/cable.xml b/model/plugin/cable.xml index 5edd11cc..4d1337ff 100644 --- a/model/plugin/cable.xml +++ b/model/plugin/cable.xml @@ -35,6 +35,7 @@ + diff --git a/model/plugin/coil.xml b/model/plugin/coil.xml index 8686941c..2709cc9a 100644 --- a/model/plugin/coil.xml +++ b/model/plugin/coil.xml @@ -70,6 +70,7 @@ + @@ -80,6 +81,7 @@ + diff --git a/plugin/elasticity/cable.cc b/plugin/elasticity/cable.cc index d7192728..08bd6884 100644 --- a/plugin/elasticity/cable.cc +++ b/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 -void LocalForce(mjtNum qfrc[3], const mjtNum stiffness[4], +// stress - local stress contribution +void LocalStress(mjtNum stress[3], + const mjtNum stiffness[4], const mjtNum quat[4], const mjtNum omega0[3], - const mjtNum xquat[4], mjtNum scl) { - mjtNum omega[3], lfrc[3]; + bool pullback = false) { + 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 - mju_addToScl3(qfrc, lfrc, scl); + // pull-back into the other body frame + 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, - d->xquat+4*(i+prev[b]), 1); + LocalStress(stress.data() + 3 * b, stiffness.data() + 4 * b, quat, + 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, - d->xquat+4*i, -1); + LocalStress(stress.data() + 3 * bn, stiffness.data() + 4 * bn, quat, + 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(d->plugin_data[instance]); elasticity->Compute(m, d, instance); }; + plugin.visualize = + +[](const mjModel* m, mjData* d, mjvScene* scn, int instance) { + auto* elasticity = reinterpret_cast(d->plugin_data[instance]); + elasticity->Visualize(m, d, scn, instance); + }; mjp_registerPlugin(&plugin); } -} // namespace mujoco::plugin::elasticity +} // namespace mujoco::plugin::elasticity \ No newline at end of file diff --git a/plugin/elasticity/cable.h b/plugin/elasticity/cable.h index 186bb127..f6f7542e 100644 --- a/plugin/elasticity/cable.h +++ b/plugin/elasticity/cable.h @@ -21,6 +21,7 @@ #include #include #include +#include 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 next; // indices of next bodies (n x 1) std::vector stiffness; // stiffness parameters (n x 4) std::vector omega0; // reference curvature (n x 3) + std::vector stress; // mechanical stress (n x 3) + mjtNum vmax; // max value in colormap private: Cable(const mjModel* m, mjData* d, int instance); diff --git a/src/engine/engine_vis_visualize.c b/src/engine/engine_vis_visualize.c index 577fd40e..31284ee4 100644 --- a/src/engine/engine_vis_visualize.c +++ b/src/engine/engine_vis_visualize.c @@ -23,6 +23,7 @@ #include #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; inplugin; 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); + } + } + } }