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

Attach multiple keyframes.

Browse Source 
Still missing: support for ctrl, mocap, and time.

PiperOrigin-RevId: 662871569
Change-Id: I8a418c4007663e3063d7146bd72db4eba78ddc5d
This commit is contained in:
Alessio Quaglino
2024-08-14 04:41:51 -07:00
committed by Copybara-Service
parent 2cb601c98d
commit afd7c73f44
6 changed files with 73 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+2
View File
@@ -9,6 +9,8 @@ General
^^^^^^^
1. Add :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.
Version 3.2.2 (Aug 8, 2024)
---------------------------
src/user/user_model.cc
+19 -12
View File
@@ -109,6 +109,9 @@ mjCModel::mjCModel() {
// this class allocated the plugins
plugin_owner = true;
// default state name
state_name_ = "state";
}
@@ -2862,14 +2865,17 @@ void mjCModel::CopyObjects(mjModel* m) {
template <class T>
void mjCModel::SaveState(const T* qpos, const T* qvel, const T* act) {
for (auto joint : joints_) {
if (qpos) mjuu_copyvec(joint->qpos, qpos + joint->qposadr_, joint->nq());
if (qvel) mjuu_copyvec(joint->qvel, qvel + joint->dofadr_, joint->nv());
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());
}
for (auto actuator : actuators_) {
if (actuator->actadr_ != -1 && act) {
actuator->act.assign(actuator->actdim_, 0);
mjuu_copyvec(actuator->act.data(), act + actuator->actadr_, actuator->actdim_);
if (actuator->actadr_ != -1 && actuator->actdim_ != -1 && act) {
actuator->act().assign(actuator->actdim_, 0);
mjuu_copyvec(actuator->act().data(), act + actuator->actadr_, actuator->actdim_);
}
}
}
@@ -2893,21 +2899,21 @@ template <class T>
void mjCModel::RestoreState(const mjtNum* pos0, T* qpos, T* qvel, T* act) {
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()[0])) {
mjuu_copyvec(qpos + joint->qposadr_, joint->qpos(), 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()[0]) && qvel) {
mjuu_copyvec(qvel + joint->dofadr_, joint->qvel(), joint->nv());
}
}
// restore act
for (auto actuator : actuators_) {
if (mjuu_defined(actuator->act[0]) && act) {
mjuu_copyvec(act + actuator->actadr_, actuator->act.data(), actuator->actdim_);
if (!actuator->act().empty() && mjuu_defined(actuator->act()[0]) && act) {
mjuu_copyvec(act + actuator->actadr_, actuator->act().data(), actuator->actdim_);
}
}
}
@@ -2941,8 +2947,8 @@ void mjCModel::StoreKeyframes() {
info.qvel = !key->spec_qvel_.empty();
info.act = !key->spec_act_.empty();
key_pending_.push_back(info);
state_name_ = info.name;
SaveState(key->spec_qpos_.data(), key->spec_qvel_.data(), key->spec_act_.data());
break; // (b/350784262) save only the first keyframe for now
}
if (resetlists) {
@@ -3512,6 +3518,7 @@ void mjCModel::ResolveKeyframes(const mjModel* m) {
if (info.qpos) key->spec_qpos_.assign(nq, 0);
if (info.qvel) key->spec_qvel_.assign(nv, 0);
if (info.act) key->spec_act_.assign(na, 0);
state_name_ = info.name;
RestoreState(m->qpos0, key->spec_qpos_.data(), key->spec_qvel_.data(), key->spec_act_.data());
}
src/user/user_model.h
+2
View File
@@ -381,5 +381,7 @@ 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_
src/user/user_objects.cc
+34 -7
View File
@@ -1408,8 +1408,8 @@ void mjCBody::ComputeBVH() {
// reset keyframe references for allowing self-attach
void mjCBody::ForgetKeyframes() const {
for (auto joint : joints) {
joint->qpos[0] = mjNAN;
joint->qvel[0] = mjNAN;
joint->qpos_.clear();
joint->qvel_.clear();
}
for (auto body : bodies) {
body->ForgetKeyframes();
@@ -1769,8 +1769,8 @@ mjCJoint::mjCJoint(mjCModel* _model, mjCDef* _def) {
CopyFromSpec();
// no previous state when a joint is created
qpos[0] = mjNAN;
qvel[0] = mjNAN;
qposadr_ = -1;
dofadr_ = -1;
}
@@ -1828,6 +1828,24 @@ 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};
}
return qpos_.at(model->state_name_).data();
}
mjtNum* mjCJoint::qvel() {
if (qvel_.find(model->state_name_) == qvel_.end()) {
qvel_[model->state_name_] = {mjNAN, 0, 0, 0, 0, 0};
}
return qvel_.at(model->state_name_).data();
}
void mjCJoint::PointToLocal() {
spec.element = static_cast<mjsElement*>(this);
spec.name = &name;
@@ -5349,7 +5367,8 @@ mjCActuator::mjCActuator(mjCModel* _model, mjCDef* _def) {
PointToLocal();
// no previous state when an actuator is created
act.push_back(mjNAN);
actadr_ = -1;
actdim_ = -1;
}
@@ -5374,8 +5393,7 @@ mjCActuator& mjCActuator::operator=(const mjCActuator& other) {
void mjCActuator::ForgetKeyframes() {
act.clear();
act.push_back(mjNAN);
act_.clear();
}
@@ -5386,6 +5404,15 @@ 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);
}
return act_.at(model->state_name_);
}
void mjCActuator::PointToLocal() {
spec.element = static_cast<mjsElement*>(this);
spec.name = &name;
src/user/user_objects.h
+13 -7
View File
@@ -16,6 +16,7 @@
#define MUJOCO_SRC_USER_USER_OBJECTS_H_
#include <cstddef>
#include <array>
#include <cstdlib>
#include <functional>
#include <map>
@@ -401,10 +402,10 @@ class mjCJoint_ : public mjCBase {
mjCBody* body; // joint's body
// variable used for temporarily storing the state of the joint
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
int qposadr_; // address of dof in data->qpos
int dofadr_; // address of dof in data->qvel
std::map<std::string, std::array<mjtNum, 7>> qpos_; // qpos at the previous step
std::map<std::string, std::array<mjtNum, 6>> qvel_; // qvel at the previous step
// variable-size data
std::vector<double> userdata_;
@@ -443,6 +444,9 @@ class mjCJoint : public mjCJoint_, private mjsJoint {
int nq() const { return nq(spec.type); }
int nv() const { return nv(spec.type); }
mjtNum* qpos();
mjtNum* qvel();
private:
int Compile(void); // compiler; return dofnum
void PointToLocal(void);
@@ -1390,9 +1394,9 @@ class mjCActuator_ : public mjCBase {
int trnid[2]; // id of transmission target
// variable used for temporarily storing the state of the actuator
int actadr_; // address of dof in data->act
int actdim_; // number of dofs in data->act
std::vector<mjtNum> act; // act at the previous step
int actadr_; // address of dof in data->act
int actdim_; // number of dofs in data->act
std::map<std::string, std::vector<mjtNum>> act_; // act at the previous step
// variable-size data
std::string plugin_name;
@@ -1431,6 +1435,8 @@ class mjCActuator : public mjCActuator_, private mjsActuator {
bool is_forcelimited() const;
bool is_actlimited() const;
std::vector<mjtNum>& act();
private:
void Compile(void); // compiler
void CopyFromSpec();
test/user/user_api_test.cc
+3
View File
@@ -504,6 +504,7 @@ TEST_F(MujocoTest, AttachSame) {
<key name="two" qpos="2 0" act="2 2 0 0"/>
<key name="three" qpos="3 0" act="3 3 0 0"/>
<key name="attached-two-1" qpos="0 2" act="0 0 2 2"/>
<key name="attached-three-1" qpos="0 3" act="0 0 3 3"/>
</keyframe>
</mujoco>)";
@@ -621,6 +622,7 @@ TEST_F(MujocoTest, AttachDifferent) {
<keyframe>
<key name="attached-two-1" qpos="0 0 0 1 0 0 0 2" act="2 2"/>
<key name="attached-three-1" qpos="0 0 0 1 0 0 0 3" act="3 3"/>
</keyframe>
</mujoco>)";
@@ -742,6 +744,7 @@ TEST_F(MujocoTest, AttachFrame) {
<keyframe>
<key name="attached-two-1" qpos="0 0 0 1 0 0 0 2" act="2 2"/>
<key name="attached-three-1" qpos="0 0 0 1 0 0 0 3" act="3 3"/>
</keyframe>
</mujoco>)";