reward ajustments
robot starts balancing without changing phase and farms alive bonus fix ->external force and bigger height penalty
This commit is contained in:
@@ -3,6 +3,7 @@ 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
|
||||
|
||||
import gymnasium as gym
|
||||
@@ -18,6 +19,14 @@ from ml.MetricsOverlay import MetricsHUD, LeaderCrown
|
||||
COLOR_FAILED = [0.3, 0.3, 0.3, 0.6] # Collapsed / Tilted Robot (Dark Semi-Transparent Gray)
|
||||
|
||||
|
||||
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 hexapods with floating text & crown overlay."""
|
||||
|
||||
@@ -44,6 +53,7 @@ class JackBotEnv(gym.Env):
|
||||
self.cumulative_reward = 0.0
|
||||
self.robot_rewards = [0.0]
|
||||
self.failed_robots_mask = [False]
|
||||
self.episode_height_sum = 0.0
|
||||
self._first_reset = True
|
||||
|
||||
# Initialize Simulation Manager
|
||||
@@ -87,12 +97,38 @@ class JackBotEnv(gym.Env):
|
||||
self.start_positions = [[0.0, 0.0, 0.0]]
|
||||
self.max_distance_from_start = [0.0]
|
||||
self.max_survival_steps = 0
|
||||
self.curriculum_phase = 0
|
||||
self.curriculum_episode_limit = 150
|
||||
|
||||
# Curriculum Initialization via Enum
|
||||
self.curriculum_phase = CurriculumPhase.STAND_ONLY
|
||||
self.curriculum_episode_limit = 500 # 500 steps limit gives headroom for 400-step requirement
|
||||
|
||||
# Gates required to unlock each target phase
|
||||
self.curriculum_stage_requirements = {
|
||||
1: {"survival_steps": 1000, "distance": 0.00, "stability_roll_pitch": 0.35},
|
||||
2: {"survival_steps": 350, "distance": 10.0, "stability_roll_pitch": 0.30},
|
||||
3: {"survival_steps": 550, "distance": 15.00, "stability_roll_pitch": 0.25},
|
||||
CurriculumPhase.FORWARD: {
|
||||
"survival_steps": 400, # Must survive ~8 seconds
|
||||
"max_displacement": 0.25, # Must remain within 0.25m radius
|
||||
"min_avg_height_ratio": 0.90, # Average height >= 90% of target
|
||||
"stability_roll_pitch": 0.18, # Max ~10 degrees tilt
|
||||
},
|
||||
CurriculumPhase.TURN_AND_DIRECTION: {
|
||||
"survival_steps": 500,
|
||||
"min_forward_distance": 2.5, # Must walk +2.5m forward (+X)
|
||||
"max_lateral_drift": 0.8, # Max 0.8m drift on Y axis
|
||||
"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,
|
||||
},
|
||||
}
|
||||
|
||||
# Floating HUD & Leader Crown Visualizers
|
||||
@@ -107,24 +143,36 @@ class JackBotEnv(gym.Env):
|
||||
"""Curriculum command sampler with survival-gated difficulty progression."""
|
||||
phase = self.curriculum_phase
|
||||
|
||||
if phase == 0:
|
||||
# Phase 1: Forward Walking Focus
|
||||
if phase == CurriculumPhase.STAND_ONLY:
|
||||
# Phase 0: Pure Standing (Zero commands)
|
||||
vx = 0.0
|
||||
vy = 0.0
|
||||
vz = 0.0
|
||||
omega = 0.0
|
||||
elif phase == CurriculumPhase.FORWARD:
|
||||
# Phase 1: Straight Forward Walking
|
||||
vx = np.random.uniform(0.5, 1.0)
|
||||
vy = 0.0
|
||||
vz = 0.0
|
||||
omega = 0.0
|
||||
elif phase == 1:
|
||||
elif phase == CurriculumPhase.TURN_AND_DIRECTION:
|
||||
# Phase 2: Forward/Backward + Turning
|
||||
vx = np.random.uniform(-1.0, 1.0)
|
||||
vy = 0.0
|
||||
vz = 0.0
|
||||
omega = np.random.uniform(-0.8, 0.8)
|
||||
else:
|
||||
elif phase == CurriculumPhase.OMNI_DIRECTION:
|
||||
# Phase 3: Full Omnidirectional Movement
|
||||
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)
|
||||
else:
|
||||
# Phase 4: Full Unconstrained Commands
|
||||
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)
|
||||
|
||||
return np.array([vx, vy, vz, omega], dtype=np.float32)
|
||||
|
||||
@@ -135,6 +183,7 @@ class JackBotEnv(gym.Env):
|
||||
self.cumulative_reward = 0.0
|
||||
self.robot_rewards = [0.0]
|
||||
self.failed_robots_mask = [False]
|
||||
self.episode_height_sum = 0.0
|
||||
|
||||
for idx, (pb_id, robot_obj) in enumerate(zip(self.pb_robots, self.robots)):
|
||||
spawn_pos = self._robot_base_position(idx, self.robot_spacing)
|
||||
@@ -152,8 +201,8 @@ class JackBotEnv(gym.Env):
|
||||
self.exploration_bonus_active = False
|
||||
|
||||
if self._first_reset:
|
||||
self.curriculum_phase = 0
|
||||
self.curriculum_episode_limit = min(self.max_episode_steps, 150)
|
||||
self.curriculum_phase = CurriculumPhase.STAND_ONLY
|
||||
self.curriculum_episode_limit = min(self.max_episode_steps, 500)
|
||||
self._first_reset = False
|
||||
|
||||
if self.random_command:
|
||||
@@ -190,6 +239,9 @@ class JackBotEnv(gym.Env):
|
||||
metrics = []
|
||||
for idx, pb_id in enumerate(self.pb_robots):
|
||||
pos, _ = self.sim_manager.get_robot_pose(pb_id)
|
||||
if idx == 0:
|
||||
self.episode_height_sum += float(pos[2])
|
||||
|
||||
start_x, start_y, _ = self.start_positions[idx]
|
||||
dist = float(np.linalg.norm(np.array([pos[0] - start_x, pos[1] - start_y], dtype=np.float32)))
|
||||
self.max_distance_from_start[idx] = max(self.max_distance_from_start[idx], dist)
|
||||
@@ -251,6 +303,13 @@ class JackBotEnv(gym.Env):
|
||||
for robot, act in zip(self.robots, action_per_robot):
|
||||
robot.apply_rl_action(act)
|
||||
|
||||
if self.step_count % 60 == 0:
|
||||
random_force = np.random.uniform(-2.0, 2.0, size=2) # X and Y push (Newtons)
|
||||
self.sim_manager.apply_external_force(
|
||||
body_id=self.pb_robots[0],
|
||||
force=[random_force[0], random_force[1], 0.0]
|
||||
)
|
||||
|
||||
self.sim_manager.step()
|
||||
|
||||
self._update_robot_failures()
|
||||
@@ -270,48 +329,72 @@ class JackBotEnv(gym.Env):
|
||||
self._update_leader_visuals()
|
||||
return obs, reward, terminated, truncated, {}
|
||||
|
||||
def _phase_progress_ready(self, phase: int) -> bool:
|
||||
if phase not in self.curriculum_stage_requirements:
|
||||
def _phase_progress_ready(self, next_phase: CurriculumPhase) -> bool:
|
||||
if next_phase not in self.curriculum_stage_requirements or not self.pb_robots:
|
||||
return False
|
||||
|
||||
if not self.pb_robots or not self.max_distance_from_start:
|
||||
return False
|
||||
|
||||
req = self.curriculum_stage_requirements[phase]
|
||||
req = self.curriculum_stage_requirements[next_phase]
|
||||
|
||||
# 1. Survival Step Check
|
||||
survival_ok = self.max_survival_steps >= req["survival_steps"]
|
||||
distance_ok = self.max_distance_from_start[0] >= req["distance"]
|
||||
|
||||
# 2. Body Posture & Tilt Check
|
||||
position, (roll, pitch, _) = self.sim_manager.get_robot_pose_and_rpy(self.pb_robots[0])
|
||||
stability_ok = abs(roll) <= req["stability_roll_pitch"] and abs(pitch) <= req["stability_roll_pitch"]
|
||||
height_ok = position[2] >= max(0.09, self.target_height * 0.85)
|
||||
|
||||
return survival_ok and distance_ok and stability_ok and height_ok
|
||||
# 3. Average Height Ratio Check Across the Episode
|
||||
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
|
||||
|
||||
# 4. Distance and Displacement Checks
|
||||
start_x, start_y, _ = self.start_positions[0]
|
||||
dx = position[0] - start_x
|
||||
dy = position[1] - start_y
|
||||
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"]
|
||||
|
||||
displacement_ok = True
|
||||
if "max_displacement" in req:
|
||||
displacement_ok = dist_2d <= req["max_displacement"]
|
||||
|
||||
drift_ok = True
|
||||
if "max_lateral_drift" in req:
|
||||
drift_ok = abs(dy) <= req["max_lateral_drift"]
|
||||
|
||||
return survival_ok and height_ok and stability_ok and displacement_ok and distance_ok and drift_ok
|
||||
|
||||
def _update_curriculum(self):
|
||||
self._curriculum_advanced = False
|
||||
|
||||
phase_labels = {
|
||||
0: "stand-and-forward",
|
||||
1: "turn-and-direction",
|
||||
2: "omni-direction",
|
||||
3: "full-command",
|
||||
}
|
||||
if self.curriculum_phase < CurriculumPhase.FORWARD and self._phase_progress_ready(CurriculumPhase.FORWARD):
|
||||
self.curriculum_phase = CurriculumPhase.FORWARD
|
||||
self.curriculum_episode_limit = min(self.max_episode_steps, 600)
|
||||
self._curriculum_advanced = True
|
||||
print(f"[Curriculum] Phase {self.curriculum_phase.name} unlocked at total step {self.total_steps}")
|
||||
|
||||
if self.curriculum_phase < 1 and self._phase_progress_ready(1):
|
||||
self.curriculum_phase = 1
|
||||
self.curriculum_episode_limit = min(self.max_episode_steps, 400)
|
||||
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_episode_limit = min(self.max_episode_steps, 800)
|
||||
self._curriculum_advanced = True
|
||||
print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}")
|
||||
elif self.curriculum_phase < 2 and self._phase_progress_ready(2):
|
||||
self.curriculum_phase = 2
|
||||
self.curriculum_episode_limit = min(self.max_episode_steps, 700)
|
||||
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_episode_limit = min(self.max_episode_steps, 1000)
|
||||
self._curriculum_advanced = True
|
||||
print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}")
|
||||
elif self.curriculum_phase < 3 and self._phase_progress_ready(3):
|
||||
self.curriculum_phase = 3
|
||||
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_episode_limit = min(self.max_episode_steps, self.max_episode_steps)
|
||||
self._curriculum_advanced = True
|
||||
print(f"[Curriculum] phase {self.curriculum_phase} ({phase_labels.get(self.curriculum_phase, 'unknown')}) unlocked at total step {self.total_steps}")
|
||||
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) -> Tuple[float, list[float]]:
|
||||
rewards = []
|
||||
@@ -361,9 +444,19 @@ class JackBotEnv(gym.Env):
|
||||
still_penalty = 0.35 + 0.85 * age_ratio
|
||||
|
||||
# 4. POSTURE & STABILITY PENALTIES
|
||||
height_penalty = 10.0 * ((self.target_height - pos[2]) ** 2) if pos[2] < self.target_height else 0.0
|
||||
height_penalty : float
|
||||
min_valid_h = self.target_height * 0.90 # 90% threshold
|
||||
|
||||
if pos[2] < min_valid_h:
|
||||
# Normalized drop below the 90% mark
|
||||
drop = (min_valid_h - pos[2]) / self.target_height
|
||||
# Linear + quadratic penalty that rapidly outweighs the +0.50 alive bonus
|
||||
height_penalty = 4.0 * drop + 20.0 * (drop ** 2)
|
||||
else:
|
||||
# Zero penalty inside the valid 90% - 110% zone!
|
||||
height_penalty = 0.0
|
||||
stability_penalty = 1.5 * (roll**2 + pitch**2)
|
||||
print(-height_penalty)
|
||||
|
||||
control_delta = np.abs(current_actions[idx] - previous_actions[idx])
|
||||
large_delta_mask = control_delta > 0.12
|
||||
large_delta_penalty = 0.002 * float(np.sum(np.square(control_delta[large_delta_mask]))) if np.any(large_delta_mask) else 0.0
|
||||
|
||||
Reference in New Issue
Block a user