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

Implement plugin mechanism for actuators and sensors.

Browse Source 
PiperOrigin-RevId: 474874088
Change-Id: I65a8ffdf845f4fa0f8266c165883a747ab5812d8
This commit is contained in:
Saran Tunyasuvunakool
2022-09-16 12:20:56 -07:00
committed by Copybara-Service
parent f556d4d94f
commit 1e2a9a53bc
29 changed files with 2121 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/user/user_model.cc
+131 -2
View File
@@ -24,11 +24,13 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjplugin.h>
#include <mujoco/mjvisualize.h>
#include "cc/array_safety.h"
#include "engine/engine_forward.h"
#include "engine/engine_io.h"
#include "engine/engine_macro.h"
#include "engine/engine_plugin.h"
#include "engine/engine_setconst.h"
#include "engine/engine_support.h"
#include "engine/engine_util_blas.h"
@@ -123,6 +125,7 @@ mjCModel::mjCModel() {
nuserdata = 0;
nkey = 0;
nmocap = 0;
nplugin = 0;
nuser_body = -1;
nuser_jnt = -1;
nuser_geom = -1;
@@ -198,6 +201,7 @@ mjCModel::~mjCModel() {
for (i=0; i<texts.size(); i++) delete texts[i];
for (i=0; i<tuples.size(); i++) delete tuples[i];
for (i=0; i<keys.size(); i++) delete keys[i];
for (i=0; i<plugins.size(); i++) delete plugins[i];
for (i=0; i<defaults.size(); i++) delete defaults[i];
// clear pointer lists created in model construction
@@ -216,6 +220,7 @@ mjCModel::~mjCModel() {
texts.clear();
tuples.clear();
keys.clear();
plugins.clear();
defaults.clear();
// clear sizes and pointer lists created in Compile
@@ -267,6 +272,7 @@ void mjCModel::Clear(void) {
nnumericdata = 0;
ntextdata = 0;
ntupledata = 0;
npluginattr = 0;
nnames = 0;
nemax = 0;
nM = 0;
@@ -283,6 +289,7 @@ void mjCModel::Clear(void) {
lights.clear();
// internal variables
hasImplicitPluginElem = false;
compiled = false;
errInfo = mjCError();
fixCount = 0;
@@ -405,6 +412,12 @@ mjCKey* mjCModel::AddKey(void) {
}
// add plugin instance
mjCPlugin* mjCModel::AddPlugin(void) {
return AddObject(plugins, "plugin");
}
//------------------------ API FOR ACCESS TO MODEL ELEMENTS ---------------------------------------
@@ -455,6 +468,8 @@ int mjCModel::NumObjects(mjtObj type) {
return (int)tuples.size();
case mjOBJ_KEY:
return (int)keys.size();
case mjOBJ_PLUGIN:
return (int)plugins.size();
default:
return 0;
}
@@ -509,6 +524,8 @@ mjCBase* mjCModel::GetObject(mjtObj type, int id) {
return tuples[id];
case mjOBJ_KEY:
return keys[id];
case mjOBJ_PLUGIN:
return plugins[id];
default:
return 0;
}
@@ -646,6 +663,8 @@ mjCBase* mjCModel::FindObject(mjtObj type, string name) {
return findobject(name, texts);
case mjOBJ_TUPLE:
return findobject(name, tuples);
case mjOBJ_PLUGIN:
return findobject(name, plugins);
default:
return 0;
}
@@ -893,6 +912,7 @@ void mjCModel::SetSizes(void) {
ntext = (int)texts.size();
ntuple = (int)tuples.size();
nkey = (int)keys.size();
nplugin = (int)plugins.size();
// nq, nv
for (i=0; i<njnt; i++) {
@@ -955,6 +975,9 @@ void mjCModel::SetSizes(void) {
// ntupledata
for (i=0; i<ntuple; i++) ntupledata += (int)tuples[i]->objtype.size();
// npluginattr
for (i=0; i<nplugin; i++) npluginattr += (int)plugins[i]->flattened_attributes.size();
// nnames
nnames = (int)modelname.size() + 1;
for (i=0; i<nbody; i++) nnames += (int)bodies[i]->name.length() + 1;
@@ -978,6 +1001,7 @@ void mjCModel::SetSizes(void) {
for (i=0; i<ntext; i++) nnames += (int)texts[i]->name.length() + 1;
for (i=0; i<ntuple; i++) nnames += (int)tuples[i]->name.length() + 1;
for (i=0; i<nkey; i++) nnames += (int)keys[i]->name.length() + 1;
for (i=0; i<nplugin; i++) nnames += (int)plugins[i]->name.length() + 1;
// nemax
for (i=0; i<neq; i++)
@@ -1234,6 +1258,7 @@ void mjCModel::CopyNames(mjModel* m) {
adr = namelist(texts, adr, m->name_textadr, m->names);
adr = namelist(tuples, adr, m->name_tupleadr, m->names);
adr = namelist(keys, adr, m->name_keyadr, m->names);
adr = namelist(plugins, adr, m->name_pluginadr, m->names);
// check size, SHOULD NOT OCCUR
if (adr != nnames) {
@@ -1776,7 +1801,7 @@ void mjCModel::CopyObjects(mjModel* m) {
m->tendon_adr[i] = adr;
m->tendon_num[i] = (int)pte->path.size();
m->tendon_matid[i] = pte->matid;
m->tendon_group[i] = pte->group;;
m->tendon_group[i] = pte->group;
m->tendon_limited[i] = pte->limited;
m->tendon_width[i] = (mjtNum)pte->width;
copyvec(m->tendon_solref_lim+mjNREF*i, pte->solref_limit, mjNREF);
@@ -2390,6 +2415,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
processlist(texts, "text");
processlist(tuples, "tuple");
processlist(keys, "key");
processlist(plugins, "plugin");
// set default names, convert names into indices
SetDefaultNames();
@@ -2477,6 +2503,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
for (int i=0; i<numerics.size(); i++) numerics[i]->Compile();
for (int i=0; i<texts.size(); i++) texts[i]->Compile();
for (int i=0; i<tuples.size(); i++) tuples[i]->Compile();
for (int i=0; i<plugins.size(); i++) plugins[i]->Compile();
// compile defaults: to enforce userdata length for writer
for (int i=0; i<defaults.size(); i++) {
@@ -2545,7 +2572,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
nhfield, nhfielddata, ntex, ntexdata, nmat, npair, nexclude,
neq, ntendon, nwrap, nsensor,
nnumeric, nnumericdata, ntext, ntextdata,
ntuple, ntupledata, nkey, nmocap,
ntuple, ntupledata, nkey, nmocap, nplugin, npluginattr,
nuser_body, nuser_jnt, nuser_geom, nuser_site, nuser_cam,
nuser_tendon, nuser_actuator, nuser_sensor, nnames);
if (!m) {
@@ -2558,6 +2585,69 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
CopyNames(m);
CopyTree(m);
// assign plugin slots and copy plugin config attributes
{
int adr = 0;
for (int i = 0; i < nplugin; ++i) {
m->plugin[i] = plugins[i]->plugin_slot;
const int size = plugins[i]->flattened_attributes.size();
std::memcpy(m->plugin_attr + adr,
plugins[i]->flattened_attributes.data(), size);
m->plugin_attradr[i] = adr;
adr += size;
}
}
// query and set plugin-related information
{
// set actuator_plugin to the plugin instance ID
for (int i = 0; i < nu; ++i) {
if (actuators[i]->is_plugin) {
m->actuator_plugin[i] = actuators[i]->plugin_instance->id;
} else {
m->actuator_plugin[i] = -1;
}
}
// set sensor_plugin to the plugin instance ID
std::vector<std::vector<int>> plugin_to_sensors(nplugin);
for (int i = 0; i < nsensor; ++i) {
if (sensors[i]->type == mjSENS_PLUGIN) {
int sensor_plugin = sensors[i]->plugin_instance->id;
m->sensor_plugin[i] = sensor_plugin;
plugin_to_sensors[sensor_plugin].push_back(i);
} else {
m->sensor_plugin[i] = -1;
}
}
// query plugin->nstate, compute and set plugin_state and plugin_stateadr
// for sensor plugins, also query plugin->nsensordata and set nsensordata
int stateadr = 0;
for (int i = 0; i < nplugin; ++i) {
const mjpPlugin* plugin = mjp_getPluginAtSlot(m->plugin[i]);
if (!plugin->nstate) {
mju_error_i("`nstate` is null for plugin at slot %d", m->plugin[i]);
}
int nstate = plugin->nstate(m, i);
m->plugin_stateadr[i] = stateadr;
stateadr += nstate;
if (plugin->type & 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]);
}
int nsensordata = plugin->nsensordata(m, i, sensor_id);
sensors[sensor_id]->dim = nsensordata;
sensors[sensor_id]->needstage =
static_cast<mjtStage>(plugin->needstage);
this->nsensordata += nsensordata;
}
}
}
m->npluginstate = stateadr;
}
// keyframe compilation needs access to nq, nv, na, nmocap, qpos0
for (int i=0; i<keys.size(); i++) {
keys[i]->Compile(m);
@@ -2899,3 +2989,42 @@ bool mjCModel::CopyBack(const mjModel* m) {
return true;
}
void mjCModel::ResolvePlugin(mjCBase* obj, const std::string& plugin_name,
const std::string& plugin_instance_name, mjCPlugin** plugin_instance) {
// if plugin_name is specified, check if it is in the list of active plugins
// (in XML, active plugins are those declared as <required>)
int plugin_slot = -1;
if (!plugin_name.empty()) {
for (int i = 0; i < active_plugins.size(); ++i) {
if (active_plugins[i].first->name == plugin_name) {
plugin_slot = active_plugins[i].second;
break;
}
}
if (plugin_slot == -1) {
throw mjCError(obj, "unrecognized plugin '%s'", plugin_name.c_str());
}
}
// implicit plugin instance
if (*plugin_instance && (*plugin_instance)->plugin_slot == -1) {
(*plugin_instance)->plugin_slot = plugin_slot;
(*plugin_instance)->parent = obj;
}
// explicit plugin instance, look up existing mjCPlugin by instance name
else if (!*plugin_instance) {
*plugin_instance =
static_cast<mjCPlugin*>(FindObject(mjOBJ_PLUGIN, plugin_instance_name));
if (!*plugin_instance) {
throw mjCError(
obj, "unrecognized name '%s' for plugin instance", plugin_instance_name.c_str());
}
if (plugin_slot != -1 && plugin_slot != (*plugin_instance)->plugin_slot) {
throw mjCError(
obj, "'plugin' attribute does not match that of the instance");
}
plugin_slot = (*plugin_instance)->plugin_slot;
}
}
src/user/user_model.h
+15
View File
@@ -16,10 +16,12 @@
#define MUJOCO_SRC_USER_USER_MODEL_H_
#include <string>
#include <utility>
#include <vector>
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjplugin.h>
#include "user/user_objects.h"
typedef enum _mjtInertiaFromGeom {
@@ -85,6 +87,7 @@ class mjCModel {
mjCText* AddText(void); // custom text
mjCTuple* AddTuple(void); // custom tuple
mjCKey* AddKey(void); // keyframe
mjCPlugin* AddPlugin(void); // plugin instance
//------------------------ API for access to model elements (outside tree)
int NumObjects(mjtObj type); // number of objects in specified list
@@ -100,6 +103,12 @@ class mjCModel {
mjCBase* FindObject(mjtObj type, std::string name); // find object given type and name
bool IsNullPose(const mjtNum* pos, const mjtNum* quat); // detect null pose
//------------------------ API for plugins
void ResolvePlugin(mjCBase* obj, // resolve plugin instance, create a new one if needed
const std::string& plugin_name,
const std::string& plugin_instance_name,
mjCPlugin** plugin_instance);
//------------------------ global data
std::string comment; // comment at top of XML
@@ -198,6 +207,7 @@ class mjCModel {
int ntuple; // number of tuple fields
int nkey; // number of keyframes
int nmocap; // number of mocap bodies
int nplugin; // number of plugin instances
// sizes computed by Compile
int nq; // number of generalized coordinates = dim(qpos)
@@ -220,6 +230,7 @@ class mjCModel {
int nnumericdata; // number of mjtNums in all custom fields
int ntextdata; // number of chars in all text fields, including 0
int ntupledata; // number of objects in all tuple fields
int npluginattr; // number of chars in all plugin config attributes
int nnames; // number of chars in all names
int nM; // number of non-zeros in sparse inertia matrix
int nD; // number of non-zeros in sparse derivative matrix
@@ -242,6 +253,9 @@ class mjCModel {
std::vector<mjCTuple*> tuples; // list of tuple fields
std::vector<mjCKey*> keys; // list of keyframe fields
std::vector<std::pair<const mjpPlugin*, int>> active_plugins; // list of active plugins
std::vector<mjCPlugin*> plugins; // list of plugin instances
// pointers to objects created inside kinematic tree
std::vector<mjCBody*> bodies; // list of bodies
std::vector<mjCJoint*> joints; // list of joints allowing motion relative to parent
@@ -251,6 +265,7 @@ class mjCModel {
std::vector<mjCLight*> lights; // list of lights
//------------------------ internal variables
bool hasImplicitPluginElem; // already encountered an implicit plugin sensor/actuator
bool compiled; // already compiled flag (cannot be compiled again)
mjCError errInfo; // last error info
int fixCount; // how many bodies have been fixed
src/user/user_objects.cc
+93 -1
View File
@@ -25,12 +25,14 @@
#include "lodepng.h"
#include <mujoco/mjmodel.h>
#include <mujoco/mjplugin.h>
#include "cc/array_safety.h"
#include "engine/engine_core_smooth.h"
#include "engine/engine_crossplatform.h"
#include "engine/engine_file.h"
#include "engine/engine_io.h"
#include "engine/engine_macro.h"
#include "engine/engine_plugin.h"
#include "engine/engine_util_errmem.h"
#include "engine/engine_util_misc.h"
#include "engine/engine_util_solve.h"
@@ -3391,6 +3393,11 @@ mjCActuator::mjCActuator(mjCModel* _model, mjCDef* _def) {
// set model, def
model = _model;
def = (_def ? _def : (_model ? _model->defaults[0] : 0));
is_plugin = false;
plugin_instance = nullptr;
plugin_name = "";
plugin_instance_name = "";
}
@@ -3553,6 +3560,21 @@ void mjCActuator::Compile(void) {
} else {
trnid[0] = ptarget->id;
}
// plugin
if (is_plugin) {
if (plugin_name.empty() && plugin_instance_name.empty()) {
throw mjCError(
this, "neither 'plugin' nor 'instance' is specified for actuator '%s', (id = %d)",
name.c_str(), id);
}
model->ResolvePlugin(this, plugin_name, plugin_instance_name, &plugin_instance);
const mjpPlugin* plugin = mjp_getPluginAtSlot(plugin_instance->plugin_slot);
if (!(plugin->type & mjPLUGIN_ACTUATOR)) {
throw mjCError(this, "plugin '%s' does not support actuators", plugin->name);
}
}
}
@@ -3580,6 +3602,10 @@ mjCSensor::mjCSensor(mjCModel* _model) {
// clear private variables
objid = -1;
refid = -1;
plugin_instance = nullptr;
plugin_name = "";
plugin_instance_name = "";
}
@@ -3624,7 +3650,7 @@ void mjCSensor::Compile(void) {
// get sensorized object id
objid = pobj->id;
} else if (type != mjSENS_CLOCK) {
} else if (type != mjSENS_CLOCK && type != mjSENS_PLUGIN) {
throw mjCError(this, "invalid type in sensor '%s' (id = %d)", name.c_str(), id);
}
@@ -3919,6 +3945,28 @@ void mjCSensor::Compile(void) {
}
break;
case mjSENS_PLUGIN:
dim = 0; // to be filled in by the plugin later
datatype = mjDATATYPE_REAL; // no noise added to plugin sensors, this attribute is unused
if (plugin_name.empty() && plugin_instance_name.empty()) {
throw mjCError(
this, "neither 'plugin' nor 'instance' is specified for sensor '%s', (id = %d)",
name.c_str(), id);
}
// resolve plugin instance, or create one if using the "plugin" attribute shortcut
{
model->ResolvePlugin(this, plugin_name, plugin_instance_name, &plugin_instance);
const mjpPlugin* plugin = mjp_getPluginAtSlot(plugin_instance->plugin_slot);
if (!(plugin->type & mjPLUGIN_SENSOR)) {
throw mjCError(this, "plugin '%s' does not support sensors", plugin->name);
}
needstage = static_cast<mjtStage>(plugin->needstage);
}
break;
default:
throw mjCError(this, "invalid type in sensor '%s' (id = %d)", name.c_str(), id);
}
@@ -4174,3 +4222,47 @@ void mjCKey::Compile(const mjModel* m) {
}
}
//------------------ class mjCPlugin implementation ------------------------------------------------
// initialize defaults
mjCPlugin::mjCPlugin(mjCModel* _model) {
name = "";
plugin_slot = -1;
nstate = 0;
parent = this;
model = _model;
}
// compiler
void mjCPlugin::Compile(void) {
const mjpPlugin* plugin = mjp_getPluginAtSlot(this->plugin_slot);
// concatenate all of the plugin's attribute values (as null-terminated strings) into
// flattened_attributes, in the order declared in the mjpPlugin
// each valid attribute found is appended to flattened_attributes and removed from xml_attributes
for (int i = 0; i < plugin->nattribute; ++i) {
std::string_view attr(plugin->attributes[i]);
auto it = config_attribs.find(attr);
if (it == config_attribs.end()) {
flattened_attributes.push_back('\0');
} else {
auto original_size = flattened_attributes.size();
flattened_attributes.resize(original_size + it->second.size() + 1);
std::memcpy(&flattened_attributes[original_size], it->second.c_str(),
it->second.size() + 1);
config_attribs.erase(it);
}
}
// anything left in xml_attributes at this stage is not a valid attribute
if (!config_attribs.empty()) {
std::string error =
"unrecognized attribute 'plugin:" + config_attribs.begin()->first +
"' for plugin " + std::string(plugin->name) + "'";
throw mjCError(parent, error.c_str());
}
}
src/user/user_objects.h
+33
View File
@@ -15,6 +15,7 @@
#ifndef MUJOCO_SRC_USER_USER_OBJECTS_H_
#define MUJOCO_SRC_USER_USER_OBJECTS_H_
#include <map>
#include <string>
#include <vector>
@@ -833,6 +834,27 @@ class mjCWrap : public mjCBase {
//------------------------- class mjCPlugin --------------------------------------------------------
// Describes an instance of a plugin
class mjCPlugin : public mjCBase {
friend class mjCModel;
friend class mjXWriter;
public:
int plugin_slot; // global registered slot number of the plugin
int nstate; // state size for the plugin instance
mjCBase* parent; // parent object (only used when generating error message)
std::map<std::string, std::string, std::less<>> config_attribs; // raw config attributes from XML
std::vector<char> flattened_attributes; // config attributes flattened in plugin-declared order
private:
mjCPlugin(mjCModel*); // constructor
void Compile(void); // compiler
};
//------------------------- class mjCActuator ------------------------------------------------------
// Describes an actuator
@@ -865,6 +887,12 @@ class mjCActuator : public mjCBase {
std::string slidersite; // site defining cylinder, for slider-crank only
std::string refsite; // reference site, for site transmission only
// plugin support
bool is_plugin;
std::string plugin_name;
std::string plugin_instance_name;
mjCPlugin* plugin_instance;
private:
mjCActuator(mjCModel* = 0, mjCDef* = 0);// constructor
void Compile(void); // compiler
@@ -896,6 +924,11 @@ class mjCSensor : public mjCBase {
double noise; // noise stdev
std::vector<double> userdata; // user data
// plugin support
std::string plugin_name;
std::string plugin_instance_name;
mjCPlugin* plugin_instance;
private:
mjCSensor(mjCModel*); // constructor
void Compile(void); // compiler