293 lines
12 KiB
Python
293 lines
12 KiB
Python
"""
|
|
Robot.py - Central Robot Control, Kinematics, State Machine & Hardware Abstraction
|
|
"""
|
|
from typing import Protocol, Optional, Union, Tuple, List
|
|
import numpy as np
|
|
import math
|
|
|
|
from states import STATE_REGISTRY
|
|
from states.State import State
|
|
import DataTypes as dt
|
|
import kinematics as kin
|
|
import robot_init as ri
|
|
from config import cfg, BackendType
|
|
|
|
from simulation import Simulation
|
|
from EspCommunication import ESP32Communication
|
|
from ArduinoCommunication import ArduinoCommunication
|
|
|
|
|
|
class RobotBackend(Protocol):
|
|
"""Protocol defining hardware abstraction for both Simulation and Hardware backends."""
|
|
def send_angles(self, rad_array: dt.RadArray) -> None: ...
|
|
def step_simulation(self) -> None: ...
|
|
def hard_reset_joints(self, target_angles: np.ndarray) -> None: ...
|
|
def reset_base(self, position: Optional[List[float]] = None, orientation: Optional[List[float]] = None) -> None: ...
|
|
def get_joint_angles(self) -> np.ndarray: ...
|
|
def get_base_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]: ...
|
|
def get_base_velocity(self) -> Tuple[List[float], List[float]]: ...
|
|
def cleanup(self) -> None: ...
|
|
|
|
|
|
class HardwareBackend:
|
|
"""Backend for physical ESP32 or Arduino microcontrollers."""
|
|
def __init__(self, comm_channel):
|
|
self.comm_channel = comm_channel
|
|
self.internal_angles = np.zeros(18, dtype=np.float32)
|
|
|
|
def send_angles(self, rad_array: dt.RadArray) -> None:
|
|
self.internal_angles = rad_array.data.flatten().copy()
|
|
if self.comm_channel:
|
|
self.comm_channel.send_motion(rad_array)
|
|
|
|
def step_simulation(self) -> None:
|
|
pass
|
|
|
|
def hard_reset_joints(self, target_angles: np.ndarray) -> None:
|
|
self.internal_angles = target_angles.flatten().copy()
|
|
|
|
def reset_base(self, position: Optional[List[float]] = None, orientation: Optional[List[float]] = None) -> None:
|
|
pass
|
|
|
|
def get_joint_angles(self) -> np.ndarray:
|
|
return self.internal_angles.copy()
|
|
|
|
def get_base_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]:
|
|
return [0.0, 0.0, 0.122], (0.0, 0.0, 0.0)
|
|
|
|
def get_base_velocity(self) -> Tuple[List[float], List[float]]:
|
|
return [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]
|
|
|
|
def cleanup(self) -> None:
|
|
if self.comm_channel and hasattr(self.comm_channel, 'close'):
|
|
self.comm_channel.close()
|
|
|
|
|
|
class PyBulletBackend:
|
|
"""Backend mapping Robot operations directly to PyBullet simulation engine."""
|
|
def __init__(self, sim_instance: Simulation):
|
|
self.sim = sim_instance
|
|
|
|
def send_angles(self, rad_array: dt.RadArray) -> None:
|
|
if self.sim:
|
|
self.sim.set_robot_joint_angles(rad_array)
|
|
|
|
def step_simulation(self) -> None:
|
|
if self.sim:
|
|
self.sim.step()
|
|
|
|
def hard_reset_joints(self, target_angles: np.ndarray) -> None:
|
|
if self.sim:
|
|
self.sim.hard_reset_joint_angles(target_angles)
|
|
|
|
def reset_base(self, position: Optional[List[float]] = None, orientation: Optional[List[float]] = None) -> None:
|
|
if self.sim:
|
|
self.sim.reset_robot_base(pos=position, orn=orientation)
|
|
|
|
def get_joint_angles(self) -> np.ndarray:
|
|
return self.sim.get_robot_joint_angles() if self.sim else np.zeros(18, dtype=np.float32)
|
|
|
|
def get_base_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]:
|
|
return self.sim.get_robot_pose_and_rpy() if self.sim else ([0, 0, 0], (0, 0, 0))
|
|
|
|
def get_base_velocity(self) -> Tuple[List[float], List[float]]:
|
|
return self.sim.get_robot_velocity() if self.sim else ([0, 0, 0], [0, 0, 0])
|
|
|
|
def cleanup(self) -> None:
|
|
if self.sim:
|
|
self.sim.disconnect()
|
|
|
|
|
|
class Robot:
|
|
"""
|
|
Unified JackBot Class.
|
|
Coordinates joint memory, IK solvers, procedural tripods, and backend communication.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
backend_type: Union[BackendType, RobotBackend] = BackendType.SIMULATION,
|
|
start_pose: str = "init_deg",
|
|
urdf_path: str = cfg.urdf_path,
|
|
mode: str = "kinematics" # "kinematics", "residual", or "direct"
|
|
):
|
|
self.urdf_path = urdf_path
|
|
self.start_pose = start_pose
|
|
self.mode = mode
|
|
|
|
# --- BACKEND INSTANTIATION ---
|
|
if isinstance(backend_type, BackendType):
|
|
if backend_type == BackendType.SIMULATION:
|
|
sim_instance = Simulation(urdf_path=self.urdf_path, use_gui=True)
|
|
sim_instance.load_scene()
|
|
self.backend: RobotBackend = PyBulletBackend(sim_instance)
|
|
elif backend_type == BackendType.ESP32:
|
|
comm = ESP32Communication(ip=cfg.esp32_ip, port=cfg.esp32_port)
|
|
self.backend = HardwareBackend(comm)
|
|
elif backend_type == BackendType.ARDUINO:
|
|
comm = ArduinoCommunication(port=cfg.port, baudrate=cfg.baudrate)
|
|
self.backend = HardwareBackend(comm)
|
|
else:
|
|
self.backend = backend_type
|
|
|
|
# Kinematics initialization
|
|
pose_deg = ri.init_deg if start_pose == "init_deg" else ri.init90_deg
|
|
self.current_rad: dt.RadArray = pose_deg.to_rad()
|
|
self.current_pos: dt.PosArray = kin.ikpyForward(self.current_rad)
|
|
self.center_points: dt.PosArray = (
|
|
ri.get_center_points() if hasattr(ri, "get_center_points") else ri.center_points
|
|
)
|
|
|
|
# RL configuration & Gait state variables
|
|
self.action_scale = 0.1 # Radian step scale for RL deltas
|
|
self.gait_phase = 0.0
|
|
self.vector_dirmov = [0.0, 0.0, 0.0] # [vx, vy, omega]
|
|
self.robot_state = "idle"
|
|
self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
|
|
|
|
# State Machine Initialization
|
|
self.current_state_key: str = "idle"
|
|
self.current_state: State = STATE_REGISTRY["idle"]
|
|
self.current_state.enter(self)
|
|
|
|
def set_joint_angles(self, target_rad: dt.RadArray) -> None:
|
|
"""Updates internal Python memory state and sends angles to active backend."""
|
|
self.current_rad = target_rad
|
|
if self.backend:
|
|
self.backend.send_angles(target_rad)
|
|
|
|
def step_sim(self) -> None:
|
|
if self.backend:
|
|
self.backend.step_simulation()
|
|
|
|
def reset_to_init(self) -> None:
|
|
"""Resets kinematics state and forces instant joint alignment in backend."""
|
|
pose_deg = ri.init_deg if self.start_pose == "init_deg" else ri.init90_deg
|
|
self.current_rad = pose_deg.to_rad()
|
|
self.current_pos = kin.ikpyForward(self.current_rad)
|
|
self.gait_phase = 0.0
|
|
self.robot_state = "idle"
|
|
self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
|
|
|
|
init_flat = self.current_rad.data.flatten()
|
|
if self.backend:
|
|
self.backend.hard_reset_joints(init_flat)
|
|
self.backend.send_angles(self.current_rad)
|
|
|
|
def compute_ik(self, target_pos: dt.PosArray) -> dt.RadArray:
|
|
return kin.ikpyInverse(target_pos, initial_rad=self.current_rad)
|
|
|
|
def compute_fk(self, target_rad: Optional[dt.RadArray] = None) -> dt.PosArray:
|
|
rads = target_rad if target_rad is not None else self.current_rad
|
|
return kin.ikpyForward(rads)
|
|
|
|
def transition_to(self, next_state_key: str) -> None:
|
|
if next_state_key in STATE_REGISTRY and next_state_key != self.current_state_key:
|
|
self.current_state.exit(self)
|
|
self.current_state_key = next_state_key
|
|
self.current_state = STATE_REGISTRY[next_state_key]
|
|
self.current_state.enter(self)
|
|
|
|
def tick(self, action: Optional[np.ndarray] = None) -> None:
|
|
"""
|
|
Unified control loop tick.
|
|
Processes commands through direct RL, residual RL, or State Machine kinematics.
|
|
"""
|
|
vx, vy, omega = self.vector_dirmov
|
|
|
|
if self.mode == "direct":
|
|
if action is not None:
|
|
self.apply_rl_action(action)
|
|
|
|
elif self.mode == "residual":
|
|
self.step_kinematic_gait(vx, vy, omega)
|
|
if action is not None:
|
|
self.apply_rl_action_delta(action)
|
|
|
|
else: # "kinematics" / standard State Machine execution
|
|
next_state_key = self.current_state.execute(self)
|
|
if next_state_key:
|
|
self.transition_to(next_state_key)
|
|
|
|
self.step_sim()
|
|
|
|
def step_kinematic_gait(self, vx: float, vy: float, omega: float) -> None:
|
|
"""Procedural Tripod Gait solver."""
|
|
cmd_mag = math.hypot(vx, vy) + abs(omega)
|
|
if cmd_mag < 0.03:
|
|
target_rad = self.compute_ik(self.center_points)
|
|
self.set_joint_angles(target_rad)
|
|
return
|
|
|
|
self.gait_phase = (self.gait_phase + 0.22) % (2.0 * math.pi)
|
|
stride_len = 0.045
|
|
step_height = 0.035
|
|
|
|
center_data = (
|
|
self.center_points.data if hasattr(self.center_points, 'data') else np.array(self.center_points)
|
|
)
|
|
|
|
target_positions = []
|
|
for leg_id in range(6):
|
|
base_pos = np.array(center_data[leg_id], dtype=np.float32)
|
|
phase_offset = 0.0 if (leg_id % 2 == 0) else math.pi
|
|
leg_phase = (self.gait_phase + phase_offset) % (2.0 * math.pi)
|
|
|
|
lx, ly = base_pos[0], base_pos[1]
|
|
rot_dx = -omega * ly
|
|
rot_dy = omega * lx
|
|
|
|
dx_dir = vx + rot_dx
|
|
dy_dir = vy + rot_dy
|
|
dir_norm = math.hypot(dx_dir, dy_dir) + 1e-6
|
|
|
|
dx_unit = dx_dir / dir_norm
|
|
dy_unit = dy_dir / dir_norm
|
|
|
|
if leg_phase < math.pi:
|
|
# Swing phase (leg lifted in air, moving forward)
|
|
progress = math.cos(leg_phase)
|
|
lift = math.sin(leg_phase) * step_height
|
|
dx = progress * stride_len * dx_unit
|
|
dy = progress * stride_len * dy_unit
|
|
dz = lift
|
|
else:
|
|
# Stance phase (leg on ground, pushing body forward)
|
|
progress = math.cos(leg_phase - math.pi)
|
|
dx = -progress * stride_len * dx_unit
|
|
dy = -progress * stride_len * dy_unit
|
|
dz = 0.0
|
|
|
|
target_positions.append(base_pos + np.array([dx, dy, dz], dtype=np.float32))
|
|
|
|
target_pos_array = dt.PosArray(np.array(target_positions))
|
|
target_rad = self.compute_ik(target_pos_array)
|
|
self.set_joint_angles(target_rad)
|
|
|
|
def apply_rl_action(self, action: np.ndarray) -> None:
|
|
"""Applies absolute action targets directly for Direct RL."""
|
|
action = np.asarray(action, dtype=np.float32)
|
|
scaled_action = np.clip(action, -1.0, 1.0) * (np.pi / 2.0)
|
|
new_rad = dt.RadArray(data=scaled_action.reshape(self.current_rad.data.shape))
|
|
self.set_joint_angles(new_rad)
|
|
|
|
def apply_rl_action_delta(self, action: np.ndarray) -> None:
|
|
"""Applies action deltas on top of joint state for Residual RL."""
|
|
action = np.asarray(action, dtype=np.float32)
|
|
scaled_action = np.clip(action, -1.0, 1.0) * self.action_scale
|
|
current_flat = self.backend.get_joint_angles().flatten()
|
|
updated_flat = np.clip(current_flat + scaled_action, -np.pi / 2, np.pi / 2)
|
|
new_rad = dt.RadArray(data=updated_flat.reshape(self.current_rad.data.shape))
|
|
self.set_joint_angles(new_rad)
|
|
|
|
def get_observation(self, command: Optional[np.ndarray] = None) -> np.ndarray:
|
|
"""Extracts joint positions directly from active backend."""
|
|
joint_angles = self.backend.get_joint_angles().flatten().astype(np.float32)
|
|
if command is not None:
|
|
cmd = np.asarray(command, dtype=np.float32).flatten()
|
|
return np.concatenate([joint_angles, cmd]).astype(np.float32)
|
|
return joint_angles
|
|
|
|
def cleanup(self) -> None:
|
|
if self.backend:
|
|
self.backend.cleanup() |