eval_sb3.py

"""
PART 2 - Evaluation script.

TASK 6: evaluate a trained model on source/target environments.
"""

import argparse
import csv
import os
import time
from datetime import datetime
from pathlib import Path

import gymnasium as gym
import numpy as np
import panda_gym  # noqa: F401
from stable_baselines3 import DDPG, PPO, SAC
from rand_wrapper import RandomizationWrapper


def evaluate(
    model_path: str,
    n_episodes: int,
    deterministic: bool,
    render: bool,
    env_type: str,
    source_mass: float,
    target_mass: float,
    results_dir: str,
) -> None:
    if not os.path.exists(model_path):
        raise FileNotFoundError(
            f"Model file not found: {model_path}. "
            "Make sure you saved your trained model with model.save(...)."
        )

    render_mode = "human" if render else "rgb_array"
    env = gym.make("PandaPush-v3", render_mode=render_mode, reward_type="dense")
    
    # === CRITICAL: ENVIROMENT TYPE AYARI (Sim-to-Real Farkı) ===
    if env_type == "source":
        eval_mass = source_mass
    elif env_type == "target":
        eval_mass = target_mass
    else:
        raise ValueError(f"Unknown environment type: {env_type}")

    env = RandomizationWrapper(env, mass_range=(eval_mass, eval_mass), mode="none")

    if "ppo" in model_path.lower():
        model = PPO.load(model_path, env=env)
    elif "sac" in model_path.lower():
        model = SAC.load(model_path, env=env)
    else:
        model = DDPG.load(model_path, env=env)

    episode_returns = []
    successes = []
    episode_rows = []

    for episode in range(1, n_episodes + 1):
        obs, info = env.reset()
        terminated = False
        truncated = False
        episode_return = 0.0
        episode_length = 0

        while not (terminated or truncated):
            action, _ = model.predict(obs, deterministic=deterministic)
            obs, reward, terminated, truncated, info = env.step(action)
            episode_return += float(reward)
            episode_length += 1

            if render:
                time.sleep(0.05)

        episode_returns.append(episode_return)

        success = float(info.get("is_success", 0.0)) if isinstance(info, dict) else 0.0
        successes.append(success)
        episode_rows.append(
            {
                "episode": episode,
                "return": episode_return,
                "length": episode_length,
                "success": success,
            }
        )

        print(f"Episode {episode:03d} | return = {episode_return:.3f}")

    env.close()

    returns = np.array(episode_returns, dtype=np.float32)
    print("\n=== Evaluation summary ===")
    print(f"Environment type: {env_type}")
    print(f"Object mass: {eval_mass:.3f} kg")
    print(f"Model path: {model_path}")
    print(f"Episodes: {n_episodes}")
    print(f"Mean return: {returns.mean():.3f}")
    print(f"Std return:  {returns.std():.3f}")
    print(f"Min return:  {returns.min():.3f}")
    print(f"Max return:  {returns.max():.3f}")

    if successes:
        success_rate = float(np.mean(successes))
        print(f"Success rate: {success_rate:.2%}")

    save_evaluation_results(
        results_dir=results_dir,
        model_path=model_path,
        env_type=env_type,
        eval_mass=eval_mass,
        deterministic=deterministic,
        episode_rows=episode_rows,
        returns=returns,
        success_rate=success_rate,
    )


def save_evaluation_results(
    results_dir: str,
    model_path: str,
    env_type: str,
    eval_mass: float,
    deterministic: bool,
    episode_rows: list[dict],
    returns: np.ndarray,
    success_rate: float,
) -> None:
    output_dir = Path(results_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    model_name = Path(model_path).stem
    run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    base_name = f"{model_name}_{env_type}_{run_id}"

    episode_path = output_dir / f"{base_name}_episodes.csv"
    with episode_path.open("w", newline="", encoding="utf-8") as file:
        writer = csv.DictWriter(
            file,
            fieldnames=["episode", "return", "length", "success"],
        )
        writer.writeheader()
        writer.writerows(episode_rows)

    summary_path = output_dir / "evaluation_runs.csv"
    write_header = not summary_path.exists()
    with summary_path.open("a", newline="", encoding="utf-8") as file:
        writer = csv.DictWriter(
            file,
            fieldnames=[
                "run_id",
                "model",
                "env_type",
                "eval_mass_kg",
                "episodes",
                "deterministic",
                "mean_return",
                "std_return",
                "min_return",
                "max_return",
                "success_rate",
                "episode_file",
            ],
        )
        if write_header:
            writer.writeheader()
        writer.writerow(
            {
                "run_id": run_id,
                "model": model_name,
                "env_type": env_type,
                "eval_mass_kg": eval_mass,
                "episodes": len(episode_rows),
                "deterministic": deterministic,
                "mean_return": float(returns.mean()),
                "std_return": float(returns.std()),
                "min_return": float(returns.min()),
                "max_return": float(returns.max()),
                "success_rate": success_rate,
                "episode_file": episode_path.name,
            }
        )

    print(f"Episode data saved as: {episode_path}")
    print(f"Evaluation summary updated: {summary_path}")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Evaluate RL model on PandaPush-v3")
    parser.add_argument(
        "--model-path",
        type=str,
        required=True,
        help="Path to a model zip file",
    )
    parser.add_argument(
        "--episodes",
        type=int,
        default=50,
        help="Number of evaluation episodes",
    )
    parser.add_argument(
        "--stochastic",
        action="store_true",
        help="Use stochastic policy sampling instead of deterministic actions",
    )
    parser.add_argument(
        "--render",
        action="store_true",
        help="Render with a window, render_mode='human'",
    )
    parser.add_argument(
        "--env-type",
        type=str,
        default="target",
        choices=["source", "target"],
        help="Type of environment to evaluate on",
    )
    parser.add_argument(
        "--source-mass",
        type=float,
        default=1.0,
        help="Object mass used for source-domain evaluation",
    )
    parser.add_argument(
        "--target-mass",
        type=float,
        default=5.0,
        help="Object mass used for target-domain evaluation",
    )
    parser.add_argument(
        "--results-dir",
        type=str,
        default="results/evaluations",
        help="Directory used to save per-episode and summary CSV files",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    evaluate(
        model_path=args.model_path,
        n_episodes=args.episodes,
        deterministic=not args.stochastic,
        render=args.render,
        env_type=args.env_type,
        source_mass=args.source_mass,
        target_mass=args.target_mass,
        results_dir=args.results_dir,
    )