plot_results.py

"""Plot training curves from runs stored under results/.

Handles both Ray Tune output (results/<run-name>/<trial-dir>/progress.csv)
and the legacy flat CSVs from the pre-Tune version of the training script.

Run with:  uv run plot_results.py                                # all runs
           uv run plot_results.py results/<run>/<trial>/progress.csv
"""

import argparse
from pathlib import Path

import matplotlib

matplotlib.use("Agg")  # no GUI needed; we save to a file
import matplotlib.pyplot as plt
import pandas as pd

RESULTS_DIR = Path(__file__).parent / "results"


def load_run(csv_path: Path) -> tuple[str, pd.DataFrame]:
    """Return (run label, dataframe with normalized column names)."""
    df = pd.read_csv(csv_path)

    if "env_runners/episode_return_mean" in df.columns:
        # Ray Tune progress.csv: label by experiment dir (results/<name>/...).
        label = csv_path.parent.parent.name
        normalized = pd.DataFrame(
            {
                "env_steps_total": df["num_env_steps_sampled_lifetime"],
                "episode_return_mean": df["env_runners/episode_return_mean"],
                "episode_len_mean": df["env_runners/episode_len_mean"],
            }
        )
    else:
        label = csv_path.stem
        normalized = df

    return label, normalized


def find_run_csvs() -> list[Path]:
    tune_csvs = sorted(RESULTS_DIR.glob("*/*/progress.csv"))
    legacy_csvs = sorted(RESULTS_DIR.glob("*.csv"))
    return legacy_csvs + tune_csvs


def main(csv_paths: list[Path], output: Path, hlines: list[str]) -> None:
    fig, (ax_return, ax_len) = plt.subplots(1, 2, figsize=(12, 4.5))

    for csv_path in csv_paths:
        label, df = load_run(csv_path)
        ax_return.plot(df["env_steps_total"], df["episode_return_mean"], marker="o", label=label)
        ax_len.plot(df["env_steps_total"], df["episode_len_mean"], marker="o", label=label)

    for hline in hlines:
        label, _, value = hline.rpartition("=")
        ax_return.axhline(float(value), color="gray", linestyle="--", linewidth=1, label=label)
    ax_return.set_xlabel("Environment steps")
    ax_return.set_ylabel("Mean episode return")
    ax_return.set_title("Episode return")
    ax_return.legend()
    ax_return.grid(alpha=0.3)

    ax_len.set_xlabel("Environment steps")
    ax_len.set_ylabel("Mean episode length")
    ax_len.set_title("Episode length")
    ax_len.legend()
    ax_len.grid(alpha=0.3)

    fig.tight_layout()
    fig.savefig(output, dpi=150)
    print(f"Saved plot to {output}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "csvs",
        nargs="*",
        type=Path,
        help="CSV files to plot (default: all runs found under results/)",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=RESULTS_DIR / "training_curves.png",
        help="Output image path",
    )
    parser.add_argument(
        "--hlines",
        nargs="*",
        default=["solved (475)=475"],
        help='Reference lines on the return plot, as "label=value"',
    )
    args = parser.parse_args()

    csv_paths = args.csvs or find_run_csvs()
    if not csv_paths:
        raise SystemExit(f"No run CSVs found under {RESULTS_DIR}. Run train_cartpole.py first.")
    main(csv_paths, args.output, args.hlines)