code/scripts/test_controller.py

"""End-to-end checks on the controller stack, with no data dependency.

Verifies the shape contract, that gradients reach the controller but never the
frozen world model, that the tanh bound maps exactly onto raw PushT limits,
and that refinement actually changes the plan.
"""

import sys
from pathlib import Path

import torch

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

from lejepa_control.controller import IterativeController  # noqa: E402
from lejepa_control.losses import (  # noqa: E402
    BehaviorDensity,
    refinement_loss,
    support_loss,
)
from lejepa_control.rollout import rollout_contexts, rollout_plan  # noqa: E402
from lejepa_control.solver import load_controller  # noqa: E402
from lejepa_control.world_model import load_lewm  # noqa: E402

# z-score stats measured on the full expert dataset
A_MEAN = torch.tensor([-0.00781276, 0.00686056])
A_STD = torch.tensor([0.20824118, 0.20649128])

# PushT checkpoint dirs under data/runs/ that predate exp7 (reacher/tworoom
# use different action_dim/frameskip and are out of scope for this regression)
LEGACY_PUSHT_CHECKPOINTS = [
    'ah_hold0.0', 'ah_hold0.5', 'ah_hold1.0',
    'abl_no_support', 'abl_terminal_only', 'controller', 'dryrun',
]


def test_checkpoint_rebuild(device):
    """exp7: adding no_latent_proj/fused to load_controller must not disturb
    any pre-existing checkpoint (``saved.get(..., False)`` keeps them at the
    split/projected architecture they were actually trained with)."""
    root = Path(__file__).resolve().parents[1] / 'data' / 'runs'
    checked = 0
    for name in LEGACY_PUSHT_CHECKPOINTS:
        path = root / name / 'controller.pt'
        if not path.exists():
            continue
        controller, _ = load_controller(str(path), device=device)
        assert controller.fused is False, f'{name}: fused should default False'
        assert controller.no_latent_proj is False, f'{name}: no_latent_proj should default False'
        checked += 1
    assert checked > 0, 'no legacy PushT checkpoints found to regression-test'
    print(f'checkpoint rebuild ok ({checked} checkpoints, strict load, flags default False)')


def test_no_latent_proj_guard():
    """--no-latent-proj with width != latent_dim must raise, clearly."""
    try:
        IterativeController(latent_dim=192, width=256, no_latent_proj=True)
    except AssertionError as e:
        assert 'no-latent-proj' in str(e)
    else:
        raise AssertionError('no_latent_proj with width != latent_dim did not raise')
    print('no_latent_proj width guard ok')


def test_param_counts():
    """Param counts for all four SPACE x OPERATOR cells (exp7 sec 1/7).

    Only the two cells the spec gives explicit numeric bands for (today's
    baseline, and the fused192 candidate) are hard-asserted; A/B are printed
    for the record.
    """
    def count(**kw):
        c = IterativeController(latent_dim=192, **kw)
        return sum(p.numel() for p in c.parameters())

    baseline = count(width=256)                                # today
    cell_a = count(width=192, no_latent_proj=True)              # 192id, split
    cell_b = count(width=256, fused=True)                       # 256+proj, fused
    cell_c = count(width=192, no_latent_proj=True, fused=True)  # candidate

    print(
        f'params baseline(256+proj/split)={baseline / 1e6:.2f}M  '
        f'A(192id/split)={cell_a / 1e6:.2f}M  '
        f'B(256+proj/fused)={cell_b / 1e6:.2f}M  '
        f'C(192id/fused)={cell_c / 1e6:.2f}M'
    )
    assert 6.5e6 <= baseline <= 7.1e6, f'baseline param count out of band: {baseline}'
    # spec sec 1 estimates "~4-4.5M"; verified by hand (net=8 layers@width192
    # is 3.56M of the 3.75M total, both numbers pinned explicitly by the spec
    # text) the true figure is ~3.75M — the spec's figure was a pre-impl
    # approximation that ran ~7-17% high, not an implementation bug. Band
    # widened to the verified value; still catches gross regressions (e.g. a
    # missing/duplicated encoder would be off by ~2x).
    assert 3.5e6 <= cell_c <= 4.0e6, f'fused192 param count out of verified band: {cell_c}'
    print('param counts in expected band (fused192: see comment on the spec\'s ~4-4.5M estimate)')


def test_fused_variant(model, device):
    """Fused-operator forward: shapes, k=0 bypasses delta_head, rows differ."""
    D = model.predictor.input_dim
    N, H, K, B = 3, 5, 3, 4

    controller = IterativeController(
        latent_dim=D,
        horizon=H,
        refinements=K,
        width=D,
        no_latent_proj=True,
        fused=True,
        action_center=(-A_MEAN / A_STD),
        action_scale=(1.0 / A_STD),
    ).to(device).eval()

    ctx = torch.randn(B, N, D, device=device)
    past = torch.randn(B, N - 1, 10, device=device) * 0.5
    goal = torch.randn(B, D, device=device)

    out = controller(model, ctx, past, goal)
    assert len(out['plans']) == K + 1
    assert out['plans'][0].shape == (B, H, 10), out['plans'][0].shape
    assert out['rollouts'][0].shape == (B, H, D)
    assert out['distances'][0].shape == (B, H)
    assert out['contexts'][0].shape == (B, H, N, D)
    print(f'fused shapes ok: plan {tuple(out["plans"][0].shape)} '
          f'rollout {tuple(out["rollouts"][0].shape)}')

    # --- k=0 output does not pass through delta_head -----------------------
    cond = controller.condition(ctx, goal)
    y0_before = controller.initial_plan(cond)
    with torch.no_grad():
        saved_w = controller.delta_head.weight.clone()
        controller.delta_head.weight.fill_(float('nan'))
    y0_after = controller.initial_plan(cond)
    assert torch.equal(y0_before, y0_after), 'k=0 output changed when delta_head was corrupted'
    assert not torch.isnan(y0_after).any(), 'k=0 output leaked delta_head NaNs'
    with torch.no_grad():
        controller.delta_head.weight.copy_(saved_w)
    print('k=0 bypasses delta_head ok')

    # --- per-slot ΔY rows actually differ across slots ----------------------
    pred = torch.randn(B, H, D, device=device)
    delta = controller.refine(y0_after, cond, pred, goal)
    row_diffs = [
        (delta[:, i] - delta[:, j]).abs().mean().item()
        for i in range(H)
        for j in range(i + 1, H)
    ]
    assert min(row_diffs) > 1e-6, f'fused refine rows too similar: {row_diffs}'
    print(f'fused per-slot rows differ ok (min pairwise |diff| {min(row_diffs):.5f})')


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

    model = load_lewm(device=device)
    D = model.predictor.input_dim
    N, H, K, B = 3, 5, 3, 4

    controller = IterativeController(
        latent_dim=D,
        horizon=H,
        refinements=K,
        action_center=(-A_MEAN / A_STD),
        action_scale=(1.0 / A_STD),
    ).to(device)
    print(f'params {sum(p.numel() for p in controller.parameters()) / 1e6:.2f}M')

    ctx = torch.randn(B, N, D, device=device)
    past = torch.randn(B, N - 1, 10, device=device) * 0.5
    goal = torch.randn(B, D, device=device)

    out = controller(model, ctx, past, goal)

    assert len(out['plans']) == K + 1, 'one plan per refinement, plus initial'
    assert out['plans'][0].shape == (B, H, 10), out['plans'][0].shape
    assert out['rollouts'][0].shape == (B, H, D)
    assert out['distances'][0].shape == (B, H)
    assert out['contexts'][0].shape == (B, H, N, D)
    print(f'shapes ok: plan {tuple(out["plans"][0].shape)} '
          f'rollout {tuple(out["rollouts"][0].shape)} '
          f'contexts {tuple(out["contexts"][0].shape)}')

    # --- tanh bound maps exactly onto raw actions in [-1, 1] --------------
    plan = out['plans'][-1].reshape(B, H, 5, 2).cpu()
    raw = plan * A_STD + A_MEAN
    assert raw.abs().max() <= 1.0 + 1e-5, f'raw action out of bounds {raw.abs().max()}'
    lo = (controller.action_center - controller.action_scale).cpu()
    hi = (controller.action_center + controller.action_scale).cpu()
    print(f'raw action range [{raw.min():.3f}, {raw.max():.3f}] within [-1, 1]')
    print(f'  normalized bounds {(lo * A_STD + A_MEAN).tolist()} '
          f'{(hi * A_STD + A_MEAN).tolist()}')

    # --- refinement changes the plan --------------------------------------
    deltas = [
        (out['plans'][k + 1] - out['plans'][k]).abs().mean().item()
        for k in range(K)
    ]
    print(f'plan change per refinement {[round(d, 5) for d in deltas]}')
    assert max(deltas) > 0, 'refinement never changed the plan'

    # --- context windows line up with the rollout -------------------------
    pred, frames = rollout_plan(
        model, ctx, past, out['plans'][-1], return_frames=True
    )
    cw = rollout_contexts(frames, N)
    assert torch.equal(cw[:, 0], ctx), 'first window must be the real context'
    assert torch.allclose(cw[:, 1, -1], pred[:, 0]), 'window must advance by 1'
    print('rollout context windows aligned')

    # --- loss and gradient flow -------------------------------------------
    density = BehaviorDensity(latent_dim=D).to(device)
    loss = refinement_loss(out['distances'], alpha=0.05)
    s_loss, viol = support_loss(
        density,
        torch.cat(out['contexts'][1:], 0).flatten(0, 1),
        torch.cat(out['plans'][1:], 0).flatten(0, 1),
        threshold=0.0,
    )
    total = loss + 0.01 * s_loss
    total.backward()

    trained = [
        n for n, p in controller.named_parameters() if p.grad is not None
    ]
    untrained = [
        n
        for n, p in controller.named_parameters()
        if p.grad is None or p.grad.abs().sum() == 0
    ]
    frozen_grads = [n for n, p in model.named_parameters() if p.grad is not None]

    print(f'refine loss {loss.item():.4f}  support {s_loss.item():.4f}  '
          f'violation rate {viol.item():.2f}')
    print(f'controller params with grad: {len(trained)}, '
          f'without: {len(untrained)}')
    assert not frozen_grads, f'world model got gradients: {frozen_grads[:3]}'
    assert not untrained, f'no gradient reached: {untrained}'

    # --- refinement loss weighting is 2^k ----------------------------------
    fake = [
        torch.full((1, H), float(v), device=device) for v in (1.0, 1.0, 1.0, 0.0)
    ]
    got = refinement_loss(fake, alpha=0.0).item()
    expect = (1 * 1 + 2 * 1 + 4 * 1 + 8 * 0) / 15
    assert abs(got - expect) < 1e-6, (got, expect)
    print(f'refinement weighting ok ({got:.4f} == {expect:.4f})')

    # --- exp7: fused192 architecture ---------------------------------------
    print()
    test_checkpoint_rebuild(device)
    test_no_latent_proj_guard()
    test_param_counts()
    test_fused_variant(model, device)

    print('ALL CONTROLLER TESTS PASSED')


if __name__ == '__main__':
    main()