Add overlap-aware turn prediction objectives

easycom_head.py

@@ -30,9 +30,22 @@ CLASS_NAMES = (
     "camera_wearer_speaking",
     "other_people_speaking",
 )
+FOUR_STATE_NAMES = (
+    "silence",
+    "wearer_only",
+    "others_only",
+    "overlap",
+)
 STEP_SECONDS = 0.2
 NUM_HORIZONS = 10
 NUM_CLASSES = 3
+NUM_FOUR_STATES = len(FOUR_STATE_NAMES)
+LOSS_NAMES = (
+    "lstr-mce",
+    "miniroad-normalized",
+    "independent-bce",
+    "four-state-ce",
+)
 DATASET_NAMES = ("easycom", "ego4d")
 LAYER_FUSION_MODES = (
     "single",
@@ -44,7 +57,7 @@ DatasetName = Literal["easycom", "ego4d"]
 
 
 class TurnPredictionHead(nn.Module):
-    """Map one Omni state to EgoSpeak-compatible ``[10, 3]`` logits."""
+    """Map one Omni state to configurable ``[horizon, output]`` logits."""
 
     def __init__(
         self,
@@ -150,18 +163,76 @@ class TurnPredictionHead(nn.Module):
         return self.layer_logits.softmax(dim=0)
 
 
-def _validate_loss_inputs(logits: torch.Tensor, targets: torch.Tensor) -> None:
+def _validate_loss_inputs(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    num_output_classes: int = NUM_CLASSES,
+) -> None:
     if logits.shape != targets.shape:
         raise ValueError(
             f"logits/targets shape mismatch: {tuple(logits.shape)} vs "
             f"{tuple(targets.shape)}"
         )
-    if logits.shape[-1] != NUM_CLASSES:
+    if logits.shape[-1] != num_output_classes:
         raise ValueError(
-            f"Expected {NUM_CLASSES} classes, got {logits.shape[-1]}"
+            f"Expected {num_output_classes} classes, got {logits.shape[-1]}"
         )
 
 
+def _validated_binary_activity_targets(targets: torch.Tensor) -> torch.Tensor:
+    """Validate the four legal EasyCom/EgoSpeak three-label activity rows."""
+
+    if targets.ndim < 1 or targets.shape[-1] != NUM_CLASSES:
+        raise ValueError(
+            f"Expected activity targets [...,{NUM_CLASSES}], got "
+            f"{tuple(targets.shape)}"
+        )
+    if not torch.isfinite(targets).all():
+        raise ValueError("Activity targets contain NaN or Inf")
+    rounded = targets.round()
+    if not torch.allclose(targets, rounded, atol=1e-6, rtol=0.0):
+        raise ValueError("Activity targets must contain only binary 0/1 values")
+    expected_background = (rounded[..., 1:].sum(dim=-1) == 0).to(rounded.dtype)
+    if not torch.equal(rounded[..., 0], expected_background):
+        raise ValueError(
+            "Activity targets must be one of [1,0,0], [0,1,0], "
+            "[0,0,1], or [0,1,1]"
+        )
+    return rounded
+
+
+def activity_targets_to_four_states(targets: torch.Tensor) -> torch.LongTensor:
+    """Map three-label activity targets to silence/wearer/others/overlap."""
+
+    targets = _validated_binary_activity_targets(targets)
+    return (targets[..., 1] + 2 * targets[..., 2]).long()
+
+
+def output_classes_for_loss(loss_name: str) -> int:
+    """Return the number of head logits required by one training objective."""
+
+    if loss_name not in LOSS_NAMES:
+        raise ValueError(f"Unknown loss {loss_name!r}; expected one of {LOSS_NAMES}")
+    return NUM_FOUR_STATES if loss_name == "four-state-ce" else NUM_CLASSES
+
+
+def output_space_for_loss(loss_name: str) -> str:
+    """Return a stable checkpoint description of the head output space."""
+
+    output_classes_for_loss(loss_name)
+    if loss_name == "four-state-ce":
+        return "four-state-softmax"
+    if loss_name == "independent-bce":
+        return "three-label-sigmoid"
+    return "three-class-softmax"
+
+
+def head_output_names_for_loss(loss_name: str) -> tuple[str, ...]:
+    output_classes_for_loss(loss_name)
+    return FOUR_STATE_NAMES if loss_name == "four-state-ce" else CLASS_NAMES
+
+
 def egospeak_multihorizon_loss(
     logits: torch.Tensor,
     targets: torch.Tensor,
@@ -211,6 +282,99 @@ def miniroad_multihorizon_loss(
     raise ValueError(f"Unknown reduction: {reduction}")
 
 
+def independent_bce_multihorizon_loss(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    reduction: Literal["mean", "sum", "none"] = "mean",
+) -> torch.Tensor:
+    """Independent sigmoid/BCE loss for the three marginal activity labels."""
+
+    _validate_loss_inputs(logits, targets)
+    targets = targets.to(
+        device=logits.device,
+        dtype=logits.dtype,
+    )
+    element_loss = nn.functional.binary_cross_entropy_with_logits(
+        logits,
+        targets,
+        reduction="none",
+    )
+    if reduction == "mean":
+        return element_loss.mean()
+    if reduction == "sum":
+        return element_loss.sum()
+    if reduction == "none":
+        return element_loss
+    raise ValueError(f"Unknown reduction: {reduction}")
+
+
+def four_state_multihorizon_loss(
+    logits: torch.Tensor,
+    targets: torch.Tensor,
+    *,
+    reduction: Literal["mean", "sum", "none"] = "mean",
+) -> torch.Tensor:
+    """Four-state CE for silence, wearer-only, others-only, and overlap."""
+
+    if logits.ndim < 1 or logits.shape[-1] != NUM_FOUR_STATES:
+        raise ValueError(
+            f"Expected four-state logits [...,{NUM_FOUR_STATES}], got "
+            f"{tuple(logits.shape)}"
+        )
+    if logits.shape[:-1] != targets.shape[:-1] or targets.shape[-1] != NUM_CLASSES:
+        raise ValueError(
+            "Four-state logits/activity target shape mismatch: "
+            f"{tuple(logits.shape)} vs {tuple(targets.shape)}"
+        )
+    state_targets = (targets[..., 1] + 2 * targets[..., 2]).long().to(
+        logits.device
+    )
+    loss = nn.functional.cross_entropy(
+        logits.reshape(-1, NUM_FOUR_STATES),
+        state_targets.reshape(-1),
+        reduction="none",
+    ).reshape(logits.shape[:-1])
+    if reduction == "mean":
+        return loss.mean()
+    if reduction == "sum":
+        return loss.sum()
+    if reduction == "none":
+        return loss
+    raise ValueError(f"Unknown reduction: {reduction}")
+
+
+def marginal_probabilities_from_logits(
+    logits: torch.Tensor,
+    loss_name: str,
+) -> torch.Tensor:
+    """Convert any supported head output to three EgoSpeak marginals.
+
+    Four-state marginals deliberately are not renormalized: wearer and others
+    may both be near one when the overlap state is likely.
+    """
+
+    expected_classes = output_classes_for_loss(loss_name)
+    if logits.ndim < 1 or logits.shape[-1] != expected_classes:
+        raise ValueError(
+            f"Loss {loss_name!r} expects {expected_classes} logits, got "
+            f"{tuple(logits.shape)}"
+        )
+    if loss_name == "independent-bce":
+        return logits.sigmoid()
+    probabilities = logits.softmax(dim=-1)
+    if loss_name != "four-state-ce":
+        return probabilities
+    return torch.stack(
+        (
+            probabilities[..., 0],
+            probabilities[..., 1] + probabilities[..., 3],
+            probabilities[..., 2] + probabilities[..., 3],
+        ),
+        dim=-1,
+    )
+
+
 def multihorizon_cross_entropy(
     logits: torch.Tensor,
     targets: torch.Tensor,
@@ -380,6 +544,7 @@ def load_session_data(
             f"{target_path}"
         )
     targets = torch.from_numpy(np.asarray(target_array, dtype=np.float32))
+    _validated_binary_activity_targets(targets)
     if len(label_indices) and int(label_indices[-1]) >= len(targets):
         raise ValueError(
             f"Last label index {int(label_indices[-1])} exceeds target length "
@@ -455,6 +620,30 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
     return config
 
 
+def checkpoint_loss_name(payload: dict[str, Any]) -> str:
+    """Read the objective from a checkpoint, defaulting legacy files to MCE."""
+
+    training_config = payload.get("training_config", {})
+    if not isinstance(training_config, dict):
+        raise ValueError("Checkpoint training_config must be a mapping")
+    loss_name = str(training_config.get("loss", "lstr-mce"))
+    expected_classes = output_classes_for_loss(loss_name)
+    actual_classes = int(checkpoint_head_config(payload)["num_classes"])
+    if actual_classes != expected_classes:
+        raise ValueError(
+            f"Checkpoint loss {loss_name!r} requires {expected_classes} head "
+            f"outputs, but head_config has {actual_classes}"
+        )
+    recorded_space = payload.get("output_space")
+    expected_space = output_space_for_loss(loss_name)
+    if recorded_space is not None and recorded_space != expected_space:
+        raise ValueError(
+            f"Checkpoint output_space is {recorded_space!r}, expected "
+            f"{expected_space!r} for loss {loss_name!r}"
+        )
+    return loss_name
+
+
 def validate_disjoint_sessions(
     first: Sequence[str],
     second: Sequence[str],