117 lines
3.6 KiB
Python
117 lines
3.6 KiB
Python
import math
|
|
import numpy as np
|
|
from typing import TypedDict
|
|
from enum import Enum
|
|
|
|
try:
|
|
from src.robot.kinematics import BodyStateT
|
|
except ImportError:
|
|
from robot.kinematics import BodyStateT
|
|
|
|
|
|
class GaitType(Enum):
|
|
TROT_GATE = 0
|
|
CRAWL_GATE = 1
|
|
|
|
|
|
default_offset = {
|
|
GaitType.TROT_GATE: [0, 0.5, 0.5, 0],
|
|
GaitType.CRAWL_GATE: [0, 1 / 4, 2 / 4, 3 / 4],
|
|
}
|
|
|
|
default_stand_frac = {
|
|
GaitType.TROT_GATE: 3 / 4,
|
|
GaitType.CRAWL_GATE: 3 / 4,
|
|
}
|
|
|
|
|
|
class GaitStateT(TypedDict):
|
|
step_height: float
|
|
step_x: float
|
|
step_z: float
|
|
step_angle: float
|
|
step_velocity: float
|
|
step_depth: float
|
|
stand_frac: float
|
|
offset: list[float]
|
|
gait_type: GaitType
|
|
|
|
|
|
length_multipliers = np.array([-1.4, -1.0, -1.5, -1.5, -1.5, 0.0, 0.0, 0.0, 1.5, 1.5, 1.4, 1.0])
|
|
height_profile = np.array([0.0, 0.0, 0.9, 0.9, 0.9, 0.9, 0.9, 1.1, 1.1, 1.1, 0.0, 0.0])
|
|
|
|
|
|
def sine_curve(length, angle, height, phase):
|
|
x, z = length * (1 - 2 * phase) * np.cos(angle), length * (1 - 2 * phase) * np.sin(angle)
|
|
y = height * np.cos(np.pi * (x + z) / (2 * length)) if length else 0
|
|
return np.array([x, z, y])
|
|
|
|
|
|
def yaw_arc(feet, current):
|
|
return (
|
|
np.pi / 2
|
|
+ np.arctan2(feet[1], feet[0])
|
|
+ np.arctan2(np.linalg.norm(current[:2] - feet[:2]), np.linalg.norm(feet[:2]))
|
|
)
|
|
|
|
|
|
def get_control_points(length, angle, height):
|
|
x_polar, z_polar = np.cos(angle), np.sin(angle)
|
|
|
|
x = length * length_multipliers * x_polar
|
|
z = length * length_multipliers * z_polar
|
|
y = height * height_profile
|
|
return np.stack([x, z, y], axis=1)
|
|
|
|
|
|
def bezier_curve(length, angle, height, phase):
|
|
ctrl = get_control_points(length, angle, height)
|
|
n = len(ctrl) - 1
|
|
coeffs = np.array([math.comb(n, i) * (phase**i) * ((1 - phase) ** (n - i)) for i in range(n + 1)])
|
|
return np.sum(ctrl * coeffs[:, None], axis=0)
|
|
|
|
|
|
class GaitController:
|
|
def __init__(self, default_position: np.ndarray):
|
|
self.default_position = default_position
|
|
self.phase = 0.0
|
|
|
|
def step(self, gait: GaitStateT, body: BodyStateT, dt: float):
|
|
step_x, step_z, angle = gait["step_x"], gait["step_z"], gait["step_angle"]
|
|
if not any((step_x, step_z, angle)):
|
|
body["feet"] = body["feet"] + (self.default_position - body["feet"]) * dt * 10
|
|
self.phase = 0.0
|
|
return
|
|
|
|
self._advance_phase(dt, gait["step_velocity"])
|
|
|
|
stand_fraction = gait["stand_frac"]
|
|
depth = gait["step_depth"]
|
|
height = gait["step_height"]
|
|
offsets = gait["offset"]
|
|
|
|
length = np.hypot(step_x, step_z)
|
|
if step_x < 0:
|
|
length = -length
|
|
turn_amplitude = np.arctan2(step_z, length) * 2
|
|
|
|
new_feet = np.zeros_like(self.default_position)
|
|
|
|
for i, (default_foot, current_foot) in enumerate(zip(self.default_position, body["feet"])):
|
|
phase = (self.phase + offsets[i]) % 1
|
|
ph_norm, curve_fn, amp = self._phase_params(phase, stand_fraction, depth, height)
|
|
delta_pos = curve_fn(length / 2, turn_amplitude, amp, ph_norm)
|
|
delta_rot = curve_fn(np.rad2deg(angle), yaw_arc(default_foot, current_foot), amp, ph_norm)
|
|
new_feet[i][:2] = default_foot[:2] + delta_pos[:2] + delta_rot[:2]
|
|
# new_feet[i][3] = 1
|
|
|
|
body["feet"] = new_feet
|
|
|
|
def _advance_phase(self, dt: float, velocity: float):
|
|
self.phase = (self.phase + dt * velocity) % 1
|
|
|
|
def _phase_params(self, phase: float, stand_frac: float, depth: float, height: float):
|
|
if phase < stand_frac:
|
|
return phase / stand_frac, sine_curve, -depth
|
|
return (phase - stand_frac) / (1 - stand_frac), bezier_curve, height
|