JackM323/JackBot
1
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects 2 Releases Wiki Activity

adjusted rewards and requirements

Browse Source
This commit is contained in:
JackM323
2026-08-05 11:08:33 +02:00
parent f684bf44b8
commit 15e0206739
2 changed files with 108 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ml/env.py
+64 -33
View File
@@ -46,12 +46,14 @@ class JackBotEnv(gym.Env):
self.start_pose = start_pose
self.max_episode_steps = max_episode_steps
self.urdf_path = urdf_path
self.max_robot_speed = 0.8
self.episode_count = 0
self.step_count = 0
self.total_steps = 0
self.cumulative_reward = 0.0
self.robot_rewards = [0.0]
self.consecutive_still_steps = [0]
self.failed_robots_mask = [False]
self.episode_height_sum = 0.0
self._first_reset = True
@@ -106,7 +108,6 @@ class JackBotEnv(gym.Env):
self.curriculum_stage_requirements = {
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
},
@@ -150,17 +151,24 @@ class JackBotEnv(gym.Env):
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
# Phase 1: Straight Forward Walking + Standing Checks
if np.random.random() < 0.30:
# 30% chance: Stand Still Command (vx = 0)
vx = 0.0
else:
# 70% chance: Forward Walk Command
vx = np.random.uniform(0.15, 0.50)
vy, vz, omega = 0.0, 0.0, 0.0
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)
# Phase 2: Forward/Backward + Turning + Standing
if np.random.random() < 0.20:
vx, omega = 0.0, 0.0
else:
vx = np.random.uniform(-1.0, 1.0)
omega = np.random.uniform(-0.8, 0.8)
vy, vz = 0.0, 0.0
elif phase == CurriculumPhase.OMNI_DIRECTION:
# Phase 3: Full Omnidirectional Movement
vx = np.random.uniform(-1.0, 1.0)
@@ -195,6 +203,7 @@ class JackBotEnv(gym.Env):
if self.use_gui:
self.sim_manager.set_robot_color(pb_id, [1.0, 1.0, 1.0, 1.0])
self.consecutive_still_steps = [0 for _ in range(len(self.pb_robots))]
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
@@ -406,56 +415,78 @@ class JackBotEnv(gym.Env):
linear_vel, angular_vel = self.sim_manager.get_robot_velocity(pb_id)
command = self.commands[idx]
cmd_vx = command[0]
cmd_vy = command[1]
cmd_vx = command[0] # Normalized -1.0 to 1.0
cmd_vy = command[1] # Normalized -1.0 to 1.0
cmd_yaw = command[3]
# 1. LINEAR VECTOR SPEED MAXIMIZATION
cmd_dir = np.array([cmd_vx, cmd_vy], dtype=np.float32)
cmd_norm = np.linalg.norm(cmd_dir)
actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)
# -----------------------------------------------------------------
# 1. NORMALIZED SPEED TRACKING (1.0 = v_max)
# -----------------------------------------------------------------
if cmd_norm > 0.05:
unit_cmd_dir = cmd_dir / cmd_norm
actual_vel_2d = np.array([linear_vel[0], linear_vel[1]], dtype=np.float32)
aligned_speed = float(np.dot(actual_vel_2d, unit_cmd_dir))
linear_speed_reward = 2.5 * aligned_speed
target_speed = cmd_norm * self.max_robot_speed # e.g., 0.35 * 0.8 = 0.28 m/s
if aligned_speed <= 0.0:
# Zero or backward movement gets ZERO reward
linear_speed_reward = 0.0
elif aligned_speed <= target_speed:
# Smoothly scales from 0.0 to +5.0 as speed approaches target
linear_speed_reward = 5.0 * (aligned_speed / target_speed)
else:
# Gently penalize overspeeding beyond target
overspeed_ratio = (aligned_speed - target_speed) / target_speed
linear_speed_reward = max(0.0, 5.0 - 2.5 * overspeed_ratio)
# Drift penalty for sideways sliding
perp_vel = actual_vel_2d - aligned_speed * unit_cmd_dir
drift_penalty = 0.35 * float(np.dot(perp_vel, perp_vel))
else:
aligned_speed = 0.0
linear_speed_reward = 0.0
drift_penalty = 1.0 * (linear_vel[0]**2 + linear_vel[1]**2)
# 2. TURNING SPEED MAXIMIZATION
# -----------------------------------------------------------------
# 2. TIME-ACCUMULATING STILLNESS PENALTY
# -----------------------------------------------------------------
actual_yaw_rate = angular_vel[2]
is_moving = abs(aligned_speed) > 0.03 or abs(actual_yaw_rate) > 0.08
if (cmd_norm > 0.05 or abs(cmd_yaw) > 0.05) and not is_moving:
# Increment counter every step the robot stalls under command
self.consecutive_still_steps[idx] += 1
else:
# Reset counter as soon as robot makes a valid move!
self.consecutive_still_steps[idx] = 0
# Penalty grows by 0.05 every step frozen, capped at -3.5 per step
still_penalty = min(3.5, 0.05 * self.consecutive_still_steps[idx])
# -----------------------------------------------------------------
# 3. TURNING & POSTURE PENALTIES
# -----------------------------------------------------------------
if abs(cmd_yaw) > 0.05:
turning_reward = 1.5 * (actual_yaw_rate * cmd_yaw)
else:
turning_reward = -0.6 * (actual_yaw_rate ** 2)
# 3. SMALL ALIVE BONUS + AGE-RAMPED STILLNESS PENALTY
alive_reward = 0.5
age_ratio = min(1.0, self.step_count / max(1, self.curriculum_episode_limit))
still_penalty = 0.0
if (cmd_norm > 0.1 or abs(cmd_yaw) > 0.1) and (abs(linear_vel[0]) < 0.02 and abs(actual_yaw_rate) < 0.05):
still_penalty = 0.35 + 0.85 * age_ratio
# 4. POSTURE & STABILITY PENALTIES
height_penalty : float
min_valid_h = self.target_height * 0.90 # 90% threshold
alive_reward = 0.20
# Height drop penalty
min_valid_h = self.target_height * 0.90
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)
stability_penalty = 8 * (roll**2 + pitch**2)
control_delta = np.abs(current_actions[idx] - previous_actions[idx])
large_delta_mask = control_delta > 0.12
ml/train.py
+43 -22
View File
@@ -12,14 +12,25 @@ class MilestoneCheckpointCallback(BaseCallback):
"""
Saves a model checkpoint the FIRST time total_timesteps
crosses every multiple of step_interval (e.g., 100,000).
Saves inside the matching PPO_X subfolder as created by TensorBoard.
"""
def __init__(self, save_path: str, name_prefix: str = "ppo_jackbot", step_interval: int = 100_000, verbose: int = 1):
super().__init__(verbose)
self.save_path = save_path
self.base_save_path = save_path
self.run_save_path = save_path
self.name_prefix = name_prefix
self.step_interval = step_interval
self.last_milestone = 0
os.makedirs(self.save_path, exist_ok=True)
def _on_training_start(self) -> None:
"""Executed right before training loop starts. Resolves TensorBoard's run folder name (e.g. PPO_1)."""
if self.logger and self.logger.dir:
run_folder_name = Path(self.logger.dir).name # Extracts "PPO_1", "PPO_2", etc.
self.run_save_path = os.path.join(self.base_save_path, run_folder_name)
else:
self.run_save_path = self.base_save_path
os.makedirs(self.run_save_path, exist_ok=True)
def _on_step(self) -> bool:
current_milestone = self.num_timesteps // self.step_interval
@@ -29,7 +40,7 @@ class MilestoneCheckpointCallback(BaseCallback):
milestone_step = current_milestone * self.step_interval
save_file = os.path.join(
self.save_path,
self.run_save_path,
f"{self.name_prefix}_{milestone_step}_steps.zip"
)
self.model.save(save_file)
@@ -39,6 +50,7 @@ class MilestoneCheckpointCallback(BaseCallback):
return True
class JackBotMetricsCallback(BaseCallback):
"""
Tracks the best current alive-robot performance for the most recent rollout,
@@ -53,11 +65,9 @@ class JackBotMetricsCallback(BaseCallback):
self.best_alive_reward = -float('inf')
def _on_step(self) -> bool:
"""Required by SB3 BaseCallback; no-op here because the rollout summary is emitted at rollout end."""
return True
def _on_rollout_end(self) -> bool:
"""Executed right before PPO outputs the log table to console."""
try:
vec_env = self.training_env
alive_metrics = vec_env.env_method("get_current_robot_metrics")
@@ -90,7 +100,6 @@ class JackBotMetricsCallback(BaseCallback):
self.logger.record("custom/best_alive_survival_steps", float(self.best_alive_survival_steps))
self.logger.record("custom/best_alive_reward", float(self.best_alive_reward) if np.isfinite(self.best_alive_reward) else 0.0)
# Backward-compatible aliases so old dashboards keep a stable field name.
self.logger.record("custom/max_speed_mps", float(self.best_alive_speed))
self.logger.record("custom/max_yaw_rate_rads", float(self.best_alive_yaw_rate))
self.logger.record("custom/max_distance_from_start_m", float(self.best_alive_distance))
@@ -101,6 +110,7 @@ class JackBotMetricsCallback(BaseCallback):
return True
def parse_args():
parser = argparse.ArgumentParser(description="Train a joint-command policy for JackBot.")
parser.add_argument("--timesteps", type=int, default=500_000, help="Total training timesteps")
@@ -117,7 +127,7 @@ def parse_args():
def make_env(robot_spacing, start_pose, use_gui, rank, seed=0):
def _init():
env = JackBotEnv(
use_gui=use_gui if rank == 0 else False, # Only rank 0 gets GUI if requested
use_gui=use_gui if rank == 0 else False,
random_command=True,
robot_spacing=robot_spacing,
start_pose=start_pose,
@@ -142,7 +152,6 @@ def train(
except ImportError as exc:
raise ImportError("stable-baselines3 is required. Install with: pip install stable-baselines3") from exc
# Create multi-process vector environment
if num_workers > 1:
env_fns = [
make_env(robot_spacing, start_pose, use_gui, rank=i, seed=seed)
@@ -179,24 +188,24 @@ def train(
device = resolve_device(device)
policy_kwargs = dict(
log_std_init=-1.5, # Sets initial std ~ 0.22 instead of 1.0
net_arch=dict(pi=[256, 256], vf=[256, 256])
)
log_std_init=-1.5,
net_arch=dict(pi=[256, 256], vf=[256, 256])
)
model = PPO(
"MlpPolicy",
env,
verbose=1,
seed=seed,
learning_rate=1.5e-4, # Cut LR in half (from 3e-4) to smooth out updates
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
clip_range=0.2, # Standard PPO clipping
target_kl=0.03, # EARLY STOPPING: Halts policy update if KL > 0.015!
ent_coef=0.03, # Entropy coefficient to encourage exploration
learning_rate=1.5e-4,
n_steps=256,
batch_size=256,
n_epochs=10,
gamma=0.99,
gae_lambda=0.95,
clip_range=0.2,
target_kl=0.03,
ent_coef=0.03,
vf_coef=0.5,
max_grad_norm=0.5,
device=device,
@@ -204,6 +213,7 @@ def train(
tensorboard_log=str(Path(__file__).resolve().parent / "tensorboard"),
)
# Base folder where model runs will be stored
save_dir = str(Path(model_path).parent)
model_prefix = Path(model_path).stem
@@ -217,8 +227,19 @@ def train(
model.learn(total_timesteps=total_timesteps, callback=[milestone_cb, metrics_callback])
Path(model_path).parent.mkdir(parents=True, exist_ok=True)
model.save(model_path)
# Save the final model inside the matching PPO_X directory as well
if model.logger and model.logger.dir:
run_folder_name = Path(model.logger.dir).name
final_dir = Path(model_path).parent / run_folder_name
else:
final_dir = Path(model_path).parent
final_dir.mkdir(parents=True, exist_ok=True)
final_save_path = final_dir / f"{model_prefix}_final.zip"
model.save(str(final_save_path))
if model.verbose > 0:
print(f"[Training Complete] Saved final model to -> {final_save_path}")
env.close()