132 lines
4.5 KiB
Python
132 lines
4.5 KiB
Python
"""Minimal training launcher for quick experiments.
|
|
|
|
Usage:
|
|
python ml/run_train.py --total-timesteps 1500000 --gui
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import re
|
|
import sys
|
|
import warnings
|
|
from pathlib import Path
|
|
from stable_baselines3 import PPO
|
|
from stable_baselines3.common.vec_env import SubprocVecEnv, DummyVecEnv
|
|
from stable_baselines3.common.callbacks import CheckpointCallback, EvalCallback
|
|
|
|
sys.path.append(str(Path(__file__).resolve().parent.parent))
|
|
from ml.env import JackBotEnv, CurriculumCallback
|
|
|
|
# Silence SB3's UserWarning about SubprocVecEnv vs DummyVecEnv
|
|
warnings.filterwarnings("ignore", category=UserWarning, module="stable_baselines3")
|
|
|
|
|
|
def get_next_run_number(save_dir: str) -> int:
|
|
"""Scans the save directory for existing ppo<number> patterns and returns the next integer."""
|
|
if not os.path.exists(save_dir):
|
|
return 1
|
|
|
|
existing_numbers = []
|
|
for item in os.listdir(save_dir):
|
|
# Match pattern ppo followed by numbers (e.g., ppo1, ppo_1, jackbot_ppo12)
|
|
matches = re.findall(r"ppo_?(\d+)", item, re.IGNORECASE)
|
|
for m in matches:
|
|
existing_numbers.append(int(m))
|
|
|
|
return max(existing_numbers, default=0) + 1
|
|
|
|
|
|
def make_env(rank: int, use_gui: bool = False, seed: int = 0):
|
|
"""Utility helper to instantiate parallel JackBot environments."""
|
|
def _init():
|
|
env = JackBotEnv(
|
|
use_gui=use_gui,
|
|
random_command=True,
|
|
)
|
|
env.reset(seed=seed + rank)
|
|
return env
|
|
return _init
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="JackBot PPO Curriculum Trainer")
|
|
parser.add_argument("--num-workers", type=int, default=16, help="Number of parallel sub-process environments")
|
|
parser.add_argument("--total-timesteps", type=int, default=1_500_000, help="Total training timesteps")
|
|
parser.add_argument("--log-dir", type=str, default="ml/logs", help="Directory for TensorBoard logs")
|
|
parser.add_argument("--save-dir", type=str, default="ml/checkpoints", help="Directory for model checkpoints")
|
|
parser.add_argument("--save-freq", type=int, default=50_000, help="Checkpoint save frequency (steps)")
|
|
parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering")
|
|
args = parser.parse_args()
|
|
|
|
os.makedirs(args.log_dir, exist_ok=True)
|
|
os.makedirs(args.save_dir, exist_ok=True)
|
|
|
|
# Automatically determine next run number (e.g. ppo1, ppo2, ppo3...)
|
|
run_num = get_next_run_number(args.save_dir)
|
|
ppo_name = f"ppo{run_num}"
|
|
|
|
print(f"[Train] Initializing Run #{run_num} ('{ppo_name}') with {args.num_workers} parallel workers...")
|
|
|
|
env_fns = [
|
|
make_env(rank=i, use_gui=(args.gui if i == 0 else False))
|
|
for i in range(args.num_workers)
|
|
]
|
|
vec_env = SubprocVecEnv(env_fns)
|
|
|
|
# Initialize PPO Policy Hyperparameters
|
|
model = PPO(
|
|
policy="MlpPolicy",
|
|
env=vec_env,
|
|
learning_rate=3e-4,
|
|
n_steps=256,
|
|
batch_size=256,
|
|
n_epochs=10,
|
|
gamma=0.99,
|
|
gae_lambda=0.95,
|
|
clip_range=0.2,
|
|
ent_coef=0.03,
|
|
target_kl=0.05,
|
|
vf_coef=0.5,
|
|
max_grad_norm=0.5,
|
|
verbose=1,
|
|
tensorboard_log=args.log_dir,
|
|
device="cpu",
|
|
)
|
|
|
|
# Setup Callbacks with ppo<number> naming
|
|
checkpoint_callback = CheckpointCallback(
|
|
save_freq=max(1, args.save_freq // args.num_workers),
|
|
save_path=args.save_dir,
|
|
name_prefix=f"jackbot_{ppo_name}",
|
|
)
|
|
curriculum_callback = CurriculumCallback()
|
|
|
|
eval_env = DummyVecEnv([lambda: JackBotEnv(use_gui=False, random_command=True)])
|
|
best_model_path = os.path.join(args.save_dir, f"best_model_{ppo_name}")
|
|
|
|
eval_callback = EvalCallback(
|
|
eval_env,
|
|
best_model_save_path=best_model_path,
|
|
log_path="ml/logs/results",
|
|
eval_freq=max(1, 10_000 // args.num_workers),
|
|
deterministic=True,
|
|
render=False,
|
|
)
|
|
|
|
print(f"[Train] Starting training for {args.total_timesteps} timesteps...")
|
|
try:
|
|
model.learn(
|
|
total_timesteps=args.total_timesteps,
|
|
callback=[checkpoint_callback, curriculum_callback, eval_callback],
|
|
progress_bar=True,
|
|
)
|
|
final_model_path = os.path.join(args.save_dir, f"jackbot_{ppo_name}_final.zip")
|
|
model.save(final_model_path)
|
|
print(f"[Train] Training complete! Saved final model to {final_model_path}")
|
|
finally:
|
|
vec_env.close()
|
|
eval_env.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |