code reduction for training and eval

reward and everything else changed again wont be the last time
2026-08-05 22:44:05 +02:00
parent 15e0206739
commit c93c524a10
10 changed files with 606 additions and 814 deletions
@@ -5,38 +5,109 @@ Usage:
 """

 import argparse
+import json
+import time
 import sys
 from pathlib import Path
-
-ROOT_DIR = Path(__file__).resolve().parent.parent
-if str(ROOT_DIR) not in sys.path:
-    sys.path.insert(0, str(ROOT_DIR))
-
-from ml.evaluate import evaluate
-
+import numpy as np
+from stable_baselines3 import PPO
+sys.path.append(str(Path(__file__).resolve().parent.parent))
+from ml.env import JackBotEnv

 def main():
-    parser = argparse.ArgumentParser(description="JackBot Curriculum Policy Evaluator Wrapper")
-    parser.add_argument("--model", type=str, required=True, help="Path to the trained model file (.zip)")
-    parser.add_argument("--episodes", type=int, default=5, help="Number of evaluation episodes to run.")
-    parser.add_argument("--gui", action="store_true", help="Show the PyBullet GUI during evaluation")
-    parser.add_argument("--robot-spacing", type=float, default=0.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.add_argument("--no-random-command", dest="random_command", action="store_false", help="Disable command sampling and lock to zero commands")
+    parser = argparse.ArgumentParser(description="JackBot Policy Evaluator")
+    parser.add_argument("--model", type=str, required=True, help="Path to trained model checkpoint (.zip)")
+    parser.add_argument("--episodes", type=int, default=5, help="Number of evaluation episodes to run")
+    parser.add_argument("--gui", action="store_true", help="Enable PyBullet 3D visual GUI rendering")
+    parser.add_argument("--no-random-command", dest="random_command", action="store_false", help="Lock commands to zero (disable random command sampling)")
    parser.set_defaults(random_command=True)
-    parser.add_argument("--save-metrics", type=str, default=None, help="Optional path to save JSON metrics report")
+    parser.add_argument("--save-metrics", type=str, default=None, help="Optional JSON path to output evaluation summary")
    args = parser.parse_args()

-    evaluate(
-        model_path=args.model,
-        episodes=args.episodes,
+    print(f"[Eval] Loading policy model from: {args.model}")
+    model = PPO.load(args.model, device="cpu")
+
+    # Instantiate single evaluation environment matching train setup
+    env = JackBotEnv(
        use_gui=args.gui,
-        robot_spacing=args.robot_spacing,
-        start_pose=args.start_pose,
        random_command=args.random_command,
-        save_json=args.save_metrics,
    )

+    episode_rewards = []
+    episode_lengths = []
+    episode_phases = []
+
+    try:
+        for ep in range(args.episodes):
+            obs, _ = env.reset()
+            done = False
+            total_reward = 0.0
+            steps = 0
+            initial_phase = env.curriculum_phase.name
+
+            print(f"\n--- Starting Evaluation Episode {ep + 1}/{args.episodes} [Phase: {initial_phase}] ---")
+
+            while not done:
+                # Deterministic prediction matches evaluation standards
+                action, _ = model.predict(obs, deterministic=True)
+                
+                prev_phase = env.curriculum_phase
+                obs, reward, terminated, truncated, _ = env.step(action)
+                done = terminated or truncated
+
+                total_reward += float(reward)
+                steps += 1
+
+                if env.curriculum_phase != prev_phase:
+                    print(f"  └─ [Eval Milestone] Reached {env.curriculum_phase.name} at step {steps}!")
+
+                if args.gui:
+                    time.sleep(1.0 / 240.0)
+
+            status_str = "FAILED (Terminated)" if terminated else "COMPLETED (Truncated)"
+            final_phase = env.curriculum_phase.name
+            
+            episode_rewards.append(total_reward)
+            episode_lengths.append(steps)
+            episode_phases.append(final_phase)
+
+            print(
+                f"Episode {ep + 1} Finished [{status_str}]: "
+                f"Phase = {final_phase} | Total Reward = {total_reward:.2f} | Steps = {steps}"
+            )
+
+    finally:
+        env.close()
+
+    # Calculate metrics report
+    mean_reward = float(np.mean(episode_rewards))
+    std_reward = float(np.std(episode_rewards))
+    mean_length = float(np.mean(episode_lengths))
+
+    print("\n" + "=" * 60)
+    print(f"EVALUATION COMPLETE ({args.episodes} Episodes)")
+    print(f"Final Reached Phase:  {env.curriculum_phase.name}")
+    print(f"Mean Reward:          {mean_reward:.2f} ± {std_reward:.2f}")
+    print(f"Mean Episode Length:  {mean_length:.1f} steps")
+    print("=" * 60)
+
+    if args.save_metrics:
+        metrics = {
+            "model_path": str(args.model),
+            "episodes_evaluated": args.episodes,
+            "final_curriculum_phase": env.curriculum_phase.name,
+            "mean_reward": mean_reward,
+            "std_reward": std_reward,
+            "mean_episode_length": mean_length,
+            "raw_rewards": episode_rewards,
+            "episode_phases": episode_phases,
+        }
+        out_path = Path(args.save_metrics)
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(out_path, "w") as f:
+            json.dump(metrics, f, indent=4)
+        print(f"[Eval] Saved evaluation report to: {out_path.resolve()}")
+

 if __name__ == "__main__":
    main()