simulation physics fixed
fuck them physics
This commit is contained in:
@@ -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:
|
||||
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()
|
||||
@@ -1,298 +0,0 @@
|
||||
"""
|
||||
ml/SimManager.py - PyBullet Simulation Manager (Single Robot Dedicated)
|
||||
"""
|
||||
from typing import List, Tuple, Optional, Union
|
||||
import pybullet as p
|
||||
import pybullet_data
|
||||
import numpy as np
|
||||
import DataTypes as dt
|
||||
|
||||
|
||||
class SimManager:
|
||||
"""Manages PyBullet simulation lifecycle for a single JackBot hexapod."""
|
||||
|
||||
def __init__(self, use_gui: bool = True):
|
||||
self.use_gui = use_gui
|
||||
self.physics_client = None
|
||||
self.plane: Optional[int] = None
|
||||
self.joint_indices: List[int] = []
|
||||
self.foot_indices: List[int] = []
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Connects to PyBullet and hides side GUI panels."""
|
||||
if self.physics_client is not None and p.isConnected(self.physics_client):
|
||||
return
|
||||
|
||||
flags = p.GUI if self.use_gui else p.DIRECT
|
||||
self.physics_client = p.connect(flags)
|
||||
|
||||
p.setAdditionalSearchPath(pybullet_data.getDataPath())
|
||||
p.setGravity(0, 0, -9.81, physicsClientId=self.physics_client)
|
||||
|
||||
if self.use_gui:
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.physics_client)
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0, physicsClientId=self.physics_client)
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
|
||||
|
||||
def load_scene(
|
||||
self,
|
||||
urdf_path: str,
|
||||
spacing: float = 0.0,
|
||||
position_func=None
|
||||
) -> Tuple[int, List[int], List[List[int]]]:
|
||||
"""Loads ground plane and the single robot URDF with high ground friction and force settings."""
|
||||
# 1. Load ground plane and set explicit friction
|
||||
self.plane = p.loadURDF("plane.urdf", physicsClientId=self.physics_client)
|
||||
|
||||
# 2. Determine starting position (matching simulation.py starting height of 0.20m)
|
||||
spawn_pos = position_func(0) if position_func is not None else [0.0, 0.0, 0.20]
|
||||
|
||||
# 3. Spawn single robot body
|
||||
self.robot_id = p.loadURDF(urdf_path, spawn_pos, physicsClientId=self.physics_client)
|
||||
|
||||
# 4. Retrieve revolute joint indices
|
||||
self.joint_indices = []
|
||||
for j in range(p.getNumJoints(self.robot_id, physicsClientId=self.physics_client)):
|
||||
info = p.getJointInfo(self.robot_id, j, physicsClientId=self.physics_client)
|
||||
if info[2] == p.JOINT_REVOLUTE:
|
||||
self.joint_indices.append(j)
|
||||
|
||||
# Return format maintains 100% compatibility with JackBotEnv unpack sequence
|
||||
return self.plane, [self.robot_id], [self.joint_indices]
|
||||
|
||||
def _resolve_body_id(self, body_id: Optional[int] = None) -> int:
|
||||
"""Internal helper to return the single active robot ID."""
|
||||
if body_id is not None:
|
||||
return body_id
|
||||
if self.robot_id is not None:
|
||||
return self.robot_id
|
||||
raise RuntimeError("No robot loaded in SimManager. Call load_scene() first.")
|
||||
|
||||
def set_robot_joint_angles(
|
||||
self,
|
||||
target_angles: Union[np.ndarray, List[float]],
|
||||
joint_indices: Optional[List[int]] = None,
|
||||
body_id: Optional[int] = None,
|
||||
) -> None:
|
||||
"""Applies position control with force=500 matching simulation.py logic."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
j_indices = joint_indices if joint_indices is not None else self.joint_indices
|
||||
radflat = target_angles.flatten() if isinstance(target_angles, np.ndarray) else target_angles
|
||||
|
||||
for joint_index, target_angle in zip(j_indices, radflat):
|
||||
p.setJointMotorControl2(
|
||||
bodyIndex=bid,
|
||||
jointIndex=joint_index,
|
||||
controlMode=p.POSITION_CONTROL,
|
||||
targetPosition=float(target_angle),
|
||||
force=500,
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
|
||||
def updatePosForBody(
|
||||
self,
|
||||
body_id_or_rad: Union[int, dt.RadArray],
|
||||
current_rad: Optional[dt.RadArray] = None
|
||||
) -> None:
|
||||
"""Updates robot joint positions using RadArray input."""
|
||||
if isinstance(body_id_or_rad, dt.RadArray):
|
||||
rad_data = body_id_or_rad
|
||||
bid = self.robot_id
|
||||
else:
|
||||
bid = self._resolve_body_id(body_id_or_rad)
|
||||
rad_data = current_rad
|
||||
|
||||
if rad_data is not None:
|
||||
self.set_robot_joint_angles(rad_data.data, body_id=bid)
|
||||
|
||||
def step(self) -> None:
|
||||
p.stepSimulation(physicsClientId=self.physics_client)
|
||||
|
||||
def set_rendering(self, enabled: bool) -> None:
|
||||
"""Toggles PyBullet 3D rendering visualizer."""
|
||||
if self.physics_client is not None and p.isConnected(self.physics_client):
|
||||
p.configureDebugVisualizer(
|
||||
p.COV_ENABLE_RENDERING,
|
||||
1 if enabled else 0,
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
|
||||
def disconnect(self) -> None:
|
||||
if self.physics_client is not None and p.isConnected(self.physics_client):
|
||||
p.disconnect(self.physics_client)
|
||||
self.physics_client = None
|
||||
|
||||
def get_contact_points(
|
||||
self,
|
||||
bodyA: int = -1,
|
||||
bodyB: int = -1,
|
||||
linkIndexA: int = -1,
|
||||
linkIndexB: int = -1,
|
||||
):
|
||||
"""Wrapper around pybullet.getContactPoints bound to this simulation client."""
|
||||
kwargs = {"physicsClientId": self.physics_client}
|
||||
if bodyA != -1:
|
||||
kwargs["bodyA"] = bodyA
|
||||
if bodyB != -1:
|
||||
kwargs["bodyB"] = bodyB
|
||||
if linkIndexA != -1:
|
||||
kwargs["linkIndexA"] = linkIndexA
|
||||
if linkIndexB != -1:
|
||||
kwargs["linkIndexB"] = linkIndexB
|
||||
|
||||
return p.getContactPoints(**kwargs)
|
||||
|
||||
# Add to ml/SimManager.py
|
||||
def hard_reset_joint_angles(self, target_angles: np.ndarray, body_id: Optional[int] = None) -> None:
|
||||
bid = self._resolve_body_id(body_id)
|
||||
radflat = target_angles.flatten()
|
||||
for joint_index, target_angle in zip(self.joint_indices, radflat):
|
||||
p.resetJointState(
|
||||
bodyUniqueId=bid,
|
||||
jointIndex=joint_index,
|
||||
targetValue=float(target_angle),
|
||||
targetVelocity=0.0,
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
|
||||
# --- SINGLE ROBOT GETTERS & SETTERS ---
|
||||
|
||||
def reset_robot_base(
|
||||
self,
|
||||
body_id_or_pos: Union[int, List[float]],
|
||||
position_or_orn: Optional[List[float]] = None,
|
||||
orientation: Optional[List[float]] = None,
|
||||
linear_velocity: Optional[List[float]] = None,
|
||||
angular_velocity: Optional[List[float]] = None
|
||||
) -> None:
|
||||
"""Resets the robot base pose and clears linear/angular velocities."""
|
||||
if isinstance(body_id_or_pos, int):
|
||||
bid = body_id_or_pos
|
||||
pos = position_or_orn if position_or_orn is not None else [0.0, 0.0, 0.20]
|
||||
orn = orientation if orientation is not None else [0.0, 0.0, 0.0, 1.0]
|
||||
else:
|
||||
bid = self.robot_id
|
||||
pos = body_id_or_pos
|
||||
orn = position_or_orn if position_or_orn is not None else [0.0, 0.0, 0.0, 1.0]
|
||||
|
||||
lin_v = linear_velocity if linear_velocity is not None else [0.0, 0.0, 0.0]
|
||||
ang_v = angular_velocity if angular_velocity is not None else [0.0, 0.0, 0.0]
|
||||
|
||||
p.resetBasePositionAndOrientation(
|
||||
bid, pos, orn, physicsClientId=self.physics_client
|
||||
)
|
||||
p.resetBaseVelocity(
|
||||
bid, linearVelocity=lin_v, angularVelocity=ang_v,
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
|
||||
def get_robot_pose(self, body_id: Optional[int] = None) -> Tuple[List[float], List[float]]:
|
||||
"""Returns base position (x, y, z) and orientation quaternion (x, y, z, w)."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
pos, orn = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client)
|
||||
return list(pos), list(orn)
|
||||
|
||||
def get_robot_rpy(self, body_id: Optional[int] = None) -> Tuple[float, float, float]:
|
||||
"""Returns roll, pitch, yaw angles in radians for the robot."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
_, orn = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client)
|
||||
roll, pitch, yaw = p.getEulerFromQuaternion(orn)
|
||||
return float(roll), float(pitch), float(yaw)
|
||||
|
||||
def get_robot_pose_and_rpy(self, body_id: Optional[int] = None) -> Tuple[List[float], Tuple[float, float, float]]:
|
||||
"""Returns base position and (roll, pitch, yaw) tuple in radians."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
pos, orn = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client)
|
||||
roll, pitch, yaw = p.getEulerFromQuaternion(orn)
|
||||
return list(pos), (float(roll), float(pitch), float(yaw))
|
||||
|
||||
def get_robot_velocity(self, body_id: Optional[int] = None) -> Tuple[List[float], List[float]]:
|
||||
"""Returns linear velocity (vx, vy, vz) and angular velocity (wx, wy, wz)."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
lin_v, ang_v = p.getBaseVelocity(bid, physicsClientId=self.physics_client)
|
||||
return list(lin_v), list(ang_v)
|
||||
|
||||
def get_robot_joint_angles(
|
||||
self,
|
||||
body_id: Optional[int] = None,
|
||||
joint_indices: Optional[List[int]] = None
|
||||
) -> np.ndarray:
|
||||
"""Returns joint angles as a 1D numpy float32 array."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
j_indices = joint_indices if joint_indices is not None else self.joint_indices
|
||||
joint_states = p.getJointStates(bid, j_indices, physicsClientId=self.physics_client)
|
||||
return np.array([state[0] for state in joint_states], dtype=np.float32)
|
||||
|
||||
def get_foot_link_indices(self, body_id: Optional[int] = None) -> List[int]:
|
||||
"""Inspects URDF joint structure to extract link IDs for leg tips and tibias."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
foot_indices = []
|
||||
num_joints = p.getNumJoints(bid, physicsClientId=self.physics_client)
|
||||
for j_idx in range(num_joints):
|
||||
info = p.getJointInfo(bid, j_idx, physicsClientId=self.physics_client)
|
||||
link_name = info[12].decode("utf-8")
|
||||
if "tip" in link_name or "tibia" in link_name:
|
||||
foot_indices.append(j_idx)
|
||||
return foot_indices
|
||||
|
||||
def set_robot_color(
|
||||
self,
|
||||
body_id_or_rgba: Union[int, List[float]],
|
||||
rgba: Optional[List[float]] = None
|
||||
) -> None:
|
||||
"""Changes visual color RGBA of base link and all joints."""
|
||||
if isinstance(body_id_or_rgba, int):
|
||||
bid = body_id_or_rgba
|
||||
color = rgba
|
||||
else:
|
||||
bid = self.robot_id
|
||||
color = body_id_or_rgba
|
||||
|
||||
if color is None:
|
||||
color = [1.0, 1.0, 1.0, 1.0]
|
||||
|
||||
num_joints = p.getNumJoints(bid, physicsClientId=self.physics_client)
|
||||
p.changeVisualShape(bid, -1, rgbaColor=color, physicsClientId=self.physics_client)
|
||||
for j in range(num_joints):
|
||||
p.changeVisualShape(bid, j, rgbaColor=color, physicsClientId=self.physics_client)
|
||||
|
||||
def measure_robot_height(self, body_id: Optional[int] = None) -> float:
|
||||
"""Gets current Z height of the robot base."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
pos, _ = p.getBasePositionAndOrientation(bid, physicsClientId=self.physics_client)
|
||||
return pos[2]
|
||||
|
||||
def settle_and_measure_height(
|
||||
self,
|
||||
robot_ids_or_steps: Union[List[int], int] = 200,
|
||||
steps: int = 200,
|
||||
fallback_height: float = 0.122
|
||||
) -> float:
|
||||
"""Steps simulation so robot settles onto floor, then calculates target standing height."""
|
||||
# Intelligently resolve whether the first argument was passed as robot_ids list or step count
|
||||
actual_steps = robot_ids_or_steps if isinstance(robot_ids_or_steps, int) else steps
|
||||
|
||||
for _ in range(actual_steps):
|
||||
self.step()
|
||||
height = self.measure_robot_height()
|
||||
return height if height > 0.0 else fallback_height
|
||||
|
||||
def apply_external_force(
|
||||
self,
|
||||
force: Union[List[float], np.ndarray],
|
||||
body_id: Optional[int] = None,
|
||||
link_index: int = -1,
|
||||
position: Union[List[float], np.ndarray] = (0.0, 0.0, 0.0),
|
||||
frame: int = p.WORLD_FRAME,
|
||||
) -> None:
|
||||
"""Applies external force vector to target link (defaults to base link)."""
|
||||
bid = self._resolve_body_id(body_id)
|
||||
p.applyExternalForce(
|
||||
objectUniqueId=bid,
|
||||
linkIndex=link_index,
|
||||
forceObj=list(force),
|
||||
posObj=list(position),
|
||||
flags=frame,
|
||||
physicsClientId=self.physics_client,
|
||||
)
|
||||
+2
-2
@@ -1,4 +1,4 @@
|
||||
from .env import JackBotEnv, CurriculumCallback
|
||||
from .env import JackBotEnv
|
||||
from .model import ActorCritic
|
||||
|
||||
__all__ = ["JackBotEnv", "CurriculumCallback", "ActorCritic"]
|
||||
__all__ = ["JackBotEnv", "ActorCritic"]
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
"""
|
||||
ml/callbacks.py - Stable-Baselines3 Custom Callbacks for Logging & Curriculum Advancement
|
||||
"""
|
||||
import os
|
||||
import numpy as np
|
||||
from stable_baselines3.common.callbacks import BaseCallback
|
||||
|
||||
|
||||
class RewardLoggerCallback(BaseCallback):
|
||||
"""Logs individual reward component averages to TensorBoard."""
|
||||
|
||||
def __init__(self, verbose: int = 0):
|
||||
super().__init__(verbose)
|
||||
|
||||
def _on_step(self) -> bool:
|
||||
# Pull component averages from the environment vector
|
||||
for env_idx, env in enumerate(self.training_env.envs):
|
||||
if hasattr(env, "get_reward_component_averages"):
|
||||
averages = env.get_reward_component_averages()
|
||||
for key, val in averages.items():
|
||||
self.logger.record(f"reward_components/{key}", val)
|
||||
return True
|
||||
|
||||
|
||||
class CurriculumCallback(BaseCallback):
|
||||
"""Monitors evaluation performance and automatically manages curriculum progression."""
|
||||
|
||||
def __init__(self, eval_freq: int = 10000, verbose: int = 1):
|
||||
super().__init__(verbose)
|
||||
self.eval_freq = eval_freq
|
||||
|
||||
def _on_step(self) -> bool:
|
||||
if self.n_calls % self.eval_freq == 0:
|
||||
for env in self.training_env.envs:
|
||||
if hasattr(env, "get_current_robot_metrics"):
|
||||
metrics = env.get_current_robot_metrics()
|
||||
if metrics:
|
||||
phase = metrics[0].get("phase_name", "UNKNOWN")
|
||||
dist = metrics[0].get("distance_from_start", 0.0)
|
||||
self.logger.record("curriculum/phase_idx", phase)
|
||||
self.logger.record("curriculum/max_distance", dist)
|
||||
if self.verbose > 0:
|
||||
print(f"[CurriculumCallback] Step {self.num_timesteps}: Current Phase = {phase}, Max Dist = {dist:.2f}m")
|
||||
return True
|
||||
@@ -10,14 +10,12 @@ from collections import defaultdict
|
||||
import gymnasium as gym
|
||||
from gymnasium import spaces
|
||||
import numpy as np
|
||||
from stable_baselines3.common.callbacks import BaseCallback
|
||||
|
||||
from config import cfg
|
||||
from simulation import Simulation
|
||||
from Robot import Robot, PyBulletBackend
|
||||
from ml.SimManager import SimManager
|
||||
from ml.MetricsOverlay import MetricsHUD
|
||||
|
||||
# Color Palette RGBA for Terminated/Failed Robots
|
||||
COLOR_FAILED = [0.3, 0.3, 0.3, 0.6]
|
||||
|
||||
|
||||
@@ -30,7 +28,7 @@ class CurriculumPhase(IntEnum):
|
||||
|
||||
|
||||
class JackBotEnv(gym.Env):
|
||||
"""Gymnasium environment wrapping a single JackBot hexapod."""
|
||||
"""Gymnasium environment wrapping JackBot hexapod simulation."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -38,7 +36,7 @@ class JackBotEnv(gym.Env):
|
||||
random_command: bool = True,
|
||||
max_episode_steps: int = 3000,
|
||||
urdf_path: str = cfg.urdf_path,
|
||||
robot_mode: str = "direct",
|
||||
robot_mode: str = "direct", # "direct", "residual", or "kinematics"
|
||||
):
|
||||
super().__init__()
|
||||
self.robot_mode = robot_mode
|
||||
@@ -52,7 +50,6 @@ class JackBotEnv(gym.Env):
|
||||
self.total_steps = 0
|
||||
self.cumulative_reward = 0.0
|
||||
self.robot_reward = 0.0
|
||||
self.consecutive_still_steps = 0
|
||||
self.is_failed = False
|
||||
|
||||
self.episode_count = 0
|
||||
@@ -62,94 +59,55 @@ class JackBotEnv(gym.Env):
|
||||
self._curriculum_advanced = False
|
||||
self._first_reset = True
|
||||
|
||||
# Reward Component Tracking Initialization
|
||||
self.last_reward_components: Dict[str, float] = {}
|
||||
self.episode_reward_components_sum: Dict[str, float] = defaultdict(float)
|
||||
|
||||
# Dynamic Command Resampling Timing (60 Hz control loop)
|
||||
self.control_freq = 60
|
||||
self.min_cmd_hold_steps = int(2.0 * self.control_freq) # 120 steps (2s)
|
||||
self.max_cmd_hold_steps = int(6.0 * self.control_freq) # 360 steps (6s)
|
||||
self.min_cmd_hold_steps = int(2.0 * self.control_freq)
|
||||
self.max_cmd_hold_steps = int(6.0 * self.control_freq)
|
||||
self.next_cmd_resample_step = 0
|
||||
self.initial_stand_steps = 120 # Mandatory 2s standing window at episode reset
|
||||
self.initial_stand_steps = 120
|
||||
|
||||
# Initialize Simulation Manager
|
||||
self.sim_manager = SimManager(use_gui=self.use_gui)
|
||||
self.sim_manager.connect()
|
||||
# --- INITIALIZE PYBULLET SIMULATION ENGINE ---
|
||||
self.sim = Simulation(urdf_path=self.urdf_path, use_gui=self.use_gui)
|
||||
self.plane_id, self.pb_robot, self.revolute_joints = self.sim.load_scene()
|
||||
|
||||
# Connect physics world & load single robot
|
||||
self.plane, pb_robots, robot_joint_indices = self.sim_manager.load_scene(
|
||||
self.urdf_path, 0.0, self._robot_base_position
|
||||
)
|
||||
self.joint_indices = robot_joint_indices[0]
|
||||
|
||||
# Instantiate Robot Python wrapper
|
||||
self.robot = Robot(
|
||||
backend_type=PyBulletBackend(self.sim_manager),
|
||||
urdf_path=self.urdf_path
|
||||
)
|
||||
# --- INITIALIZE ROBOT WITH BACKEND ---
|
||||
self.backend = PyBulletBackend(self.sim)
|
||||
self.robot = Robot(backend_type=self.backend, urdf_path=self.urdf_path, mode=self.robot_mode)
|
||||
|
||||
action_dim = 18
|
||||
obs_dim = 18 + 4
|
||||
obs_dim = 18 + 3 # 18 Joint Angles + 3 Command Inputs [vx, vy, omega]
|
||||
|
||||
self.action_space = spaces.Box(-1.0, 1.0, shape=(action_dim,), dtype=np.float32)
|
||||
self.observation_space = spaces.Box(-np.inf, np.inf, shape=(obs_dim,), dtype=np.float32)
|
||||
|
||||
self.command = np.zeros(4, dtype=np.float32)
|
||||
self.command = np.zeros(3, dtype=np.float32) # [vx, vy, omega]
|
||||
self.last_action = np.zeros(action_dim, dtype=np.float32)
|
||||
self.target_height = 0.122
|
||||
self.collapse_height_fraction = 0.55
|
||||
self.tilt_failure_rad = 0.9
|
||||
|
||||
self.foot_link_indices = self._find_foot_link_indices()
|
||||
self.start_position = [0.0, 0.0, 0.0]
|
||||
self.max_distance_from_start = 0.0
|
||||
self.max_survival_steps = 0
|
||||
self.default_joint_angles = np.zeros(18, dtype=np.float32)
|
||||
|
||||
# Curriculum Initialization via Enum
|
||||
self.curriculum_phase = CurriculumPhase.STAND_ONLY
|
||||
|
||||
self.curriculum_stage_requirements = {
|
||||
CurriculumPhase.FORWARD: {
|
||||
"survival_steps": 300,
|
||||
"min_avg_height_ratio": 0.88,
|
||||
"max_avg_roll_pitch": 0.18,
|
||||
},
|
||||
CurriculumPhase.TURN_AND_DIRECTION: {
|
||||
"survival_steps": 500,
|
||||
"min_forward_distance": 2.5,
|
||||
"max_lateral_drift": 0.8,
|
||||
"min_avg_height_ratio": 0.85,
|
||||
"stability_roll_pitch": 0.25,
|
||||
},
|
||||
CurriculumPhase.OMNI_DIRECTION: {
|
||||
"survival_steps": 600,
|
||||
"min_distance": 5.0,
|
||||
"min_avg_height_ratio": 0.85,
|
||||
"stability_roll_pitch": 0.25,
|
||||
},
|
||||
CurriculumPhase.FULL_COMMAND: {
|
||||
"survival_steps": 750,
|
||||
"min_distance": 8.0,
|
||||
"min_avg_height_ratio": 0.85,
|
||||
"stability_roll_pitch": 0.20,
|
||||
},
|
||||
CurriculumPhase.FORWARD: {"survival_steps": 300, "min_avg_height_ratio": 0.88, "max_avg_roll_pitch": 0.18},
|
||||
CurriculumPhase.TURN_AND_DIRECTION: {"survival_steps": 500, "min_forward_distance": 2.5, "max_lateral_drift": 0.8, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.25},
|
||||
CurriculumPhase.OMNI_DIRECTION: {"survival_steps": 600, "min_distance": 5.0, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.25},
|
||||
CurriculumPhase.FULL_COMMAND: {"survival_steps": 750, "min_distance": 8.0, "min_avg_height_ratio": 0.85, "stability_roll_pitch": 0.20},
|
||||
}
|
||||
|
||||
self.hud = MetricsHUD(physics_client_id=self.sim_manager.physics_client)
|
||||
self.hud = MetricsHUD(physics_client_id=self.sim.physics_client)
|
||||
self.last_time = time.time()
|
||||
|
||||
def _robot_base_position(self, robot_id: int, spacing: float = 0.0) -> list[float]:
|
||||
return [0.0, 0.0, 0.2]
|
||||
|
||||
def _find_foot_link_indices(self) -> list:
|
||||
return self.sim_manager.get_foot_link_indices(self.pb_robot)
|
||||
|
||||
def sample_command(self) -> np.ndarray:
|
||||
"""Samples a command vector [vx, vy, omega] based on active curriculum phase."""
|
||||
phase = self.curriculum_phase
|
||||
|
||||
stand_probabilities = {
|
||||
stand_probs = {
|
||||
CurriculumPhase.STAND_ONLY: 1.0,
|
||||
CurriculumPhase.FORWARD: 0.25,
|
||||
CurriculumPhase.TURN_AND_DIRECTION: 0.20,
|
||||
@@ -157,28 +115,19 @@ class JackBotEnv(gym.Env):
|
||||
CurriculumPhase.FULL_COMMAND: 0.15,
|
||||
}
|
||||
|
||||
if np.random.random() < stand_probabilities.get(phase, 0.15):
|
||||
return np.zeros(4, dtype=np.float32)
|
||||
if np.random.random() < stand_probs.get(phase, 0.15):
|
||||
return np.zeros(3, dtype=np.float32)
|
||||
|
||||
if phase == CurriculumPhase.FORWARD:
|
||||
vx = np.random.uniform(0.15, 0.50)
|
||||
vy, vz, omega = 0.0, 0.0, 0.0
|
||||
vx, vy, omega = np.random.uniform(0.15, 0.50), 0.0, 0.0
|
||||
elif phase == CurriculumPhase.TURN_AND_DIRECTION:
|
||||
vx = np.random.uniform(-0.8, 0.8)
|
||||
vy, vz = 0.0, 0.0
|
||||
omega = np.random.uniform(-0.8, 0.8)
|
||||
vx, vy, omega = np.random.uniform(-0.8, 0.8), 0.0, np.random.uniform(-0.8, 0.8)
|
||||
elif phase == CurriculumPhase.OMNI_DIRECTION:
|
||||
vx = np.random.uniform(-0.8, 0.8)
|
||||
vy = np.random.uniform(-0.5, 0.5)
|
||||
vz = 0.0
|
||||
omega = np.random.uniform(-0.8, 0.8)
|
||||
vx, vy, omega = np.random.uniform(-0.8, 0.8), np.random.uniform(-0.5, 0.5), np.random.uniform(-0.8, 0.8)
|
||||
else:
|
||||
vx = np.random.uniform(-1.0, 1.0)
|
||||
vy = np.random.uniform(-1.0, 1.0)
|
||||
vz = 0.0
|
||||
omega = np.random.uniform(-1.0, 1.0)
|
||||
vx, vy, omega = np.random.uniform(-1.0, 1.0), np.random.uniform(-1.0, 1.0), np.random.uniform(-1.0, 1.0)
|
||||
|
||||
return np.array([vx, vy, vz, omega], dtype=np.float32)
|
||||
return np.array([vx, vy, omega], dtype=np.float32)
|
||||
|
||||
def reset(self, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None):
|
||||
super().reset(seed=seed)
|
||||
@@ -191,24 +140,21 @@ class JackBotEnv(gym.Env):
|
||||
self.episode_height_sum = 0.0
|
||||
self.episode_roll_sum = 0.0
|
||||
self.episode_pitch_sum = 0.0
|
||||
|
||||
# Reset Component Tracking Dictionary
|
||||
self.last_reward_components = {}
|
||||
self.episode_reward_components_sum = defaultdict(float)
|
||||
|
||||
spawn_pos = self._robot_base_position(0)
|
||||
spawn_pos = [0.0, 0.0, 0.20]
|
||||
spawn_orn = [0.0, 0.0, 0.0, 1.0]
|
||||
|
||||
self.sim_manager.reset_robot_base(self.pb_robot, spawn_pos, spawn_orn)
|
||||
# 1. Reset base pose and velocities
|
||||
self.sim.reset_robot_base(spawn_pos, spawn_orn)
|
||||
|
||||
# 2. Reset internal kinematics & hard reset joints in PyBullet
|
||||
self.robot.reset_to_init()
|
||||
|
||||
init_angles = self.robot.current_rad.data.flatten()
|
||||
self.sim_manager.hard_reset_joint_angles(init_angles, self.pb_robot)
|
||||
|
||||
if self.use_gui:
|
||||
self.sim_manager.set_robot_color(self.pb_robot, [1.0, 1.0, 1.0, 1.0])
|
||||
self.sim.set_robot_color([1.0, 1.0, 1.0, 1.0])
|
||||
|
||||
self.consecutive_still_steps = 0
|
||||
self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32)
|
||||
self.max_distance_from_start = 0.0
|
||||
self.max_survival_steps = 0
|
||||
@@ -217,21 +163,15 @@ class JackBotEnv(gym.Env):
|
||||
self.curriculum_phase = CurriculumPhase.STAND_ONLY
|
||||
self._first_reset = False
|
||||
|
||||
self.command = np.zeros(4, dtype=np.float32)
|
||||
self.command = np.zeros(3, dtype=np.float32)
|
||||
self.next_cmd_resample_step = self.initial_stand_steps
|
||||
|
||||
pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
|
||||
self.start_position = [float(pos[0]), float(pos[1]), float(pos[2])]
|
||||
pos, _ = self.sim.get_robot_pose()
|
||||
self.start_position = list(pos)
|
||||
|
||||
# Fixed single-robot settlement call
|
||||
self.target_height = self.sim_manager.settle_and_measure_height(
|
||||
steps=200, fallback_height=0.122
|
||||
)
|
||||
|
||||
self.default_joint_angles = np.array(
|
||||
self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices),
|
||||
dtype=np.float32
|
||||
)
|
||||
# Drop settlement
|
||||
self.target_height = self.sim.settle_and_measure_height(steps=200, fallback_height=0.122)
|
||||
self.default_joint_angles = self.sim.get_robot_joint_angles()
|
||||
|
||||
if self.use_gui:
|
||||
self.hud.reset()
|
||||
@@ -239,36 +179,8 @@ class JackBotEnv(gym.Env):
|
||||
|
||||
return self._get_obs(), {}
|
||||
|
||||
def get_reward_component_averages(self) -> Dict[str, float]:
|
||||
"""Calculates step-averaged scores for each sub-reward component."""
|
||||
steps = max(1, self.step_count)
|
||||
return {k: float(v / steps) for k, v in self.episode_reward_components_sum.items()}
|
||||
|
||||
def get_current_robot_metrics(self) -> list:
|
||||
"""Returns metric summary for callbacks."""
|
||||
if self.is_failed:
|
||||
return []
|
||||
|
||||
pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
|
||||
linear_vel, angular_vel = self.sim_manager.get_robot_velocity(self.pb_robot)
|
||||
start_x, start_y, _ = self.start_position
|
||||
dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
|
||||
speed = float(np.linalg.norm(np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)))
|
||||
yaw_rate = float(abs(angular_vel[2]))
|
||||
|
||||
return [{
|
||||
"reward": float(self.robot_reward),
|
||||
"distance_from_start": dist,
|
||||
"speed": speed,
|
||||
"yaw_rate": yaw_rate,
|
||||
"alive": True,
|
||||
"survival_steps": int(self.step_count),
|
||||
"phase_name": self.curriculum_phase.name,
|
||||
}]
|
||||
|
||||
def _get_obs(self) -> np.ndarray:
|
||||
joint_angles = self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices)
|
||||
return np.concatenate([joint_angles, self.command]).astype(np.float32)
|
||||
return self.robot.get_observation(command=self.command)
|
||||
|
||||
def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
|
||||
self.step_count += 1
|
||||
@@ -276,28 +188,25 @@ class JackBotEnv(gym.Env):
|
||||
previous_action = self.last_action.copy()
|
||||
self.last_action = action.copy()
|
||||
|
||||
# Command resampling
|
||||
if self.random_command and (self.step_count >= self.next_cmd_resample_step or self._curriculum_advanced):
|
||||
self.command = self.sample_command()
|
||||
random_interval = np.random.randint(self.min_cmd_hold_steps, self.max_cmd_hold_steps + 1)
|
||||
self.next_cmd_resample_step = self.step_count + random_interval
|
||||
|
||||
self.robot.step_with_command(command=self.command, action=action, mode=self.robot_mode)
|
||||
# Extract [vx, vy, omega]
|
||||
cmd_vx, cmd_vy, cmd_omega = self.command
|
||||
|
||||
if self.robot_mode != "kinematics_only" and self.step_count % 60 == 0:
|
||||
# Mirror main.py input resolution logic: update robot_state and vector_dirmov directly
|
||||
self.robot.robot_state = "walking" if (abs(cmd_vx) > 0.01 or abs(cmd_vy) > 0.01 or abs(cmd_omega) > 0.01) else "idle"
|
||||
self.robot.vector_dirmov = [float(cmd_vx), float(cmd_vy), float(cmd_omega)]
|
||||
|
||||
# Delegate execution tick to Robot instance
|
||||
self.robot.tick(action=action)
|
||||
|
||||
if self.robot_mode != "kinematics" and self.step_count % 60 == 0:
|
||||
random_force = np.random.uniform(-2.0, 2.0, size=2)
|
||||
self.sim_manager.apply_external_force(
|
||||
body_id=self.pb_robot,
|
||||
force=[random_force[0], random_force[1], 0.0]
|
||||
)
|
||||
|
||||
render_freq = 10
|
||||
if self.use_gui and self.step_count % render_freq != 0:
|
||||
self.sim_manager.set_rendering(False)
|
||||
|
||||
self.sim_manager.step()
|
||||
|
||||
if self.use_gui and self.step_count % render_freq == 0:
|
||||
self.sim_manager.set_rendering(True)
|
||||
self.sim.apply_external_force(force=[random_force[0], random_force[1], 0.0])
|
||||
|
||||
self._update_robot_failure()
|
||||
self._update_distance_metrics()
|
||||
@@ -311,19 +220,15 @@ class JackBotEnv(gym.Env):
|
||||
|
||||
terminated = self.is_failed
|
||||
truncated = self.step_count >= self.max_episode_steps
|
||||
|
||||
info = {
|
||||
"reward_components": self.last_reward_components.copy()
|
||||
}
|
||||
info = {"reward_components": self.last_reward_components.copy()}
|
||||
|
||||
if self.step_count % 120 == 0 and self.use_gui:
|
||||
self._update_hud()
|
||||
|
||||
# Gymnasium standard 5-tuple return
|
||||
return obs, reward, terminated, truncated, info
|
||||
|
||||
def _update_distance_metrics(self):
|
||||
pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
|
||||
pos, _ = self.sim.get_robot_pose()
|
||||
self.episode_height_sum += float(pos[2])
|
||||
start_x, start_y, _ = self.start_position
|
||||
dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
|
||||
@@ -336,7 +241,6 @@ class JackBotEnv(gym.Env):
|
||||
|
||||
req = self.curriculum_stage_requirements[next_phase]
|
||||
survival_ok = self.max_survival_steps >= req["survival_steps"]
|
||||
|
||||
avg_roll = self.episode_roll_sum / max(1, self.step_count)
|
||||
avg_pitch = self.episode_pitch_sum / max(1, self.step_count)
|
||||
max_allowed_angle = req.get("max_avg_roll_pitch", 0.20)
|
||||
@@ -346,10 +250,8 @@ class JackBotEnv(gym.Env):
|
||||
required_min_avg_height = self.target_height * req.get("min_avg_height_ratio", 0.85)
|
||||
height_ok = avg_height >= required_min_avg_height
|
||||
|
||||
pos, _ = self.sim_manager.get_robot_pose(self.pb_robot)
|
||||
start_x, start_y, _ = self.start_position
|
||||
dx = pos[0] - start_x
|
||||
dy = pos[1] - start_y
|
||||
pos, _ = self.sim.get_robot_pose()
|
||||
dx, dy = pos[0] - self.start_position[0], pos[1] - self.start_position[1]
|
||||
dist_2d = math.hypot(dx, dy)
|
||||
|
||||
distance_ok = True
|
||||
@@ -358,10 +260,7 @@ class JackBotEnv(gym.Env):
|
||||
elif "min_distance" in req:
|
||||
distance_ok = dist_2d >= req["min_distance"]
|
||||
|
||||
drift_ok = True
|
||||
if "max_lateral_drift" in req:
|
||||
drift_ok = abs(dy) <= req["max_lateral_drift"]
|
||||
|
||||
drift_ok = abs(dy) <= req["max_lateral_drift"] if "max_lateral_drift" in req else True
|
||||
return survival_ok and height_ok and stability_ok and distance_ok and drift_ok
|
||||
|
||||
def _update_curriculum(self):
|
||||
@@ -371,174 +270,117 @@ class JackBotEnv(gym.Env):
|
||||
if self.curriculum_phase < CurriculumPhase.FORWARD and (self._phase_progress_ready(CurriculumPhase.FORWARD) or forced_forward):
|
||||
self.curriculum_phase = CurriculumPhase.FORWARD
|
||||
self._curriculum_advanced = True
|
||||
reason = "FORCED (2500 steps)" if forced_forward else "MET"
|
||||
print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked [{reason}] at step {self.step_count}")
|
||||
|
||||
elif self.curriculum_phase < CurriculumPhase.TURN_AND_DIRECTION and self._phase_progress_ready(CurriculumPhase.TURN_AND_DIRECTION):
|
||||
self.curriculum_phase = CurriculumPhase.TURN_AND_DIRECTION
|
||||
self._curriculum_advanced = True
|
||||
print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
|
||||
|
||||
elif self.curriculum_phase < CurriculumPhase.OMNI_DIRECTION and self._phase_progress_ready(CurriculumPhase.OMNI_DIRECTION):
|
||||
self.curriculum_phase = CurriculumPhase.OMNI_DIRECTION
|
||||
self._curriculum_advanced = True
|
||||
print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
|
||||
|
||||
elif self.curriculum_phase < CurriculumPhase.FULL_COMMAND and self._phase_progress_ready(CurriculumPhase.FULL_COMMAND):
|
||||
self.curriculum_phase = CurriculumPhase.FULL_COMMAND
|
||||
self._curriculum_advanced = True
|
||||
print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
|
||||
|
||||
def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> float:
|
||||
"""
|
||||
Calculates task rewards using normalized Exponential Kernels and tracks component terms.
|
||||
"""
|
||||
pos, (roll, pitch, yaw) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot)
|
||||
linear_vel, angular_vel = self.sim_manager.get_robot_velocity(self.pb_robot)
|
||||
current_joints = np.array(
|
||||
self.sim_manager.get_robot_joint_angles(self.pb_robot, self.joint_indices),
|
||||
dtype=np.float32
|
||||
)
|
||||
pos, (roll, pitch, yaw) = self.sim.get_robot_pose_and_rpy()
|
||||
linear_vel, angular_vel = self.sim.get_robot_velocity()
|
||||
current_joints = self.sim.get_robot_joint_angles()
|
||||
|
||||
cmd_vx, cmd_vy, _, cmd_yaw = self.command
|
||||
cmd_vx, cmd_vy, cmd_yaw = self.command
|
||||
cmd_norm = math.hypot(cmd_vx, cmd_vy)
|
||||
|
||||
VEL_DEADBAND = 0.04
|
||||
YAW_DEADBAND = 0.05
|
||||
|
||||
raw_speed = math.hypot(linear_vel[0], linear_vel[1])
|
||||
if raw_speed < VEL_DEADBAND:
|
||||
filtered_vx, filtered_vy = 0.0, 0.0
|
||||
filtered_speed = 0.0
|
||||
else:
|
||||
filtered_vx, filtered_vy = linear_vel[0], linear_vel[1]
|
||||
filtered_speed = raw_speed
|
||||
filtered_vx, filtered_vy = (linear_vel[0], linear_vel[1]) if raw_speed >= 0.04 else (0.0, 0.0)
|
||||
filtered_speed = raw_speed if raw_speed >= 0.04 else 0.0
|
||||
|
||||
raw_yaw_rate = abs(angular_vel[2])
|
||||
if raw_yaw_rate < YAW_DEADBAND:
|
||||
filtered_yaw_rate = 0.0
|
||||
else:
|
||||
filtered_yaw_rate = angular_vel[2]
|
||||
filtered_yaw_rate = angular_vel[2] if raw_yaw_rate >= 0.05 else 0.0
|
||||
|
||||
height_error = pos[2] - self.target_height
|
||||
r_height = math.exp(-150.0 * (height_error ** 2))
|
||||
r_stability = math.exp(-25.0 * (roll**2 + pitch**2))
|
||||
r_pose = math.exp(-2.0 * np.mean(np.square(current_joints - self.default_joint_angles)))
|
||||
r_smoothness = math.exp(-0.1 * np.mean(np.square(action - previous_action)))
|
||||
|
||||
orientation_error = roll**2 + pitch**2
|
||||
r_stability = math.exp(-25.0 * orientation_error)
|
||||
|
||||
joint_error = np.mean(np.square(current_joints - self.default_joint_angles))
|
||||
r_pose = math.exp(-2.0 * joint_error)
|
||||
|
||||
action_delta = np.mean(np.square(action - previous_action))
|
||||
r_smoothness = math.exp(-0.1 * action_delta)
|
||||
|
||||
r_lin_vel = 0.0
|
||||
r_ang_vel = 0.0
|
||||
stillness_penalty = 0.0
|
||||
gated_zero = 0.0
|
||||
r_lin_vel, r_ang_vel, stillness_penalty = 0.0, 0.0, 0.0
|
||||
|
||||
if cmd_norm < 0.05 and abs(cmd_yaw) < 0.05:
|
||||
# STANDING MODE
|
||||
w_height = 0.35
|
||||
w_stability = 0.35
|
||||
w_pose = 0.20
|
||||
w_smoothness = 0.10
|
||||
w_lin_vel = 0.0
|
||||
w_ang_vel = 0.0
|
||||
|
||||
total_reward = (
|
||||
(w_height * r_height)
|
||||
+ (w_stability * r_stability)
|
||||
+ (w_pose * r_pose)
|
||||
+ (w_smoothness * r_smoothness)
|
||||
)
|
||||
w_height, w_stability, w_pose, w_smoothness = 0.35, 0.35, 0.20, 0.10
|
||||
total_reward = (w_height * r_height) + (w_stability * r_stability) + (w_pose * r_pose) + (w_smoothness * r_smoothness)
|
||||
else:
|
||||
# WALKING / TURNING MODE
|
||||
is_moving = (filtered_speed > 0.0) or (abs(filtered_yaw_rate) > 0.0)
|
||||
|
||||
target_vx = cmd_vx * self.max_robot_speed
|
||||
target_vy = cmd_vy * self.max_robot_speed
|
||||
target_vx, target_vy = cmd_vx * self.max_robot_speed, cmd_vy * self.max_robot_speed
|
||||
target_speed = math.hypot(target_vx, target_vy)
|
||||
|
||||
if not is_moving:
|
||||
gated_zero = 1.0
|
||||
total_reward = 0.0
|
||||
w_lin_vel, w_ang_vel, w_height, w_stability, w_pose, w_smoothness = 0, 0, 0, 0, 0, 0
|
||||
else:
|
||||
lin_vel_error = (filtered_vx - target_vx)**2 + (filtered_vy - target_vy)**2
|
||||
r_lin_vel = math.exp(-25.0 * lin_vel_error)
|
||||
|
||||
ang_vel_error = (filtered_yaw_rate - cmd_yaw)**2
|
||||
r_ang_vel = math.exp(-15.0 * ang_vel_error)
|
||||
r_ang_vel = math.exp(-15.0 * ((filtered_yaw_rate - cmd_yaw)**2))
|
||||
|
||||
if target_speed > 0.08 and raw_speed < 0.03:
|
||||
r_lin_vel = 0.0
|
||||
stillness_penalty = -0.25
|
||||
|
||||
w_lin_vel = 0.55
|
||||
w_ang_vel = 0.15
|
||||
w_height = 0.10
|
||||
w_stability = 0.12
|
||||
w_pose = 0.0
|
||||
w_smoothness = 0.08
|
||||
|
||||
w_lin_vel, w_ang_vel, w_height, w_stability, w_smoothness = 0.55, 0.15, 0.10, 0.12, 0.08
|
||||
total_reward = (
|
||||
(w_lin_vel * r_lin_vel)
|
||||
+ (w_ang_vel * r_ang_vel)
|
||||
+ (w_height * r_height)
|
||||
+ (w_stability * r_stability)
|
||||
+ (w_smoothness * r_smoothness)
|
||||
+ stillness_penalty
|
||||
(w_lin_vel * r_lin_vel) + (w_ang_vel * r_ang_vel) + (w_height * r_height)
|
||||
+ (w_stability * r_stability) + (w_smoothness * r_smoothness) + stillness_penalty
|
||||
)
|
||||
|
||||
final_reward = float(total_reward / 10.0)
|
||||
|
||||
comp = {
|
||||
"lin_vel": float((w_lin_vel * r_lin_vel) / 10.0),
|
||||
"ang_vel": float((w_ang_vel * r_ang_vel) / 10.0),
|
||||
"height": float((w_height * r_height) / 10.0),
|
||||
"stability": float((w_stability * r_stability) / 10.0),
|
||||
"pose": float((w_pose * r_pose) / 10.0),
|
||||
"smoothness": float((w_smoothness * r_smoothness) / 10.0),
|
||||
"stillness_penalty": float(stillness_penalty / 10.0),
|
||||
"gated_zero": gated_zero,
|
||||
"total_step_reward": final_reward,
|
||||
self.last_reward_components = {
|
||||
"height": float(r_height),
|
||||
"stability": float(r_stability),
|
||||
"pose": float(r_pose),
|
||||
"smoothness": float(r_smoothness),
|
||||
"lin_vel": float(r_lin_vel),
|
||||
"ang_vel": float(r_ang_vel),
|
||||
"total": final_reward,
|
||||
}
|
||||
|
||||
self.last_reward_components = comp
|
||||
for key, val in comp.items():
|
||||
self.episode_reward_components_sum[key] += val
|
||||
for k, v in self.last_reward_components.items():
|
||||
self.episode_reward_components_sum[k] += v
|
||||
|
||||
return final_reward
|
||||
|
||||
def get_reward_component_averages(self) -> Dict[str, float]:
|
||||
steps = max(1, self.step_count)
|
||||
return {k: v / steps for k, v in self.episode_reward_components_sum.items()}
|
||||
|
||||
def get_current_robot_metrics(self) -> List[Dict[str, Any]]:
|
||||
linear_vel, angular_vel = self.sim.get_robot_velocity()
|
||||
speed = float(math.hypot(linear_vel[0], linear_vel[1]))
|
||||
yaw_rate = float(abs(angular_vel[2]))
|
||||
|
||||
metrics = {
|
||||
"alive": not self.is_failed,
|
||||
"phase_name": self.curriculum_phase.name,
|
||||
"reward": float(self.cumulative_reward),
|
||||
"speed": speed,
|
||||
"yaw_rate": yaw_rate,
|
||||
"distance_from_start": float(self.max_distance_from_start),
|
||||
"survival_steps": int(self.max_survival_steps),
|
||||
}
|
||||
return [metrics]
|
||||
|
||||
def _update_hud(self):
|
||||
if not self.use_gui:
|
||||
return
|
||||
|
||||
now = time.time()
|
||||
fps = 1.0 / max(now - self.last_time, 1e-5)
|
||||
self.last_time = now
|
||||
|
||||
pos, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot)
|
||||
|
||||
pos, (roll, pitch, _) = self.sim.get_robot_pose_and_rpy()
|
||||
self.hud.update(
|
||||
episode=self.episode_count,
|
||||
step=self.total_steps,
|
||||
robot_rewards=[self.robot_reward],
|
||||
cmd_vel=self.command,
|
||||
fps=fps,
|
||||
height=pos[2],
|
||||
roll_pitch=(math.degrees(roll), math.degrees(pitch))
|
||||
episode=self.episode_count, step=self.total_steps, robot_rewards=[self.robot_reward],
|
||||
cmd_vel=self.command, fps=fps, height=pos[2], roll_pitch=(math.degrees(roll), math.degrees(pitch))
|
||||
)
|
||||
|
||||
def _update_robot_failure(self):
|
||||
if self.is_failed:
|
||||
if self.is_failed or self.step_count < 15:
|
||||
return
|
||||
|
||||
if self.step_count < 15:
|
||||
return
|
||||
|
||||
position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robot)
|
||||
position, (roll, pitch, _) = self.sim.get_robot_pose_and_rpy()
|
||||
collapse_threshold = max(0.04, self.collapse_height_fraction * self.target_height)
|
||||
is_tilted = abs(roll) > self.tilt_failure_rad or abs(pitch) > self.tilt_failure_rad
|
||||
is_collapsed = position[2] < collapse_threshold
|
||||
@@ -546,105 +388,7 @@ class JackBotEnv(gym.Env):
|
||||
if is_tilted or is_collapsed:
|
||||
self.is_failed = True
|
||||
if self.use_gui:
|
||||
self.sim_manager.set_robot_color(self.pb_robot, COLOR_FAILED)
|
||||
self.sim.set_robot_color(COLOR_FAILED)
|
||||
|
||||
def close(self):
|
||||
self.sim_manager.disconnect()
|
||||
|
||||
|
||||
class CurriculumCallback(BaseCallback):
|
||||
"""Logs curriculum phase breakdown and best performance metrics to TensorBoard."""
|
||||
|
||||
def __init__(self, verbose=0):
|
||||
super().__init__(verbose)
|
||||
self.best_speed = 0.0
|
||||
self.best_yaw_rate = 0.0
|
||||
self.best_distance = 0.0
|
||||
self.best_survival_steps = 0.0
|
||||
self.best_reward = -float('inf')
|
||||
|
||||
def _on_step(self) -> bool:
|
||||
return True
|
||||
|
||||
def _on_rollout_end(self) -> bool:
|
||||
try:
|
||||
vec_env = self.training_env
|
||||
alive_metrics = vec_env.env_method("get_current_robot_metrics")
|
||||
|
||||
self.best_speed = 0.0
|
||||
self.best_yaw_rate = 0.0
|
||||
self.best_distance = 0.0
|
||||
self.best_survival_steps = 0.0
|
||||
self.best_reward = -float('inf')
|
||||
|
||||
phase_counts = {
|
||||
"stand_only": 0,
|
||||
"forward": 0,
|
||||
"turn_and_direction": 0,
|
||||
"omni_direction": 0,
|
||||
"full_command": 0,
|
||||
}
|
||||
|
||||
for worker_res in alive_metrics:
|
||||
for metrics in worker_res:
|
||||
if not metrics.get("alive", False):
|
||||
continue
|
||||
|
||||
phase_key = metrics.get("phase_name", "STAND_ONLY").lower()
|
||||
if phase_key in phase_counts:
|
||||
phase_counts[phase_key] += 1
|
||||
|
||||
if metrics["reward"] > self.best_reward:
|
||||
self.best_reward = float(metrics["reward"])
|
||||
if metrics["speed"] > self.best_speed:
|
||||
self.best_speed = float(metrics["speed"])
|
||||
if metrics["yaw_rate"] > self.best_yaw_rate:
|
||||
self.best_yaw_rate = float(metrics["yaw_rate"])
|
||||
if metrics["distance_from_start"] > self.best_distance:
|
||||
self.best_distance = float(metrics["distance_from_start"])
|
||||
if metrics["survival_steps"] > self.best_survival_steps:
|
||||
self.best_survival_steps = float(metrics["survival_steps"])
|
||||
|
||||
for phase_name, count in phase_counts.items():
|
||||
self.logger.record(f"phase/{phase_name}", count)
|
||||
|
||||
self.logger.record("custom/best_reward", float(self.best_reward) if np.isfinite(self.best_reward) else 0.0)
|
||||
self.logger.record("custom/best_survival_steps", float(self.best_survival_steps))
|
||||
self.logger.record("custom/best_distance_from_start_m", float(self.best_distance))
|
||||
self.logger.record("custom/best_speed_mps", float(self.best_speed))
|
||||
self.logger.record("custom/best_yaw_rate_rads", float(self.best_yaw_rate))
|
||||
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return True
|
||||
|
||||
|
||||
class RewardLoggerCallback(BaseCallback):
|
||||
"""
|
||||
Logs step-averaged individual reward components to TensorBoard during PPO training.
|
||||
"""
|
||||
|
||||
def __init__(self, verbose=0):
|
||||
super().__init__(verbose)
|
||||
|
||||
def _on_step(self) -> bool:
|
||||
return True
|
||||
|
||||
def _on_rollout_end(self) -> bool:
|
||||
try:
|
||||
vec_env = self.training_env
|
||||
all_comp_averages = vec_env.env_method("get_reward_component_averages")
|
||||
|
||||
if not all_comp_averages:
|
||||
return True
|
||||
|
||||
keys = all_comp_averages[0].keys()
|
||||
for key in keys:
|
||||
avg_val = np.mean([env_comp.get(key, 0.0) for env_comp in all_comp_averages])
|
||||
self.logger.record(f"reward_components/{key}", float(avg_val))
|
||||
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return True
|
||||
self.sim.disconnect()
|
||||
+6
-13
@@ -6,31 +6,27 @@ import numpy as np
|
||||
from ml.env import JackBotEnv, CurriculumPhase
|
||||
|
||||
def evaluate_kinematics(episode_length: int = 1000):
|
||||
# Use kinematics_only mode so inverse kinematics generates gait motion from commands
|
||||
env = JackBotEnv(
|
||||
use_gui=True,
|
||||
random_command=False,
|
||||
max_episode_steps=episode_length,
|
||||
robot_mode="kinematics_only"
|
||||
robot_mode="kinematics"
|
||||
)
|
||||
|
||||
print("\n" + "=" * 70)
|
||||
print(" RUNNING MULTI-PHASE REWARD BENCHMARK (KINEMATICS MODE)")
|
||||
print("=" * 70 + "\n")
|
||||
|
||||
# Define test suite covering every curriculum stage
|
||||
phase_configs = [
|
||||
(CurriculumPhase.STAND_ONLY, "STAND ONLY", np.array([0.0, 0.0, 0.0, 0.0], dtype=np.float32)),
|
||||
(CurriculumPhase.FORWARD, "FORWARD GAIT", np.array([0.1, 0.0, 0.0, 0.0], dtype=np.float32)),
|
||||
(CurriculumPhase.TURN_AND_DIRECTION, "FORWARD + YAW TURN", np.array([0.3, 0.0, 0.0, 0.4], dtype=np.float32)),
|
||||
(CurriculumPhase.OMNI_DIRECTION, "STRIDE LATERAL", np.array([0.2, 0.3, 0.0, 0.0], dtype=np.float32)),
|
||||
(CurriculumPhase.FULL_COMMAND, "FULL OMNI COMBINATION",np.array([0.3, 0.2, 0.0, 0.3], dtype=np.float32)),
|
||||
(CurriculumPhase.FORWARD, "FORWARD GAIT", np.array([1.0, 0.0, 0.0], dtype=np.float32)),
|
||||
(CurriculumPhase.TURN_AND_DIRECTION, "FORWARD + YAW TURN", np.array([0.5, 0.0, 0.4], dtype=np.float32)),
|
||||
(CurriculumPhase.OMNI_DIRECTION, "STRIDE LATERAL", np.array([0.5, 0.5, 0.0], dtype=np.float32)),
|
||||
(CurriculumPhase.FULL_COMMAND, "FULL OMNI COMBINATION",np.array([0.3, 0.2, 0.3], dtype=np.float32)),
|
||||
]
|
||||
|
||||
for phase_enum, label, cmd in phase_configs:
|
||||
obs, _ = env.reset()
|
||||
|
||||
# Force specific curriculum phase & target command
|
||||
env.curriculum_phase = phase_enum
|
||||
env.command = cmd.copy()
|
||||
|
||||
@@ -39,9 +35,7 @@ def evaluate_kinematics(episode_length: int = 1000):
|
||||
step_count = 0
|
||||
|
||||
while not done:
|
||||
# Action array is unused in kinematics_only mode
|
||||
dummy_action = np.zeros(18, dtype=np.float32)
|
||||
|
||||
obs, reward, terminated, truncated, _ = env.step(dummy_action)
|
||||
total_reward += reward
|
||||
step_count += 1
|
||||
@@ -49,11 +43,10 @@ def evaluate_kinematics(episode_length: int = 1000):
|
||||
|
||||
time.sleep(1.0 / 60.0)
|
||||
|
||||
# Retrieve detailed component averages
|
||||
comp_averages = env.get_reward_component_averages()
|
||||
|
||||
print(f"\n--- Episode Stage: [{phase_enum.name}] ({label}) ---")
|
||||
print(f"Command Applied: vx={cmd[0]:.2f}, vy={cmd[1]:.2f}, vz={cmd[2]:.2f}, yaw={cmd[3]:.2f}")
|
||||
print(f"Command Applied: vx={cmd[0]:.2f}, vy={cmd[1]:.2f}, yaw={cmd[2]:.2f}")
|
||||
print(f"Total Episode Reward: {total_reward:.4f}")
|
||||
print("Component Step Averages:")
|
||||
for name, value in comp_averages.items():
|
||||
|
||||
+2
-1
@@ -15,7 +15,8 @@ from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv
|
||||
from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
|
||||
|
||||
sys.path.append(str(Path(__file__).resolve().parent.parent))
|
||||
from ml.env import JackBotEnv, CurriculumCallback
|
||||
from ml.env import JackBotEnv
|
||||
from ml.callbacks import CurriculumCallback
|
||||
|
||||
# Silence SB3's UserWarning about SubprocVecEnv vs DummyVecEnv
|
||||
warnings.filterwarnings("ignore", category=UserWarning, module="stable_baselines3")
|
||||
|
||||
+154
-43
@@ -1,89 +1,200 @@
|
||||
"""
|
||||
simulation.py - PyBullet Simulation Interface & Standalone Runner
|
||||
simulation.py - PyBullet Simulation Interface & Physics Engine
|
||||
Consolidates scene management, physics queries, motor control, and rendering.
|
||||
"""
|
||||
import time
|
||||
import math
|
||||
from typing import List, Tuple, Optional, Union
|
||||
import numpy as np
|
||||
import pybullet as p
|
||||
import pybullet_data # Added missing import
|
||||
import pybullet_data
|
||||
|
||||
from config import cfg
|
||||
import DataTypes as dt
|
||||
|
||||
|
||||
class Simulation:
|
||||
def __init__(self, urdf_path: str = cfg.urdf_path):
|
||||
"""Single PyBullet simulation manager providing getters/setters for JackBot."""
|
||||
|
||||
def __init__(self, urdf_path: str = cfg.urdf_path, use_gui: bool = True):
|
||||
self.urdf_path = urdf_path
|
||||
self.use_gui = use_gui
|
||||
self.physics_client: Optional[int] = None
|
||||
self.plane_id: Optional[int] = None
|
||||
self.robot_id: Optional[int] = None
|
||||
self.revolute_joints: List[int] = []
|
||||
|
||||
self.connect()
|
||||
|
||||
def connect(self) -> None:
|
||||
"""Establishes connection to PyBullet GUI or DIRECT mode."""
|
||||
if self.physics_client is not None and p.isConnected(self.physics_client):
|
||||
return
|
||||
|
||||
flags = p.GUI if self.use_gui else p.DIRECT
|
||||
self.physics_client = p.connect(flags)
|
||||
|
||||
# Connect to PyBullet GUI
|
||||
self.physicsClient = p.connect(p.GUI)
|
||||
p.setAdditionalSearchPath(pybullet_data.getDataPath())
|
||||
p.setGravity(0, 0, -9.81)
|
||||
p.setGravity(0, 0, -9.81, physicsClientId=self.physics_client)
|
||||
|
||||
# Load plane and robot URDF
|
||||
self.planeId = p.loadURDF("plane.urdf")
|
||||
self.robot = p.loadURDF(self.urdf_path, [0, 0, 0.2])
|
||||
if self.use_gui:
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.physics_client)
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0, physicsClientId=self.physics_client)
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
|
||||
p.resetDebugVisualizerCamera(
|
||||
cameraDistance=1.0, cameraYaw=50, cameraPitch=-35, cameraTargetPosition=[0, 0, 0],
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
|
||||
def load_scene(self, spawn_pos: Optional[List[float]] = None) -> Tuple[int, int, List[int]]:
|
||||
"""Loads plane and robot URDF, discovering revolute joint indices dynamically."""
|
||||
if spawn_pos is None:
|
||||
spawn_pos = [0.0, 0.0, 0.20]
|
||||
|
||||
self.plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client)
|
||||
self.robot_id = p.loadURDF(self.urdf_path, spawn_pos, physicsClientId=self.physics_client)
|
||||
|
||||
# Discover revolute joint indices dynamically
|
||||
self.revolute_joints = []
|
||||
for j in range(p.getNumJoints(self.robot)):
|
||||
joint_info = p.getJointInfo(self.robot, j)
|
||||
if joint_info[2] == p.JOINT_REVOLUTE:
|
||||
for j in range(p.getNumJoints(self.robot_id, physicsClientId=self.physics_client)):
|
||||
info = p.getJointInfo(self.robot_id, j, physicsClientId=self.physics_client)
|
||||
if info[2] == p.JOINT_REVOLUTE:
|
||||
self.revolute_joints.append(j)
|
||||
|
||||
self.set_all_joints_to_90()
|
||||
p.resetDebugVisualizerCamera(
|
||||
cameraDistance=1.0,
|
||||
cameraYaw=50,
|
||||
cameraPitch=-35,
|
||||
cameraTargetPosition=[0, 0, 0],
|
||||
)
|
||||
return self.plane_id, self.robot_id, self.revolute_joints
|
||||
|
||||
def set_all_joints_to_90(self):
|
||||
for joint_index in self.revolute_joints:
|
||||
p.resetJointState(self.robot, joint_index, math.radians(90))
|
||||
# --- ACTUATION (SETTERS) ---
|
||||
|
||||
def set_robot_joint_angles(
|
||||
self, target_angles: Union[np.ndarray, List[float], dt.RadArray]
|
||||
) -> None:
|
||||
"""Applies motor torque to pull joints toward target position angles."""
|
||||
if isinstance(target_angles, dt.RadArray):
|
||||
radflat = target_angles.data.flatten()
|
||||
elif isinstance(target_angles, np.ndarray):
|
||||
radflat = target_angles.flatten()
|
||||
else:
|
||||
radflat = target_angles
|
||||
|
||||
def updatePos(self, current_rad: dt.RadArray):
|
||||
radflat = current_rad.data.flatten()
|
||||
for joint_index, target_angle in zip(self.revolute_joints, radflat):
|
||||
p.setJointMotorControl2(
|
||||
bodyIndex=self.robot,
|
||||
bodyIndex=self.robot_id,
|
||||
jointIndex=joint_index,
|
||||
controlMode=p.POSITION_CONTROL,
|
||||
targetPosition=target_angle,
|
||||
targetPosition=float(target_angle),
|
||||
force=500,
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
|
||||
def step(self):
|
||||
p.stepSimulation()
|
||||
def hard_reset_joint_angles(self, target_angles: Union[np.ndarray, dt.RadArray]) -> None:
|
||||
"""Instantly teleports joint angles to target positions, clearing velocity state."""
|
||||
radflat = target_angles.data.flatten() if isinstance(target_angles, dt.RadArray) else target_angles.flatten()
|
||||
for joint_index, target_angle in zip(self.revolute_joints, radflat):
|
||||
p.resetJointState(
|
||||
bodyUniqueId=self.robot_id,
|
||||
jointIndex=joint_index,
|
||||
targetValue=float(target_angle),
|
||||
targetVelocity=0.0,
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
|
||||
def disconnect(self):
|
||||
if p.isConnected(self.physicsClient):
|
||||
p.disconnect(self.physicsClient)
|
||||
def reset_robot_base(
|
||||
self,
|
||||
pos: Optional[List[float]] = None,
|
||||
orn: Optional[List[float]] = None,
|
||||
linear_velocity: Optional[List[float]] = None,
|
||||
angular_velocity: Optional[List[float]] = None
|
||||
) -> None:
|
||||
"""Resets root torso position, orientation quaternion, and clears base velocities."""
|
||||
if pos is None:
|
||||
pos = [0.0, 0.0, 0.20]
|
||||
if orn is None:
|
||||
orn = [0.0, 0.0, 0.0, 1.0]
|
||||
|
||||
def close(self):
|
||||
"""Cleanup wrapper for Robot backend compatibility."""
|
||||
lin_v = linear_velocity if linear_velocity is not None else [0.0, 0.0, 0.0]
|
||||
ang_v = angular_velocity if angular_velocity is not None else [0.0, 0.0, 0.0]
|
||||
|
||||
p.resetBasePositionAndOrientation(self.robot_id, pos, orn, physicsClientId=self.physics_client)
|
||||
p.resetBaseVelocity(self.robot_id, linearVelocity=lin_v, angularVelocity=ang_v, physicsClientId=self.physics_client)
|
||||
|
||||
# --- TELEMETRY (GETTERS) ---
|
||||
|
||||
def get_robot_pose(self) -> Tuple[List[float], List[float]]:
|
||||
pos, orn = p.getBasePositionAndOrientation(self.robot_id, physicsClientId=self.physics_client)
|
||||
return list(pos), list(orn)
|
||||
|
||||
def get_robot_rpy(self) -> Tuple[float, float, float]:
|
||||
_, orn = p.getBasePositionAndOrientation(self.robot_id, physicsClientId=self.physics_client)
|
||||
roll, pitch, yaw = p.getEulerFromQuaternion(orn)
|
||||
return float(roll), float(pitch), float(yaw)
|
||||
|
||||
def get_robot_pose_and_rpy(self) -> Tuple[List[float], Tuple[float, float, float]]:
|
||||
pos, orn = p.getBasePositionAndOrientation(self.robot_id, physicsClientId=self.physics_client)
|
||||
roll, pitch, yaw = p.getEulerFromQuaternion(orn)
|
||||
return list(pos), (float(roll), float(pitch), float(yaw))
|
||||
|
||||
def get_robot_velocity(self) -> Tuple[List[float], List[float]]:
|
||||
lin_v, ang_v = p.getBaseVelocity(self.robot_id, physicsClientId=self.physics_client)
|
||||
return list(lin_v), list(ang_v)
|
||||
|
||||
def get_robot_joint_angles(self) -> np.ndarray:
|
||||
joint_states = p.getJointStates(self.robot_id, self.revolute_joints, physicsClientId=self.physics_client)
|
||||
return np.array([state[0] for state in joint_states], dtype=np.float32)
|
||||
|
||||
# --- SIMULATION LIFECYCLE CONTROLS ---
|
||||
|
||||
def step(self) -> None:
|
||||
"""Advances physics simulation by 1 time step."""
|
||||
p.stepSimulation(physicsClientId=self.physics_client)
|
||||
|
||||
def set_robot_color(self, rgba: List[float]) -> None:
|
||||
num_joints = p.getNumJoints(self.robot_id, physicsClientId=self.physics_client)
|
||||
p.changeVisualShape(self.robot_id, -1, rgbaColor=rgba, physicsClientId=self.physics_client)
|
||||
for j in range(num_joints):
|
||||
p.changeVisualShape(self.robot_id, j, rgbaColor=rgba, physicsClientId=self.physics_client)
|
||||
|
||||
def settle_and_measure_height(self, steps: int = 200, fallback_height: float = 0.122) -> float:
|
||||
for _ in range(steps):
|
||||
self.step()
|
||||
pos, _ = self.get_robot_pose()
|
||||
return pos[2] if pos[2] > 0.0 else fallback_height
|
||||
|
||||
def apply_external_force(
|
||||
self, force: Union[List[float], np.ndarray], link_index: int = -1, position: Tuple[float, float, float] = (0.0, 0.0, 0.0)
|
||||
) -> None:
|
||||
p.applyExternalForce(
|
||||
objectUniqueId=self.robot_id,
|
||||
linkIndex=link_index,
|
||||
forceObj=list(force),
|
||||
posObj=list(position),
|
||||
flags=p.WORLD_FRAME,
|
||||
physicsClientId=self.physics_client,
|
||||
)
|
||||
|
||||
def disconnect(self) -> None:
|
||||
if self.physics_client is not None and p.isConnected(self.physics_client):
|
||||
p.disconnect(self.physics_client)
|
||||
self.physics_client = None
|
||||
|
||||
def close(self) -> None:
|
||||
self.disconnect()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from Robot import Robot, PyBulletBackend
|
||||
|
||||
# 1. Initialize PyBullet simulation environment
|
||||
sim_instance = Simulation()
|
||||
backend = PyBulletBackend(sim_instance)
|
||||
sim_instance = Simulation(use_gui=True)
|
||||
sim_instance.load_scene()
|
||||
|
||||
# 2. Instantiate Robot with simulation backend
|
||||
robot = Robot(backend_type=backend)
|
||||
backend = PyBulletBackend(sim_instance)
|
||||
robot = Robot(backend_type=backend, mode="kinematics")
|
||||
robot.reset_to_init()
|
||||
|
||||
# 3. Command forward movement [vx, vy, omega]
|
||||
robot.vector_dirmov = [1.0, 0.0, 0.0]
|
||||
# Forward velocity command
|
||||
robot.vector_dirmov = [0.3, 0.0, 0.0]
|
||||
|
||||
# 4. Main test execution loop
|
||||
try:
|
||||
while True:
|
||||
# Executes state machine logic
|
||||
robot.tick()
|
||||
time.sleep(1.0 / 60.0)
|
||||
except KeyboardInterrupt:
|
||||
|
||||
Reference in New Issue
Block a user