Training uses new Robot.py
new Metrics and SimManager for live training viewing Robot.py usage added to machinelearning
This commit is contained in:
@@ -0,0 +1,137 @@
|
||||
"""
|
||||
ml/MetricsOverlay.py - Camera-Facing (Billboard) 3D Floating Text Overlay
|
||||
"""
|
||||
from typing import List, Tuple, Optional
|
||||
import numpy as np
|
||||
import pybullet as p
|
||||
|
||||
|
||||
class MetricsHUD:
|
||||
"""Renders real-time telemetry as black text floating in 3D, always facing the active camera."""
|
||||
|
||||
def __init__(self, physics_client_id: int = 0):
|
||||
self.client_id = physics_client_id
|
||||
self._text_id: Optional[int] = None
|
||||
|
||||
def _get_camera_facing_orientation(self) -> List[float]:
|
||||
"""Calculates a quaternion that points the text towards the current GUI camera."""
|
||||
try:
|
||||
cam_info = p.getDebugVisualizerCamera(physicsClientId=self.client_id)
|
||||
# cam_info index 8: yaw, index 9: pitch
|
||||
yaw = cam_info[8]
|
||||
pitch = cam_info[9]
|
||||
|
||||
# Orient the text normal toward the camera view direction
|
||||
# PyBullet text default faces local +Z/-Y depending on roll,
|
||||
# converting visualizer yaw/pitch to Euler angles (roll, pitch, yaw in radians)
|
||||
roll_rad = 0.0
|
||||
pitch_rad = np.radians(pitch + 90.0)
|
||||
yaw_rad = np.radians(yaw)
|
||||
|
||||
text_orientation = p.getQuaternionFromEuler(
|
||||
[pitch_rad, roll_rad, yaw_rad],
|
||||
physicsClientId=self.client_id
|
||||
)
|
||||
return text_orientation
|
||||
except Exception:
|
||||
# Fallback default orientation if camera info call fails
|
||||
return [0.0, 0.0, 0.0, 1.0]
|
||||
|
||||
def update(
|
||||
self,
|
||||
episode: int,
|
||||
step: int,
|
||||
robot_rewards: List[float],
|
||||
cmd_vel: np.ndarray,
|
||||
fps: float = 0.0,
|
||||
avg_height: float = 0.0,
|
||||
roll_pitch: Tuple[float, float] = (0.0, 0.0)
|
||||
) -> None:
|
||||
"""Updates floating black text block in 3D space with billboarding."""
|
||||
sorted_rewards = sorted(robot_rewards, reverse=True)
|
||||
top1 = f"{sorted_rewards[0]:+.2f}" if len(sorted_rewards) > 0 else "0.00"
|
||||
top2 = f"{sorted_rewards[1]:+.2f}" if len(sorted_rewards) > 1 else "0.00"
|
||||
top3 = f"{sorted_rewards[2]:+.2f}" if len(sorted_rewards) > 2 else "0.00"
|
||||
|
||||
vx = cmd_vel[0] if len(cmd_vel) > 0 else 0.0
|
||||
vy = cmd_vel[1] if len(cmd_vel) > 1 else 0.0
|
||||
omega = cmd_vel[3] if len(cmd_vel) > 3 else 0.0
|
||||
|
||||
hud_text = (
|
||||
f"=== JACKBOT METRICS ===\n"
|
||||
f"Episode: {episode}\n"
|
||||
f"Global Step: {step}\n"
|
||||
f"FPS: {fps:.1f}\n"
|
||||
f"----------------------\n"
|
||||
f"Top Rewards: [{top1}, {top2}, {top3}]\n"
|
||||
f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]\n"
|
||||
f"Height: {avg_height:.3f} m\n"
|
||||
f"Roll/Pitch: {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°"
|
||||
)
|
||||
|
||||
# Position above origin in simulation world
|
||||
text_position = [-0.8, -0.8, 1.2]
|
||||
text_color = [0, 0, 0] # Pure black
|
||||
|
||||
# Calculate dynamic orientation to align text flat against camera plane
|
||||
text_orientation = self._get_camera_facing_orientation()
|
||||
|
||||
if self._text_id is None:
|
||||
self._text_id = p.addUserDebugText(
|
||||
text=hud_text,
|
||||
textPosition=text_position,
|
||||
textColorRGB=text_color,
|
||||
textSize=0.1,
|
||||
textOrientation=text_orientation,
|
||||
physicsClientId=self.client_id
|
||||
)
|
||||
else:
|
||||
self._text_id = p.addUserDebugText(
|
||||
text=hud_text,
|
||||
textPosition=text_position,
|
||||
textColorRGB=text_color,
|
||||
textSize=0.1,
|
||||
textOrientation=text_orientation,
|
||||
replaceItemUniqueId=self._text_id,
|
||||
physicsClientId=self.client_id
|
||||
)
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Clears text reference on environment reset."""
|
||||
self._text_id = None
|
||||
|
||||
|
||||
class LeaderCrown:
|
||||
"""Renders a floating crown or star emoji above the leading robot in PyBullet."""
|
||||
|
||||
def __init__(self, physics_client_id: int = 0):
|
||||
self.client_id = physics_client_id
|
||||
self._text_id = None
|
||||
|
||||
def update(self, leader_pos: list[float]):
|
||||
"""Positions a floating crown ~0.35m directly above the lead robot's base."""
|
||||
crown_pos = [leader_pos[0], leader_pos[1], leader_pos[2] + 0.35]
|
||||
|
||||
# You can use "👑 CROWN", "⭐ LEADER", or "★ TOP1"
|
||||
crown_text = "👑"
|
||||
|
||||
if self._text_id is None:
|
||||
self._text_id = p.addUserDebugText(
|
||||
text=crown_text,
|
||||
textPosition=crown_pos,
|
||||
textColorRGB=[1.0, 0.84, 0.0],
|
||||
textSize=2.0,
|
||||
physicsClientId=self.client_id
|
||||
)
|
||||
else:
|
||||
self._text_id = p.addUserDebugText(
|
||||
text=crown_text,
|
||||
textPosition=crown_pos,
|
||||
textColorRGB=[1.0, 0.84, 0.0],
|
||||
textSize=2.0,
|
||||
replaceItemUniqueId=self._text_id,
|
||||
physicsClientId=self.client_id
|
||||
)
|
||||
|
||||
def reset(self):
|
||||
self._text_id = None
|
||||
@@ -0,0 +1,88 @@
|
||||
"""
|
||||
ml/SimManager.py - PyBullet Simulation & Multi-Body Manager
|
||||
"""
|
||||
from typing import Dict, List, Tuple
|
||||
import pybullet as p
|
||||
import pybullet_data
|
||||
import DataTypes as dt
|
||||
|
||||
|
||||
class SimManager:
|
||||
"""Manages PyBullet simulation lifecycle and multi-robot physics."""
|
||||
|
||||
def __init__(self, use_gui: bool = True):
|
||||
self.use_gui = use_gui
|
||||
self.physics_client = None
|
||||
self.robot_joints: Dict[int, List[int]] = {}
|
||||
|
||||
def connect(self):
|
||||
"""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:
|
||||
# Disable PyBullet side panel and preview windows
|
||||
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, num_robots: int, robot_spacing: float, base_pos_fn
|
||||
) -> Tuple[int, List[int], List[List[int]]]:
|
||||
"""Loads plane and hexapod bodies into the simulation scene."""
|
||||
plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client)
|
||||
robots = []
|
||||
robot_joint_indices = []
|
||||
self.robot_joints.clear()
|
||||
|
||||
for r_id in range(num_robots):
|
||||
base_pos = base_pos_fn(r_id, num_robots, robot_spacing)
|
||||
robot = p.loadURDF(
|
||||
urdf_path,
|
||||
basePosition=base_pos,
|
||||
useFixedBase=False,
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
robots.append(robot)
|
||||
|
||||
joint_indices = [
|
||||
i for i in range(p.getNumJoints(robot, physicsClientId=self.physics_client))
|
||||
if p.getJointInfo(robot, i, physicsClientId=self.physics_client)[2] == p.JOINT_REVOLUTE
|
||||
]
|
||||
robot_joint_indices.append(joint_indices)
|
||||
self.robot_joints[robot] = joint_indices
|
||||
|
||||
return plane_id, robots, robot_joint_indices
|
||||
|
||||
def updatePosForBody(self, body_id: int, current_rad: dt.RadArray):
|
||||
"""Sets joint motor position targets on individual robot bodies."""
|
||||
if body_id not in self.robot_joints:
|
||||
return
|
||||
|
||||
joint_indices = self.robot_joints[body_id]
|
||||
radflat = current_rad.data.flatten()
|
||||
|
||||
for joint_index, target_angle in zip(joint_indices, radflat):
|
||||
p.setJointMotorControl2(
|
||||
bodyIndex=body_id,
|
||||
jointIndex=joint_index,
|
||||
controlMode=p.POSITION_CONTROL,
|
||||
targetPosition=float(target_angle),
|
||||
force=250,
|
||||
physicsClientId=self.physics_client
|
||||
)
|
||||
|
||||
def step(self):
|
||||
p.stepSimulation(physicsClientId=self.physics_client)
|
||||
|
||||
def disconnect(self):
|
||||
if self.physics_client is not None and p.isConnected(self.physics_client):
|
||||
p.disconnect(self.physics_client)
|
||||
self.physics_client = None
|
||||
@@ -1,130 +1,104 @@
|
||||
# ml/env.py
|
||||
"""
|
||||
ml/env.py - Gymnasium Environment for JackBot Hexapod RL Training
|
||||
"""
|
||||
import time
|
||||
import math
|
||||
from typing import Optional, Tuple, Dict, Any, List
|
||||
|
||||
import gymnasium as gym
|
||||
from gymnasium import spaces
|
||||
import numpy as np
|
||||
import math
|
||||
import pybullet as p
|
||||
|
||||
from config import cfg
|
||||
from Robot import Robot, PyBulletBackend
|
||||
from ml.sim_manager import SimManager
|
||||
from ml.SimManager import SimManager
|
||||
from ml.MetricsOverlay import MetricsHUD, LeaderCrown
|
||||
|
||||
# Color Palette RGBA for Terminated/Failed Robots
|
||||
COLOR_FAILED = [0.3, 0.3, 0.3, 0.6] # Collapsed / Tilted Robot (Dark Semi-Transparent Gray)
|
||||
|
||||
|
||||
class JackBotEnv(gym.Env):
|
||||
def __init__(self, use_gui: bool = False, num_robots: int = 1):
|
||||
self.sim_manager = SimManager(use_gui=use_gui)
|
||||
"""Gymnasium environment wrapping JackBot hexapods with floating text & crown overlay."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
use_gui: bool = True,
|
||||
random_command: bool = True,
|
||||
num_robots: int = 1,
|
||||
robot_spacing: float = 0.5,
|
||||
start_pose: str = "init_deg",
|
||||
max_episode_steps: int = 3000,
|
||||
urdf_path: str = cfg.urdf_path,
|
||||
termination_threshold: float = 0.001, # Termination ratio threshold
|
||||
):
|
||||
super().__init__()
|
||||
self.use_gui = use_gui
|
||||
self.random_command = random_command
|
||||
self.num_robots = num_robots
|
||||
self.robot_spacing = robot_spacing
|
||||
self.start_pose = start_pose
|
||||
self.max_episode_steps = max_episode_steps
|
||||
self.urdf_path = urdf_path
|
||||
self.termination_threshold = termination_threshold
|
||||
|
||||
self.episode_count = 0
|
||||
self.step_count = 0
|
||||
self.total_steps = 0
|
||||
self.cumulative_reward = 0.0
|
||||
self.robot_rewards = [0.0 for _ in range(self.num_robots)]
|
||||
self.failed_robots_mask = [False for _ in range(self.num_robots)]
|
||||
self._first_reset = True
|
||||
|
||||
# Initialize Simulation Manager
|
||||
self.sim_manager = SimManager(use_gui=self.use_gui)
|
||||
self.sim_manager.connect()
|
||||
|
||||
# Load simulation bodies
|
||||
self.plane, self.pb_robots, self.joint_indices = self.sim_manager.load_scene(
|
||||
cfg.urdf_path, num_robots=num_robots
|
||||
)
|
||||
|
||||
# Instantiate dedicated Robot Python object for EACH spawned robot
|
||||
self.robots = [
|
||||
Robot(backend=PyBulletBackend(self.sim_manager, body_id=pb_id))
|
||||
for pb_id in self.pb_robots
|
||||
]
|
||||
|
||||
# Action & Observation Spaces
|
||||
action_dim = num_robots * 18
|
||||
obs_dim = num_robots * (18 + 4)
|
||||
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)
|
||||
def _robot_base_position(self, robot_id: int, num_robots: int = 1, spacing: float = 0.5) -> list[float]:
|
||||
"""Calculates grid coordinates for spawning multiple robots in PyBullet."""
|
||||
cols = int(math.sqrt(num_robots - 1)) + 1
|
||||
row = robot_id // cols
|
||||
col = robot_id % cols
|
||||
x = (col - (cols - 1) / 2.0) * spacing
|
||||
y = (row - (cols - 1) / 2.0) * spacing
|
||||
return [x, y, 0.1]
|
||||
|
||||
def _connect_sim(self):
|
||||
self.sim_manager.connect()
|
||||
self.plane, self.robots, self.robot_joint_indices = self.sim_manager.load_scene(
|
||||
# Connect physics world
|
||||
self.plane, self.pb_robots, self.robot_joint_indices = self.sim_manager.load_scene(
|
||||
self.urdf_path, self.num_robots, self.robot_spacing, self._robot_base_position
|
||||
)
|
||||
|
||||
joint_count = len(self.robot_joint_indices[0])
|
||||
if any(len(indices) != joint_count for indices in self.robot_joint_indices):
|
||||
raise ValueError("All robots must have the same number of revolute joints.")
|
||||
|
||||
lower_limits, upper_limits = [], []
|
||||
for joint_index in self.robot_joint_indices[0]:
|
||||
info = p.getJointInfo(self.robots[0], joint_index)
|
||||
lower, upper = info[8], info[9]
|
||||
if lower >= upper:
|
||||
lower, upper = -math.pi, math.pi
|
||||
lower_limits.append(lower)
|
||||
upper_limits.append(upper)
|
||||
|
||||
limits = np.array([lower_limits, upper_limits], dtype=np.float32)
|
||||
self.joint_lower = np.tile(limits[0], (self.num_robots, 1))
|
||||
self.joint_upper = np.tile(limits[1], (self.num_robots, 1))
|
||||
|
||||
pose = ri.init_deg if self.start_pose == "init_deg" else ri.init90_deg
|
||||
initial_pose = pose.to_rad().data.flatten()
|
||||
self.initial_angles = np.tile(initial_pose, (self.num_robots, 1)).astype(np.float32)
|
||||
self.initial_angles = np.clip(self.initial_angles, self.joint_lower, self.joint_upper)
|
||||
|
||||
self.last_action = np.zeros(self.num_robots * joint_count, dtype=np.float32)
|
||||
self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
|
||||
|
||||
def reset(self, command: np.ndarray | None = None, seed: int | None = None, options: dict | None = None):
|
||||
self.episode_count += 1
|
||||
self.step_count = 0
|
||||
self.cumulative_reward = 0.0
|
||||
|
||||
if not self._first_reset:
|
||||
# Reset simulation bodies without reconnecting PyBullet
|
||||
p.resetSimulation()
|
||||
p.setGravity(0, 0, -9.81)
|
||||
self.sim_manager.reset_hud()
|
||||
self.plane, self.robots, self.robot_joint_indices = self.sim_manager.load_scene(
|
||||
self.urdf_path, self.num_robots, self.robot_spacing, self._robot_base_position
|
||||
# Instantiate Robot Python wrappers per PyBullet body ID
|
||||
self.robots = [
|
||||
Robot(
|
||||
backend_type=PyBulletBackend(self.sim_manager, body_id=pb_id),
|
||||
start_pose=self.start_pose,
|
||||
urdf_path=self.urdf_path
|
||||
)
|
||||
else:
|
||||
self._first_reset = False
|
||||
for pb_id in self.pb_robots
|
||||
]
|
||||
|
||||
self.last_action = np.zeros(self.num_robots * len(self.robot_joint_indices[0]), dtype=np.float32)
|
||||
# Action (18 joint deltas per robot) & Observation (18 angles + 4 command dims per robot)
|
||||
action_dim = self.num_robots * 18
|
||||
obs_dim = self.num_robots * (18 + 4)
|
||||
|
||||
if seed is not None:
|
||||
self.seed(seed)
|
||||
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)
|
||||
|
||||
if command is not None:
|
||||
self.commands = np.array(command, dtype=np.float32).reshape(self.num_robots, 4)
|
||||
elif self.random_command:
|
||||
self.commands = np.stack([self.sample_command() for _ in range(self.num_robots)])
|
||||
else:
|
||||
self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
|
||||
self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
|
||||
self.last_action = np.zeros(action_dim, dtype=np.float32)
|
||||
|
||||
self._reset_pose()
|
||||
self._update_gui_hud(reward=0.0)
|
||||
# Floating HUD & Leader Crown Visualizers
|
||||
self.hud = MetricsHUD(physics_client_id=self.sim_manager.physics_client)
|
||||
self.leader_crown = LeaderCrown(physics_client_id=self.sim_manager.physics_client)
|
||||
self.last_time = time.time()
|
||||
|
||||
return self._get_obs(), {}
|
||||
def _set_robot_color(self, pb_id: int, rgba: List[float]):
|
||||
"""Helper to change the visual color of a robot body and all its links."""
|
||||
num_joints = p.getNumJoints(pb_id, physicsClientId=self.sim_manager.physics_client)
|
||||
p.changeVisualShape(pb_id, -1, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client)
|
||||
for j in range(num_joints):
|
||||
p.changeVisualShape(pb_id, j, rgbaColor=rgba, physicsClientId=self.sim_manager.physics_client)
|
||||
|
||||
def seed(self, seed: int | None = None):
|
||||
if seed is None:
|
||||
seed = int(np.random.randint(0, 2**31 - 1))
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
self._seed = seed
|
||||
return [seed]
|
||||
|
||||
def _reset_pose(self):
|
||||
for robot, joint_indices, init_angles in zip(self.robots, self.robot_joint_indices, self.initial_angles):
|
||||
for joint_index, target_angle in zip(joint_indices, init_angles):
|
||||
p.resetJointState(robot, joint_index, target_angle)
|
||||
p.setJointMotorControl2(
|
||||
bodyIndex=robot,
|
||||
jointIndex=joint_index,
|
||||
controlMode=p.POSITION_CONTROL,
|
||||
targetPosition=target_angle,
|
||||
force=250,
|
||||
)
|
||||
|
||||
for _ in range(400):
|
||||
p.stepSimulation()
|
||||
|
||||
self.joint_angles = self.initial_angles.copy()
|
||||
def _robot_base_position(self, robot_id: int, num_robots: int = 1, spacing: float = 0.5) -> list[float]:
|
||||
cols = int(math.sqrt(num_robots - 1)) + 1
|
||||
row = robot_id // cols
|
||||
col = robot_id % cols
|
||||
x = (col - (cols - 1) / 2.0) * spacing
|
||||
y = (row - (cols - 1) / 2.0) * spacing
|
||||
return [x, y, 0.2]
|
||||
|
||||
def sample_command(self) -> np.ndarray:
|
||||
vx = np.random.uniform(-1.0, 1.0)
|
||||
@@ -133,72 +107,173 @@ class JackBotEnv(gym.Env):
|
||||
omega = np.random.uniform(-1.0, 1.0)
|
||||
return np.array([vx, vy, vz, omega], dtype=np.float32)
|
||||
|
||||
def reset(self, seed: Optional[int] = None, options: Optional[Dict[str, Any]] = None):
|
||||
super().reset(seed=seed)
|
||||
self.episode_count += 1
|
||||
self.step_count = 0
|
||||
self.cumulative_reward = 0.0
|
||||
self.robot_rewards = [0.0 for _ in range(self.num_robots)]
|
||||
self.failed_robots_mask = [False for _ in range(self.num_robots)]
|
||||
|
||||
if not self._first_reset:
|
||||
p.resetSimulation(physicsClientId=self.sim_manager.physics_client)
|
||||
p.setGravity(0, 0, -9.81, physicsClientId=self.sim_manager.physics_client)
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.sim_manager.physics_client)
|
||||
self.hud.reset()
|
||||
self.leader_crown.reset()
|
||||
self.plane, self.pb_robots, self.robot_joint_indices = self.sim_manager.load_scene(
|
||||
self.urdf_path, self.num_robots, self.robot_spacing, self._robot_base_position
|
||||
)
|
||||
for robot_obj, pb_id in zip(self.robots, self.pb_robots):
|
||||
robot_obj.backend = PyBulletBackend(self.sim_manager, body_id=pb_id)
|
||||
else:
|
||||
self._first_reset = False
|
||||
|
||||
self.last_action = np.zeros(self.action_space.shape[0], dtype=np.float32)
|
||||
|
||||
if self.random_command:
|
||||
self.commands = np.stack([self.sample_command() for _ in range(self.num_robots)])
|
||||
else:
|
||||
self.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
|
||||
|
||||
for robot in self.robots:
|
||||
robot.reset_to_init()
|
||||
|
||||
for _ in range(100):
|
||||
self.sim_manager.step()
|
||||
|
||||
self._update_hud()
|
||||
return self._get_obs(), {}
|
||||
|
||||
def _get_obs(self) -> np.ndarray:
|
||||
return np.concatenate([self.joint_angles.flatten(), self.commands.flatten()]).astype(np.float32)
|
||||
obs_list = []
|
||||
for idx, (pb_id, joint_indices) in enumerate(zip(self.pb_robots, self.robot_joint_indices)):
|
||||
joint_states = p.getJointStates(
|
||||
pb_id,
|
||||
joint_indices,
|
||||
physicsClientId=self.sim_manager.physics_client
|
||||
)
|
||||
joint_angles = np.array([state[0] for state in joint_states], dtype=np.float32)
|
||||
robot_obs = np.concatenate([joint_angles, self.commands[idx]])
|
||||
obs_list.append(robot_obs)
|
||||
|
||||
def step(self, action: np.ndarray):
|
||||
action_per_robot = action.reshape(len(self.robots), 18)
|
||||
return np.concatenate(obs_list).astype(np.float32)
|
||||
|
||||
def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
|
||||
self.step_count += 1
|
||||
self.total_steps += 1
|
||||
self.last_action = action.copy()
|
||||
|
||||
action_per_robot = action.reshape(self.num_robots, 18)
|
||||
|
||||
# Apply RL actions independently to each Robot object instance
|
||||
for robot, act in zip(self.robots, action_per_robot):
|
||||
robot.apply_rl_action(act)
|
||||
|
||||
# Step PyBullet physics engine once
|
||||
self.sim_manager.step()
|
||||
|
||||
# Gather observations across all robot objects
|
||||
obs = np.concatenate([
|
||||
robot.get_observation(command=np.zeros(4))
|
||||
for robot in self.robots
|
||||
])
|
||||
# Update failure status & gray coloring
|
||||
self._update_robot_failures()
|
||||
|
||||
reward = self._compute_reward()
|
||||
done = False
|
||||
return obs, reward, done, False, {}
|
||||
obs = self._get_obs()
|
||||
reward, per_robot_step_rewards = self._compute_reward()
|
||||
|
||||
self.cumulative_reward += reward
|
||||
for idx, r_step in enumerate(per_robot_step_rewards):
|
||||
self.robot_rewards[idx] += r_step
|
||||
|
||||
def _update_gui_hud(self, reward: float):
|
||||
"""Passes current state metrics to the SimManager HUD renderer."""
|
||||
linear_vel, _ = p.getBaseVelocity(self.robots[0])
|
||||
cmd = self.commands[0]
|
||||
terminated = self._is_done()
|
||||
truncated = self.step_count >= self.max_episode_steps
|
||||
|
||||
stats = {
|
||||
"Episode": self.episode_count,
|
||||
"Step": f"{self.step_count} / {self.max_episode_steps}",
|
||||
"Step Reward": reward,
|
||||
"Total Reward": self.cumulative_reward,
|
||||
"Target Cmd (vx,vy,w)": [cmd[0], cmd[1], cmd[3]],
|
||||
"Actual Vel (vx,vy)": [linear_vel[0], linear_vel[1]],
|
||||
}
|
||||
self.sim_manager.update_hud(stats)
|
||||
self._update_hud()
|
||||
self._update_leader_visuals()
|
||||
return obs, reward, terminated, truncated, {}
|
||||
|
||||
def _compute_reward(self) -> float:
|
||||
def _compute_reward(self) -> Tuple[float, list[float]]:
|
||||
rewards = []
|
||||
for robot_id, robot in enumerate(self.robots):
|
||||
linear_vel, angular_vel = p.getBaseVelocity(robot)
|
||||
_, orientation = p.getBasePositionAndOrientation(robot)
|
||||
for idx, pb_id in enumerate(self.pb_robots):
|
||||
linear_vel, angular_vel = p.getBaseVelocity(pb_id, physicsClientId=self.sim_manager.physics_client)
|
||||
_, orientation = p.getBasePositionAndOrientation(pb_id, physicsClientId=self.sim_manager.physics_client)
|
||||
roll, pitch, _ = p.getEulerFromQuaternion(orientation)
|
||||
command = self.commands[robot_id]
|
||||
command = self.commands[idx]
|
||||
|
||||
forward_reward = command[0] * linear_vel[0] + command[1] * linear_vel[1]
|
||||
rotation_reward = command[3] * angular_vel[2]
|
||||
stability_penalty = abs(roll) + abs(pitch)
|
||||
action_penalty = float(np.sum(np.square(self.last_action.reshape(self.num_robots, -1)[robot_id]))) * 0.01
|
||||
action_penalty = float(np.sum(np.square(self.last_action.reshape(self.num_robots, -1)[idx]))) * 0.01
|
||||
|
||||
rewards.append(0.1 + forward_reward + rotation_reward - 0.2 * stability_penalty - action_penalty)
|
||||
r_step = 0.1 + forward_reward + rotation_reward - 0.2 * stability_penalty - action_penalty
|
||||
rewards.append(r_step)
|
||||
|
||||
return float(np.sum(rewards))
|
||||
return float(np.sum(rewards)), rewards
|
||||
|
||||
def _is_done(self) -> bool:
|
||||
for robot in self.robots:
|
||||
_, orientation = p.getBasePositionAndOrientation(robot)
|
||||
roll, pitch, _ = p.getEulerFromQuaternion(orientation)
|
||||
if abs(roll) > 0.7 or abs(pitch) > 0.7:
|
||||
return True
|
||||
position, _ = p.getBasePositionAndOrientation(robot)
|
||||
if position[2] < 0.02:
|
||||
return True
|
||||
"""Returns True only when the percentage of failed robots exceeds the threshold."""
|
||||
failed_count = sum(self.failed_robots_mask)
|
||||
failure_ratio = failed_count / self.num_robots
|
||||
return failure_ratio >= self.termination_threshold
|
||||
|
||||
return self.step_count >= self.max_episode_steps
|
||||
def _update_hud(self):
|
||||
if not self.use_gui or not self.pb_robots:
|
||||
return
|
||||
|
||||
now = time.time()
|
||||
fps = 1.0 / max(now - self.last_time, 1e-5)
|
||||
self.last_time = now
|
||||
|
||||
heights = []
|
||||
rolls = []
|
||||
pitches = []
|
||||
for pb_id in self.pb_robots:
|
||||
pos, orient = p.getBasePositionAndOrientation(pb_id, physicsClientId=self.sim_manager.physics_client)
|
||||
roll, pitch, _ = p.getEulerFromQuaternion(orient)
|
||||
heights.append(pos[2])
|
||||
rolls.append(math.degrees(roll))
|
||||
pitches.append(math.degrees(pitch))
|
||||
|
||||
avg_height = float(np.mean(heights))
|
||||
avg_roll_pitch = (float(np.mean(rolls)), float(np.mean(pitches)))
|
||||
|
||||
self.hud.update(
|
||||
episode=self.episode_count,
|
||||
step=self.total_steps,
|
||||
robot_rewards=self.robot_rewards,
|
||||
cmd_vel=self.commands[0],
|
||||
fps=fps,
|
||||
avg_height=avg_height,
|
||||
roll_pitch=avg_roll_pitch
|
||||
)
|
||||
|
||||
def _update_leader_visuals(self):
|
||||
"""Positions floating crown above top robot without altering active materials."""
|
||||
if not self.use_gui or self.num_robots <= 1:
|
||||
return
|
||||
|
||||
best_idx = int(np.argmax(self.robot_rewards))
|
||||
leader_pb_id = self.pb_robots[best_idx]
|
||||
leader_pos, _ = p.getBasePositionAndOrientation(
|
||||
leader_pb_id, physicsClientId=self.sim_manager.physics_client
|
||||
)
|
||||
self.leader_crown.update(leader_pos)
|
||||
|
||||
def _update_robot_failures(self):
|
||||
"""Checks failure condition for each robot and turns failed ones gray."""
|
||||
for idx, pb_id in enumerate(self.pb_robots):
|
||||
if self.failed_robots_mask[idx]:
|
||||
continue # Already marked failed
|
||||
|
||||
position, orientation = p.getBasePositionAndOrientation(
|
||||
pb_id, physicsClientId=self.sim_manager.physics_client
|
||||
)
|
||||
roll, pitch, _ = p.getEulerFromQuaternion(orientation)
|
||||
|
||||
is_tilted = abs(roll) > 0.7 or abs(pitch) > 0.7
|
||||
is_collapsed = position[2] < 0.05
|
||||
|
||||
if is_tilted or is_collapsed:
|
||||
self.failed_robots_mask[idx] = True
|
||||
if self.use_gui:
|
||||
# Turn failed robot semi-transparent dark gray
|
||||
self._set_robot_color(pb_id, COLOR_FAILED)
|
||||
|
||||
def close(self):
|
||||
self.sim_manager.disconnect()
|
||||
@@ -1,105 +0,0 @@
|
||||
from typing import Dict, Any, List, Tuple
|
||||
import pybullet as p
|
||||
import pybullet_data
|
||||
import numpy as np
|
||||
|
||||
|
||||
class SimManager:
|
||||
"""
|
||||
Manages the PyBullet simulation lifecycle and live HUD overlays.
|
||||
"""
|
||||
|
||||
def __init__(self, use_gui: bool = True):
|
||||
self.use_gui = use_gui
|
||||
self.physics_client = None
|
||||
self.debug_text_ids: Dict[str, int] = {}
|
||||
|
||||
def connect(self):
|
||||
"""Connects to PyBullet and sets up the basic physics world."""
|
||||
if self.physics_client is not None:
|
||||
p.disconnect(self.physics_client)
|
||||
|
||||
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)
|
||||
|
||||
if self.use_gui:
|
||||
# Disable unnecessary side panels for a clean UI
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0)
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0)
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0)
|
||||
p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0)
|
||||
|
||||
def load_scene(
|
||||
self, urdf_path: str, num_robots: int, robot_spacing: float, base_pos_fn
|
||||
) -> Tuple[int, List[int], List[List[int]]]:
|
||||
"""Loads plane and robots into the active PyBullet simulation."""
|
||||
plane_id = p.loadURDF("plane.urdf")
|
||||
robots = []
|
||||
robot_joint_indices = []
|
||||
|
||||
for robot_id in range(num_robots):
|
||||
base_pos = base_pos_fn(robot_id, num_robots, robot_spacing)
|
||||
robot = p.loadURDF(urdf_path, basePosition=base_pos, useFixedBase=False)
|
||||
robots.append(robot)
|
||||
|
||||
joint_indices = [
|
||||
i
|
||||
for i in range(p.getNumJoints(robot))
|
||||
if p.getJointInfo(robot, i)[2] == p.JOINT_REVOLUTE
|
||||
]
|
||||
robot_joint_indices.append(joint_indices)
|
||||
|
||||
return plane_id, robots, robot_joint_indices
|
||||
|
||||
def disconnect(self):
|
||||
"""Safely disconnects from the simulation."""
|
||||
if self.physics_client is not None:
|
||||
p.disconnect(self.physics_client)
|
||||
self.physics_client = None
|
||||
self.debug_text_ids.clear()
|
||||
|
||||
def update_hud(self, stats: Dict[str, Any]):
|
||||
"""
|
||||
Renders live training metrics onto the PyBullet 3D viewport.
|
||||
Uses replaceItemUniqueId to prevent flicker.
|
||||
"""
|
||||
if not self.use_gui or self.physics_client is None:
|
||||
return
|
||||
|
||||
# Fixed position near the top-left of the origin in 3D world coordinates
|
||||
x_pos, y_pos, z_start = -1.2, -1.2, 1.6
|
||||
line_height = 0.10
|
||||
|
||||
for i, (label, value) in enumerate(stats.items()):
|
||||
if isinstance(value, float):
|
||||
display_text = f"{label}: {value:.3f}"
|
||||
elif isinstance(value, (list, np.ndarray)):
|
||||
formatted_vals = ", ".join(f"{v:.2f}" for v in np.atleast_1d(value))
|
||||
display_text = f"{label}: [{formatted_vals}]"
|
||||
else:
|
||||
display_text = f"{label}: {value}"
|
||||
|
||||
color = [0, 0, 0] # Black text for clear visibility against the light plane
|
||||
|
||||
if label in self.debug_text_ids:
|
||||
p.addUserDebugText(
|
||||
display_text,
|
||||
[x_pos, y_pos, z_start - i * line_height],
|
||||
textColorRGB=color,
|
||||
textSize=1.1,
|
||||
replaceItemUniqueId=self.debug_text_ids[label],
|
||||
)
|
||||
else:
|
||||
self.debug_text_ids[label] = p.addUserDebugText(
|
||||
display_text,
|
||||
[x_pos, y_pos, z_start - i * line_height],
|
||||
textColorRGB=color,
|
||||
textSize=1.1,
|
||||
)
|
||||
|
||||
def reset_hud(self):
|
||||
"""Clears debug text tracking."""
|
||||
self.debug_text_ids.clear()
|
||||
+1
-1
@@ -110,4 +110,4 @@ if __name__ == "__main__":
|
||||
num_robots=args.num_robots,
|
||||
robot_spacing=args.robot_spacing,
|
||||
start_pose=args.start_pose,
|
||||
)
|
||||
)
|
||||
Reference in New Issue
Block a user