| """Plots from results.json + live rollouts of the trained policy. |
| |
| Outputs (plots/): |
| trap_trajectories.png - one wall-trap layout, RL vs classical |
| four_agent.png - one 4-agent layout, RL vs classical (2 panels) |
| clean_comparison.png - bar chart: clean success / collision, RL vs |
| classical, single-agent and 4-agent |
| training_curve.png - training success/collision over iterations |
| """ |
| import argparse |
| import json |
|
|
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
| import torch |
|
|
| from sim import generate_layout, NavSim |
| from policies import ActorCritic, classical_action, rl_action |
|
|
|
|
| def draw(ax, layout, traj, goals, title): |
| ax.set_xlim(0, 12); ax.set_ylim(0, 12) |
| ax.set_aspect("equal"); ax.set_title(title, fontsize=10) |
| ax.set_xlabel("x"); ax.set_ylabel("y") |
| for (cx, cy, r) in layout["circles"]: |
| ax.add_patch(plt.Circle((cx, cy), r, color="0.7", zorder=1)) |
| ax.add_patch(plt.Circle((cx, cy), r, fill=False, color="k", lw=0.8, zorder=2)) |
| for (x1, y1, x2, y2) in layout["segments"]: |
| ax.plot([x1, x2], [y1, y2], color="k", lw=4, solid_capstyle="butt", zorder=2) |
| colors = ["#1f77b4", "#d62728", "#2ca02c", "#9467bd"] |
| for i, tr in enumerate(traj): |
| tr = np.array(tr) |
| ax.plot(tr[:, 0], tr[:, 1], color=colors[i % 4], lw=1.5, zorder=3) |
| ax.scatter([tr[0, 0]], [tr[0, 1]], marker="o", s=45, color=colors[i % 4], zorder=4) |
| ax.scatter([goals[i, 0]], [goals[i, 1]], marker="*", s=160, color=colors[i % 4], |
| edgecolor="k", linewidth=0.5, zorder=4) |
| ax.add_patch(plt.Rectangle((0, 0), 12, 12, fill=False, color="k", lw=1.2)) |
|
|
|
|
| def rollout(policy, net, layout, n_agents): |
| sim = NavSim(n_agents, seed=0) |
| obs = sim.reset(layout) |
| while True: |
| act = classical_action(obs) if policy == "classical" else rl_action(net, obs, deterministic=True) |
| obs, _, done, _ = sim.step(act) |
| if done: |
| break |
| return sim.traj, sim.success, sim.collision |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--weights", default="ppo_policy.pt") |
| ap.add_argument("--results", default="results.json") |
| ap.add_argument("--curve", default="training_curve.csv") |
| ap.add_argument("--outdir", default="plots") |
| args = ap.parse_args() |
| import os |
| os.makedirs(args.outdir, exist_ok=True) |
| net = ActorCritic() |
| net.load_state_dict(torch.load(args.weights, map_location="cpu", weights_only=True)) |
| net.eval() |
|
|
| |
| rng = np.random.default_rng(1001) |
| trap_layout = generate_layout(rng, 1, trap=True) |
| fig, axes = plt.subplots(1, 2, figsize=(10, 5)) |
| for ax, policy in zip(axes, ["classical", "rl"]): |
| traj, succ, coll = rollout(policy, net, trap_layout, 1) |
| draw(ax, trap_layout, traj, trap_layout["goals"], |
| f"{policy}: {'SUCCESS' if succ[0] else 'COLLISION' if coll[0] else 'TIMEOUT'}") |
| fig.tight_layout() |
| fig.savefig(f"{args.outdir}/trap_trajectories.png", dpi=150) |
| plt.close(fig) |
|
|
| |
| rng = np.random.default_rng(2001) |
| lay4 = generate_layout(rng, 4, trap=True) |
| fig, axes = plt.subplots(1, 2, figsize=(10, 5)) |
| for ax, policy in zip(axes, ["classical", "rl"]): |
| traj, succ, coll = rollout(policy, net, lay4, 4) |
| draw(ax, lay4, traj, lay4["goals"], |
| f"{policy}: {int(succ.sum())}/4 success, {int(coll.sum())} collisions") |
| fig.tight_layout() |
| fig.savefig(f"{args.outdir}/four_agent.png", dpi=150) |
| plt.close(fig) |
|
|
| |
| with open(args.results) as f: |
| res = json.load(f) |
| clean = res["conditions"]["clean"] |
| labels = ["1-agent\nclassical", "1-agent\nRL", "4-agent\nclassical", "4-agent\nRL"] |
| succ = [clean["1_agent"]["classical"]["success_rate"], |
| clean["1_agent"]["rl"]["success_rate"], |
| clean["4_agent"]["classical"]["success_rate"], |
| clean["4_agent"]["rl"]["success_rate"]] |
| coll = [clean["1_agent"]["classical"]["collision_rate"], |
| clean["1_agent"]["rl"]["collision_rate"], |
| clean["4_agent"]["classical"]["collision_rate"], |
| clean["4_agent"]["rl"]["collision_rate"]] |
| x = np.arange(len(labels)); w = 0.35 |
| fig, ax = plt.subplots(figsize=(8, 4.5)) |
| ax.bar(x - w / 2, succ, w, label="success rate") |
| ax.bar(x + w / 2, coll, w, label="collision rate") |
| ax.set_xticks(x); ax.set_xticklabels(labels) |
| ax.set_ylim(0, 1.05); ax.set_ylabel("rate"); ax.set_title("Clean evaluation: RL vs classical") |
| ax.legend() |
| for i, v in enumerate(succ): |
| ax.text(i - w / 2, v + 0.02, f"{v:.2f}", ha="center", fontsize=9) |
| for i, v in enumerate(coll): |
| ax.text(i + w / 2, v + 0.02, f"{v:.2f}", ha="center", fontsize=9) |
| fig.tight_layout() |
| fig.savefig(f"{args.outdir}/clean_comparison.png", dpi=150) |
| plt.close(fig) |
|
|
| |
| try: |
| import csv |
| iters, s, c = [], [], [] |
| with open(args.curve) as f: |
| for row in csv.DictReader(f): |
| if row["success_rate_finished"] and row["success_rate_finished"] != "nan": |
| iters.append(int(row["iteration"])); s.append(float(row["success_rate_finished"])) |
| c.append(float(row["collision_rate_finished"])) |
| fig, ax = plt.subplots(figsize=(8, 4)) |
| ax.plot(iters, s, label="training success rate (finished episodes)") |
| ax.plot(iters, c, label="training collision rate") |
| ax.set_xlabel("iteration"); ax.set_ylabel("rate"); ax.legend() |
| ax.set_title("PPO training curve (last training run only)") |
| fig.tight_layout() |
| fig.savefig(f"{args.outdir}/training_curve.png", dpi=150) |
| plt.close(fig) |
| except FileNotFoundError: |
| pass |
| print("plots written to", args.outdir) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
