From ba8a2d4c8ad5ae74df8a96253225864d9ecad00e Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Thu, 6 Jun 2024 03:29:56 -0700 Subject: [PATCH] Dynamic model editing. PiperOrigin-RevId: 640838852 Change-Id: I559b1d60f2a3a2d866788d4006be31757454c619 --- src/engine/engine_io.c | 6 +- src/engine/engine_io.h | 3 + src/user/user_api.cc | 10 ++++ src/user/user_api.h | 4 ++ src/user/user_model.cc | 83 ++++++++++++++++++++++++-- src/user/user_model.h | 10 ++-- src/user/user_objects.cc | 7 +++ src/user/user_objects.h | 10 +++- test/user/user_api_test.cc | 119 +++++++++++++++++++++++++++++++++++++ 9 files changed, 239 insertions(+), 13 deletions(-) diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index ee836546..c41932eb 100644 --- a/src/engine/engine_io.c +++ b/src/engine/engine_io.c @@ -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 diff --git a/src/engine/engine_io.h b/src/engine/engine_io.h index 1c733a8a..da4d6082 100644 --- a/src/engine/engine_io.h +++ b/src/engine/engine_io.h @@ -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 diff --git a/src/user/user_api.cc b/src/user/user_api.cc index 34cfa704..1c6384fc 100644 --- a/src/user/user_api.cc +++ b/src/user/user_api.cc @@ -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(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) { diff --git a/src/user/user_api.h b/src/user/user_api.h index 3f6ae15e..8cb70d8f 100644 --- a/src/user/user_api.h +++ b/src/user/user_api.h @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -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); diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 7ff79c2c..f2bcbe5a 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -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 @@ -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(&m), *const_cast(&data), vfs); + TryCompile(*const_cast(&model), *const_cast(&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; } diff --git a/src/user/user_model.h b/src/user/user_model.h index 135510af..95939964 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -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 void DeleteMaterial(std::vector& 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 defaults_; @@ -275,6 +276,7 @@ class mjCModel : public mjCModel_, private mjSpec { std::vector> 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 diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 034abae6..d59e4950 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -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); } diff --git a/src/user/user_objects.h b/src/user/user_objects.h index d5773515..8860a6c9 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -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 act; // act at the previous step }; diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index e6b63761..2bc661bc 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -747,5 +747,124 @@ TEST_F(MujocoTest, DetachBody) { TestDetachBody(/*compile=*/true); } +TEST_F(MujocoTest, PreserveState) { + std::array er; + std::string field = ""; + + static constexpr char xml_full[] = R"( + + + + + + + + + + + + + + + + )"; + + static constexpr char xml_expected[] = R"( + + + + + + + + + + + + + + + )"; + + // 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