Add mocap state to SaveState/RestoreState and attached keyframes.

PiperOrigin-RevId: 662975257
Change-Id: I3edb44f34426525813033902caf31db16ad7b2ab
This commit is contained in:
Alessio Quaglino
2024-08-14 10:41:40 -07:00
committed by Copybara-Service
parent 5927b3d4d9
commit 4c64e6f036
8 changed files with 264 additions and 67 deletions
-1
View File
@@ -3808,7 +3808,6 @@ all attachments will appear in the saved XML file.
- An entire model cannot be attached (i.e. including all elements, referenced or not).
- All assets from the child model will be copied in, whether they are referenced or not.
- Self-attach or circular references are not checked for and will lead to infinite loops.
- :ref:`Keyframes<keyframe>` are not yet supported. When attaching, all keyframes will be deleted.
.. _body-attach-model:
+2 -3
View File
@@ -7,10 +7,9 @@ Upcoming version (not yet released)
General
^^^^^^^
1. Add :ref:`mjSpec` option for creating a texture from a buffer.
1. Added :ref:`mjSpec` option for creating a texture from a buffer.
2. :ref:`shellinertia <body-geom-shellinertia>` is now supported by all geom types.
3. Add support for :ref:`attaching<meAttachment>` keyframes. Note: this only supports keyframe containing qpos, qvel,
and act.
3. Added support for :ref:`attaching<meAttachment>` keyframes.
Version 3.2.2 (Aug 8, 2024)
---------------------------
+4 -2
View File
@@ -93,15 +93,17 @@ mjModel* mj_compile(mjSpec* s, const mjVFS* vfs) {
// recompile spec into existing model and data while preserving the state
void mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d) {
mjCModel* modelC = static_cast<mjCModel*>(s->element);
std::string state_name = "state";
mjtNum time = 0;
if (d) {
time = d->time;
modelC->SaveState(d->qpos, d->qvel, d->act, d->ctrl);
modelC->SaveState(state_name, d->qpos, d->qvel, d->act, d->ctrl, d->mocap_pos, d->mocap_quat);
}
modelC->Compile(vfs, &m);
if (d) {
modelC->MakeData(m, &d);
modelC->RestoreState(m->qpos0, d->qpos, d->qvel, d->act, d->ctrl);
modelC->RestoreState(state_name, m->qpos0, m->body_pos, m->body_quat, d->qpos, d->qvel,
d->act, d->ctrl, d->mocap_pos, d->mocap_quat);
d->time = time;
}
}
+109 -36
View File
@@ -109,9 +109,6 @@ mjCModel::mjCModel() {
// this class allocated the plugins
plugin_owner = true;
// default state name
state_name_ = "state";
}
@@ -216,6 +213,7 @@ void mjCModel::SaveDofOffsets() {
int qposadr = 0;
int dofadr = 0;
int actadr = 0;
int nmocap = 0;
for (auto joint : joints_) {
joint->qposadr_ = qposadr;
@@ -233,6 +231,15 @@ void mjCModel::SaveDofOffsets() {
actuator->actadr_ = actuator->actdim_ ? actadr : -1;
actadr += actuator->actdim_;
}
for (mjCBody* body : bodies_) {
if (body->spec.mocap) {
body->mocapid = nmocap;
nmocap++;
} else {
body->mocapid = -1;
}
}
}
@@ -2863,23 +2870,36 @@ void mjCModel::CopyObjects(mjModel* m) {
// save the current state
template <class T>
void mjCModel::SaveState(const T* qpos, const T* qvel, const T* act, const T* ctrl) {
void mjCModel::SaveState(const std::string& state_name, const T* qpos, const T* qvel, const T* act,
const T* ctrl, const T* mpos, const T* mquat) {
for (auto joint : joints_) {
if (joint->qposadr_ == -1 || joint->dofadr_ == -1) {
throw mjCError(NULL, "SaveState: joint %s has no address", joint->name.c_str());
}
if (qpos) mjuu_copyvec(joint->qpos(), qpos + joint->qposadr_, joint->nq());
if (qvel) mjuu_copyvec(joint->qvel(), qvel + joint->dofadr_, joint->nv());
if (qpos) mjuu_copyvec(joint->qpos(state_name), qpos + joint->qposadr_, joint->nq());
if (qvel) mjuu_copyvec(joint->qvel(state_name), qvel + joint->dofadr_, joint->nv());
}
for (unsigned int i=0; i<actuators_.size(); i++) {
auto actuator = actuators_[i];
if (actuator->actadr_ != -1 && actuator->actdim_ != -1 && act) {
actuator->act().assign(actuator->actdim_, 0);
mjuu_copyvec(actuator->act().data(), act + actuator->actadr_, actuator->actdim_);
actuator->act(state_name).assign(actuator->actdim_, 0);
mjuu_copyvec(actuator->act(state_name).data(), act + actuator->actadr_, actuator->actdim_);
}
if (ctrl) {
actuator->ctrl() = ctrl[i];
actuator->ctrl(state_name) = ctrl[i];
}
}
for (auto body : bodies_) {
if (!body->spec.mocap) {
continue;
}
if (mpos) {
mjuu_copyvec(body->mpos(state_name), mpos + 3*body->mocapid, 3);
}
if (mquat) {
mjuu_copyvec(body->mquat(state_name), mquat + 4*body->mocapid, 4);
}
}
}
@@ -2900,42 +2920,63 @@ void mjCModel::MakeData(const mjModel* m, mjData** dest) {
// restore the previous state
template <class T>
void mjCModel::RestoreState(const mjtNum* pos0, T* qpos, T* qvel, T* act, T* ctrl) {
void mjCModel::RestoreState(const std::string& state_name, const mjtNum* pos0,
const mjtNum* mpos0, const mjtNum* mquat0, T* qpos,
T* qvel, T* act, T* ctrl, T* mpos, T* mquat) {
for (auto joint : joints_) {
if (qpos) {
if (mjuu_defined(joint->qpos()[0])) {
mjuu_copyvec(qpos + joint->qposadr_, joint->qpos(), joint->nq());
if (mjuu_defined(joint->qpos(state_name)[0])) {
mjuu_copyvec(qpos + joint->qposadr_, joint->qpos(state_name), joint->nq());
} else {
mjuu_copyvec(qpos + joint->qposadr_, pos0 + joint->qposadr_, joint->nq());
}
}
if (mjuu_defined(joint->qvel()[0]) && qvel) {
mjuu_copyvec(qvel + joint->dofadr_, joint->qvel(), joint->nv());
if (mjuu_defined(joint->qvel(state_name)[0]) && qvel) {
mjuu_copyvec(qvel + joint->dofadr_, joint->qvel(state_name), joint->nv());
}
}
// restore act
for (unsigned int i=0; i<actuators_.size(); i++) {
auto actuator = actuators_[i];
if (!actuator->act().empty() && mjuu_defined(actuator->act()[0]) && act) {
mjuu_copyvec(act + actuator->actadr_, actuator->act().data(), actuator->actdim_);
if (!actuator->act(state_name).empty() && mjuu_defined(actuator->act(state_name)[0]) && act) {
mjuu_copyvec(act + actuator->actadr_, actuator->act(state_name).data(), actuator->actdim_);
}
if (ctrl) {
ctrl[i] = mjuu_defined(actuator->ctrl()) ? actuator->ctrl() : 0;
ctrl[i] = mjuu_defined(actuator->ctrl(state_name)) ? actuator->ctrl(state_name) : 0;
}
}
for (unsigned int i=0; i<bodies_.size(); i++) {
auto body = bodies_[i];
if (!body->mocap) {
continue;
}
if (mpos) {
if (mjuu_defined(body->mpos(state_name)[0])) {
mjuu_copyvec(mpos + 3*body->mocapid, body->mpos(state_name), 3);
} else {
mjuu_copyvec(mpos + 3*body->mocapid, mpos0 + 3*i, 3);
}
}
if (mquat) {
if (mjuu_defined(body->mquat(state_name)[0])) {
mjuu_copyvec(mquat + 4*body->mocapid, body->mquat(state_name), 4);
} else {
mjuu_copyvec(mquat + 4*body->mocapid, mquat0 + 4*i, 4);
}
}
}
}
// force explicit instantiations
template void mjCModel::SaveState<mjtNum>(const mjtNum* qpos,
const mjtNum* qvel,
const mjtNum* act,
const mjtNum* ctrl);
template void mjCModel::SaveState<mjtNum>(
const std::string& name, const mjtNum* qpos, const mjtNum* qvel, const mjtNum* act,
const mjtNum* ctrl, const mjtNum* mpos, const mjtNum* mquat);
template void mjCModel::RestoreState<mjtNum>(const mjtNum* qpos0, mjtNum* qpos,
mjtNum* qvel, mjtNum* act, mjtNum* ctrl);
template void mjCModel::RestoreState<mjtNum>(
const std::string& name, const mjtNum* qpos0, const mjtNum* mpos0, const mjtNum* mquat0,
mjtNum* qpos, mjtNum* qvel, mjtNum* act, mjtNum* ctrl, mjtNum* mpos, mjtNum* mquat);
@@ -2957,10 +2998,12 @@ void mjCModel::StoreKeyframes() {
info.qvel = !key->spec_qvel_.empty();
info.act = !key->spec_act_.empty();
info.ctrl = !key->spec_ctrl_.empty();
info.mpos = !key->spec_mpos_.empty();
info.mquat = !key->spec_mquat_.empty();
key_pending_.push_back(info);
state_name_ = info.name;
SaveState(key->spec_qpos_.data(), key->spec_qvel_.data(),
key->spec_act_.data(), key->spec_ctrl_.data());
SaveState(info.name, key->spec_qpos_.data(), key->spec_qvel_.data(),
key->spec_act_.data(), key->spec_ctrl_.data(),
key->spec_mpos_.data(), key->spec_mquat_.data());
}
if (resetlists) {
@@ -3502,7 +3545,10 @@ void mjCModel::ResolveKeyframes(const mjModel* m) {
return;
}
// resize non-pending keyframes to the new number of dofs
// store dof offsets in joints and actuators
SaveDofOffsets();
// resize existing keyframes to the new state, fill in missing default values
for (unsigned int i = 0; i < nkey - key_pending_.size(); i++) {
mjCKey* key = keys_[i];
if (!key->spec_qpos_.empty()) {
@@ -3521,12 +3567,36 @@ void mjCModel::ResolveKeyframes(const mjModel* m) {
if (!key->spec_ctrl_.empty()) {
key->spec_ctrl_.resize(nu);
}
if (!key->spec_mpos_.empty()) {
int nmocap0 = key->spec_mpos_.size() / 3;
key->spec_mpos_.resize(3*nmocap);
for (unsigned int j = 0; j < bodies_.size(); j++) {
if (bodies_[j]->mocapid < nmocap0) {
continue;
}
int i = bodies_[j]->mocapid;
key->spec_mpos_[3*i+0] = (double)m->body_pos[3*j+0];
key->spec_mpos_[3*i+1] = (double)m->body_pos[3*j+1];
key->spec_mpos_[3*i+2] = (double)m->body_pos[3*j+2];
}
}
if (!key->spec_mquat_.empty()) {
int nmocap0 = key->spec_mquat_.size() / 4;
key->spec_mquat_.resize(4*nmocap);
for (unsigned int j = 0; j < bodies_.size(); j++) {
if (bodies_[j]->mocapid < nmocap0) {
continue;
}
int i = bodies_[j]->mocapid;
key->spec_mquat_[4*i+0] = (double)m->body_quat[4*j+0];
key->spec_mquat_[4*i+1] = (double)m->body_quat[4*j+1];
key->spec_mquat_[4*i+2] = (double)m->body_quat[4*j+2];
key->spec_mquat_[4*i+3] = (double)m->body_quat[4*j+3];
}
}
}
// store dof offsets in joints and actuators
SaveDofOffsets();
// copy state stored in joints and actuators to keyframes
// create new keyframes, fill in missing default values
for (const auto& info : key_pending_) {
mjCKey* key = (mjCKey*)FindObject(mjOBJ_KEY, info.name);
key->name = info.name;
@@ -3535,9 +3605,12 @@ void mjCModel::ResolveKeyframes(const mjModel* m) {
if (info.qvel) key->spec_qvel_.assign(nv, 0);
if (info.act) key->spec_act_.assign(na, 0);
if (info.ctrl) key->spec_ctrl_.assign(nu, 0);
state_name_ = info.name;
RestoreState(m->qpos0, key->spec_qpos_.data(), key->spec_qvel_.data(),
key->spec_act_.data(), key->spec_ctrl_.data());
if (info.mpos) key->spec_mpos_.assign(3*nmocap, 0);
if (info.mquat) key->spec_mquat_.assign(4*nmocap, 0);
RestoreState(info.name, m->qpos0, m->body_pos, m->body_quat,
key->spec_qpos_.data(), key->spec_qvel_.data(),
key->spec_act_.data(), key->spec_ctrl_.data(),
key->spec_mpos_.data(), key->spec_mquat_.data());
}
// the attached keyframes have been copied into the model
+11 -5
View File
@@ -41,6 +41,8 @@ typedef struct mjKeyInfo_ {
bool qvel;
bool act;
bool ctrl;
bool mpos;
bool mquat;
} mjKeyInfo;
class mjCModel_ : public mjsElement {
@@ -284,9 +286,15 @@ class mjCModel : public mjCModel_, private mjSpec {
template <class T> void DeleteMaterial(std::vector<T*>& list,
std::string_view name = "");
// save/restore the current state
template <class T> void SaveState(const T* qpos, const T* qvel, const T* act, const T* ctrl);
template <class T> void RestoreState(const mjtNum* pos0, T* qpos, T* qvel, T* act, T* ctrl);
// save the current state
template <class T>
void SaveState(const std::string& state_name, const T* qpos, const T* qvel, const T* act,
const T* ctrl, const T* mpos, const T* mquat);
// restore the previously saved state
template <class T>
void RestoreState(const std::string& state_name, const mjtNum* pos0, const mjtNum* mpos0,
const mjtNum* mquat0, T* qpos, T* qvel, T* act, T* ctrl, T* mpos, T* mquat);
// clear existing data
void MakeData(const mjModel* m, mjData** dest);
@@ -383,7 +391,5 @@ class mjCModel : public mjCModel_, private mjSpec {
mjCError errInfo; // last error info
bool plugin_owner; // this class allocated the plugins
std::vector<mjKeyInfo> key_pending_; // attached keyframes
std::string state_name_;
};
#endif // MUJOCO_SRC_USER_USER_MODEL_H_
+36 -16
View File
@@ -1411,6 +1411,8 @@ void mjCBody::ForgetKeyframes() const {
joint->qpos_.clear();
joint->qvel_.clear();
}
model->FindBody((mjCBody*)this, name)->mpos_.clear(); // this is a hack to avoid const
model->FindBody((mjCBody*)this, name)->mquat_.clear(); // this is a hack to avoid const
for (auto body : bodies) {
body->ForgetKeyframes();
}
@@ -1418,6 +1420,24 @@ void mjCBody::ForgetKeyframes() const {
mjtNum* mjCBody::mpos(const std::string& state_name) {
if (mpos_.find(state_name) == mpos_.end()) {
mpos_[state_name] = {mjNAN, 0, 0};
}
return mpos_.at(state_name).data();
}
mjtNum* mjCBody::mquat(const std::string& state_name) {
if (mquat_.find(state_name) == mquat_.end()) {
mquat_[state_name] = {mjNAN, 0, 0, 0};
}
return mquat_.at(state_name).data();
}
// compiler
void mjCBody::Compile(void) {
CopyFromSpec();
@@ -1828,20 +1848,20 @@ int mjCJoint::nv(mjtJoint joint_type) {
mjtNum* mjCJoint::qpos() {
if (qpos_.find(model->state_name_) == qpos_.end()) {
qpos_[model->state_name_] = {mjNAN, 0, 0, 0, 0, 0, 0};
mjtNum* mjCJoint::qpos(const std::string& state_name) {
if (qpos_.find(state_name) == qpos_.end()) {
qpos_[state_name] = {mjNAN, 0, 0, 0, 0, 0, 0};
}
return qpos_.at(model->state_name_).data();
return qpos_.at(state_name).data();
}
mjtNum* mjCJoint::qvel() {
if (qvel_.find(model->state_name_) == qvel_.end()) {
qvel_[model->state_name_] = {mjNAN, 0, 0, 0, 0, 0};
mjtNum* mjCJoint::qvel(const std::string& state_name) {
if (qvel_.find(state_name) == qvel_.end()) {
qvel_[state_name] = {mjNAN, 0, 0, 0, 0, 0};
}
return qvel_.at(model->state_name_).data();
return qvel_.at(state_name).data();
}
@@ -5405,20 +5425,20 @@ bool mjCActuator::is_actlimited() const { return islimited(actlimited, actrange)
std::vector<mjtNum>& mjCActuator::act() {
if (act_.find(model->state_name_) == act_.end()) {
act_[model->state_name_] = std::vector<mjtNum>(model->nu, mjNAN);
std::vector<mjtNum>& mjCActuator::act(const std::string& state_name) {
if (act_.find(state_name) == act_.end()) {
act_[state_name] = std::vector<mjtNum>(model->nu, mjNAN);
}
return act_.at(model->state_name_);
return act_.at(state_name);
}
mjtNum& mjCActuator::ctrl() {
if (ctrl_.find(model->state_name_) == ctrl_.end()) {
ctrl_[model->state_name_] = mjNAN;
mjtNum& mjCActuator::ctrl(const std::string& state_name) {
if (ctrl_.find(state_name) == ctrl_.end()) {
ctrl_[state_name] = mjNAN;
}
return ctrl_.at(model->state_name_);
return ctrl_.at(state_name);
}
+12 -4
View File
@@ -255,6 +255,10 @@ class mjCBody_ : public mjCBase {
std::string plugin_instance_name;
std::vector<double> userdata_;
std::vector<double> spec_userdata_;
// variables used for temporarily storing the state of the mocap bodies
std::map<std::string, std::array<mjtNum, 3>> mpos_; // saved mocap_pos
std::map<std::string, std::array<mjtNum, 4>> mquat_; // saved mocap_quat
};
class mjCBody : public mjCBody_, private mjsBody {
@@ -323,6 +327,10 @@ class mjCBody : public mjCBody_, private mjsBody {
// reset keyframe references for allowing self-attach
void ForgetKeyframes() const;
// get mocap position and quaternion
mjtNum* mpos(const std::string& state_name);
mjtNum* mquat(const std::string& state_name);
private:
mjCBody(const mjCBody& other, mjCModel* _model); // copy constructor
mjCBody& operator=(const mjCBody& other); // copy assignment
@@ -444,8 +452,8 @@ class mjCJoint : public mjCJoint_, private mjsJoint {
int nq() const { return nq(spec.type); }
int nv() const { return nv(spec.type); }
mjtNum* qpos();
mjtNum* qvel();
mjtNum* qpos(const std::string& state_name);
mjtNum* qvel(const std::string& state_name);
private:
int Compile(void); // compiler; return dofnum
@@ -1436,8 +1444,8 @@ class mjCActuator : public mjCActuator_, private mjsActuator {
bool is_forcelimited() const;
bool is_actlimited() const;
std::vector<mjtNum>& act();
mjtNum& ctrl();
std::vector<mjtNum>& act(const std::string& state_name);
mjtNum& ctrl(const std::string& state_name);
private:
void Compile(void); // compiler
+90
View File
@@ -870,6 +870,8 @@ TEST_F(MujocoTest, PreserveState) {
<joint type="slide" axis="0 0 1" name="slide"/>
<geom type="sphere" size=".2"/>
</body>
<body name="mocap_detach" mocap="true"/>
<body name="mocap" mocap="true"/>
</worldbody>
<actuator>
<position name="hinge" joint="hinge" timeconst=".01"/>
@@ -888,6 +890,7 @@ TEST_F(MujocoTest, PreserveState) {
<joint type="slide" axis="0 0 1"/>
<geom type="sphere" size=".3"/>
</body>
<body name="mocap" mocap="true"/>
</worldbody>
<actuator>
<position name="slide" joint="slide" timeconst=".01"/>
@@ -915,6 +918,14 @@ TEST_F(MujocoTest, PreserveState) {
data->ctrl[1] = 2;
d_expected->ctrl[0] = 2;
// set mocap
data->mocap_pos[3] = 1;
data->mocap_quat[4] = 0;
data->mocap_quat[5] = 1;
d_expected->mocap_pos[0] = 1;
d_expected->mocap_quat[0] = 0;
d_expected->mocap_quat[1] = 1;
// step models
mj_step(model, data);
mj_step(m_expected, d_expected);
@@ -925,6 +936,11 @@ TEST_F(MujocoTest, PreserveState) {
EXPECT_THAT(body, NotNull());
EXPECT_THAT(mjs_detachBody(spec, body), 0);
// detach mocap
mjsBody* mocap_body = mjs_findBody(spec, "mocap_detach");
EXPECT_THAT(mocap_body, NotNull());
EXPECT_THAT(mjs_detachBody(spec, mocap_body), 0);
// add body
mjsBody* newbody = mjs_addBody(mjs_findBody(spec, "world"), 0);
EXPECT_THAT(newbody, NotNull());
@@ -968,6 +984,17 @@ TEST_F(MujocoTest, PreserveState) {
EXPECT_EQ(data->act[i], d_expected->act[i]) << i;
}
// compare mocap
EXPECT_EQ(model->nmocap, m_expected->nmocap);
for (int i = 0; i < model->nmocap; ++i) {
for (int j = 0; j < 3; ++j) {
EXPECT_EQ(data->mocap_pos[3*i+j], d_expected->mocap_pos[3*i+j]) << i;
}
for (int j = 0; j < 4; ++j) {
EXPECT_EQ(data->mocap_quat[4*i+j], d_expected->mocap_quat[4*i+j]) << i;
}
}
// check that the function is callable with no data
mj_deleteData(data);
mj_recompile(spec, 0, model, nullptr);
@@ -979,5 +1006,68 @@ TEST_F(MujocoTest, PreserveState) {
mj_deleteModel(m_expected);
}
TEST_F(MujocoTest, AttachMocap) {
std::array<char, 1000> er;
mjtNum tol = 0;
std::string field = "";
static constexpr char xml[] = R"(
<mujoco>
<worldbody>
<body pos="1 1 1" quat="0 1 0 0" name="mocap" mocap="true"/>
</worldbody>
<keyframe>
<key name="key" time="1" mpos="2 2 2" mquat="1 0 0 0"/>
</keyframe>
</mujoco>)";
static constexpr char xml_expected[] = R"(
<mujoco>
<worldbody>
<body pos="1 1 1" quat="0 1 0 0" name="mocap" mocap="true"/>
<body pos="3 3 3" quat="0 0 1 0" name="attached-mocap-1" mocap="true"/>
</worldbody>
<keyframe>
<key name="key" time="1" mpos="2 2 2 3 3 3" mquat="1 0 0 0 0 0 1 0"/>
<key name="attached-key-1" time="1" mpos="1 1 1 2 2 2" mquat="0 1 0 0 1 0 0 0"/>
</keyframe>
</mujoco>)";
mjSpec* spec = mj_parseXMLString(xml, 0, er.data(), er.size());
EXPECT_THAT(spec, NotNull()) << er.data();
mjsBody* body = mjs_findBody(spec, "mocap");
EXPECT_THAT(body, NotNull());
mjsBody* world = mjs_findBody(spec, "world");
EXPECT_THAT(world, NotNull());
mjsFrame* frame = mjs_addFrame(world, NULL);
mjs_attachBody(frame, body, "attached-", "-1");
mjsBody* attached_body = mjs_findBody(spec, "attached-mocap-1");
EXPECT_THAT(attached_body, NotNull());
attached_body->pos[0] = 3;
attached_body->pos[1] = 3;
attached_body->pos[2] = 3;
attached_body->quat[0] = 0;
attached_body->quat[1] = 0;
attached_body->quat[2] = 1;
attached_body->quat[3] = 0;
mjModel* model = mj_compile(spec, 0);
EXPECT_THAT(model, NotNull());
mjModel* m_expected = LoadModelFromString(xml_expected, er.data(), er.size());
EXPECT_THAT(m_expected, NotNull()) << er.data();
EXPECT_LE(CompareModel(model, m_expected, field), tol)
<< "Expected and attached models are different!\n"
<< "Different field: " << field << '\n';
mj_deleteSpec(spec);
mj_deleteModel(model);
mj_deleteModel(m_expected);
}
} // namespace
} // namespace mujoco