From 1e2a9a53bcca2e2bf6758b350db32346d132098d Mon Sep 17 00:00:00 2001 From: Saran Tunyasuvunakool Date: Fri, 16 Sep 2022 12:20:56 -0700 Subject: [PATCH] Implement plugin mechanism for actuators and sensors. PiperOrigin-RevId: 474874088 Change-Id: I65a8ffdf845f4fa0f8266c165883a747ab5812d8 --- CMakeLists.txt | 1 + doc/changelog.rst | 8 + include/mujoco/mjdata.h | 8 + include/mujoco/mjmodel.h | 18 +- include/mujoco/mjplugin.h | 61 +++ include/mujoco/mjxmacro.h | 14 +- include/mujoco/mujoco.h | 28 ++ introspect/enums.py | 13 +- introspect/functions.py | 97 +++++ src/engine/CMakeLists.txt | 2 + src/engine/engine_forward.c | 44 ++ src/engine/engine_io.c | 110 ++++- src/engine/engine_io.h | 7 +- src/engine/engine_plugin.cc | 431 +++++++++++++++++++ src/engine/engine_plugin.h | 62 +++ src/engine/engine_sensor.c | 88 ++++ src/engine/engine_support.c | 25 +- src/engine/engine_util_misc.c | 7 + src/user/user_model.cc | 133 +++++- src/user/user_model.h | 15 + src/user/user_objects.cc | 94 +++- src/user/user_objects.h | 33 ++ src/xml/xml_native_reader.cc | 164 ++++++- src/xml/xml_native_reader.h | 1 + src/xml/xml_native_writer.cc | 148 ++++++- src/xml/xml_native_writer.h | 1 + test/engine/CMakeLists.txt | 8 + test/engine/engine_plugin_test.cc | 524 +++++++++++++++++++++++ unity/Runtime/Bindings/MujocoBindings.cs | 27 +- 29 files changed, 2121 insertions(+), 51 deletions(-) create mode 100644 include/mujoco/mjplugin.h create mode 100644 src/engine/engine_plugin.cc create mode 100644 src/engine/engine_plugin.h create mode 100644 test/engine/engine_plugin_test.cc diff --git a/CMakeLists.txt b/CMakeLists.txt index ee665ca0..a8407f44 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -56,6 +56,7 @@ set(MUJOCO_HEADERS include/mujoco/mjdata.h include/mujoco/mjexport.h include/mujoco/mjmodel.h + include/mujoco/mjplugin.h include/mujoco/mjrender.h include/mujoco/mjtnum.h include/mujoco/mjui.h diff --git a/doc/changelog.rst b/doc/changelog.rst index bc81bf37..e9cde2aa 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -28,6 +28,14 @@ General - Added :ref:`mju_mulVecMatVec` to multiply a square matrix :math:`M` with vectors :math:`x` and :math:`y` on both sides. The function returns :math:`x^TMy`. +- Added new plugin API. Plugins allow developers to extend MuJoCo's capability without modifying core engine code. + The plugin mechanism is intended to replace the existing callbacks, though these will remain for the time being as an + option for simple use cases and backward compatibility. The new mechanism manages stateful plugins and supports + multiple plugins from different sources, allowing MuJoCo extensions to be introduced in a modular fashion, rather than + as global overrides. Note the new mechanism is currently undocumented except in code, as we test it internally. + If you are interested in using the pluging mechanism, please get in touch first. + + Version 2.2.2 (September 7, 2022) --------------------------------- diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index d84d5f85..5ef00a87 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -15,6 +15,8 @@ #ifndef MUJOCO_MJDATA_H_ #define MUJOCO_MJDATA_H_ +#include + #include #include @@ -124,6 +126,7 @@ struct mjData_ { // constant sizes int nstack; // number of mjtNums that can fit in stack int nbuffer; // size of main buffer in bytes + int nplugin; // number of plugin instances // stack pointer int pstack; // first available mjtNum address in stack @@ -164,6 +167,7 @@ struct mjData_ { mjtNum* qvel; // velocity (nv x 1) mjtNum* act; // actuator activation (na x 1) mjtNum* qacc_warmstart; // acceleration used for warmstart (nv x 1) + mjtNum* plugin_state; // plugin state (npluginstate x 1) // control mjtNum* ctrl; // control (nu x 1) @@ -184,6 +188,10 @@ struct mjData_ { // sensors mjtNum* sensordata; // sensor data array (nsensordata x 1) + // plugins + int* plugin; // copy of m->plugin, required for deletion (nplugin x 1) + uintptr_t* plugin_data; // pointer to plugin-managed data structure (nplugin x 1) + //-------------------------------- POSITION dependent // computed by mj_fwdPosition/mj_kinematics diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index b44e7817..02541969 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -239,7 +239,8 @@ typedef enum mjtObj_ { // type of MujoCo object mjOBJ_NUMERIC, // numeric mjOBJ_TEXT, // text mjOBJ_TUPLE, // tuple - mjOBJ_KEY // keyframe + mjOBJ_KEY, // keyframe + mjOBJ_PLUGIN // plugin instance } mjtObj; @@ -315,6 +316,9 @@ typedef enum mjtSensor_ { // type of sensor // global sensors mjSENS_CLOCK, // simulation time + // plugin-controlled sensors + mjSENS_PLUGIN, // plugin-controlled + // user-defined sensor mjSENS_USER // sensor data provided by mjcb_sensor callback } mjtSensor; @@ -565,6 +569,8 @@ struct mjModel_ { int ntupledata; // number of objects in all tuple fields int nkey; // number of keyframes int nmocap; // number of mocap bodies + int nplugin; // number of plugin instances + int npluginattr; // number of chars in all plugin config attributes int nuser_body; // number of mjtNums in body_user int nuser_jnt; // number of mjtNums in jnt_user int nuser_geom; // number of mjtNums in geom_user @@ -584,6 +590,7 @@ struct mjModel_ { int nstack; // number of fields in mjData stack int nuserdata; // number of extra fields in mjData int nsensordata; // number of fields in sensor data vector + int npluginstate; // number of fields in the plugin state vector int nbuffer; // number of bytes in buffer @@ -856,6 +863,7 @@ struct mjModel_ { mjtNum* actuator_length0; // actuator length in qpos0 (nu x 1) mjtNum* actuator_lengthrange; // feasible actuator length range (nu x 2) mjtNum* actuator_user; // user data (nu x nuser_actuator) + int* actuator_plugin; // plugin instance id; -1: not a plugin actuator (nu x 1) // sensors int* sensor_type; // sensor type (mjtSensor) (nsensor x 1) @@ -870,6 +878,13 @@ struct mjModel_ { mjtNum* sensor_cutoff; // cutoff for real and positive; 0: ignore (nsensor x 1) mjtNum* sensor_noise; // noise standard deviation (nsensor x 1) mjtNum* sensor_user; // user data (nsensor x nuser_sensor) + int* sensor_plugin; // plugin instance id; -1: not a plugin sensor (nsensor x 1) + + // plugin instances + int* plugin; // globally registered plugin slot number (nplugin x 1) + int* plugin_stateadr; // address in the plugin state array (nplugin x 1) + char* plugin_attr; // config attributes of plugin instances (npluginattr x 1) + int* plugin_attradr; // address to each instance's config attrib (nplugin x 1) // custom numeric fields int* numeric_adr; // address of field in numeric_data (nnumeric x 1) @@ -919,6 +934,7 @@ struct mjModel_ { int* name_textadr; // text name pointers (ntext x 1) int* name_tupleadr; // tuple name pointers (ntuple x 1) int* name_keyadr; // keyframe name pointers (nkey x 1) + int* name_pluginadr; // plugin instance name pointers (nplugin x 1) char* names; // names of all objects, 0-terminated (nnames x 1) }; typedef struct mjModel_ mjModel; diff --git a/include/mujoco/mjplugin.h b/include/mujoco/mjplugin.h new file mode 100644 index 00000000..be303364 --- /dev/null +++ b/include/mujoco/mjplugin.h @@ -0,0 +1,61 @@ +// Copyright 2022 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_INCLUDE_MJPLUGIN_H_ +#define MUJOCO_INCLUDE_MJPLUGIN_H_ + +#include +#include + +typedef enum mjtPluginTypeBit_ { + mjPLUGIN_ACTUATOR = 1<<0, + mjPLUGIN_SENSOR = 1<<1 +} mjtPluginTypeBit; + +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 needstage; // an mjtStage enum value specifying the sensor computation stage + + // number of mjtNums needed to store the state of a plugin instance (required) + int (*nstate)(const mjModel* m, int instance); + + // dimension of the specified sensor's output (required only for sensor plugins) + int (*nsensordata)(const mjModel* m, int instance, int sensor_id); + + // called when a new mjData is being created (required) + void (*init)(const mjModel* m, mjData* d, int instance); + + // called when an mjData is being freed (optional) + void (*destroy)(mjData* d, int instance); + + // called when an mjData is being copied (optional) + void (*copy)(mjData* dest, const mjModel* m, const mjData* src, int instance); + + // called when an mjData is being reset (required) + 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); + + // called when time integration occurs (optional) + void (*advance)(const mjModel* m, mjData* d, int instance); +}; +typedef struct mjpPlugin_ mjpPlugin; + +#endif // MUJOCO_INCLUDE_MJPLUGIN_H_ diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index cec1ee7b..f6d53c09 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -102,6 +102,8 @@ X( ntupledata ) \ X( nkey ) \ X( nmocap ) \ + X( nplugin ) \ + X( npluginattr ) \ X( nuser_body ) \ X( nuser_jnt ) \ X( nuser_geom ) \ @@ -119,6 +121,7 @@ X( nstack ) \ X( nuserdata ) \ X( nsensordata ) \ + X( npluginstate ) \ X( nbuffer ) @@ -138,7 +141,6 @@ int nu = m->nu; \ int nmocap = m->nmocap; - // macro for annotating that an array size in an X macro is a member of mjModel // by default this macro does nothing, but users can redefine it as necessary #define MJ_M(n) n @@ -367,6 +369,7 @@ X( mjtNum, actuator_length0, nu, 1 ) \ X( mjtNum, actuator_lengthrange, nu, 2 ) \ X( mjtNum, actuator_user, nu, MJ_M(nuser_actuator) ) \ + X( int, actuator_plugin, nu, 1 ) \ X( int, sensor_type, nsensor, 1 ) \ X( int, sensor_datatype, nsensor, 1 ) \ X( int, sensor_needstage, nsensor, 1 ) \ @@ -379,6 +382,11 @@ X( mjtNum, sensor_cutoff, nsensor, 1 ) \ X( mjtNum, sensor_noise, nsensor, 1 ) \ X( mjtNum, sensor_user, nsensor, MJ_M(nuser_sensor) ) \ + X( int, sensor_plugin, nsensor, 1 ) \ + X( int, plugin, nplugin, 1 ) \ + X( int, plugin_stateadr, nplugin, 1 ) \ + X( char, plugin_attr, npluginattr, 1 ) \ + X( int, plugin_attradr, nplugin, 1 ) \ X( int, numeric_adr, nnumeric, 1 ) \ X( int, numeric_size, nnumeric, 1 ) \ X( mjtNum, numeric_data, nnumericdata, 1 ) \ @@ -418,6 +426,7 @@ X( int, name_textadr, ntext, 1 ) \ X( int, name_tupleadr, ntuple, 1 ) \ X( int, name_keyadr, nkey, 1 ) \ + X( int, name_pluginadr, nplugin, 1 ) \ X( char, names, nnames, 1 ) @@ -435,6 +444,7 @@ X( mjtNum, qvel, nv, 1 ) \ X( mjtNum, act, na, 1 ) \ X( mjtNum, qacc_warmstart, nv, 1 ) \ + X( mjtNum, plugin_state, npluginstate, 1 ) \ X( mjtNum, ctrl, nu, 1 ) \ X( mjtNum, qfrc_applied, nv, 1 ) \ X( mjtNum, xfrc_applied, nbody, 6 ) \ @@ -444,6 +454,8 @@ X( mjtNum, act_dot, na, 1 ) \ X( mjtNum, userdata, nuserdata, 1 ) \ X( mjtNum, sensordata, nsensordata, 1 ) \ + X( int, plugin, nplugin, 1 ) \ + X( uintptr_t, plugin_data, nplugin, 1 ) \ X( mjtNum, xpos, nbody, 3 ) \ X( mjtNum, xquat, nbody, 4 ) \ X( mjtNum, xmat, nbody, 9 ) \ diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 41e464c6..059cd2de 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -34,6 +34,7 @@ extern "C" { // type definitions #include #include +#include #include #include #include @@ -453,6 +454,10 @@ MJAPI mjtNum mj_getTotalmass(const mjModel* m); // Scale body masses and inertias to achieve specified total mass. MJAPI void mj_setTotalmass(mjModel* m, mjtNum newmass); +// Return a config attribute value of a plugin instance; +// NULL: invalid plugin instance ID or attribute name +MJAPI const char* mj_getPluginConfig(const mjModel* m, int plugin_id, const char* attrib); + // Return version number: 1.0.2 is encoded as 102. MJAPI int mj_version(void); @@ -1135,6 +1140,29 @@ MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte cen +//---------------------- Plugins ------------------------------------------------------------------- + +// Set default plugin definition. +MJAPI void mjp_defaultPlugin(mjpPlugin* plugin); + +// Globally register a plugin. This function is thread-safe. +// If an identical mjpPlugin is already registered, this function does nothing. +// If a non-identical mjpPlugin with the same name is already registered, an mju_error is raised. +// Two mjpPlugins are considered identical if all member function pointers and numbers are equal, +// and the name and attribute strings are all identical, however the char pointers to the strings +// need not be the same. +MJAPI int mjp_registerPlugin(const mjpPlugin* plugin); + +// Return the number of globally registered plugins. +MJAPI int mjp_pluginCount(); + +// Look up a plugin by name. If slot is not NULL, also write its registered slot number into it. +MJAPI const mjpPlugin* mjp_getPlugin(const char* name, int* slot); + +// Look up a plugin by the registered slot number that was returned by mjp_registerPlugin. +MJAPI const mjpPlugin* mjp_getPluginAtSlot(int slot); + + #if defined(__cplusplus) } #endif diff --git a/introspect/enums.py b/introspect/enums.py index 123b0cb3..afc479ae 100755 --- a/introspect/enums.py +++ b/introspect/enums.py @@ -263,6 +263,7 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjOBJ_TEXT', 21), ('mjOBJ_TUPLE', 22), ('mjOBJ_KEY', 23), + ('mjOBJ_PLUGIN', 24), ]), )), ('mjtConstraint', @@ -333,7 +334,8 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjSENS_SUBTREELINVEL', 33), ('mjSENS_SUBTREEANGMOM', 34), ('mjSENS_CLOCK', 35), - ('mjSENS_USER', 36), + ('mjSENS_PLUGIN', 36), + ('mjSENS_USER', 37), ]), )), ('mjtStage', @@ -406,6 +408,15 @@ ENUMS: Mapping[str, EnumDecl] = dict([ ('mjNTIMER', 13), ]), )), + ('mjtPluginTypeBit', + EnumDecl( + name='mjtPluginTypeBit', + declname='enum mjtPluginTypeBit_', + values=dict([ + ('mjPLUGIN_ACTUATOR', 1), + ('mjPLUGIN_SENSOR', 2), + ]), + )), ('mjtGridPos', EnumDecl( name='mjtGridPos', diff --git a/introspect/functions.py b/introspect/functions.py index 5ffb6f50..e304f040 100644 --- a/introspect/functions.py +++ b/introspect/functions.py @@ -2672,6 +2672,32 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Scale body masses and inertias to achieve specified total mass.', )), + ('mj_getPluginConfig', + FunctionDecl( + name='mj_getPluginConfig', + return_type=PointerType( + inner_type=ValueType(name='char', is_const=True), + ), + parameters=( + FunctionParameterDecl( + name='m', + type=PointerType( + inner_type=ValueType(name='mjModel', is_const=True), + ), + ), + FunctionParameterDecl( + name='plugin_id', + type=ValueType(name='int'), + ), + FunctionParameterDecl( + name='attrib', + type=PointerType( + inner_type=ValueType(name='char', is_const=True), + ), + ), + ), + doc='Return a config attribute value of a plugin instance; NULL: invalid plugin instance ID or attribute name', # pylint: disable=line-too-long + )), ('mj_version', FunctionDecl( name='mj_version', @@ -7147,4 +7173,75 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([ ), doc='Finite differenced transition matrices (control theory notation) d(x_next) = A*dx + B*du d(sensor) = C*dx + D*du required output matrix dimensions: A: (2*nv+na x 2*nv+na) B: (2*nv+na x nu) D: (nsensordata x 2*nv+na) C: (nsensordata x nu)', # pylint: disable=line-too-long )), + ('mjp_defaultPlugin', + FunctionDecl( + name='mjp_defaultPlugin', + return_type=ValueType(name='void'), + parameters=( + FunctionParameterDecl( + name='plugin', + type=PointerType( + inner_type=ValueType(name='mjpPlugin'), + ), + ), + ), + doc='Set default plugin definition.', + )), + ('mjp_registerPlugin', + FunctionDecl( + name='mjp_registerPlugin', + return_type=ValueType(name='int'), + parameters=( + FunctionParameterDecl( + name='plugin', + type=PointerType( + inner_type=ValueType(name='mjpPlugin', is_const=True), + ), + ), + ), + doc='Globally register a plugin. This function is thread-safe. If an identical mjpPlugin is already registered, this function does nothing. If a non-identical mjpPlugin with the same name is already registered, an mju_error is raised. Two mjpPlugins are considered identical if all member function pointers and numbers are equal, and the name and attribute strings are all identical, however the char pointers to the strings need not be the same.', # pylint: disable=line-too-long + )), + ('mjp_pluginCount', + FunctionDecl( + name='mjp_pluginCount', + return_type=ValueType(name='int'), + parameters=(), + doc='Return the number of globally registered plugins.', + )), + ('mjp_getPlugin', + FunctionDecl( + name='mjp_getPlugin', + return_type=PointerType( + inner_type=ValueType(name='mjpPlugin', is_const=True), + ), + parameters=( + FunctionParameterDecl( + name='name', + type=PointerType( + inner_type=ValueType(name='char', is_const=True), + ), + ), + FunctionParameterDecl( + name='slot', + type=PointerType( + inner_type=ValueType(name='int'), + ), + ), + ), + doc='Look up a plugin by name. If slot is not NULL, also write its registered slot number into it.', # pylint: disable=line-too-long + )), + ('mjp_getPluginAtSlot', + FunctionDecl( + name='mjp_getPluginAtSlot', + return_type=PointerType( + inner_type=ValueType(name='mjpPlugin', is_const=True), + ), + parameters=( + FunctionParameterDecl( + name='slot', + type=ValueType(name='int'), + ), + ), + doc='Look up a plugin by the registered slot number that was returned by mjp_registerPlugin.', # pylint: disable=line-too-long + )), ]) diff --git a/src/engine/CMakeLists.txt b/src/engine/CMakeLists.txt index 2103f657..6d6b1fce 100644 --- a/src/engine/CMakeLists.txt +++ b/src/engine/CMakeLists.txt @@ -39,6 +39,8 @@ set(MUJOCO_ENGINE_SRCS engine_io.c engine_io.h engine_macro.h + engine_plugin.cc + engine_plugin.h engine_print.c engine_print.h engine_ray.c diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index da25dd65..996741d2 100644 --- a/src/engine/engine_forward.c +++ b/src/engine/engine_forward.c @@ -19,6 +19,7 @@ #include #include +#include #include "engine/engine_callback.h" #include "engine/engine_collision_driver.h" #include "engine/engine_core_constraint.h" @@ -27,6 +28,7 @@ #include "engine/engine_inverse.h" #include "engine/engine_io.h" #include "engine/engine_macro.h" +#include "engine/engine_plugin.h" #include "engine/engine_sensor.h" #include "engine/engine_solver.h" #include "engine/engine_support.h" @@ -194,6 +196,11 @@ void mj_fwdActuation(const mjModel* m, mjData* d) { // force = gain .* [ctrl/act] + bias for (int i=0; iactuator_plugin[i] >= 0) { + continue; + } + // extract gain info prm = m->actuator_gainprm + mjNGAIN*i; @@ -262,6 +269,24 @@ void mj_fwdActuation(const mjModel* m, mjData* d) { force[i] += bias; } + // handle actuator plugins + if (m->nplugin) { + const int nslot = mjp_pluginCount(); + 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->type & mjPLUGIN_ACTUATOR) { + if (!plugin->compute) { + mju_error_i("`compute` is a null function pointer for plugin at slot %d", slot); + } + plugin->compute(m, d, i, mjPLUGIN_ACTUATOR); + } + } + } + // clamp actuator_force for (int i=0; iactuator_forcelimited[i]) { @@ -275,6 +300,10 @@ void mj_fwdActuation(const mjModel* m, mjData* d) { // act_dot for stateful actuators for (int i=nu-na; iactuator_plugin[i] >= 0) { + continue; + } + // extract info prm = m->actuator_dynprm + i*mjNDYN; int j = i-(nu-na); @@ -481,6 +510,21 @@ static void mj_advance(const mjModel* m, mjData* d, // advance time d->time += m->opt.timestep; + + // advance plugin states + 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->advance) { + plugin->advance(m, d, i); + } + } + } } // Euler integrator, semi-implicit in velocity, possibly skipping factorisation diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index d0d7baa3..46758776 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -22,8 +22,10 @@ #include #include +#include #include #include "engine/engine_macro.h" +#include "engine/engine_plugin.h" #include "engine/engine_util_blas.h" #include "engine/engine_util_errmem.h" #include "engine/engine_vfs.h" @@ -396,9 +398,10 @@ mjModel* mj_makeModel(int nq, int nv, int nu, int na, int nbody, int njnt, int ntex, int ntexdata, int nmat, int npair, int nexclude, int neq, int ntendon, int nwrap, int nsensor, int nnumeric, int nnumericdata, int ntext, int ntextdata, - int ntuple, int ntupledata, int nkey, int nmocap, - int nuser_body, int nuser_jnt, int nuser_geom, int nuser_site, int nuser_cam, - int nuser_tendon, int nuser_actuator, int nuser_sensor, int nnames) { + int ntuple, int ntupledata, int nkey, int nmocap, int nplugin, + int npluginattr, int nuser_body, int nuser_jnt, int nuser_geom, + int nuser_site, int nuser_cam, int nuser_tendon, int nuser_actuator, + int nuser_sensor, int nnames) { intptr_t offset = 0; // allocate mjModel @@ -449,6 +452,8 @@ mjModel* mj_makeModel(int nq, int nv, int nu, int na, int nbody, int njnt, m->ntupledata = ntupledata; m->nkey = nkey; m->nmocap = nmocap; + m->nplugin = nplugin; + m->npluginattr = npluginattr; m->nuser_body = nuser_body; m->nuser_jnt = nuser_jnt; m->nuser_geom = nuser_geom; @@ -533,10 +538,10 @@ mjModel* mj_copyModel(mjModel* dest, const mjModel* src) { src->ntex, src->ntexdata, src->nmat, src->npair, src->nexclude, src->neq, src->ntendon, src->nwrap, src->nsensor, src->nnumeric, src->nnumericdata, src->ntext, src->ntextdata, - src->ntuple, src->ntupledata, src->nkey, src->nmocap, - src->nuser_body, src->nuser_jnt, src->nuser_geom, src->nuser_site, - src->nuser_cam, src->nuser_tendon, src->nuser_actuator, src->nuser_sensor, - src->nnames); + src->ntuple, src->ntupledata, src->nkey, src->nmocap, src->nplugin, + src->npluginattr, src->nuser_body, src->nuser_jnt, src->nuser_geom, + src->nuser_site, src->nuser_cam, src->nuser_tendon, src->nuser_actuator, + src->nuser_sensor, src->nnames); } if (!dest) { mju_error("Failed to make mjModel. Invalid sizes."); @@ -705,7 +710,8 @@ mjModel* mj_loadModel(const char* filename, const mjVFS* vfs) { info[21], info[22], info[23], info[24], info[25], info[26], info[27], info[28], info[29], info[30], info[31], info[32], info[33], info[34], info[35], info[36], info[37], info[38], info[39], info[40], info[41], - info[42], info[43], info[44], info[45], info[46], info[47], info[48]); + info[42], info[43], info[44], info[45], info[46], info[47], info[48], + info[49], info[50]); if (!m || m->nbuffer!=info[getnint()-1]) { if (fp) { fclose(fp); @@ -884,6 +890,17 @@ static mjData* _makeData(const mjModel* m) { // set pointers into buffer, reset data mj_setPtrData(m, d); + // copy plugins into d, required for deletion + d->nplugin = m->nplugin; + for (int i = 0; i < m->nplugin; ++i) { + d->plugin[i] = m->plugin[i]; + const mjpPlugin* plugin = mjp_getPluginAtSlot(m->plugin[i]); + if (!plugin->init) { + mju_error_i("`init` is a null function pointer for plugin at slot %d", m->plugin[i]); + } + plugin->init(m, d, i); + } + return d; } @@ -927,6 +944,17 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) { dest->stack = save_stack; mj_setPtrData(m, dest); + // save plugin_data, since the X macro copying block below will override it + const size_t plugin_data_size = sizeof(*dest->plugin_data) * dest->nplugin; + uintptr_t* save_plugin_data = NULL; + if (plugin_data_size) { + save_plugin_data = (uintptr_t*)mju_malloc(plugin_data_size); + if (!save_plugin_data) { + mju_error("failed to allocate temporary memory for plugin_data"); + } + memcpy(save_plugin_data, dest->plugin_data, plugin_data_size); + } + // copy buffer { MJDATA_POINTERS_PREAMBLE(m) @@ -936,6 +964,22 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) { #undef X } + // restore plugin_data + if (plugin_data_size) { + memcpy(dest->plugin_data, save_plugin_data, plugin_data_size); + free(save_plugin_data); + save_plugin_data = NULL; + } + + // copy plugin instances + dest->nplugin = m->nplugin; + for (int i = 0; i < m->nplugin; ++i) { + const mjpPlugin* plugin = mjp_getPluginAtSlot(m->plugin[i]); + if (plugin->copy) { + plugin->copy(dest, m, src, i); + } + } + return dest; } @@ -966,6 +1010,12 @@ mjtNum* mj_stackAlloc(mjData* d, int size) { // clear data, set defaults static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { + //------------------------------ save plugin state and data + mjtNum* plugin_state = mju_malloc(sizeof(mjtNum) * m->npluginstate); + memcpy(plugin_state, d->plugin_state, sizeof(mjtNum) * m->npluginstate); + uintptr_t* plugindata = mju_malloc(sizeof(uintptr_t) * m->nplugin); + memcpy(plugindata, d->plugin_data, sizeof(uintptr_t) * m->nplugin); + //------------------------------ clear header // clear stack pointer @@ -1048,6 +1098,22 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) { d->mocap_quat[4*i] = 1.0; } } + + // restore pluginstate and plugindata + memcpy(d->plugin_state, plugin_state, sizeof(mjtNum) * m->npluginstate); + mju_free(plugin_state); + memcpy(d->plugin_data, plugindata, sizeof(uintptr_t) * m->nplugin); + mju_free(plugindata); + + // restore the plugin array back into d and reset the instances + for (int i = 0; i < m->nplugin; ++i) { + d->plugin[i] = m->plugin[i]; + const mjpPlugin* plugin = mjp_getPluginAtSlot(m->plugin[i]); + if (!plugin->reset) { + mju_error_i("`reset` is a null function pointer for plugin at slot %d", m->plugin[i]); + } + plugin->reset(m, d, i); + } } @@ -1087,6 +1153,13 @@ void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key) { // de-allocate mjData void mj_deleteData(mjData* d) { if (d) { + // destroy plugin instances + for (int i = 0; i < d->nplugin; ++i) { + const mjpPlugin* plugin = mjp_getPluginAtSlot(d->plugin[i]); + if (plugin->destroy) { + plugin->destroy(d, i); + } + } mju_free(d->buffer); mju_free(d->stack); mju_free(d); @@ -1145,6 +1218,9 @@ static int sensorSize(mjtSensor sensor_type, int nuser_sensor) { case mjSENS_USER: return nuser_sensor; + case mjSENS_PLUGIN: + return -1; + // don't use a 'default' case, so compiler warns about missing values } return -1; @@ -1202,6 +1278,8 @@ static int numObjects(const mjModel* m, mjtObj objtype) { return m->ntuple; case mjOBJ_KEY: return m->nkey; + case mjOBJ_PLUGIN: + return m->nplugin; } return -2; } @@ -1251,6 +1329,10 @@ const char* mj_validateReferences(const mjModel* m) { X(skin_bonevertid, nskinbonevert, nskinvert , 0 ) \ X(pair_geom1, npair, ngeom , 0 ) \ X(pair_geom2, npair, ngeom , 0 ) \ + X(actuator_plugin, nu, nplugin , 0 ) \ + X(sensor_plugin, nsensor, nplugin , 0 ) \ + X(plugin_stateadr, nplugin, npluginstate , 0 ) \ + X(plugin_attradr, nplugin, npluginattr , 0 ) \ X(tendon_adr, ntendon, nwrap , m->tendon_num ) \ X(tendon_matid, ntendon, nmat , 0 ) \ X(numeric_adr, nnumeric, nnumericdata , m->numeric_size ) \ @@ -1459,7 +1541,17 @@ const char* mj_validateReferences(const mjModel* m) { } for (int i=0; insensor; i++) { mjtSensor sensor_type = m->sensor_type[i]; - int sensor_size = sensorSize(sensor_type, m->nuser_sensor); + int sensor_size; + if (sensor_type == mjSENS_PLUGIN) { + const mjpPlugin* plugin = mjp_getPluginAtSlot(m->plugin[m->sensor_plugin[i]]); + if (!plugin->nsensordata) { + mju_error_i("`nsensordata` is a null function pointer for plugin at slot %d", + m->plugin[m->sensor_plugin[i]]); + } + sensor_size = plugin->nsensordata(m, m->sensor_plugin[i], i); + } else { + sensor_size = sensorSize(sensor_type, m->nuser_sensor); + } if (sensor_size < 0) { return "Invalid model: Bad sensor_type."; } diff --git a/src/engine/engine_io.h b/src/engine/engine_io.h index ebb29ce1..c0d18e2f 100644 --- a/src/engine/engine_io.h +++ b/src/engine/engine_io.h @@ -52,9 +52,10 @@ mjModel* mj_makeModel(int nq, int nv, int nu, int na, int nbody, int njnt, int ntex, int ntexdata, int nmat, int npair, int nexclude, int neq, int ntendon, int nwrap, int nsensor, int nnumeric, int nnumericdata, int ntext, int ntextdata, - int ntuple, int ntupledata, int nkey, int nmocap, - int nuser_body, int nuser_jnt, int nuser_geom, int nuser_site, int nuser_cam, - int nuser_tendon, int nuser_actuator, int nuser_sensor, int nnames); + int ntuple, int ntupledata, int nkey, int nmocap, int nplugin, + int npluginattr, int nuser_body, int nuser_jnt, int nuser_geom, + int nuser_site, int nuser_cam, int nuser_tendon, int nuser_actuator, + int nuser_sensor, int nnames); // copy mjModel; allocate new if dest is NULL MJAPI mjModel* mj_copyModel(mjModel* dest, const mjModel* src); diff --git a/src/engine/engine_plugin.cc b/src/engine/engine_plugin.cc new file mode 100644 index 00000000..73823a40 --- /dev/null +++ b/src/engine/engine_plugin.cc @@ -0,0 +1,431 @@ +// Copyright 2022 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Plugin registration is implemented in C++, unlike the rest of the engine code which is in C. +// This is because C++ provides a standard cross-platform mutex, which we use to guard the global +// plugin table in order to make the API thread-safe. We only expose a C API externally, and this +// entire file can in principle be re-implemented in C if necessary, without breaking any external +// or internal MuJoCo code elsewhere. + +#include "engine/engine_plugin.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef __APPLE__ +#include +#if !defined(MAC_OS_X_VERSION_MIN_REQUIRED) && defined(__MAC_OS_X_VERSION_MIN_REQUIRED) +#define MAC_OS_X_VERSION_MIN_REQUIRED __MAC_OS_X_VERSION_MIN_REQUIRED +#endif +#endif + +#include +#include "engine/engine_util_errmem.h" + +// set default plugin definition +void mjp_defaultPlugin(mjpPlugin* plugin) { + std::memset(plugin, 0, sizeof(*plugin)); +} + +namespace { +constexpr int kMaxNameLength = 1024; +constexpr int kMaxAttributes = 255; + +constexpr int kCacheLine = 64; + +// A table of registered plugins, implemented as a linked list of array "blocks". +// This is a compromise that maintains a good degree of memory locality while not invalidating +// existing pointers when growing the table. It is expected that for most users, the number of +// plugins loaded into a program will be small enough to fit in the initial block, and so the global +// table will behave like an array. Since pointers are never invalidated, we do not need to apply a +// read lock on the global table when resolving a plugin. +struct alignas(kCacheLine) PluginTable { + static constexpr int kBlockSize = 15; + + PluginTable() { + for (int i = 0; i < kBlockSize; ++i) { + mjp_defaultPlugin(&plugins[i]); + } + } + + mjpPlugin plugins[kBlockSize]; + PluginTable* next = nullptr; +}; + +static_assert( + sizeof(PluginTable) / kCacheLine == + sizeof(PluginTable::plugins) / kCacheLine + (sizeof(PluginTable::plugins) % kCacheLine > 0), + "PluginTable::next doesn't fit in the same cache line as the end of PluginTable::plugins"); + +using Mutex = std::shared_mutex; + +class Global { + public: + Global() { + new(mutex_) Mutex; + } + PluginTable& table() { + return table_; + } + std::atomic_int& count() { + return count_; + } + Mutex& mutex() { + return *std::launder(reinterpret_cast(&mutex_)); + } + + private: + PluginTable table_; + std::atomic_int count_; + + // A mutex whose destructor is never run. + // When a C++ program terminates, the destructors for function static objects and globals will be + // executed by whichever thread started that termination but there is no guarantee that other + // threads have terminated. In other words, a static object may be accessed by another thread + // after it is deleted. We avoid destruction issues by never running the destructor. + alignas(Mutex) unsigned char mutex_[sizeof(Mutex)]; +}; + +Global& GetGlobal() { + static Global global; + static_assert(std::is_trivially_destructible_v); + return global; +} + +// return the length of a null-terminated string, or -1 if it is not terminated after kMaxNameLength +int strnlen(const char* s) { + for (int i = 0; i < kMaxNameLength; ++i) { + if (!s[i]) { + return i; + } + } + return -1; +} + +// copy a null-terminated string into a new heap-allocated char array managed by a unique_ptr +std::unique_ptr CopyName(const char* s) { + int len = strnlen(s); + if (len == -1) { + return nullptr; + } + std::unique_ptr out(new(std::nothrow) char[len + 1]); + if (!out) { + return nullptr; + } + std::strncpy(out.get(), s, len); + out.get()[len] = '\0'; + return out; +} + +// check if two plugins are identical +bool PluginsAreIdentical(const mjpPlugin& plugin1, const mjpPlugin& plugin2) { + if (plugin1.name && !plugin2.name) { + return false; + } + if (plugin2.name && !plugin1.name) { + return false; + } + if (plugin1.name && plugin2.name && + std::strncmp(plugin1.name, plugin2.name, kMaxNameLength)) { + return false; + } + + if (plugin1.nattribute != plugin2.nattribute) { + return false; + } + for (int i = 0; i < plugin1.nattribute; ++i) { + if (plugin1.attributes[i] && !plugin2.attributes[i]) { + return false; + } + if (plugin2.attributes[i] && !plugin1.attributes[i]) { + return false; + } + if (plugin1.attributes[i] && plugin2.attributes[i] && + std::strncmp(plugin1.attributes[i], plugin2.attributes[i], + kMaxNameLength)) { + return false; + } + } + + const char* ptr1 = reinterpret_cast(&plugin1.attributes) + + sizeof(plugin1.attributes); + const char* ptr2 = reinterpret_cast(&plugin2.attributes) + + sizeof(plugin2.attributes); + std::size_t remaining_size = + sizeof(mjpPlugin) - (ptr1 - reinterpret_cast(&plugin1)); + return !std::memcmp(ptr1, ptr2, remaining_size); +} +} // namespace + +// globally register a plugin (thread-safe), return new slot id +int mjp_registerPlugin(const mjpPlugin* plugin) { + if (!plugin->name) { + mju_error("plugin->name is a null pointer"); + } else if (plugin->name[0] == '\0') { + mju_error("plugin->name is an empty string"); + } else if (plugin->nattribute < 0) { + mju_error("plugin->nattribute is negative"); + } else if (plugin->nattribute > kMaxAttributes) { + mju_error_i("plugin->nattribute exceeds the maximum limit of ", + kMaxAttributes); + } + + char err[512]; + err[0] = '\0'; + + // ========= ATTENTION! ========================================================================== + // Do not handle objects with nontrivial destructors outside of this lambda. + // Do not call mju_error inside this lambda. + int slot = [&]() -> int { + // check and copy the plugin name + std::unique_ptr name = CopyName(plugin->name); + if (!name) { + if (strnlen(plugin->name) == -1) { + std::snprintf(err, sizeof(err), + "plugin->name length exceeds the maximum limit of %d", kMaxNameLength); + } else { + std::snprintf(err, sizeof(err), "failed to allocate memory for plugin name"); + } + return -1; + } + + // check and copy plugin attributes + std::vector> attributes_vec; + if (plugin->nattribute) { + attributes_vec.reserve(plugin->nattribute); + for (int i = 0; i < plugin->nattribute; ++i) { + std::unique_ptr attr = CopyName(plugin->attributes[i]); + if (!attr) { + if (strnlen(plugin->attributes[i]) == -1) { + std::snprintf( + err, sizeof(err), + "plugin->attributes[%d] exceeds the maximum limit of %d", i, kMaxAttributes); + } else { + std::snprintf(err, sizeof(err), "failed to allocate memory for plugin attribute"); + } + return -1; + } + attributes_vec.emplace_back(std::move(attr)); + } + } + + // exclusively lock the global plugin table + Global& global = GetGlobal(); + std::unique_lock lock(global.mutex()); + + int count = global.count().load(std::memory_order_acquire); + int local_idx = 0; + PluginTable* table = &global.table(); + + // check if a non-identical plugin with the same name has already been registered + for (int i = 0; i < count; ++i, ++local_idx) { + if (local_idx == PluginTable::kBlockSize) { + local_idx = 0; + table = table->next; + } + mjpPlugin& existing = table->plugins[local_idx]; + if (std::strcmp(plugin->name, existing.name) == 0) { + if (PluginsAreIdentical(*plugin, existing)) { + return i; + } else { + std::snprintf(err, sizeof(err), "plugin '%s' is already registered", plugin->name); + return -1; + } + } + } + + // allocate a new block of PluginTable if the last allocated block is full + if (local_idx == PluginTable::kBlockSize) { + local_idx = 0; +#if defined(MAC_OS_X_VERSION_MIN_REQUIRED) && MAC_OS_X_VERSION_MIN_REQUIRED < MAC_OS_X_VERSION_10_14 + // aligned nothrow new is not available until macOS 10.14 + posix_memalign(reinterpret_cast(&table->next), + alignof(PluginTable), sizeof(PluginTable)); + if (table->next) new(table->next) PluginTable; +#else + table->next = new(std::nothrow) PluginTable; +#endif + if (!table->next) { + std::snprintf(err, sizeof(err), "failed to allocate memory for the global plugin table"); + return -1; + } + table = table->next; + } + + // release the attribute names from unique_ptr into a plain array + const char** attributes = nullptr; + if (plugin->nattribute) { + attributes = new(std::nothrow) const char*[plugin->nattribute]; + if (!attributes) { + std::snprintf(err, sizeof(err), "failed to allocate memory for plugin attribute array"); + return -1; + } + for (int i = 0; i < plugin->nattribute; ++i) { + attributes[i] = attributes_vec[i].release(); + } + } + + // all checked passed, actually register the plugin into the global table + mjpPlugin& registered_plugin = table->plugins[local_idx]; + registered_plugin = *plugin; + registered_plugin.name = name.release(); + registered_plugin.attributes = attributes; + + // increment the global plugin count with a release memory barrier + global.count().store(count + 1, std::memory_order_release); + + return count; + }(); + + // ========= ATTENTION! ========================================================================== + // End of safe lambda, do not handle objects with non-trivial destructors beyond this point. + + // plugin registration failed, throw an mju_error + if (slot < 0) { + err[sizeof(err) - 1] = '\0'; + mju_error(err); + } + + return slot; +} + +// look up plugin by slot number, assuming that mjp_pluginCount has already been called +const mjpPlugin* mjp_getPluginAtSlotUnsafe(int slot, int nslot) { + if (slot < 0 || slot >= nslot) { + return nullptr; + } + + Global& global = GetGlobal(); + PluginTable* table = &global.table(); + + // iterate over blocks in the global table until the local index is less the block size + int local_idx = slot; + while (local_idx >= PluginTable::kBlockSize) { + local_idx -= PluginTable::kBlockSize; + table = table->next; + if (!table) { + return nullptr; + } + } + + // local_idx is now a valid index into the current block + const mjpPlugin& plugin = table->plugins[local_idx]; + if (!plugin.name) { + return nullptr; + } + return &plugin; +} + +// look up plugin by name, assuming that mjp_pluginCount has already been called +const mjpPlugin* mjp_getPluginUnsafe(const char* name, int* slot, int nslot) { + if (slot) *slot = -1; + + if (!name || !name[0]) { + return nullptr; + } + + Global& plugin = GetGlobal(); + PluginTable* table = &plugin.table(); + int found_slot = 0; + while (table) { + for (int i = 0; + i < PluginTable::kBlockSize && found_slot < nslot; + ++i, ++found_slot) { + const mjpPlugin& plugin = table->plugins[i]; + + // reached an uninitialized plugin, which means that iterated beyond the plugin count + // this should never happen if `count` was actually returned by mjp_pluginCount + if (!plugin.name) { + return nullptr; + } + if (std::strcmp(plugin.name, name) == 0) { + if (slot) *slot = found_slot; + return &plugin; + } + } + table = table->next; + } + + return nullptr; +} + +// return the number of globally registered plugins +int mjp_pluginCount() { + return GetGlobal().count().load(std::memory_order_acquire); +} + + +// look up a plugin by slot number +const mjpPlugin* mjp_getPluginAtSlot(int slot) { + const int count = mjp_pluginCount(); + + // mjp_pluginCount uses memory_order_acquire which acts as a barrier that guarantees that all + // plugins up to `count` have been completely inserted + return mjp_getPluginAtSlotUnsafe(slot, count); +} + +// look up a plugin by name, optionally also get its registered slot number +const mjpPlugin* mjp_getPlugin(const char* name, int* slot) { + const int count = mjp_pluginCount(); + int found_slot = -1; + const mjpPlugin* plugin = mjp_getPluginUnsafe(name, &found_slot, count); + + if (slot) *slot = found_slot; + return plugin; +} + +namespace { +// seek the nth config attrib of a plugin instance by counting null terminators +const char* PluginAttrSeek(const mjModel* m, int plugin_id, int attrib_id) { + const char* ptr = m->plugin_attr + m->plugin_attradr[plugin_id]; + for (int i = 0; i < attrib_id; ++i) { + while (*ptr) { + ++ptr; + } + ++ptr; + } + return ptr; +} +} // namespace + + +// return a config attribute of a plugin instance +// NULL: invalid plugin instance ID or attribute name +const char* mj_getPluginConfig(const mjModel* m, int plugin_id, const char* attrib) { + if (plugin_id < 0 || plugin_id >= m->nplugin || attrib == nullptr) { + return nullptr; + } + + const mjpPlugin* plugin = mjp_getPluginAtSlot(m->plugin[plugin_id]); + if (!plugin) { + return nullptr; + } + + for (int i = 0; i < plugin->nattribute; ++i) { + if (std::strcmp(plugin->attributes[i], attrib) == 0) { + return PluginAttrSeek(m, plugin_id, i); + } + } + + return nullptr; +} diff --git a/src/engine/engine_plugin.h b/src/engine/engine_plugin.h new file mode 100644 index 00000000..1b9dbd44 --- /dev/null +++ b/src/engine/engine_plugin.h @@ -0,0 +1,62 @@ +// Copyright 2022 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MUJOCO_SRC_ENGINE_ENGINE_PLUGIN_H_ +#define MUJOCO_SRC_ENGINE_ENGINE_PLUGIN_H_ + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// set default plugin definition +MJAPI void mjp_defaultPlugin(mjpPlugin* plugin); + +// globally register a plugin (thread-safe), return new slot id +MJAPI int mjp_registerPlugin(const mjpPlugin* plugin); + +// return the number of globally registered plugins +MJAPI int mjp_pluginCount(); + +// look up a plugin by name, optionally also get its registered slot number +MJAPI const mjpPlugin* mjp_getPlugin(const char* name, int* slot); + +// look up a plugin by slot number +MJAPI const mjpPlugin* mjp_getPluginAtSlot(int slot); + +// return a config attribute of a plugin instance +// NULL: invalid plugin instance ID or attribute name +MJAPI const char* mj_getPluginConfig(const mjModel* m, int plugin_id, const char* attrib); + +// ================================================================================================= +// MuJoCo-internal functions beyond this point. +// "Unsafe" suffix indicates that improper use of these functions may result in data races. +// +// The unsafe functions assume that called mjp_pluginCount has already been called, and that it is +// safe to assume that all plugins up to `count` have been completely written into the global table. +// ================================================================================================= + +// look up a plugin by name, assuming that mjp_pluginCount has already been called +const mjpPlugin* mjp_getPluginUnsafe(const char* name, int* slot, int nslot); + +// look up a plugin by slot number, assuming that mjp_pluginCount has already been called +const mjpPlugin* mjp_getPluginAtSlotUnsafe(int slot, int nslot); + +#ifdef __cplusplus +} +#endif + +#endif // MUJOCO_SRC_ENGINE_ENGINE_PLUGIN_H_ diff --git a/src/engine/engine_sensor.c b/src/engine/engine_sensor.c index 48dcc7b5..9efd8634 100644 --- a/src/engine/engine_sensor.c +++ b/src/engine/engine_sensor.c @@ -18,10 +18,12 @@ #include #include +#include #include "engine/engine_callback.h" #include "engine/engine_core_smooth.h" #include "engine/engine_io.h" #include "engine/engine_macro.h" +#include "engine/engine_plugin.h" #include "engine/engine_ray.h" #include "engine/engine_support.h" #include "engine/engine_util_blas.h" @@ -200,6 +202,11 @@ void mj_sensorPos(const mjModel* m, mjData* d) { // process sensors matching stage for (int i=0; insensor; i++) { + // skip sensor plugins -- these are handled after builtin sensor types + if (m->sensor_type[i] == mjSENS_PLUGIN) { + continue; + } + if (m->sensor_needstage[i]==mjSTAGE_POS) { // get sensor info objtype = m->sensor_objtype[i]; @@ -342,6 +349,25 @@ void mj_sensorPos(const mjModel* m, mjData* d) { add_noise(m, d, mjSTAGE_POS); } + // compute plugin sensor values + if (m->nplugin) { + const int nslot = mjp_pluginCount(); + 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->type & 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); + } + plugin->compute(m, d, i, mjPLUGIN_SENSOR); + } + } + } + // cutoff apply_cutoff(m, d, mjSTAGE_POS); } @@ -362,6 +388,11 @@ void mj_sensorVel(const mjModel* m, mjData* d) { // process sensors matching stage int subtreeVel = 0; for (int i=0; insensor; i++) { + // skip sensor plugins -- these are handled after builtin sensor types + if (m->sensor_type[i] == mjSENS_PLUGIN) { + continue; + } + if (m->sensor_needstage[i]==mjSTAGE_VEL) { // get sensor info type = m->sensor_type[i]; @@ -499,6 +530,32 @@ void mj_sensorVel(const mjModel* m, mjData* d) { add_noise(m, d, mjSTAGE_VEL); } + // trigger computation of plugins + if (m->nplugin) { + const int nslot = mjp_pluginCount(); + 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->type & mjPLUGIN_SENSOR) && plugin->needstage==mjSTAGE_VEL) { + if (!plugin->compute) { + mju_error_i("`compute` is null for plugin at slot %d", slot); + } + if (subtreeVel == 0) { + // compute subtree_linvel, subtree_angmom + // TODO(b/247107630): add a flag to allow plugin to specify whether it actually needs this + mj_subtreeVel(m, d); + + // mark computed + subtreeVel = 1; + } + plugin->compute(m, d, i, mjPLUGIN_SENSOR); + } + } + } + // cutoff apply_cutoff(m, d, mjSTAGE_VEL); } @@ -520,6 +577,11 @@ void mj_sensorAcc(const mjModel* m, mjData* d) { // process sensors matching stage int rnePost = 0; for (int i=0; insensor; i++) { + // skip sensor plugins -- these are handled after builtin sensor types + if (m->sensor_type[i] == mjSENS_PLUGIN) { + continue; + } + if (m->sensor_needstage[i]==mjSTAGE_ACC) { // get sensor info type = m->sensor_type[i]; @@ -677,6 +739,32 @@ void mj_sensorAcc(const mjModel* m, mjData* d) { add_noise(m, d, mjSTAGE_ACC); } + // trigger computation of plugins + if (m->nplugin) { + const int nslot = mjp_pluginCount(); + 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->type & mjPLUGIN_SENSOR) && plugin->needstage==mjSTAGE_ACC) { + if (!plugin->compute) { + mju_error_i("`compute` is null for plugin at slot %d", slot); + } + if (rnePost == 0) { + // compute cacc, cfrc_int, cfrc_ext + // TODO(b/247107630): add a flag to allow plugin to specify whether it actually needs this + mj_rnePostConstraint(m, d); + + // mark computed + rnePost = 1; + } + plugin->compute(m, d, i, mjPLUGIN_SENSOR); + } + } + } + // cutoff apply_cutoff(m, d, mjSTAGE_ACC); } diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c index 5bf9071e..9cf0ebf9 100644 --- a/src/engine/engine_support.c +++ b/src/engine/engine_support.c @@ -440,107 +440,90 @@ static int _getnumadr(const mjModel* m, mjtObj type, int** padr) { case mjOBJ_XBODY: *padr = m->name_bodyadr; return m->nbody; - break; case mjOBJ_JOINT: *padr = m->name_jntadr; return m->njnt; - break; case mjOBJ_GEOM: *padr = m->name_geomadr; return m->ngeom; - break; case mjOBJ_SITE: *padr = m->name_siteadr; return m->nsite; - break; case mjOBJ_CAMERA: *padr = m->name_camadr; return m->ncam; - break; case mjOBJ_LIGHT: *padr = m->name_lightadr; return m->nlight; - break; case mjOBJ_MESH: *padr = m->name_meshadr; return m->nmesh; - break; case mjOBJ_SKIN: *padr = m->name_skinadr; return m->nskin; - break; case mjOBJ_HFIELD: *padr = m->name_hfieldadr; return m->nhfield; - break; case mjOBJ_TEXTURE: *padr = m->name_texadr; return m->ntex; - break; case mjOBJ_MATERIAL: *padr = m->name_matadr; return m->nmat; - break; case mjOBJ_PAIR: *padr = m->name_pairadr; return m->npair; - break; case mjOBJ_EXCLUDE: *padr = m->name_excludeadr; return m->nexclude; - break; case mjOBJ_EQUALITY: *padr = m->name_eqadr; return m->neq; - break; case mjOBJ_TENDON: *padr = m->name_tendonadr; return m->ntendon; - break; case mjOBJ_ACTUATOR: *padr = m->name_actuatoradr; return m->nu; - break; case mjOBJ_SENSOR: *padr = m->name_sensoradr; return m->nsensor; - break; case mjOBJ_NUMERIC: *padr = m->name_numericadr; return m->nnumeric; - break; case mjOBJ_TEXT: *padr = m->name_textadr; return m->ntext; - break; case mjOBJ_TUPLE: *padr = m->name_tupleadr; return m->ntuple; - break; case mjOBJ_KEY: *padr = m->name_keyadr; return m->nkey; - break; + + case mjOBJ_PLUGIN: + *padr = m->name_pluginadr; + return m->nplugin; default: *padr = 0; diff --git a/src/engine/engine_util_misc.c b/src/engine/engine_util_misc.c index 31ea1be1..ecb728d2 100644 --- a/src/engine/engine_util_misc.c +++ b/src/engine/engine_util_misc.c @@ -831,6 +831,9 @@ const char* mju_type2Str(int type) { case mjOBJ_KEY: return "key"; + case mjOBJ_PLUGIN: + return "plugin"; + default: return 0; } @@ -932,6 +935,10 @@ int mju_str2Type(const char* str) { return mjOBJ_KEY; } + else if (!strcmp(str, "plugin")) { + return mjOBJ_PLUGIN; + } + else { return mjOBJ_UNKNOWN; } diff --git a/src/user/user_model.cc b/src/user/user_model.cc index e2db338a..7a875ee6 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -24,11 +24,13 @@ #include #include +#include #include #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; iobjtype.size(); + // npluginattr + for (i=0; iflattened_attributes.size(); + // nnames nnames = (int)modelname.size() + 1; for (i=0; iname.length() + 1; @@ -978,6 +1001,7 @@ void mjCModel::SetSizes(void) { for (i=0; iname.length() + 1; for (i=0; iname.length() + 1; for (i=0; iname.length() + 1; + for (i=0; iname.length() + 1; // nemax for (i=0; iname_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; iCompile(); for (int i=0; iCompile(); for (int i=0; iCompile(); + for (int i=0; iCompile(); // compile defaults: to enforce userdata length for writer for (int i=0; iplugin[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> 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(plugin->needstage); + this->nsensordata += nsensordata; + } + } + } + m->npluginstate = stateadr; + } + // keyframe compilation needs access to nq, nv, na, nmocap, qpos0 for (int i=0; iCompile(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 ) + 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(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; + } +} diff --git a/src/user/user_model.h b/src/user/user_model.h index f73cb685..bf976e60 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -16,10 +16,12 @@ #define MUJOCO_SRC_USER_USER_MODEL_H_ #include +#include #include #include #include +#include #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 tuples; // list of tuple fields std::vector keys; // list of keyframe fields + std::vector> active_plugins; // list of active plugins + std::vector plugins; // list of plugin instances + // pointers to objects created inside kinematic tree std::vector bodies; // list of bodies std::vector joints; // list of joints allowing motion relative to parent @@ -251,6 +265,7 @@ class mjCModel { std::vector 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 diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index af7c2d4c..5e0b7767 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -25,12 +25,14 @@ #include "lodepng.h" #include +#include #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(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()); + } +} diff --git a/src/user/user_objects.h b/src/user/user_objects.h index 0b457960..35d62796 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -15,6 +15,7 @@ #ifndef MUJOCO_SRC_USER_USER_OBJECTS_H_ #define MUJOCO_SRC_USER_USER_OBJECTS_H_ +#include #include #include @@ -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> config_attribs; // raw config attributes from XML + std::vector 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 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 diff --git a/src/xml/xml_native_reader.cc b/src/xml/xml_native_reader.cc index 20665c04..ce48762c 100644 --- a/src/xml/xml_native_reader.cc +++ b/src/xml/xml_native_reader.cc @@ -17,32 +17,65 @@ #include #include #include +#include +#include +#include #include #include +#include +#include #include #include #include #include "engine/engine_macro.h" +#include "engine/engine_plugin.h" #include "engine/engine_util_errmem.h" #include "engine/engine_util_misc.h" #include "user/user_composite.h" #include "user/user_model.h" +#include "user/user_objects.h" #include "user/user_util.h" +#include "xml/xml_util.h" #include "tinyxml2.h" namespace { - using std::string; using std::vector; using tinyxml2::XMLElement; +void ReadPluginConfigs(tinyxml2::XMLElement* elem, mjCPlugin* pp) { + std::map> config_attribs; + XMLElement* child = elem->FirstChildElement(); + while (child) { + std::string_view name = child->Value(); + if (name == "config") { + std::string key, value; + mjXUtil::ReadAttrTxt(child, "key", key, /* required = */ true); + if (config_attribs.find(key) != config_attribs.end()) { + std::string err = "duplicate config key: " + key; + throw mjXError(child, err.c_str()); + } + mjXUtil::ReadAttrTxt(child, "value", value, /* required = */ true); + config_attribs[key] = value; + } + child = child->NextSiblingElement(); + } + + if (!pp && !config_attribs.empty()) { + throw mjXError(elem, + "plugin configuration attributes cannot be used in an " + "element that references a predefined plugin instance"); + } else if (pp) { + pp->config_attribs = std::move(config_attribs); + } +} } // namespace //---------------------------------- MJCF schema --------------------------------------------------- -static const int nMJCF = 165; +static const int nMJCF = 183; static const char* MJCF[nMJCF][mjXATTRNUM] = { {"mujoco", "!", "1", "model"}, {"<"}, @@ -150,6 +183,17 @@ static const char* MJCF[nMJCF][mjXATTRNUM] = { "gain", "user", "group"}, {">"}, + {"extension", "*", "0"}, + {"<"}, + {"required", "*", "1", "plugin"}, + {"<"}, + {"instance", "*", "1", "name"}, + {"<"}, + {"config", "*", "2", "key", "value"}, + {">"}, + {">"}, + {">"}, + {"custom", "*", "0"}, {"<"}, {"numeric", "*", "3", "name", "size", "data"}, @@ -308,6 +352,13 @@ static const char* MJCF[nMJCF][mjXATTRNUM] = { "lmin", "lmax", "vmax", "fpmax", "fvmax"}, {"adhesion", "*", "9", "name", "class", "group", "forcelimited", "ctrlrange", "forcerange", "user", "body", "gain"}, + {"plugin", "*", "19", "name", "class", "plugin", "instance", "group", + "ctrllimited", "forcelimited", "ctrlrange", "forcerange", + "lengthrange", "gear", "cranklength", "joint", "jointinparent", + "site", "tendon", "cranksite", "slidersite", "user"}, + {"<"}, + {"config", "*", "2", "key", "value"}, + {">"}, {">"}, {"sensor", "*", "0"}, @@ -350,6 +401,11 @@ static const char* MJCF[nMJCF][mjXATTRNUM] = { {"clock", "*", "4", "name", "cutoff", "noise", "user"}, {"user", "*", "9", "name", "objtype", "objname", "datatype", "needstage", "dim", "cutoff", "noise", "user"}, + {"plugin", "*", "9", "name", "plugin", "instance", "cutoff", "objtype", "objname", "reftype", "refname", + "user"}, + {"<"}, + {"config", "*", "2", "key", "value"}, + {">"}, {">"}, {"keyframe", "*", "0"}, @@ -711,6 +767,11 @@ void mjXReader::Parse(XMLElement* root) { } readingdefaults = false; + for (section = root->FirstChildElement("extension"); section; + section = section->NextSiblingElement("extension")) { + Extension(section); + } + for (section = root->FirstChildElement("custom"); section; section = section->NextSiblingElement("custom")) { Custom(section); @@ -1627,6 +1688,18 @@ void mjXReader::OneActuator(XMLElement* elem, mjCActuator* pact) { pact->biastype = mjBIAS_NONE; } + else if (type == "plugin") { + pact->is_plugin = true; + ReadAttrTxt(elem, "plugin", pact->plugin_name); + ReadAttrTxt(elem, "instance", pact->plugin_instance_name); + if (pact->plugin_instance_name.empty()) { + pact->plugin_instance = model->AddPlugin(); + } else { + model->hasImplicitPluginElem = true; + } + ReadPluginConfigs(elem, pact->plugin_instance); + } + else { // SHOULD NOT OCCUR throw mjXError(elem, "unrecognized actuator type: %s", type.c_str()); } @@ -1908,6 +1981,61 @@ void mjXReader::Default(XMLElement* section, int parentid) { +// extension section parser +void mjXReader::Extension(XMLElement* section) { + XMLElement* elem = section->FirstChildElement(); + while (elem) { + // get sub-element name + std::string_view name = elem->Value(); + + if (name == "required") { + std::string plugin_name; + int plugin_slot = -1; + ReadAttrTxt(elem, "plugin", plugin_name, /* required = */ true); + const mjpPlugin* plugin = mjp_getPlugin(plugin_name.c_str(), &plugin_slot); + if (!plugin) { + throw mjXError(elem, "unknown plugin '%s'", plugin_name.c_str()); + } + + bool already_declared = false; + for (const auto& [existing_plugin, existing_slot] : model->active_plugins) { + if (plugin == existing_plugin) { + already_declared = true; + break; + } + } + if (!already_declared) { + model->active_plugins.emplace_back(std::make_pair(plugin, plugin_slot)); + } + + XMLElement* child = elem->FirstChildElement(); + while (child) { + if (std::string(child->Value())=="instance") { + if (model->hasImplicitPluginElem) { + throw mjXError( + child, "explicit plugin instance must appear before implicit plugin elements"); + } + mjCPlugin* pp = model->AddPlugin(); + GetXMLPos(child, pp); + ReadAttrTxt(child, "name", pp->name, /* required = */ true); + if (pp->name.empty()) { + throw mjXError(child, "plugin instance must have a name"); + } + ReadPluginConfigs(child, pp); + pp->plugin_slot = plugin_slot; + pp->nstate = -1; // actual value to be filled in by the plugin later + } + child = child->NextSiblingElement(); + } + } + + // advance to next element + elem = elem->NextSiblingElement(); + } +} + + + // custom section parser void mjXReader::Custom(XMLElement* section) { string text, name; @@ -2800,6 +2928,38 @@ void mjXReader::Sensor(XMLElement* section) { psen->datatype = (mjtDataType)n; } + else if (type=="plugin") { + psen->type = mjSENS_PLUGIN; + ReadAttrTxt(elem, "plugin", psen->plugin_name); + ReadAttrTxt(elem, "instance", psen->plugin_instance_name); + if (psen->plugin_instance_name.empty()) { + psen->plugin_instance = model->AddPlugin(); + } else { + model->hasImplicitPluginElem = true; + } + ReadPluginConfigs(elem, psen->plugin_instance); + ReadAttrTxt(elem, "objtype", text); + psen->objtype = (mjtObj)mju_str2Type(text.c_str()); + ReadAttrTxt(elem, "objname", psen->objname); + if (psen->objtype != mjOBJ_UNKNOWN && psen->objname.empty()) { + throw mjXError(elem, "objtype is specified but objname is not"); + } + if (psen->objtype == mjOBJ_UNKNOWN && !psen->objname.empty()) { + throw mjXError(elem, "objname is specified but objtype is not"); + } + ReadAttrTxt(elem, "reftype", text); + psen->reftype = (mjtObj)mju_str2Type(text.c_str()); + ReadAttrTxt(elem, "refname", psen->refname); + if (psen->reftype != mjOBJ_UNKNOWN && psen->refname.empty()) { + throw mjXError(elem, "reftype is specified but refname is not"); + } + if (psen->reftype == mjOBJ_UNKNOWN && !psen->refname.empty()) { + throw mjXError(elem, "refname is specified but reftype is not"); + } + } + + GetXMLPos(elem, psen); + // advance to next element elem = elem->NextSiblingElement(); } diff --git a/src/xml/xml_native_reader.h b/src/xml/xml_native_reader.h index cfd68f69..849649b1 100644 --- a/src/xml/xml_native_reader.h +++ b/src/xml/xml_native_reader.h @@ -37,6 +37,7 @@ class mjXReader : public mjXBase { private: // XML section specific to MJCF void Default(tinyxml2::XMLElement* section, int parentid); // default section + void Extension(tinyxml2::XMLElement* section); // extension section void Custom(tinyxml2::XMLElement* section); // custom section void Visual(tinyxml2::XMLElement* section); // visual section void Statistic(tinyxml2::XMLElement* section); // statistic section diff --git a/src/xml/xml_native_writer.cc b/src/xml/xml_native_writer.cc index 0d4f8a2e..e9db89d7 100644 --- a/src/xml/xml_native_writer.cc +++ b/src/xml/xml_native_writer.cc @@ -18,10 +18,15 @@ #include #include #include +#include +#include +#include #include "engine/engine_io.h" +#include "engine/engine_plugin.h" #include "engine/engine_util_errmem.h" #include "engine/engine_util_misc.h" +#include "user/user_objects.h" #include "user/user_util.h" #include "xml/xml_util.h" #include "tinyxml2.h" @@ -604,12 +609,38 @@ void mjXWriter::OneActuator(XMLElement* elem, mjCActuator* pact, mjCDef* def) { WriteAttr(elem, "lengthrange", 2, pact->lengthrange, def->actuator.lengthrange); WriteAttr(elem, "gear", 6, pact->gear, def->actuator.gear); WriteAttr(elem, "cranklength", 1, &pact->cranklength, &def->actuator.cranklength); - WriteAttrKey(elem, "dyntype", dyn_map, dyn_sz, pact->dyntype, def->actuator.dyntype); - WriteAttrKey(elem, "gaintype", gain_map, gain_sz, pact->gaintype, def->actuator.gaintype); - WriteAttrKey(elem, "biastype", bias_map, bias_sz, pact->biastype, def->actuator.biastype); - WriteAttr(elem, "dynprm", mjNDYN, pact->dynprm, def->actuator.dynprm); - WriteAttr(elem, "gainprm", mjNGAIN, pact->gainprm, def->actuator.gainprm); - WriteAttr(elem, "biasprm", mjNBIAS, pact->biasprm, def->actuator.biasprm); + + // plugins: write config attributes + if (pact->is_plugin) { + if (!pact->plugin_instance_name.empty()) { + WriteAttrTxt(elem, "instance", pact->plugin_instance_name); + } else { + WriteAttrTxt(elem, "plugin", pact->plugin_name); + const mjpPlugin* plugin = mjp_getPluginAtSlot( + pact->plugin_instance->plugin_slot); + const char* c = &pact->plugin_instance->flattened_attributes[0]; + for (int i = 0; i < plugin->nattribute; ++i) { + std::string value(c); + if (!value.empty()) { + XMLElement* config_elem = InsertEnd(elem, "config"); + WriteAttrTxt(config_elem, "key", plugin->attributes[i]); + WriteAttrTxt(config_elem, "value", value); + c += value.size(); + } + ++c; + } + } + } + + // non-plugins: write actuator parameters + else { + WriteAttrKey(elem, "dyntype", dyn_map, dyn_sz, pact->dyntype, def->actuator.dyntype); + WriteAttrKey(elem, "gaintype", gain_map, gain_sz, pact->gaintype, def->actuator.gaintype); + WriteAttrKey(elem, "biastype", bias_map, bias_sz, pact->biastype, def->actuator.biastype); + WriteAttr(elem, "dynprm", mjNDYN, pact->dynprm, def->actuator.dynprm); + WriteAttr(elem, "gainprm", mjNGAIN, pact->gainprm, def->actuator.gainprm); + WriteAttr(elem, "biasprm", mjNBIAS, pact->biasprm, def->actuator.biasprm); + } // userdata if (writingdefaults) { @@ -659,6 +690,7 @@ void mjXWriter::Write(FILE* fp) { writingdefaults = true; Default(root, model->defaults[0]); writingdefaults = false; + Extension(root); Custom(root); Asset(root); Body(InsertEnd(root, "worldbody"), model->GetWorld()); @@ -1020,6 +1052,69 @@ void mjXWriter::Default(XMLElement* root, mjCDef* def) { +// extension section +void mjXWriter::Extension(XMLElement* root) { + // skip section if there is no required plugin + if (model->active_plugins.empty()) { + return; + } + + // create section + XMLElement* section = InsertEnd(root, "extension"); + + // keep track of plugins whose section have been created + std::unordered_set seen_plugins; + + // write all plugins + const mjpPlugin* last_plugin = nullptr; + XMLElement* required_elem = nullptr; + for (int i = 0; i < model->plugins.size(); ++i) { + mjCPlugin* pp = static_cast(model->GetObject(mjOBJ_PLUGIN, i)); + + if (pp->name.empty()) { + // reached the first unnamed plugin instance, meaning that it was created through an + // "implicit" plugin element, e.g. sensor or actuator + break; + } + + // check if we need to open a new section + const mjpPlugin* plugin = mjp_getPluginAtSlot(pp->plugin_slot); + if (plugin != last_plugin) { + required_elem = InsertEnd(section, "required"); + WriteAttrTxt(required_elem, "plugin", plugin->name); + seen_plugins.insert(plugin); + last_plugin = plugin; + } + + // write instance element + XMLElement* elem = InsertEnd(required_elem, "instance"); + WriteAttrTxt(elem, "name", pp->name); + + // write plugin config attributes + const char* c = &pp->flattened_attributes[0]; + for (int i = 0; i < plugin->nattribute; ++i) { + std::string value(c); + if (!value.empty()) { + XMLElement* config_elem = InsertEnd(elem, "config"); + WriteAttrTxt(config_elem, "key", plugin->attributes[i]); + WriteAttrTxt(config_elem, "value", value); + c += value.size(); + } + ++c; + } + } + + // write elements for plugins without explicit instances + for (const auto& [plugin, slot] : model->active_plugins) { + if (seen_plugins.find(plugin) == seen_plugins.end()) { + required_elem = InsertEnd(section, "required"); + WriteAttrTxt(required_elem, "plugin", plugin->name); + } + } +} + + + // custom section void mjXWriter::Custom(XMLElement* root) { XMLElement* elem; @@ -1400,7 +1495,12 @@ void mjXWriter::Actuator(XMLElement* root) { // write all actuators for (int i=0; iGetObject(mjOBJ_ACTUATOR, i); - XMLElement* elem = InsertEnd(section, "general"); + XMLElement* elem; + if (pact->is_plugin) { + elem = InsertEnd(section, "plugin"); + } else { + elem = InsertEnd(section, "general"); + } OneActuator(elem, pact, pact->def); } } @@ -1593,6 +1693,34 @@ void mjXWriter::Sensor(XMLElement* root) { elem = InsertEnd(section, "clock"); break; + + // plugin-controlled sensor + case mjSENS_PLUGIN: + elem = InsertEnd(section, "plugin"); + if (psen->objtype != mjOBJ_UNKNOWN) { + WriteAttrTxt(elem, "objtype", mju_type2Str(psen->objtype)); + WriteAttrTxt(elem, "objname", psen->objname); + } + if (!psen->plugin_instance_name.empty()) { + WriteAttrTxt(elem, "instance", psen->plugin_instance_name); + } else { + WriteAttrTxt(elem, "plugin", psen->plugin_name); + const mjpPlugin* plugin = mjp_getPluginAtSlot( + psen->plugin_instance->plugin_slot); + const char* c = &psen->plugin_instance->flattened_attributes[0]; + for (int i = 0; i < plugin->nattribute; ++i) { + std::string value(c); + if (!value.empty()) { + XMLElement* config_elem = InsertEnd(elem, "config"); + WriteAttrTxt(config_elem, "key", plugin->attributes[i]); + WriteAttrTxt(config_elem, "value", value); + c += value.size(); + } + ++c; + } + } + break; + // user-defined sensor case mjSENS_USER: elem = InsertEnd(section, "user"); @@ -1610,11 +1738,13 @@ void mjXWriter::Sensor(XMLElement* root) { // write name, noise, userdata WriteAttrTxt(elem, "name", psen->name); WriteAttr(elem, "cutoff", 1, &psen->cutoff, &zero); - WriteAttr(elem, "noise", 1, &psen->noise, &zero); + if (psen->type != mjSENS_PLUGIN) { + WriteAttr(elem, "noise", 1, &psen->noise, &zero); + } WriteVector(elem, "user", psen->userdata); // add reference if present - if (psen->reftype > 0) { + if (psen->reftype != mjOBJ_UNKNOWN) { WriteAttrTxt(elem, "reftype", mju_type2Str(psen->reftype)); WriteAttrTxt(elem, "refname", psen->refname); } diff --git a/src/xml/xml_native_writer.h b/src/xml/xml_native_writer.h index 7e349a6f..3cc6c492 100644 --- a/src/xml/xml_native_writer.h +++ b/src/xml/xml_native_writer.h @@ -37,6 +37,7 @@ class mjXWriter : public mjXBase { void Visual(tinyxml2::XMLElement* root); // visual section void Statistic(tinyxml2::XMLElement* root); // statistic section void Default(tinyxml2::XMLElement* root, mjCDef* def); // default section + void Extension(tinyxml2::XMLElement* root); // extension section void Custom(tinyxml2::XMLElement* root); // custom section void Asset(tinyxml2::XMLElement* root); // asset section void Body(tinyxml2::XMLElement* elem, mjCBody* body); // body/world section diff --git a/test/engine/CMakeLists.txt b/test/engine/CMakeLists.txt index 2cc20693..c536572e 100644 --- a/test/engine/CMakeLists.txt +++ b/test/engine/CMakeLists.txt @@ -38,6 +38,14 @@ target_link_libraries( absl::str_format ) +mujoco_test(engine_plugin_test) +target_link_libraries( + engine_plugin_test + fixture + gmock + absl::str_format +) + mujoco_test(engine_print_test) target_link_libraries(engine_print_test fixture gmock) diff --git a/test/engine/engine_plugin_test.cc b/test/engine/engine_plugin_test.cc new file mode 100644 index 00000000..2b8a870a --- /dev/null +++ b/test/engine/engine_plugin_test.cc @@ -0,0 +1,524 @@ +// Copyright 2022 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Tests for plugin-related functionalities. +// TODO(b/247110452) add more comments to this file, or add sample plugins + +#include "src/engine/engine_plugin.h" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include "test/fixture.h" + +namespace mujoco { +namespace { +using PluginTest = MujocoTest; +using ::testing::HasSubstr; +using ::testing::NotNull; + +constexpr int kNumFakePlugins = 30; + +class BaseTestPlugin { + public: + static constexpr int kDefaultStride = 1; + + BaseTestPlugin(const mjModel* m, mjData* d, int instance) + : reset_count(d->plugin_state[m->plugin_stateadr[instance]]), + compute_count(d->plugin_state[m->plugin_stateadr[instance] + 1]), + advance_count(d->plugin_state[m->plugin_stateadr[instance] + 2]) { + { + const char* s = mj_getPluginConfig(m, instance, "stride"); + if (*s) { + std::stringstream(s) >> stride; + } else { + stride = kDefaultStride; + } + } + + reset_count = 0; + compute_count = 0; + advance_count = 0; + } + + void Reset() { + reset_count += stride; + compute_count = 0; + advance_count = 0; + } + + void Compute() { + compute_count += stride; + } + + void Advance() { + advance_count += stride; + } + + protected: + int stride; + + mjtNum& reset_count; + mjtNum& compute_count; + mjtNum& advance_count; +}; + +class TestSensor : public BaseTestPlugin { + public: + TestSensor(const mjModel* m, mjData* d, int instance) + : BaseTestPlugin(m, d, instance) { + for (int i = 0; i < m->nsensor; ++i) { + if (m->sensor_type[i] == mjSENS_PLUGIN && + m->sensor_plugin[i] == instance) { + sensors.push_back( + reinterpret_cast(&d->sensordata[m->sensor_adr[i]])); + } + } + } + + static int& InitCount() { + static int counter = 0; + return counter; + } + + static int& DestroyCount() { + static int counter = 0; + return counter; + } + + void Reset() { + BaseTestPlugin::Reset(); + WriteSensorData(); + } + + void Compute() { + BaseTestPlugin::Compute(); + WriteSensorData(); + } + + void Advance() { + BaseTestPlugin::Advance(); + WriteSensorData(); + } + + private: + std::vector sensors; + + void WriteSensorData() { + for (auto* sensordata_ptr : sensors) { + auto& sensordata = *sensordata_ptr; + sensordata[0] = reset_count; + sensordata[1] = compute_count; + sensordata[2] = advance_count; + } + } +}; + +class TestActuator : public BaseTestPlugin { + public: + static constexpr mjtNum kDefaultMultiplier = 1.0; + + TestActuator(const mjModel* m, mjData* d, int instance) + : BaseTestPlugin(m, d, instance) { + const char* s = mj_getPluginConfig(m, instance, "multiplier"); + if (*s) { + std::stringstream(s) >> multiplier; + } else { + multiplier = kDefaultMultiplier; + } + for (int i = 0; i < m->nu; ++i) { + if (m->actuator_plugin[i] == instance) { + actuators.push_back(&d->actuator_force[i]); + } + } + } + + static int& InitCount() { + static int counter = 0; + return counter; + } + + static int& DestroyCount() { + static int counter = 0; + return counter; + } + + void Reset() { + BaseTestPlugin::Reset(); + WriteActuatorForce(); + } + + void Compute() { + BaseTestPlugin::Compute(); + WriteActuatorForce(); + } + + void Advance() { + BaseTestPlugin::Advance(); + WriteActuatorForce(); + } + + private: + mjtNum multiplier; + std::vector actuators; + + void WriteActuatorForce() { + for (mjtNum* actuator_force : actuators) { + *actuator_force = advance_count * multiplier; + } + } +}; + +int RegisterSensorPlugin() { + mjpPlugin plugin; + mjp_defaultPlugin(&plugin); + + plugin.name = "mujoco.test.sensor"; + + const char* attributes[] = {"stride"}; + plugin.nattribute = sizeof(attributes) / sizeof(*attributes); + plugin.attributes = attributes; + + plugin.type |= mjPLUGIN_SENSOR; + + plugin.nstate = +[](const mjModel* m, int instance) { return 3; }; + plugin.nsensordata = + +[](const mjModel* m, int instance, int sensor_id) { return 3; }; + + plugin.init = +[](const mjModel* m, mjData* d, int instance) { + auto* sensor = new TestSensor(m, d, instance); + d->plugin_data[instance] = reinterpret_cast(sensor); + TestSensor::InitCount()++; + }; + plugin.destroy = +[](mjData* d, int instance) { + delete reinterpret_cast(d->plugin_data[instance]); + d->plugin_data[instance] = 0; + TestSensor::DestroyCount()++; + }; + + plugin.reset = +[](const mjModel* m, mjData* d, int instance) { + auto sensor = reinterpret_cast(d->plugin_data[instance]); + sensor->Reset(); + }; + plugin.compute = +[](const mjModel* m, mjData* d, int instance, int type) { + auto sensor = reinterpret_cast(d->plugin_data[instance]); + sensor->Compute(); + }; + plugin.advance = +[](const mjModel* m, mjData* d, int instance) { + auto sensor = reinterpret_cast(d->plugin_data[instance]); + sensor->Advance(); + }; + + return mjp_registerPlugin(&plugin); +} + +int RegisterActuatorPlugin() { + mjpPlugin plugin; + mjp_defaultPlugin(&plugin); + + plugin.name = "mujoco.test.actuator"; + + const char* attributes[] = {"stride", "multiplier"}; + plugin.nattribute = sizeof(attributes) / sizeof(*attributes); + plugin.attributes = attributes; + + plugin.type |= mjPLUGIN_ACTUATOR; + + plugin.nstate = +[](const mjModel* m, int instance) { return 3; }; + + plugin.init = +[](const mjModel* m, mjData* d, int instance) { + auto* actuator = new TestActuator(m, d, instance); + d->plugin_data[instance] = reinterpret_cast(actuator); + TestActuator::InitCount()++; + }; + plugin.destroy = +[](mjData* d, int instance) { + delete reinterpret_cast(d->plugin_data[instance]); + d->plugin_data[instance] = 0; + TestActuator::DestroyCount()++; + }; + + plugin.reset = +[](const mjModel* m, mjData* d, int instance) { + auto actuator = reinterpret_cast(d->plugin_data[instance]); + actuator->Reset(); + }; + plugin.compute = +[](const mjModel* m, mjData* d, int instance, int type) { + auto actuator = reinterpret_cast(d->plugin_data[instance]); + actuator->Compute(); + }; + plugin.advance = +[](const mjModel* m, mjData* d, int instance) { + auto actuator = reinterpret_cast(d->plugin_data[instance]); + actuator->Advance(); + }; + + return mjp_registerPlugin(&plugin); +} + +__attribute__((constructor)) void RegisterAllPlugins() { + RegisterSensorPlugin(); + + for (int i = 1; i <= kNumFakePlugins; ++i) { + mjpPlugin plugin; + mjp_defaultPlugin(&plugin); + std::string name = absl::StrFormat("mujoco.test.fake%u", i); + plugin.name = name.c_str(); + mjp_registerPlugin(&plugin); + } + + RegisterActuatorPlugin(); +} + +constexpr char xml[] = R"( + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +)"; + +TEST_F(PluginTest, MultiplePluginTableBlocks) { + EXPECT_EQ(mjp_pluginCount(), kNumFakePlugins + 2); + + const mjpPlugin* last_plugin = nullptr; + int table_count = 0; + for (int i = 1; i <= kNumFakePlugins; ++i) { + int slot; + std::string name = absl::StrFormat("mujoco.test.fake%u", i); + const mjpPlugin* plugin = mjp_getPlugin(name.c_str(), &slot); + EXPECT_EQ(slot, i); + EXPECT_THAT(plugin, NotNull()); + if (plugin - last_plugin != 1) { + ++table_count; + } + last_plugin = plugin; + } + + // Make sure that there are enough fake plugins to fill multiple table blocks. + EXPECT_GT(table_count, 1); + + // Make sure that a block contains multiple plugins. + EXPECT_LT(table_count, kNumFakePlugins); +} + +TEST_F(PluginTest, RegisterIdenticalPlugin) { + EXPECT_EQ(RegisterSensorPlugin(), 0); + EXPECT_EQ(RegisterActuatorPlugin(), kNumFakePlugins + 1); + EXPECT_EQ(mjp_pluginCount(), kNumFakePlugins + 2); +} + +TEST_F(PluginTest, SaveXml) { + char error[1024] = {0}; + + mjModel* m = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(m, testing::NotNull()) << error; + + std::string saved_xml = SaveAndReadXml(m); + std::string_view expected_xml(xml); + + const std::string extension_open = ""; + const std::string extension_close = ""; + int extension_start = expected_xml.find(extension_open); + ASSERT_NE(extension_start, std::string::npos); + int extension_end = + expected_xml.find(extension_close) + extension_close.size(); + ASSERT_NE(extension_end, std::string::npos); + ASSERT_LE(extension_start, extension_end); + auto expected_extension_section = + expected_xml.substr(extension_start, extension_end - extension_start); + + const std::string sensor_open = ""; + const std::string sensor_close = ""; + int sensor_start = expected_xml.find(sensor_open); + ASSERT_NE(sensor_start, std::string::npos); + int sensor_end = expected_xml.find(sensor_close) + sensor_close.size(); + ASSERT_NE(sensor_end, std::string::npos); + ASSERT_LE(sensor_start, sensor_end); + auto expected_sensor_section = + expected_xml.substr(sensor_start, sensor_end - sensor_start); + + const std::string actuator_open = ""; + const std::string actuator_close = ""; + int actuator_start = expected_xml.find(actuator_open); + ASSERT_NE(actuator_start, std::string::npos); + int actuator_end = expected_xml.find(actuator_close) + actuator_close.size(); + ASSERT_NE(actuator_end, std::string::npos); + ASSERT_LE(actuator_start, actuator_end); + auto expected_actuator_section = + expected_xml.substr(actuator_start, actuator_end - actuator_start); + + EXPECT_THAT(saved_xml, HasSubstr(expected_extension_section)); + EXPECT_THAT(saved_xml, HasSubstr(expected_sensor_section)); + EXPECT_THAT(saved_xml, HasSubstr(expected_actuator_section)); + + mj_deleteModel(m); + + // make sure that the saved XML can still be compiled + mjModel* m2 = LoadModelFromString(saved_xml, error, sizeof(error)); + ASSERT_THAT(m, testing::NotNull()) << error; + mj_deleteModel(m2); +} + +TEST_F(PluginTest, SensorPlugin) { + int expected_init_count = TestSensor::InitCount(); + int expected_destroy_count = TestSensor::DestroyCount(); + EXPECT_EQ(expected_init_count, expected_destroy_count); + + char error[1024] = {0}; + mjModel* m = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(m, testing::NotNull()) << error; + + // mj_makeModel calls mj_makeData and mj_deleteData internally + expected_init_count += 3; + expected_destroy_count += 3; + EXPECT_EQ(TestSensor::InitCount(), expected_init_count); + EXPECT_EQ(TestSensor::DestroyCount(), expected_destroy_count); + + EXPECT_EQ(m->nplugin, 5); + EXPECT_EQ(mj_name2id(m, mjOBJ_PLUGIN, "twosensors"), 0); + EXPECT_EQ(mj_name2id(m, mjOBJ_PLUGIN, "threesensors"), 1); + + mjData* d = mj_makeData(m); + expected_init_count += 3; + EXPECT_EQ(TestSensor::InitCount(), expected_init_count); + EXPECT_EQ(TestSensor::DestroyCount(), expected_destroy_count); + + for (int i = 0; i < 5; ++i) { + for (int j = 0; j < 10; ++j) { + EXPECT_THAT(*reinterpret_cast(d->plugin_state + + m->plugin_stateadr[0]), + testing::ElementsAreArray({i+1, 2*j, j})); + EXPECT_THAT(*reinterpret_cast(d->plugin_state + + m->plugin_stateadr[1]), + testing::ElementsAreArray({3*(i+1), 6*j, 3*j})); + EXPECT_THAT(*reinterpret_cast(d->plugin_state + + m->plugin_stateadr[4]), + testing::ElementsAreArray({5*(i+1), 10*j, 5*j})); + EXPECT_THAT(*reinterpret_cast(d->sensordata), + testing::ElementsAreArray({ i+1, 2*j, j, + 5*(i+1), 10*j, 5*j, + 3*(i+1), 6*j, 3*j, + i+1, 2*j, j, + 3*(i+1), 6*j, 3*j, + 3*(i+1), 6*j, 3*j})); + mj_step(m, d); + mj_forward(m, d); + } + mj_resetData(m, d); + } + + mj_deleteData(d); + expected_destroy_count += 3; + EXPECT_EQ(TestSensor::InitCount(), expected_init_count); + EXPECT_EQ(TestSensor::DestroyCount(), expected_destroy_count); + + mj_deleteModel(m); + EXPECT_EQ(TestSensor::InitCount(), expected_init_count); + EXPECT_EQ(TestSensor::DestroyCount(), expected_destroy_count); +} + +TEST_F(PluginTest, ActuatorPlugin) { + int expected_init_count = TestActuator::InitCount(); + int expected_destroy_count = TestActuator::DestroyCount(); + EXPECT_EQ(expected_init_count, expected_destroy_count); + + char error[1024] = {0}; + mjModel* m = LoadModelFromString(xml, error, sizeof(error)); + ASSERT_THAT(m, testing::NotNull()) << error; + + // mj_makeModel calls mj_makeData and mj_deleteData internally + expected_init_count += 2; + expected_destroy_count += 2; + EXPECT_EQ(TestActuator::InitCount(), expected_init_count); + EXPECT_EQ(TestActuator::DestroyCount(), expected_destroy_count); + + EXPECT_EQ(m->nplugin, 5); + EXPECT_EQ(mj_name2id(m, mjOBJ_PLUGIN, "actuator2"), 2); + + mjData* d = mj_makeData(m); + expected_init_count += 2; + EXPECT_EQ(TestActuator::InitCount(), expected_init_count); + EXPECT_EQ(TestActuator::DestroyCount(), expected_destroy_count); + + for (int i = 0; i < 5; ++i) { + for (int j = 0; j < 10; ++j) { + EXPECT_THAT(*reinterpret_cast(d->plugin_state + + m->plugin_stateadr[2]), + testing::ElementsAreArray({2*(i+1), 4*j, 2*j})); + EXPECT_THAT(*reinterpret_cast(d->plugin_state + + m->plugin_stateadr[3]), + testing::ElementsAreArray({4*(i+1), 8*j, 4*j})); + EXPECT_THAT(*reinterpret_cast(d->actuator_force), + testing::ElementsAreArray({0.125*j, 0.25*j})); + mj_step(m, d); + mj_forward(m, d); + } + mj_resetData(m, d); + } + + mj_deleteData(d); + expected_destroy_count += 2; + EXPECT_EQ(TestActuator::InitCount(), expected_init_count); + EXPECT_EQ(TestActuator::DestroyCount(), expected_destroy_count); + + mj_deleteModel(m); + EXPECT_EQ(TestActuator::InitCount(), expected_init_count); + EXPECT_EQ(TestActuator::DestroyCount(), expected_destroy_count); +} + +} // namespace +} // namespace mujoco diff --git a/unity/Runtime/Bindings/MujocoBindings.cs b/unity/Runtime/Bindings/MujocoBindings.cs index 1c34567f..f3e330da 100644 --- a/unity/Runtime/Bindings/MujocoBindings.cs +++ b/unity/Runtime/Bindings/MujocoBindings.cs @@ -49,6 +49,7 @@ public const int mjNFLUID = 12; public const int mjNREF = 2; public const int mjNIMP = 5; public const int mjNSOLVER = 1000; +public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJPLUGIN_H_ = true; public const bool THIRD_PARTY_MUJOCO_MJRENDER_H_ = true; public const int mjNAUX = 10; public const int mjMAXTEXTURE = 1000; @@ -282,6 +283,7 @@ public enum mjtObj : int{ mjOBJ_TEXT = 21, mjOBJ_TUPLE = 22, mjOBJ_KEY = 23, + mjOBJ_PLUGIN = 24, } public enum mjtConstraint : int{ mjCNSTR_EQUALITY = 0, @@ -337,7 +339,8 @@ public enum mjtSensor : int{ mjSENS_SUBTREELINVEL = 33, mjSENS_SUBTREEANGMOM = 34, mjSENS_CLOCK = 35, - mjSENS_USER = 36, + mjSENS_PLUGIN = 36, + mjSENS_USER = 37, } public enum mjtStage : int{ mjSTAGE_NONE = 0, @@ -357,6 +360,10 @@ public enum mjtLRMode : int{ mjLRMODE_MUSCLEUSER = 2, mjLRMODE_ALL = 3, } +public enum mjtPluginTypeBit : int{ + mjPLUGIN_ACTUATOR = 1, + mjPLUGIN_SENSOR = 2, +} public enum mjtGridPos : int{ mjGRID_TOPLEFT = 0, mjGRID_TOPRIGHT = 1, @@ -541,6 +548,7 @@ public unsafe struct mjSolverStat_ { public unsafe struct mjData_ { public int nstack; public int nbuffer; + public int nplugin; public int pstack; public int maxuse_stack; public int maxuse_con; @@ -1581,6 +1589,7 @@ public unsafe struct mjData_ { public double* qvel; public double* act; public double* qacc_warmstart; + public double* plugin_state; public double* ctrl; public double* qfrc_applied; public double* xfrc_applied; @@ -1590,6 +1599,8 @@ public unsafe struct mjData_ { public double* act_dot; public double* userdata; public double* sensordata; + public int* plugin; + public UIntPtr* plugin_data; public double* xpos; public double* xquat; public double* xmat; @@ -1883,6 +1894,8 @@ public unsafe struct mjModel_ { public int ntupledata; public int nkey; public int nmocap; + public int nplugin; + public int npluginattr; public int nuser_body; public int nuser_jnt; public int nuser_geom; @@ -1900,6 +1913,7 @@ public unsafe struct mjModel_ { public int nstack; public int nuserdata; public int nsensordata; + public int npluginstate; public int nbuffer; public mjOption_ opt; public mjVisual_ vis; @@ -2126,6 +2140,7 @@ public unsafe struct mjModel_ { public double* actuator_length0; public double* actuator_lengthrange; public double* actuator_user; + public int* actuator_plugin; public int* sensor_type; public int* sensor_datatype; public int* sensor_needstage; @@ -2138,6 +2153,11 @@ public unsafe struct mjModel_ { public double* sensor_cutoff; public double* sensor_noise; public double* sensor_user; + public int* sensor_plugin; + public int* plugin; + public int* plugin_stateadr; + public char* plugin_attr; + public int* plugin_attradr; public int* numeric_adr; public int* numeric_size; public double* numeric_data; @@ -2177,6 +2197,7 @@ public unsafe struct mjModel_ { public int* name_textadr; public int* name_tupleadr; public int* name_keyadr; + public int* name_pluginadr; public char* names; } @@ -3000,6 +3021,10 @@ public static unsafe extern double mj_getTotalmass(mjModel_* m); [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern void mj_setTotalmass(mjModel_* m, double newmass); +[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] +[return: MarshalAs(UnmanagedType.LPStr)] +public static unsafe extern string mj_getPluginConfig(mjModel_* m, int plugin_id, [MarshalAs(UnmanagedType.LPStr)]string attrib); + [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)] public static unsafe extern int mj_version();