added ability to start from mujoco key frames

Caltech-AMBER · Jul 26, 2024 · 8522fbc · 8522fbc
1 parent 9d4ed48
commit 8522fbc
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 0 deletions.
diff --git a/obelisk/cpp/obelisk_cpp/include/obelisk_mujoco_sim_robot.h b/obelisk/cpp/obelisk_cpp/include/obelisk_mujoco_sim_robot.h
@@ -17,6 +17,9 @@ namespace obelisk {
       public:
         explicit ObeliskMujocoRobot(const std::string& name) : ObeliskSimRobot<ControlMessageT>(name) {
             this->template declare_parameter<std::string>("mujoco_setting", "");
+            this->template declare_parameter<std::string>(
+                "ic_keyframe",
+                "ic"); // Parameter is the name of a mujoco key frame used as the initial condition. Defualts to "ic".
 
             mujoco_sim_instance_    = this;
             activation_complete_    = false;
@@ -199,6 +202,21 @@ namespace obelisk {
             }
             RCLCPP_INFO_STREAM(this->get_logger(), "Starting Mujoco simulation loop.");
 
+            // Set the initial condition based on a keyframe
+            RCLCPP_INFO_STREAM(this->get_logger(), "Found " << this->model_->nkey << " Mujoco keyframes.");
+            for (int i = 0; i < this->model_->nkey; i++) {
+                std::string potential_keyframe(this->model_->names + this->model_->name_keyadr[i]);
+                if (potential_keyframe == this->get_parameter("ic_keyframe").as_string()) {
+                    RCLCPP_INFO_STREAM(this->get_logger(),
+                                       "Setting initial condition to keyframe: " << potential_keyframe);
+                    mju_copy(this->data_->qpos, &this->model_->key_qpos[i * this->model_->nq], this->model_->nq);
+                    mju_copy(this->data_->qvel, &this->model_->key_qvel[i * this->model_->nv], this->model_->nv);
+                    this->data_->time = this->model_->key_time[i];
+
+                    mju_copy(this->data_->ctrl, &this->model_->key_ctrl[i * this->model_->nu], this->model_->nu);
+                }
+            }
+
             while (!this->stop_thread_) {
                 if (!glfwWindowShouldClose(window_)) {
                     auto start_time = std::chrono::steady_clock::now();

diff --git a/obelisk_ws/src/obelisk_ros/config/dummy_cpp.yaml b/obelisk_ws/src/obelisk_ros/config/dummy_cpp.yaml
@@ -62,6 +62,8 @@ onboard:
     - is_simulated: True
       pkg: obelisk_sim_cpp
       executable: obelisk_mujoco_robot
+      params:
+        ic_keyframe: ic
       # callback_groups:
       # publishers:
       subscribers:

diff --git a/obelisk_ws/src/robots/dummy_bot/mujoco/dummy.xml b/obelisk_ws/src/robots/dummy_bot/mujoco/dummy.xml
@@ -50,4 +50,9 @@
         <framepos name="tip_pos_sensor" objtype="site" objname="tip"/>
         <framequat name="tip_orientation_sensor" objtype="site" objname="tip"/>
     </sensor>
+
+    <keyframe>
+        <key name="other" qpos="-1" ctrl="-1" time="1"/>
+        <key name="ic" qpos="1" ctrl="1" time="3" />
+    </keyframe>
 </mujoco>