Johnny Fernandes
247 lines
8.6 KiB
Python
247 lines
8.6 KiB
Python
"""
|
|
PPO training script for the herding task.
|
|
|
|
Usage examples
|
|
--------------
|
|
# Proper 5-sheep curriculum, 1 M steps per stage:
|
|
python train.py --curriculum --steps-per-stage 1000000 --total-steps 5000000
|
|
|
|
# Success-rate curriculum (advances when 70 % success over 100 episodes):
|
|
python train.py --curriculum --threshold 0.70
|
|
|
|
# Resume from checkpoint at stage 3:
|
|
python train.py --resume runs/ppo_herding/ckpt_3000000_steps.zip --n-sheep 3 \
|
|
--curriculum --steps-per-stage 1000000 --total-steps 5000000
|
|
|
|
# Quick smoke-test:
|
|
python train.py --n-envs 1 --total-steps 50000
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
|
|
import numpy as np
|
|
from stable_baselines3 import PPO
|
|
from stable_baselines3.common.callbacks import (
|
|
BaseCallback,
|
|
CallbackList,
|
|
CheckpointCallback,
|
|
EvalCallback,
|
|
)
|
|
from stable_baselines3.common.vec_env import SubprocVecEnv, VecNormalize
|
|
|
|
from herding_env import HerdingEnv
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Curriculum callback
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class CurriculumCallback(BaseCallback):
|
|
"""
|
|
Advances n_sheep on both training and eval envs.
|
|
|
|
Two modes (mutually exclusive):
|
|
steps_per_stage — advance every N environment steps regardless of
|
|
success rate (recommended for reliability).
|
|
threshold — advance when rolling success rate exceeds this value
|
|
(requires the policy to actually reach the threshold).
|
|
"""
|
|
|
|
def __init__(self, start_sheep: int, max_sheep: int,
|
|
eval_env=None,
|
|
steps_per_stage: int = None,
|
|
threshold: float = 0.75,
|
|
window: int = 100,
|
|
min_episodes: int = 50,
|
|
verbose: int = 1):
|
|
super().__init__(verbose)
|
|
self.max_sheep = max_sheep
|
|
self.eval_env = eval_env
|
|
self.steps_per_stage = steps_per_stage
|
|
self.threshold = threshold
|
|
self.window = window
|
|
self.min_episodes = min_episodes
|
|
self._cur_sheep = start_sheep
|
|
self._successes = []
|
|
self._stage_start = 0
|
|
|
|
def _advance(self):
|
|
self._cur_sheep += 1
|
|
self.training_env.env_method("set_n_sheep", self._cur_sheep)
|
|
if self.eval_env is not None:
|
|
self.eval_env.env_method("set_n_sheep", self._cur_sheep)
|
|
self._stage_start = self.num_timesteps
|
|
self._successes.clear()
|
|
if self.verbose:
|
|
print(f"\n[Curriculum] → {self._cur_sheep} sheep "
|
|
f"at step {self.num_timesteps:,}\n")
|
|
|
|
def _on_step(self) -> bool:
|
|
if self._cur_sheep >= self.max_sheep:
|
|
return True
|
|
|
|
if self.steps_per_stage is not None:
|
|
# Time-based: advance every steps_per_stage env steps
|
|
if self.num_timesteps - self._stage_start >= self.steps_per_stage:
|
|
self._advance()
|
|
else:
|
|
# Success-rate based
|
|
for info, done in zip(self.locals["infos"], self.locals["dones"]):
|
|
if done:
|
|
truncated = info.get("TimeLimit.truncated", False)
|
|
self._successes.append(0 if truncated else 1)
|
|
if len(self._successes) > self.window:
|
|
self._successes.pop(0)
|
|
|
|
if (len(self._successes) >= self.min_episodes
|
|
and np.mean(self._successes) >= self.threshold):
|
|
self._advance()
|
|
|
|
return True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Environment factory
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def make_env(n_sheep: int, seed: int, max_steps: int):
|
|
def _init():
|
|
env = HerdingEnv(n_sheep=n_sheep, max_steps=max_steps)
|
|
env.reset(seed=seed)
|
|
return env
|
|
return _init
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def parse_args():
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--n-sheep", type=int, default=1,
|
|
help="Starting sheep count")
|
|
p.add_argument("--max-sheep", type=int, default=5,
|
|
help="Final sheep count for curriculum")
|
|
p.add_argument("--n-envs", type=int, default=8,
|
|
help="Parallel training environments")
|
|
p.add_argument("--total-steps", type=int, default=5_000_000)
|
|
p.add_argument("--max-steps", type=int, default=2000,
|
|
help="Episode step limit")
|
|
p.add_argument("--curriculum", action="store_true",
|
|
help="Enable curriculum advancement")
|
|
p.add_argument("--steps-per-stage", type=int, default=None,
|
|
help="Advance curriculum every N steps (overrides --threshold)")
|
|
p.add_argument("--threshold", type=float, default=0.75,
|
|
help="Success-rate threshold to advance (used without --steps-per-stage)")
|
|
p.add_argument("--resume", type=str, default=None,
|
|
help="Checkpoint .zip to resume from")
|
|
p.add_argument("--run-dir", type=str, default="runs/ppo_herding")
|
|
p.add_argument("--save-freq", type=int, default=100_000)
|
|
p.add_argument("--eval-freq", type=int, default=50_000)
|
|
p.add_argument("--eval-eps", type=int, default=20)
|
|
return p.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
os.makedirs(args.run_dir, exist_ok=True)
|
|
ckpt_dir = os.path.join(args.run_dir, "checkpoints")
|
|
best_dir = os.path.join(args.run_dir, "best_model")
|
|
norm_path = os.path.join(args.run_dir, "vecnorm.pkl")
|
|
os.makedirs(ckpt_dir, exist_ok=True)
|
|
|
|
# Training envs
|
|
train_env = SubprocVecEnv([
|
|
make_env(args.n_sheep, seed=i, max_steps=args.max_steps)
|
|
for i in range(args.n_envs)
|
|
])
|
|
if args.resume and os.path.exists(norm_path):
|
|
train_env = VecNormalize.load(norm_path, train_env)
|
|
train_env.training = True
|
|
train_env.norm_reward = True
|
|
else:
|
|
train_env = VecNormalize(train_env, norm_obs=True, norm_reward=True,
|
|
clip_obs=10.0)
|
|
|
|
# Eval env — starts at same difficulty, advances with curriculum callback
|
|
eval_env = SubprocVecEnv([
|
|
make_env(args.n_sheep, seed=1000 + i, max_steps=args.max_steps)
|
|
for i in range(2)
|
|
])
|
|
eval_env = VecNormalize(eval_env, norm_obs=True, norm_reward=False,
|
|
clip_obs=10.0, training=False)
|
|
|
|
# Callbacks
|
|
checkpoint_cb = CheckpointCallback(
|
|
save_freq=max(args.save_freq // args.n_envs, 1),
|
|
save_path=ckpt_dir,
|
|
name_prefix="ckpt",
|
|
save_vecnormalize=True,
|
|
)
|
|
eval_cb = EvalCallback(
|
|
eval_env,
|
|
best_model_save_path=best_dir,
|
|
log_path=args.run_dir,
|
|
eval_freq=max(args.eval_freq // args.n_envs, 1),
|
|
n_eval_episodes=args.eval_eps,
|
|
deterministic=True,
|
|
verbose=1,
|
|
)
|
|
callbacks = [checkpoint_cb, eval_cb]
|
|
|
|
if args.curriculum:
|
|
cur_cb = CurriculumCallback(
|
|
start_sheep=args.n_sheep,
|
|
max_sheep=args.max_sheep,
|
|
eval_env=eval_env,
|
|
steps_per_stage=args.steps_per_stage,
|
|
threshold=args.threshold,
|
|
)
|
|
callbacks.append(cur_cb)
|
|
|
|
callback_list = CallbackList(callbacks)
|
|
|
|
# Model
|
|
ppo_kwargs = dict(
|
|
policy = "MlpPolicy",
|
|
env = train_env,
|
|
learning_rate = 3e-4,
|
|
n_steps = 2048,
|
|
batch_size = 256,
|
|
n_epochs = 10,
|
|
gamma = 0.995,
|
|
gae_lambda = 0.95,
|
|
clip_range = 0.2,
|
|
ent_coef = 0.005,
|
|
vf_coef = 0.5,
|
|
max_grad_norm = 0.5,
|
|
policy_kwargs = dict(net_arch=[256, 256]),
|
|
tensorboard_log = args.run_dir,
|
|
verbose = 1,
|
|
)
|
|
|
|
if args.resume:
|
|
print(f"Resuming from {args.resume}")
|
|
model = PPO.load(args.resume, env=train_env, **{
|
|
k: v for k, v in ppo_kwargs.items()
|
|
if k not in ("policy", "env")
|
|
})
|
|
else:
|
|
model = PPO(**ppo_kwargs)
|
|
|
|
model.learn(
|
|
total_timesteps=args.total_steps,
|
|
callback=callback_list,
|
|
reset_num_timesteps=args.resume is None,
|
|
tb_log_name="ppo",
|
|
)
|
|
|
|
model.save(os.path.join(args.run_dir, "final_model"))
|
|
train_env.save(norm_path)
|
|
print(f"\nTraining complete. Artefacts saved to {args.run_dir}/")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|