JackM323
acb3d671be
new reward/penalty system learning phases with curriculum learning new training parameters cleanup of old code better logging while training multiple environments instead of robots (they could bumb into each other)
92 lines
2.8 KiB
Python
92 lines
2.8 KiB
Python
"""
|
|
ml/evaluate.py - Evaluation routine for trained JackBot PPO policies.
|
|
"""
|
|
import argparse
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Optional, Dict, List, Any
|
|
import numpy as np
|
|
|
|
|
|
def evaluate(
|
|
model_path: str,
|
|
episodes: int = 5,
|
|
use_gui: bool = True,
|
|
robot_spacing: float = 0.5,
|
|
start_pose: str = "init_deg",
|
|
random_command: bool = False,
|
|
save_json: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
try:
|
|
from stable_baselines3 import PPO
|
|
except ImportError as exc:
|
|
raise ImportError(
|
|
"stable-baselines3 is required for evaluation. Install with: pip install stable-baselines3"
|
|
) from exc
|
|
|
|
from .env import JackBotEnv
|
|
|
|
print(f"[Eval] Loading policy model from: {model_path}")
|
|
# Force device="cpu" to prevent AMD ROCm/hipBLASLt matrix multiplication crashes
|
|
model = PPO.load(model_path, device="cpu")
|
|
|
|
# Initialize standard environment (returns single array of shape (352,))
|
|
env = JackBotEnv(
|
|
use_gui=use_gui,
|
|
random_command=random_command,
|
|
robot_spacing=robot_spacing,
|
|
start_pose=start_pose,
|
|
)
|
|
|
|
episode_rewards: List[float] = []
|
|
episode_lengths: List[int] = []
|
|
|
|
for ep in range(episodes):
|
|
obs, _ = env.reset()
|
|
done = False
|
|
total_reward = 0.0
|
|
steps = 0
|
|
|
|
print(f"\n--- Starting Evaluation Episode {ep + 1}/{episodes} ---")
|
|
|
|
while not done:
|
|
action, _ = model.predict(obs, deterministic=True)
|
|
obs, reward, terminated, truncated, _ = env.step(action)
|
|
|
|
done = terminated or truncated
|
|
total_reward += float(reward)
|
|
steps += 1
|
|
|
|
if use_gui:
|
|
time.sleep(1.0 / 240.0)
|
|
|
|
episode_rewards.append(total_reward)
|
|
episode_lengths.append(steps)
|
|
print(f"Episode {ep + 1} Finished: Total Reward = {total_reward:.2f} | Steps = {steps}")
|
|
|
|
env.close()
|
|
|
|
metrics = {
|
|
"model_path": str(model_path),
|
|
"episodes_evaluated": episodes,
|
|
"mean_reward": float(np.mean(episode_rewards)),
|
|
"std_reward": float(np.std(episode_rewards)),
|
|
"mean_episode_length": float(np.mean(episode_lengths)),
|
|
"raw_rewards": episode_rewards,
|
|
}
|
|
|
|
print("\n" + "=" * 50)
|
|
print(f"EVALUATION COMPLETE ({episodes} Episodes)")
|
|
print(f"Mean Reward: {metrics['mean_reward']:.2f} ± {metrics['std_reward']:.2f}")
|
|
print(f"Mean Episode Length: {metrics['mean_episode_length']:.1f} steps")
|
|
print("=" * 50)
|
|
|
|
if save_json:
|
|
out_path = Path(save_json)
|
|
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 metrics to: {out_path.resolve()}")
|
|
|
|
return metrics |