code/scripts/train_decoder.py

"""Train the latent->pixel decoder (LeWM appendix D), for visualization only.

    python tools/make_decoder_cache.py --frames 8000
    python scripts/train_decoder.py --steps 2000

The world model is **not** touched: the decoder reads cached latents, so no
gradient can reach the encoder even by accident. That matches how LeWM used it
(Fig. 8 is a read-only probe) and matters, because the paper's own appendix G
ablation shows that letting a reconstruction loss into LeWM training *hurts*
control — PushT success 96.0 ± 2.83 without it, 86.0 ± 7.54 with it.

Presets pick the speed/detail trade-off; all were measured on this CPU box at
batch 32:

    paper   224px, 16px patches, W384 d4   7.6M params   ~6.0 s/step
    cpu     224px, 16px patches, W256 d3   2.7M params   ~1.4 s/step  (default)
    fast    224px, 28px patches, W256 d3   3.0M params   ~0.5 s/step

``fast`` keeps the full 224 output but tiles it 8x8, so a 10px pusher lands
inside one patch — fine for "where is the block", poor for fine pose. ``cpu``
tiles 14x14 and is the recommended default here.
"""

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

import numpy as np
import torch

REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO))
os.environ.setdefault('STABLEWM_HOME', str(REPO / 'data' / 'swm_home'))

from lejepa_control.decoder import LatentDecoder, reconstruction_loss  # noqa: E402
from tools.paths import artifact_dir  # noqa: E402

PRESETS = {
    'paper': dict(patch_size=16, hidden_dim=384, depth=4, heads=6),
    'cpu': dict(patch_size=16, hidden_dim=256, depth=3, heads=4),
    'fast': dict(patch_size=28, hidden_dim=256, depth=3, heads=4),
}


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument('--cache', default=None,
                   help='default: $LEJEPA_DATA/decoder_cache, else data/decoder_cache')
    p.add_argument('--out', default='data/runs/decoder')
    p.add_argument('--preset', choices=list(PRESETS), default='cpu')
    p.add_argument('--steps', type=int, default=2000)
    p.add_argument('--batch-size', type=int, default=32)
    p.add_argument('--lr', type=float, default=3e-4)
    p.add_argument('--weight-decay', type=float, default=0.05)
    p.add_argument('--warmup', type=int, default=100)
    p.add_argument('--val-frac', type=float, default=0.05)
    p.add_argument('--log-every', type=int, default=50)
    p.add_argument('--preview-every', type=int, default=500)
    p.add_argument('--seed', type=int, default=0)
    return p.parse_args()


def save_preview(decoder, z, target, path, device):
    """Side-by-side target/reconstruction strip — the thing you actually look at."""
    from PIL import Image

    decoder.eval()
    with torch.no_grad():
        pred = decoder(z.to(device)).clamp(0, 1).cpu()
    decoder.train()

    n = min(6, len(z))
    rows = []
    for tensor in (target[:n], pred[:n]):
        row = tensor.permute(0, 2, 3, 1).numpy()
        rows.append(np.concatenate(list(row), axis=1))
    strip = (np.concatenate(rows, axis=0) * 255).astype(np.uint8)
    path.parent.mkdir(parents=True, exist_ok=True)
    Image.fromarray(strip).save(path)


def main():
    args = parse_args()
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    torch.manual_seed(args.seed)

    cache = artifact_dir('decoder_cache', args.cache)
    meta = json.loads((cache / 'meta.json').read_text())
    images = np.load(cache / 'images.npy', mmap_mode='r')   # (N,3,H,W) uint8
    latents = np.load(cache / 'latents.npy')                # (N,192) fp16
    assert len(images) == len(latents), 'cache is inconsistent'
    print(f'cache: {len(images)} pairs, latent={meta["latent_source"]}, '
          f'source={meta["source"]}, image={meta["image_size"]}px')

    # Held-out split so the reported loss is not just memorisation. The decoder
    # has 2.7M params against 8k images and will happily overfit.
    rng = np.random.default_rng(args.seed)
    perm = rng.permutation(len(images))
    n_val = max(1, int(len(images) * args.val_frac))
    val_idx, train_idx = np.sort(perm[:n_val]), perm[n_val:]
    print(f'  train {len(train_idx)} / val {len(val_idx)}')

    cfg = dict(
        latent_dim=int(meta['latent_dim']),
        image_size=int(meta['image_size']),
        **PRESETS[args.preset],
    )
    decoder = LatentDecoder(**cfg)
    n_params = sum(p.numel() for p in decoder.parameters())
    print(f'decoder[{args.preset}]: {n_params / 1e6:.2f}M params, '
          f'P={decoder.num_patches} patches of {decoder.patch_size}px')

    def batch(idx_pool, size):
        idx = np.sort(rng.choice(idx_pool, size, replace=False))
        z = torch.from_numpy(latents[idx].astype(np.float32))
        x = torch.from_numpy(np.asarray(images[idx], dtype=np.float32) / 255.0)
        return z, x

    # Baseline: predicting the dataset mean image. Any decoder that does not
    # beat this has learned nothing about the latent. On PushT it is a strict
    # bar, not a weak one — the frames are mostly identical white background,
    # so a constant image is already a decent predictor and only the agent,
    # block and target carry error.
    mean_img = torch.from_numpy(
        np.asarray(images[train_idx[:1000]], dtype=np.float32).mean(0) / 255.0
    )
    zv, xv = batch(val_idx, min(64, len(val_idx)))
    baseline = float(((mean_img.unsqueeze(0) - xv) ** 2).mean())
    print(f'  mean-image baseline val MSE = {baseline:.5f}')

    decoder.init_output_at(mean_img.mean(dim=(1, 2)))
    decoder = decoder.to(device)
    print(f'  output head parked on the mean colour '
          f'{[round(float(v), 3) for v in mean_img.mean(dim=(1, 2))]}\n')

    # Built after the init and the device move so AdamW never sees stale params.
    opt = torch.optim.AdamW(
        decoder.parameters(), lr=args.lr, weight_decay=args.weight_decay
    )
    sched = torch.optim.lr_scheduler.LambdaLR(
        opt,
        lambda s: min(1.0, (s + 1) / max(1, args.warmup))
        * (0.5 * (1 + np.cos(np.pi * min(1.0, s / args.steps)))),
    )

    out_dir = REPO / args.out
    out_dir.mkdir(parents=True, exist_ok=True)
    history = []
    t0 = time.perf_counter()

    for step in range(1, args.steps + 1):
        z, x = batch(train_idx, args.batch_size)
        loss = reconstruction_loss(decoder(z.to(device)), x.to(device))
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(decoder.parameters(), 1.0)
        opt.step()
        sched.step()

        if step % args.log_every == 0 or step == 1:
            decoder.eval()
            with torch.no_grad():
                val = float(reconstruction_loss(decoder(zv.to(device)), xv.to(device)))
            decoder.train()
            rate = step / (time.perf_counter() - t0)
            train_loss = float(loss.detach())
            history.append({'step': step, 'train': train_loss, 'val': val})
            print(f'  step {step:5}/{args.steps}  train {train_loss:.5f}  '
                  f'val {val:.5f}  ({val / baseline:.2f}x baseline)  '
                  f'{rate:.2f} it/s  eta {(args.steps - step) / rate / 60:.1f} min',
                  flush=True)

        if step % args.preview_every == 0 or step == args.steps:
            save_preview(decoder, zv, xv, out_dir / f'preview_{step:05d}.png', device)

    torch.save({
        'state_dict': decoder.state_dict(),
        'config': cfg,
        'preset': args.preset,
        'meta': meta,
        'step': args.steps,
        'history': history,
        'baseline_val_mse': baseline,
    }, out_dir / 'decoder.pt')

    print(f'\nsaved -> {out_dir / "decoder.pt"}')
    print(f'  {(time.perf_counter() - t0) / 60:.1f} min')
    print(f'  final val MSE {history[-1]["val"]:.5f} vs baseline {baseline:.5f}')
    print(f'  previews: {out_dir}/preview_*.png  (top row target, bottom row decoded)')


if __name__ == '__main__':
    main()