Add transition-aware R1 and R2 objectives

easycom_head.py

@@ -46,6 +46,11 @@ LOSS_NAMES = (
     "independent-bce",
     "four-state-ce",
 )
+TRANSITION_OBJECTIVE_NAMES = (
+    "none",
+    "r1-supervised",
+    "r2-expected-reward",
+)
 DATASET_NAMES = ("easycom", "ego4d")
 LAYER_FUSION_MODES = (
     "single",
@@ -68,6 +73,7 @@ class TurnPredictionHead(nn.Module):
         num_input_layers: int = 1,
         learnable_layer_mixture: bool = False,
         layer_fusion: str | None = None,
+        transition_auxiliary: bool = False,
     ) -> None:
         super().__init__()
         if num_input_layers <= 0:
@@ -101,6 +107,7 @@ class TurnPredictionHead(nn.Module):
         self.num_input_layers = num_input_layers
         self.layer_fusion = layer_fusion
         self.learnable_layer_mixture = layer_fusion == "scalar-mixture"
+        self.transition_auxiliary = transition_auxiliary
         if self.learnable_layer_mixture:
             # Zero initialization exactly reproduces a uniform layer mean at
             # step zero while allowing training to select useful query layers.
@@ -123,8 +130,15 @@ class TurnPredictionHead(nn.Module):
             nn.Dropout(dropout),
             nn.Linear(hidden_size // 2, num_horizons * num_classes),
         )
+        self.transition_head: nn.Module | None
+        if transition_auxiliary:
+            # Share the activity MLP trunk and add only a small binary readout
+            # for each future horizon.
+            self.transition_head = nn.Linear(hidden_size // 2, num_horizons)
+        else:
+            self.transition_head = None
 
-    def forward(self, states: torch.Tensor) -> torch.Tensor:
+    def _fuse_layers(self, states: torch.Tensor) -> torch.Tensor:
         if states.shape[-1] != self.hidden_size:
             raise ValueError(
                 f"Expected hidden size {self.hidden_size}, got {states.shape[-1]}"
@@ -150,11 +164,40 @@ class TurnPredictionHead(nn.Module):
             assert self.layer_fusion_net is not None
             states = states.flatten(start_dim=-2)
             states = self.layer_fusion_net(states)
-        logits = self.net(states)
+        return states
+
+    def _shared_features(self, states: torch.Tensor) -> torch.Tensor:
+        """Return the activity MLP representation shared by both readouts."""
+
+        states = self._fuse_layers(states)
+        return self.net[:-1](states)
+
+    def forward(self, states: torch.Tensor) -> torch.Tensor:
+        features = self._shared_features(states)
+        logits = self.net[-1](features)
         return logits.view(
-            *states.shape[:-1], self.num_horizons, self.num_classes
+            *features.shape[:-1], self.num_horizons, self.num_classes
         )
 
+    def forward_with_transition(
+        self,
+        states: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return activity logits and binary transition logits."""
+
+        if self.transition_head is None:
+            raise RuntimeError(
+                "This head was created without transition_auxiliary=True"
+            )
+        features = self._shared_features(states)
+        activity_logits = self.net[-1](features).view(
+            *features.shape[:-1],
+            self.num_horizons,
+            self.num_classes,
+        )
+        transition_logits = self.transition_head(features)
+        return activity_logits, transition_logits
+
     def layer_mixture_weights(self) -> torch.Tensor | None:
         """Return normalized oldest-to-newest layer weights, if enabled."""
 
@@ -209,6 +252,133 @@ def activity_targets_to_four_states(targets: torch.Tensor) -> torch.LongTensor:
     return (targets[..., 1] + 2 * targets[..., 2]).long()
 
 
+def transition_targets_from_activity(
+    activity_sequence: torch.Tensor,
+) -> torch.Tensor:
+    """Mark every activity-state change between adjacent 5 Hz target rows.
+
+    The input ends in [time, 3] and the output ends in [time - 1].
+    A silence/onset, offset, speaker handover, or overlap-state change is one
+    positive transition.
+    """
+
+    activity_sequence = _validated_binary_activity_targets(activity_sequence)
+    if activity_sequence.ndim < 2 or activity_sequence.shape[-2] < 2:
+        raise ValueError(
+            "Activity sequence must contain at least two target rows, got "
+            f"{tuple(activity_sequence.shape)}"
+        )
+    changed = activity_sequence[..., 1:, :] != activity_sequence[..., :-1, :]
+    return changed.any(dim=-1).to(activity_sequence.dtype)
+
+
+def weighted_transition_bce_loss(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    pos_weight: float = 3.0,
+) -> torch.Tensor:
+    """Binary transition loss with an explicit positive-class weight."""
+
+    if logits.shape != targets.shape:
+        raise ValueError(
+            f"Transition logits/targets mismatch: {tuple(logits.shape)} vs "
+            f"{tuple(targets.shape)}"
+        )
+    if pos_weight <= 0:
+        raise ValueError("Transition pos_weight must be positive")
+    targets = targets.to(device=logits.device, dtype=logits.dtype)
+    if not torch.isfinite(targets).all() or not torch.all(
+        (targets == 0) | (targets == 1)
+    ):
+        raise ValueError("Transition targets must contain only finite 0/1 values")
+    return nn.functional.binary_cross_entropy_with_logits(
+        logits,
+        targets,
+        pos_weight=logits.new_tensor(pos_weight),
+    )
+
+
+def expected_transition_reward(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    pos_weight: float = 3.0,
+) -> torch.Tensor:
+    """Exact expected weighted correctness of a Bernoulli transition policy.
+
+    Correct transition/non-transition actions receive reward one and incorrect
+    actions zero. The expectation is evaluated analytically rather than by
+    sampling, and normalized by the sum of example weights to remain in [0, 1].
+    """
+
+    if logits.shape != targets.shape:
+        raise ValueError(
+            f"Transition logits/targets mismatch: {tuple(logits.shape)} vs "
+            f"{tuple(targets.shape)}"
+        )
+    if pos_weight <= 0:
+        raise ValueError("Transition pos_weight must be positive")
+    targets = targets.to(device=logits.device, dtype=logits.dtype)
+    if not torch.isfinite(targets).all() or not torch.all(
+        (targets == 0) | (targets == 1)
+    ):
+        raise ValueError("Transition targets must contain only finite 0/1 values")
+    probability = logits.sigmoid()
+    probability_correct = (
+        targets * probability + (1.0 - targets) * (1.0 - probability)
+    )
+    weights = 1.0 + (pos_weight - 1.0) * targets
+    return (weights * probability_correct).sum() / weights.sum()
+
+
+def transition_aware_objective(
+    activity_loss: torch.Tensor,
+    transition_logits: torch.Tensor,
+    transition_targets: torch.Tensor,
+    *,
+    objective: str,
+    transition_weight: float = 0.2,
+    pos_weight: float = 3.0,
+    reward_weight: float = 0.1,
+) -> dict[str, torch.Tensor]:
+    """Compose the R1 or R2 transition-aware training objective."""
+
+    if objective not in TRANSITION_OBJECTIVE_NAMES[1:]:
+        raise ValueError(
+            f"Transition objective must be R1 or R2, got {objective!r}"
+        )
+    if transition_weight < 0 or reward_weight < 0:
+        raise ValueError("Transition and reward weights must be non-negative")
+    transition_loss = weighted_transition_bce_loss(
+        transition_logits,
+        transition_targets,
+        pos_weight=pos_weight,
+    )
+    expected_reward = expected_transition_reward(
+        transition_logits,
+        transition_targets,
+        pos_weight=pos_weight,
+    )
+    applied_reward = (
+        expected_reward
+        if objective == "r2-expected-reward"
+        else expected_reward.detach().new_zeros(())
+    )
+    total = (
+        activity_loss
+        + transition_weight * transition_loss
+        - reward_weight * applied_reward
+    )
+    return {
+        "total": total,
+        "activity": activity_loss,
+        "transition": transition_loss,
+        "expected_reward": expected_reward,
+        "applied_reward": applied_reward,
+    }
+
+
 def output_classes_for_loss(loss_name: str) -> int:
     """Return the number of head logits required by one training objective."""
 
@@ -609,6 +779,7 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
     config.setdefault("dropout", 0.1)
     config.setdefault("num_input_layers", 1)
     config.setdefault("learnable_layer_mixture", False)
+    config.setdefault("transition_auxiliary", False)
     config.setdefault(
         "layer_fusion",
         (