Blame for applications/avoiding/avoiding_baselines/DiRecT/eval.py

import os
import csv
import json
import shutil

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

import utils
import tqdm
import tyro
import gym

import numpy as np
import torch
from dataclasses import dataclass, field

from datasets.sequence import SequenceDataset
from model import TemporalUnet
from obstacles import (
    CircularObstacle, PlanarObstacle,
    PILLARS, NOVEL_CIRCULAR, NOVEL_PLANAR, CONSTRAINT_SETS,
    count_constraints, check_violation,
)
from policy.unconditional import UnconditionalPolicy
from policy.gradient_guidance import GradientGuidancePolicy
from policy.classifier_guidance import ClassifierGuidancePolicy
from policy.projected_diffusion import ProjectedDiffusionPolicy
from policy.value_model import ProxyValueModel
from utils.training import deterministic, save_config, set_cuda_visible_device

import d3il


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

@dataclass
class EvaluationConfig:
    seed: int = 0
    random_repeat: int = 50
    device: str = "cuda:0"
    log_folder: str = "logs"
    exp_name: str = "diffusion"

    render: bool = False

    controller: str = "rh"
    replan_steps: int = 8

    env: str = "avoiding-v0"
    horizon: int = 16
    normalizer: str = "LimitsNormalizer"
    preprocess_fns: list = field(default_factory=lambda: [])
    max_path_length: int = 200
    max_n_episodes: int = 500
    termination_penalty: float = 0
    max_episode_length: int = 100

    state_dim: int = 4
    action_dim: int = 2

    diffusion_exp_name: str = "diffusion"
    diffusion_cp: str = "0"

    batch_size: int = 1
    n_sampling_steps: int = 10

    constraints: bool = False
    obstacle_margin: float = 0.02
    draw_obstacle_margin: bool = False

    # guidance
    guidance_method: str = "no"  # "no", "gradient", "classifier", "projection", "safediffuser", "primal-dual", "augmented-lagrangian", "direct"
    guidance_lr: float = 0.05
    guidance_steps: int = 5
    gradient_start: float = 0.0
    projection_start: float = 0.5

    # value model weights
    value_objective: str = ""  # "" or "distance"
    objective_weight: float = 1.0
    obstacle_weight: float = 100.0
    dynamics_weight: float = 0.0

    # safediffuser CBF
    cbf_algorithm: str = "RoS"  # "RoS", "ReS", "TVS"
    cbf_coefficient: float = 1.0
    guidance_start_fraction: float = 0.0  # fraction of steps before guidance begins
    guidance_last_steps: int = 0  # apply guidance only to last N steps (0 = use guidance_start_fraction)
    cbf_slack_weight: float = 1000.0
    cbf_relaxation_end: float = 0.9
    cbf_sigmoid_bias: float = 0.9
    cbf_sigmoid_input: float = 32.0
    cbf_sigmoid_output: float = 1.0

    # primal-dual / augmented Lagrangian
    pd_learning_rate: float = 0.456
    pd_initial_penalty: float = 2.5e-4
    pd_final_penalty: float = 0.00154
    pd_use_cbfs: bool = True
    pd_cbf_alpha: float = 0.9

    # DiRecT (prediction in clean space)
    pred_control_weight: float = 1.0
    pred_objective: str = ""           # "" or "distance"
    pred_objective_scale: float = 100.0
    max_action_delta: float = 0.01  # trajectory curvature limit (second difference on positions, 0 = disabled)

    # dynamics (always loaded; relaxation controls equality vs inequality tolerance)
    dynamics_relaxation: float = 0.0  # 0 = equality, >0 = inequality with tolerance

    # PID tracking controller
    pd_kp: float = 5.0  # proportional gain on position error
    pd_kd: float = 1.0  # derivative gain on velocity error
    pd_ki: float = 0.0  # integral gain on accumulated position error

    no_images: bool = False
    no_trajectories: bool = False
    no_early_stop: bool = False




# ---------------------------------------------------------------------------
# Batched episode runner
# ---------------------------------------------------------------------------

def run_batched_eval(envs, policy, cfg, normalizer=None, save_path=None):
    """Run random_repeat episodes using len(envs) env slots, batching inference.

    Slots are reused: when an episode finishes, the env is reset for the next
    episode until all random_repeat episodes are complete.
    """
    B = len(envs)  # number of env slots (= batch_size)
    total_episodes = cfg.random_repeat
    constraint_name = "novel" if cfg.constraints else ""

    assert (
        cfg.replan_steps < cfg.horizon
    ), f"replan steps ({cfg.replan_steps}) must be smaller than horizon ({cfg.horizon})"

    # results indexed by episode id
    all_rollouts = [None] * total_episodes
    all_successes = [False] * total_episodes
    all_violations = [0.0] * total_episodes
    all_rewards = [0.0] * total_episodes
    all_comp_times = [[] for _ in range(total_episodes)]
    all_planned_obs = [[] for _ in range(total_episodes)]  # list of predicted obs per replan
    all_plan_feasibility = [[] for _ in range(total_episodes)]  # True/False per replan
    all_ipopt_iters = [[] for _ in range(total_episodes)]  # total IPOPT iters per replan
    all_ipopt_failures = [[] for _ in range(total_episodes)]  # failed IPOPT solves per replan
    all_ipopt_solves = [[] for _ in range(total_episodes)]  # IPOPT solve count per replan

    # per-slot state
    slot_episode = list(range(min(B, total_episodes)))  # which episode each slot runs
    next_episode = len(slot_episode)  # next episode id to assign
    slot_active = [True] * B
    observations = [None] * B
    rollouts = [None] * B
    planned_actions = [None] * B
    planned_obs = [None] * B  # predicted observations for PD tracking
    action_indices = [0] * B
    slot_steps = [0] * B  # track steps per slot for max_episode_length
    integral_error = [np.zeros(2) for _ in range(B)]  # accumulated position error for I term
    last_action_norm = [None] * B  # last executed action in normalized space (for action rate constraint)
    use_pd = cfg.pd_kp > 0 or cfg.pd_kd > 0 or cfg.pd_ki > 0

    # deactivate unused slots if random_repeat < batch_size
    for i in range(B):
        if i >= total_episodes:
            slot_active[i] = False
        else:
            obs = envs[i].reset()
            observations[i] = obs
            rollouts[i] = [obs.copy()]

    episodes_done = 0
    pbar = tqdm.tqdm(desc=f"Eval (0/{total_episodes} episodes done)")
    while episodes_done < total_episodes:
        # 1. find slots that need replanning
        need_replan = [
            i for i in range(B)
            if slot_active[i] and (planned_actions[i] is None or action_indices[i] >= cfg.replan_steps)
        ]

        if need_replan:
            # 2. batch observations through policy
            batch_obs = np.stack([observations[i] for i in need_replan])
            batch_conditions = {0: batch_obs}

            # Pass last executed action for action rate constraint
            policy_kwargs = {}
            if hasattr(policy, 'max_action_delta') and policy.max_action_delta > 0:
                ap = last_action_norm[need_replan[0]]
                if ap is not None:
                    policy_kwargs["a_prev"] = ap

            _, batch_trajs, _, _, info = policy(
                batch_conditions, batch_size=len(need_replan),
                **policy_kwargs,
            )

            for j, i in enumerate(need_replan):
                planned_actions[i] = batch_trajs.actions[j]
                planned_obs[i] = batch_trajs.observations[j]
                action_indices[i] = 0
                integral_error[i] = np.zeros(2)  # reset I term at replan
                ep = slot_episode[i]
                all_planned_obs[ep].append(batch_trajs.observations[j])

                # Check feasibility of planned trajectory
                obs_plan = batch_trajs.observations[j]  # (H, 4)
                n_viol = sum(1 for t in range(len(obs_plan)) if check_violation(obs_plan[t], constraint_name) > 0)
                all_plan_feasibility[ep].append(n_viol == 0)
                if "computation_time" in info:
                    all_comp_times[ep].append(info["computation_time"] / len(need_replan))
                if "ipopt_solves" in info:
                    all_ipopt_iters[ep].append(float(info["ipopt_iters"][j]))
                    all_ipopt_failures[ep].append(float(info["ipopt_failures"][j]))
                    all_ipopt_solves[ep].append(float(info["ipopt_solves"][j]))

        # 3. step each active slot
        for i in range(B):
            if not slot_active[i]:
                continue

            ep = slot_episode[i]
            action = planned_actions[i][action_indices[i]].copy()

            # PID tracking correction (in normalized space for comparable scales)
            if use_pd and planned_obs[i] is not None and normalizer is not None:
                obs_actual_n = normalizer.normalize(
                    observations[i].reshape(1, -1), "observations"
                ).flatten()
                obs_planned_n = normalizer.normalize(
                    planned_obs[i][action_indices[i]].reshape(1, -1), "observations"
                ).flatten()
                action_n = normalizer.normalize(
                    action.reshape(1, -1), "actions"
                ).flatten()
                pos_error = obs_planned_n[2:4] - obs_actual_n[2:4]
                vel_error = obs_planned_n[0:2] - obs_actual_n[0:2]
                integral_error[i] += pos_error
                action_n += (cfg.pd_kp * pos_error
                           + cfg.pd_kd * vel_error
                           + cfg.pd_ki * integral_error[i])
                action = normalizer.unnormalize(
                    action_n.reshape(1, -1), "actions"
                ).flatten()

            action_indices[i] += 1
            slot_steps[i] += 1

            obs, reward, terminated, info_env = envs[i].step(action)
            observations[i] = obs
            rollouts[i].append(obs.copy())
            all_rewards[ep] += reward

            # Track last action in normalized space for action rate constraint
            if normalizer is not None:
                last_action_norm[i] = normalizer.normalize(
                    action.reshape(1, -1), "actions"
                ).flatten()

            success = info_env[1]
            violation = check_violation(obs, constraint_name)
            if violation > 0 or (terminated and not success):
                violation = 1.0
            else:
                violation = 0.0
            all_violations[ep] += violation

            episode_done = False
            if terminated or success:
                episode_done = True
            elif violation > 0 and not cfg.no_early_stop:
                episode_done = True
            elif slot_steps[i] >= cfg.max_episode_length:
                episode_done = True

            if episode_done:
                # save episode results
                all_rollouts[ep] = np.array(rollouts[i])
                all_successes[ep] = success
                episodes_done += 1

                # Print episode feasibility summary
                feas = all_plan_feasibility[ep]
                n_feasible = sum(feas) if feas else 0
                n_total = len(feas) if feas else 0
                feas_str = f"plans: {n_feasible}/{n_total} feasible" if n_total > 0 else ""
                status_str = "OK" if success else "FAIL"
                safe_str = "safe" if all_violations[ep] == 0 else f"violations={int(all_violations[ep])}"
                pbar.write(f"  ep {ep}: {status_str} | {safe_str} | {feas_str} | {slot_steps[i]} steps")

                # Save per-episode prediction overlay plot
                if save_path is not None and not cfg.no_images and all_planned_obs[ep]:
                    pred_dir = os.path.join(save_path, "predictions")
                    os.makedirs(pred_dir, exist_ok=True)
                    ep_plotter = utils.AvoidingTrajectoryPlotter(
                        constraint=("novel" if cfg.constraints else ""),
                        obstacle_margin=cfg.obstacle_margin,
                        draw_obstacle_margin=cfg.draw_obstacle_margin,
                    )
                    ep_fig, ep_ax = ep_plotter.plot_overlaid_predictions_and_real(
                        predicted_trajectories=all_planned_obs[ep],
                        real_trajectory=all_rollouts[ep],
                        title=f"Episode {ep}",
                    )
                    ep_plotter.save_figure(ep_fig, os.path.join(pred_dir, f"ep_{ep:03d}.png"))
                    plt.close(ep_fig)

                pbar.update(1)
                pbar.set_description(f"Eval ({episodes_done}/{total_episodes} episodes done)")

                # reuse slot for next episode
                if next_episode < total_episodes:
                    slot_episode[i] = next_episode
                    next_episode += 1
                    obs = envs[i].reset()
                    observations[i] = obs
                    rollouts[i] = [obs.copy()]
                    planned_actions[i] = None
                    planned_obs[i] = None
                    action_indices[i] = 0
                    integral_error[i] = np.zeros(2)
                    last_action_norm[i] = None
                    slot_steps[i] = 0
                else:
                    slot_active[i] = False

    pbar.close()
    return all_rollouts, all_successes, all_violations, all_rewards, all_comp_times, all_planned_obs, all_plan_feasibility, all_ipopt_iters, all_ipopt_failures, all_ipopt_solves


# ---------------------------------------------------------------------------
# Main evaluation
# ---------------------------------------------------------------------------

def evaluate(cfg: EvaluationConfig):

    # get dataset (only for its normalizer)
    dataset = SequenceDataset(
        env=cfg.env,
        horizon=cfg.horizon,
        normalizer=cfg.normalizer,
        preprocess_fns=cfg.preprocess_fns,
        max_path_length=cfg.max_path_length,
        max_n_episodes=cfg.max_n_episodes,
        termination_penalty=0,
        seed=cfg.seed,
    )
    normalizer = dataset.normalizer

    model = TemporalUnet(
        horizon=cfg.horizon,
        transition_dim=cfg.state_dim + cfg.action_dim,
        cond_dim=cfg.state_dim,
        dim=32,
        dim_mults=(1, 4, 8),
        attention=False,
    ).to(cfg.device)
    model.load_state_dict(
        torch.load(
            os.path.join(
                cfg.log_folder,
                cfg.env,
                "diffusion",
                cfg.diffusion_exp_name,
                f"model_ema_{cfg.diffusion_cp}.pth",
            )
        )
    )
    model.eval()

    # --- load dynamics model (always on) ---
    dynamics_path = os.path.join(cfg.log_folder, cfg.env, "dynamics", "linear_model.npz")
    print(f"Loading dynamics model from: {dynamics_path}")
    dynamics_model = dict(np.load(dynamics_path, allow_pickle=True))

    # --- build policy ---
    if cfg.guidance_method == "gradient":
        value_model = ProxyValueModel(
            horizon=cfg.horizon,
            action_dim=cfg.action_dim,
            state_dim=cfg.state_dim,
            normalizer=normalizer,
            objective=cfg.value_objective,
            constraints=cfg.constraints,
            dynamics=True,
            objective_weight=cfg.objective_weight,
            obstacle_weight=cfg.obstacle_weight,
            dynamics_weight=cfg.dynamics_weight,
            obstacle_margin=cfg.obstacle_margin,
            dynamics_model=dynamics_model,
        ).to(cfg.device)

        policy = GradientGuidancePolicy(
            model=model,
            normalizer=normalizer,
            value_model=value_model,
            action_dim=cfg.action_dim,
            state_dim=cfg.state_dim,
            horizon=cfg.horizon,
            n_sampling_steps=cfg.n_sampling_steps,
            guidance_lr=cfg.guidance_lr,
            guidance_steps=cfg.guidance_steps,
            gradient_start=cfg.gradient_start,
            device=cfg.device,
        )
    elif cfg.guidance_method == "classifier":
        value_model = ProxyValueModel(
            horizon=cfg.horizon,
            action_dim=cfg.action_dim,
            state_dim=cfg.state_dim,
            normalizer=normalizer,
            objective=cfg.value_objective,
            constraints=cfg.constraints,
            dynamics=True,
            objective_weight=cfg.objective_weight,
            obstacle_weight=cfg.obstacle_weight,
            dynamics_weight=cfg.dynamics_weight,
            obstacle_margin=cfg.obstacle_margin,
            dynamics_model=dynamics_model,
        ).to(cfg.device)

        policy = ClassifierGuidancePolicy(
            model=model,
            normalizer=normalizer,
            value_model=value_model,
            action_dim=cfg.action_dim,
            state_dim=cfg.state_dim,
            horizon=cfg.horizon,
            n_sampling_steps=cfg.n_sampling_steps,
            guidance_lr=cfg.guidance_lr,
            gradient_start=cfg.gradient_start,
            device=cfg.device,
        )
    elif cfg.guidance_method == "projection":
        from projector import IpoptProjector
        constraint_name = "novel" if cfg.constraints else ""
        circular_obs = CONSTRAINT_SETS[constraint_name]["circular"] if constraint_name else []
        planar_obs = CONSTRAINT_SETS[constraint_name]["planar"] if constraint_name else []

        projector = IpoptProjector(
            circular_obstacles=circular_obs,
            planar_obstacles=planar_obs,
            horizon=cfg.horizon,
            action_dim=cfg.action_dim,
            state_dim=cfg.state_dim,
            normalizer=normalizer,
            dynamics=dynamics_model,
            dynamics_relaxation=cfg.dynamics_relaxation,
            obstacle_margin=cfg.obstacle_margin,
            max_action_delta=cfg.max_action_delta,
        )

        policy = ProjectedDiffusionPolicy(
            model=model,
            normalizer=normalizer,
            projector=projector,
            action_dim=cfg.action_dim,
            state_dim=cfg.state_dim,
            horizon=cfg.horizon,
            n_sampling_steps=cfg.n_sampling_steps,
            projection_start=cfg.projection_start,
            n_projection_workers=cfg.batch_size,
            device=cfg.device,
        )
    elif cfg.guidance_method == "safediffuser":
        from policy.safediffuser import SafeDiffuserPolicy

        policy = SafeDiffuserPolicy(
            model=model,
            normalizer=normalizer,
            action_dim=cfg.action_dim,
            state_dim=cfg.state_dim,
            horizon=cfg.horizon,
            n_sampling_steps=cfg.n_sampling_steps,
            algorithm_type=cfg.cbf_algorithm,
            cbf_coefficient=cfg.cbf_coefficient,
            guidance_start=cfg.guidance_start_fraction,
            guidance_last_steps=cfg.guidance_last_steps,
            obstacle_margin=cfg.obstacle_margin,
            slack_weight=cfg.cbf_slack_weight,
            relaxation_end_fraction=cfg.cbf_relaxation_end,
            sigmoid_bias_fraction=cfg.cbf_sigmoid_bias,
            sigmoid_input_scale=cfg.cbf_sigmoid_input,
            sigmoid_output_scale=cfg.cbf_sigmoid_output,
            dynamics_model=dynamics_model,
            device=cfg.device,
        )
    elif cfg.guidance_method in ("primal-dual", "augmented-lagrangian"):
        from policy.augmented_lagrangian import AugmentedLagrangianPolicy

        policy = AugmentedLagrangianPolicy(
            model=model,
            normalizer=normalizer,
            action_dim=cfg.action_dim,
            state_dim=cfg.state_dim,
            horizon=cfg.horizon,
            n_sampling_steps=cfg.n_sampling_steps,
            algorithm_type=cfg.guidance_method,
            guidance_start=cfg.guidance_start_fraction,
            guidance_last_steps=cfg.guidance_last_steps,
            obstacle_margin=cfg.obstacle_margin,
            learning_rate=cfg.pd_learning_rate,
            initial_penalty=cfg.pd_initial_penalty,
            final_penalty=cfg.pd_final_penalty,
            use_cbfs=cfg.pd_use_cbfs,
            cbf_alpha=cfg.pd_cbf_alpha,
            dynamics_model=dynamics_model,
            device=cfg.device,
        )
    elif cfg.guidance_method == "direct":
        from policy.direct import DiReCTPolicy

        policy = DiReCTPolicy(
            model=model,
            normalizer=normalizer,
            action_dim=cfg.action_dim,
            state_dim=cfg.state_dim,
            horizon=cfg.horizon,
            n_sampling_steps=cfg.n_sampling_steps,
            guidance_start=cfg.guidance_start_fraction,
            guidance_last_steps=cfg.guidance_last_steps,
            obstacle_margin=cfg.obstacle_margin,
            control_penalty_weight=cfg.pred_control_weight,
            prediction_objective=cfg.pred_objective,
            prediction_objective_scale=cfg.pred_objective_scale,
            max_action_delta=cfg.max_action_delta,
            dynamics_model=dynamics_model,
            dynamics_relaxation=cfg.dynamics_relaxation,
            device=cfg.device,
        )
    else:
        policy = UnconditionalPolicy(
            model=model,
            normalizer=normalizer,
            action_dim=cfg.action_dim,
            state_dim=cfg.state_dim,
            horizon=cfg.horizon,
            n_sampling_steps=cfg.n_sampling_steps,
            device=cfg.device,
        )

    # create batch_size env instances (reused across episodes)
    n_envs = min(cfg.batch_size, cfg.random_repeat)
    print(f"Creating {n_envs} environment instances for {cfg.random_repeat} episodes...")
    envs = []
    for _ in range(n_envs):
        env = gym.make("avoiding-v0", render=cfg.render)
        env.set_seed(cfg.seed)
        env.start()
        envs.append(env)

    save_path = os.path.join(cfg.log_folder, cfg.env, "eval", cfg.exp_name)
    if os.path.exists(save_path):
        shutil.rmtree(save_path)
    os.makedirs(save_path)

    # run batched evaluation
    (
        real_trajectories,
        successes,
        total_violations,
        total_rewards,
        computation_times,
        planned_observations,
        all_plan_feasibility,
        ipopt_iters,
        ipopt_failures,
        ipopt_solves,
    ) = run_batched_eval(envs, policy, cfg, normalizer=normalizer, save_path=save_path)

    # --- per-run metrics ---
    n_constraints = count_constraints(("novel" if cfg.constraints else ""))
    trajectory_data = []

    for i in range(cfg.random_repeat):
        steps = len(real_trajectories[i]) - 1
        safety = total_violations[i] == 0
        comp_times = computation_times[i]
        n_solves = sum(ipopt_solves[i])

        trajectory_data.append({
            "run_id": i,
            "steps": steps,
            "total_violations": float(total_violations[i]),
            "total_rewards": float(total_rewards[i]),
            "success": bool(successes[i]),
            "safety": bool(safety),
            "n_constraints": n_constraints,
            "constraint_set": ("novel" if cfg.constraints else ""),
            "average_computation_time": (
                float(np.mean(comp_times)) if comp_times else None
            ),
            "avg_ipopt_iters": (
                sum(ipopt_iters[i]) / n_solves if n_solves > 0 else None
            ),
            "ipopt_failure_rate": (
                sum(ipopt_failures[i]) / n_solves if n_solves > 0 else None
            ),
        })

    # --- always save per-run metrics CSV ---
    csv_path = os.path.join(save_path, "runs.csv")
    with open(csv_path, "w", newline="") as csvfile:
        fieldnames = list(trajectory_data[0].keys())
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        for data in trajectory_data:
            writer.writerow(data)
    print(f"Saved per-run metrics to: {csv_path}")

    # --- always save experiment summary ---
    n_runs = len(trajectory_data)
    n_success = sum(1 for d in trajectory_data if d["success"])
    n_safe = sum(1 for d in trajectory_data if d["safety"])
    avg_steps = np.mean([d["steps"] for d in trajectory_data])
    avg_violations = np.mean([d["total_violations"] for d in trajectory_data])
    avg_rewards = np.mean([d["total_rewards"] for d in trajectory_data])
    comp_times_all = [d["average_computation_time"] for d in trajectory_data if d["average_computation_time"] is not None]

    # global IPOPT statistics across all episodes (one "projection" = one solve)
    total_ipopt_solves = sum(sum(s) for s in ipopt_solves)
    total_ipopt_iters = sum(sum(s) for s in ipopt_iters)
    total_ipopt_failures = sum(sum(s) for s in ipopt_failures)

    summary = {
        "exp_name": cfg.exp_name,
        "constraint_set": ("novel" if cfg.constraints else ""),
        "n_constraints": n_constraints,
        "n_runs": n_runs,
        "success_rate": n_success / n_runs,
        "safety_rate": n_safe / n_runs,
        "avg_steps": float(avg_steps),
        "avg_violations": float(avg_violations),
        "avg_rewards": float(avg_rewards),
        "avg_computation_time": float(np.mean(comp_times_all)) if comp_times_all else None,
        "avg_ipopt_iters": (
            total_ipopt_iters / total_ipopt_solves if total_ipopt_solves > 0 else None
        ),
        "ipopt_failure_rate": (
            total_ipopt_failures / total_ipopt_solves if total_ipopt_solves > 0 else None
        ),
    }

    summary_path = os.path.join(save_path, "summary.json")
    with open(summary_path, "w") as f:
        json.dump(summary, f, indent=2)
    print(f"Saved experiment summary to: {summary_path}")

    # --- always save all-trajectories overlay ---
    plotter = utils.AvoidingTrajectoryPlotter(
        constraint=("novel" if cfg.constraints else ""),
        obstacle_margin=cfg.obstacle_margin,
        draw_obstacle_margin=cfg.draw_obstacle_margin,
    )
    fig, ax = plotter.plot_multiple_real_trajectories(real_trajectories)
    plotter.save_figure(fig, os.path.join(save_path, "all_real.png"))

    # --- save trajectory data (unless --no-trajectories) ---
    if not cfg.no_trajectories:
        traj_path = os.path.join(save_path, "trajectories.npz")
        np.savez(
            traj_path,
            **{f"run_{i}": traj for i, traj in enumerate(real_trajectories)},
        )
        print(f"Saved trajectory data to: {traj_path}")

    # print summary
    print(f"\n{'='*50}")
    print(f"Success rate: {summary['success_rate']:.1%} ({n_success}/{n_runs})")
    print(f"Safety rate:  {summary['safety_rate']:.1%} ({n_safe}/{n_runs})")
    print(f"Avg steps:    {summary['avg_steps']:.1f}")
    if summary["avg_ipopt_iters"] is not None:
        print(f"Avg IPOPT iters/projection: {summary['avg_ipopt_iters']:.1f}")
        print(f"IPOPT failure rate: {summary['ipopt_failure_rate']:.1%} ({int(total_ipopt_failures)}/{int(total_ipopt_solves)})")

    # Plan feasibility summary
    total_plans = sum(len(f) for f in all_plan_feasibility)
    feasible_plans = sum(sum(f) for f in all_plan_feasibility)
    if total_plans > 0:
        print(f"Plan feasibility: {feasible_plans}/{total_plans} ({feasible_plans/total_plans:.1%})")
        for ep_i, feas in enumerate(all_plan_feasibility):
            if feas and not all(feas):
                n_infeas = sum(1 for f in feas if not f)
                print(f"  ep {ep_i}: {n_infeas}/{len(feas)} infeasible plans")
    print(f"{'='*50}")


if __name__ == "__main__":
    cfg = tyro.cli(EvaluationConfig)
    set_cuda_visible_device(cfg)
    deterministic(cfg.seed)

    log_subfolder = os.path.join(cfg.log_folder, cfg.env, "eval", cfg.exp_name)
    save_config(cfg, log_subfolder)

    evaluate(cfg)