PaperlayerView exact revision on arXiv

Source-linked AI summary

DT-JEPA: Learning to Predict Future Representations for Decision-Making

Changxu Ni, Yiyang Ni, Jun Li, Feiran You, Hongyang Du, Haitao Zhao, Shi Jin

arXiv:2610.32657v1cs.LG

TL;DR

Offline RL seeks reusable policies from heterogeneous trajectories, while conventional world-model reconstruction can model details not directly relevant to decisions. DT-JEPA predicts probabilistic future representations from reward-free offline data and fine-tunes with a critic; it reports strong benchmark performance and efficient adaptation to unseen rewards.

  • Problem

    Offline RL seeks reusable policies from heterogeneous multi-task trajectories, while reconstruction-based world models may model details not directly relevant to downstream decisions.

    Source: §1 · pp. 1–2

  • Method

    DT-JEPA pretrains a probabilistic predictor of future latent representations from reward-free trajectories, uses them to guide DT action generation, and fine-tunes with a critic on new-task trajectories.

    Source: §1 · p. 2

  • Results

    DT-JEPA achieves strong performance and efficient adaptation, and experiments report superiority over existing DT variants and competitive offline RL methods, including adaptation to unseen rewards with limited interactions.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    The experiments support predictive future representations as effective conditioning information for generalizable sequence decision making.

    Source: §6 · p. 9

    The finite-sample bound assumes bounded occupancy mismatch, a finite policy class, bounded log-likelihood variation, and approximation error.

    Source: §3.4 · p. 6

Abstract

from arXiv · show

Decision Transformer (DT) has demonstrated strong performance in various offline reinforcement learning (RL) tasks by conditioning action generation on expected future returns and historical trajectories. However, treating expected future returns as the condition provides an incomplete description of the future, as trajectories with different state transitions and behaviors may share similar returns. Moreover, expected future returns inherently depend on task-specific reward functions, which makes policy learning reliant on reward-labeled datasets and restricts its scalability to tasks with similar but unseen rewards. Recently, Joint Embedding Predictive Architectures (JEPA) have emerged as a promising paradigm of self-supervised learning-based world models, which learn predictive representations of future dynamics in latent space rather than reconstructing future states in the original space. Inspired by JEPA, we propose DT-JEPA, an offline pretraining and online fine-tuning framework that learns general knowledge from reward-free offline trajectories and achieves excellent adaptation to unseen tasks. Instead of predicting a single deterministic future embedding, we design a novel probabilistic predictor to model the conditional distribution of future representations in the latent space, thereby capturing both predictive information and uncertainty about future dynamics. Theoretically, we show that 1) the JEPA objective seeks to maximize a variational lower bound on the mutual information between historical representations and future trajectories; and 2) predictive future-conditioned policy learning admits a finite-sample performance bound with an O(N-1/4) estimation term. Experiments on D4RL benchmarks demonstrate that DT-JEPA achieves efficient adaptation across varying offline RL tasks, outperforming traditional DT-based and other offline RL methods.

1 INTRODUCTION

Offline RL aims to learn reusable decision knowledge from heterogeneous trajectories, but return-conditioned policies can miss behavioral differences and depend on reward-labeled data. DT-JEPA learns probabilistic future representations from reward-free trajectories and reports strong benchmark performance and adaptation to unseen rewards.

1 INTRODUCTION

  • Return-conditioned policies can conflate distinct future behaviors with similar returns and require reward-annotated trajectories, limiting pretraining scalability and transferability.The introduction frames learning future-dynamics distributions from observable history as an open question.Source: §1 · p. 2
  • JEPA predicts latent representations of future observations from observed context rather than reconstructing observations, targeting predictive information without fine-grained reconstruction.The paper positions this as a way to learn decision-relevant dynamics from reward-free trajectories.Source: §1 · p. 2
  • DT-JEPA combines reward-free offline pretraining with online fine-tuning, using a probabilistic latent-future predictor, AdaLN to guide actions, and a critic to adapt to unseen rewards.The critic estimates action values under the target reward and guides policy adaptation using newly collected trajectories.Source: §1 · p. 2
  • The JEPA objective has a variational lower bound on mutual information, and future-conditioned policy learning has a finite-sample bound with an O(N^-1/4) estimation term.On D4RL, DT-JEPA outperforms existing DT variants and competitive offline RL methods and adapts efficiently to unseen rewards with limited interactions.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries define offline RL and DT as sequence modeling conditioned on return-to-go, then introduce JEPA's latent prediction and AdaLN's conditional feature modulation.

2.1 OFFLINE REINFORCEMENT LEARNING AND DECISION TRANSFORMER

  • Offline RL learns from a fixed dataset without environment interaction during training; DT predicts actions from trajectory history and a desired return-to-go.DT uses supervised action prediction over reward-augmented trajectories, with return-to-go as its conditioning signal.Source: §2.1 · p. 3

2.2 JOINT-EMBEDDING PREDICTIVE ARCHITECTURE

  • JEPA predicts target-data embeddings from context representations instead of reconstructing the original observations, encouraging representations predictive of the target.The target encoder is updated by an exponential moving average of the context encoder.Source: §2.2 · p. 3

2.3 ADAPTIVE LAYER NORMALIZATION

  • AdaLN maps a conditional representation to scale and bias parameters that modulate normalized features, injecting conditioning without adding it to the input sequence.The normalization uses feature-wise mean and variance, with a constant for numerical stability.Source: §2.3 · pp. 3–4

3 METHOD

DT-JEPA conditions a Decision Transformer on predicted future representations, then adapts the pretrained policy to unseen reward objectives with task-specific critics. Its design addresses the ambiguity of return-to-go conditioning by learning predictive latent representations from offline trajectories.

3.1 MOTIVATION

  • Different future trajectories can share similar returns, motivating a condition that captures more about future dynamics than return-to-go alone.The method frames this as reducing uncertainty about future trajectories using a predictive representation.Source: §3.1 · p. 4

3.2 VARIATIONAL JOINT-EMBEDDING FUTURE CONDITIONING

  • A history encoder maps recent state-action context to a historical representation, while a target encoder represents future trajectories from offline data.The target encoder is updated via exponential moving average.Source: §3.2 · p. 5
  • A probabilistic predictor samples future latents conditioned on history, modeling a distribution that captures predictive information and uncertainty.The JEPA objective is motivated by a variational lower bound on mutual information between historical representations and future trajectories.Source: §3.2 · p. 5

3.3 ADALN-BASED LATENT CONDITIONED ACTION GENERATION

  • AdaLN uses the predicted future latent to scale and shift the transformer's historical representation before the policy predicts an action.Two lightweight MLPs map the future latent to dimension-wise modulation parameters.Source: §3.3 · p. 5

3.4 REWARD-FREE PRETRAINING AND REWARD-ORIENTED FINETUNING

  • Reward-free pretraining combines behavior cloning on multi-task trajectories with JEPA prediction loss.The loss balances the behavior-cloning and JEPA terms.Source: §3.4 · p. 6
  • Under the stated assumptions, the learned predictive-conditioned policy approaches its population counterpart as offline data grows, up to approximation error.The finite-sample analysis assumes bounded occupancy mismatch, a finite policy class, bounded log-likelihood variation, and bounded approximation error.Source: §3.4 · p. 6
  • During online adaptation, task-specific critics guide policy optimization toward higher expected long-term return under the new reward objective.The pretrained policy first collects task-specific trajectories, including a high-return subset for critic learning.Source: §3.4 · pp. 6–7

4 EXPERIMENTS

The experiments assess whether reward-free offline data supports transferable policy adaptation and transfer to unseen tasks, and examine the roles of latent future representations and online fine-tuning. Evaluation uses D4RL tasks and compares DT-JEPA with offline RL and DT-based methods.

4 EXPERIMENTS

  • The evaluation asks whether reward-free datasets support downstream adaptation, whether pretrained knowledge transfers to unseen tasks, and how latent futures and online fine-tuning contribute.These are the three stated research questions guiding the experiments.Source: §4 · p. 7

4.1 BASELINES

  • DT-JEPA is compared with offline RL and DT-based baselines, including TD3-BC, TD-JEPA, QT, ODT, Prompt DT, VDT, LSDT, and PDT.The baselines cover value-guided, online-adapted, demonstration-conditioned, and trajectory-modeling approaches.Source: §4.1 · p. 7

4.2 BENCHMARK DATASETS AND EVALUATION PROTOCOL

  • The benchmark protocol covers D4RL locomotion and navigation datasets, reporting normalized return means and standard deviations across five random seeds.Locomotion includes Hopper, Walker2d, HalfCheetah, and Ant; navigation includes AntMaze and Maze2D.Source: §4.2 · p. 8, Table 1 · p. 8

5 EXPERIMENTAL RESULTS

DT-JEPA performs strongly across offline benchmarks and adapts efficiently to unseen reward objectives. Ablations favor probabilistic future prediction, AdaLN modulation, and online fine-tuning.

5 EXPERIMENTAL RESULTS

  • DT-JEPA ranks first on 10 of 12 Gym MuJoCo datasets, beats PDT on 18 of 19 datasets, and consistently surpasses TD-JEPA.It remains competitive on navigation tasks and ranks first on three of four medium-replay datasets.Source: §5 · p. 8
  • DT-JEPA converges faster than Prompt DT, ODT, and PDT, and outperforms competitors across four MuJoCo tasks averaged over three unseen reward objectives.Figure 3 evaluates fine-tuning performance every 10 episodes, averaging three seeds with 95% confidence intervals; Figure 4 reports average normalized return on three unseen objectives, with 95% confidence intervals.Source: §5 · p. 8, Figure 3 · p. 9, Figure 4 · p. 9
  • Future-conditioned variants outperform scalar return conditioning, probabilistic JEPA improves over deterministic prediction, and AdaLN performs better overall than token injection.Removing critic guidance or freezing the pretrained backbone also degrades performance.Source: §5 · pp. 8–9

6 CONCLUSION

DT-JEPA combines JEPA with Decision Transformer for transferable decision making, while the related work situates this approach among sequence-modeling, latent-prediction, and reward-free methods. The appendix proof connects the JEPA objective to mutual information between historical and future representations.

6 CONCLUSION

  • DT-JEPA learns probabilistic future representations from offline trajectories, conditions action generation through AdaLN, and uses task-specific critics for adaptation to unseen reward objectives.The paper also analyzes the predictive-learning objective and derives a finite-sample performance bound for future-conditioned policy learning.Source: §6 · p. 9

A.1 DECISION TRANSFORMERS AND CONDITIONAL SEQUENCE MODELING

  • Decision Transformer conditions actions on desired returns, while later methods extend sequence-policy conditioning to demonstrations, future states, trajectory statistics, or task representations.These studies frame conditioning information as central to sequence-based policy generalization.Source: §A.1 · p. 13

A.2 JEPA-BASED REPRESENTATION LEARNING FOR RL

  • JEPA predicts target latent representations from observed context rather than reconstructing inputs, and RL work applies latent prediction to policy-conditioned dynamics and actor–critic learning.The cited approaches include visual-observation learning and masked spatiotemporal prediction.Source: §A.2 · p. 13

A.3 REWARD-FREE PRETRAINING AND ZERO-SHOT RL

  • Reward-free methods including Forward-Backward, HILP, and RLDP extract reusable behavioral knowledge from offline data without task-specific reward supervision.They pursue downstream control through occupancy representations, temporally structured latent spaces, or reusable policy representations.Source: §A.3 · p. 13

B.1 PROOF OF THEOREM 1

  • The JEPA proof rewrites mutual information as future-representation entropy plus expected conditional log-likelihood, then uses a predictive distribution to approximate the true conditional.The true distribution p(Ft | Ht) is described as generally difficult to obtain directly.Source: §B.1 · p. 14

B.2 PROOF OF THEOREM 2

Theorem 2 analyzes how a learned predictive-conditioned policy compares with its population counterpart from finite offline data. Its argument uses likelihood estimation and occupancy-distribution discrepancy under stated regularity assumptions.

B.2 PROOF OF THEOREM 2

  • Theorem 2 bounds the learned policy’s performance relative to the population predictive-conditioned policy under bounded occupancy mismatch, a finite policy class, bounded log-likelihood variation, and approximation error.The analysis considers a finite-horizon MDP with rewards in [0, 1].Source: §B.2 · p. 15
  • For continuous-control tasks, the policy’s action distribution is treated as a fixed-covariance Gaussian, making negative log-likelihood minimization equivalent to behavior cloning.The proof conditions on a given predictive model and samples predicted future representations from it.Source: §B.2 · pp. 15–16
  • The proof relates policy discrepancy to occupancy mismatch, then bounds return differences using bounded rewards and empirical predictive-conditioned risk from offline trajectory segments.It assumes the population occupancy is absolutely continuous with respect to the behavior distribution and has finite mismatch coefficient.Source: §B.2 · pp. 17–19

C.1 EXPERIMENTAL TASKS AND DATASETS

The evaluation spans MuJoCo locomotion and AntMaze/Maze2D navigation settings that vary in dynamics, maze complexity, and offline-data diversity. Baselines cover conventional offline RL, sequence modeling, value guidance, and reward-free future-conditioned pretraining.

C.1 EXPERIMENTAL TASKS AND DATASETS

  • The benchmarks cover four MuJoCo locomotion environments and navigation settings whose maze scale, horizon, and trajectory diversity vary.MuJoCo includes Hopper, HalfCheetah, Walker2d, and Ant; AntMaze and Maze2D provide progressively larger navigation configurations.Source: §C.1 · pp. 20–22
  • Offline locomotion datasets are medium, medium-replay, and medium-expert, differing in behavior quality and diversity from coherent suboptimal data to mixed-quality trajectories.AntMaze configurations also vary in maze layout and navigation behavior diversity, while Maze2D configurations mainly vary in maze scale and navigation complexity.Source: §C.1 · pp. 21–22

C.2 BASELINES

  • The baseline set spans behavior-regularized actor–critic TD3-BC, offline-to-online ODT, value-guided QT and VDT, and trajectory-prompted Prompt DT.These methods provide comparisons across non-sequence offline RL, sequence modeling, online adaptation, and value-guided policy learning.Source: §C.2 · pp. 22–23
  • Reward-free TD-JEPA and future-conditioned PDT provide related comparisons for latent predictive representation learning and reward-free sequence-model pretraining.PDT uses a future representation to condition action prediction; TD-JEPA learns policy-conditioned latent dynamics and maps downstream rewards to its task space.Source: §C.2 · pp. 22–23

D DT-JEPA TRAINING AND DEPLOYMENT ALGORITHMS

DT-JEPA pretrains a future-conditioned policy and predictive latent model from reward-free offline trajectories, then adapts the policy online to unseen tasks using task-specific critics. The appendix also specifies shared experimental hyperparameters.

D DT-JEPA TRAINING AND DEPLOYMENT ALGORITHMS

  • Reward-free pretraining encodes history and its succeeding trajectory segment, then predicts a distribution over future latents conditioned on the history.The predicted latent conditions action generation through AdaLN; predictive and behavior-cloning losses are jointly optimized, with an EMA-updated target encoder.Source: D DT-JEPA TRAINING AND DEPLOYMENT… · pp. 23–24
  • Online adaptation collects trajectories for the unseen task, selects a high-return subset, and trains task-specific critics to guide policy updates toward higher expected long-term return.Twin critics use temporal-difference targets based on the minimum target-critic value, and target networks are maintained with EMA.Source: D DT-JEPA TRAINING AND DEPLOYMENT… · p. 24

E HYPERPARAMETERS

F CONSTRUCTION OF UNSEEN REWARD TASKS

The unseen-task tests preserve each MuJoCo environment’s dynamics while changing its reward, introducing forward, backward, and jump objectives only during adaptation. Additional analyses examine how sampled future latents affect actions and how future conditioning and AdaLN relate to attention patterns.

F CONSTRUCTION OF UNSEEN REWARD TASKS

  • Unseen downstream tasks modify only reward in four MuJoCo environments, leaving state and action spaces, dynamics, and termination conditions unchanged.Each environment is evaluated on forward, backward, and jump objectives that are withheld during reward-free pretraining and introduced during adaptation.Source: F CONSTRUCTION OF UNSEEN REWARD TASKS · pp. 25–26
  • The jump objective removes horizontal-velocity reward and encourages upward displacement, while forward and backward objectives reward opposite horizontal locomotion directions.Walker2d’s jump reward uses a horizontal-velocity penalty and height term; Ant retains a contact penalty across its objectives.Source: F CONSTRUCTION OF UNSEEN REWARD TASKS · p. 26

G.1 EFFECT OF PREDICTED FUTURE REPRESENTATIONS ON THE ACTION DISTRIBUTIONS.

  • With history held fixed, different sampled future latents induce distinguishable action distributions across action dimensions and timesteps, indicating that the policy uses the predicted latent.The analysis samples three predicted futures on Walker2d and conditions actions through AdaLN.Source: §G.1 · p. 27, Figure 7 · p. 28

G.2 EFFECT OF PREDICTED FUTURE REPRESENTATIONS ON THE ATTENTION SCORE.

  • DT-JEPA has the most concentrated diagonal attention pattern among the compared variants, with future conditioning and AdaLN described as complementary contributors.The comparison includes DT, RTG-AdaLN, JEPA-Token, and DT-JEPA; the diagonal pattern is interpreted as greater emphasis on temporally adjacent, dynamically relevant information.Source: §G.2 · pp. 27–28, Figure 8 · p. 29

G.3.1 SENSITIVITY TO λJ DURING TRAINING

The preferred JEPA-loss weight varies by environment, although performance remains relatively stable over a moderate range. Deployment-stage sensitivity likewise shows task-dependent trade-offs between predictive learning and critic guidance.

G.3.1 SENSITIVITY TO λJ DURING TRAINING

  • Hopper performs best around λJ = 0.1, while Walker2d, AntMaze, and Maze2D-Large favor around λJ = 1; HalfCheetah and Maze2D-Umaze do better with smaller weights.Performance remains relatively stable over a moderate λJ range, but excessively large weights may reduce emphasis on accurate action generation.Source: §G.3.1 · p. 29, Figure 9 · p. 30

G.3.2 SENSITIVITY TO λ′J AND λQ DURING DEPLOYMENT

  • During deployment, Hopper favors λJ = 0.1, while Walker2d, HalfCheetah, and Ant perform well around λJ = 1 or 0.1; critic-weight effects vary with adaptation stage and environment.Larger λQ can accelerate early adaptation in some environments, whereas smaller values may yield more stable later improvements; Walker2d favors λQ = 0.1 late and Ant benefits from λQ = 0.01.Source: §G.3.2 · pp. 29–30, Figure 10 · p. 30

G.4 ANALYSIS ON TRAJECTORY STITCHING ABILITY

The analysis examines trajectory stitching and how predictor depth and sequence length affect performance across locomotion and navigation tasks. DT-JEPA retains strong performance on fragmented datasets, while the sensitivity analyses show that useful capacity and context lengths vary by task.

G.4 ANALYSIS ON TRAJECTORY STITCHING ABILITY

  • DT-JEPA maintains strong performance on fragmented datasets, with future representations providing predictive structure for connecting effective behavior segments from different trajectories.The advantage is especially evident in long-horizon navigation, though QT remains competitive on several datasets.Source: §G.4 · p. 31

G.5 SENSITIVITY TO THE JEPA PREDICTOR DEPTH

  • Average normalized return rises from 79.7 with one predictor block to 87.4 with three, then slips to 86.9 with four.The benefit of added predictor capacity saturates at moderate depth.Source: §G.5 · p. 32
  • Predictor-depth effects differ by task family: locomotion generally improves through three blocks, while AntMaze and Maze2D show larger gains from deeper predictors.HalfCheetah reaches comparable or slightly better performance with two blocks, and four blocks add only marginal gains over three on Maze2D.Source: §G.5 · pp. 32–33

G.6 SENSITIVITY TO THE SEQUENCE LENGTH

  • Increasing sequence length from K = 5 to K = 20 consistently improves performance across almost all environments.The average normalized return increases as context grows from K = 5 to K = 10 and again to K = 20.Source: §G.6 · p. 33
  • K = 20 is best on almost all locomotion datasets, while most AntMaze tasks and larger Maze2D configurations slightly favor K = 30; K = 40 causes a performance drop.The benefit of longer context saturates, and further increases are not uniformly useful.Source: §G.6 · pp. 33–34

G.7 ADDITIONAL ABLATION RESULTS

Even with AdaLN and critic-guided adaptation, PDT remains behind full DT-JEPA on all evaluated unseen tasks. The remaining advantage points to the future-conditioning representation itself.

G.7 ADDITIONAL ABLATION RESULTS

  • PDT-AdaLN-CG remains behind full DT-JEPA, pointing to history-grounded prediction of future latent representations—not only modulation or critic supervision—as a distinctive ingredient.DT-JEPA conditions on future representations predicted from trajectory history and models their conditional distribution.Source: §G.7 · pp. 34–35

Related papers

Loading 2610.32657v1…