From cab6196d5ce72f4fa13abffeb8c112268ae85299 Mon Sep 17 00:00:00 2001
From: Zachary <zach.olkin@gmail.com>
Date: Fri, 26 Jul 2024 17:29:39 +0000
Subject: [PATCH] added ic keyframe support in python

---
 obelisk/python/obelisk_py/core/sim/mujoco.py  | 19 +++++++++++++++++++
 .../src/obelisk_ros/config/dummy_py.yaml      |  2 ++
 2 files changed, 21 insertions(+)

diff --git a/obelisk/python/obelisk_py/core/sim/mujoco.py b/obelisk/python/obelisk_py/core/sim/mujoco.py
index d2fabcbc..fd1168bc 100644
--- a/obelisk/python/obelisk_py/core/sim/mujoco.py
+++ b/obelisk/python/obelisk_py/core/sim/mujoco.py
@@ -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,24 @@ 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}")
+                mujoco.mju_copy(
+                    self.mj_data.qpos, self.mj_model.key_qpos[i * self.mj_model.nq : (i + 1) * self.mj_model.nq]
+                )
+                mujoco.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]
+
+                mujoco.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
diff --git a/obelisk_ws/src/obelisk_ros/config/dummy_py.yaml b/obelisk_ws/src/obelisk_ros/config/dummy_py.yaml
index 61457aab..6104e96a 100644
--- a/obelisk_ws/src/obelisk_ros/config/dummy_py.yaml
+++ b/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: