""" Robot.py - Unified Robot Class for JackBot Handles state, kinematics, backends (Hardware/Simulation), and motion execution. """ from typing import Protocol, Optional, Union import numpy as np import math import pybullet as p 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 # Import communications and simulation modules from simulation import Simulation from EspCommunication import ESP32Communication from ArduinoCommunication import ArduinoCommunication class RobotBackend(Protocol): """Abstraction layer for hardware vs simulation output.""" def send_angles(self, rad_array: dt.RadArray) -> None: ... def step_simulation(self) -> None: ... def cleanup(self) -> None: ... class HardwareBackend: """Backend for physical ESP32 or Arduino robot.""" def __init__(self, comm_channel): self.comm_channel = comm_channel def send_angles(self, rad_array: dt.RadArray) -> None: if self.comm_channel: self.comm_channel.send_motion(rad_array) def step_simulation(self) -> None: pass # Physical hardware steps in real-time def cleanup(self) -> None: if self.comm_channel and hasattr(self.comm_channel, 'close'): self.comm_channel.close() class PyBulletBackend: """Backend for PyBullet simulation execution.""" def __init__(self, sim_instance, body_id: Optional[int] = None): self.sim = sim_instance self.body_id = body_id def send_angles(self, rad_array: dt.RadArray) -> None: if self.sim: # If body_id is set, target that specific robot body if self.body_id is not None and hasattr(self.sim, 'updatePosForBody'): self.sim.updatePosForBody(self.body_id, rad_array) else: self.sim.updatePos(rad_array) def step_simulation(self) -> None: if self.sim: self.sim.step() def cleanup(self) -> None: if self.sim and hasattr(self.sim, 'disconnect'): self.sim.disconnect() class Robot: """ Encapsulates a single JackBot hexapod instance. Maintains joint states, leg positions, kinematics, and backend control. """ def __init__( self, backend_type: Union[BackendType, RobotBackend] = BackendType.SIMULATION, start_pose: str = "init_deg", urdf_path: str = cfg.urdf_path ): self.urdf_path = urdf_path self.start_pose = start_pose # --- BACKEND FACTORY CREATION --- if isinstance(backend_type, BackendType): if backend_type == BackendType.SIMULATION: # Launch PyBullet 3D Simulation GUI sim_instance = Simulation(urdf_path=self.urdf_path) 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 and position 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 self.action_scale = 0.1 # Joint delta step size (radians) # Gait / motion variables self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"]) self.robot_state = "idle" self.vector_dirmov = [0.0, 0.0, 0.0] # [vx, vy, omega] # State machine initialization self.current_state_key: str = "idle" self.current_state: State = STATE_REGISTRY["idle"] self.current_state.enter(self) def change_state(self, new_state: State) -> None: if self.current_state: self.current_state.exit(self) self.current_state = new_state self.current_state.enter(self) def update(self) -> None: if self.current_state: self.current_state.execute(self) self.step_sim() def set_joint_angles(self, target_rad: dt.RadArray) -> None: 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: 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.set_joint_angles(self.current_rad) self.step_sim() 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) -> None: next_state_key = self.current_state.execute(self) if next_state_key: self.transition_to(next_state_key) self.step_sim() # --- RL METHODS --- def apply_rl_action(self, action: np.ndarray) -> None: """Applies continuous RL action deltas [-1, 1] to current joint angles.""" action = np.asarray(action, dtype=np.float32) scaled_action = np.clip(action, -1.0, 1.0) * self.action_scale current_flat = self.current_rad.data.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: """ Returns observation vector [18 joint angles] + [optional 4 command dimensions]. Queries SimManager helper if PyBulletBackend is used; falls back to internal state otherwise. """ if isinstance(self.backend, PyBulletBackend) and self.backend.sim: body_id = self.backend.body_id if self.backend.body_id is not None else 0 if hasattr(self.backend.sim, 'get_robot_joint_angles'): joint_angles = self.backend.sim.get_robot_joint_angles(body_id) elif hasattr(self.backend.sim, 'physics_client'): physics_client = self.backend.sim.physics_client joint_indices = self.backend.sim.robot_joints.get(body_id, list(range(18))) if hasattr(self.backend.sim, 'robot_joints') else list(range(18)) joint_states = p.getJointStates(body_id, joint_indices, physicsClientId=physics_client) joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32) else: joint_angles = self.current_rad.data.flatten().astype(np.float32) else: joint_angles = self.current_rad.data.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.astype(np.float32) def cleanup(self) -> None: if self.backend and hasattr(self.backend, 'cleanup'): self.backend.cleanup()