robots/g1/controllers/keyframe.py

"""Keyframe Animation Controller for G1 robot."""
import numpy as np
from .base import BaseController


class KeyframeController(BaseController):
    """
    Controller that plays back keyframe animations with PD tracking.

    Supports multiple animations that can be selected at runtime.
    """

    # Joint indices
    LEFT_ARM = slice(15, 22)
    RIGHT_ARM = slice(22, 29)

    def __init__(self, num_joints: int = 29):
        super().__init__(num_joints)

        # Base standing pose - matched to RL policy default angles
        self.base_pose = np.zeros(num_joints)
        # Legs - use same default angles as RL policy for stability
        leg_pose = [-0.1, 0.0, 0.0, 0.3, -0.2, 0.0]
        self.base_pose[0:6] = leg_pose   # Left leg
        self.base_pose[6:12] = leg_pose  # Right leg
        # Arms slightly out
        self.base_pose[16] = 0.3    # left_shoulder_roll
        self.base_pose[18] = 0.3    # left_elbow
        self.base_pose[23] = -0.3   # right_shoulder_roll
        self.base_pose[25] = 0.3    # right_elbow

        # PD gains matched to RL policy (from g1.yaml) for stability
        self.kp = np.zeros(num_joints)
        self.kd = np.zeros(num_joints)

        # Legs - use same gains as RL policy
        # Order: hip_pitch, hip_roll, hip_yaw, knee, ankle_pitch, ankle_roll
        leg_kp = [100, 100, 100, 150, 40, 40]
        leg_kd = [2, 2, 2, 4, 2, 2]
        self.kp[0:6] = leg_kp   # Left leg
        self.kd[0:6] = leg_kd
        self.kp[6:12] = leg_kp  # Right leg
        self.kd[6:12] = leg_kd
        # Waist
        self.kp[12:15] = 100.0
        self.kd[12:15] = 5.0
        # Arms (moderate for smooth animation)
        self.kp[15:29] = 50.0
        self.kd[15:29] = 3.0

        # Current animation
        self.current_animation = "wave"
        self.animation_speed = 1.0

        # Define animations as keyframe sequences
        # Each animation is a list of (time, joint_offsets) tuples
        self.animations = {
            "wave": self._create_wave_animation(),
            "squat": self._create_squat_animation(),
            "arms_up": self._create_arms_up_animation(),
            "idle": self._create_idle_animation(),
        }

    def _create_wave_animation(self):
        """Wave right hand animation - subtle to maintain balance."""
        # Returns list of (time, pose_offset) keyframes
        # Pose offset is added to base_pose
        keyframes = []
        duration = 4.0  # Slower animation for stability

        # Keyframe times and right arm positions (more subtle movements)
        # Right arm: indices 22-28 (shoulder_pitch, shoulder_roll, shoulder_yaw, elbow, wrist_roll, wrist_pitch, wrist_yaw)
        wave_poses = [
            (0.0, {22: 0.0, 23: -0.3, 25: 0.3}),                    # Start
            (0.8, {22: -0.8, 23: -0.6, 24: 0.0, 25: 0.8}),          # Arm up (gentler)
            (1.2, {22: -0.8, 23: -0.6, 24: 0.2, 25: 0.8}),          # Wave right
            (1.6, {22: -0.8, 23: -0.6, 24: -0.2, 25: 0.8}),         # Wave left
            (2.0, {22: -0.8, 23: -0.6, 24: 0.2, 25: 0.8}),          # Wave right
            (2.4, {22: -0.8, 23: -0.6, 24: -0.2, 25: 0.8}),         # Wave left
            (3.2, {22: 0.0, 23: -0.3, 25: 0.3}),                    # Return
            (4.0, {22: 0.0, 23: -0.3, 25: 0.3}),                    # End
        ]

        for t, offsets in wave_poses:
            pose = np.zeros(self.num_joints)
            for idx, val in offsets.items():
                pose[idx] = val
            keyframes.append((t, pose))

        return {"keyframes": keyframes, "duration": duration, "loop": True}

    def _create_squat_animation(self):
        """Simple squat animation - subtle for stability."""
        keyframes = []
        duration = 6.0  # Slower for balance

        # Leg joint offsets for squat (more gentle)
        # Left leg: 0-5, Right leg: 6-11
        # hip_pitch, hip_roll, hip_yaw, knee, ankle_pitch, ankle_roll

        squat_poses = [
            (0.0, {}),  # Standing
            (2.0, {0: -0.3, 3: 0.6, 4: -0.3,    # Left leg bent (gentler)
                   6: -0.3, 9: 0.6, 10: -0.3}),  # Right leg bent
            (4.0, {0: -0.3, 3: 0.6, 4: -0.3,
                   6: -0.3, 9: 0.6, 10: -0.3}),  # Hold
            (6.0, {}),  # Return to standing
        ]

        for t, offsets in squat_poses:
            pose = np.zeros(self.num_joints)
            for idx, val in offsets.items():
                pose[idx] = val
            keyframes.append((t, pose))

        return {"keyframes": keyframes, "duration": duration, "loop": True}

    def _create_arms_up_animation(self):
        """Raise both arms animation - symmetric for balance."""
        keyframes = []
        duration = 4.0  # Slower

        arms_poses = [
            (0.0, {}),
            (1.5, {15: -1.2, 16: 0.3, 18: 0.0,      # Left arm up (gentler)
                   22: -1.2, 23: -0.3, 25: 0.0}),    # Right arm up (symmetric)
            (2.5, {15: -1.2, 16: 0.3, 18: 0.0,
                   22: -1.2, 23: -0.3, 25: 0.0}),    # Hold
            (4.0, {}),
        ]

        for t, offsets in arms_poses:
            pose = np.zeros(self.num_joints)
            for idx, val in offsets.items():
                pose[idx] = val
            keyframes.append((t, pose))

        return {"keyframes": keyframes, "duration": duration, "loop": True}

    def _create_idle_animation(self):
        """Subtle idle breathing motion."""
        keyframes = []
        duration = 4.0

        idle_poses = [
            (0.0, {}),
            (2.0, {14: 0.05}),   # Slight waist pitch forward
            (4.0, {}),
        ]

        for t, offsets in idle_poses:
            pose = np.zeros(self.num_joints)
            for idx, val in offsets.items():
                pose[idx] = val
            keyframes.append((t, pose))

        return {"keyframes": keyframes, "duration": duration, "loop": True}

    def set_animation(self, name: str):
        """Set current animation by name."""
        if name in self.animations:
            self.current_animation = name
            self.time = 0.0

    def get_animation_names(self):
        """Get list of available animation names."""
        return list(self.animations.keys())

    def _interpolate_keyframes(self, keyframes, t):
        """Interpolate between keyframes at time t."""
        # Find surrounding keyframes
        prev_kf = keyframes[0]
        next_kf = keyframes[-1]

        for i, (kf_time, kf_pose) in enumerate(keyframes):
            if kf_time <= t:
                prev_kf = (kf_time, kf_pose)
                if i + 1 < len(keyframes):
                    next_kf = keyframes[i + 1]
            else:
                next_kf = (kf_time, kf_pose)
                break

        prev_t, prev_pose = prev_kf
        next_t, next_pose = next_kf

        # Interpolation factor
        if next_t - prev_t > 0:
            alpha = (t - prev_t) / (next_t - prev_t)
            # Smooth interpolation using smoothstep
            alpha = alpha * alpha * (3 - 2 * alpha)
        else:
            alpha = 1.0

        return prev_pose + alpha * (next_pose - prev_pose)

    def compute_action(self, obs: np.ndarray, data) -> np.ndarray:
        """Compute control action from keyframe animation."""
        joint_pos = obs[13:13 + self.num_joints]
        joint_vel = obs[13 + self.num_joints:13 + 2 * self.num_joints]

        # Get current animation
        anim = self.animations[self.current_animation]
        keyframes = anim["keyframes"]
        duration = anim["duration"]

        # Loop animation
        anim_time = (self.time * self.animation_speed) % duration

        # Interpolate target pose
        pose_offset = self._interpolate_keyframes(keyframes, anim_time)
        target_pose = self.base_pose + pose_offset

        # PD control
        pos_error = target_pose - joint_pos
        vel_error = -joint_vel

        # Smooth startup
        gain_scale = min(1.0, self.time / 1.0)

        torques = gain_scale * (self.kp * pos_error + self.kd * vel_error)

        return torques

    @property
    def name(self) -> str:
        return f"Animation: {self.current_animation}"