Init

behaviors/custom/keyframe/keyframe.py

@@ -0,0 +1,89 @@
+from behaviors.behavior import Behavior
+import yaml
+
+
+class KeyframeSkill(Behavior):
+    def __init__(self, agent, file: str):
+        super().__init__(agent)
+
+        self.keyframe_step: int = 0
+        self.keyframe_start_time: float = None
+
+        with open(f"{file}", "r") as f:
+            self.skill_description = yaml.safe_load(f)
+
+        self.has_symmetry: bool = self.skill_description["symmetry"]
+        self.keyframes: list = self.skill_description["keyframes"]
+        self.kp: float = self.skill_description["kp"]
+        self.kd: float = self.skill_description["kd"]
+
+    def execute(self, reset, *args, **kwargs):
+        if reset:
+            self.keyframe_step = 0
+            self.keyframe_start_time = self.agent.world.server_time
+
+        current_keyframe = self.keyframes[self.keyframe_step]
+
+        raw_motor_positions = current_keyframe["motor_positions"]
+
+        current_keyframe_kp = current_keyframe.get("kp", self.kp)
+        current_keyframe_kd = current_keyframe.get("kd", self.kd)
+
+        p_gains = current_keyframe.get("p_gains", None)
+        d_gains = current_keyframe.get("d_gains", None)
+
+        for readable_motor_name, position in raw_motor_positions.items():
+
+            p_gain = (
+                p_gains.get(readable_motor_name, current_keyframe_kp)
+                if p_gains
+                else current_keyframe_kp
+            )
+            d_gain = (
+                d_gains.get(readable_motor_name, current_keyframe_kd)
+                if d_gains
+                else current_keyframe_kd
+            )
+
+            if self.has_symmetry:
+                motor_names, is_inverse_direction = self.agent.robot.MOTOR_SYMMETRY[
+                    readable_motor_name
+                ]
+                for motor_name in motor_names:
+                    server_motor_name = self.agent.robot.MOTOR_FROM_READABLE_TO_SERVER[
+                        motor_name
+                    ]
+
+                    self.agent.robot.set_motor_target_position(
+                        motor_name=server_motor_name,
+                        target_position=position if is_inverse_direction else -position,
+                        kp=p_gain,
+                        kd=d_gain,
+                    )
+                    if is_inverse_direction:
+                        is_inverse_direction = False  # Only inverts one joint
+            else:
+                server_motor_name = self.agent.robot.MOTOR_FROM_READABLE_TO_SERVER[
+                    readable_motor_name
+                ]
+                self.agent.robot.set_motor_target_position(
+                    motor_name=server_motor_name,
+                    target_position=position,
+                    kp=self.kp,
+                    kd=self.kd,
+                )
+
+        keyframe_time: float = current_keyframe["delta"]
+        elapsed_keyframe_time = self.agent.world.server_time - self.keyframe_start_time
+
+        if elapsed_keyframe_time >= keyframe_time:
+            self.keyframe_start_time = self.agent.world.server_time
+            self.keyframe_step += 1
+
+            if self.keyframe_step >= len(self.keyframes):
+                return True
+
+        return False
+
+    def is_ready(self, *args):
+        return True