Spaces:
Running
Running
File size: 4,452 Bytes
e9ca231 | 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 | """Generate reward plots from episode JSON logs.
Usage:
python scripts/plot_rewards.py --logs outputs/logs --out outputs/plots
Produces:
- outputs/plots/reward_curve.png (mean per-step reward curve)
- outputs/plots/reward_hist.png (histogram / boxplot of total rewards)
"""
from __future__ import annotations
import argparse
from pathlib import Path
import json
import math
import matplotlib.pyplot as plt
try:
import numpy as np
except Exception:
np = None
def load_episodes(log_dir: Path):
files = sorted(log_dir.glob("episode_*.json"))
episodes = []
for p in files:
try:
d = json.loads(p.read_text(encoding="utf-8"))
except Exception:
continue
# try common shapes
total = d.get("total_reward") if isinstance(d, dict) else None
if total is None:
total = d.get("reward") if isinstance(d, dict) else None
steps = d.get("steps") or d.get("episode") or d.get("trajectory") or []
if total is None and steps:
total = sum(s.get("reward", 0) for s in steps if isinstance(s, dict))
# extract per-step rewards
per_step = []
if steps and isinstance(steps, list):
for s in steps:
if isinstance(s, dict):
per_step.append(float(s.get("reward", 0)))
else:
# if step is numeric
try:
per_step.append(float(s))
except Exception:
pass
episodes.append({"name": p.name, "total": total, "per_step": per_step})
return episodes
def ensure_out(dirpath: Path):
dirpath.mkdir(parents=True, exist_ok=True)
def plot_reward_curve(episodes, out_path: Path):
# build aligned per-step arrays
per_steps = [e["per_step"] for e in episodes if e["per_step"]]
if not per_steps:
print("No per-step data found to plot reward curve.")
return
max_len = max(len(p) for p in per_steps)
if np is not None:
arr = np.full((len(per_steps), max_len), np.nan, dtype=float)
for i, p in enumerate(per_steps):
arr[i, : len(p)] = p
mean = np.nanmean(arr, axis=0)
std = np.nanstd(arr, axis=0)
x = list(range(1, len(mean) + 1))
else:
# fallback: compute mean manually
sums = [0.0] * max_len
counts = [0] * max_len
for p in per_steps:
for i, v in enumerate(p):
sums[i] += v
counts[i] += 1
mean = [sums[i] / counts[i] if counts[i] else 0.0 for i in range(max_len)]
std = [0.0] * max_len
x = list(range(1, max_len + 1))
plt.figure(figsize=(9, 4))
plt.plot(x, mean, label="mean reward per step")
try:
plt.fill_between(x, [m - s for m, s in zip(mean, std)], [m + s for m, s in zip(mean, std)], alpha=0.2)
except Exception:
pass
plt.xlabel("Step")
plt.ylabel("Reward")
plt.title("Mean per-step reward (shaded = std)")
plt.grid(alpha=0.2)
plt.legend()
plt.tight_layout()
plt.savefig(out_path / "reward_curve.png", dpi=150)
plt.close()
def plot_total_hist(episodes, out_path: Path):
totals = [e["total"] for e in episodes if e["total"] is not None]
if not totals:
print("No total reward data found to plot histogram.")
return
plt.figure(figsize=(6, 4))
plt.subplot(1, 2, 1)
plt.hist(totals, bins=min(30, max(5, len(totals)//2)), alpha=0.7)
plt.xlabel("Total reward")
plt.ylabel("Count")
plt.title("Total reward distribution")
plt.subplot(1, 2, 2)
plt.boxplot(totals, vert=False)
plt.title("Total reward (box)")
plt.tight_layout()
plt.savefig(out_path / "reward_hist.png", dpi=150)
plt.close()
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--logs", type=Path, default=Path("outputs/logs"))
parser.add_argument("--out", type=Path, default=Path("outputs/plots"))
args = parser.parse_args()
if not args.logs.exists():
print(f"Logs directory not found: {args.logs}")
return
ensure_out(args.out)
episodes = load_episodes(args.logs)
print(f"Found {len(episodes)} episode files in {args.logs}")
plot_reward_curve(episodes, args.out)
plot_total_hist(episodes, args.out)
print(f"Saved plots to {args.out}")
if __name__ == "__main__":
main()
|