Recoleta
GitHub
Source note

AdaJEPA: An Adaptive Latent World Model

Authors
Ying Wang, Oumayma Bounou, Yann LeCun, Mengye Ren
Source
arXiv
Published
Collection
Embodied AI
Latent World ModelTest Time AdaptationModel Predictive ControlJepaRobot Planning
Open arXivMarkdown source

Summary

AdaJEPA adapts a JEPA latent world model during MPC using the transition just observed after each executed action. The paper reports higher goal-reaching success under shape, visual, dynamics, and maze-layout shifts with one gradient step per replan.

Problem

  • Frozen latent world models can give MPC wrong rollouts after visual or dynamics shifts, so the planner optimizes actions that fail in the real environment.
  • This matters for robot planning because small one-step prediction errors can compound across the MPC horizon and lower goal-reaching success.
  • The paper targets adaptation without reward labels, expert demonstrations, or a separate target-domain data collection phase.

Approach

  • AdaJEPA runs a plan, act, adapt, replan loop inside MPC.
  • After executing the first action or action chunk, it stores the observed transition (o_t, a_t, o_{t+1}) in a small online buffer.
  • It applies the same JEPA latent prediction loss used in training: predict the next latent state from the current latent state and action, then compare it to the encoded next observation with stop-gradient on the target.
  • At test time it updates only a small parameter subset, usually final encoder and predictor layers, then immediately uses the updated model for the next MPC plan.
  • The default experiment uses one gradient step per MPC replan, a buffer of the 5 most recent transitions, a maximum of 20 MPC steps, and 50 episodes across each of 3 test-data seeds.

Results

  • On PushObj in-distribution training shapes, the paper reports the largest gain as over 20% relative to the frozen model when adapting to the current shape.
  • On unseen PushObj shapes, the text says AdaJEPA nearly doubles planning success versus the frozen model; exact percentages are not provided in the excerpt.
  • On PointMaze-Medium with low-mass dynamics, predlast + enclast improves GD success from 77.3 ± 8.2 to 80.0 ± 3.3 and CEM success from 82.0 ± 2.8 to 86.7 ± 2.5.
  • On PointMaze-Medium with high damping, predfirst + enclast improves CEM success from 76.0 ± 2.8 to 82.0 ± 3.3, while GD rises from 77.3 ± 5.0 to 78.7 ± 4.7.
  • On unseen PointMaze layouts, predfirst + enclast improves GD success from 53.3 ± 8.2 to 78.7 ± 5.0 and CEM success from 49.3 ± 6.2 to 70.7 ± 3.8.
  • Latency increases are small in the reported table: for the global-feature Temporal Straightening world model, GD time changes from 3.14 s to 3.17 s and CEM time from 0.24 s to 0.27 s; for the spatial-feature variant, GD changes from 3.37 s to 3.38 s and CEM from 5.37 s to 5.39 s.

Link