reward ajustments

robot starts balancing without changing phase and farms alive bonus
fix ->external force and bigger height penalty
This commit is contained in:
2026-08-04 22:17:33 +02:00
parent 402c20dfb5
commit 101004ead1
4 changed files with 163 additions and 45 deletions
+26
View File
@@ -164,3 +164,29 @@ class SimManager:
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
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,
)
+131 -38
View File
@@ -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[next_phase]
req = self.curriculum_stage_requirements[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
View File
@@ -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
View File
@@ -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