tests/test_trainer.py

from __future__ import annotations

from pathlib import Path
from types import SimpleNamespace

from dovla_cil.data.schema import RewardInfo
from dovla_cil.generation.pipeline import generate_cil_dataset
from dovla_cil.tasks.library import built_in_toy_tasks
from dovla_cil.training.losses import InterventionalLossWeights
from dovla_cil.training.trainer import (
    DoVLATrainer,
    TrainerConfig,
    _best_records_by_group,
    _coerce_policy_target_action,
    _cross_state_pair_indices,
    _load_policy_target_action_map,
    _reward_utility_values,
)
from dovla_cil.utils.io import read_json


def test_trainer_runs_one_epoch_and_writes_checkpoints(tmp_path: Path) -> None:
    dataset_dir = tmp_path / "cil"
    run_dir = tmp_path / "run"
    generate_cil_dataset(
        backend="toy",
        tasks=built_in_toy_tasks()[:3],
        out_dir=dataset_dir,
        num_states_per_task=2,
        k=4,
        seed=5,
        shard_size=8,
        inline_observations=True,
    )

    result = DoVLATrainer(
        TrainerConfig(
            dataset_dir=dataset_dir,
            output_dir=run_dir,
            epochs=1,
            batch_groups=2,
            records_per_group=4,
            hidden_dim=64,
            learning_rate=1e-3,
            seed=5,
            device="cpu",
        )
    ).train()

    assert (run_dir / "latest.pt").exists()
    assert (run_dir / "best.pt").exists()
    assert (run_dir / "best_policy.pt").exists()
    assert "rank_acc" in result["history"][0]["val"]
    assert "bc_loss" in result["best_policy"]
    metrics = read_json(run_dir / "metrics.json")
    assert "rank_acc" in metrics["history"][0]["val"]
    assert "best_policy" in metrics


def test_trainer_can_supervise_typed_proposal_head(tmp_path: Path) -> None:
    dataset_dir = tmp_path / "cil"
    run_dir = tmp_path / "run"
    generate_cil_dataset(
        backend="toy",
        tasks=built_in_toy_tasks()[:2],
        out_dir=dataset_dir,
        num_states_per_task=2,
        k=4,
        seed=9,
        shard_size=8,
        inline_observations=True,
    )

    result = DoVLATrainer(
        TrainerConfig(
            dataset_dir=dataset_dir,
            output_dir=run_dir,
            epochs=1,
            batch_groups=2,
            records_per_group=4,
            hidden_dim=64,
            learning_rate=1e-3,
            seed=9,
            device="cpu",
            proposal_types=("expert", "near_miss"),
            losses=InterventionalLossWeights(proposal=1.0),
        )
    ).train()

    assert "proposal_loss" in result["history"][0]["val"]
    resolved = read_json(run_dir / "resolved_config.json")
    assert resolved["proposal_types"] == ["expert", "near_miss"]


def test_field_utility_includes_terminal_success_bonus() -> None:
    records = [
        SimpleNamespace(
            reward=RewardInfo(
                progress=0.4,
                success=False,
                terminal_success=False,
            )
        ),
        SimpleNamespace(
            reward=RewardInfo(
                progress=0.4,
                success=True,
                terminal_success=True,
            )
        ),
    ]

    assert _reward_utility_values(records) == [0.4, 1.4]


def test_cross_state_pairs_preserve_task_and_reward_order() -> None:
    records = [
        SimpleNamespace(
            task_id="pick",
            group_id=f"g{group}",
            reward=SimpleNamespace(score=reward),
        )
        for group, reward in ((0, 0.1), (0, 0.9), (1, 0.2), (1, 0.8), (2, 0.4))
    ]

    pairs = _cross_state_pair_indices(records, pair_count=12, seed=7)

    assert len(pairs) == 12
    for better, worse in pairs:
        assert records[better].task_id == records[worse].task_id
        assert records[better].group_id != records[worse].group_id
        assert records[better].reward.score > records[worse].reward.score


def test_cross_state_scope_rejects_lattice_field_objective(tmp_path: Path) -> None:
    try:
        TrainerConfig(
            dataset_dir=tmp_path,
            output_dir=tmp_path / "out",
            objective="lattice_field",
            pair_scope="cross_state",
        )
    except ValueError as exc:
        assert "legacy objective" in str(exc)
    else:  # pragma: no cover - protects baseline semantics
        raise AssertionError(
            "cross-state pairs cannot silently leave same-state field edges active"
        )


def test_policy_target_type_filter_selects_best_allowed_candidate() -> None:
    records = [
        SimpleNamespace(
            group_id="g0",
            candidate_type="expert",
            reward=SimpleNamespace(score=2.0),
            rank_within_group=0,
            record_id="expert",
        ),
        SimpleNamespace(
            group_id="g0",
            candidate_type="near_miss",
            reward=SimpleNamespace(score=1.5),
            rank_within_group=1,
            record_id="near",
        ),
        SimpleNamespace(
            group_id="g1",
            candidate_type="expert",
            reward=SimpleNamespace(score=1.0),
            rank_within_group=0,
            record_id="fallback",
        ),
    ]

    selected = _best_records_by_group(records, candidate_types=("near_miss",))

    assert {record.group_id: record.record_id for record in selected} == {
        "g0": "near",
        "g1": "fallback",
    }


def test_policy_target_map_overrides_group_target_with_fallback() -> None:
    records = [
        SimpleNamespace(
            group_id="g0",
            candidate_type="expert",
            reward=SimpleNamespace(score=2.0),
            rank_within_group=0,
            record_id="expert",
        ),
        SimpleNamespace(
            group_id="g0",
            candidate_type="near_miss",
            reward=SimpleNamespace(score=1.5),
            rank_within_group=1,
            record_id="field_choice",
        ),
        SimpleNamespace(
            group_id="g1",
            candidate_type="near_miss",
            reward=SimpleNamespace(score=0.7),
            rank_within_group=1,
            record_id="fallback",
        ),
    ]

    selected = _best_records_by_group(
        records,
        candidate_types=("near_miss",),
        target_record_ids={"g0": "field_choice"},
    )

    assert {record.group_id: record.record_id for record in selected} == {
        "g0": "field_choice",
        "g1": "fallback",
    }


def test_policy_target_action_map_loads_continuous_targets(tmp_path: Path) -> None:
    path = tmp_path / "targets.json"
    path.write_text(
        """
        {
          "targets": {
            "g0": {"record_id": "r0", "action_values": [[0.1, 0.2], [0.3, 0.4]]},
            "g1": "legacy_record_id"
          }
        }
        """
    )

    loaded = _load_policy_target_action_map(path)

    assert loaded == {"g0": [[0.1, 0.2], [0.3, 0.4]]}
    assert _coerce_policy_target_action(
        loaded["g0"],
        action_dim=3,
        action_horizon=3,
    ) == [[0.1, 0.2, 0.0], [0.3, 0.4, 0.0], [0.0, 0.0, 0.0]]