Skills › Scientific Agent Skills › stable-baselines3 › train_rl_agent.py
train_rl_agent.py
99 lines · skills/stable-baselines3/scripts/train_rl_agent.py at 5f2bc03
1"""CPU-friendly training with held-out monitoring and matching normalization files.
3Run SubprocVecEnv callers from a file under an if __name__ == '__main__' guard.
4Evaluation during training selects checkpoints; use new seeds for final testing.
5"""
7from pathlib import Path
9from stable_baselines3 import PPO
10from stable_baselines3.common.callbacks import BaseCallback, CheckpointCallback, EvalCallback
11from stable_baselines3.common.env_util import make_vec_env
12from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
15class SaveBestNormalization(BaseCallback):
16 """Pair each EvalCallback best model with statistics from that same step."""
18 def __init__(self, path):
19 super().__init__()
20 self.path = Path(path)
22 def _on_step(self):
23 normalizer = self.model.get_vec_normalize_env()
24 if normalizer is not None:
25 normalizer.save(str(self.path))
26 return True
29def train_agent(
30 env_id="CartPole-v1", algorithm=PPO, policy="MlpPolicy", n_envs=4,
31 total_timesteps=100000, eval_freq=10000, save_freq=10000,
32 log_dir="./logs/", save_path="./models/", *, seed=0, eval_seed=10000,
33 n_eval_episodes=10, normalize=False, vec_env_cls=DummyVecEnv,
34 algorithm_kwargs=None, tensorboard=True, device="cpu",
35 save_replay_buffer=False, env_kwargs=None,
36):
37 """Train and save a core SB3 model; return it with its environments closed.
39 Frequencies are requested total transitions (rounded down to a vector step,
40 at least one call). total_timesteps is a lower bound due to rollout collection.
41 algorithm_kwargs supplies algorithm-specific parameters, e.g. PPO n_steps.
42 normalize writes final_model_vecnormalize.pkl and best_model/vecnormalize.pkl,
43 and matching statistics for periodic checkpoints. Use a fresh output directory
44 per run. For off-policy continuation, opt in to replay-buffer saving.
45 """
46 for name, value in (("n_envs", n_envs), ("total_timesteps", total_timesteps),
47 ("eval_freq", eval_freq), ("save_freq", save_freq),
48 ("n_eval_episodes", n_eval_episodes)):
49 if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
50 raise ValueError(f"{name} must be a positive integer")
51 log_dir, save_path = Path(log_dir), Path(save_path)
52 log_dir.mkdir(parents=True, exist_ok=True)
53 save_path.mkdir(parents=True, exist_ok=True)
54 best_path = save_path / "best_model"
55 env = eval_env = None
56 try:
57 env = make_vec_env(env_id, n_envs=n_envs, seed=seed,
58 vec_env_cls=vec_env_cls, env_kwargs=env_kwargs)
59 eval_env = make_vec_env(env_id, n_envs=1, seed=eval_seed, env_kwargs=env_kwargs)
60 params = dict(verbose=1, tensorboard_log=str(log_dir) if tensorboard else None,
61 seed=seed, device=device)
62 params.update(algorithm_kwargs or {})
63 if normalize:
64 gamma = params.get("gamma", 0.99)
65 env = VecNormalize(env, gamma=gamma)
66 eval_env = VecNormalize(eval_env, training=False, norm_reward=False, gamma=gamma)
67 evaluation = EvalCallback(
68 eval_env, best_model_save_path=str(best_path), log_path=str(log_dir / "eval"),
69 eval_freq=max(eval_freq // n_envs, 1), n_eval_episodes=n_eval_episodes,
70 deterministic=True,
71 callback_on_new_best=SaveBestNormalization(best_path / "vecnormalize.pkl"),
72 )
73 checkpoint = CheckpointCallback(
74 save_freq=max(save_freq // n_envs, 1), save_path=str(save_path),
75 name_prefix="rl_model", save_replay_buffer=save_replay_buffer,
76 save_vecnormalize=normalize,
77 )
78 model = algorithm(policy, env, **params)
79 model.learn(total_timesteps=total_timesteps, callback=[evaluation, checkpoint],
80 tb_log_name=f"{algorithm.__name__}_{env_id}")
81 model.save(str(save_path / "final_model"))
82 if normalize:
83 env.save(str(save_path / "final_model_vecnormalize.pkl"))
84 if save_replay_buffer and getattr(model, "replay_buffer", None) is not None:
85 model.save_replay_buffer(str(save_path / "final_model_replay_buffer.pkl"))
86 print(f"Training complete! Actual transitions: {model.num_timesteps}")
87 print(f"Final model: {save_path / 'final_model.zip'}")
88 if (best_path / "best_model.zip").exists():
89 print(f"Best validation model: {best_path / 'best_model.zip'}")
90 return model
91 finally:
92 if eval_env is not None:
93 eval_env.close()
94 if env is not None:
95 env.close()
98if __name__ == "__main__":
99 train_agent()