Add overlap-aware turn prediction objectives

test_easycom_head_objectives.py

@@ -0,0 +1,109 @@
+"""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_loss_name,
+    four_state_multihorizon_loss,
+    independent_bce_multihorizon_loss,
+    marginal_probabilities_from_logits,
+    output_classes_for_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_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")
+
+    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()