| """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
|
| from lejepa_control.decoder import load_decoder
|
| from lejepa_control.rollout import rollout_plan
|
| from lejepa_control.solver import load_controller
|
| from lejepa_control.world_model import load_lewm
|
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1] / 'tools'))
|
| from decode_rollout import panel, to_uint8
|
|
|
|
|
| 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)
|
|
|
| ys.append(tokens.clone())
|
| xhats.append(pred.clone())
|
| dists.append((tokens - pred).norm(dim=-1))
|
| coss.append(F.cosine_similarity(tokens, pred, dim=-1))
|
|
|
| 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()
|
|
|
