From 5480b3cf05c9e4b1cf604d6cbea1f8188a492ccc Mon Sep 17 00:00:00 2001 From: Josh Creek <8179928+jcreek@users.noreply.github.com> Date: Sun, 19 Jul 2026 10:41:11 +0100 Subject: [PATCH] feat(*): Add --n-steps and --batch-size flags to train.py, applied on resume as well --- training/train.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/training/train.py b/training/train.py index 04f43cbd..995a30b5 100644 --- a/training/train.py +++ b/training/train.py @@ -40,6 +40,8 @@ def parse_args(): parser.add_argument("--seed", type=int, default=0) parser.add_argument("--resume", default=None, help="Checkpoint .zip to resume from") parser.add_argument("--ent-coef", type=float, default=0.0001, help="Entropy bonus coefficient (applied on resume too)") + parser.add_argument("--n-steps", type=int, default=256, help="Rollout length per env between updates (applied on resume too)") + parser.add_argument("--batch-size", type=int, default=256, help="PPO minibatch size (applied on resume too)") parser.add_argument( "--reset-std", type=float, @@ -74,8 +76,18 @@ def main(): env = VecMonitor(env) if args.resume: - model = PPO.load(args.resume, env=env, tensorboard_log=str(log_dir), ent_coef=args.ent_coef) - print(f"Resumed from {args.resume} at {model.num_timesteps} timesteps (ent_coef={args.ent_coef})") + model = PPO.load( + args.resume, + env=env, + tensorboard_log=str(log_dir), + ent_coef=args.ent_coef, + n_steps=args.n_steps, + batch_size=args.batch_size, + ) + print( + f"Resumed from {args.resume} at {model.num_timesteps} timesteps " + f"(ent_coef={args.ent_coef}, n_steps={args.n_steps}, batch_size={args.batch_size})" + ) if args.reset_std is not None: import math @@ -90,8 +102,8 @@ def main(): env, verbose=1, ent_coef=args.ent_coef, - n_steps=256, - batch_size=256, + n_steps=args.n_steps, + batch_size=args.batch_size, learning_rate=3e-4, tensorboard_log=str(log_dir), )