AdaJEPA: An Adaptive Latent World Model

1 Introduction

• MPC repeatedly rolls out a world model, optimizes a short-horizon action sequence, executes its first action, and replans from the next observation.The paper presents world models with MPC as a standard recipe for goal-conditioned control and planning.Source: §1 · p. 1
• Frozen world-model errors can compound across the planning horizon, making imagined-effective actions fail under real dynamics, especially with visual or dynamics shifts.The introduction argues that even high-capacity models can degrade substantially under small test-environment changes.Source: §1 · p. 1
• AdaJEPA uses each executed action's observed transition as a self-supervised signal to update the model before replanning.Each update changes only a small parameter subset and can use a single gradient step on the latest transition.Source: §1 · p. 2, Figure 1 · p. 2
• AdaJEPA significantly outperforms the frozen model on in- and out-of-distribution goal-reaching tasks, with particularly strong gains when training data is limited.The approach couples the model's action predictions with transition-based updates during deployment.Source: §1 · p. 2

2 Related Work

• Prior JEPA world models are kept frozen at test time, leaving potential performance degradation under distribution shift.JEPAs encode observations into a compact latent space and learn dynamics by predicting future latent targets.Source: §2 · p. 2
• Parthasarathy et al. (2025) reduce the train-test gap through training-time data synthesis, but keep the model fixed during planning and evaluate only in-distribution.Source: §2 · p. 3
• Test-time adaptation commonly uses an unsupervised or self-supervised objective computed from test inputs; AdaJEPA applies this principle to latent planning.Each executed transition supplies the next observation as a self-supervised prediction target.Source: §2 · p. 3
• Unlike approaches requiring target-domain finetuning data, extra rollouts, or outer self-improvement loops, AdaJEPA adapts inside MPC from each executed transition.The update uses no reward labels, expert labels, or separate data-collection phase.Source: §2 · p. 3

3.1 Background: JEPA World Models

• JEPA world models encode observations and actions into latent representations and predict future latent targets, while MPC plans by minimizing distance to the goal latent.MPC executes the first action and replans from the next observation.Source: §3.1 · pp. 3–4

3.2 Closed-Loop Plan-and-Adapt

• At each replanning step, AdaJEPA plans with the current model, executes the first action, observes the transition, adapts, and replans with the updated model.Each episode starts from the pretrained model and maintains its own model copy and transition buffer.Source: §3 · p. 3, §3.2 · p. 4
• The online buffer stores recent transitions and is capped at N; recent-N favors local conditions, while hard-N retains transitions with the largest prediction errors.Source: §3.2 · p. 4
• AdaJEPA minimizes the latent prediction loss used in pretraining, with stop-gradient on the target representation as the default anti-collapse stabilizer.For longer histories, action chunks, or frameskips, the loss averages over valid prediction windows.Source: §3.2 · p. 4
• After each MPC step, AdaJEPA performs U gradient updates on selected parameters, and the adapted model is immediately reused for the next planning problem.The updated parameters can be restricted to a small subset of encoder or predictor parameters.Source: §3.2 · pp. 4–5

4.1 Setup

• The experiments cover PushT and PointMaze, evaluating unseen in-distribution trajectories and out-of-distribution shifts; adaptation uses one gradient step per MPC replanning step by default.The setup uses GD or CEM trajectory optimization and updates only the final visual-encoder and predictor layers unless specified otherwise.Source: §4.1 · pp. 5–6, §4 · p. 5

4.2 Results

• In-distribution, adaptation improves PushObj performance by over 20% in some cases and preserves the strong PointMaze baseline.It yields gains when the frozen model is suboptimal and does no harm when it is near-optimal.Source: §4.2 · p. 6
• Across shape, visual, dynamics, and layout shifts, AdaJEPA improves planning success; gains grow over replanning steps and maze trajectories move closer to shortest paths.Unseen-shape success nearly doubles, while gains are modest under red-anchor and red-block shifts; the frozen PointMaze model is already strong under dynamics shifts.Source: §4.2 · p. 6, Figure 2 · p. 7, Figure 4 · p. 8

4.3 Ablations

• Adaptation targets are environment-dependent: predictor-only updates are less effective for visual and layout shifts, while shape-shift variants perform similarly even with a frozen encoder.Selected-layer updates are already strong across settings, and LoRA does not consistently outperform direct updates.Source: §4.3 · pp. 8–9
• One gradient step at the training learning rate is a strong practical default; recent-transition replay gives the most stable gains, though stronger adaptation can help under shifts.Larger learning rates can be less stable over multiple steps, while smaller ones may require more updates and increase replanning cost.Source: §4.3 · p. 9
• AdaJEPA improves planning across JEPA implementations, representations, architectures, planners, and training objectives, adding 0.01–0.03s per replanning step.The reported comparison uses PushT validation trajectories and one H200.Source: §4.3 · p. 9, Table 2 · p. 9

4.4 Discussion

• With the same 16k-trajectory budget, training on four shapes yields 51.9% unseen-shape success with AdaJEPA, versus 45.8% when all trajectories use one shape.Larger, more diverse data improves both frozen and adapted models; adaptation also raises seen-shape success from 28.1% to 60.8% for K=1, N=1k, exceeding the 43.5% frozen result with 16× more trajectories per shape.Source: §4.4 · pp. 9–10
• Lightweight adaptation is bounded by pretrained-representation coverage and may not fully close the gap when test environments require features absent from training.The authors identify combining adaptation with continual and active learning as a next step.Source: §4.4 · p. 10

• A single gradient step per transition is sufficient to recover substantial planning performance under visual and dynamics shifts.Source: §5 · p. 11
• AdaJEPA uses a plan–execute–adapt–replan loop, reusing each newly observed transition to update the model before the next planning step.Source: §5 · p. 11
• The conclusion argues that latent world models should continue to be trained during deployment rather than kept frozen.Source: §5 · p. 11

A.1 PushT (DINO-WM)

• PushT uses a circular pusher to move a T-shaped block toward a target pose, with validation trajectories from DINO-WM used for comparison (Zhou et al., 2025).Source: §A.1 · p. 14

A.2 PushObj

• PushObj replaces the T block with six shapes; contact-biased data and contact-containing test trajectories address sparse interactions from replaying original actions.The dataset starts from 18,500 PushT training trajectories and generates 16,000 training trajectories per shape; the TLZ+ and T-only subsets support shape- and visual-shift experiments.Source: §A.2 · p. 14

A.3 PointMaze-Medium

• PointMaze-Medium uses 26 open grid cells and resamples distinct start–goal cells until their Euclidean distance exceeds 3 cell units.The setup uses Medium-Maze data and a pretrained ResNet global-feature checkpoint for dynamics-shift experiments.Source: §A.3 · p. 14

• The PointMaze setup uses 25 random 8×8 layouts for training and five held-out layouts for evaluation, with 50 episodes and goals 3–5 cells away by shortest-path distance.Each training layout contributes 2,000 trajectories; evaluation uses 10 episodes per held-out layout.Source: §A.4 · p. 14
• Test-time adaptation improves planning success across distribution shifts and is largely insensitive to the choice of adapted layers.Figure 8 reports per-shift success rates averaged over setups within each shift.Source: Figure 8 · p. 15

B.1 Hyperparameters

• Table 3 lists training hyperparameters.Source: Table 3 · p. 15
• Training and planning hyperparameters are documented separately in Tables 3 and 4, respectively.Source: Table 3 · p. 15, Table 4 · p. 15
• Table 4 lists planning hyperparameters.Source: Table 4 · p. 15

Setup

• The ablations use one gradient step per replan, a five-transition buffer, predictor learning rate 5 × 10−4, and encoder learning rate 10−5 by default.The predictor and encoder are a ViT-style module and a small ResNet, respectively; MPC is capped at 20 steps in these experiments.Source: §B.2 · p. 15

What to adapt

• All tested adaptation targets improve over the frozen model across shape, visual, dynamics, and layout shifts.Adapting the predictor’s last block and encoder’s last stage is consistently competitive; earlier predictor layers can further help under layout shifts, while LoRA does not consistently outperform direct adaptation.Source: §B.2 · p. 16

Learning rates, optimization steps, and replay buffers

• 88% on T and 51% on square, up from 50% and 20% respectively, with adaptation across the tested hyperparameter and replay-buffer settings.Figure 9 presents the adaptation hyperparameter and replay-buffer success-rate comparisons.Source: §B.2 · p. 16, Figure 9 · p. 16
• A large learning rate works well with one step but can overshoot with multiple steps, while a small rate is more stable but needs more updates.The default is the training learning rate with one gradient step per MPC replan, a setting reported as competitive on both shapes at minimal cost.Source: §B.2 · p. 16
• Replay-buffer choice has a smaller effect: success ranges from 81%–87% on T and 35%–44% on square, and every variant beats the frozen model.The main experiments use a recent sliding-window buffer for its most stable gains.Source: §B.2 · pp. 16–17

Setup

• The study varies training shapes K∈{1, 2, 4} and trajectories per shape N∈{1k, 2k, 4k, 8k, 16k}, evaluating each model on all seven shapes.Success is averaged over three seeds and over shapes grouped as seen or unseen.Source: §B.3 · p. 17

Results

• Adaptation improves average success by over 30% on seen shapes and over 15% on unseen shapes, and can offset large reductions in training data.A one-shape model trained with 1k trajectories reaches 61% on seen shapes, exceeding the largest frozen model’s 54%; on unseen shapes, an adapted single-shape model surpasses the best four-shape frozen model once N≥2k.Source: §B.3 · p. 17
• With 16k trajectories total, distributing data across four shapes beats concentrating it on one: adapted success is 81% vs. 77% seen and 52% vs. 46% unseen.The corresponding frozen comparisons are 51% vs. 44% on seen shapes and 36% vs. 29% on unseen shapes.Source: §B.3 · p. 17
• From K=1, N=1k to K=4, N=16k, frozen-model success rises from 28% to 54% on seen shapes and from 20% to 38% on unseen shapes.Both more trajectories and greater shape diversity are beneficial in this comparison.Source: §B.3 · p. 17
• Test-time adaptation is most valuable when data is scarce, with benefits persisting as data scales.Source: §B.3 · p. 17

Planning trajectories

• Examples compare executed trajectories and decoded imagined rollouts for both methods under identical start and goal conditions.The frozen run is shown for 13 steps, while AdaJEPA steps after reaching the goal are grayed out; check marks and crosses indicate success or failure.Source: C Visualization of Planning Trajectories · p. 18
• AdaJEPA reduces latent prediction loss by refining predictions against observations reached through executed actions; the frozen model often fails with inaccurate predictions.The bottom panel tracks loss at each replanning step, comparing a short predicted latent rollout with latents encoded from reached observations.Source: C Visualization of Planning Trajectories · p. 18

Shape shifts

• The seen-shape example trains on T, L, +, and Z and tests on +.Source: Figure 11 · p. 18
• The shape-shift example tests an unseen smallT shape after training on T, L, Z, and +; its decoder aims for the closest seen shapes.Source: Figure 12 · p. 18

Visual shifts

• Under salt-and-pepper noise and a red-agent test shift, AdaJEPA consistently decreases prediction loss; visual-shift adaptation also leads to better planning.The decoder, trained only on default PushT data, tends to reconstruct the original colors and scenes.Source: C Visualization of Planning Trajectories · p. 19, Figure 13 · p. 19, Figure 14 · p. 19