code/scripts/test_planner.py

"""Tests for the recursive planner. Run: ``python -m pytest`` or directly.

The first three tests cover failures that are *silent* — the loss still falls,
training still looks healthy, and the resulting controller is quietly wrong:

* A stray ``detach()`` on the ``h``-chain turns the planner into a greedy
  one-step policy. Caught by ``test_h_chain_is_differentiable``.
* A leaked gradient into the frozen world model would let the planner "fix"
  the simulator instead of the plan. Caught by ``test_world_model_stays_frozen``.
* A missing Change-1 detach schedule triples retained activations and makes the
  90-deep tied recursion untrainable. Caught by
  ``test_retained_activation_count``.

The remaining tests pin the action/frame alignment against the phase-1 rollout
(an off-by-one there trains on misaligned transitions and also fails silently)
and check that one optimizer step actually reduces the loss.

These run on a tiny random stand-in for LeWM by default so they need no
checkpoint and no GPU. Pass ``--real`` to run the alignment test against the
actual frozen checkpoint as well.
"""

import sys
from pathlib import Path

import pytest
import torch
from torch import nn

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

from lejepa_control.rollout import rollout_contexts, rollout_plan  # noqa: E402
from lejepa_control_2.losses import DistanceScale, planner_loss  # noqa: E402
from lejepa_control_2.planner import RecursivePlanner  # noqa: E402

D, N, A, W = 192, 3, 10, 32
INNER, CYCLES, HORIZON, BATCH = 6, 3, 5, 2


class ToyWorldModel(nn.Module):
    """Stand-in with LeWM's exact interface: ``action_encoder`` + ``predict``.

    Deliberately not an identity map — it mixes the whole context window and
    the whole action window, so an alignment error changes the output instead
    of cancelling out.
    """

    def __init__(self, latent_dim=D, num_context=N, action_dim=A, seed=0):
        super().__init__()
        torch.manual_seed(seed)
        self.action_encoder = nn.Sequential(
            nn.Linear(action_dim, latent_dim), nn.SiLU()
        )
        self.mix_emb = nn.Linear(num_context * latent_dim, latent_dim)
        self.mix_act = nn.Linear(num_context * latent_dim, latent_dim)

        class _P:
            num_frames = num_context
            input_dim = latent_dim

        self.predictor = _P()

    def predict(self, emb, act_emb):
        h = torch.tanh(self.mix_emb(emb.flatten(1)) + self.mix_act(act_emb.flatten(1)))
        # LeWM returns one prediction per context position; only [:, -1] is used
        return h.unsqueeze(1).expand(-1, emb.size(1), -1)


def make_planner(**kw):
    torch.manual_seed(0)
    opts = dict(
        latent_dim=D,
        num_context=N,
        width=W,
        hidden=2 * W,
        inner=INNER,
        cycles=CYCLES,
        horizon=HORIZON,
    )
    opts.update(kw)
    return RecursivePlanner(**opts)


def make_batch(batch=BATCH, seed=1):
    g = torch.Generator().manual_seed(seed)
    return {
        'context': torch.randn(batch, N, D, generator=g),
        'past_actions': torch.randn(batch, N - 1, A, generator=g) * 0.3,
        'goal': torch.randn(batch, D, generator=g),
        'goal_offset': torch.randint(1, HORIZON + 1, (batch,), generator=g),
    }


# --------------------------------------------------------------------------
# invariant 1: the h-chain stays differentiable
# --------------------------------------------------------------------------


def test_h_chain_is_differentiable():
    """``d[-1]`` must carry grad and the graph must reach ``blocks[0]``.

    If a stray detach cuts the state chain the loss still trains — it just
    silently becomes a greedy one-step policy, because credit for the first
    block can no longer arrive from the terminal distance.
    """
    model, planner = ToyWorldModel(), make_planner()
    batch = make_batch()

    out = planner(
        model, batch['context'], batch['past_actions'], batch['goal']
    )
    d, blocks = out['distance_seq'], out['block_seq']
    terminal = d[-1]

    assert terminal.requires_grad, 'terminal distance carries no gradient'

    # the real check: does the terminal distance depend on the FIRST block?
    grad = torch.autograd.grad(
        terminal.sum(), blocks[0], retain_graph=True, allow_unused=True
    )[0]
    assert grad is not None, (
        'terminal distance does not reach blocks[0] — the h-chain is cut and '
        'the planner has collapsed to a greedy one-step policy'
    )
    assert grad.abs().sum() > 0, 'blocks[0] receives an all-zero gradient'

    # and every intermediate step should be reachable too
    for k in range(HORIZON):
        g = torch.autograd.grad(
            terminal.sum(), blocks[k], retain_graph=True, allow_unused=True
        )[0]
        assert g is not None and g.abs().sum() > 0, f'block {k} unreachable'

    # gradient magnitude should not collapse across the horizon; a chain that
    # is intact but vanishing is nearly as bad as one that is cut
    norms = [
        torch.autograd.grad(
            terminal.sum(), blocks[k], retain_graph=True
        )[0].norm().item()
        for k in range(HORIZON)
    ]
    assert min(norms) > 0.01 * max(norms), f'gradient collapses: {norms}'


def test_planner_parameters_receive_gradient():
    """``f``, ``g``, ``phi`` and ``psi`` must all be trained by the loss."""
    model, planner = ToyWorldModel(), make_planner()
    batch = make_batch()
    scale = DistanceScale()
    scale.update(batch['context'], batch['goal'])

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

    # y0 is unused when warm start is on. z0 is consumed by cycle 1, which is
    # inside the no_grad region whenever T > 1 — so under the Change-1 schedule
    # the initial scratchpad is untrainable by construction, not by mistake.
    # test_initial_scratchpad_trains_at_one_cycle pins that it is wired at all.
    untrained = {'y0', 'z0'}
    for name, p in planner.named_parameters():
        if name in untrained:
            continue
        assert p.grad is not None, f'{name} received no gradient'
        assert torch.isfinite(p.grad).all(), f'{name} has non-finite gradient'

    assert planner.f.body.out.weight.grad.abs().sum() > 0, 'f is not training'
    assert planner.g.body.out.weight.grad.abs().sum() > 0, 'g is not training'
    assert (
        planner.action_embed.encode_net[0].weight.grad.abs().sum() > 0
    ), 'phi is not training (the anchor should reach it)'


def test_initial_scratchpad_trains_at_one_cycle():
    """``z0`` must be in the graph when there is no no_grad region to hide it."""
    model, planner = ToyWorldModel(), make_planner(cycles=1)
    batch = make_batch()
    scale = DistanceScale()
    scale.update(batch['context'], batch['goal'])
    out = planner(
        model, batch['context'], batch['past_actions'], batch['goal'], cycles=1
    )
    loss, _ = planner_loss(
        out, batch['goal_offset'], scale, planner.action_embed
    )
    loss.backward()
    assert planner.z0.grad is not None and planner.z0.grad.abs().sum() > 0


# --------------------------------------------------------------------------
# invariant 2: the world model is frozen but transparent
# --------------------------------------------------------------------------


def test_world_model_stays_frozen():
    """Gradients flow THROUGH the world model, never INTO it."""
    model = ToyWorldModel()
    model.requires_grad_(False)
    planner = make_planner()
    batch = make_batch()
    scale = DistanceScale()
    scale.update(batch['context'], batch['goal'])

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

    assert next(model.parameters()).grad is None, (
        'the frozen world model accumulated a gradient'
    )
    for name, p in model.named_parameters():
        assert p.grad is None, f'world model {name} accumulated a gradient'

    # transparent, though: the plan still got its signal
    assert planner.f.body.out.weight.grad.abs().sum() > 0


# --------------------------------------------------------------------------
# invariant 3: the Change-1 detach schedule is actually in effect
# --------------------------------------------------------------------------


def test_retained_activation_count():
    """Grad-enabled ``f``/``g`` applications must be ``H(n+1)``, not ``H*T*(n+1)``.

    ``H(n+1) = 35`` at ``H=5, n=6``; without the detach schedule it would be
    ``H*T*(n+1) = 105``. That 3x is the memory and the vanishing/exploding
    gradient the whole of Change 1 exists to avoid.
    """
    model, planner = ToyWorldModel(), make_planner()
    batch = make_batch()

    planner.reset_call_counts()
    planner(model, batch['context'], batch['past_actions'], batch['goal'])

    expected = HORIZON * (INNER + 1)
    unstaged = HORIZON * CYCLES * (INNER + 1)
    assert expected == 35 and unstaged == 105, 'test constants drifted'
    assert planner.grad_calls == expected, (
        f'retained recursion applications = {planner.grad_calls}, expected '
        f'{expected} = H(n+1). Full backprop would be {unstaged}; a count '
        f'near that means the Change-1 no_grad region is not in effect.'
    )


def test_no_grad_cycles_produce_no_graph():
    """The lookahead distances from cycles 1..T-1 must be constants."""
    model, planner = ToyWorldModel(), make_planner()
    batch = make_batch()
    out = planner(
        model, batch['context'], batch['past_actions'], batch['goal']
    )
    cycles = out['cycle_distances']
    assert cycles.shape == (BATCH, HORIZON, CYCLES)
    # only the last cycle of each horizon step is the committed transition
    assert out['distances'].allclose(cycles[:, :, -1])


def test_horizon_carry_is_detached():
    """``y`` and ``z`` must not backprop across horizon steps."""
    model, planner = ToyWorldModel(), make_planner(horizon=2)
    batch = make_batch()
    out = planner(
        model, batch['context'], batch['past_actions'], batch['goal']
    )
    # step 1's answer must not depend on step 0's answer through the recursion
    grad = torch.autograd.grad(
        out['answers'][:, 1].sum(),
        out['answers'][:, 0],
        retain_graph=True,
        allow_unused=True,
    )[0]
    assert grad is None or grad.abs().sum() == 0, (
        'the recursion carry is not detached across horizon steps'
    )


# --------------------------------------------------------------------------
# action/frame alignment
# --------------------------------------------------------------------------


def test_alignment_matches_phase1_rollout():
    """The driver's ``m_step`` chain must equal ``rollout_plan`` exactly.

    Block ``k`` is the block leaving context frame ``k``; with ``N`` context
    frames there are ``N-1`` past blocks between them and the current frame
    pairs with the first block of the plan. Getting this off by one trains on
    misaligned transitions and fails silently.
    """
    model, planner = ToyWorldModel(), make_planner()
    batch = make_batch()

    out = planner(
        model, batch['context'], batch['past_actions'], batch['goal']
    )
    # replay the planner's own committed blocks through the phase-1 rollout
    pred, frames = rollout_plan(
        model,
        batch['context'],
        batch['past_actions'],
        out['blocks'].detach(),
        return_frames=True,
    )

    assert torch.allclose(pred, out['frames'][:, N:].detach(), atol=1e-5), (
        'planner rollout disagrees with rollout_plan — action/frame alignment '
        'is off by one'
    )
    assert torch.allclose(frames, out['frames'].detach(), atol=1e-5)
    assert torch.allclose(
        rollout_contexts(frames, N), out['contexts'].detach(), atol=1e-5
    ), 'per-step context windows disagree with rollout_contexts'


def test_alignment_is_order_sensitive():
    """A permuted plan must change the rollout, or the test above is vacuous."""
    model, planner = ToyWorldModel(), make_planner()
    batch = make_batch()
    out = planner(
        model, batch['context'], batch['past_actions'], batch['goal']
    )
    blocks = out['blocks'].detach()
    shifted = torch.roll(blocks, 1, dims=1)
    a = rollout_plan(model, batch['context'], batch['past_actions'], blocks)
    b = rollout_plan(model, batch['context'], batch['past_actions'], shifted)
    assert not torch.allclose(a, b, atol=1e-4), (
        'the toy world model is insensitive to block order; the alignment '
        'test would pass even with a broken driver'
    )


@pytest.mark.parametrize('offset', [1, 3, 5])
def test_arrival_term_follows_the_offset(offset):
    """The arrival term must be indexed by ``q``, not pinned at ``H``."""
    model, planner = ToyWorldModel(), make_planner()
    batch = make_batch()
    scale = DistanceScale()
    scale.update(batch['context'], batch['goal'])
    out = planner(
        model, batch['context'], batch['past_actions'], batch['goal']
    )
    q = torch.full((BATCH,), offset, dtype=torch.long)
    _, metrics = planner_loss(out, q, scale, planner.action_embed)

    # d_q + w_hold * mean(d_{q+1..H}); at q = H there is nothing left to hold
    d = scale.normalize(out['distances'])
    hold = (
        d[:, offset:].mean(dim=1)
        if offset < HORIZON
        else torch.zeros(BATCH)
    )
    expected = d[:, offset - 1] + 0.5 * hold
    assert torch.allclose(metrics['arrival'], expected.mean(), atol=1e-5)


# --------------------------------------------------------------------------
# end to end
# --------------------------------------------------------------------------


def test_one_optimizer_step_reduces_loss():
    """A tiny config must be able to overfit a fixed batch."""
    torch.manual_seed(0)
    model = ToyWorldModel()
    model.requires_grad_(False)
    planner = make_planner(horizon=3, cycles=2)
    batch = make_batch()
    scale = DistanceScale()

    opt = torch.optim.AdamW(planner.parameters(), lr=1e-2, weight_decay=1e-4)

    losses = []
    for _ in range(20):
        scale.update(batch['context'], batch['goal'])
        out = planner(
            model,
            batch['context'],
            batch['past_actions'],
            batch['goal'],
            horizon=3,
        )
        loss, _ = planner_loss(
            out, batch['goal_offset'], scale, planner.action_embed
        )
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(planner.parameters(), 1.0)
        opt.step()
        losses.append(loss.item())

    assert losses[-1] < losses[0], (
        f'loss did not fall over 20 steps: {losses[0]:.4f} -> {losses[-1]:.4f}'
    )


def test_shapes_and_finiteness():
    model, planner = ToyWorldModel(), make_planner()
    batch = make_batch()
    out = planner(
        model, batch['context'], batch['past_actions'], batch['goal']
    )
    assert out['distances'].shape == (BATCH, HORIZON)
    assert out['cycle_distances'].shape == (BATCH, HORIZON, CYCLES)
    assert out['blocks'].shape == (BATCH, HORIZON, A)
    assert out['raw'].shape == (BATCH, HORIZON, A)
    assert out['contexts'].shape == (BATCH, HORIZON, N, D)
    assert out['answers'].shape == (BATCH, HORIZON, W)
    assert out['frames'].shape == (BATCH, N + HORIZON, D)
    for k, v in out.items():
        if torch.is_tensor(v):
            assert torch.isfinite(v).all(), f'{k} is not finite'


def test_actions_respect_the_tanh_bound():
    """``psi`` must keep blocks inside the normalized action range."""
    center, scale_ = -0.5, 4.8
    model = ToyWorldModel()
    planner = make_planner(action_center=center, action_scale=scale_)
    batch = make_batch()
    out = planner(
        model, batch['context'], batch['past_actions'], batch['goal']
    )
    b = out['blocks']
    assert (b <= center + scale_ + 1e-4).all()
    assert (b >= center - scale_ - 1e-4).all()


def test_feedback_off_produces_no_cycle_distances():
    """Without Change 2 there is no lookahead, so nothing to deep-supervise."""
    model = ToyWorldModel()
    planner = make_planner(use_feedback=False, cycles=1)
    batch = make_batch()
    out = planner(
        model, batch['context'], batch['past_actions'], batch['goal']
    )
    assert out['cycle_distances'] is None


def test_warm_start_uses_the_last_executed_block():
    model = ToyWorldModel()
    planner = make_planner(warm_start=True)
    batch = make_batch()
    expected = planner.action_embed.encode(batch['past_actions'][:, -1])
    got = planner._initial_answer(batch['past_actions'], BATCH)
    assert torch.allclose(got, expected)


if __name__ == '__main__':
    sys.exit(pytest.main([__file__, '-v', '--no-header', '-x']))