code/scripts/test_training.py

"""Tests for the training script's staging, curricula and checkpoint logic.

Separate from ``test_planner.py``, which covers the recursion itself. These
cover the wiring around it — the places where a run can be configured into
something other than what the stage table says, silently.
"""

import sys
from pathlib import Path

import pytest
import torch

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

from lejepa_control.losses import arrival_hold_loss  # noqa: E402
from lejepa_control_2.scripts.train_planner import (  # noqa: E402
    STAGES,
    current_value,
    parse_args,
    parse_curriculum,
)


def test_stage_presets_match_the_bring_up_table():
    """Each stage turns on exactly one more mechanism than the last."""
    a = parse_args(['--stage', 'A'])
    assert a.cycles == 1 and a.use_feedback is False
    assert (a.lambda_cycle, a.lambda_anchor, a.lambda_sat, a.lambda_support) \
        == (0.0, 0.0, 0.0, 0.0)

    b = parse_args(['--stage', 'B'])
    assert b.cycles == 3 and b.use_feedback is True and b.lambda_cycle == 0.3
    assert (b.lambda_anchor, b.lambda_sat, b.lambda_support) == (0.0, 0.0, 0.0)

    c = parse_args(['--stage', 'C'])
    assert (c.lambda_anchor, c.lambda_sat) == (0.05, 1e-3)
    assert c.lambda_support == 0.0

    d = parse_args(['--stage', 'D'])
    assert d.lambda_support == 0.01
    assert max(v for _, v in parse_curriculum(d.horizon_curriculum, 100)) == 5

    # E is D with warm start off — the cold-start control, not a new mechanism
    e = parse_args(['--stage', 'E'])
    assert e.warm_start is False
    for key in ('cycles', 'lambda_cycle', 'lambda_support', 'lambda_anchor',
                'lambda_sat', 'horizon_curriculum'):
        assert getattr(e, key) == getattr(d, key), f'E differs from D in {key}'


def test_explicit_flags_override_the_preset():
    """A preset must not silently win over something typed on the CLI."""
    args = parse_args(['--stage', 'A', '--cycles', '5', '--lambda-cycle', '0.7'])
    assert args.cycles == 5 and args.lambda_cycle == 0.7
    # and untouched keys still come from the preset
    assert args.use_feedback is False


def test_default_run_uses_the_section_10_hyperparameters():
    args = parse_args([])
    assert (args.inner, args.cycles) == (6, 3)
    assert (args.hold_weight, args.alpha) == (0.5, 0.05)
    assert (args.lambda_cycle, args.lambda_support) == (0.3, 0.01)
    assert (args.lambda_anchor, args.lambda_sat) == (0.05, 1e-3)
    assert (args.lr, args.weight_decay, args.grad_clip) == (1e-4, 1e-4, 1.0)
    assert (args.batch_size, args.steps) == (32, 20000)
    assert args.curriculum == '0:2,0.25:3,0.5:5'
    assert parse_curriculum(args.horizon_curriculum, 20000) == [
        (0, 3), (10000, 5)
    ]


def test_curriculum_advances_at_the_right_steps():
    stages = parse_curriculum('0:2,0.25:3,0.5:5', 20000)
    assert stages == [(0, 2), (5000, 3), (10000, 5)]
    assert current_value(stages, 0) == 2
    assert current_value(stages, 4999) == 2
    assert current_value(stages, 5000) == 3
    assert current_value(stages, 19999) == 5


@pytest.mark.parametrize(
    'horizon_spec,step,want_h,want_q',
    [
        # stage D: the two curricula advance together, so the clamp is inert
        ('0:3,0.5:5', 0, 3, 2),
        ('0:3,0.5:5', 6000, 3, 3),
        ('0:3,0.5:5', 12000, 5, 5),
        # stages B and C: H is pinned at 3 while the offset curriculum still
        # climbs to 5. This is where the clamp actually bites.
        ('0:3', 12000, 3, 3),
    ],
)
def test_goal_offset_is_clamped_to_the_live_horizon(
    horizon_spec, step, want_h, want_q
):
    """The offset curriculum must not outrun the horizon being rolled out.

    Stages B and C hold ``H = 3`` while the shared offset curriculum advances
    to 5 at the halfway point. Sampling ``q = 5`` against a 3-step rollout
    would be silently clamped down to 3 inside ``arrival_hold_loss``, so the
    deadline would stop meaning what the curriculum says it means — and the
    dataset would be relabeling goals from 5 transitions ahead that the
    planner has no steps left to reach. The clamp belongs at the sampler.
    """
    horizon = current_value(parse_curriculum(horizon_spec, 20000), step)
    offset_stages = parse_curriculum('0:2,0.25:3,0.5:5', 20000)
    offset = min(current_value(offset_stages, step), horizon)
    assert (horizon, offset) == (want_h, want_q)
    assert offset <= horizon


def test_arrival_term_is_undistorted_once_the_offset_is_clamped():
    """With the clamp in place ``q`` always indexes a step that exists."""
    d = torch.tensor([[0.9, 0.6, 0.3]])
    for q in (1, 2, 3):
        offset = torch.tensor([q])
        got = arrival_hold_loss(d, offset, 0.5).mean()
        hold = d[:, q:].mean(dim=1) if q < 3 else torch.zeros(1)
        assert torch.allclose(got, (d[:, q - 1] + 0.5 * hold).mean())


def test_ablation_switches_are_reachable_but_off_by_default():
    """Section 12's ablations are wired behind flags, not on by default."""
    args = parse_args([])
    assert args.detach_schedule == 'last-cycle'  # ablation 4
    assert args.terminal_only is False           # ablation 5
    assert args.path_weighting == 'late'         # ablation 6
    assert args.use_feedback is True             # ablation 2

    assert parse_args(['--terminal-only']).terminal_only is True
    assert parse_args(['--no-feedback']).use_feedback is False
    assert parse_args(['--path-weighting', 'discount']).path_weighting \
        == 'discount'
    for schedule in ('last-cycle', 'one-step', 'full'):
        assert parse_args(['--detach-schedule', schedule]).detach_schedule \
            == schedule


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