chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add updateScene callback for plugins.

Browse Source 
PiperOrigin-RevId: 492497689
Change-Id: I4a15cc5fa39dc7c9b85d132727f82033d4aec2c8
This commit is contained in:
Alessio Quaglino
2022-12-02 10:10:57 -08:00
committed by Copybara-Service
parent 9fcf3b0c2e
commit 0d52feaa94
9 changed files with 214 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+2
View File
@@ -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
+5
View File
@@ -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
+101 -101
View File
@@ -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
+1
View File
@@ -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
+1
View File
@@ -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
+2
View File
@@ -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
+81 -25
View File
@@ -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<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
+4
View File
@@ -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
+17
View File
@@ -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);
}
}
}
}