From efce7f39e5fa3cead4416c9063781dc488286847 Mon Sep 17 00:00:00 2001 From: Masaki Murooka Date: Sat, 24 Jun 2023 10:31:21 +0900 Subject: [PATCH] Change the way data->qpos and data->qvel are initialized. Whereas previously the joint ids was obtained implicitly by push_back in sequence, this commit uses mj_jnt_ids and mj_jnt_to_rjo to obtain the ids of the joint. This allows handling of cases where the number of joints differs between the mc_rtc robot and mjData. Such a case occurs in the simulation of a soft body (composite body). --- src/mj_sim.cpp | 94 ++++++++++++++++++++--------------------------- src/mj_sim_impl.h | 6 --- src/mj_utils.cpp | 19 ---------- src/mj_utils.h | 3 -- 4 files changed, 39 insertions(+), 83 deletions(-) diff --git a/src/mj_sim.cpp b/src/mj_sim.cpp index 84586cc..c88e639 100644 --- a/src/mj_sim.cpp +++ b/src/mj_sim.cpp @@ -281,6 +281,12 @@ void MjRobot::initialize(mjModel * model, const mc_rbdyn::Robot & robot) { root_body_id = mj_name2id(model, mjOBJ_BODY, root_body.c_str()); } + if(root_joint.size()) + { + auto root_joint_id = mj_name2id(model, mjOBJ_JOINT, root_joint.c_str()); + root_qpos_idx = model->jnt_qposadr[root_joint_id]; + root_qvel_idx = model->jnt_dofadr[root_joint_id]; + } auto init_sensor_id = [&](const char * mj_name, const char * mc_name, const std::string & sensor_name, const char * suffix, mjtSensor type, std::unordered_map & mapping) { auto mj_sensor = prefixed(fmt::format("{}_{}", sensor_name, suffix)); @@ -386,31 +392,23 @@ void MjSimImpl::setSimulationInitialState() { if(controller) { - qInit.resize(0); - alphaInit.resize(0); - for(auto & o : objects) { o.initialize(model); - auto start_q = 0; - auto start_dof = 0; if(o.root_qpos_idx != -1 && o.root_joint_type == mjJNT_FREE) { - sva::PTransformd pose = o.init_pose; - const auto & t = pose.translation(); - for(size_t i = 0; i < 3; ++i) - { - qInit.push_back(t[i]); - alphaInit.push_back(0); - alphaInit.push_back(0); - } - Eigen::Quaterniond q = Eigen::Quaterniond(pose.rotation()).inverse(); - qInit.push_back(q.w()); - qInit.push_back(q.x()); - qInit.push_back(q.y()); - qInit.push_back(q.z()); - start_q = 7; - start_dof = 6; + const auto & t = o.init_pose.translation(); + data->qpos[o.root_qpos_idx + 0] = t.x(); + data->qpos[o.root_qpos_idx + 1] = t.y(); + data->qpos[o.root_qpos_idx + 2] = t.z(); + Eigen::Quaterniond q = Eigen::Quaterniond(o.init_pose.rotation()).inverse(); + data->qpos[o.root_qpos_idx + 3] = q.w(); + data->qpos[o.root_qpos_idx + 4] = q.x(); + data->qpos[o.root_qpos_idx + 5] = q.y(); + data->qpos[o.root_qpos_idx + 6] = q.z(); + // push linear/angular velocities + mju_zero3(&data->qvel[o.root_qvel_idx]); + mju_zero3(&data->qvel[o.root_qvel_idx + 3]); } else if(o.root_body_id != -1) { @@ -424,40 +422,26 @@ void MjSimImpl::setSimulationInitialState() model->body_quat[4 * o.root_body_id + 2] = q.y(); model->body_quat[4 * o.root_body_id + 3] = q.z(); } - for(int i = start_q; i < o.nq; ++i) - { - qInit.push_back(0.0); - } - for(int i = start_dof; i < o.ndof; ++i) - { - alphaInit.push_back(0.0); - } } for(auto & r : robots) { const auto & robot = controller->robots().robot(r.name); r.initialize(model, robot); - if(r.root_joint.size()) + if(r.root_qpos_idx != -1) { - r.root_qpos_idx = qInit.size(); - r.root_qvel_idx = alphaInit.size(); - if(robot.mb().joint(0).dof() == 6) - { - const auto & t = robot.posW().translation(); - for(size_t i = 0; i < 3; ++i) - { - qInit.push_back(t[i]); - // push linear/angular velocities - alphaInit.push_back(0); - alphaInit.push_back(0); - } - Eigen::Quaterniond q = Eigen::Quaterniond(robot.posW().rotation()).inverse(); - qInit.push_back(q.w()); - qInit.push_back(q.x()); - qInit.push_back(q.y()); - qInit.push_back(q.z()); - } + const auto & t = robot.posW().translation(); + data->qpos[r.root_qpos_idx + 0] = t.x(); + data->qpos[r.root_qpos_idx + 1] = t.y(); + data->qpos[r.root_qpos_idx + 2] = t.z(); + Eigen::Quaterniond q = Eigen::Quaterniond(robot.posW().rotation()).inverse(); + data->qpos[r.root_qpos_idx + 3] = q.w(); + data->qpos[r.root_qpos_idx + 4] = q.x(); + data->qpos[r.root_qpos_idx + 5] = q.y(); + data->qpos[r.root_qpos_idx + 6] = q.z(); + // push linear/angular velocities + mju_zero3(&data->qvel[r.root_qvel_idx]); + mju_zero3(&data->qvel[r.root_qvel_idx + 3]); } else if(r.root_body_id != -1) { @@ -471,18 +455,18 @@ void MjSimImpl::setSimulationInitialState() model->body_quat[4 * r.root_body_id + 2] = q.y(); model->body_quat[4 * r.root_body_id + 3] = q.z(); } - for(size_t i = 0; i < r.mj_jnt_names.size(); ++i) + for(size_t i = 0; i < r.mj_jnt_ids.size(); ++i) { - qInit.push_back(r.encoders[r.mj_jnt_to_rjo[i]]); - alphaInit.push_back(r.alphas[r.mj_jnt_to_rjo[i]]); + if(r.mj_jnt_to_rjo[i] == -1) + { + continue; + } + data->qpos[model->jnt_qposadr[r.mj_jnt_ids[i]]] = r.encoders[r.mj_jnt_to_rjo[i]]; + data->qvel[model->jnt_dofadr[r.mj_jnt_ids[i]]] = r.alphas[r.mj_jnt_to_rjo[i]]; } } } - // set initial qpos, qvel in mujoco - if(!mujoco_set_const(model, data, qInit, alphaInit)) - { - mc_rtc::log::error_and_throw("[mc_mujoco] Set inital state failed."); - } + mj_forward(model, data); } diff --git a/src/mj_sim_impl.h b/src/mj_sim_impl.h index 064ef93..900a643 100644 --- a/src/mj_sim_impl.h +++ b/src/mj_sim_impl.h @@ -200,12 +200,6 @@ struct MjSimImpl /** MuJoCo data */ mjData * data = nullptr; - /** Initial state */ - std::vector qInit; - - /** Initial velocity */ - std::vector alphaInit; - #ifndef USE_UI_ADAPTER /** GLFW window, might be null if the visualization is disabled */ GLFWwindow * window = nullptr; diff --git a/src/mj_utils.cpp b/src/mj_utils.cpp index fb32504..0071be5 100644 --- a/src/mj_utils.cpp +++ b/src/mj_utils.cpp @@ -477,25 +477,6 @@ void mujoco_create_window(MjSimImpl * mj_sim) ImGui_ImplOpenGL3_Init(glsl_version); } -bool mujoco_set_const(mjModel * m, mjData * d, const std::vector & qpos, const std::vector & qvel) -{ - if(qpos.size() != m->nq || qvel.size() != m->nv) - { - std::cerr << "qpos size: " << qpos.size() << ". Should be: " << m->nq << std::endl; - std::cerr << "qvel size: " << qvel.size() << ". Should be: " << m->nv << std::endl; - return false; - } - - mj_setConst(m, d); - const double * qpos_init = &qpos[0]; - const double * qvel_init = &qvel[0]; - mju_copy(d->qpos, qpos_init, m->nq); - mju_copy(d->qvel, qvel_init, m->nv); - d->time = 0.0; - mj_forward(m, d); - return true; -} - void mujoco_cleanup(MjSimImpl * mj_sim) { if(mj_sim->config.with_visualization) diff --git a/src/mj_utils.h b/src/mj_utils.h index 3b925a0..3376fdd 100644 --- a/src/mj_utils.h +++ b/src/mj_utils.h @@ -41,9 +41,6 @@ bool mujoco_init(MjSimImpl * mj_sim, /*! Create GLFW window */ void mujoco_create_window(MjSimImpl * mj_sim); -/*! Sets initial qpos and qvel in mjData */ -bool mujoco_set_const(mjModel * m, mjData * d, const std::vector & qpos, const std::vector & qvel); - /** Returns a sensor id from name, -1 if the type does not match or the sensor does not exist */ int mujoco_get_sensor_id(const mjModel & m, const std::string & name, mjtSensor type);