287 lines
11 KiB
Python
287 lines
11 KiB
Python
"""
|
|
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 pybullet as p
|
|
|
|
from config import cfg
|
|
from Robot import Robot, PyBulletBackend
|
|
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):
|
|
"""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()
|
|
|
|
# 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
|
|
)
|
|
|
|
# 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
|
|
)
|
|
for pb_id in self.pb_robots
|
|
]
|
|
|
|
# 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)
|
|
|
|
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.commands = np.zeros((self.num_robots, 4), dtype=np.float32)
|
|
self.last_action = np.zeros(action_dim, dtype=np.float32)
|
|
|
|
# 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()
|
|
|
|
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 _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)
|
|
vy = np.random.uniform(-0.5, 0.5)
|
|
vz = 0.0
|
|
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)]
|
|
|
|
for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
|
|
# Get default spawn position
|
|
spawn_pos = self._robot_base_position(idx, self.num_robots, self.robot_spacing)
|
|
spawn_orn = [0, 0, 0, 1]
|
|
|
|
# Teleport base back to start
|
|
p.resetBasePositionAndOrientation(
|
|
pb_id, spawn_pos, spawn_orn, physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
p.resetBaseVelocity(
|
|
pb_id, linearVelocity=[0, 0, 0], angularVelocity=[0, 0, 0],
|
|
physicsClientId=self.sim_manager.physics_client
|
|
)
|
|
|
|
# Reset joint angles directly without reloading URDF
|
|
robot_obj.reset_to_init()
|
|
|
|
# Restore original default visual color (clears failure dark gray)
|
|
if self.use_gui:
|
|
self._set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0])
|
|
|
|
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 _ in range(15):
|
|
self.sim_manager.step()
|
|
|
|
if self.use_gui:
|
|
self.hud.reset()
|
|
self.leader_crown.reset()
|
|
self._update_hud()
|
|
|
|
return self._get_obs(), {}
|
|
|
|
def _get_obs(self) -> np.ndarray:
|
|
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)
|
|
|
|
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)
|
|
|
|
for robot, act in zip(self.robots, action_per_robot):
|
|
robot.apply_rl_action(act)
|
|
|
|
self.sim_manager.step()
|
|
|
|
# Update failure status & gray coloring
|
|
self._update_robot_failures()
|
|
|
|
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
|
|
|
|
terminated = self._is_done()
|
|
truncated = self.step_count >= self.max_episode_steps
|
|
|
|
self._update_hud()
|
|
self._update_leader_visuals()
|
|
return obs, reward, terminated, truncated, {}
|
|
|
|
def _compute_reward(self) -> Tuple[float, list[float]]:
|
|
rewards = []
|
|
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[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)[idx]))) * 0.01
|
|
|
|
r_step = 0.1 + forward_reward + rotation_reward - 0.2 * stability_penalty - action_penalty
|
|
rewards.append(r_step)
|
|
|
|
return float(np.sum(rewards)), rewards
|
|
|
|
def _is_done(self) -> bool:
|
|
"""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
|
|
|
|
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() |