Add transition-aware R1 and R2 objectives

test_easycom_head_objectives.py

@@ -10,11 +10,16 @@ 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,
 )
 
 
@@ -71,6 +76,81 @@ class ObjectiveTests(unittest.TestCase):
         )
         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)
@@ -95,6 +175,7 @@ class ObjectiveTests(unittest.TestCase):
 
     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 = {