World models building blocks, starting with LeWM / Sub-JEPA

Nobody has claimed this yet.

Following up on the discussion in [the comment thread here](https://github.com/lightly-ai/lightly/pull/1924#issuecomment-4579124158), opening this to discuss/propose a direction for world models in lightly and refine the scope.

The goal is to add reusable building blocks only, consistent with the library philosophy and the removal of the `LeJEPAEncoder` wrapper in #1926. No world-model wrapper class, training loop, environment, or dataset code. SIGReg (already in `lightly/loss/lejepa_loss.py`) is reused as the regularizer, so the immediate next steps stay well defined: a prediction loss and a predictor module, with a Sub-JEPA regularizer as a follow-up.

For context: LeWM (LeWorldModel) frames a world model as next-latent prediction along a trajectory. An encoder embeds each observation, a predictor maps the current latent, conditioned on an action a_t, to the next-step latent, and SIGReg keeps the latent distribution from collapsing. No pixel decoder, EMA target, or stop-gradient in training.

Proposed, as separate PRs:

- [ ] A next-embedding prediction loss for latent world models (LeWM core objective): `L_pred = mean || pred(z_t, a_t) - z_{t+1} ||^2`, composing with the existing SIGReg as `L_pred + lambda * SIGReg`. Functional, mirroring `lejepa_invariance_loss`. **Fit**: useful broadly, since written over `(predicted, target)` it is modality-agnostic and also serves image masked-token prediction. Its value is the composed recipe, not a bare MSE.
- [ ] An action-conditioned transformer predictor (Adaptive LayerNorm conditioning, temporal causal mask), as a standalone `nn.Module` in `lightly/models/modules/`, following the `IJEPAPredictor` / `IJEPAPredictorTIMM` torch-and-timm pattern. Note: the existing `IJEPAPredictor` is a masked-patch spatial predictor, so this is a new capability rather than a change to I-JEPA. **Fit**: this is the only block that needs sequence and action data, which lightly does not ship, so it is a bring-your-own-data module (users supply their own trajectories and actions).
- [ ] (follow-up) `MultiSubspaceSIGReg` from Sub-JEPA, a drop-in alternative regularizer that runs SIGReg over K random subspaces. Reuses `SIGReg` unchanged. **Fit**: modality-agnostic like SIGReg, so it adds no data dependency.
- [ ] (optional) A minimal runnable demo assembling these blocks on a tiny synthetic dynamical system generated inline (`z_{t+1} = f(z_t, a_t)` for a simple transition `f`), so the predictor has real structure to learn rather than noise. This is example code, not a shipped dataset. If a non-synthetic demo is preferred, an optional control-environment dependency used only in the example is an alternative.

Adjacent directions considered:

The same building blocks generalize across the world-model / JEPA family, so beyond LeWM and Sub-JEPA:

- Video and temporal masked latent prediction in the V-JEPA / V-JEPA 2 spirit (latent prediction under multi-block spatiotemporal masking, with an L1 rather than MSE loss, per V-JEPA section 3.1 Eq. 1 and section 3.2, and V-JEPA 2 section 3). **Considered, deferred**: lightly has no video or temporal infrastructure today (no temporal axis in the transforms, and `data/_video.py` only extracts frames as images), so this is a new subsystem rather than a building block. The in-scope sliver is already covered by the prediction loss above, which generalizes to image masked-token prediction (the I-JEPA case lightly supports).
- A non-conditioned temporal predictor variant for plain sequence prediction. **Considered, dropped**: lightly has no temporal/sequence predictor to vary (its predictors are MLP heads and the spatial I-JEPA predictor) and no sequence data for a user to run it on, so there is no clear in-library use today.
- World-model evaluation following LeWM's protocol (latent-probe MSE and correlation, plus planning success, per LeWM sections 4 and 5.1). **Considered, deferred**: that protocol needs control/trajectory data (Two-Room, Reacher, PushT in the Sub-JEPA reproductions), which lightly does not have, and ImageNet does not fit it. lightly's existing kNN/linear harness (#1932) evaluates the encoder, not predictor rollouts.

Happy to revisit the deferred ones if the relevant infrastructure lands, or take other directions you have in mind, based on what is most useful for lightly.

References: [LeWM](https://arxiv.org/abs/2603.19312), [Sub-JEPA](https://arxiv.org/abs/2605.09241), [LeJEPA](https://arxiv.org/abs/2511.08544), [I-JEPA](https://arxiv.org/abs/2301.08243), [V-JEPA](https://arxiv.org/abs/2404.08471), [V-JEPA 2](https://arxiv.org/abs/2506.09985).

Questions to consider for scope refinement:
- Confirm both a torch-native and a timm-backed predictor variant are wanted, matching I-JEPA.
- Confirm the prediction loss should be generic over (predicted, target) tensors, as the draft assumes, so it also serves image masked-token prediction (rather than being specific to the temporal next-step case).

Start by reading lightly/loss/lejepa_loss.py for SIGReg and the IJEPAPredictor/IJEPAPredictorTIMM implementations under lightly/models/modules/. Before coding, resolve which proposed block and predictor variants are in scope; done means the agreed building block is specified sufficiently for a separate PR, with its data assumptions and validation path defined.

Written by the indexing model from the issue text.