Sheep training flock of 10 fix?
This commit is contained in:
Johnny Fernandes
@@ -0,0 +1,316 @@
|
||||
"""
|
||||
Single-episode visualization for the herding policy.
|
||||
|
||||
Outputs (all saved to --out-dir):
|
||||
trajectory.png — full field view: dog path + every sheep path
|
||||
timeseries.png — radius, per-sheep pen distance, action magnitude, reward
|
||||
episode.gif — animated replay (slow enough to read)
|
||||
|
||||
Run with no model to watch a RANDOM policy (useful baseline):
|
||||
python visualize.py --random --n-sheep 3 --out-dir vis_random/
|
||||
|
||||
Usage:
|
||||
python visualize.py \\
|
||||
--model runs/ppo_consolidation/final_model.zip \\
|
||||
--vecnorm runs/ppo_consolidation/vecnorm.pkl \\
|
||||
--n-sheep 3 --out-dir vis_out/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import math
|
||||
import numpy as np
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.patches as mpatches
|
||||
import matplotlib.animation as animation
|
||||
from matplotlib.collections import LineCollection
|
||||
from stable_baselines3 import PPO
|
||||
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
|
||||
from herding_env import HerdingEnv
|
||||
|
||||
|
||||
# ── colours ──────────────────────────────────────────────────────────────────
|
||||
SHEEP_COLORS = [
|
||||
"#e41a1c", "#377eb8", "#4daf4a", "#984ea3", "#ff7f00",
|
||||
"#a65628", "#f781bf", "#999999", "#66c2a5", "#fc8d62",
|
||||
]
|
||||
DOG_COLOR = "#4e342e"
|
||||
PEN_COLOR = "#ffe082"
|
||||
FIELD_COLOR = "#dcedc8"
|
||||
|
||||
|
||||
def make_env(n_sheep, max_steps, seed=42):
|
||||
def _init():
|
||||
env = HerdingEnv(n_sheep=n_sheep, max_steps=max_steps)
|
||||
env.reset(seed=seed)
|
||||
return env
|
||||
return _init
|
||||
|
||||
|
||||
def run_episode(model, env, n_sheep, max_steps):
|
||||
"""Run one deterministic episode; return recorded history."""
|
||||
obs = env.reset()
|
||||
inner = env.envs[0]
|
||||
done = False
|
||||
|
||||
dog_xs, dog_ys = [], []
|
||||
sheep_xs = [[] for _ in range(n_sheep)]
|
||||
sheep_ys = [[] for _ in range(n_sheep)]
|
||||
radii = []
|
||||
pen_dists = [[] for _ in range(n_sheep)]
|
||||
action_mags = []
|
||||
rewards = []
|
||||
penned_at = [None] * n_sheep # step when each sheep was penned
|
||||
|
||||
step = 0
|
||||
while not done:
|
||||
if model is None:
|
||||
action = env.action_space.sample()[np.newaxis]
|
||||
else:
|
||||
action, _ = model.predict(obs, deterministic=True)
|
||||
|
||||
obs, reward, dones, infos = env.step(action)
|
||||
done = dones[0]
|
||||
step += 1
|
||||
|
||||
dx, dy = float(inner.dog_pos[0]), float(inner.dog_pos[1])
|
||||
dog_xs.append(dx); dog_ys.append(dy)
|
||||
|
||||
com, radius, _ = inner._flock_stats()
|
||||
radii.append(radius)
|
||||
rewards.append(float(reward[0]))
|
||||
|
||||
act = action[0]
|
||||
action_mags.append(float(np.linalg.norm(act)))
|
||||
|
||||
for i in range(n_sheep):
|
||||
sx, sy = float(inner.sheep_pos[i][0]), float(inner.sheep_pos[i][1])
|
||||
sheep_xs[i].append(sx)
|
||||
sheep_ys[i].append(sy)
|
||||
pen_dists[i].append(float(np.linalg.norm(inner.sheep_pos[i] - inner.PEN_CENTER)))
|
||||
if inner.penned[i] and penned_at[i] is None:
|
||||
penned_at[i] = step
|
||||
|
||||
info = infos[0]
|
||||
n_penned = info.get("n_penned", 0)
|
||||
success = n_penned == n_sheep
|
||||
|
||||
return dict(
|
||||
dog_xs=dog_xs, dog_ys=dog_ys,
|
||||
sheep_xs=sheep_xs, sheep_ys=sheep_ys,
|
||||
radii=radii, pen_dists=pen_dists,
|
||||
action_mags=action_mags, rewards=rewards,
|
||||
penned_at=penned_at,
|
||||
n_penned=n_penned, n_sheep=n_sheep,
|
||||
success=success, steps=step,
|
||||
)
|
||||
|
||||
|
||||
# ── plot helpers ─────────────────────────────────────────────────────────────
|
||||
|
||||
def draw_field(ax):
|
||||
ax.set_xlim(-16, 16); ax.set_ylim(-16, 16)
|
||||
ax.set_aspect("equal"); ax.set_facecolor(FIELD_COLOR)
|
||||
ax.add_patch(mpatches.Rectangle((-15,-15), 30, 30,
|
||||
fill=False, edgecolor="#795548", lw=2))
|
||||
ax.add_patch(mpatches.Rectangle((10,-15), 3, 7,
|
||||
facecolor=PEN_COLOR, edgecolor="#795548", lw=2))
|
||||
ax.text(11.5, -11.5, "pen", ha="center", va="center",
|
||||
fontsize=8, color="#795548")
|
||||
|
||||
|
||||
def faded_path(ax, xs, ys, color, lw=1.5, label=None):
|
||||
"""Draw a path with alpha fading from start (transparent) to end (opaque)."""
|
||||
n = len(xs)
|
||||
if n < 2:
|
||||
return
|
||||
points = np.array([xs, ys]).T.reshape(-1, 1, 2)
|
||||
segs = np.concatenate([points[:-1], points[1:]], axis=1)
|
||||
alphas = np.linspace(0.15, 1.0, len(segs))
|
||||
colors = [(*matplotlib.colors.to_rgb(color), a) for a in alphas]
|
||||
lc = LineCollection(segs, colors=colors, linewidth=lw)
|
||||
ax.add_collection(lc)
|
||||
if label:
|
||||
ax.plot([], [], color=color, lw=lw, label=label)
|
||||
|
||||
|
||||
# ── main plots ────────────────────────────────────────────────────────────────
|
||||
|
||||
def plot_trajectory(hist, out_path):
|
||||
fig, ax = plt.subplots(figsize=(7, 7))
|
||||
draw_field(ax)
|
||||
|
||||
# Sheep paths
|
||||
for i in range(hist["n_sheep"]):
|
||||
c = SHEEP_COLORS[i % len(SHEEP_COLORS)]
|
||||
xs, ys = hist["sheep_xs"][i], hist["sheep_ys"][i]
|
||||
faded_path(ax, xs, ys, c, lw=1.2, label=f"sheep {i+1}")
|
||||
ax.plot(xs[0], ys[0], "o", color=c, ms=7, zorder=4)
|
||||
pa = hist["penned_at"][i]
|
||||
end = pa if pa is not None else -1
|
||||
ax.plot(xs[end], ys[end], "*", color=c, ms=11, zorder=5)
|
||||
|
||||
# Dog path
|
||||
faded_path(ax, hist["dog_xs"], hist["dog_ys"], DOG_COLOR, lw=2.0, label="dog")
|
||||
ax.plot(hist["dog_xs"][0], hist["dog_ys"][0], "s", color=DOG_COLOR, ms=10, zorder=5)
|
||||
ax.plot(hist["dog_xs"][-1], hist["dog_ys"][-1], "D", color=DOG_COLOR, ms=10, zorder=5)
|
||||
|
||||
result = "SUCCESS" if hist["success"] else f"FAIL ({hist['n_penned']}/{hist['n_sheep']} penned)"
|
||||
ax.set_title(f"Trajectory — {result} — {hist['steps']} steps", fontsize=12)
|
||||
ax.legend(loc="upper left", fontsize=8)
|
||||
plt.tight_layout()
|
||||
fig.savefig(out_path, dpi=120)
|
||||
plt.close(fig)
|
||||
print(f" saved {out_path}")
|
||||
|
||||
|
||||
def plot_timeseries(hist, out_path):
|
||||
t = np.arange(hist["steps"])
|
||||
fig, axes = plt.subplots(4, 1, figsize=(12, 10), sharex=True)
|
||||
|
||||
# 1. Flock radius
|
||||
axes[0].plot(t, hist["radii"], color="steelblue")
|
||||
axes[0].axhline(5.0, color="orange", ls="--", lw=1, label="compact threshold (5m)")
|
||||
axes[0].set_ylabel("flock radius (m)")
|
||||
axes[0].legend(fontsize=8)
|
||||
axes[0].set_title("Flock radius — goal: get below 5m")
|
||||
|
||||
# 2. Per-sheep distance to pen
|
||||
for i in range(hist["n_sheep"]):
|
||||
c = SHEEP_COLORS[i % len(SHEEP_COLORS)]
|
||||
axes[1].plot(t, hist["pen_dists"][i], color=c, lw=1, label=f"sheep {i+1}")
|
||||
pa = hist["penned_at"][i]
|
||||
if pa is not None:
|
||||
axes[1].axvline(pa, color=c, ls=":", lw=1)
|
||||
axes[1].set_ylabel("dist to pen (m)")
|
||||
axes[1].legend(fontsize=7, ncol=min(hist["n_sheep"], 5))
|
||||
axes[1].set_title("Per-sheep distance to pen — goal: all reach 0")
|
||||
|
||||
# 3. Action magnitude (how fast dog is moving)
|
||||
axes[2].plot(t, hist["action_mags"], color="tomato", lw=1)
|
||||
axes[2].axhline(1.0, color="gray", ls="--", lw=1, label="max")
|
||||
axes[2].set_ylabel("action ||(vx,vy)||")
|
||||
axes[2].set_ylim(0, 1.5)
|
||||
axes[2].set_title("Dog action magnitude — 0=stopped, 1=full speed")
|
||||
axes[2].legend(fontsize=8)
|
||||
|
||||
# 4. Reward per step
|
||||
axes[3].plot(t, hist["rewards"], color="purple", lw=1, alpha=0.7)
|
||||
axes[3].axhline(0, color="black", lw=0.5)
|
||||
axes[3].set_ylabel("reward")
|
||||
axes[3].set_xlabel("step")
|
||||
axes[3].set_title("Reward per step")
|
||||
|
||||
result = "SUCCESS" if hist["success"] else f"FAIL ({hist['n_penned']}/{hist['n_sheep']} penned)"
|
||||
fig.suptitle(f"n_sheep={hist['n_sheep']} {result} {hist['steps']} steps", fontsize=13)
|
||||
plt.tight_layout()
|
||||
fig.savefig(out_path, dpi=120)
|
||||
plt.close(fig)
|
||||
print(f" saved {out_path}")
|
||||
|
||||
|
||||
def save_gif(hist, out_path, fps=15, skip=5):
|
||||
"""Animated replay, every `skip` steps."""
|
||||
n = hist["n_sheep"]
|
||||
idxs = list(range(0, hist["steps"], skip))
|
||||
|
||||
fig, ax = plt.subplots(figsize=(6, 6))
|
||||
|
||||
def _frame(k):
|
||||
ax.clear()
|
||||
draw_field(ax)
|
||||
t = idxs[k]
|
||||
|
||||
for i in range(n):
|
||||
c = SHEEP_COLORS[i % len(SHEEP_COLORS)]
|
||||
s0 = max(0, t - 30)
|
||||
ax.plot(hist["sheep_xs"][i][s0:t+1],
|
||||
hist["sheep_ys"][i][s0:t+1],
|
||||
color=c, lw=0.8, alpha=0.5)
|
||||
color = "#ff69b4" if (hist["penned_at"][i] is not None
|
||||
and t >= hist["penned_at"][i]) else c
|
||||
ax.plot(hist["sheep_xs"][i][t], hist["sheep_ys"][i][t],
|
||||
"o", color=color, ms=10, zorder=4,
|
||||
markeredgecolor="#555", markeredgewidth=1)
|
||||
|
||||
s0 = max(0, t - 30)
|
||||
ax.plot(hist["dog_xs"][s0:t+1], hist["dog_ys"][s0:t+1],
|
||||
color=DOG_COLOR, lw=1.5, alpha=0.6)
|
||||
ax.plot(hist["dog_xs"][t], hist["dog_ys"][t],
|
||||
"s", color=DOG_COLOR, ms=13, zorder=5,
|
||||
markeredgecolor="black", markeredgewidth=1.5)
|
||||
|
||||
r = hist["radii"][t]
|
||||
ax.set_title(f"step {t}/{hist['steps']} radius={r:.1f}m "
|
||||
f"penned={hist['n_penned'] if t==hist['steps']-1 else '?'}/{n}",
|
||||
fontsize=10)
|
||||
|
||||
ani = animation.FuncAnimation(fig, _frame, frames=len(idxs), interval=1000//fps)
|
||||
ani.save(out_path, writer="pillow", fps=fps)
|
||||
plt.close(fig)
|
||||
print(f" saved {out_path}")
|
||||
|
||||
|
||||
# ── entry point ───────────────────────────────────────────────────────────────
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--model", default=None, help="Model .zip (omit for random policy)")
|
||||
p.add_argument("--vecnorm", default=None)
|
||||
p.add_argument("--n-sheep", type=int, default=3)
|
||||
p.add_argument("--max-steps", type=int, default=2000)
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument("--out-dir", default="vis_out")
|
||||
p.add_argument("--random", action="store_true",
|
||||
help="Use random policy (baseline comparison)")
|
||||
p.add_argument("--gif-fps", type=int, default=15)
|
||||
p.add_argument("--gif-skip", type=int, default=5,
|
||||
help="Render every Nth step in the GIF")
|
||||
p.add_argument("--no-gif", action="store_true")
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
os.makedirs(args.out_dir, exist_ok=True)
|
||||
|
||||
raw = DummyVecEnv([make_env(args.n_sheep, args.max_steps, args.seed)])
|
||||
|
||||
if args.random or args.model is None:
|
||||
print("Using RANDOM policy")
|
||||
env = raw
|
||||
model = None
|
||||
else:
|
||||
if args.vecnorm:
|
||||
env = VecNormalize.load(args.vecnorm, raw)
|
||||
env.training = False
|
||||
env.norm_reward = False
|
||||
else:
|
||||
env = raw
|
||||
model = PPO.load(args.model, env=env)
|
||||
print(f"Loaded model: {args.model}")
|
||||
|
||||
print(f"Running episode n_sheep={args.n_sheep} seed={args.seed} ...")
|
||||
hist = run_episode(model, env, args.n_sheep, args.max_steps)
|
||||
|
||||
result = "SUCCESS" if hist["success"] else f"FAIL ({hist['n_penned']}/{hist['n_sheep']} penned)"
|
||||
print(f"Episode done: {result} steps={hist['steps']}")
|
||||
print(f" min radius : {min(hist['radii']):.2f} m")
|
||||
print(f" mean reward: {np.mean(hist['rewards']):.4f}")
|
||||
print(f" mean action: {np.mean(hist['action_mags']):.3f}")
|
||||
|
||||
env.close()
|
||||
|
||||
plot_trajectory(hist, os.path.join(args.out_dir, "trajectory.png"))
|
||||
plot_timeseries(hist, os.path.join(args.out_dir, "timeseries.png"))
|
||||
if not args.no_gif:
|
||||
save_gif(hist, os.path.join(args.out_dir, "episode.gif"),
|
||||
fps=args.gif_fps, skip=args.gif_skip)
|
||||
|
||||
print(f"\nAll outputs saved to {args.out_dir}/")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user