394 lines
16 KiB
Python
394 lines
16 KiB
Python
"""
|
|
ml/env.py - Gymnasium Environment for JackBot Hexapod RL Training
|
|
"""
|
|
import time
|
|
import math
|
|
from enum import IntEnum
|
|
from typing import Optional, Tuple, Dict, Any, List
|
|
from collections import defaultdict
|
|
|
|
import gymnasium as gym
|
|
from gymnasium import spaces
|
|
import numpy as np
|
|
|
|
from config import cfg
|
|
from simulation import Simulation
|
|
from Robot import Robot, PyBulletBackend
|
|
from ml.MetricsOverlay import MetricsHUD
|
|
|
|
COLOR_FAILED = [0.3, 0.3, 0.3, 0.6]
|
|
|
|
|
|
class CurriculumPhase(IntEnum):
|
|
STAND_ONLY = 0
|
|
FORWARD = 1
|
|
TURN_AND_DIRECTION = 2
|
|
OMNI_DIRECTION = 3
|
|
FULL_COMMAND = 4
|
|
|
|
|
|
class JackBotEnv(gym.Env):
|
|
"""Gymnasium environment wrapping JackBot hexapod simulation."""
|
|
|
|
def __init__(
|
|
self,
|
|
use_gui: bool = True,
|
|
random_command: bool = True,
|
|
max_episode_steps: int = 3000,
|
|
urdf_path: str = cfg.urdf_path,
|
|
robot_mode: str = "direct", # "direct", "residual", or "kinematics"
|
|
):
|
|
super().__init__()
|
|
self.robot_mode = robot_mode
|
|
self.use_gui = use_gui
|
|
self.random_command = random_command
|
|
self.max_episode_steps = max_episode_steps
|
|
self.urdf_path = urdf_path
|
|
self.max_robot_speed = 0.6
|
|
|
|
self.step_count = 0
|
|
self.total_steps = 0
|
|
self.cumulative_reward = 0.0
|
|
self.robot_reward = 0.0
|
|
self.is_failed = False
|
|
|
|
self.episode_count = 0
|
|
self.episode_height_sum = 0.0
|
|
self.episode_roll_sum = 0.0
|
|
self.episode_pitch_sum = 0.0
|
|
self._curriculum_advanced = False
|
|
self._first_reset = True
|
|
|
|
self.last_reward_components: Dict[str, float] = {}
|
|
self.episode_reward_components_sum: Dict[str, float] = defaultdict(float)
|
|
|
|
self.control_freq = 60
|
|
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
|
|
|
|
# --- 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()
|
|
|
|
# --- 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 + 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(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.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)
|
|
|
|
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},
|
|
}
|
|
|
|
self.hud = MetricsHUD(physics_client_id=self.sim.physics_client)
|
|
self.last_time = time.time()
|
|
|
|
def sample_command(self) -> np.ndarray:
|
|
"""Samples a command vector [vx, vy, omega] based on active curriculum phase."""
|
|
phase = self.curriculum_phase
|
|
stand_probs = {
|
|
CurriculumPhase.STAND_ONLY: 1.0,
|
|
CurriculumPhase.FORWARD: 0.25,
|
|
CurriculumPhase.TURN_AND_DIRECTION: 0.20,
|
|
CurriculumPhase.OMNI_DIRECTION: 0.15,
|
|
CurriculumPhase.FULL_COMMAND: 0.15,
|
|
}
|
|
|
|
if np.random.random() < stand_probs.get(phase, 0.15):
|
|
return np.zeros(3, dtype=np.float32)
|
|
|
|
if phase == CurriculumPhase.FORWARD:
|
|
vx, vy, omega = np.random.uniform(0.15, 0.50), 0.0, 0.0
|
|
elif phase == CurriculumPhase.TURN_AND_DIRECTION:
|
|
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, 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, 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, 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_reward = 0.0
|
|
self.is_failed = False
|
|
|
|
self.episode_height_sum = 0.0
|
|
self.episode_roll_sum = 0.0
|
|
self.episode_pitch_sum = 0.0
|
|
self.last_reward_components = {}
|
|
self.episode_reward_components_sum = defaultdict(float)
|
|
|
|
spawn_pos = [0.0, 0.0, 0.20]
|
|
spawn_orn = [0.0, 0.0, 0.0, 1.0]
|
|
|
|
# 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()
|
|
|
|
if self.use_gui:
|
|
self.sim.set_robot_color([1.0, 1.0, 1.0, 1.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
|
|
|
|
if self._first_reset:
|
|
self.curriculum_phase = CurriculumPhase.STAND_ONLY
|
|
self._first_reset = False
|
|
|
|
self.command = np.zeros(3, dtype=np.float32)
|
|
self.next_cmd_resample_step = self.initial_stand_steps
|
|
|
|
pos, _ = self.sim.get_robot_pose()
|
|
self.start_position = list(pos)
|
|
|
|
# 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()
|
|
self._update_hud()
|
|
|
|
return self._get_obs(), {}
|
|
|
|
def _get_obs(self) -> np.ndarray:
|
|
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
|
|
self.total_steps += 1
|
|
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
|
|
|
|
# Extract [vx, vy, omega]
|
|
cmd_vx, cmd_vy, cmd_omega = self.command
|
|
|
|
# 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.apply_external_force(force=[random_force[0], random_force[1], 0.0])
|
|
|
|
self._update_robot_failure()
|
|
self._update_distance_metrics()
|
|
self._update_curriculum()
|
|
|
|
obs = self._get_obs()
|
|
reward = self._compute_reward(action, previous_action)
|
|
|
|
self.cumulative_reward += reward
|
|
self.robot_reward += reward
|
|
|
|
terminated = self.is_failed
|
|
truncated = self.step_count >= self.max_episode_steps
|
|
info = {"reward_components": self.last_reward_components.copy()}
|
|
|
|
if self.step_count % 120 == 0 and self.use_gui:
|
|
self._update_hud()
|
|
|
|
return obs, reward, terminated, truncated, info
|
|
|
|
def _update_distance_metrics(self):
|
|
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)))
|
|
self.max_distance_from_start = max(self.max_distance_from_start, dist)
|
|
self.max_survival_steps = max(self.max_survival_steps, self.step_count)
|
|
|
|
def _phase_progress_ready(self, next_phase: CurriculumPhase) -> bool:
|
|
if next_phase not in self.curriculum_stage_requirements:
|
|
return False
|
|
|
|
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)
|
|
stability_ok = (avg_roll <= max_allowed_angle) and (avg_pitch <= max_allowed_angle)
|
|
|
|
avg_height = self.episode_height_sum / max(1, self.step_count)
|
|
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.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
|
|
if "min_forward_distance" in req:
|
|
distance_ok = dx >= req["min_forward_distance"]
|
|
elif "min_distance" in req:
|
|
distance_ok = dist_2d >= req["min_distance"]
|
|
|
|
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):
|
|
self._curriculum_advanced = False
|
|
forced_forward = (self.curriculum_phase < CurriculumPhase.FORWARD) and (self.step_count >= 2500)
|
|
|
|
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
|
|
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
|
|
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
|
|
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
|
|
|
|
def _compute_reward(self, action: np.ndarray, previous_action: np.ndarray) -> float:
|
|
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_norm = math.hypot(cmd_vx, cmd_vy)
|
|
|
|
raw_speed = math.hypot(linear_vel[0], linear_vel[1])
|
|
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])
|
|
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)))
|
|
|
|
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:
|
|
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:
|
|
is_moving = (filtered_speed > 0.0) or (abs(filtered_yaw_rate) > 0.0)
|
|
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:
|
|
total_reward = 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)
|
|
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, 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
|
|
)
|
|
|
|
final_reward = float(total_reward / 10.0)
|
|
|
|
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,
|
|
}
|
|
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.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))
|
|
)
|
|
|
|
def _update_robot_failure(self):
|
|
if self.is_failed or self.step_count < 15:
|
|
return
|
|
|
|
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
|
|
|
|
if is_tilted or is_collapsed:
|
|
self.is_failed = True
|
|
if self.use_gui:
|
|
self.sim.set_robot_color(COLOR_FAILED)
|
|
|
|
def close(self):
|
|
self.sim.disconnect() |