Johnny Fernandes
154 lines
5.7 KiB
Python
154 lines
5.7 KiB
Python
"""
|
|
Render Webots-side debug trajectory from debug.csv.
|
|
|
|
The shepherd_dog_rl controller writes per-step state to debug.csv when
|
|
DOG_DEBUG=1. This script reads it and produces:
|
|
|
|
trajectory.png — dog path + sheep paths overlaid on the field
|
|
obs_drift.png — normalized observation distribution over time
|
|
actions.png — vx, vy time series
|
|
|
|
Run:
|
|
python plot_debug.py # uses debug.csv next to this file
|
|
python plot_debug.py --csv path/to.csv --out-dir somewhere/
|
|
"""
|
|
import argparse
|
|
import csv
|
|
import os
|
|
import sys
|
|
|
|
import matplotlib
|
|
matplotlib.use("Agg")
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.patches as mpatches
|
|
import numpy as np
|
|
|
|
|
|
def load_csv(path):
|
|
rows = []
|
|
with open(path) as f:
|
|
rd = csv.DictReader(f)
|
|
for r in rd:
|
|
rows.append(r)
|
|
if not rows:
|
|
sys.exit(f"empty CSV: {path}")
|
|
return rows
|
|
|
|
|
|
def parse_floats(s):
|
|
return [float(x) for x in s.split(";") if x]
|
|
|
|
|
|
def plot_trajectory(rows, out_path):
|
|
fig, ax = plt.subplots(figsize=(7, 7))
|
|
ax.set_xlim(-16, 16); ax.set_ylim(-16, 16); ax.set_aspect("equal")
|
|
ax.set_facecolor("#dcedc8")
|
|
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="#ffe082", edgecolor="#795548", lw=2))
|
|
ax.text(11.5, -11.5, "pen", ha="center", va="center", fontsize=8)
|
|
|
|
dog_x = [float(r["dog_x"]) for r in rows]
|
|
dog_y = [float(r["dog_y"]) for r in rows]
|
|
ax.plot(dog_x, dog_y, color="#4e342e", lw=1.5, alpha=0.7, label="dog")
|
|
ax.plot(dog_x[0], dog_y[0], "s", color="#4e342e", ms=10)
|
|
ax.plot(dog_x[-1], dog_y[-1], "D", color="#4e342e", ms=10)
|
|
|
|
# Sheep — re-shape into per-sheep tracks
|
|
sx_all = [parse_floats(r["sheep_xs"]) for r in rows]
|
|
sy_all = [parse_floats(r["sheep_ys"]) for r in rows]
|
|
if sx_all and sx_all[-1]:
|
|
n_sheep = len(sx_all[-1])
|
|
palette = ["#e41a1c","#377eb8","#4daf4a","#984ea3","#ff7f00",
|
|
"#a65628","#f781bf","#999999","#66c2a5","#fc8d62"]
|
|
for i in range(n_sheep):
|
|
xs = [r[i] if i < len(r) else None for r in sx_all]
|
|
ys = [r[i] if i < len(r) else None for r in sy_all]
|
|
xs = [x for x in xs if x is not None]
|
|
ys = [y for y in ys if y is not None]
|
|
if xs:
|
|
c = palette[i % len(palette)]
|
|
ax.plot(xs, ys, color=c, lw=0.8, alpha=0.6, label=f"sheep {i+1}")
|
|
ax.plot(xs[0], ys[0], "o", color=c, ms=6)
|
|
ax.plot(xs[-1], ys[-1], "*", color=c, ms=10)
|
|
|
|
n_in_pen = int(rows[-1]["n_penned"])
|
|
ax.set_title(f"Webots trajectory {len(rows)} steps penned={n_in_pen}",
|
|
fontsize=12)
|
|
ax.legend(loc="upper left", fontsize=7, ncol=2)
|
|
plt.tight_layout()
|
|
fig.savefig(out_path, dpi=120)
|
|
plt.close(fig)
|
|
|
|
|
|
def plot_actions(rows, out_path):
|
|
t = np.arange(len(rows))
|
|
vx = np.array([float(r["vx"]) for r in rows])
|
|
vy = np.array([float(r["vy"]) for r in rows])
|
|
mag = np.sqrt(vx ** 2 + vy ** 2)
|
|
|
|
fig, axes = plt.subplots(3, 1, figsize=(12, 7), sharex=True)
|
|
axes[0].plot(t, vx, color="tab:red", lw=0.8); axes[0].set_ylabel("vx")
|
|
axes[0].axhline(0, color="black", lw=0.4); axes[0].set_ylim(-1.1, 1.1)
|
|
axes[1].plot(t, vy, color="tab:blue", lw=0.8); axes[1].set_ylabel("vy")
|
|
axes[1].axhline(0, color="black", lw=0.4); axes[1].set_ylim(-1.1, 1.1)
|
|
axes[2].plot(t, mag, color="tab:purple", lw=0.8); axes[2].set_ylabel("||action||")
|
|
axes[2].axhline(np.sqrt(2), color="orange", ls="--", lw=1, label="saturated √2")
|
|
axes[2].axhline(1.0, color="gray", ls="--", lw=1)
|
|
axes[2].set_xlabel("step"); axes[2].legend(fontsize=8)
|
|
fig.suptitle("Webots action time series")
|
|
plt.tight_layout()
|
|
fig.savefig(out_path, dpi=120)
|
|
plt.close(fig)
|
|
|
|
|
|
def plot_obs(rows, out_path):
|
|
norm = np.array([parse_floats(r["norm_obs"]) for r in rows])
|
|
raw = np.array([parse_floats(r["raw_obs"]) for r in rows])
|
|
if norm.size == 0:
|
|
return
|
|
n_dims = norm.shape[1]
|
|
labels = [
|
|
"dog_x", "dog_y", "com-dog_x", "com-dog_y",
|
|
"far1-com_x", "far1-com_y", "far2-com_x", "far2-com_y",
|
|
"far3-com_x", "far3-com_y", "pen-com_x", "pen-com_y",
|
|
"pen-far1_x", "pen-far1_y", "radius", "frac_active",
|
|
][:n_dims]
|
|
|
|
t = np.arange(norm.shape[0])
|
|
fig, axes = plt.subplots(n_dims, 1, figsize=(11, 1.0 * n_dims), sharex=True)
|
|
if n_dims == 1: axes = [axes]
|
|
for i in range(n_dims):
|
|
axes[i].plot(t, raw[:, i], color="tab:gray", lw=0.6, alpha=0.6, label="raw")
|
|
axes[i].plot(t, norm[:, i], color="tab:red", lw=0.8, label="normalised")
|
|
axes[i].set_ylabel(labels[i], fontsize=8)
|
|
axes[i].tick_params(labelsize=7)
|
|
if i == 0:
|
|
axes[i].legend(fontsize=7, loc="upper right")
|
|
axes[-1].set_xlabel("step")
|
|
fig.suptitle("Observation values over time (raw vs VecNormalize-normalised)")
|
|
plt.tight_layout()
|
|
fig.savefig(out_path, dpi=110)
|
|
plt.close(fig)
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser()
|
|
here = os.path.dirname(os.path.abspath(__file__))
|
|
p.add_argument("--csv", default=os.path.join(here, "debug.csv"))
|
|
p.add_argument("--out-dir", default=os.path.join(here, "debug_out"))
|
|
args = p.parse_args()
|
|
|
|
rows = load_csv(args.csv)
|
|
os.makedirs(args.out_dir, exist_ok=True)
|
|
print(f"loaded {len(rows)} rows from {args.csv}")
|
|
plot_trajectory(rows, os.path.join(args.out_dir, "trajectory.png"))
|
|
plot_actions(rows, os.path.join(args.out_dir, "actions.png"))
|
|
plot_obs(rows, os.path.join(args.out_dir, "obs.png"))
|
|
print(f"saved trajectory.png + actions.png + obs.png to {args.out_dir}/")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|