File size: 8,195 Bytes
ef118cf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 | """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}"
|