feat(jepa): add frozen/EMA target encoder + compute_jepa_loss for FutureState

FutureState (autowarefoundation#80) predicts future BEV features and FeatureReconstructionLoss
scores them, but the target side was missing. Per @ryotayamada's 'make it
useful' list (autowarefoundation#56/autowarefoundation#13), add JepaTargetEncoder: a stop-gradient target encoder
supporting both modes he named ('frozen or EMA'), plus compute_jepa_loss to
fold the term into a training step. Additive and optional; does not touch
AutoE2E.forward or the training loop. Frequency/horizons/weighting/data-pipeline
remain open design decisions (Zain, 17-06 action item).

Signed-off-by: GABRIELA CORDOVA <100548769@alumnos.uc3m.es>

Author: gcordova10

Model/model_components/jepa_target_encoder.py

@@ -0,0 +1,133 @@
+"""Frozen / EMA target encoder for the JEPA feature-reconstruction objective.
+
+Background
+----------
+Feature B (merged, #80) added the *prediction* side of the JEPA objective:
+
+* ``FutureState`` predicts future BEV feature maps — a list of
+  ``num_future_steps`` tensors, each ``[B, C, H, W]``.
+* ``losses.FeatureReconstructionLoss`` scores those predictions against
+  *target* feature maps of identical shape. Its own docstring states the
+  targets are "extracted by a frozen copy of the image backbone (no gradient)
+  applied to the future frames at +1.6s, +3.2s, +4.8s and +6.4s".
+
+That **target encoder** was the missing piece. In Joint-Embedding Predictive
+Architectures (I-JEPA / V-JEPA) the targets come from an encoder that is *not*
+trained by backprop from the loss — it is either a **frozen** copy of the
+online encoder or an **exponential-moving-average (EMA)** of it, and its output
+is detached (stop-gradient). This is what prevents representational collapse
+(the predictor cannot win by driving every feature to zero, because the target
+encoder is not pulled along).
+
+This module implements exactly that, supporting BOTH modes @RyotaYamada listed
+in #56 / #13 ("frozen or EMA") so the choice stays a *configuration* rather than
+a hard-coded decision.
+
+Deliberately left open (owned by Zain — 17-06 action item, #13)
+---------------------------------------------------------------
+This module does NOT decide: the input frequency (1 Hz / TBD), the predictor
+space (BEV vs feature), the exact future horizons, the JEPA-vs-imitation loss
+weighting, or the data pipeline that supplies the future frames. It only
+provides the collapse-safe target generator those decisions will plug into.
+"""
+
+import copy
+
+import torch
+import torch.nn as nn
+
+
+class JepaTargetEncoder(nn.Module):
+    """Produce stop-gradient target features for the JEPA loss.
+
+    Wraps a *copy* of an online encoder (e.g. the image backbone, or the
+    backbone+fusion path) and runs it without gradients to generate the target
+    feature maps that ``FeatureReconstructionLoss`` compares ``FutureState``'s
+    predictions against.
+
+    Args:
+        encoder: the online encoder to mirror. A deep copy is taken at
+            construction; the original is never modified by this module.
+        mode: ``"frozen"`` (a fixed copy, never updated) or ``"ema"`` (updated
+            from the online encoder via :meth:`update`). Default ``"ema"``.
+        ema_decay: EMA momentum in ``[0, 1]``; only used in ``"ema"`` mode.
+            ``target ← decay * target + (1 - decay) * online``.
+
+    The wrapped encoder's parameters always have ``requires_grad=False`` and
+    :meth:`forward` returns **detached** outputs, so no gradient ever reaches
+    the target branch regardless of how the caller composes the loss.
+    """
+
+    def __init__(self, encoder: nn.Module, mode: str = "ema",
+                 ema_decay: float = 0.999):
+        super().__init__()
+        if mode not in ("frozen", "ema"):
+            raise ValueError(f"mode must be 'frozen' or 'ema', got {mode!r}")
+        if not 0.0 <= ema_decay <= 1.0:
+            raise ValueError(f"ema_decay must be in [0, 1], got {ema_decay}")
+        self.mode = mode
+        self.ema_decay = ema_decay
+
+        # A detached, non-trainable mirror of the online encoder.
+        self.encoder = copy.deepcopy(encoder)
+        self.encoder.requires_grad_(False)
+        self.encoder.eval()
+
+    @torch.no_grad()
+    def update(self, online_encoder: nn.Module) -> None:
+        """EMA-update the target weights toward ``online_encoder``.
+
+        No-op in ``"frozen"`` mode. Buffers (e.g. BatchNorm running stats) are
+        copied directly rather than averaged. Call once per optimizer step,
+        after the online encoder has been updated.
+        """
+        if self.mode != "ema":
+            return
+        d = self.ema_decay
+        for t_p, o_p in zip(self.encoder.parameters(),
+                            online_encoder.parameters()):
+            t_p.mul_(d).add_(o_p.detach(), alpha=1.0 - d)
+        for t_b, o_b in zip(self.encoder.buffers(), online_encoder.buffers()):
+            t_b.copy_(o_b)
+
+    @torch.no_grad()
+    def forward(self, future_observations):
+        """Encode each future observation into a detached target feature map.
+
+        Args:
+            future_observations: an iterable of ``num_future_steps`` tensors,
+                each a valid input to the wrapped encoder (e.g. future camera
+                frames or future fused features), ordered by horizon.
+
+        Returns:
+            list of ``num_future_steps`` detached feature maps ``[B, C, H, W]``,
+            ready to pass as ``target_features`` to
+            :class:`FeatureReconstructionLoss`.
+        """
+        self.encoder.eval()
+        return [self.encoder(obs).detach() for obs in future_observations]
+
+
+def compute_jepa_loss(predicted_features, future_observations, target_encoder,
+                      loss_fn, weight: float = 1.0):
+    """Glue helper: targets ← target_encoder, then weighted JEPA loss.
+
+    This is the single call a training loop adds to fold the JEPA objective in
+    alongside the trajectory imitation loss, e.g.::
+
+        total = imitation_loss + compute_jepa_loss(
+            future, future_frames, target_encoder, recon_loss, weight=0.1)
+
+    Args:
+        predicted_features: ``FutureState`` output (list of ``[B, C, H, W]``).
+        future_observations: inputs for ``target_encoder`` (one per horizon).
+        target_encoder: a :class:`JepaTargetEncoder`.
+        loss_fn: a :class:`FeatureReconstructionLoss` (or compatible callable).
+        weight: scalar coefficient for the JEPA term (``lambda``).
+
+    Returns:
+        ``weight * loss_fn(predicted_features, targets)`` — a scalar tensor
+        differentiable w.r.t. ``predicted_features`` only (targets are detached).
+    """
+    targets = target_encoder(future_observations)
+    return weight * loss_fn(predicted_features, targets)

Model/tests/test_jepa_target_encoder.py

@@ -0,0 +1,170 @@
+"""Unit tests for the JEPA target encoder (frozen / EMA, stop-gradient).
+
+Verifies the collapse-safety guarantees of the target branch and that it
+plugs into the already-merged ``FutureState`` + ``FeatureReconstructionLoss``:
+  - the target encoder never carries gradient (frozen params, detached output),
+  - "frozen" mode is a stable deep copy independent of the online encoder,
+  - "ema" mode moves the target toward the online encoder on ``update()``,
+  - ``compute_jepa_loss`` is differentiable w.r.t. the prediction ONLY,
+  - end-to-end with the real FutureState world-model head.
+"""
+
+import pytest
+import torch
+import torch.nn as nn
+
+from model_components.future_state import FutureState
+from model_components.jepa_target_encoder import (
+    JepaTargetEncoder,
+    compute_jepa_loss,
+)
+from model_components.losses.feature_reconstruction_loss import (
+    FeatureReconstructionLoss,
+)
+
+EMBED_DIM = 256
+NUM_FUTURE = 4
+
+
+class _TinyEncoder(nn.Module):
+    """Maps a feature map [B, C, H, W] -> [B, C, H, W] (stand-in for backbone)."""
+
+    def __init__(self, c=EMBED_DIM):
+        super().__init__()
+        self.conv = nn.Conv2d(c, c, 3, padding=1)
+        self.bn = nn.BatchNorm2d(c)
+
+    def forward(self, x):
+        return self.bn(self.conv(x))
+
+
+def _future_obs(batch, device, h=8, w=8):
+    return [torch.randn(batch, EMBED_DIM, h, w, device=device)
+            for _ in range(NUM_FUTURE)]
+
+
+def test_invalid_mode_and_decay():
+    enc = _TinyEncoder()
+    with pytest.raises(ValueError, match="mode must be"):
+        JepaTargetEncoder(enc, mode="bogus")
+    with pytest.raises(ValueError, match="ema_decay"):
+        JepaTargetEncoder(enc, mode="ema", ema_decay=1.5)
+
+
+def test_target_params_are_frozen(device):
+    target = JepaTargetEncoder(_TinyEncoder(), mode="frozen").to(device)
+    assert all(not p.requires_grad for p in target.parameters())
+
+
+def test_forward_output_is_detached(device):
+    target = JepaTargetEncoder(_TinyEncoder(), mode="ema").to(device)
+    outs = target(_future_obs(2, device))
+    assert len(outs) == NUM_FUTURE
+    for o in outs:
+        assert not o.requires_grad
+        assert o.grad_fn is None
+        assert o.shape == (2, EMBED_DIM, 8, 8)
+
+
+def test_frozen_is_independent_deep_copy(device):
+    online = _TinyEncoder().to(device)
+    target = JepaTargetEncoder(online, mode="frozen").to(device)
+    before = [p.clone() for p in target.encoder.parameters()]
+
+    # Mutate the ONLINE encoder; the frozen target must not change.
+    with torch.no_grad():
+        for p in online.parameters():
+            p.add_(1.0)
+    target.update(online)  # no-op in frozen mode
+
+    for b, p in zip(before, target.encoder.parameters()):
+        assert torch.equal(b, p)
+
+
+def test_ema_update_moves_target_toward_online(device):
+    torch.manual_seed(0)
+    online = _TinyEncoder().to(device)
+    target = JepaTargetEncoder(online, mode="ema", ema_decay=0.9).to(device)
+
+    # Push the online encoder far away, then EMA-update once.
+    with torch.no_grad():
+        for p in online.parameters():
+            p.add_(10.0)
+
+    t_before = [p.clone() for p in target.encoder.parameters()]
+    target.update(online)
+
+    for tb, t_after, o_p in zip(t_before, target.encoder.parameters(),
+                                online.parameters()):
+        moved = (t_after - tb).abs().sum().item()
+        assert moved > 0.0, "ema target did not move"
+        # decay=0.9 -> target should move ~10% of the way, not all the way.
+        assert not torch.allclose(t_after, o_p), "target jumped fully to online"
+        expected = 0.9 * tb + 0.1 * o_p
+        assert torch.allclose(t_after, expected, atol=1e-5)
+
+
+def test_compute_jepa_loss_zero_when_prediction_matches_target(device):
+    target = JepaTargetEncoder(_TinyEncoder(), mode="frozen").to(device)
+    loss_fn = FeatureReconstructionLoss(num_future_steps=NUM_FUTURE).to(device)
+    obs = _future_obs(2, device)
+    targets = target(obs)
+    predicted = [t.clone().requires_grad_(True) for t in targets]
+    loss = compute_jepa_loss(predicted, obs, target, loss_fn, weight=1.0)
+    assert loss.ndim == 0
+    assert loss.item() == pytest.approx(0.0, abs=1e-6)
+
+
+def test_compute_jepa_loss_weight_scales_linearly(device):
+    target = JepaTargetEncoder(_TinyEncoder(), mode="frozen").to(device)
+    loss_fn = FeatureReconstructionLoss(num_future_steps=NUM_FUTURE).to(device)
+    obs = _future_obs(3, device)
+    predicted = [torch.randn(3, EMBED_DIM, 8, 8, device=device)
+                 for _ in range(NUM_FUTURE)]
+    l1 = compute_jepa_loss(predicted, obs, target, loss_fn, weight=1.0)
+    l2 = compute_jepa_loss(predicted, obs, target, loss_fn, weight=0.25)
+    assert l2.item() == pytest.approx(0.25 * l1.item(), rel=1e-5)
+
+
+def test_gradient_flows_to_prediction_only(device):
+    target = JepaTargetEncoder(_TinyEncoder(), mode="ema").to(device)
+    loss_fn = FeatureReconstructionLoss(num_future_steps=NUM_FUTURE).to(device)
+    obs = _future_obs(2, device)
+    predicted = [torch.randn(2, EMBED_DIM, 8, 8, device=device,
+                             requires_grad=True) for _ in range(NUM_FUTURE)]
+
+    loss = compute_jepa_loss(predicted, obs, target, loss_fn, weight=0.5)
+    loss.backward()
+
+    assert all(p.grad is not None for p in predicted), \
+        "prediction must receive gradient"
+    assert all(p.grad is None for p in target.parameters()), \
+        "target encoder must NOT receive gradient (stop-gradient)"
+
+
+def test_end_to_end_with_future_state(device):
+    """FutureState (predictor) + JepaTargetEncoder (targets) + recon loss.
+
+    Mirrors the intended training step: the world-model head predicts future
+    features, targets come from the frozen/EMA encoder on future observations,
+    and the JEPA loss back-props into FutureState but never into the target."""
+    torch.manual_seed(0)
+    future_state = FutureState(embed_dim=EMBED_DIM,
+                               ego_hidden_dim=EMBED_DIM).to(device)
+    target = JepaTargetEncoder(_TinyEncoder(), mode="ema").to(device)
+    loss_fn = FeatureReconstructionLoss(num_future_steps=NUM_FUTURE).to(device)
+
+    fused = torch.randn(2, EMBED_DIM, 8, 8, device=device)
+    ego_hidden = torch.randn(2, EMBED_DIM, device=device)
+    predicted = future_state(fused, ego_hidden)
+    assert len(predicted) == NUM_FUTURE
+
+    future_obs = _future_obs(2, device)
+    loss = compute_jepa_loss(predicted, future_obs, target, loss_fn, weight=0.1)
+    assert loss.ndim == 0 and torch.isfinite(loss)
+
+    loss.backward()
+    assert any(p.grad is not None for p in future_state.parameters()), \
+        "FutureState must receive gradient from the JEPA loss"
+    assert all(p.grad is None for p in target.parameters()), \
+        "target encoder must stay gradient-free"