Proposal: Feature Reconstruction Loss for FutureState (Auxiliary Self-Supervised Task)

[riita10069]

Summary

The FutureState module currently outputs predicted future visual features ([B, 1440, 7, 7] × 4) during the forward pass, but these predictions are never used — there is no loss function that compares them against actual future observations. This issue proposes implementing the Feature Reconstruction Loss to activate FutureState as an auxiliary training signal and to enable introspection of model uncertainty.

Current State

• FutureState.forward() predicts features at t+1.6s, t+3.2s, t+4.8s, t+6.4s
• FrozenBackbone exists in the codebase but is unused
• The predicted future features are returned from AutoE2E.forward() but discarded

Proposed: Feature Reconstruction Loss

How it works

During training with sequential frame data, we compare what the model predicted the future would look like against what the future actually looked like:

Frame t (current):
  image_t → Backbone → Fusion → FutureState → predicted_features[t+1.6s]

Frame t+16 (actual 1.6s later, at 10Hz):
  image_{t+16} → FrozenBackbone → Fusion → target_features[t+1.6s]

feature_reconstruction_loss = MSE(predicted_features, target_features.detach())

Why predict in feature space, not pixel space

Predicting future images (pixels) forces the model to reconstruct irrelevant details like exact texture, shadow positions, and compression artifacts. JEPA (Joint Embedding Predictive Architecture) [1] proposes predicting in a learned latent space instead, where only semantically meaningful changes are captured. MILE [2] demonstrated this principle specifically for autonomous driving — their latent world model outperformed pixel-prediction baselines.

Why use a FrozenBackbone for targets

If both the predictor and the target encoder are trainable, the network can trivially minimize loss by collapsing all representations to a constant (representation collapse). Using a frozen or slowly-updating target encoder prevents this. This asymmetric design is established in BYOL [3] and JEPA [1].

Two update strategies for the frozen encoder:

1. Keep at pretrained initialization (simplest)
2. Exponential Moving Average from trainable Backbone:

   for p_frozen, p_train in zip(frozen.parameters(), trainable.parameters()):
       p_frozen.data = τ * p_frozen.data + (1 - τ) * p_train.data  # τ ≈ 0.996

How this helps DrivingPolicy (auxiliary task)

Training only with trajectory loss provides a weak gradient signal to the Backbone and View Fusion modules — the loss is far downstream and must propagate through many layers. Feature reconstruction loss provides a direct training signal to these early modules: "extract features that contain enough information to predict how the scene will change." This richer gradient signal improves the quality of the shared representation, which in turn benefits trajectory prediction.

UniAD [4] uses a similar multi-task design where perception auxiliary losses (detection, tracking, map segmentation) improve the shared BEV representation quality, leading to better final planning performance.

Integration with total loss

total_loss = trajectory_loss + λ * feature_reconstruction_loss

The weight λ balances the two objectives. Typical range: 0.5 – 1.0. If λ is too large, the model focuses on future prediction at the expense of trajectory accuracy. If too small, the auxiliary signal is negligible.

Proposed: Introspection via Prediction Error

Uncertainty estimation at inference time

After training, the feature reconstruction loss can be computed at inference time (when sequential frames are available) as a measure of model uncertainty:

prediction_error = MSE(predicted_features_t_plus_1_6, actual_features_t_plus_1_6)

A high prediction error at a given frame indicates the model encountered a situation it did not expect — the scene evolved in a way that differs from its internal world model. This can signal:

• Unusual driving scenarios (construction zones, accidents)
• Out-of-distribution situations (weather the model has not seen)
• Sudden events (pedestrian stepping onto road)

This prediction error can be monitored as a real-time safety signal: if the model's world model is consistently wrong, the system should reduce confidence in its trajectory predictions or request human intervention.

GAIA-1 [5] demonstrates that generative world models trained on driving data develop internal representations of scene dynamics that can be used for anomaly detection in exactly this way.

Implementation Scope

Item	Notes
Wire FrozenBackbone output as target for FutureState loss	Requires sequential frame pairs in DataLoader
Add feature_reconstruction_loss computation	MSE between predicted and actual features
Add λ hyperparameter	Config-level parameter
Add EMA update option for FrozenBackbone	Post optimizer.step() hook
Add prediction_error logging at inference	For introspection/monitoring

Out of scope (separate issues)

• Training loop infrastructure (optimizer, scheduler, DataLoader)
• Trajectory loss implementation
• Dataset selection
• BEV segmentation auxiliary loss

References

1. LeCun, Y. "A Path Towards Autonomous Machine Intelligence." 2022. https://openreview.net/pdf?id=BZ5a1r-kVsf
2. Hu, A., et al. "Model-Based Imitation Learning for Urban Driving (MILE)." NeurIPS 2022. https://arxiv.org/abs/2210.07729
3. Grill, J.B., et al. "Bootstrap Your Own Latent (BYOL)." NeurIPS 2020. https://arxiv.org/abs/2006.07733
4. Hu, Y., et al. "Planning-oriented Autonomous Driving (UniAD)." CVPR 2023. https://arxiv.org/abs/2212.10156
5. Hu, A., et al. "GAIA-1: A Generative World Model for Autonomous Driving." 2023. https://arxiv.org/abs/2309.17080

[riita10069]

Following up on this — now that the forward pass is complete and FutureState outputs predicted future features, the missing piece to activate it is wiring up the FrozenBackbone as a target and computing the MSE between predicted and actual features.

Without this, the FutureState module runs during forward but its output is discarded — it doesn't contribute to learning, and the introspection capability (monitoring prediction error at inference time as an uncertainty signal) is also not yet functional.

I think once the trajectory loss (#15) and dataset (#16) are decided, this could be implemented relatively quickly since the FrozenBackbone already exists in the codebase. The main design questions are probably:

• Should the FrozenBackbone stay at pretrained initialization, or use EMA updates from the trainable Backbone?
• Should the loss apply to all 4 time horizons equally, or weight near-future more?

Would appreciate any thoughts from the team on priority and preferred approach. Happy to take a pass at this when the time is right.

[m-zain-khawaja]

Thanks @riita10069:

• I think we can use frozen backbone with pre-trained initialization to calculate the ground truth future state.
• For now, we can do equal weights to all for time horizons and see what the training results look like, we can then consider applying weights to specific timesteps if it helps the model converge better

[intisar1020]

I would like to share a small incremental idea which almost alligns with @riita10069 san proposal.

I recently came across a paper, Temporal Difference in Vision (TDV) (just yesterday), which proposes a self-supervised learning approach that trains purely from video data using a single causal assumption -> the past causes the future. Rather than leaning on the usual inductive biases (augmentation, masking, cropping that we see in DINO v1/v2/v3, iBOT, byol etc.), it enforces a clean additive relationship in latent space. It seems at least loosely aligned with what we're trying to do here, so I thought it might be worth a look. Following I am sharing detail.

Core idea of TDV

The paper relies on a simple causal assumption
current state + encoded motion ≈ future state

Architecturally, instead of predicting absolute future features from scratch, it enforces an additive temporal relationship in latent space:

$$z_t + \Delta z_{\text{motion}} = z_{t+1}$$

where $z_t$ is the current frame representation and $\Delta z_{\text{motion}}$ is the learned temporal difference. The authors report this matches SOTA recipes on dense spatial tasks, which feels potentially relevant to our shared BEV representation goals (though I haven't verified how well the claims hold up).

How we might apply this

1. Predicting deltas via residual learning. Instead of having FutureState reconstruct absolute feature maps from scratch at each horizon, which arguably can comsume capacity re-learning static BEV geometry (roads, barriers), we could frame it as residual learning by adding the current frame's fused_features to the chunked FutureState outputs:

$$\hat{z}_{t+\Delta t} = z_t + \mathrm{FutureState}(z_t)$$

The idea is to anchor the prediction to the current spatial context, so the network mostly models what changes (ego-motion transitions and dynamic actors).

2. Likely zero downstream changes. Since the shape of $\hat{z}_{t+\Delta t}$ would stay identical to the current implementation, the loss function, training loop, and dataloader logic hopefully shouldn't need to change. We compute the standard MSE loss against the full future feature targets from the FrozenBackbone as @riita10069 showed earllier.

prediction_error = MSE(z_t + delta_motion, actual_features_t_plus_1_6)
where delta_motion = FutureState(z_t, ego_hidden)

3. targets. to compute the groudn truth latent trargets as @m-zain-khawaja proposed we can use the frozen backbones output as it is. The additive constraint + frozen backbone feature could act as a good regularizer and help guard against the representation collapse usually feared in BYOL/JEPA-style setups.

Proposed code change

The nice part is that this is a minimal, drop-in adjustment at the very end of FutureState.forward(), leaving the rest of the architecture untouched:

def forward(self, fused_features, ego_hidden):
    # fused_features: [B, C, H, W]; ego_hidden: [B, ego_hidden_dim]

    # Inject planner intent as a per-channel bias
    bias = self.ego_proj(ego_hidden).view(-1, self.embed_dim, 1, 1)
    conditioned = fused_features + bias

    # Predict future channels
    future_features = self.predict_future_1(conditioned)
    future_features = self.activation(future_features)
    future_features = self.predict_future_2(future_features)

    # first we split into 4 future delta feature vectors where each [B, C, H, W]
    future_deltas = torch.chunk(future_features, chunks=4, dim=1)

    # Second, we apply additive residual skip connection from the current frame
    future_visual_features = tuple(fused_features + delta for delta in future_deltas)

    return future_visual_features

I look forward to your opinion.

@misc{daithankar2026dontneedstrongassumptions,
  title         = {You Don't Need Strong Assumptions: Visual Representation Learning via Temporal Differences},
  author        = {Ninad Daithankar and Alexi Gladstone and Yann LeCun and Heng Ji},
  year          = {2026},
  eprint        = {2606.15956},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CV},
  url           = {https://arxiv.org/abs/2606.15956},
}

[gcordova10]

Hi @m-zain-khawaja — following the priority order in #56 (riita10069 put FutureState + JEPA as #1), I've prepared the loss side of Feature B and I'd like to align with you before wiring it into training, since FutureState is your module and you have the memory-lightening action item from 17/06.

What's ready on my side (feat/future-state-jepa)

• losses/feature_reconstruction_loss.py: a JEPA-style feature-reconstruction loss over FutureState's predicted future feature maps, with optional temporal weighting and a zero-sum-weights guard. Additive, tested.
• Optional action conditioning in future_state.py (action_dim=None ⇒ byte-identical to current default).

What I need to align with you (the "make it useful" parts riita10069 listed)

1. Future-frame targets — which future features do we reconstruct (horizons +1.6/3.2/4.8/6.4 s?) and how is the target produced from the data pipeline?
2. Target encoder — frozen backbone vs EMA of the online encoder? (avoids representational collapse)
3. Memory budget — JEPA adds a forecast at BEV resolution (450×300). How does this fit your feature-state lightening? Options: bf16/fp16, downsample features before the forecast, fewer horizons.
4. Loss weighting — JEPA loss vs trajectory imitation loss (and vs any RL term).
5. Integration point — I assume this hooks into the training loop from feat(platform): MLOps training platform on EKS — Flyte + MLflow training loop (#14) #78/feat: training loop, lerobot 0.5.x compat, and e2e dataset tests #67 rather than a standalone loop. Confirm?

Proposal: I open PR #80 with the loss + action-conditioning as the additive first step, keeping the future_state.py diff minimal/optional, and we iterate on (1)–(5) together. Happy to follow your lead on the memory and target-encoder design since it's your module. Thanks!