VLA-JEPA default world-model loss leaks future frames into context via shared V-JEPA2 encoder pass

[ravediamond]

Ticket Type

🐛 Bug Report (Something isn't working)

Environment & System Info

Not runtime-environment specific — this is a logic bug in the loss computation, reproducible on any install. Found while reading src/lerobot/policies/vla_jepa/modeling_vla_jepa.py against the VLA-JEPA paper and the authors' reference implementation.

Description

The default (always-on) world-model loss for VLA-JEPA leaks future frames into the context it predicts from, because context and target are both sliced out of a single shared encoder pass instead of being encoded separately.

• VLAJEPAModel._world_model_loss (src/lerobot/policies/vla_jepa/modeling_vla_jepa.py:197-242) runs the frozen V-JEPA2 encoder once over the whole video clip, then splits that one output tensor by index into input_states (context, fed to the predictor) and gt_states (the target it's supposed to predict).
• torch.no_grad() around the call keeps the encoder frozen, which is correct, but it only blocks gradients — it doesn't stop information flowing forward through attention during that single pass.
• V-JEPA2's attention blocks are non-causal (is_causal = False, checked directly in transformers/models/vjepa2/modeling_vjepa2.py), so every token in the output, including the ones that end up in input_states, already attended to every frame in the clip — including the ones that end up in gt_states. The slicing happens after the fact, so it can't undo that.

This contradicts the paper's own design (Sun et al., arXiv:2602.10098v2):

• Eq. 1 defines the target state as a function of the frame at a single timestep, and the abstract is explicit that "a target encoder produces latent representations from future frames, while the student pathway sees only the current observation—future information is used solely as supervision targets, never as input."
• As implemented, the context isn't leakage-free.

It affects every VLA-JEPA training run by default — not a rare edge case, not something you have to opt into.

Context & Reproduction

Relevant code, src/lerobot/policies/vla_jepa/modeling_vla_jepa.py:217-232:

​```python
with torch.no_grad():
video_embeddings = self.video_encoder.get_vision_features(pixel_values_videos=video_pixels)
video_embeddings = torch.cat(torch.chunk(video_embeddings, chunks=v, dim=0), dim=2)

tubelet_size = self.video_encoder.config.tubelet_size
t_enc_total = self.config.num_video_frames // tubelet_size
t_enc_ctx = t_enc_total - 1
tokens_per_frame = video_embeddings.shape[1] // t_enc_total
input_states = video_embeddings[:, : tokens_per_frame * t_enc_ctx, :]
gt_states = video_embeddings[:, tokens_per_frame:, :]
​```

Also checked this against the paper authors' own reference implementation (ginwind/VLA-JEPA, starVLA/model/framework/VLA_JEPA.py:224-232) — same bug is present there too, same single-pass-then-slice pattern. So this isn't something the LeRobot port introduced; it's a faithful port of a bug in the original code.

Filed upstream as ginwind/VLA-JEPA#29.

Relevant logs or stack trace

Checklist

• I have searched existing tickets to ensure this isn't a duplicate.
• I am using the latest version of the main branch.
• I have verified this is not an environment-specific problem.

Additional Info / Workarounds

There's already precedent for exactly this type of fix in this codebase: PR #4116 ("Add opt-in multi-horizon training objective for VLA-JEPA world model") ran into the same problem in a newer, opt-in code path and fixed it there by encoding one causal raw-frame prefix per temporal position instead of one shared pass. But it's disabled by default and explicitly preserves "the legacy state dict, inference path, and single full-clip encoder call when the feature is disabled" — so the base path this issue is about was left untouched.

Suggested fix:

• mirror feat(vla_jepa): add multi-horizon world model training #4116's approach — encode context and target from separate forward passes instead of one shared call, or at minimum make sure input_states comes from a raw-frame window that can't see any frame contributing to gt_states.
• Will cost more encoder compute (per feat(vla_jepa): add multi-horizon world model training #4116's own numbers, roughly 4 encoder calls instead of 1 for the default 8-frame / tubelet-size-2 setup).

[ravediamond]

I'll keep this updated as the upstream issue (ginwind/VLA-JEPA#29) progresses — will comment here once the paper authors respond or land a fix.