A delivery robot at a Bengaluru fulfillment centre watches a camera feed of a conveyor belt: packages sliding past under flickering tube lights, a moth occasionally crossing the frame, dust motes catching the light, the belt's rubber texture shifting subtly as it wears. Somewhere in that video stream is the one thing the robot's controller actually needs to predict — where a package's edge will be half a second from now, so a suction gripper can intercept it. Everything else in the frame — the moth, the flicker, the dust, the exact shade of grey on the belt — is noise the robot must learn to ignore, yet a world model that predicts raw pixels is forced to spend capacity modelling all of it, because pixel-reconstruction loss penalizes every pixel equally, relevant or not.
If you have already studied a recurrent latent-dynamics model that compresses observations into a state and imagines forward through it to train a policy, you have seen one answer to "don't work in pixel space": compress first, then predict the compressed representation, then decode when you need to. That family still asks the encoder to preserve enough information to reconstruct something pixel-like, because the training signal ultimately traces back to a reconstruction or reward-prediction loss on decoded outputs. This chapter covers a different, more radical answer, developed by Yann LeCun and collaborators at Meta AI FAIR: architectures that never reconstruct an observation at all, that are trained end-to-end to predict only in representation space, and whose "imagination" is used not to train a reusable policy network but to run an explicit search over action sequences at the moment of decision. The family is called JEPA — Joint-Embedding Predictive Architecture — and it is the theoretical basis for I-JEPA (images), V-JEPA (video), and their action-conditioned descendant V-JEPA 2, which Meta has used for zero-shot robot manipulation planning.
Why predicting pixels is the wrong objective
Formally, an observation o_t (a camera frame) contains far more information than the "world state" that actually matters for control. Write o_t = s_t + n_t informally, where s_t is the causally relevant, predictable part (package position, belt speed) and n_t is high-entropy nuisance detail (specular highlights, sensor noise, a moth). A generative world model trained with a reconstruction loss like ||decode(predict(z_t)) − o_{t+1}||² is graded on reconstructing n_t too, even though n_t is close to unpredictable by definition — its conditional entropy given the past is high. Two costs follow. First, capacity: model parameters and gradient signal get spent fitting noise that carries zero information about the task. Second, and more subtly, a generative model facing unpredictable detail is pulled toward one of two bad habits — blurring the unpredictable region (averaging over its possibilities, which is safe under an L2 loss but throws away sharp task-relevant edges nearby) or hallucinating a single plausible instance of it with false confidence. Neither failure mode touches control quality, yet both cost training compute.
LeCun's 2022 position paper "A Path Towards Autonomous Machine Intelligence" names this directly: a world model should be trained to predict a representation of the future, not the future itself, and should be evaluated with an energy function — a scalar compatibility score between a predicted representation and an observed one — rather than a reconstruction loss. JEPA is the architecture that instantiates this idea.
The JEPA architecture: two encoders and a predictor, no decoder
A JEPA has three trainable-or-derived pieces. A context encoder E_θ maps a visible portion of the input (a masked image, or past video frames) to a latent embedding z_x. A target encoder E_ema maps the portion to be predicted (the masked-out patches, or future frames) to a latent embedding z_y — its weights are not learned by gradient descent at all; they are an exponential moving average (EMA) of the context encoder's weights, updated as θ_ema ← τ·θ_ema + (1−τ)·θ with momentum τ typically around 0.996–0.999. A predictor P_φ, a small transformer, takes z_x plus positional or action tokens describing what is being asked ("predict the patch at this location," or "predict the state after applying this action") and outputs ẑ_y, a guess at the target embedding. Training minimizes the energy D(ẑ_y, z_y) = ||ẑ_y − z_y||², but — and this is the architectural crux — gradients from that loss update only θ and φ. The target-encoder branch is under stop_gradient; it is never differentiated through. It only ever changes via the slow EMA copy.
There is no decoder anywhere in this graph. The model is never asked to produce a viewable image. Its entire competence is "given a context, predict what a slowly-averaged version of my own encoder would output for the target" — a purely self-referential, representation-space objective.
Why the asymmetry stops the model from cheating
A tempting simplification would be to drop the EMA target encoder and just share weights symmetrically: run the same encoder on both context and target, backpropagate the energy loss into both branches. This collapses. The optimizer discovers, almost immediately, that the loss ||ẑ_y − z_y||² is trivially driven to zero by making both encoders output the same constant vector for every input, regardless of content. There is no term in the objective — no reconstruction loss, no classification label — that forces the representation to retain any information about the input at all, so the path of least resistance is to discard all of it. This is the fundamental failure mode of any purely self-referential predictive objective, and it is why naive "predict my own future embedding" ideas do not work out of the box.
JEPA prevents it with two complementary mechanisms, both load-bearing in the real systems. First, the architectural asymmetry itself: because the target branch is stop-gradient and only updated by a slow momentum average of the context branch, the two branches cannot collude to collapse in a single gradient step — the target the predictor is chasing moves independently and slowly, and empirically this asymmetry (borrowed from earlier self-supervised methods such as BYOL) is sufficient to keep the representation non-trivial when combined with sensible initialization and learning-rate schedules. Second, several JEPA variants add an explicit anti-collapse regularizer rather than relying on the EMA trick alone: VICReg (Bardes, Ponce and LeCun, ICLR 2022) adds a variance term that penalizes any embedding dimension whose standard deviation across a batch falls below a threshold, plus a covariance term that decorrelates different dimensions of the embedding — directly outlawing the constant-vector solution by construction, independent of the momentum dynamics.
Worked example: planning with the energy function via CEM
Once a JEPA-style predictor is trained, it can be used to plan without training any separate policy network at all — a genuinely different mode of "imagination" from rolling out a learned policy. The method is Model Predictive Control (MPC) using the cross-entropy method (CEM) for the inner search (Rubinstein, 1997; used with learned dynamics models in robotics by Chua et al., NeurIPS 2018). At each control step: sample a batch of candidate action sequences, roll each one forward through the predictor entirely in latent space, score each resulting latent state against a goal latent state using the energy function, keep the lowest-cost ("elite") sequences, refit the sampling distribution toward them, and execute only the first action of the best sequence before re-observing and replanning. No policy weights are ever updated — the "imagination" is a one-shot search performed fresh at every timestep.
Take our fulfillment-centre robot. Suppose its encoder (already trained, not shown here — encode() is an assumed helper) maps the current camera frame to a 2-D latent gripper-position embedding z0 = (2.0, 3.0), and the encoded goal frame (package correctly positioned under the gripper) is z* = (5.0, 5.0). The predictor has learned a linear latent transition (the true predictor is a transformer, but a trained one often behaves near-linearly in a local latent region, which is exactly what makes it plannable): ẑ_{t+1} = z_t + W·a_t, with W = [[0.9, 0.1], [-0.1, 0.9]] — a small cross-axis coupling reflecting actual actuator crosstalk in the arm.
CEM samples four candidate two-step action sequences (each action a 2-D velocity command):
import numpy as np
W = np.array([[0.9, 0.1],
[-0.1, 0.9]]) # latent predictor transition weights
z0 = np.array([2.0, 3.0]) # encode(current_frame) -- assumed helper, not shown
z_star = np.array([5.0, 5.0]) # encode(goal_frame) -- assumed helper, not shown
candidates = [
[(1, 1), (1, 1)],
[(2, 0), (1, 1)],
[(0, 1), (1, 2)],
[(1, 0), (0, 1)],
]
def predict(z, actions):
z = z.copy()
for a in actions:
z = z + W @ np.array(a)
return z
def energy(z_pred, z_goal):
diff = z_pred - z_goal
return float(diff @ diff)
costs = [energy(predict(z0, seq), z_star) for seq in candidates]
print([round(c, 2) for c in costs])Trace it by hand for candidate 1, [(1,1), (1,1)]: W·(1,1) = (0.9·1+0.1·1, −0.1·1+0.9·1) = (1.0, 0.8). After the first action, z1 = (2.0+1.0, 3.0+0.8) = (3.0, 3.8). The second action adds the same (1.0, 0.8) again: z2 = (4.0, 4.6). Cost is squared distance to (5.0, 5.0): (4.0−5.0)² + (4.6−5.0)² = 1.00 + 0.16 = 1.16. Repeating for the other three candidates gives z2 = (4.8, 3.6), cost 2.00; z2 = (3.2, 5.6), cost 3.60; and z2 = (3.0, 3.8), cost 5.44. The printed output is exactly [1.16, 2.0, 3.6, 5.44].
CEM keeps the k=2 lowest-cost sequences — candidates 1 and 2 — and refits the sampling mean by averaging their actions:
k = 2
elite_idx = np.argsort(costs)[:k]
elite_seqs = [candidates[i] for i in elite_idx]
new_mean_a0 = np.mean([s[0] for s in elite_seqs], axis=0)
new_mean_a1 = np.mean([s[1] for s in elite_seqs], axis=0)
print(new_mean_a0, new_mean_a1)argsort([1.16, 2.00, 3.60, 5.44]) is already [0, 1, 2, 3], so the elites are candidates 1 and 2. Their first actions are (1,1) and (2,0), averaging to (1.5, 0.5); their second actions are (1,1) and (1,1), averaging to (1.0, 1.0). The output is [1.5 0.5] [1. 1.]. Under receding-horizon MPC, the robot executes only the first action, (1.5, 0.5), as its next velocity command, re-encodes the new camera frame, and replans from scratch — the second planned action is discarded and never trusted, because by the time it would be needed the world has already been re-observed.
Common misconception
Having just studied a world model that generates imagined rollouts to train a policy, the natural assumption is that "imagination-based" always means the same thing: a generative simulator producing observation-like trajectories that a policy network learns from offline. It is tempting to conclude that a model which never produces a viewable frame — JEPA's predictor only ever outputs latent vectors — cannot be "imagining" anything and cannot be used for control. This is wrong on both counts. The worked example above is planning: the predictor's forward rolls through latent space are exactly the imagined trajectories, evaluated against a goal with no policy network anywhere in the loop. What differs from policy-based imagination is not whether imagination happens, but where the resulting competence is stored — in a reusable set of policy weights trained once and evaluated cheaply forever after, versus a fresh CEM search re-run at every single control step, with no weights dedicated to "how to act" at all. Meta's action-conditioned V-JEPA 2 (2025) demonstrated real robot arms performing pick-and-place manipulation in previously unseen environments using precisely this CEM-over-latent-energy loop, with zero task-specific policy training — a capability that a purely policy-based imagination approach cannot offer, because a policy network is only as good as the tasks it was trained to imagine toward.
Production considerations
The compute case for JEPA-style world models rests on three concrete savings. First, dimensionality: V-JEPA operates on patch-token embeddings (each covering a spatial-temporal tube of the video, typically summarized in a few hundred dimensions) rather than reconstructing full-resolution pixel grids — a 224×224×3 frame has roughly 150,000 target values to fit under a pixel loss, against a few hundred per patch embedding under JEPA's energy loss, a reduction of two to three orders of magnitude in what the loss actually has to explain. Second, no decoder network exists at all in the trained model, which removes an entire sub-network's worth of parameters and forward/backward compute from every training step — pixel-generative world models typically spend a comparable or larger parameter budget on the decoder as on the encoder. Third, the EMA target encoder does add one extra forward pass per step (it must encode the target region too), but never a backward pass — it carries no optimizer state, no gradient buffers, roughly halving the backward-pass memory that branch would otherwise cost on a training GPU.
The trade-off appears at inference time instead of training time. A policy network trained via imagined rollouts costs a single forward pass to act. CEM-based planning over a JEPA predictor costs one predictor forward pass per candidate sequence per CEM iteration — commonly tens to low hundreds of rollouts, several iterations, every control step. This is real-time-serving-relevant: it moves cost from an offline training budget (amortized once, in a data-centre, before deployment) to an online inference budget (paid every control cycle, on whatever compute rides on the robot or is reachable with acceptable latency). For a warehouse robot with a control loop running at tens of hertz, this bounds how large the predictor and how wide the CEM search can be, and is precisely why production deployments favour predictors that stay small and close to linear in the relevant latent region, and why receding-horizon replanning (rather than trusting a long open-loop plan) is standard — it lets the search stay short per step while still correcting for accumulated model error.
Active recall
Attempt each question before reading its answer.
1. Why does JEPA's predictor operate on encoder outputs z rather than being trained to reconstruct the target pixels y directly?
2. A colleague proposes simplifying JEPA by removing the EMA target encoder and the stop-gradient, sharing one encoder for both branches and backpropagating through both. What happens to training, and why?
3. Suppose an actuator recalibration fixes a sign error in the predictor's coupling term, changing W from [[0.9, 0.1], [-0.1, 0.9]] to W' = [[0.9, 0.1], [0.1, 0.9]]. Recompute the costs for all four candidate sequences from the worked example. Does the elite set change? Does the executed action change?
4. How does CEM-based planning over a JEPA predictor differ from training a policy network on imagined rollouts, in terms of what is learned once versus computed at every control step?
5. Name one concrete way VICReg's regularization prevents representational collapse without relying on the EMA/stop-gradient trick at all.
6. Give one quantitative reason a JEPA-style world model is cheaper to train per step than a pixel-generative world model with a comparably sized encoder.
Answers.
1. Most of a raw pixel target's information is high-entropy nuisance detail (lighting, texture noise, irrelevant background motion) that is nearly unpredictable from the context and carries no information about the task. A reconstruction loss forces the model to spend capacity fitting that detail anyway, and tends to produce blurred or hallucinated outputs in exactly those regions. Predicting in representation space lets the encoder discard nuisance detail during encoding, so the predictor's job is only to model the causally relevant, genuinely predictable part of the future.
2. Training collapses: both encoders converge to output a constant vector for every input, driving the energy ||ẑ_y − z_y||² to zero trivially, because nothing in the objective penalizes discarding all input information. The asymmetry — a slowly-drifting, non-differentiated target — is what prevents the two branches from colluding into this degenerate solution in a single step; removing it removes the only thing keeping the representation non-trivial (short of adding an explicit regularizer like VICReg instead).
3. Recomputing with W': candidate 1 gives z2=(4.0,5.0), cost 1.00; candidate 2 gives z2=(4.8,4.2), cost 0.68; candidate 3 gives z2=(3.2,5.8), cost 3.88; candidate 4 gives z2=(3.0,4.0), cost 5.00. The elite set is still {candidate 1, candidate 2} — unchanged — but their order flips: candidate 2 is now cheaper than candidate 1, where before candidate 1 was cheaper. An unweighted top-k average (as coded above) is blind to internal ordering, so it still outputs the same first action, (1.5, 0.5) — deceptively "unaffected" by the recalibration. But a cost-weighted CEM variant, which many real implementations use, weights each elite's contribution by its rank or inverse cost; since candidate 2 is now the stronger elite, that variant's mean action would shift measurably toward candidate 2's first action, (2, 0). The ripple is real even though the naive unweighted average masks it — a caution about which CEM implementation detail actually matters after a dynamics-model update.
4. Policy-based imagination trains reusable weights once, offline, by rolling the world model forward many times and backpropagating a value or policy-gradient signal into a policy network; acting later costs one forward pass through that network. CEM-over-JEPA trains no policy at all — every control step re-runs a fresh search, rolling the predictor forward for every candidate sequence and scoring against the current goal's energy, paying that search cost online instead of amortizing it in training.
5. VICReg adds an explicit variance term to the loss that penalizes any embedding dimension whose standard deviation across a batch drops below a fixed threshold, directly forbidding the constant-vector (collapsed) solution regardless of what the EMA momentum dynamics are doing — it works as a standalone anti-collapse force, which is why some JEPA variants use it even alongside the EMA trick, as extra insurance.
6. A 224×224×3 pixel frame has on the order of 150,000 target values for a reconstruction loss to fit, while a JEPA target embedding per patch is typically only a few hundred dimensions — two to three orders of magnitude fewer values the loss has to explain — and there is no decoder sub-network's parameters, activations, or gradients to compute at all.
Practice Exercises
Now it is time to practice! Complete these challenges to solidify your understanding:
- Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
- Exercise 2: Find a real-world example where world models and imagination-based learning is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
- Exercise 3: Create a mind-map connecting world models and imagination-based learning to at least 3 other topics you have studied.
Key Takeaways — Summary and Recap
Let us recap what we covered: the core ideas behind world models and imagination-based learning, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.