reward ajustments
robot starts balancing without changing phase and farms alive bonus fix ->external force and bigger height penalty
This commit is contained in:
+27
-1
@@ -163,4 +163,30 @@ class SimManager:
|
||||
self.step()
|
||||
heights = self.measure_robot_heights(robot_ids)
|
||||
mean_height = float(np.mean(heights)) if heights else fallback_height
|
||||
return mean_height if mean_height > 0.0 else fallback_height
|
||||
return mean_height if mean_height > 0.0 else fallback_height
|
||||
|
||||
def apply_external_force(
|
||||
self,
|
||||
body_id: int,
|
||||
force: list[float] | np.ndarray,
|
||||
link_index: int = -1,
|
||||
position: list[float] | np.ndarray = (0.0, 0.0, 0.0),
|
||||
frame: int = p.WORLD_FRAME,
|
||||
):
|
||||
"""
|
||||
Applies a 3D force vector (in Newtons) to a robot link.
|
||||
|
||||
:param body_id: PyBullet body ID.
|
||||
:param force: [fx, fy, fz] force vector in Newtons.
|
||||
:param link_index: Target link index (-1 refers to the base/torso).
|
||||
:param position: Offset [x, y, z] relative to link center where force is applied.
|
||||
:param frame: p.WORLD_FRAME (global axes) or p.LINK_FRAME (robot's body axes).
|
||||
"""
|
||||
p.applyExternalForce(
|
||||
objectUniqueId=body_id,
|
||||
linkIndex=link_index,
|
||||
forceObj=list(force),
|
||||
posObj=list(position),
|
||||
flags=frame,
|
||||
physicsClientId=self.physics_client,
|
||||
)
|
||||
@@ -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
|
||||
|
||||
+2
-3
@@ -44,17 +44,16 @@ def evaluate(
|
||||
|
||||
for ep in range(episodes):
|
||||
obs, _ = env.reset()
|
||||
done = False
|
||||
terminated = False
|
||||
total_reward = 0.0
|
||||
steps = 0
|
||||
|
||||
print(f"\n--- Starting Evaluation Episode {ep + 1}/{episodes} ---")
|
||||
|
||||
while not done:
|
||||
while not terminated:
|
||||
action, _ = model.predict(obs, deterministic=True)
|
||||
obs, reward, terminated, truncated, _ = env.step(action)
|
||||
|
||||
done = terminated or truncated
|
||||
total_reward += float(reward)
|
||||
steps += 1
|
||||
|
||||
|
||||
+2
-2
@@ -189,8 +189,8 @@ def train(
|
||||
verbose=1,
|
||||
seed=seed,
|
||||
learning_rate=1.5e-4, # Cut LR in half (from 3e-4) to smooth out updates
|
||||
n_steps=1024, # Larger rollout buffer per env for stable gradients
|
||||
batch_size=128, # Larger minibatches reduce noise
|
||||
n_steps=256, # Larger rollout buffer per env for stable gradients
|
||||
batch_size=256, # Larger minibatches reduce noise
|
||||
n_epochs=10, # Number of epoch updates per rollout
|
||||
gamma=0.99, # Discount factor
|
||||
gae_lambda=0.95, # GAE smoothing
|
||||
|
||||
Reference in New Issue
Block a user