"""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}"