File size: 6,331 Bytes
dc9f917 | 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 143 144 145 146 147 148 149 150 151 152 153 154 155 | """E3 -- token-latent geometry probe (exp7 explainability battery, sec 5).
Only meaningful in the 192-identity cells (A, C): with ``width ==
latent_dim``, plan tokens ``y_j`` live in literally the same coordinate
system as world-model latents, so they can be compared to ``x_hat_j`` (the
state block ``j`` actually causes) with a plain distance/cosine, and decoded
*directly* through the pixel decoder -- no analogous move exists in any
256-d variant, which is exactly what "identity embedding buys" (sec 1).
"""
import argparse
import json
import sys
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from lejepa_control.data import LatentGoalDataset, split_episodes # noqa: E402
from lejepa_control.decoder import load_decoder # noqa: E402
from lejepa_control.rollout import rollout_plan # noqa: E402
from lejepa_control.solver import load_controller # noqa: E402
from lejepa_control.world_model import load_lewm # noqa: E402
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'tools'))
from decode_rollout import panel, to_uint8 # noqa: E402
def parse_args():
p = argparse.ArgumentParser()
p.add_argument('--controller', default='data/runs/exp7/fused192/controller.pt')
p.add_argument('--decoder', default='data/runs/decoder/decoder.pt')
p.add_argument('--latents', default='data/latents')
p.add_argument('--samples', type=int, default=256)
p.add_argument('--horizon', type=int, default=5)
p.add_argument('--decode-episodes', type=int, default=3)
p.add_argument('--tag', default='main')
p.add_argument('--out', default='data/runs/probe_geometry')
return p.parse_args()
@torch.no_grad()
def geometry_trace(controller, model, ctx, past, goal):
"""Per-refinement ``y_k`` (raw plan tokens), ``x_hat_k`` (the latent each
block causes), ``||y - x_hat||`` and cosine. ``K+1`` entries, same
recursion as ``IterativeController.forward``.
"""
cond = controller.condition(ctx, goal)
tokens = controller.initial_plan(cond)
ys, xhats, dists, coss = [], [], [], []
for k in range(controller.refinements + 1):
actions = controller.to_actions(tokens)
pred = rollout_plan(model, ctx, past, actions) # x_hat_j, (B, H, D)
ys.append(tokens.clone())
xhats.append(pred.clone())
dists.append((tokens - pred).norm(dim=-1)) # (B, H)
coss.append(F.cosine_similarity(tokens, pred, dim=-1)) # (B, H)
if k == controller.refinements:
break
delta = controller.refine(tokens, cond, pred, goal)
idx = min(k, controller.step_logit.numel() - 1)
tokens = tokens + torch.sigmoid(controller.step_logit[idx]) * delta
return torch.stack(ys), torch.stack(xhats), torch.stack(dists), torch.stack(coss)
@torch.no_grad()
def decode_panel(decoder, ys, sample_idx, out_path):
"""One panel per sample: rows = refinement k, cols = horizon slot j,
each cell = ``decoder(y_k[sample, j])`` -- what that plan slot "intends".
"""
rows = [to_uint8(decoder(ys[k, sample_idx])) for k in range(ys.size(0))]
labels = [f'k={k}' for k in range(ys.size(0))]
titles = [f'slot {j}' for j in range(ys.size(2))]
return panel(rows, labels, titles, out_path)
def main():
args = parse_args()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.manual_seed(0)
stats = json.loads((Path(args.latents) / 'stats.json').read_text())
model = load_lewm(device=device)
controller, ckpt = load_controller(
args.controller, latent_dim=stats['latent_dim'], device=device
)
controller.eval()
if not controller.no_latent_proj:
print(f'{args.controller}: no_latent_proj=False (width != latent_dim) '
'-- E3 is only defined for the 192-identity cells (A, C). Skipping.')
return
print(f'controller step {ckpt["step"]}, K={controller.refinements}, '
f'fused={controller.fused}')
_, val_eps = split_episodes(stats['n_episodes'])
val = LatentGoalDataset(
args.latents, max_offset=5, episodes=val_eps, horizon=args.horizon
)
idx = np.random.default_rng(0).choice(len(val), args.samples, replace=False)
batch = {
k: torch.stack([val[int(i)][k] for i in idx]).to(device)
for k in ('context', 'past_actions', 'goal')
}
ys, xhats, dists, coss = geometry_trace(
controller, model, batch['context'], batch['past_actions'], batch['goal']
)
per_k_dist = dists.mean(dim=(1, 2))
per_k_cos = coss.mean(dim=(1, 2))
print('\n=== E3: plan token vs. its own consequence, ||y_j - x_hat_j|| ===')
print(f'{"k":>3}{"mean ||.||":>12}{"mean cos":>11}')
for k in range(dists.size(0)):
print(f'{k:>3}{per_k_dist[k].item():>12.4f}{per_k_cos[k].item():>11.4f}')
report = {
'tag': args.tag,
'checkpoint': args.controller,
'step': int(ckpt['step']),
'fused': bool(controller.fused),
'mean_dist_per_k': [round(float(v), 5) for v in per_k_dist],
'mean_cos_per_k': [round(float(v), 5) for v in per_k_cos],
}
out = Path(args.out)
out.mkdir(parents=True, exist_ok=True)
with (out / 'probe_geometry.jsonl').open('a') as f:
f.write(json.dumps(report) + '\n')
print(f'wrote {out / "probe_geometry.jsonl"}')
if args.decode_episodes > 0 and Path(args.decoder).exists():
decoder, dec_ckpt = load_decoder(args.decoder, device=device)
latent_kind = dec_ckpt['meta']['latent_source']
if latent_kind != 'emb':
print(f"decoder trained on '{latent_kind}' latents, not 'emb' "
'-- skipping the direct plan-token decode.')
else:
for i in range(min(args.decode_episodes, ys.size(1))):
path = decode_panel(
decoder, ys, i, out / f'{args.tag}_sample{i}.png'
)
print(f' -> {path}')
elif args.decode_episodes > 0:
print(f'no decoder at {args.decoder} -- skipping the direct plan-token decode.')
if __name__ == '__main__':
main()
|