Skip to content

Commit

Permalink
[state] Changed getRootDim API and fixed details in wb tests
Browse files Browse the repository at this point in the history
  • Loading branch information
cmastalli committed Apr 3, 2024
1 parent de6ffe1 commit 96b3125
Show file tree
Hide file tree
Showing 8 changed files with 50 additions and 34 deletions.
30 changes: 21 additions & 9 deletions include/crocoddyl_msgs/conversions.h
Original file line number Diff line number Diff line change
Expand Up @@ -90,12 +90,24 @@ static inline std::size_t getRootJointId(
}

/**
* @brief Return the root dimension
* @brief Return the root nq dimension
*
* @param return Root joint dimension
* @param return Root joint nq dimension
*/
template <int Options, template <typename, int> class JointCollectionTpl>
static inline std::size_t getRootDim(
static inline std::size_t getRootNq(
const pinocchio::ModelTpl<double, Options, JointCollectionTpl> &model) {
const std::size_t root_joint_id = getRootJointId(model);
return model.frames[root_joint_id].name != "universe" ? model.joints[root_joint_id].nq() : 0;
}

/**
* @brief Return the root nv dimension
*
* @param return Root joint nv dimension
*/
template <int Options, template <typename, int> class JointCollectionTpl>
static inline std::size_t getRootNv(
const pinocchio::ModelTpl<double, Options, JointCollectionTpl> &model) {
const std::size_t root_joint_id = getRootJointId(model);
return model.frames[root_joint_id].name != "universe" ? model.joints[root_joint_id].nv() : 0;
Expand Down Expand Up @@ -215,8 +227,8 @@ static inline void toMsg(
" but received " + std::to_string(a.size()));
}
const std::size_t root_joint_id = getRootJointId(model);
const std::size_t nq_root = model.frames[root_joint_id].name != "universe" ? model.joints[root_joint_id].nq() : 0;
const std::size_t nv_root = model.frames[root_joint_id].name != "universe" ? model.joints[root_joint_id].nv() : 0;
const std::size_t nq_root = getRootNq(model);
const std::size_t nv_root = getRootNv(model);
const std::size_t njoints = model.nv - nv_root;
if (tau.size() != static_cast<int>(njoints) && tau.size() != 0) {
throw std::invalid_argument("Expected tau to be 0 or " +
Expand Down Expand Up @@ -689,8 +701,8 @@ static inline void fromReduced(
const Eigen::Ref<const Eigen::VectorXd> &qref,
const std::vector<pinocchio::JointIndex> &locked_joint_ids) {
const std::size_t root_joint_id = getRootJointId(model);
const std::size_t nq_root = model.frames[root_joint_id].name != "universe" ? model.joints[root_joint_id].nq() : 0;
const std::size_t nv_root = model.frames[root_joint_id].name != "universe" ? model.joints[root_joint_id].nv() : 0;
const std::size_t nq_root = getRootNq(model);
const std::size_t nv_root = getRootNv(model);
const std::size_t njoints = model.nv - nv_root;
const std::size_t njoints_reduced = reduced_model.nv - nv_root;
if (q_out.size() != model.nq) {
Expand Down Expand Up @@ -770,8 +782,8 @@ toReduced(const pinocchio::ModelTpl<double, Options, JointCollectionTpl> &model,
const Eigen::Ref<const Eigen::VectorXd> &v_in,
const Eigen::Ref<const Eigen::VectorXd> &tau_in) {
const std::size_t root_joint_id = getRootJointId(model);
const std::size_t nq_root = model.frames[root_joint_id].name != "universe" ? model.joints[root_joint_id].nq() : 0;
const std::size_t nv_root = model.frames[root_joint_id].name != "universe" ? model.joints[root_joint_id].nv() : 0;
const std::size_t nq_root = getRootNq(model);
const std::size_t nv_root = getRootNv(model);
const std::size_t njoints = model.nv - nv_root;
const std::size_t njoints_reduced = reduced_model.nv - nv_root;
if (q_out.size() != reduced_model.nq) {
Expand Down
2 changes: 1 addition & 1 deletion include/crocoddyl_msgs/whole_body_state_publisher.h
Original file line number Diff line number Diff line change
Expand Up @@ -183,7 +183,7 @@ class WholeBodyStateRosPublisher {
}
pinocchio::buildReducedModel(model_, joint_ids_, qref_, reduced_model_);
// Initialize the vectors and dimensions
const std::size_t nv_root = getRootDim(model_);
const std::size_t nv_root = getRootNv(model_);
qfull_ = Eigen::VectorXd::Zero(model_.nq);
vfull_ = Eigen::VectorXd::Zero(model_.nv);
ufull_ = Eigen::VectorXd::Zero(model_.nv - nv_root);
Expand Down
2 changes: 1 addition & 1 deletion include/crocoddyl_msgs/whole_body_state_subscriber.h
Original file line number Diff line number Diff line change
Expand Up @@ -221,7 +221,7 @@ class WholeBodyStateRosSubscriber {
f_tmp_;

void init(const std::vector<std::string> &locked_joints = DEFAULT_VECTOR) {
const std::size_t nv_root = getRootDim(model_);
const std::size_t nv_root = getRootNv(model_);
if (locked_joints.size() != 0) {
// Check the size of the reference configuration
if (qref_.size() != model_.nq) {
Expand Down
2 changes: 1 addition & 1 deletion include/crocoddyl_msgs/whole_body_trajectory_publisher.h
Original file line number Diff line number Diff line change
Expand Up @@ -253,7 +253,7 @@ class WholeBodyTrajectoryRosPublisher {

// Initialize the vectors and dimensions
const std::size_t root_joint_id = getRootJointId(model_);
const std::size_t nv_root = getRootDim(model_);
const std::size_t nv_root = getRootNv(model_);
qfull_ = Eigen::VectorXd::Zero(model_.nq);
vfull_ = Eigen::VectorXd::Zero(model_.nv);
ufull_ = Eigen::VectorXd::Zero(model_.nv - nv_root);
Expand Down
2 changes: 1 addition & 1 deletion include/crocoddyl_msgs/whole_body_trajectory_subscriber.h
Original file line number Diff line number Diff line change
Expand Up @@ -254,7 +254,7 @@ class WholeBodyTrajectoryRosSubscriber {
spinner_.start();
#endif

const std::size_t nv_root = getRootDim(model_);
const std::size_t nv_root = getRootNv(model_);
if (locked_joints.size() != 0) {
// Check the size of the reference configuration
if (qref_.size() != model_.nq) {
Expand Down
11 changes: 8 additions & 3 deletions src/crocoddyl_ros.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -250,10 +250,15 @@ PYBIND11_MODULE(crocoddyl_ros, m) {
":param model: Pinocchio model\n"
":return root joint id");

m.def("getRootDim", &getRootDim<0, pinocchio::JointCollectionDefaultTpl>,
"Return the root joint dimension.\n\n"
m.def("getRootNq", &getRootNq<0, pinocchio::JointCollectionDefaultTpl>,
"Return the root joint nq dimension.\n\n"
":param model: Pinocchio model\n"
":return root joint dimension");
":return root joint nq dimension");

m.def("getRootNv", &getRootNv<0, pinocchio::JointCollectionDefaultTpl>,
"Return the root joint nv dimension.\n\n"
":param model: Pinocchio model\n"
":return root joint nv dimension");

m.def(
"fromReduced",
Expand Down
18 changes: 8 additions & 10 deletions unittest/test_whole_body_state.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
WholeBodyStateRosPublisher,
WholeBodyStateRosSubscriber,
toReduced,
getRootDim,
getRootNv,
)


Expand Down Expand Up @@ -89,7 +89,7 @@ def test_publisher_without_contact(self):
pub = WholeBodyStateRosPublisher(self.MODEL, "whole_body_state_without_contact")
time.sleep(1)
# publish whole-body state messages
nv_root = getRootDim(self.MODEL)
nv_root = getRootNv(self.MODEL)
q = pinocchio.randomConfiguration(self.MODEL)
v = np.random.rand(self.MODEL.nv)
tau = np.random.rand(self.MODEL.nv - nv_root)
Expand All @@ -111,7 +111,7 @@ def test_communication(self):
pub = WholeBodyStateRosPublisher(self.MODEL, "whole_body_state")
time.sleep(1)
# publish whole-body state messages
nv_root = getRootDim(self.MODEL)
nv_root = getRootNv(self.MODEL)
q = pinocchio.randomConfiguration(self.MODEL)
v = np.random.rand(self.MODEL.nv)
tau = np.random.rand(self.MODEL.nv - nv_root)
Expand Down Expand Up @@ -158,19 +158,17 @@ def test_communication(self):
)

def test_communication_with_reduced_model(self):
# locked_joints = ["larm_elbow_joint", "rarm_elbow_joint"]
locked_joints = ["wrist2_joint"]
qref = pinocchio.randomConfiguration(self.MODEL)
reduced_model = pinocchio.buildReducedModel(
self.MODEL, [self.MODEL.getJointId(name) for name in locked_joints], qref
self.MODEL, [self.MODEL.getJointId(name) for name in self.LOCKED_JOINTS], qref
)
sub = WholeBodyStateRosSubscriber(
self.MODEL, locked_joints, qref, "reduced_whole_body_state"
self.MODEL, self.LOCKED_JOINTS, qref, "reduced_whole_body_state"
)
pub = WholeBodyStateRosPublisher(self.MODEL, locked_joints, qref, "reduced_whole_body_state")
pub = WholeBodyStateRosPublisher(self.MODEL, self.LOCKED_JOINTS, qref, "reduced_whole_body_state")
time.sleep(1)
# publish whole-body state messages
nv_root = getRootDim(self.MODEL)
nv_root = getRootNv(self.MODEL)
q = pinocchio.randomConfiguration(self.MODEL)
v = np.random.rand(self.MODEL.nv)
tau = np.random.rand(self.MODEL.nv - nv_root)
Expand Down Expand Up @@ -229,7 +227,7 @@ def test_communication_with_non_locked_joints(self):
pub = WholeBodyStateRosPublisher(self.MODEL, locked_joints, qref, "non_locked_whole_body_state")
time.sleep(1)
# publish whole-body state messages
nv_root = getRootDim(self.MODEL)
nv_root = getRootNv(self.MODEL)
q = pinocchio.randomConfiguration(self.MODEL)
v = np.random.rand(self.MODEL.nv)
tau = np.random.rand(self.MODEL.nv - nv_root)
Expand Down
17 changes: 9 additions & 8 deletions unittest/test_whole_body_trajectory.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
WholeBodyTrajectoryRosPublisher,
WholeBodyTrajectoryRosSubscriber,
toReduced,
getRootDim,
getRootNv,
)


Expand Down Expand Up @@ -103,7 +103,7 @@ def test_publisher_without_contact(self):
# publish whole-body trajectory messages
N = len(self.ts)
xs = []
nv_root = getRootDim(self.MODEL)
nv_root = getRootNv(self.MODEL)
for _ in range(N):
q = pinocchio.randomConfiguration(self.MODEL)
q[:3] = np.random.rand(3)
Expand Down Expand Up @@ -135,7 +135,7 @@ def test_communication(self):
q[:3] = np.random.rand(3)
v = np.random.rand(self.MODEL.nv)
xs.append(np.hstack([q, v]))
nv_root = getRootDim(self.MODEL)
nv_root = getRootNv(self.MODEL)
us = [np.random.rand(self.MODEL.nv - nv_root) for _ in range(N)]
while True:
pub.publish(self.ts, xs, us, self.ps, self.pds, self.fs, self.ss)
Expand Down Expand Up @@ -200,7 +200,7 @@ def test_communication_with_reduced_model(self):
# publish whole-body trajectory messages
N = len(self.ts)
xs, us = [], []
nv_root = getRootDim(self.MODEL)
nv_root = getRootNv(self.MODEL)
for _ in range(N):
q = pinocchio.randomConfiguration(self.MODEL)
q[:3] = np.random.rand(3)
Expand Down Expand Up @@ -258,21 +258,22 @@ def test_communication_with_reduced_model(self):
)

def test_communication_with_non_locked_joints(self):
locked_joints = []
qref = pinocchio.randomConfiguration(self.MODEL)
reduced_model = pinocchio.buildReducedModel(
self.MODEL, [self.MODEL.getJointId(name) for name in self.LOCKED_JOINTS], qref
self.MODEL, [self.MODEL.getJointId(name) for name in locked_joints], qref
)
sub = WholeBodyTrajectoryRosSubscriber(
self.MODEL, self.LOCKED_JOINTS, qref, "non_locked_whole_body_trajectory"
self.MODEL, locked_joints, qref, "non_locked_whole_body_trajectory"
)
pub = WholeBodyTrajectoryRosPublisher(
self.MODEL, self.LOCKED_JOINTS, qref, "non_locked_whole_body_trajectory"
self.MODEL, locked_joints, qref, "non_locked_whole_body_trajectory"
)
time.sleep(1)
# publish whole-body trajectory messages
N = len(self.ts)
xs, us = [], []
nv_root = getRootDim(self.MODEL)
nv_root = getRootNv(self.MODEL)
for _ in range(N):
q = pinocchio.randomConfiguration(self.MODEL)
q[:3] = np.random.rand(3)
Expand Down

0 comments on commit 96b3125

Please sign in to comment.