simulation physics fixed

fuck them physics
This commit is contained in:
2026-08-06 15:18:06 +02:00
parent 523a4aea89
commit 346ec9e949
8 changed files with 457 additions and 900 deletions
+121 -159
View File
@@ -1,12 +1,9 @@
"""
Robot.py - Unified Robot Class for JackBot
Handles state, kinematics, backends (Hardware/Simulation), and motion execution.
Robot.py - Central Robot Control, Kinematics, State Machine & Hardware Abstraction
"""
from typing import Protocol, Optional, Union
from typing import Protocol, Optional, Union, Tuple, List
import numpy as np
import math
import pybullet as p
from states import STATE_REGISTRY
from states.State import State
@@ -15,33 +12,51 @@ 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:
...
"""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 robot."""
"""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 # Physical hardware steps in real-time
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'):
@@ -49,52 +64,62 @@ class HardwareBackend:
class PyBulletBackend:
"""Backend for PyBullet simulation execution."""
def __init__(self, sim_instance, body_id: Optional[int] = None):
"""Backend mapping Robot operations directly to PyBullet simulation engine."""
def __init__(self, sim_instance: Simulation):
self.sim = sim_instance
self.body_id = body_id
def send_angles(self, rad_array: dt.RadArray) -> None:
if not self.sim:
return
if self.body_id is not None and hasattr(self.sim, 'updatePosForBody'):
self.sim.updatePosForBody(self.body_id, rad_array)
elif hasattr(self.sim, 'updatePos'):
self.sim.updatePos(rad_array)
elif hasattr(self.sim, 'set_robot_joint_angles') and self.body_id is not None:
joint_indices = getattr(self.sim, 'joint_indices', list(range(18)))
self.sim.set_robot_joint_angles(self.body_id, joint_indices, rad_array.data.flatten())
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 and hasattr(self.sim, 'disconnect'):
if self.sim:
self.sim.disconnect()
class Robot:
"""
Encapsulates a single JackBot hexapod instance.
Maintains joint states, leg positions, kinematics, and backend control.
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
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 FACTORY CREATION ---
# --- BACKEND INSTANTIATION ---
if isinstance(backend_type, BackendType):
if backend_type == BackendType.SIMULATION:
# Launch PyBullet 3D Simulation GUI
sim_instance = Simulation(urdf_path=self.urdf_path)
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)
@@ -105,7 +130,7 @@ class Robot:
else:
self.backend = backend_type
# Kinematics and position initialization
# 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)
@@ -113,32 +138,20 @@ class Robot:
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)
# RL configuration & Gait state variables
self.action_scale = 0.1 # Radian step scale for RL deltas
self.gait_phase = 0.0
# 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]
self.robot_state = "idle"
self.leg_state = np.array(["step", "drag", "step", "drag", "step", "drag"])
# State machine initialization
# 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:
"""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)
@@ -148,12 +161,18 @@ class Robot:
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.set_joint_angles(self.current_rad)
self.step_sim()
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)
@@ -169,41 +188,51 @@ class Robot:
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)
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 IK solver for kinematic execution & evaluation."""
"""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
# Advance gait step cycle phase
self.gait_phase = (self.gait_phase + 0.22) % (2.0 * math.pi)
stride_len = 0.045 # 4.5 cm maximum stride
step_height = 0.035 # 3.5 cm foot clearance height
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)
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)
# Tripod leg grouping phase offset (even vs odd leg IDs)
phase_offset = 0.0 if (leg_id % 2 == 0) else math.pi
leg_phase = (self.gait_phase + phase_offset) % (2.0 * math.pi)
# Directional motion unit vector calculation
lx, ly = base_pos[0], base_pos[1]
rot_dx = -omega * ly
rot_dy = omega * lx
@@ -216,116 +245,49 @@ class Robot:
dy_unit = dy_dir / dir_norm
if leg_phase < math.pi:
# Swing Phase (Leg lifted & stepping forward)
# 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 grounded & propelling torso)
progress = math.cos(leg_phase - math.pi)
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_leg_pos = base_pos + np.array([dx, dy, dz], dtype=np.float32)
target_positions.append(target_leg_pos)
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 step_with_command(
self,
command: np.ndarray,
action: Optional[np.ndarray] = None,
mode: str = "direct"
) -> None:
"""
Unified motion execution method supporting Direct RL, Residual RL, and Pure Kinematics.
Args:
command: np.ndarray [vx, vy, vz, omega] from RL environment
action: np.ndarray [18,] RL action deltas from neural network ([-1, 1])
mode: "kinematics_only" | "residual" | "direct"
"""
# 1. Map RL command [vx, vy, vz, omega] to Robot motion vector [vx, vy, omega]
cmd_vx, cmd_vy, _, cmd_omega = command
self.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
# 2. Automatically trigger state transitions based on command magnitude
cmd_magnitude = math.hypot(cmd_vx, cmd_vy) + abs(cmd_omega)
if cmd_magnitude > 0.05 and self.current_state_key == "idle":
# Transition to walking state if registered in state machine
target_state = "walk" if "walk" in STATE_REGISTRY else "move"
if target_state in STATE_REGISTRY:
self.transition_to(target_state)
elif cmd_magnitude <= 0.05 and self.current_state_key != "idle":
self.transition_to("idle")
# 3. Execute according to chosen mode
if mode == "kinematics_only":
self.step_kinematic_gait(cmd_vx, cmd_vy, cmd_omega)
elif mode == "residual":
if self.current_state_key != "idle" and self.current_state and self.current_state_key in STATE_REGISTRY:
self.current_state.execute(self)
else:
self.step_kinematic_gait(cmd_vx, cmd_vy, cmd_omega)
if action is not None:
action_flat = np.clip(np.asarray(action, dtype=np.float32), -1.0, 1.0) * self.action_scale
kin_flat = self.current_rad.data.flatten()
final_flat = np.clip(kin_flat + action_flat, -np.pi / 2, np.pi / 2)
self.set_joint_angles(dt.RadArray(data=final_flat.reshape(self.current_rad.data.shape)))
elif mode == "direct":
if action is not None:
self.apply_rl_action(action)
# --- RL METHODS ---
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
# Sync with actual PyBullet state if backend supports it
if isinstance(self.backend, PyBulletBackend) and self.backend.sim:
actual_angles = self.backend.sim.get_robot_joint_angles(self.backend.body_id)
current_flat = actual_angles.flatten()
else:
current_flat = self.current_rad.data.flatten()
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:
"""
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)
"""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.astype(np.float32)
return joint_angles
def cleanup(self) -> None:
if self.backend and hasattr(self.backend, 'cleanup'):
if self.backend:
self.backend.cleanup()