code/scripts/smoke_planner.py

"""Smoke test the planner against the real frozen LeWM before any long run.

Section 5 of the brief: tiny batch, ~50 steps, H=3, T=1. Verify shapes, verify
both silent-failure invariants, verify one optimizer step reduces the loss, and
measure memory and it/s — ``H*T*n`` applications of ``f`` per sample is
expensive and the throughput here is what the time estimate is built from.

The unit tests cover the same invariants against a toy stand-in so they need no
checkpoint. This runs them against the actual world model, where the predictor
is a real 6-layer transformer and the rollout is where the memory goes.
"""

import argparse
import json
import sys
import time
from pathlib import Path

import torch
from torch.utils.data import DataLoader

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from lejepa_control.data import LatentGoalDataset, split_episodes  # noqa: E402
from lejepa_control.world_model import load_lewm  # noqa: E402
from lejepa_control_2.losses import DistanceScale, planner_loss  # noqa: E402
from lejepa_control_2.scripts.train_planner import (  # noqa: E402
    build_planner,
    parse_args,
)


def check_invariants(planner, model, batch, scale, device, horizon, cycles):
    """The two failures that are silent, plus the activation count."""
    planner.reset_call_counts()
    out = planner(
        model,
        batch['context'],
        batch['past_actions'],
        batch['goal'],
        horizon=horizon,
        cycles=cycles,
    )

    B = batch['context'].size(0)
    N = planner.num_context
    shapes = {
        'distances': (B, horizon),
        'blocks': (B, horizon, planner.block_dim),
        'raw': (B, horizon, planner.block_dim),
        'contexts': (B, horizon, N, planner.latent_dim),
        'answers': (B, horizon, planner.width),
        'frames': (B, N + horizon, planner.latent_dim),
    }
    for key, want in shapes.items():
        got = tuple(out[key].shape)
        assert got == want, f'{key}: shape {got}, expected {want}'
        assert torch.isfinite(out[key]).all(), f'{key} is not finite'
    print(f'  shapes ok ({len(shapes)} tensors, all finite)')

    # invariant 1: the h-chain is differentiable all the way back to block 0
    terminal = out['distance_seq'][-1]
    assert terminal.requires_grad, 'terminal distance carries no gradient'
    norms = []
    for k in range(horizon):
        g = torch.autograd.grad(
            terminal.sum(), out['block_seq'][k], retain_graph=True,
            allow_unused=True,
        )[0]
        assert g is not None and g.abs().sum() > 0, (
            f'terminal distance does not reach block {k} — the h-chain is cut '
            f'and the planner has collapsed to a greedy one-step policy'
        )
        norms.append(round(g.norm().item(), 6))
    print(f'  invariant 1 ok: d[-1] reaches every block, norms {norms}')

    # invariant 3: the Change-1 detach schedule is actually in effect
    expected = horizon * (planner.inner + 1)
    unstaged = horizon * cycles * (planner.inner + 1)
    assert planner.grad_calls == expected, (
        f'grad-enabled f/g applications = {planner.grad_calls}, expected '
        f'{expected} = H(n+1). Full backprop would be {unstaged}.'
    )
    print(
        f'  invariant 3 ok: {planner.grad_calls} grad-enabled applications '
        f'= H(n+1); full backprop would be {unstaged}'
    )

    loss, _ = planner_loss(out, batch['goal_offset'], scale,
                           planner.action_embed)
    loss.backward()

    # invariant 2: the world model is frozen but transparent
    leaked = [n for n, p in model.named_parameters() if p.grad is not None]
    assert not leaked, f'frozen world model accumulated gradients: {leaked[:5]}'
    assert planner.f.body.out.weight.grad.abs().sum() > 0, (
        'no gradient reached f — the world model is opaque, not just frozen'
    )
    print(
        f'  invariant 2 ok: 0/{sum(1 for _ in model.parameters())} world-model '
        f'params have grads, and f still trains'
    )
    planner.zero_grad(set_to_none=True)


def main():
    argv = sys.argv[1:]
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument('--smoke-steps', type=int, default=50)
    smoke_args, rest = parser.parse_known_args(argv)
    # tiny by default; anything on the command line still wins
    defaults = ['--stage', 'A', '--batch-size', '8', '--pretrain-steps', '20']
    args = parse_args(defaults + rest)

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    torch.manual_seed(args.seed)

    stats = json.loads((Path(args.latents) / 'stats.json').read_text())
    model = load_lewm(device=device)
    planner = build_planner(args, stats, stats['latent_dim'], device)
    scale = DistanceScale().to(device)

    horizon, cycles = 3, args.cycles
    print(
        f'smoke: H={horizon} T={cycles} n={args.inner} B={args.batch_size} '
        f'on {device}, planner '
        f'{sum(p.numel() for p in planner.parameters()) / 1e6:.2f}M params'
    )

    train_eps, _ = split_episodes(stats['n_episodes'])
    dataset = LatentGoalDataset(
        args.latents, episodes=train_eps, horizon=horizon, max_offset=3
    )
    loader = DataLoader(
        dataset, batch_size=args.batch_size, shuffle=True, drop_last=True
    )
    raw = next(iter(loader))
    batch = {k: v.to(device) for k, v in raw.items()}
    scale.update(batch['context'], batch['goal'])

    print('invariants:')
    check_invariants(planner, model, batch, scale, device, horizon, cycles)

    # does one optimizer step actually reduce the loss?
    opt = torch.optim.AdamW(
        planner.parameters(), lr=args.lr, weight_decay=args.weight_decay
    )
    if device == 'cuda':
        torch.cuda.reset_peak_memory_stats()

    losses = []
    t0 = time.perf_counter()
    for _ in range(smoke_args.smoke_steps):
        scale.update(batch['context'], batch['goal'])
        out = planner(
            model, batch['context'], batch['past_actions'], batch['goal'],
            horizon=horizon, cycles=cycles,
        )
        loss, metrics = planner_loss(
            out, batch['goal_offset'], scale, planner.action_embed,
            lambda_cycle=args.lambda_cycle,
            lambda_support=args.lambda_support,
            lambda_anchor=args.lambda_anchor,
            lambda_sat=args.lambda_sat,
        )
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(planner.parameters(), args.grad_clip)
        opt.step()
        losses.append(loss.item())
    if device == 'cuda':
        torch.cuda.synchronize()
    elapsed = time.perf_counter() - t0

    print(
        f'overfit {smoke_args.smoke_steps} steps on a fixed batch: '
        f'{losses[0]:.4f} -> {losses[-1]:.4f}'
    )
    assert losses[-1] < losses[0], 'loss did not fall on a fixed batch'

    rate = smoke_args.smoke_steps / elapsed
    peak = (
        torch.cuda.max_memory_allocated() / 2**30 if device == 'cuda' else 0.0
    )
    print(
        f'throughput {rate:.2f} it/s at B={args.batch_size} H={horizon} '
        f'T={cycles}, peak VRAM {peak:.2f} GiB'
    )
    # what the real run costs: the driver's cost scales with H*T, and stages
    # B-E run at B=32
    scale_factor = (args.batch_size / 32) * (3 / horizon) * (1 / max(cycles, 1))
    print(
        f'  extrapolated to B=32 H=3 T=3: '
        f'~{rate * scale_factor * 3:.2f} it/s '
        f'({20000 / max(rate * scale_factor * 3, 1e-6) / 60:.0f} min for 20k)'
    )
    print('SMOKE TEST PASSED')


if __name__ == '__main__':
    main()