JackM323
acb3d671be
new reward/penalty system learning phases with curriculum learning new training parameters cleanup of old code better logging while training multiple environments instead of robots (they could bumb into each other)
219 lines
8.0 KiB
Python
219 lines
8.0 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
|
|
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 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() |