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

Implement plugin mechanism for actuators and sensors.

Browse Source 
PiperOrigin-RevId: 474874088
Change-Id: I65a8ffdf845f4fa0f8266c165883a747ab5812d8
This commit is contained in:
Saran Tunyasuvunakool
2022-09-16 12:20:56 -07:00
committed by Copybara-Service
parent f556d4d94f
commit 1e2a9a53bc
29 changed files with 2121 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_io.c
+101 -9
View File
@@ -22,8 +22,10 @@
#include <string.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjplugin.h>
#include <mujoco/mjxmacro.h>
#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; i<m->nsensor; 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.";
}