code reduction for training and eval

reward and everything else changed
again
wont be the last time

ml/run_train.py

@@ -1,47 +1,132 @@
 """Minimal training launcher for quick experiments.
 
 Usage:
-    python ml/run_train.py --timesteps 500000 --model ml/checkpoints/ppo_joint_command
-
-This is a convenience wrapper around `ml.train.train` with friendly defaults
-for interactive experimentation.
+    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
 
-ROOT_DIR = Path(__file__).resolve().parent.parent
-if str(ROOT_DIR) not in sys.path:
-    sys.path.insert(0, str(ROOT_DIR))
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+from ml.env import JackBotEnv, CurriculumCallback
 
-from ml.train import train
+# 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()
-    parser.add_argument("--timesteps", type=int, default=500000, help="Total number of 'practice steps'.")
-    parser.add_argument("--model", type=str, default="ml/checkpoints/ppo_joint_command", help="Path to save the trained model")
-    parser.add_argument("--seed", type=int, default=0, help="Random seed for reproducibility")
-    parser.add_argument("--device", type=str, default="auto", help="Device to use: 'cpu', 'cuda', or 'auto'")
-    parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during training")
-    parser.add_argument("--num-workers", type=int, default=1, help="Number of training environment workers")
-    parser.add_argument("--robot-spacing", type=float, default=1.5, help="Spacing between robots in meters")
-    parser.add_argument("--start-pose", type=str, choices=["init_deg", "init90_deg"], default="init_deg", help="Initial robot pose at reset")
+    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()
 
-    Path(args.model).parent.mkdir(parents=True, exist_ok=True)
-    train(
-        total_timesteps=args.timesteps,
-        model_path=args.model,
-        seed=args.seed,
-        device=args.device,
-        use_gui=args.gui,
-        num_workers=args.num_workers,
-        robot_spacing=args.robot_spacing,
-        start_pose=args.start_pose,
+    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()
+    main()