Files
JackBot/ml/run_eval.py
T
JackM323 c93c524a10 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

113 lines
4.1 KiB
Python

"""Run a trained policy in the PyBullet sim with full curriculum progression.
Usage:
python ml/run_eval.py --model ml/checkpoints/ppo_joint_command.zip --episodes 5 --gui
"""
import argparse
import json
import time
import sys
from pathlib import Path
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 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 JSON path to output evaluation summary")
args = parser.parse_args()
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,
random_command=args.random_command,
)
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()