Added ability to give Mujoco initial conditions from the keyframes.

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();

obelisk/python/obelisk_py/core/sim/mujoco.py

@@ -9,7 +9,7 @@
 import obelisk_sensor_msgs.msg as osm
 import rclpy
 from ament_index_python.packages import get_package_share_directory
-from mujoco import MjData, MjModel, mj_forward, mj_name2id, mj_step  # type: ignore
+from mujoco import MjData, MjModel, mj_forward, mj_name2id, mj_step, mju_copy  # type: ignore
 from rclpy.callback_groups import ReentrantCallbackGroup
 from rclpy.lifecycle import LifecycleState, TransitionCallbackReturn
 from rclpy.publisher import Publisher
@@ -25,6 +25,7 @@ def __init__(self, node_name: str = "obelisk_mujoco_robot") -> None:
         """Initialize the mujoco simulator."""
         super().__init__(node_name)
         self.declare_parameter("mujoco_setting", rclpy.Parameter.Type.STRING)
+        self.declare_parameter("ic_keyframe", "ic")
 
     def _get_msg_type_from_string(self, msg_type_str: str) -> Type[ObeliskSensorMsg]:
         """Get the message type from a string.
@@ -421,6 +422,18 @@ def publish_true_sim_state(self) -> osm.TrueSimState:
 
     def run_simulator(self) -> None:
         """Run the mujoco simulator."""
+        # Set the initial condition based on a keyframe
+        self.get_logger().info(f"Found {self.mj_model.nkey} Mujoco keyframes.")
+        for i in range(self.mj_model.nkey):
+            potential_keyframe = self.mj_model.key(i).name
+            if potential_keyframe == self.get_parameter("ic_keyframe").get_parameter_value().string_value:
+                self.get_logger().info(f"Setting initial condition to keyframe: {potential_keyframe}")
+                mju_copy(self.mj_data.qpos, self.mj_model.key_qpos[i * self.mj_model.nq : (i + 1) * self.mj_model.nq])
+                mju_copy(self.mj_data.qvel, self.mj_model.key_qvel[i * self.mj_model.nv : (i + 1) * self.mj_model.nv])
+                self.mj_data.time = self.mj_model.key_time[i]
+
+                mju_copy(self.mj_data.ctrl, self.mj_model.key_ctrl[i * self.mj_model.nu : (i + 1) * self.mj_model.nu])
+
         with mujoco.viewer.launch_passive(self.mj_model, self.mj_data) as viewer:
             while viewer.is_running() and self.is_sim_running.value:
                 # simulate at realtime rate

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:

obelisk_ws/src/obelisk_ros/config/dummy_py.yaml

@@ -54,6 +54,8 @@ onboard:
     - is_simulated: True
       pkg: obelisk_sim_py
       executable: obelisk_mujoco_robot
+      params:
+        ic_keyframe: other
       # callback_groups:
       # publishers:
       subscribers:

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>