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/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()