"""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()