""" 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 = cfg.step_length step_height = cfg.step_height 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: action = np.asarray(action, dtype=np.float32) new_rad = dt.RadArray(data=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()