icassp2027-experiment/test_easycom_head_objectives.py

"""Focused tests for EasyCom loss/output-space ablations."""

from __future__ import annotations

import math
import unittest

import torch

from easycom_head import (
    TurnPredictionHead,
    activity_targets_to_four_states,
    checkpoint_head_config,
    checkpoint_loss_name,
    expected_transition_reward,
    four_state_multihorizon_loss,
    independent_bce_multihorizon_loss,
    marginal_probabilities_from_logits,
    output_classes_for_loss,
    transition_aware_objective,
    transition_targets_from_activity,
    weighted_transition_bce_loss,
)


class ObjectiveTests(unittest.TestCase):
    def setUp(self) -> None:
        self.targets = torch.tensor(
            [
                [1.0, 0.0, 0.0],
                [0.0, 1.0, 0.0],
                [0.0, 0.0, 1.0],
                [0.0, 1.0, 1.0],
            ]
        )

    def test_activity_target_mapping(self) -> None:
        actual = activity_targets_to_four_states(self.targets)
        self.assertTrue(torch.equal(actual, torch.tensor([0, 1, 2, 3])))

    def test_invalid_activity_target_is_rejected(self) -> None:
        with self.assertRaisesRegex(ValueError, "must be one of"):
            activity_targets_to_four_states(torch.tensor([[0.0, 0.0, 0.0]]))

    def test_independent_bce_uses_standard_mean(self) -> None:
        logits = torch.zeros_like(self.targets)
        loss = independent_bce_multihorizon_loss(logits, self.targets)
        self.assertAlmostEqual(float(loss), math.log(2.0), places=6)

    def test_four_state_loss_accepts_three_label_targets(self) -> None:
        logits = torch.zeros(4, 4)
        loss = four_state_multihorizon_loss(logits, self.targets)
        self.assertAlmostEqual(float(loss), math.log(4.0), places=6)

    def test_four_state_probabilities_are_marginalized(self) -> None:
        state_probabilities = torch.tensor([[0.1, 0.2, 0.3, 0.4]])
        logits = state_probabilities.log()
        actual = marginal_probabilities_from_logits(logits, "four-state-ce")
        expected = torch.tensor([[0.1, 0.6, 0.7]])
        self.assertTrue(torch.allclose(actual, expected, atol=1e-6))
        self.assertGreater(float(actual.sum()), 1.0)

    def test_bce_probabilities_use_sigmoid(self) -> None:
        actual = marginal_probabilities_from_logits(
            torch.zeros(2, 3),
            "independent-bce",
        )
        self.assertTrue(torch.equal(actual, torch.full((2, 3), 0.5)))

    def test_four_state_head_shape(self) -> None:
        head = TurnPredictionHead(
            hidden_size=8,
            num_horizons=2,
            num_classes=output_classes_for_loss("four-state-ce"),
            dropout=0.0,
        )
        self.assertEqual(tuple(head(torch.randn(3, 8)).shape), (3, 2, 4))

    def test_transition_targets_mark_adjacent_activity_changes(self) -> None:
        sequence = torch.tensor(
            [
                [1.0, 0.0, 0.0],
                [1.0, 0.0, 0.0],
                [0.0, 1.0, 0.0],
                [0.0, 1.0, 0.0],
                [0.0, 1.0, 1.0],
                [0.0, 0.0, 1.0],
            ]
        )
        actual = transition_targets_from_activity(sequence)
        expected = torch.tensor([0.0, 1.0, 0.0, 1.0, 1.0])
        self.assertTrue(torch.equal(actual, expected))

    def test_transition_auxiliary_head_shapes(self) -> None:
        head = TurnPredictionHead(
            hidden_size=8,
            num_horizons=2,
            dropout=0.0,
            transition_auxiliary=True,
        )
        activity, transition = head.forward_with_transition(torch.randn(3, 8))
        self.assertEqual(tuple(activity.shape), (3, 2, 3))
        self.assertEqual(tuple(transition.shape), (3, 2))
        self.assertEqual(tuple(head(torch.randn(3, 8)).shape), (3, 2, 3))

    def test_transition_bce_uses_positive_weight_three(self) -> None:
        logits = torch.zeros(2)
        targets = torch.tensor([0.0, 1.0])
        actual = weighted_transition_bce_loss(logits, targets, pos_weight=3.0)
        self.assertAlmostEqual(float(actual), 2.0 * math.log(2.0), places=6)

    def test_expected_transition_reward_is_exact_weighted_correctness(self) -> None:
        logits = torch.tensor([2.0, -2.0])
        targets = torch.tensor([1.0, 0.0])
        actual = expected_transition_reward(logits, targets, pos_weight=3.0)
        self.assertAlmostEqual(
            float(actual),
            float(torch.sigmoid(torch.tensor(2.0))),
            places=6,
        )

    def test_r1_and_r2_objective_formulas(self) -> None:
        activity_loss = torch.tensor(2.0)
        logits = torch.zeros(2, requires_grad=True)
        targets = torch.tensor([0.0, 1.0])
        transition_loss = weighted_transition_bce_loss(
            logits,
            targets,
            pos_weight=3.0,
        )
        r1 = transition_aware_objective(
            activity_loss,
            logits,
            targets,
            objective="r1-supervised",
        )
        r2 = transition_aware_objective(
            activity_loss,
            logits,
            targets,
            objective="r2-expected-reward",
        )
        self.assertAlmostEqual(
            float(r1["total"].detach()),
            2.0 + 0.2 * float(transition_loss.detach()),
            places=6,
        )
        self.assertAlmostEqual(
            float(r2["total"].detach()),
            2.0 + 0.2 * float(transition_loss.detach()) - 0.1 * 0.5,
            places=6,
        )

    def test_all_new_objectives_backpropagate_to_marginals(self) -> None:
        states = torch.randn(4, 8)
        future_targets = self.targets[:, None, :].repeat(1, 2, 1)
        for loss_name, loss_function in (
            ("independent-bce", independent_bce_multihorizon_loss),
            ("four-state-ce", four_state_multihorizon_loss),
        ):
            with self.subTest(loss=loss_name):
                head = TurnPredictionHead(
                    hidden_size=8,
                    num_horizons=2,
                    num_classes=output_classes_for_loss(loss_name),
                    dropout=0.0,
                )
                logits = head(states)
                loss = loss_function(logits, future_targets)
                loss.backward()
                scores = marginal_probabilities_from_logits(logits, loss_name)
                self.assertEqual(tuple(scores.shape), (4, 2, 3))
                self.assertTrue(torch.isfinite(scores).all())
                self.assertIsNotNone(head.net[-1].weight.grad)

    def test_legacy_checkpoint_defaults_to_lstr(self) -> None:
        self.assertEqual(checkpoint_loss_name({}), "lstr-mce")
        self.assertFalse(checkpoint_head_config({})["transition_auxiliary"])

    def test_checkpoint_rejects_loss_output_mismatch(self) -> None:
        payload = {
            "training_config": {"loss": "four-state-ce"},
            "head_config": {"num_classes": 3},
        }
        with self.assertRaisesRegex(ValueError, "requires 4"):
            checkpoint_loss_name(payload)


if __name__ == "__main__":
    unittest.main()