PaperlayerView exact revision on arXiv

Source-linked AI summary

JEPA-TTT: Persistent Test-Time Training of Latent World Models for Planning under Dynamics Shifts

Zheyuan Zhang, Suyu Ye, Nakul Agarwal, Hossein Nourkhiz Mahjoub, Ehsan Moradi Pari, Daniel Khashabi, Tianmin Shu, Vaishnav Tadiparthi

arXiv:2610.00722v1cs.LG

TL;DR

Test-time dynamics shifts can make pretrained world models predict transitions incorrectly, impairing planning. JEPA-TTT persistently adapts a pretrained JEPA latent dynamics predictor from self-supervised online experience; across eight shifts it improves planning on every shift.

  • Problem

    When test-time transition dynamics differ from training, pretrained world models can make systematic prediction errors that lead to poor planning.

    Source: §1 · p. 1

  • Method

    JEPA-TTT updates a pretrained JEPA predictor from self-supervised latent prediction targets and carries predictor, optimizer, and replay-buffer state across episodes.

    Source: §1 · p. 2, §3.3 · pp. 5–6

  • Results

    Across eight dynamics shifts, JEPA-TTT improves planning on every shift; mean normalized held-out planning-score AUC is 0.571 versus 0.267 for the frozen JEPA world model.

    Source: §1 · p. 2, §5 · p. 13

  • Takeaways & Limitations

    Ablations support dense temporal coverage, with a small additional benefit from replay.

    Source: §4.4 · p. 12, §5 · p. 13

    The study holds one dynamics change fixed throughout test time and does not address simultaneous visual or reward shifts.

    Source: §6 · p. 13

Abstract

from arXiv · show

World models enable agents to plan by predicting future states of the environment, but their predictions can become unreliable when test-time dynamics differ from those seen during training. We present JEPA-TTT, which adapts the latent dynamics predictor of a pretrained action-conditioned Joint-Embedding Predictive Architecture world model throughout test time. Self-supervised updates accumulate across episodes, while the visual encoder and reward head remain fixed, preserving the pretrained representation and task objective. Planning requires neither a goal image nor online environment reward. JEPA-TTT uses dense replay, which forms prediction windows at every temporal offset, retains them in a growing buffer, and samples minibatches from that buffer for predictor updates. Across eight dynamics shifts in four continuous-control environments, JEPA-TTT improves planning on every shift. After 500 test-time episodes, it reduces autoregressive latent prediction error by 83% on average and improves planning performance by 153% over the frozen JEPA world model. These results show that persistent self-supervised test-time training can adapt a pretrained latent world model under changed dynamics.

1. Introduction

Pretrained world models can mispredict when test-time transitions differ from training dynamics, undermining planning. JEPA-TTT adapts the latent predictor from online experience and improves planning across the evaluated shifts.

1. Introduction

  • A mismatch between training and test transitions can produce systematic prediction errors and poor planning.The shifts include changes to actuator responses, contact dynamics, or object motion.Source: §1 · p. 1
  • JEPA-TTT adapts a pretrained JEPA latent dynamics predictor through self-supervised updates that accumulate across episodes.Dense replay forms prediction windows at every temporal offset and samples minibatches from a growing buffer; the visual encoder and offline-trained reward head remain fixed.Source: §1 · p. 2
  • Planning uses the adapted predictor and fixed reward head without requiring goal images or online environment rewards.The offline-trained reward head preserves the task objective while the visual representation remains fixed.Source: §1 · p. 2
  • JEPA-TTT improves planning on all eight shifts; mean normalized held-out planning-score AUC is 0.571 versus 0.267 for the frozen JEPA world model.After 500 test-time episodes, it reduces autoregressive latent prediction error by 83% on average.Source: §1 · p. 2

2. Related Work

Related work covers latent world models, action-conditioned JEPA planning, test-time training, and online dynamics adaptation. JEPA-TTT focuses on persistent adaptation of a pretrained JEPA world model during deployment.

2. Related Work

  • Prior JEPA world-model work primarily studies model training and planning, whereas JEPA-TTT investigates persistent adaptation to dynamics shifts during deployment.Action-conditioned planning examples include V-JEPA 2-AC [Assran et al., 2025], DINO-WM [Zhou et al., 2024], and PLDM [Sobal et al., 2025a].Source: §2 · p. 3
  • Test-time training commonly adapts model parameters from inference-time data using self-supervised objectives or entropy minimization.Online variants can carry adaptation state across sequential inputs.Source: §2 · p. 3
  • Unlike concurrent AdaJEPA [Wang et al., 2026a], JEPA-TTT retains adaptation across episodes, uses a frozen reward head without a goal image, and applies dense replay over accumulated transitions.The passage describes AdaJEPA as resetting its model and buffer each episode and using a short within-episode history.Source: §2 · p. 4

3. Persistent Test-Time Training of Latent World Models

The paper studies planning after a persistent change in transition dynamics while the task remains fixed. JEPA-TTT adapts a latent dynamics predictor from online experience while keeping the visual encoder and reward head fixed.

3.1. Problem Setting

  • JEPA-TTT evaluates finite-horizon control after one transition-dynamics change that remains fixed at test time, while the task and its score function stay unchanged.Source: §3.1 · p. 4

3.2. Action-Conditioned JEPA World Model and Planning

  • The action-conditioned JEPA predictor rolls candidate action sequences forward in latent space, and a frozen reward head scores predicted trajectories for receding-horizon planning.The encoder and reward head remain fixed, so test-time adaptation changes planning through dynamics prediction.Source: §3.2 · pp. 4–5

3.3. JEPA-TTT: Persistent Self-Supervised Test-Time Training

  • At test time, JEPA-TTT trains the predictor to match predicted latents to frozen-encoder targets, while predictor parameters, optimizer state, and the online buffer persist across episodes.The qualitative Figure 3 examples show the adapted predictor recovering the correct ranking of candidate plans under dynamics shift while both planners use the same reward head.Source: §3.3 · pp. 5–6, Figure 3 · p. 6

3.4. Dense Replay

  • Dense replay forms prediction windows at every valid temporal offset, retains them in a growing buffer, and samples minibatches uniformly for predictor updates.The comparisons distinguish added optimization, broader temporal coverage, and replay from the accumulated dense-window set under matched update budgets.Source: §3.4 · pp. 6–7

4. Experiments

Across eight controlled dynamics shifts, JEPA-TTT improves held-out planning over the frozen JEPA model and reduces autoregressive latent prediction error. Ablations indicate benefits from temporal coverage and persistent updates.

4.1. Experimental Setup

  • Across four continuous-control environments, two shifts per task preserve observation mapping, task objective, action space, and episode horizon; deployments run 500 episodes with three baselines.The baselines are frozen JEPA, PPO-TTT, and AdaJEPA.Source: §4.1 · pp. 8–9

4.2. Test-Time Training Improves Planning under Dynamics Shifts

  • 153% higher best score and 113% higher mean AUC than frozen JEPA, with JEPA-TTT outperforming all three baselines on both metrics across all eight shifts.Figure 5 summarizes best held-out score and normalized held-out planning-score AUC from episode 0 to 500.Source: Figure 5 · p. 9, §4.2 · p. 11
  • On six of eight shifts, JEPA-TTT’s best score is very close to or slightly above the offline reference after 500 online episodes; that reference is neither deployable nor an upper bound.Figure 6 tracks held-out performance over time on a fixed set of 100 episodes.Source: Figure 6 · p. 10, §4.2 · p. 11

4.3. Test-Time Training Improves Dynamics Prediction

  • 83% lower autoregressive latent prediction MSE on average after 500 test-time episodes, with reductions on every shift in evaluations using identical held-out trajectories.The comparison predicts five latent states without intermediate observations.Source: §4.3 · p. 11, Figure 7 · p. 11

4.4. Ablation Studies

  • Dense replay achieves the highest aggregate best score and mean AUC; versus the dense stream, mean AUC rises from 0.548 to 0.571 while best scores are similar.The ablation varies update count, temporal coverage, and replay.Source: §4.4 · p. 12, Table 1 · p. 12
  • With the same number of predictor updates, retaining adaptation across episodes increases best score from 0.289 to 0.729 and mean AUC from 0.273 to 0.637.The matched control uses episode-local dense windows.Source: §4.4 · p. 12

5. Conclusion

JEPA-TTT adapts a pretrained world model during deployment by updating only its latent dynamics predictor and retaining it across episodes. It improves planning across eight dynamics shifts, with lower latent prediction error and ablations supporting dense temporal windows.

  • JEPA-TTT improves planning on every one of eight dynamics shifts, with gains accompanied by lower latent prediction error.The predictor is retained across episodes so adaptation can accumulate over time.Source: §5 · p. 13
  • JEPA-TTT updates only the latent dynamics predictor from observed transitions while keeping the visual encoder and reward head fixed.The adapted predictor is retained across episodes.Source: §5 · p. 13
  • Ablations show the benefit of dense temporal windows, with a small further gain from replay.Source: §5 · p. 13

6. Limitations

The evaluation studies one dynamics change that remains fixed throughout test time, while keeping the visual observation mapping, task score, action interface, and episode horizon unchanged. It does not assess sequences of dynamics changes, return to training dynamics, or retention and forgetting in continual learning.

Scope of dynamics shifts

  • The study covers a single dynamics change that remains fixed throughout test time, with the observation mapping, task score, action interface, and episode horizon unchanged.Source: §6 · p. 13
  • Because the encoder and reward head remain frozen, JEPA-TTT is not designed for simultaneous visual or reward shifts.Source: §6 · p. 13

Continual adaptation and catastrophic forgetting

  • The study does not test sequences of dynamics changes or a return to training dynamics, and persistent updates may degrade predictions for previously learned dynamics.Source: §6 · p. 13
  • Retention, forward or backward transfer, and catastrophic forgetting are not measured in continual learning settings.Source: §6 · p. 13

A1.1. World-Model Pretraining

The pretrained world models use a shared LeWorldModel configuration with a ViT-Tiny encoder and autoregressive predictor. Each task has 3,000 training-dynamics episodes for collection, and the reward head is fitted after pretraining on an episode-disjoint split.

A1.1. World-Model Pretraining

  • All four world models share a LeWorldModel configuration with a ViT-Tiny image encoder and a six-layer autoregressive predictor.The encoder uses 14 × 14 patches and 192-dimensional outputs; the predictor has 16 attention heads and hidden dimension 192.Source: A1.1. World-Model Pretraining · p. 18
  • Each task uses 3,000 complete episodes under training dynamics for data collection, followed by scratch training for 20 epochs with a random 90/10 window split.The epoch-20 checkpoint is predeclared rather than selected using test performance.Source: A1.1. World-Model Pretraining · p. 18
  • The reward head is fitted after pretraining on an episode-disjoint split collected under training dynamics.The optional decoder is also trained but neither it nor its output is used for planning.Source: A1.1. World-Model Pretraining · p. 18

A1.2. Planning and Test-Time Training

The shared test-time protocol compares dense and sparse training-window construction, with data collection continuing to the fixed episode horizon and updates occurring only when full batches are available. Dense replay retains valid windows for later updates, and the evaluation runs share one pretrained model.

A1.2. Planning and Test-Time Training

  • Dense construction makes a window start valid after a clip spanning (C+1)K raw observation rows is available, whereas sparse construction waits for C complete action blocks.Source: A1.2. Planning and Test-Time Training · p. 18
  • Test-time interaction continues to the fixed episode horizon after success, and the sparse and dense streams update only when a full batch is available.Held-out evaluation episodes use normal task termination.Source: A1.2. Planning and Test-Time Training · p. 18
  • The shared test-time planning and training settings define the protocol for all eight shifts.Three deployment runs start from the same pretrained model; reported standard deviations reflect interaction and adaptation variability, not variation across independently pretrained models.Source: Table A1 · p. 19, A1.2. Planning and Test-Time Training · p. 20
  • Dense replay receives the same number of updates as the dense stream, with all stored windows remaining valid for later updates.Each replay batch samples without replacement, and incomplete batches remain pending at scheduled evaluations and episode 500.Source: A1.2. Planning and Test-Time Training · p. 20

A2. Tasks and Dynamics Shifts

The task setup holds the goal, score, action bounds, and initial-state distribution fixed; each shift changes only how commanded actions produce the next state.

A2. Tasks and Dynamics Shifts

  • Across the shifts, only the action-to-next-state dynamics change; the visual goal, score, action bounds, and initial-state distribution stay fixed.These conditions define the scope of the task shifts.Source: A2. Tasks and Dynamics Shifts · p. 20

A2.1. PushT

The appendix details JEPA-TTT’s evaluation across four tasks and reports planning gains, lower latent prediction error, and controls on persistence and continued training. It also describes diagnostic limits and the scope of the offline reference.

A2. Tasks and Dynamics Shifts

A3. Baseline Details

  • AdaJEPA adaptation resets each episode, unlike PPO-TTT’s persistent policy and optimizer; the two AdaJEPA trainable-component settings differ by 0.0017, and the default is pred-last+enc-last.The AdaJEPA comparison uses the same pretrained model, reward head, and CEM planner, without a goal image.Source: A3. Baseline Details · pp. 21–22, Table A3 · p. 25

A4. Complete Planning Results

  • JEPA-TTT outperforms all three baselines on every shift under both planning metrics; dense replay has the highest aggregate best score and AUC among its ablations.Figure A1 shows the ablation comparison; Table A2 reports the baseline comparison.Source: Table A2 · p. 23, Figure A1 · p. 24

A5.1. Fixed-Trajectory Latent Error

  • Five-block autoregressive latent MSE falls after dense replay on all eight shifts, with reported reductions ranging from 45% to 99%.The fixed-trajectory evaluation gives the frozen and adapted models identical inputs; raw MSE magnitudes should not be compared across tasks.Source: A5.1. Fixed-Trajectory Latent Error · p. 25, Table A4 · p. 26

A5.2. Decoded Predictions

A6. Matched Controls

A6.3. Training-Dynamics Checks

  • PushT reward-head diagnostics show sparse high-score validation data and systematic underprediction; a score-balanced head still yields no strict CEM successes.The reported limitation is specific to PushT under the training dynamics.Source: A6.3. Training-Dynamics Checks · p. 30

A6.4. Offline Reference under the Test-Time Dynamics

Loading 2610.00722v1…