tests/test_losses.py

from __future__ import annotations

import pytest

from dovla_cil.training.losses import (
    CompositeLoss,
    behavior_cloning_loss,
    causal_contrastive_loss,
    effect_prediction_loss,
    language_minimal_pair_loss,
    lattice_cycle_residual,
    lattice_field_loss,
    pairwise_ranking_loss,
    progress_loss,
    regret_loss,
    regret_targets,
    same_state_pairwise_ranking_loss,
    success_loss,
)


def test_ranking_loss_prefers_correct_order() -> None:
    rewards = [1.0, 0.0, -1.0]
    good_scores = [2.0, 1.0, 0.0]
    bad_scores = [0.0, 1.0, 2.0]
    good_loss = same_state_pairwise_ranking_loss(good_scores, rewards)
    bad_loss = same_state_pairwise_ranking_loss(bad_scores, rewards)
    assert good_loss < bad_loss


def test_ranking_loss_ignores_ties() -> None:
    assert same_state_pairwise_ranking_loss([0.0, 1.0], [1.0, 1.0]) == 0.0


def test_ranking_loss_checks_shape() -> None:
    with pytest.raises(ValueError):
        same_state_pairwise_ranking_loss([0.0], [0.0, 1.0])


def test_regret_targets() -> None:
    assert regret_targets([1.0, 0.25, -1.0]) == [0.0, 0.75, 2.0]


def test_pairwise_ranking_loss_lower_when_order_is_correct() -> None:
    good = pairwise_ranking_loss([2.0, 1.0], [0.0, 0.0], [1.0, 1.0], [0.0, -1.0])
    bad = pairwise_ranking_loss([0.0, 0.0], [2.0, 1.0], [1.0, 1.0], [0.0, -1.0])
    assert good < bad


def test_regret_loss_zero_when_exact() -> None:
    assert regret_loss([0.0, 0.5, 1.0], [0.0, 0.5, 1.0]) == 0.0


def test_behavior_cloning_loss_works() -> None:
    assert behavior_cloning_loss([1.0, 2.0], [1.0, 4.0]) == pytest.approx(2.0)


def test_effect_prediction_loss_combines_continuous_and_binary_terms() -> None:
    loss = effect_prediction_loss(
        {"continuous": [0.0, 1.0], "binary_logits": [0.0]},
        {"continuous": [0.0, 3.0], "binary": [1.0]},
    )
    assert loss > 0.0


def test_success_loss_prefers_correct_logits() -> None:
    assert success_loss([3.0], [1.0]) < success_loss([-3.0], [1.0])


def test_progress_loss_works() -> None:
    assert progress_loss([0.5], [0.5]) == pytest.approx(0.0)


def test_contrastive_loss_finite() -> None:
    loss = causal_contrastive_loss(
        [[1.0, 0.0]],
        [[1.0, 0.0]],
        [[0.0, 1.0]],
        temperature=0.1,
    )
    assert float(loss) >= 0.0


def test_language_minimal_pair_loss_pushes_and_pulls() -> None:
    close_different = language_minimal_pair_loss([[0.0, 0.0]], [[0.1, 0.0]], [True], margin=1.0)
    far_different = language_minimal_pair_loss([[0.0, 0.0]], [[2.0, 0.0]], [True], margin=1.0)
    same_identical = language_minimal_pair_loss([[0.0, 0.0]], [[0.0, 0.0]], [False], margin=1.0)
    same_apart = language_minimal_pair_loss([[0.0, 0.0]], [[1.0, 0.0]], [False], margin=1.0)
    assert far_different < close_different
    assert same_identical < same_apart


def test_composite_returns_components_and_total() -> None:
    output = CompositeLoss()(
        predictions={
            "pred_action": [1.0, 2.0],
            "pred_regret": [0.0, 1.0],
            "pred_scores_i": [2.0],
            "pred_scores_j": [0.0],
        },
        targets={
            "target_action": [1.0, 3.0],
            "target_regret": [0.0, 1.0],
            "rewards_i": [1.0],
            "rewards_j": [0.0],
        },
    )
    assert set(output) >= {"total", "bc", "rank", "regret"}
    assert float(output["total"]) >= 0.0


def test_lattice_field_loss_is_invariant_to_state_reward_offsets() -> None:
    torch = pytest.importorskip("torch")
    potential = torch.tensor([0.2, -0.1, 0.7, 0.0])
    utility = torch.tensor([0.8, 0.3, 0.6, 0.1])
    effect = torch.tensor([[0.0], [0.2], [0.5], [0.1]])
    target_effect = torch.tensor([[0.1], [0.4], [0.6], [0.0]])
    group_ids = ["state-a", "state-a", "state-b", "state-b"]

    base = lattice_field_loss(potential, utility, effect, target_effect, group_ids)
    shifted = lattice_field_loss(
        potential,
        utility + torch.tensor([17.0, 17.0, -9.0, -9.0]),
        effect,
        target_effect,
        group_ids,
    )

    assert torch.allclose(base["potential"], shifted["potential"])
    assert base["edge_count"] == shifted["edge_count"] == 2


def test_lattice_field_is_zero_under_groupwise_gauge_shifts() -> None:
    torch = pytest.importorskip("torch")
    utility = torch.tensor([0.8, 0.3, 0.6, 0.1])
    target_effect = torch.tensor([[0.1, 0.2], [0.4, 0.0], [0.6, -0.2], [0.0, 0.3]])
    group_ids = ["state-a", "state-a", "state-b", "state-b"]
    potential = utility + torch.tensor([5.0, 5.0, -2.0, -2.0])
    predicted_effect = target_effect + torch.tensor(
        [[1.0, -3.0], [1.0, -3.0], [-4.0, 2.0], [-4.0, 2.0]]
    )

    loss = lattice_field_loss(
        potential,
        utility,
        predicted_effect,
        target_effect,
        group_ids,
    )

    assert float(loss["potential"]) == pytest.approx(0.0, abs=1e-7)
    assert float(loss["effect"]) == pytest.approx(0.0, abs=1e-7)


def test_lattice_field_orientation_penalizes_reversed_edge_order() -> None:
    torch = pytest.importorskip("torch")
    utility = torch.tensor([1.0, 0.0])
    effect = torch.zeros((2, 2))
    correct = lattice_field_loss(
        torch.tensor([1.0, 0.0]),
        utility,
        effect,
        effect,
        ["state", "state"],
    )
    reversed_order = lattice_field_loss(
        torch.tensor([0.0, 1.0]),
        utility,
        effect,
        effect,
        ["state", "state"],
    )

    assert float(correct["potential"]) == pytest.approx(0.0, abs=1e-7)
    assert float(reversed_order["orientation"]) > 0.0
    assert float(reversed_order["potential"]) > float(correct["potential"])
    assert float(reversed_order["preference"]) > float(correct["preference"])


def test_lattice_field_preference_is_group_offset_invariant() -> None:
    torch = pytest.importorskip("torch")
    potential = torch.tensor([0.4, -0.2, 2.0, 1.7])
    utility = torch.tensor([0.9, 0.1, 0.8, 0.2])
    effect = torch.zeros((4, 2))
    group_ids = ["a", "a", "b", "b"]

    base = lattice_field_loss(potential, utility, effect, effect, group_ids)
    shifted = lattice_field_loss(
        potential,
        utility + torch.tensor([10.0, 10.0, -4.0, -4.0]),
        effect,
        effect,
        group_ids,
    )

    assert torch.allclose(base["preference"], shifted["preference"])


def test_scalar_potential_has_zero_cycle_residual() -> None:
    torch = pytest.importorskip("torch")
    residual = lattice_cycle_residual(torch.tensor([0.3, -1.2, 2.4]), [[0, 1, 2]])
    assert float(residual) == pytest.approx(0.0, abs=1e-7)