JackM323
61cb0150f0
new Metrics and SimManager for live training viewing Robot.py usage added to machinelearning
217 lines
7.9 KiB
Python
217 lines
7.9 KiB
Python
"""
|
|
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
|
|
|
|
# --- 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:
|
|
self.current_rad = ri.init_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 PyBullet if backend is PyBulletBackend; otherwise falls back to internal state.
|
|
"""
|
|
if isinstance(self.backend, PyBulletBackend) and self.backend.sim and hasattr(self.backend.sim, 'physics_client'):
|
|
physics_client = self.backend.sim.physics_client
|
|
body_id = self.backend.body_id if self.backend.body_id is not None else 0
|
|
|
|
# Retrieve joint mapping from SimManager/Simulation if available
|
|
if hasattr(self.backend.sim, 'robot_joints') and body_id in self.backend.sim.robot_joints:
|
|
joint_indices = self.backend.sim.robot_joints[body_id]
|
|
else:
|
|
joint_indices = 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)
|
|
|
|
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() |