File size: 6,058 Bytes
6953bd1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | """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()
# trap layout: seed 1001 (second single-agent layout, trap)
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)
# 4-agent layout: seed 2001 (second 4-agent layout, trap)
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)
# bar chart from results.json
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)
# training curve (concatenated curves passed in --curve if present)
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()
|