Leplanner / code /report /make_figures.py
nottygian's picture
Push code package
872cf4d verified
Raw
History Blame Contribute Delete
21.6 kB
"""Generate every figure in the report from the raw result files.
Reads only ``data/runs/eval/results.jsonl``, ``data/runs/diagnostics/*`` and
the saved checkpoints' validation profiles, so the figures cannot drift from
the numbers they are supposed to show. Writes PDF (for LaTeX) and PNG (for
quick viewing) side by side into ``report/figures/``.
Run from the repo root:
python report/make_figures.py
"""
import json
import sys
from pathlib import Path
import matplotlib
import numpy as np
matplotlib.use('Agg')
import matplotlib.pyplot as plt # noqa: E402
ROOT = Path(__file__).resolve().parents[1]
FIG = Path(__file__).resolve().parent / 'figures'
FIG.mkdir(parents=True, exist_ok=True)
plt.rcParams.update({
'figure.dpi': 140,
'savefig.dpi': 140,
'font.size': 9,
'axes.grid': True,
'grid.alpha': 0.25,
'grid.linewidth': 0.6,
'axes.spines.top': False,
'axes.spines.right': False,
'axes.titlesize': 10,
'legend.frameon': False,
'legend.fontsize': 8,
})
# consistent colour per variant across every figure
C = {
'abl_terminal_only': '#b2182b',
'original': '#ef8a62',
'abl_no_support': '#d6604d',
'ah_hold0.0': '#92c5de',
'ah_hold0.5': '#2166ac',
'ah_hold1.0': '#4393c3',
'cem': '#4d4d4d',
}
LABEL = {
'abl_terminal_only': r'terminal-only ($\alpha{=}0$)',
'original': 'original',
'abl_no_support': 'no support',
'ah_hold0.0': r'arrival ($\lambda_h{=}0$)',
'ah_hold0.5': r'arrival+hold ($\lambda_h{=}0.5$)',
'ah_hold1.0': r'arrival+hold ($\lambda_h{=}1$)',
'cem': 'CEM',
}
def save(fig, name):
for ext in ('pdf', 'png'):
fig.savefig(FIG / f'{name}.{ext}', bbox_inches='tight')
plt.close(fig)
print(f' wrote figures/{name}.pdf + .png')
# --------------------------------------------------------------------------
# data loading
# --------------------------------------------------------------------------
def load_eval():
rows = [
json.loads(x)
for x in (ROOT / 'data/runs/eval/results.jsonl').read_text().splitlines()
if x.strip()
]
for r in rows:
ckpt = r.get('checkpoint') or ''
if 'ah_hold0.5' in ckpt:
r['variant'] = 'ah_hold0.5'
elif 'ah_hold0.0' in ckpt:
r['variant'] = 'ah_hold0.0'
elif 'ah_hold1.0' in ckpt:
r['variant'] = 'ah_hold1.0'
elif 'abl_no_support' in ckpt:
r['variant'] = 'abl_no_support'
elif 'abl_terminal_only' in ckpt:
r['variant'] = 'abl_terminal_only'
elif 'controller/controller.pt' in ckpt.replace('\\', '/'):
r['variant'] = 'original'
else:
r['variant'] = 'cem'
r['K'] = int(r['planner'].split('_K')[-1]) if '_K' in r['planner'] else None
r['m'] = r.get('receding_horizon', 1)
return rows
def load_diag():
return {
json.loads(x)['tag']: json.loads(x)
for x in (ROOT / 'data/runs/diagnostics/diagnostics.jsonl')
.read_text().splitlines() if x.strip()
}
def load_profiles():
"""Per-q distance profiles from each checkpoint's saved validation state.
The original controller predates the profile logging, so its entry comes
from ``report/recover_profiles.py``, which recomputes it with the same
``evaluate()`` on the same held-out split.
"""
import torch
out = {}
for name in ('controller', 'abl_terminal_only', 'abl_no_support',
'ah_hold0.0', 'ah_hold0.5', 'ah_hold1.0'):
p = ROOT / f'data/runs/{name}/controller.pt'
if not p.exists():
continue
ck = torch.load(p, map_location='cpu', weights_only=False)
prof = ck.get('val_profile')
key = 'original' if name == 'controller' else name
if prof:
out[key] = {int(q): np.array(v) for q, v in prof.items()}
posthoc = ROOT / 'data/runs/diagnostics/profiles_posthoc.json'
if posthoc.exists():
for name, rec in json.loads(posthoc.read_text()).items():
key = 'original' if name == 'controller' else name
out.setdefault(key, {int(q): np.array(v)
for q, v in rec['profile'].items()})
return out
def latest(rows, variant, m, K=3):
"""Most recent row for a variant/schedule (replicates append)."""
sel = [r for r in rows
if r['variant'] == variant and r['m'] == m
and (K is None or r['K'] == K)]
return sel[-1] if sel else None
# --------------------------------------------------------------------------
# Figure 1 — the pathology: success vs execution length
# --------------------------------------------------------------------------
def fig_execution_sweep(rows):
fig, ax = plt.subplots(figsize=(4.4, 3.0))
ms = [1, 2, 3, 4, 5]
succ = []
for m in ms:
r = latest(rows, 'original', m, K=3)
succ.append(r['success_rate'] if r else np.nan)
ax.plot(ms, succ, 'o-', color=C['original'], lw=2, ms=6,
label='original controller (K=3)')
ax.axhline(succ[0], color=C['original'], ls=':', lw=1, alpha=0.6)
ax.annotate('', xy=(4.55, succ[-1]), xytext=(4.55, succ[0]),
arrowprops=dict(arrowstyle='<->', color='0.35', lw=1.2))
ax.text(4.45, (succ[0] + succ[-1]) / 2, f'{succ[-1] - succ[0]:+.0f} pts',
ha='right', va='center', fontsize=8.5, color='0.25')
# CEM at both schedules: the same pathology in a gradient-free planner
ax.plot([1, 5],
[latest(rows, 'cem', 1, K=None)['success_rate'],
latest(rows, 'cem', 5, K=None)['success_rate']],
's--', color=C['cem'], lw=1.6, ms=5, label='CEM (300$\\times$30)')
ax.set_xlabel('blocks executed before replanning, $m$')
ax.set_ylabel('success rate (\\%)' if plt.rcParams['text.usetex']
else 'success rate (%)')
ax.set_xticks(ms)
ax.set_ylim(0, 100)
ax.set_title('Executing more of the plan helps — backwards from MPC theory')
ax.legend(loc='lower right')
save(fig, 'fig1_execution_sweep')
# --------------------------------------------------------------------------
# Figure 2 — the mechanism: per-q distance profiles
# --------------------------------------------------------------------------
def fig_profiles(profiles):
show = ['abl_terminal_only', 'original', 'ah_hold0.5']
fig, axes = plt.subplots(1, len(show), figsize=(9.6, 2.9), sharey=True)
blocks = np.arange(1, 6)
cmap = plt.get_cmap('viridis')
for ax, name in zip(axes, show):
prof = profiles.get(name)
if prof is None:
continue
for q in sorted(prof):
d = prof[q]
col = cmap((q - 1) / 4 * 0.85)
ax.plot(blocks, d, 'o-', color=col, ms=4, lw=1.4,
label=f'$q={q}$')
j = int(np.argmin(d))
ax.plot(blocks[j], d[j], '*', color=col, ms=13,
markeredgecolor='k', markeredgewidth=0.4, zorder=5)
ax.set_yscale('log')
ax.set_xticks(blocks)
ax.set_xlabel('plan block $j$')
ax.set_title(LABEL[name])
axes[0].set_ylabel(r'predicted $d_j$ (log)')
axes[-1].legend(loc='upper right', ncol=1)
fig.suptitle(r'Stars mark $\arg\min_j d_j$. Left and middle: the minimum '
r'is pinned at block 5 for every goal offset $q$. '
r'Right: it tracks $q$.',
y=1.04, fontsize=9)
save(fig, 'fig2_arrival_profiles')
# --------------------------------------------------------------------------
# Figure 3 — closed-loop contraction traces and fixed points
# --------------------------------------------------------------------------
def fig_contraction(diag):
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8.4, 3.1))
order = ['abl_terminal_only', 'original', 'ah_hold0.0',
'ah_hold0.5', 'ah_hold1.0']
for name in order:
d = diag.get(name if name != 'original' else 'original')
if d is None:
continue
tr = np.array(d['contraction']['exec1']['mean_trace'])
fp = d['contraction']['exec1']['fixed_point']
n = np.arange(1, len(tr) + 1)
ax1.plot(n, tr, 'o-', color=C[name], ms=3.5, lw=1.6, label=LABEL[name])
ax1.axhline(fp, color=C[name], ls=':', lw=1, alpha=0.55)
ax1.set_xlabel('replan $n$')
ax1.set_ylabel(r'mean latent goal distance $D_n$')
ax1.set_yscale('log')
ax1.set_ylim(4.5e-3, 0.85)
ax1.set_title(r'Closed-loop trace, $m=1$ (dotted: fitted $D^\ast$)')
ax1.legend(loc='lower left', fontsize=7.2, ncol=2)
# right panel: the recursion itself
names, cs, bs, fps = [], [], [], []
for name in order:
d = diag.get(name)
if d is None:
continue
f = d['contraction']['exec1']
names.append(name)
cs.append(f['c'])
bs.append(f['b'])
fps.append(f['fixed_point'])
x = np.arange(len(names))
ax2.bar(x, fps, color=[C[n] for n in names], width=0.62)
for xi, (fp, c, b) in enumerate(zip(fps, cs, bs)):
ax2.text(xi, fp + 0.006, f'{fp:.3f}', ha='center', fontsize=8)
ax2.text(xi, 0.004, f'$c$={c:.2f}\n$b$={b:.3f}', ha='center',
fontsize=6.8, color='w', va='bottom')
ax2.set_xticks(x)
ax2.set_xticklabels([LABEL[n].replace(' (', '\n(') for n in names],
fontsize=6.4, rotation=22, ha='right')
ax2.set_ylabel(r'$D^\ast = b/(1-c)$')
ax2.set_ylim(0, 0.235)
ax2.set_title(r'Fixed point of $D_{n+1}=cD_n+b$')
save(fig, 'fig3_contraction')
# --------------------------------------------------------------------------
# Figure 4 — the ablation matrix
# --------------------------------------------------------------------------
def fig_ablation(rows):
order = ['abl_terminal_only', 'abl_no_support', 'original',
'ah_hold0.0', 'ah_hold1.0', 'ah_hold0.5']
m1 = [latest(rows, v, 1)['success_rate'] for v in order]
m5 = [latest(rows, v, 5)['success_rate'] for v in order]
fig, ax = plt.subplots(figsize=(6.6, 3.2))
x = np.arange(len(order))
w = 0.38
ax.bar(x - w / 2, m1, w, label='$m=1$ (replan every block)',
color=[C[v] for v in order], edgecolor='k', linewidth=0.4)
ax.bar(x + w / 2, m5, w, label='$m=5$ (execute full plan)',
color=[C[v] for v in order], edgecolor='k', linewidth=0.4,
alpha=0.42, hatch='///')
for xi, (a, b) in enumerate(zip(m1, m5)):
ax.text(xi - w / 2, a + 1.5, f'{a:.0f}', ha='center', fontsize=8)
ax.text(xi + w / 2, b + 1.5, f'{b:.0f}', ha='center', fontsize=8)
gap = a - b
ax.text(xi, -12, f'{gap:+.0f}', ha='center', fontsize=8,
color='#b2182b' if gap < -10 else '#1a6b3c',
fontweight='bold')
ax.set_xticks(x)
ax.set_xticklabels([LABEL[v].replace(' (', '\n(') for v in order],
fontsize=7.5)
ax.set_ylabel('success rate (%)')
ax.set_ylim(-16, 122)
ax.axhline(0, color='k', lw=0.8)
ax.text(-0.72, -12, 'gap', fontsize=8, fontweight='bold', color='0.3')
ax.set_title('Objective ablations at both execution schedules '
'(gap $=$ $m{=}1$ $-$ $m{=}5$)')
ax.legend(loc='upper left', ncol=2, fontsize=7.5)
save(fig, 'fig4_ablation')
# --------------------------------------------------------------------------
# Figure 5 — refinement past trained depth
# --------------------------------------------------------------------------
def fig_refinement(diag):
from matplotlib.lines import Line2D
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8.6, 3.1))
for name in ('original', 'ah_hold0.5', 'abl_terminal_only'):
d = diag.get(name)
if d is None:
continue
rf = d['refinement']
k = np.arange(len(rf['terminal']))
ax1.plot(k, rf['terminal'], 'o-', color=C[name], ms=4)
ax1.plot(k, rf['arrival'], 's--', color=C[name], ms=4, alpha=0.65)
ax2.plot(k[1:], rf['plan_change'][1:], 'o-', color=C[name], ms=4,
label=LABEL[name])
ax1.axvspan(3.5, 8.5, color='0.85', alpha=0.5, zorder=0)
ax1.text(6, 0.38, 'beyond trained\ndepth $K=3$', ha='center',
fontsize=7.5, color='0.35')
ax1.set_xlabel('refinement $k$')
ax1.set_ylabel('predicted distance (log)')
ax1.set_yscale('log')
ax1.set_ylim(8e-3, 0.75)
ax1.set_title('Terminal vs arrival cost by refinement')
ax1.legend(handles=[
Line2D([], [], color='0.3', marker='o', ls='-',
label=r'terminal $d_H$'),
Line2D([], [], color='0.3', marker='s', ls='--',
label=r'arrival $d_q$'),
], loc='lower left', fontsize=7.5)
ax2.set_xlabel('refinement $k$')
ax2.set_ylabel(r'mean $|b^{(k)}-b^{(k-1)}|$')
ax2.set_ylim(0, 0.30)
ax2.set_title('Plan movement never settles to zero')
ax2.legend(fontsize=7.5)
save(fig, 'fig5_refinement')
# --------------------------------------------------------------------------
# Figure 6 — cost/accuracy Pareto front
# --------------------------------------------------------------------------
def fig_pareto(rows):
fig, ax = plt.subplots(figsize=(6.2, 3.7))
# (x, y, label, colour, marker, label offset in points)
pts = []
off = {0: (9, -3), 1: (-6, 6), 2: (3, -13), 3: (-1, -14), 5: (-9, 7)}
for K in (0, 1, 2, 3, 5):
r = latest(rows, 'original', 5, K=K)
if r:
pts.append((r['predictor_rows_per_episode'], r['success_rate'],
f'$K$={K}', C['original'], 'o', off[K]))
r = latest(rows, 'original', 1, K=3)
pts.append((r['predictor_rows_per_episode'], r['success_rate'],
'original, $m$=1', C['original'], 'X', (10, -3)))
for m, o in ((1, (-10, 5)), (5, (-10, 5))):
r = latest(rows, 'cem', m, K=None)
pts.append((r['predictor_rows_per_episode'], r['success_rate'],
f'CEM, $m$={m}', C['cem'], 's', o))
for xx, yy, lab, col, mk, o in pts:
ax.scatter(xx, yy, marker=mk, s=48, color=col,
edgecolor='k', linewidth=0.5, zorder=3)
ax.annotate(lab, (xx, yy), textcoords='offset points',
xytext=o, fontsize=7.2,
ha='right' if o[0] < 0 else 'left')
# the three corrected controllers sit on top of each other — label the
# cluster once rather than three overlapping annotations
corrected = []
for v, mk in (('ah_hold0.5', '*'), ('ah_hold0.0', 'D'), ('ah_hold1.0', 'v')):
r = latest(rows, v, 1)
if r:
corrected.append((r['predictor_rows_per_episode'],
r['success_rate'], v, mk))
ax.scatter(r['predictor_rows_per_episode'], r['success_rate'],
marker=mk, s=170 if mk == '*' else 50, color=C[v],
edgecolor='k', linewidth=0.5, zorder=4,
label=LABEL[v] + ', $m$=1')
if corrected:
cx = float(np.mean([p[0] for p in corrected]))
ax.annotate('corrected controllers,\n$m=1$ (replanning every block)',
xy=(cx * 1.35, 94.6), xytext=(900, 101),
fontsize=7.4, ha='left', va='center', color='0.2',
arrowprops=dict(arrowstyle='->', color='0.45', lw=1,
connectionstyle='arc3,rad=0.25'))
ax.set_xscale('log')
ax.set_xlim(4, 4e6)
ax.set_xlabel('predictor rows per episode (log scale)')
ax.set_ylabel('success rate (%)')
ax.set_ylim(25, 106)
ax.set_title('Success against planning cost')
ax.legend(loc='lower left', fontsize=7.2)
save(fig, 'fig6_pareto')
# --------------------------------------------------------------------------
# Figure 7 — the survivorship confound
# --------------------------------------------------------------------------
def fig_survivorship(rows):
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8.2, 3.0))
paired = [r for r in rows if r.get('episode_successes')]
s = np.array([r['success_rate'] for r in paired])
mt = np.array([r['mean_terminal_distance'] for r in paired])
r_coef = np.corrcoef(s, mt)[0, 1]
ax1.scatter(s, mt, s=32, color='#2166ac', alpha=0.75,
edgecolor='k', linewidth=0.4)
fit = np.polyfit(s, mt, 1)
xs = np.linspace(s.min(), s.max(), 20)
ax1.plot(xs, np.polyval(fit, xs), '--', color='#b2182b', lw=1.4)
ax1.set_xlabel('success rate (%)')
ax1.set_ylabel('mean terminal distance')
ax1.set_title(f'Better controllers report higher mean cost\n'
f'$r={r_coef:+.3f}$ over {len(paired)} paired rows')
order = ['abl_terminal_only', 'original', 'abl_no_support',
'ah_hold0.0', 'ah_hold1.0', 'ah_hold0.5']
uniq = [latest(rows, v, 1, K=3) for v in order]
uniq = [r for r in uniq if r]
uniq.sort(key=lambda r: r['success_rate'])
names = [r['variant'] for r in uniq]
alive = [r['predictor_rows_per_episode'] * 50 / r['predictor_calls']
for r in uniq]
succ = [r['success_rate'] for r in uniq]
x = np.arange(len(names))
ax2.bar(x, alive, color=[C[n] for n in names], edgecolor='k', linewidth=0.4)
for xi, (a, sv) in enumerate(zip(alive, succ)):
ax2.text(xi, a + 0.8, f'{a:.1f}', ha='center', fontsize=8)
ax2.text(xi, 1.2, f'{sv:.0f}%', ha='center', fontsize=7.5, color='w')
ax2.set_xticks(x)
ax2.set_xticklabels([LABEL[n].replace(' (', '\n(') for n in names],
fontsize=6.2, rotation=30, ha='right')
ax2.set_ylabel('mean episodes still alive')
ax2.set_ylim(0, 53)
ax2.set_title('rows/call = mean surviving episodes\n'
'(all six share identical per-decision cost)')
save(fig, 'fig7_survivorship')
# --------------------------------------------------------------------------
# Figure 8 — training-time early warning: terminal vs arrival
# --------------------------------------------------------------------------
def fig_training_signal():
import re
logs = {
'original': ROOT / 'data/runs/controller/train.log',
}
job = Path('C:/Users/omnap/.claude/jobs/e23cbae5/tmp')
spans = {
'ah_hold0.5': (job / 'sweep2.log', 2, 104),
'abl_terminal_only': (job / 'sweep4.log', 1, 133),
}
def parse(lines):
step, term, arr = [], [], []
pat = re.compile(
r'^step\s+(\d+).*?terminal\s+([\d.]+)(?:\s+arrival\s+([\d.]+))?')
for ln in lines:
m = pat.match(ln)
if m:
step.append(int(m.group(1)))
term.append(float(m.group(2)))
arr.append(float(m.group(3)) if m.group(3) else np.nan)
return np.array(step), np.array(term), np.array(arr)
series = {}
for name, p in logs.items():
if p.exists():
series[name] = parse(p.read_text(errors='ignore').splitlines())
for name, (p, lo, hi) in spans.items():
if p.exists():
series[name] = parse(
p.read_text(errors='ignore').splitlines()[lo - 1:hi])
fig, ax = plt.subplots(figsize=(5.9, 3.3))
for name in ('abl_terminal_only', 'ah_hold0.5'):
if name not in series:
continue
st, term, arr = series[name]
ax.plot(st, term, '-', color=C[name], lw=1.6,
label=LABEL[name] + r' — terminal $d_H$')
if not np.all(np.isnan(arr)):
ax.plot(st, arr, '--', color=C[name], lw=1.6, alpha=0.75,
label=LABEL[name] + r' — arrival $d_q$')
if 'abl_terminal_only' in series:
st, term, arr = series['abl_terminal_only']
ratio = arr[-1] / term[-1]
ax.annotate(f'{ratio:.0f}$\\times$ gap', xy=(st[-1], arr[-1]),
xytext=(-6, -2), textcoords='offset points',
ha='right', va='top', fontsize=7.5, color=C['abl_terminal_only'])
ax.set_yscale('log')
ax.set_ylim(8e-3, 1.4)
ax.set_xlabel('training step')
ax.set_ylabel('running mean distance (log)')
ax.set_title('The pathology is visible during training,\nwith no rollout '
'needed')
ax.legend(fontsize=7, loc='upper right', ncol=1)
save(fig, 'fig8_training_signal')
def main():
print('loading results...')
rows = load_eval()
diag = load_diag()
print(f' {len(rows)} eval rows, {len(diag)} diagnostics records')
try:
profiles = load_profiles()
print(f' {len(profiles)} checkpoint profiles')
except Exception as e: # torch missing / checkpoints moved
print(f' !! could not load checkpoints ({e}); skipping fig2')
profiles = {}
print('rendering...')
fig_execution_sweep(rows)
if profiles:
fig_profiles(profiles)
fig_contraction(diag)
fig_ablation(rows)
fig_refinement(diag)
fig_pareto(rows)
fig_survivorship(rows)
fig_training_signal()
print('done.')
if __name__ == '__main__':
sys.exit(main())