90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
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
|