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

Dynamic model editing.

Browse Source 
PiperOrigin-RevId: 640838852
Change-Id: I559b1d60f2a3a2d866788d4006be31757454c619
This commit is contained in:
Alessio Quaglino
2024-06-06 03:29:56 -07:00
committed by Copybara-Service
parent 40755e165d
commit ba8a2d4c8a
9 changed files with 239 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_io.c
+3 -3
View File
@@ -1093,7 +1093,7 @@ static void mj_setPtrData(const mjModel* m, mjData* d) {
// initialize plugins, copy into d (required for deletion)
static void _initPlugin(const mjModel* m, mjData* d) {
void mj_initPlugin(const mjModel* m, mjData* d) {
d->nplugin = m->nplugin;
for (int i = 0; i < m->nplugin; ++i) {
d->plugin[i] = m->plugin[i];
@@ -1203,7 +1203,7 @@ mjData* mj_makeData(const mjModel* m) {
mjData* d = NULL;
mj_makeRawData(&d, m);
if (d) {
_initPlugin(m, d);
mj_initPlugin(m, d);
mj_resetData(m, d);
}
return d;
@@ -1219,7 +1219,7 @@ mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src) {
// allocate new data if needed
if (!dest) {
mj_makeRawData(&dest, m);
_initPlugin(m, dest);
mj_initPlugin(m, dest);
}
// check sizes
src/engine/engine_io.h
+3
View File
@@ -110,6 +110,9 @@ MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
// mjData arena allocate
MJAPI void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment);
// init plugins
MJAPI void mj_initPlugin(const mjModel* m, mjData* d);
#ifndef ADDRESS_SANITIZER
// mjData mark stack frame
src/user/user_api.cc
+10
View File
@@ -83,6 +83,16 @@ mjModel* mjs_compile(mjSpec* s, const mjVFS* vfs) {
// recompile spec into existing model and data while preserving the state
void mjs_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d) {
mjCModel* modelC = static_cast<mjCModel*>(s->element);
modelC->SaveState(d);
modelC->Compile(vfs, &m);
modelC->RestoreState(m, &d);
}
// attach body to a frame of the parent
int mjs_attachBody(mjsFrame* parent, const mjsBody* child,
const char* prefix, const char* suffix) {
src/user/user_api.h
+4
View File
@@ -17,6 +17,7 @@
#include <math.h>
#include <stddef.h>
#include <mujoco/mjdata.h>
#include <mujoco/mjexport.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjtnum.h>
@@ -754,6 +755,9 @@ MJAPI mjSpec* mjs_createSpec(void);
// Compile spec to model.
MJAPI mjModel* mjs_compile(mjSpec* s, const mjVFS* vfs);
// Recompile spec to model preserving the current state.
MJAPI void mjs_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d);
// Copy spec.
MJAPI mjSpec* mjs_copySpec(const mjSpec* s);
src/user/user_model.cc
+78 -5
View File
@@ -1843,6 +1843,8 @@ void mjCModel::CopyTree(mjModel* m) {
int jid = pj->id;
// set joint fields
pj->qposadr_ = qposadr;
pj->dofadr_ = dofadr;
m->jnt_type[jid] = pj->type;
m->jnt_group[jid] = pj->group;
m->jnt_limited[jid] = (mjtByte)pj->is_limited();
@@ -2584,6 +2586,8 @@ void mjCModel::CopyObjects(mjModel* m) {
m->actuator_trnid[2*i+1] = pac->trnid[1];
m->actuator_actnum[i] = pac->actdim + pac->plugin_actdim;
m->actuator_actadr[i] = m->actuator_actnum[i] ? adr : -1;
pac->actadr_ = m->actuator_actadr[i];
pac->actnum_ = m->actuator_actnum[i];
adr += m->actuator_actnum[i];
m->actuator_group[i] = pac->group;
m->actuator_ctrllimited[i] = (mjtByte)pac->is_ctrllimited();
@@ -2716,6 +2720,75 @@ void mjCModel::CopyObjects(mjModel* m) {
// save the current state
void mjCModel::SaveState(const mjData* d) {
for (auto joint : joints_) {
switch (joint->type) {
case mjJNT_FREE:
mju_copy(joint->qpos, d->qpos + joint->qposadr_, 7);
mju_copy(joint->qvel, d->qvel + joint->dofadr_, 6);
break;
case mjJNT_BALL:
mju_copy(joint->qpos, d->qpos + joint->qposadr_, 4);
mju_copy(joint->qvel, d->qvel + joint->dofadr_, 3);
break;
case mjJNT_HINGE:
case mjJNT_SLIDE:
mju_copy(joint->qpos, d->qpos + joint->qposadr_, 1);
mju_copy(joint->qvel, d->qvel + joint->dofadr_, 1);
break;
}
}
for (auto actuator : actuators_) {
if (actuator->actadr_ != -1) {
actuator->act.assign(actuator->actnum_, 0);
mju_copy(actuator->act.data(), d->act + actuator->actadr_, actuator->actnum_);
}
}
}
// restore the previous state
void mjCModel::RestoreState(const mjModel* m, mjData** dest) {
mj_makeRawData(dest, m);
mjData* d = *dest;
if (d) {
mj_initPlugin(m, d);
mj_resetData(m, d);
}
for (auto joint : joints_) {
if (!mjuu_defined(joint->qpos[0]) || !mjuu_defined(joint->qvel[0])) {
continue;
}
switch (joint->type) {
case mjJNT_FREE:
mju_copy(d->qpos + joint->qposadr_, joint->qpos, 7);
mju_copy(d->qvel + joint->dofadr_, joint->qvel, 6);
break;
case mjJNT_BALL:
mju_copy(d->qpos + joint->qposadr_, joint->qpos, 4);
mju_copy(d->qvel + joint->dofadr_, joint->qvel, 3);
break;
case mjJNT_HINGE:
case mjJNT_SLIDE:
mju_copy(d->qpos + joint->qposadr_, joint->qpos, 1);
mju_copy(d->qvel + joint->dofadr_, joint->qvel, 1);
break;
}
}
for (auto actuator : actuators_) {
if (mjuu_defined(actuator->act[0])) {
mju_copy(d->act + actuator->actadr_, actuator->act.data(), actuator->actnum_);
}
}
}
//------------------------------- FUSE STATIC ------------------------------------------------------
template <class T>
@@ -3075,7 +3148,7 @@ static void warninghandler(const char* msg) {
// compiler
mjModel* mjCModel::Compile(const mjVFS* vfs) {
mjModel* mjCModel::Compile(const mjVFS* vfs, mjModel** m) {
if (compiled) {
// clear kinematic tree
for (int i=0; i<bodies_.size(); i++) {
@@ -3096,7 +3169,7 @@ mjModel* mjCModel::Compile(const mjVFS* vfs) {
// setjmp returns, and therefore to pass nullptr directly to the
// mj_deleteModel and mj_deleteData calls in the subsequent catch block,
// without ever reading the actual pointer values.
mjModel* volatile m = nullptr;
mjModel* volatile model = (m && *m) ? *m : nullptr;
mjData* volatile data = nullptr;
// save error and warning handlers
@@ -3118,10 +3191,10 @@ mjModel* mjCModel::Compile(const mjVFS* vfs) {
// TryCompile resulted in an mju_error which was converted to a longjmp.
throw mjCError(0, "engine error: %s", errortext);
}
TryCompile(*const_cast<mjModel**>(&m), *const_cast<mjData**>(&data), vfs);
TryCompile(*const_cast<mjModel**>(&model), *const_cast<mjData**>(&data), vfs);
} catch (mjCError err) {
// deallocate everything allocated in Compile
mj_deleteModel(m);
mj_deleteModel(model);
mj_deleteData(data);
mjCBody* world = bodies_[0];
Clear();
@@ -3140,7 +3213,7 @@ mjModel* mjCModel::Compile(const mjVFS* vfs) {
_mjPRIVATE__set_tls_error_fn(save_error);
_mjPRIVATE__set_tls_warning_fn(save_warning);
compiled = true;
return m;
return model;
}
src/user/user_model.h
+6 -4
View File
@@ -168,7 +168,7 @@ class mjCModel : public mjCModel_, private mjSpec {
mjSpec spec;
mjModel* Compile(const mjVFS* vfs = nullptr); // construct mjModel
mjModel* Compile(const mjVFS* vfs = nullptr, mjModel** m = nullptr); // construct mjModel
bool CopyBack(const mjModel*); // DECOMPILER: copy numeric back
void FuseStatic(); // fuse static bodies with parent
void FuseReindex(mjCBody* body); // reindex elements during fuse
@@ -254,9 +254,6 @@ class mjCModel : public mjCModel_, private mjSpec {
const std::string& plugin_instance_name,
mjCPlugin** plugin_instance);
void TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs);
mjModel* _Compile(const mjVFS* vfs);
// clear objects allocated by Compile
void Clear();
@@ -267,6 +264,10 @@ class mjCModel : public mjCModel_, private mjSpec {
template <class T> void DeleteMaterial(std::vector<T*>& list,
std::string_view name = "");
// save/restore the current state
void SaveState(const mjData* d);
void RestoreState(const mjModel* m, mjData** dest);
private:
// settings for each defaults class
std::vector<mjCDef*> defaults_;
@@ -275,6 +276,7 @@ class mjCModel : public mjCModel_, private mjSpec {
std::vector<std::pair<const mjpPlugin*, int>> active_plugins_;
// compile phases
void TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs);
void MakeLists(mjCBody* body); // make lists of bodies, geoms, joints, sites
void SetNuser(); // set nuser fields
void IndexAssets(bool discard); // convert asset names into indices
src/user/user_objects.cc
+7
View File
@@ -1733,6 +1733,10 @@ mjCJoint::mjCJoint(mjCModel* _model, mjCDef* _def) {
// in case this joint is not compiled
CopyFromSpec();
// no previous state when a joint is created
qpos[0] = mjNAN;
qvel[0] = mjNAN;
}
@@ -5003,6 +5007,9 @@ mjCActuator::mjCActuator(mjCModel* _model, mjCDef* _def) {
// point to local
PointToLocal();
// no previous state when an actuator is created
act.push_back(mjNAN);
}
src/user/user_objects.h
+9 -1
View File
@@ -423,10 +423,14 @@ class mjCJoint : public mjCJoint_, private mjsJoint {
bool is_limited() const;
bool is_actfrclimited() const;
private:
int Compile(void); // compiler; return dofnum
void PointToLocal(void);
int qposadr_; // address of dof in data->qpos
int dofadr_; // address of dof in data->qvel
mjtNum qpos[7]; // qpos at the previous step
mjtNum qvel[6]; // qvel at the previous step
};
@@ -1412,6 +1416,10 @@ class mjCActuator : public mjCActuator_, private mjsActuator {
void NameSpace(const mjCModel* m);
mjCBase* ptarget; // transmission target
int actadr_; // address of dof in data->act
int actnum_; // number of dofs in data->act
std::vector<mjtNum> act; // act at the previous step
};
test/user/user_api_test.cc
+119
View File
@@ -747,5 +747,124 @@ TEST_F(MujocoTest, DetachBody) {
TestDetachBody(/*compile=*/true);
}
TEST_F(MujocoTest, PreserveState) {
std::array<char, 1000> er;
std::string field = "";
static constexpr char xml_full[] = R"(
<mujoco>
<worldbody>
<body name="detachable" pos="1 0 0">
<joint type="hinge" axis="0 0 1" name="hinge"/>
<geom type="sphere" size=".1"/>
</body>
<body name="persistent">
<joint type="slide" axis="0 0 1" name="slide"/>
<geom type="sphere" size=".2"/>
</body>
</worldbody>
<actuator>
<position name="hinge" joint="hinge" timeconst=".01"/>
<position name="slide" joint="slide" timeconst=".01"/>
</actuator>
</mujoco>)";
static constexpr char xml_expected[] = R"(
<mujoco>
<worldbody>
<body name="persistent">
<joint type="slide" axis="0 0 1" name="slide"/>
<geom type="sphere" size=".2"/>
</body>
<body name="newbody" pos="2 0 0">
<joint type="slide" axis="0 0 1"/>
<geom type="sphere" size=".3"/>
</body>
</worldbody>
<actuator>
<position name="slide" joint="slide" timeconst=".01"/>
</actuator>
</mujoco>)";
// load spec
mjSpec* spec = mj_parseXMLString(xml_full, 0, er.data(), er.size());
EXPECT_THAT(spec, NotNull()) << er.data();
// compile models
mjModel* model = mjs_compile(spec, 0);
EXPECT_THAT(model, NotNull());
mjModel* m_expected = LoadModelFromString(xml_expected, er.data(), er.size());
EXPECT_THAT(m_expected, NotNull());
// create data
mjData* data = mj_makeData(model);
EXPECT_THAT(data, NotNull());
mjData* d_expected = mj_makeData(m_expected);
EXPECT_THAT(d_expected, NotNull());
// set ctrl
data->ctrl[0] = 1;
data->ctrl[1] = 2;
d_expected->ctrl[0] = 2;
// step models
mj_step(model, data);
mj_step(m_expected, d_expected);
// detach subtree
mjsBody* body = mjs_findBody(spec, "detachable");
EXPECT_THAT(body, NotNull());
EXPECT_THAT(mjs_detachBody(spec, body), 0);
// add body
mjsBody* newbody = mjs_addBody(mjs_findBody(spec, "world"), 0);
EXPECT_THAT(newbody, NotNull());
// add geom and joint
mjsGeom* geom = mjs_addGeom(newbody, 0);
mjsJoint* joint = mjs_addJoint(newbody, 0);
// set properties
newbody->pos[0] = 2;
geom->size[0] = .3;
joint->type = mjJNT_SLIDE;
joint->axis[0] = 0;
joint->axis[1] = 0;
joint->axis[2] = 1;
joint->ref = d_expected->qpos[m_expected->nq-1];
// compile new model
mjs_recompile(spec, 0, model, data);
EXPECT_THAT(model, NotNull());
// compare qpos
EXPECT_EQ(model->nq, m_expected->nq);
for (int i = 0; i < model->nq; ++i) {
EXPECT_EQ(data->qpos[i], d_expected->qpos[i]) << i;
}
// compare qvel
EXPECT_EQ(model->nv, m_expected->nv);
for (int i = 0; i < model->nv-1; ++i) {
EXPECT_EQ(data->qvel[i], d_expected->qvel[i]) << i;
}
// second body was added after stepping so qvel should be zero
EXPECT_EQ(data->qvel[model->nv-1], 0);
// compare act
EXPECT_EQ(model->na, m_expected->na);
for (int i = 0; i < model->na; ++i) {
EXPECT_EQ(data->act[i], d_expected->act[i]) << i;
}
// destroy everything
mj_deleteData(data);
mj_deleteData(d_expected);
mjs_deleteSpec(spec);
mj_deleteModel(model);
mj_deleteModel(m_expected);
}
} // namespace
} // namespace mujoco