code/scripts/train_planner.py

"""Train the recursive latent planner against the frozen LeWM.

Staged bring-up, per ``recursive_planner_design.pdf`` and the task brief's
section 5. Each stage turns on one more mechanism and has a gate that must
clear before the next one starts::

    A  arrival+hold + late path, H=3, T=1, no feedback
       gate: loss falls, real-env success beats random
    B  + consequence feedback (2) + deep supervision (3), T=3
       gate: per-cycle distances strictly decreasing  <-- the important one
    C  + manifold anchor (7) + saturation barrier (6)
       gate: anchor ratio < 0.05 stable, |r| > 2 fraction < 10%
    D  + support hinge (8), H -> 5
       gate: violation fraction < 5% and flat, real-env success still rising
    E  warm-start verification (cold-start control run)
       gate: equal-or-better success, per-cycle curve flattens earlier

Stage presets only fill arguments left unset, so any flag given explicitly on
the command line wins.

Every run writes ``config.json``, ``metrics.jsonl`` (one record per log
interval, with the full diagnostic set from section 11) and a resumable
``planner.pt``.
"""

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

import numpy as np
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.losses import BehaviorDensity  # 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.planner import RecursivePlanner  # noqa: E402

# Arguments a stage preset is allowed to set. Anything the user passes
# explicitly overrides the preset, so `--stage B --cycles 5` does what it says.
STAGES = {
    'A': dict(
        cycles=1,
        horizon_curriculum='0:3',
        use_feedback=False,
        lambda_cycle=0.0,
        lambda_anchor=0.0,
        lambda_sat=0.0,
        lambda_support=0.0,
    ),
    'B': dict(
        cycles=3,
        horizon_curriculum='0:3',
        use_feedback=True,
        lambda_cycle=0.3,
        lambda_anchor=0.0,
        lambda_sat=0.0,
        lambda_support=0.0,
    ),
    'C': dict(
        cycles=3,
        horizon_curriculum='0:3',
        use_feedback=True,
        lambda_cycle=0.3,
        lambda_anchor=0.05,
        lambda_sat=1e-3,
        lambda_support=0.0,
    ),
    'D': dict(
        cycles=3,
        horizon_curriculum='0:3,0.5:5',
        use_feedback=True,
        lambda_cycle=0.3,
        lambda_anchor=0.05,
        lambda_sat=1e-3,
        lambda_support=0.01,
    ),
    # E is D's configuration with warm start off: the cold-start control the
    # stage-E gate is a paired comparison against.
    'E': dict(
        cycles=3,
        horizon_curriculum='0:3,0.5:5',
        use_feedback=True,
        lambda_cycle=0.3,
        lambda_anchor=0.05,
        lambda_sat=1e-3,
        lambda_support=0.01,
        warm_start=False,
    ),
}


def parse_args(argv=None):
    p = argparse.ArgumentParser()
    p.add_argument('--latents', default='data/latents')
    p.add_argument('--density', default='data/runs/density/density.pt')
    p.add_argument('--out', default='data/runs/planner')
    p.add_argument('--stage', choices=sorted(STAGES), default=None)
    p.add_argument('--resume', action='store_true')
    p.add_argument('--seed', type=int, default=0)

    # -- shapes (section 10; fixed by LeWM) --------------------------------
    p.add_argument('--width', type=int, default=256)
    p.add_argument('--hidden', type=int, default=512)

    # -- loops --------------------------------------------------------------
    p.add_argument('--inner', type=int, default=6)
    p.add_argument('--cycles', type=int, default=None)
    p.add_argument(
        '--horizon-curriculum',
        default=None,
        help='fraction_of_training:horizon, comma separated (base: 0:3,0.5:5)',
    )
    p.add_argument(
        '--curriculum',
        default='0:2,0.25:3,0.5:5',
        help='fraction_of_training:max_goal_offset, comma separated',
    )

    # -- loss weights (section 10) -----------------------------------------
    p.add_argument('--hold-weight', type=float, default=0.5)
    p.add_argument('--alpha', type=float, default=0.05)
    p.add_argument('--lambda-cycle', type=float, default=None)
    p.add_argument('--lambda-support', type=float, default=None)
    p.add_argument('--lambda-anchor', type=float, default=None)
    p.add_argument('--lambda-sat', type=float, default=None)
    p.add_argument('--sat-limit', type=float, default=2.0)

    # -- optimization -------------------------------------------------------
    p.add_argument('--steps', type=int, default=20000)
    p.add_argument('--batch-size', type=int, default=32)
    p.add_argument('--lr', type=float, default=1e-4)
    p.add_argument('--weight-decay', type=float, default=1e-4)
    p.add_argument('--grad-clip', type=float, default=1.0)
    p.add_argument('--workers', type=int, default=0)

    # -- phi/psi autoencoder pre-training (Change 7) ------------------------
    p.add_argument('--pretrain-steps', type=int, default=400)
    p.add_argument('--pretrain-lr', type=float, default=1e-3)

    # -- structural switches (Changes 9 and 11 are structural, not staged) --
    p.add_argument('--no-warm-start', dest='warm_start', action='store_false')
    p.add_argument('--lambda-z', type=float, default=0.0)
    p.add_argument('--learn-lambda-z', action='store_true')
    p.set_defaults(warm_start=None)

    # -- section 12 ablations: wired, not run as part of this task ----------
    # 1: --cycles / --eval-cycles     2: --no-feedback     3: --inner/--cycles
    p.add_argument('--no-feedback', dest='use_feedback', action='store_false')
    p.set_defaults(use_feedback=None)
    # 4: detach schedule. 'full' is also the --cycle-grad-boundary option --
    # it keeps every cycle in the graph so deep supervision gets real gradient.
    p.add_argument(
        '--detach-schedule',
        choices=['last-cycle', 'one-step', 'full'],
        default='last-cycle',
    )
    # 5: arrival+hold vs a fixed terminal d_H
    p.add_argument('--terminal-only', action='store_true')
    # 6: late-rising vs discounted path weighting
    p.add_argument(
        '--path-weighting', choices=['late', 'discount'], default='late'
    )
    p.add_argument('--gamma', type=float, default=0.9)
    # 7 is a sweep over --lambda-support. 8 (fixed vs growing n) is skipped.

    # -- logging ------------------------------------------------------------
    p.add_argument('--log-every', type=int, default=100)
    p.add_argument('--val-every', type=int, default=1000)
    p.add_argument('--val-batches', type=int, default=20)
    p.add_argument('--in-memory', action='store_true', default=True)
    p.add_argument('--mmap', dest='in_memory', action='store_false')

    args = p.parse_args(argv)
    return apply_stage(args)


def apply_stage(args):
    """Fill unset arguments from the stage preset; explicit flags win."""
    preset = STAGES.get(args.stage, {})
    fallback = dict(
        cycles=3,
        horizon_curriculum='0:3,0.5:5',
        use_feedback=True,
        warm_start=True,
        lambda_cycle=0.3,
        lambda_support=0.01,
        lambda_anchor=0.05,
        lambda_sat=1e-3,
    )
    for key, default in fallback.items():
        if getattr(args, key) is None:
            setattr(args, key, preset.get(key, default))
    return args


def parse_curriculum(spec, total_steps):
    stages = []
    for part in spec.split(','):
        frac, value = part.split(':')
        stages.append((int(float(frac) * total_steps), int(value)))
    return sorted(stages)


def current_value(stages, step):
    value = stages[0][1]
    for start, v in stages:
        if step >= start:
            value = v
    return value


def seed_everything(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def build_planner(args, stats, latent_dim, device):
    a_mean = torch.tensor(stats['action_mean'])
    a_std = torch.tensor(stats['action_std'])
    return RecursivePlanner(
        latent_dim=latent_dim,
        width=args.width,
        hidden=args.hidden,
        inner=args.inner,
        cycles=args.cycles,
        horizon=max(v for _, v in parse_curriculum(
            args.horizon_curriculum, args.steps
        )),
        use_feedback=args.use_feedback,
        warm_start=args.warm_start,
        lambda_z=args.lambda_z,
        learn_lambda_z=args.learn_lambda_z,
        detach_schedule=args.detach_schedule,
        # tanh bounds live in the normalized action units the world model was
        # trained on: raw PushT actions are in [-1, 1], so the bound is
        # (-mean/std, 1/std)
        action_center=(-a_mean / a_std),
        action_scale=(1.0 / a_std),
    ).to(device)


def pretrain_action_embedding(planner, loader, args, device, log=print):
    """Change 7's prerequisite: fit ``phi``/``psi`` as a plain autoencoder.

    The anchor pulls ``y`` toward whatever ``phi(psi(y))`` happens to be. If
    the pair starts random that is a meaningless target and the anchor spends
    the early run fighting the goal loss over where the manifold even is. A
    few hundred steps on real dataset blocks settles it first.
    """
    if args.pretrain_steps <= 0:
        return None
    embed = planner.action_embed
    opt = torch.optim.AdamW(embed.parameters(), lr=args.pretrain_lr)
    losses, step = [], 0
    while step < args.pretrain_steps:
        for batch in loader:
            # every real block in the sample: the two past blocks and the
            # block leaving the current frame
            blocks = torch.cat(
                [
                    batch['past_actions'].flatten(0, 1),
                    batch['real_action'],
                ],
                dim=0,
            ).to(device)
            recon = embed.decode(embed.encode(blocks))
            loss = (recon - blocks).pow(2).mean()
            opt.zero_grad(set_to_none=True)
            loss.backward()
            opt.step()
            losses.append(loss.item())
            step += 1
            if step >= args.pretrain_steps:
                break
    log(
        f'phi/psi pretrain: {losses[0]:.4f} -> '
        f'{sum(losses[-20:]) / min(20, len(losses)):.4f} '
        f'over {step} steps'
    )
    return {'first': losses[0], 'last': sum(losses[-20:]) / min(20, len(losses))}


@torch.no_grad()
def evaluate(planner, model, loader, scale, args, device, horizon):
    """Held-out imagined metrics. Real-env success is a separate script.

    ``arrival`` is the distance at each sample's own relabel offset ``q``,
    which is what receding-horizon execution depends on. ``terminal`` is the
    distance at step ``H`` regardless of ``q`` — a planner that defers arrival
    scores well on terminal and badly on arrival, and the gap between them is
    the Change-4 diagnostic.
    """
    planner.eval()
    total = {'terminal': 0.0, 'arrival': 0.0}
    per_step = torch.zeros(horizon, device=device)
    per_cycle = None
    batches = 0
    for batch in loader:
        q = batch['goal_offset'].to(device).clamp(1, horizon)
        out = planner(
            model,
            batch['context'].to(device),
            batch['past_actions'].to(device),
            batch['goal'].to(device),
            horizon=horizon,
        )
        d = scale.normalize(out['distances'])
        total['terminal'] += d[:, -1].mean().item()
        total['arrival'] += (
            d.gather(1, (q - 1).unsqueeze(1)).squeeze(1).mean().item()
        )
        per_step += d.mean(dim=0)
        if out['cycle_distances'] is not None:
            c = scale.normalize(out['cycle_distances']).mean(dim=(0, 1))
            per_cycle = c if per_cycle is None else per_cycle + c
        batches += 1
        if batches >= args.val_batches:
            break
    planner.train()
    result = {k: v / batches for k, v in total.items()}
    result['per_step'] = (per_step / batches).tolist()
    if per_cycle is not None:
        result['per_cycle'] = (per_cycle / batches).tolist()
    return result


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

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    ckpt_path = out_dir / 'planner.pt'
    log_path = out_dir / 'metrics.jsonl'

    stats = json.loads((Path(args.latents) / 'stats.json').read_text())
    latent_dim = stats['latent_dim']

    model = load_lewm(device=device)  # frozen, eval, requires_grad_(False)
    planner = build_planner(args, stats, latent_dim, device)
    scale = DistanceScale().to(device)

    horizon_stages = parse_curriculum(args.horizon_curriculum, args.steps)
    offset_stages = parse_curriculum(args.curriculum, args.steps)
    max_horizon = max(v for _, v in horizon_stages)

    n_params = sum(p.numel() for p in planner.parameters())
    print(
        f'stage {args.stage or "custom"}  planner {n_params / 1e6:.2f}M params  '
        f'n={args.inner} T={args.cycles} H={horizon_stages} '
        f'feedback={args.use_feedback} warm_start={args.warm_start}'
    )

    # The support model is loaded whenever it exists, even at stages that do
    # not pay for it: the violation fraction is an early-warning diagnostic and
    # it leads real-env degradation by 1-2k steps (section 11).
    density, c95 = None, None
    if Path(args.density).exists():
        d_ckpt = torch.load(args.density, map_location=device)
        density = BehaviorDensity(
            latent_dim=latent_dim, components=d_ckpt['components']
        ).to(device)
        density.load_state_dict(d_ckpt['state_dict'])
        density.eval().requires_grad_(False)
        c95 = d_ckpt['c95']
        print(
            f'support model loaded, c95={c95:.4f} '
            f'(lambda_support={args.lambda_support})'
        )
    else:
        print(f'no support model at {args.density} — support term disabled')

    train_eps, val_eps = split_episodes(stats['n_episodes'])
    train_set = LatentGoalDataset(
        args.latents,
        episodes=train_eps,
        horizon=max_horizon,
        in_memory=args.in_memory,
    )
    val_set = LatentGoalDataset(
        args.latents,
        max_offset=5,
        episodes=val_eps,
        horizon=max_horizon,
        in_memory=args.in_memory,
    )
    gen = torch.Generator().manual_seed(args.seed)
    loader = DataLoader(
        train_set,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.workers,
        drop_last=True,
        persistent_workers=args.workers > 0,
        pin_memory=True,
        generator=gen,
    )
    val_loader = DataLoader(
        val_set,
        batch_size=args.batch_size,
        shuffle=True,
        generator=torch.Generator().manual_seed(args.seed + 1),
    )

    opt = torch.optim.AdamW(
        planner.parameters(), lr=args.lr, weight_decay=args.weight_decay
    )
    sched = torch.optim.lr_scheduler.OneCycleLR(
        opt, max_lr=args.lr, total_steps=args.steps, pct_start=0.05
    )

    step = 0
    pretrain = None
    if args.resume and ckpt_path.exists():
        ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
        saved = ckpt['args']
        # OneCycle's shape is defined by its total, and both curricula are
        # indexed by fractions of it, so a resume with a different budget is a
        # different run — fail loudly rather than silently changing the schedule
        for key in ('steps', 'curriculum', 'horizon_curriculum', 'lr'):
            if saved[key] != getattr(args, key):
                raise SystemExit(
                    f'cannot resume: --{key.replace("_", "-")} was '
                    f'{saved[key]!r} in the checkpoint, now '
                    f'{getattr(args, key)!r}. Start a new run instead.'
                )
        planner.load_state_dict(ckpt['state_dict'])
        scale.load_state_dict(ckpt['scale'])
        opt.load_state_dict(ckpt['optimizer'])
        sched.load_state_dict(ckpt['scheduler'])
        step = ckpt['step']
        torch.set_rng_state(ckpt['rng_torch'].cpu().to(torch.uint8))
        np.random.set_state(ckpt['rng_numpy'])
        random.setstate(ckpt['rng_python'])
        if ckpt.get('rng_cuda') is not None and torch.cuda.is_available():
            torch.cuda.set_rng_state_all(
                [s.cpu().to(torch.uint8) for s in ckpt['rng_cuda']]
            )
        print(f'resumed from {ckpt_path} at step {step}')
    else:
        pretrain = pretrain_action_embedding(planner, loader, args, device)
        (out_dir / 'config.json').write_text(json.dumps(vars(args), indent=2))
        log_path.write_text('')

    def write_record(record):
        with log_path.open('a') as fh:
            fh.write(json.dumps(record) + '\n')

    def save(extra=None):
        torch.save(
            {
                'state_dict': planner.state_dict(),
                'scale': scale.state_dict(),
                'optimizer': opt.state_dict(),
                'scheduler': sched.state_dict(),
                'step': step,
                'args': vars(args),
                'action_mean': stats['action_mean'],
                'action_std': stats['action_std'],
                'latent_dim': latent_dim,
                'horizon': planner.horizon,
                'rng_torch': torch.get_rng_state(),
                'rng_numpy': np.random.get_state(),
                'rng_python': random.getstate(),
                'rng_cuda': (
                    torch.cuda.get_rng_state_all()
                    if torch.cuda.is_available()
                    else None
                ),
                **(extra or {}),
            },
            ckpt_path,
        )

    running, t0 = {}, time.perf_counter()
    planner.train()
    while step < args.steps:
        for batch in loader:
            horizon = current_value(horizon_stages, step)
            # clamp the offset curriculum to the horizon actually being rolled
            # out; otherwise arrival_hold_loss silently clamps q down to H and
            # the deadline stops meaning what the curriculum says it means
            offset = min(current_value(offset_stages, step), horizon)
            if train_set.max_offset != offset:
                train_set.set_max_offset(offset)

            ctx = batch['context'].to(device, non_blocking=True)
            past = batch['past_actions'].to(device, non_blocking=True)
            goal = batch['goal'].to(device, non_blocking=True)
            q = batch['goal_offset'].to(device, non_blocking=True)

            capture = (step + 1) % args.log_every == 0
            if capture:
                planner.start_capture()

            scale.update(ctx, goal)
            out = planner(model, ctx, past, goal, horizon=horizon)
            loss, metrics = planner_loss(
                out,
                q,
                scale,
                planner.action_embed,
                density=density,
                c95=c95,
                hold_weight=args.hold_weight,
                alpha=args.alpha,
                lambda_cycle=args.lambda_cycle,
                lambda_support=args.lambda_support,
                lambda_anchor=args.lambda_anchor,
                lambda_sat=args.lambda_sat,
                sat_limit=args.sat_limit,
                terminal_only=args.terminal_only,
                path_weighting=args.path_weighting,
                gamma=args.gamma,
            )

            opt.zero_grad(set_to_none=True)
            loss.backward()
            grad = torch.nn.utils.clip_grad_norm_(
                planner.parameters(), args.grad_clip
            )
            opt.step()
            sched.step()
            step += 1

            for key in ('loss', 'arrival', 'path', 'cycle', 'support',
                        'violation', 'anchor', 'anchor_ratio', 'sat',
                        'sat_fraction'):
                if key in metrics:
                    running[key] = running.get(key, 0.0) + metrics[key].item()
            running['grad'] = running.get('grad', 0.0) + grad.item()
            running['n'] = running.get('n', 0) + 1

            if capture:
                planner.stop_capture()
                taps = planner.taps
                n = running.pop('n')
                rate = step / (time.perf_counter() - t0)
                record = {
                    'step': step,
                    'horizon': horizon,
                    'max_offset': offset,
                    'lr': sched.get_last_lr()[0],
                    'it_per_s': round(rate, 3),
                    'scale': scale.scale.item(),
                    **{k: v / n for k, v in running.items()},
                    # section 11's diagnostic set
                    'per_step': [round(v, 5) for v in
                                 metrics['per_step'].tolist()],
                    'gates': planner.gate_values(),
                    # ||z^(i)|| across i: flat is healthy, RMSNorm makes it so
                    'z_norms': [round(v, 4) for v in taps['z_norms'][:32]],
                    # first vs last f application: within 10x, or the backprop
                    # depth is larger than the detach schedule intends
                    'grad_first': taps['grad_first'][:1],
                    'grad_last': taps['grad_last'][:1],
                }
                if 'per_cycle' in metrics:
                    record['per_cycle'] = [
                        round(v, 5) for v in metrics['per_cycle'].tolist()
                    ]
                write_record(record)

                msg = (
                    f'step {step:6d}  H={horizon} q<={offset}  '
                    f'loss {record["loss"]:.4f}  '
                    f'arrival {record["arrival"]:.4f}  '
                    f'grad {record["grad"]:.2f}  {rate:.2f} it/s'
                )
                if 'per_cycle' in record:
                    msg += f'  cycles {record["per_cycle"]}'
                if 'violation' in record:
                    msg += f'  viol {record["violation"]:.3f}'
                if 'anchor_ratio' in record:
                    msg += f'  anchor {record["anchor_ratio"]:.4f}'
                if 'sat_fraction' in record:
                    msg += f'  sat {record["sat_fraction"]:.3f}'
                print(msg, flush=True)
                running = {}

            if step % args.val_every == 0 or step >= args.steps:
                val = evaluate(
                    planner, model, val_loader, scale, args, device, horizon
                )
                print(
                    f'  [val] arrival {val["arrival"]:.4f}  '
                    f'terminal {val["terminal"]:.4f}  '
                    f'per_step {[round(v, 4) for v in val["per_step"]]}'
                    + (
                        f'  per_cycle '
                        f'{[round(v, 4) for v in val["per_cycle"]]}'
                        if 'per_cycle' in val
                        else ''
                    ),
                    flush=True,
                )
                write_record({'step': step, 'split': 'val', **val})
                save({'val': val, 'pretrain': pretrain})

            if step >= args.steps:
                break

    save({'pretrain': pretrain})
    print(
        f'done in {(time.perf_counter() - t0) / 60:.1f} min -> {out_dir}',
        flush=True,
    )


if __name__ == '__main__':
    main()