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

Move keyframe resizing from compilation to attach (mjCModel::operator+=).

Browse Source 
This allows to resize parent keyframes also when the child has no keyframes.

PiperOrigin-RevId: 691045299
Change-Id: I3ab6509b13b5754e937ec77daf9efe697cabf1bc
This commit is contained in:
Alessio Quaglino
2024-10-29 09:16:54 -07:00
committed by Copybara-Service
parent e889cffe23
commit b941e994d9
3 changed files with 96 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/user/user_model.cc
+21 -21
View File
@@ -401,6 +401,16 @@ mjCModel& mjCModel::operator+=(const mjCModel& other) {
}
}
// resize keyframes in the parent model
if (!keys_.empty()) {
SaveDofOffsets(/*computesize=*/true);
ComputeReference();
for (auto* key : keys_) {
ResizeKeyframe(key, qpos0.data(), body_pos0.data(), body_quat0.data());
}
nq = nv = na = nu = nmocap = 0;
}
// restore to the original state
if (!compiled) {
ResetTreeLists();
@@ -3173,10 +3183,7 @@ void mjCModel::StoreKeyframes(mjCModel* dest) {
// do not change compilation quantities in case the user wants to recompile preserving the state
if (!compiled) {
SaveDofOffsets(/*computesize=*/true);
qpos0.resize(nq);
body_pos0.resize(3*bodies_.size());
body_quat0.resize(4*bodies_.size());
ComputeReference(qpos0, body_pos0, body_quat0);
ComputeReference();
}
// save keyframe info and resize keyframes
@@ -3729,26 +3736,28 @@ void mjCModel::CompileMeshes(const mjVFS* vfs) {
// compute qpos0
template <class T>
void mjCModel::ComputeReference(std::vector<T>& q0, std::vector<T>& bpos, std::vector<T>& bquat) {
void mjCModel::ComputeReference() {
int b = 0;
qpos0.resize(nq);
body_pos0.resize(3*bodies_.size());
body_quat0.resize(4*bodies_.size());
for (auto body : bodies_) {
mjuu_copyvec(bpos.data()+3*b, body->spec.pos, 3);
mjuu_copyvec(bquat.data()+4*b, body->spec.quat, 4);
mjuu_copyvec(body_pos0.data()+3*b, body->spec.pos, 3);
mjuu_copyvec(body_quat0.data()+4*b, body->spec.quat, 4);
for (auto joint : body->joints) {
switch (joint->type) {
case mjJNT_FREE:
mjuu_copyvec(q0.data()+joint->qposadr_, body->spec.pos, 3);
mjuu_copyvec(q0.data()+joint->qposadr_+3, body->spec.quat, 4);
mjuu_copyvec(qpos0.data()+joint->qposadr_, body->spec.pos, 3);
mjuu_copyvec(qpos0.data()+joint->qposadr_+3, body->spec.quat, 4);
break;
case mjJNT_BALL:
mjuu_setvec(q0.data()+joint->qposadr_, 1, 0, 0, 0);
mjuu_setvec(qpos0.data()+joint->qposadr_, 1, 0, 0, 0);
break;
case mjJNT_SLIDE:
case mjJNT_HINGE:
q0[joint->qposadr_] = (T)joint->spec.ref;
qpos0[joint->qposadr_] = (mjtNum)joint->spec.ref;
break;
default:
@@ -3811,18 +3820,9 @@ void mjCModel::ResizeKeyframe(mjCKey* key, const mjtNum* qpos0_,
// convert pending keyframes info to actual keyframes
void mjCModel::ResolveKeyframes(const mjModel* m) {
if (key_pending_.empty()) {
return;
}
// store dof offsets in joints and actuators
SaveDofOffsets();
// resize existing keyframes to the new state, fill in missing default values
for (auto* key : keys_) {
ResizeKeyframe(key, m->qpos0, m->body_pos, m->body_quat);
}
// create new keyframes, fill in missing default values
for (const auto& info : key_pending_) {
mjCKey* key = (mjCKey*)FindObject(mjOBJ_KEY, info.name);
src/user/user_model.h
+1 -3
View File
@@ -393,9 +393,7 @@ class mjCModel : public mjCModel_, private mjSpec {
void ResizeKeyframe(mjCKey* key, const mjtNum* qpos0_, const mjtNum* bpos, const mjtNum* bquat);
// compute qpos0
template <class T>
void ComputeReference(std::vector<T>& q0, std::vector<T>& bpos,
std::vector<T>& bquat);
void ComputeReference();
mjListKeyMap ids; // map from object names to ids
mjCError errInfo; // last error info
test/user/user_api_test.cc
+74 -11
View File
@@ -1489,7 +1489,7 @@ TEST_F(MujocoTest, AttachMocap) {
<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"/>
<key name="key" time="1" mpos="2 2 2" mquat="0 0 0 1"/>
</keyframe>
</mujoco>)";
@@ -1497,11 +1497,11 @@ TEST_F(MujocoTest, AttachMocap) {
<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"/>
<body pos="1 1 1" quat="0 1 0 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"/>
<key name="key" time="1" mpos="2 2 2 1 1 1" mquat="0 0 0 1 0 1 0 0"/>
<key name="attached-key-1" time="1" mpos="1 1 1 2 2 2" mquat="0 1 0 0 0 0 0 1"/>
</keyframe>
</mujoco>)";
@@ -1519,13 +1519,6 @@ TEST_F(MujocoTest, AttachMocap) {
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());
@@ -1839,6 +1832,76 @@ TEST_F(MujocoTest, RepeatedAttachKeyframe) {
mj_deleteModel(model_2);
}
TEST_F(MujocoTest, ResizeParentKeyframe) {
static constexpr char xml_parent[] = R"(
<mujoco model="MuJoCo Model">
<worldbody>
<frame name="frame"/>
<body name="body">
<joint/>
<geom size="0.1"/>
</body>
</worldbody>
<keyframe>
<key name="home" qpos="1"/>
</keyframe>
</mujoco>)";
static constexpr char xml_child[] = R"(
<mujoco model="MuJoCo Model">
<worldbody>
<body name="body">
<joint/>
<geom size="0.1"/>
</body>
</worldbody>
</mujoco>)";
static constexpr char xml_expected[] = R"(
<mujoco model="MuJoCo Model">
<worldbody>
<body name="body">
<joint/>
<geom size="0.1"/>
</body>
<frame name="frame">
<body name="child-body">
<joint/>
<geom size="0.1"/>
</body>
</frame>
</worldbody>
<keyframe>
<key name="home" qpos="1 0"/>
</keyframe>
</mujoco>)";
std::array<char, 1000> er;
mjSpec* parent = mj_parseXMLString(xml_parent, 0, er.data(), er.size());
EXPECT_THAT(parent, NotNull()) << er.data();
mjSpec* child = mj_parseXMLString(xml_child, 0, er.data(), er.size());
EXPECT_THAT(child, NotNull()) << er.data();
mjs_attachBody(mjs_findFrame(parent, "frame"), mjs_findBody(child, "body"),
"child-", "");
mjModel* model = mj_compile(parent, 0);
EXPECT_THAT(model, NotNull());
mjtNum tol = 0;
std::string field = "";
mjModel* expected = LoadModelFromString(xml_expected, er.data(), er.size());
EXPECT_THAT(expected, NotNull()) << er.data();
EXPECT_LE(CompareModel(model, expected, field), tol)
<< "Expected and attached models are different!\n"
<< "Different field: " << field << '\n';
mj_deleteSpec(parent);
mj_deleteSpec(child);
mj_deleteModel(model);
mj_deleteModel(expected);
}
TEST_F(MujocoTest, DifferentUnitsAllowed) {
mjSpec* child = mj_makeSpec();
child->compiler.degree = 1;