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

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

cs.LGOct 23, 2026

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

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

Organizations: Southeast University · Jiangsu Second Normal University · The University of Hong Kong · National University of Defense Technology

Abstract

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.

Figures & tables

Figure 1: (a) RTG summarizes future with a ambiguous target return, while DT-JEPA instead predicts latent future dynamics from historical context and effectively guide current decision making. (b) From an information-theoretic perspective, a more informative future condition reduces uncertainty about future trajectories and preserves richer future information.

Figure 2: Overview of the proposed DT-JEPA framework . DT-JEPA learns to make decisions using predicted future-dynamics information. During pretraining, a probabilistic JEPA predictor generates latent future representations to modulate DT through AdaLN. During online fine-tuning, specific critics provide value guidance to adapt the pretrained policy to unseen reward objectives.

Dataset	TD3-BC	TD-JEPA	ODT	QT	Prompt DT	VDT	LSDT	PDT	DT-JEPA
Hopper-m	$82.6\pm 3.8$	$85.7\pm 3.4$	$79.6\pm 3.5$	$95.3\pm 3.6$	$90.6\pm 2.9$	$94.1\pm 2.5$	$89.4\pm 1.9$	$\mathbf{96.2\pm 1.8}$	$\underline{95.4\pm 1.7}$
Hopper-m-r	$64.3\pm 5.6$	$68.9\pm 5.0$	$74.8\pm 5.1$	$\underline{97.3\pm 3.1}$	$86.1\pm 3.4$	$96.1\pm 3.5$	$93.5\pm 2.9$	$82.9\pm 3.5$	$\mathbf{98.3\pm 2.1}$
Hopper-m-e	$96.8\pm 2.7$	$99.1\pm 2.4$	$98.9\pm 2.5$	$102.0\pm 1.8$	$100.7\pm 2.1$	$100.4\pm 1.7$	$\underline{104.1\pm 1.7}$	$102.4\pm 1.8$	$\mathbf{105.2\pm 1.3}$
Walker2d-m	$78.9\pm 4.7$	$81.2\pm 4.2$	$74.6\pm 5.3$	$82.6\pm 3.4$	$79.8\pm 4.1$	$81.7\pm 3.6$	$\underline{83.4\pm 3.1}$	$79.3\pm 3.9$	$\mathbf{86.1\pm 2.8}$
Walker2d-m-r	$88.3\pm 3.7$	$73.6\pm 5.4$	$72.5\pm 4.9$	$\mathbf{99.0\pm 3.9}$	$72.9\pm 5.6$	$75.1\pm 4.8$	$80.4\pm 3.9$	$64.6\pm 4.3$	$\underline{94.4\pm 2.3}$
Walker2d-m-e	$104.8\pm 3.0$	$106.8\pm 2.6$	$106.2\pm 2.7$	$110.5\pm 2.0$	$107.7\pm 2.4$	$109.1\pm 2.1$	$\underline{111.8\pm 1.7}$	$108.6\pm 2.0$	$\mathbf{112.3\pm 1.5}$

Table 1: Normalized returns on D4RL locomotion and navigation tasks. We report the mean performance together with the standard deviation over 5 random seeds.

Figure 3: Fine-tuning performance and convergence comparison of DT-JEPA and competing methods on unseen Gym MuJoCO tasks. Performance is evaluated every 10 episodes. We report results averaged over three random seeds, with shaded regions indicating 95

\%

confidence intervals.

Figure 4: Generalization performance of DT-JEPA and competing methods on unseen downstream tasks. The results report the averaged normalized return achieved on 3 unseen downstream tasks, and the error bars represent 95

\%

confidence intervals.

Method	Condition	Hopper	Walker2d	HalfCheetah	Ant	Avg.
DT	RTG	68.4	64.9	72.3	60.8	66.6
DT-RTG-AdaLN	RTG	71.0	67.8	71.9	64.6	68.8
DT + Deterministic Future Predictor	Future	76.8	74.5	80.2	71.6	75.8
DT-JEPA w/o Critic Guidance	Future	82.1	78.6	84.8	76.0	80.4
DT-JEPA w/ Frozen Backbone	Future	83.5	80.9	84.4	76.7	81.4
DT-JEPA-Token	Future	84.7	80.1	87.9	79.3	83.0

Table 2: Ablation study of DT-JEPA on unseen downstream tasks. We evaluate the effects on four MuJoCo environments with unseen reward objectives during pretraining. Results are reported in terms of normalized return.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Illustration of four MuJoCo environments utilized in the paper.

Figure 6: Illustration of four AntMaze environments utilized in the paper.

\hat{Q}_{m}^{j,(i)}=\sum_{k=m}^{t-1}\gamma^{k-m}r_{k}^{j,(i)}+\gamma^{t-m}\min_{d\in\{1,2\}}Q_{\bar{\omega}_{d}^{j}}^{j}\left(s_{t}^{j,(i)},\hat{a}_{t}^{j,(i)}\right)

Appendix

Algorithm 2 Critic-Guided Online Fine-Tuning of DT-JEPA

Hyperparameter	Value
JEPA encoder architecture	lightweight Transformer
JEPA encoder layers	2
future latent dimension $d_{F}$	64
context length $K$	20
EMA coefficient	0.99
DT Transformer layers	4

Appendix

Table 3: Common hyperparameters used for DT-JEPA.

Figure 7: Action distributions produced by DT-JEPA under different latent conditions of an episode. Each subplot corresponds to one dimension of the continuous action space. The horizontal axis denotes the action value, while the vertical axis represents the normalized probability density of the histogram, computed as

\hat{p}_{i}=N_{i}/(N\Delta a)

, where

N_{i}

is the number of action samples falling into the

i

-th bin,

N

is the total number of sampled actions, and

\Delta a

is the bin width.

Figure 8: Attention score visualization on HalfCheetah-medium-replay. Compared with DT, RTG-AdaLN, and JEPA-Token, DT-JEPA exhibits a more concentrated attention pattern around the main diagonal. Since HalfCheetah is dominated by strong local temporal dependencies, this diagonal structure indicates that the model places greater emphasis on temporally adjacent and dynamically relevant trajectory information. The comparison further highlights the complementary effects of predictive future conditioning and AdaLN modulation.

Figure 9: Sensitivity analysis of the training-stage JEPA loss weight

\lambda_{J}

across different control and navigation environments. We vary

\lambda_{J}

over

\{0.001,0.01,0.1,1,10\}

and report the normalized return. Error bars indicate the standard deviation over multiple random seeds.

Figure 10: Sensitivity analysis of $\lambda_{J}$ and $\lambda_{Q}$ on four MuJoCo Tasks. The upper panel varies the weight of the JEPA predictive objective $\lambda_{J}$ , while the below panel varies the critic-guidance coefficient $\lambda_{Q}$ .

Dataset	QT	Prompt DT	VDT	LSDT	PDT	DT-JEPA
Hopper-m	$\underline{84.8}$	$78.3$	$77.8$	$77.9$	$77.6$	$\mathbf{86.6}$
Hopper-m-r	$\mathbf{76.8}$	$60.2$	$63.4$	$67.9$	$65.2$	$\underline{75.9}$
Hopper-m-e	$98.7$	$91.5$	$94.2$	$\underline{99.4}$	$96.8$	$\mathbf{101.2}$
Walker2d-m	$80.5$	$73.2$	$73.5$	$\underline{80.8}$	$73.4$	$\mathbf{83.2}$
Walker2d-m-r	$\underline{74.3}$	$64.9$	$68.0$	$72.6$	$60.7$	$\mathbf{76.8}$
Walker2d-m-e	$106.2$	$99.3$	$102.0$	$\mathbf{109.0}$	$103.7$	$\underline{108.6}$

Appendix

Table 4: Trajectory stitching ability on degraded D4RL locomotion and navigation datasets. For each dataset, the top

40\%

highest-return complete trajectories are removed before training, resulting in offline datasets that contain more fragmented and suboptimal behaviors. All methods are trained using the same degraded datasets without additional environment interaction. The best result is shown in bold , while the second-best result is underlined .

Dataset	$L_{P}=1$	$L_{P}=2$	$L_{P}=3$	$L_{P}=4$
Hopper-m	$89.6\pm 2.5$	$94.1\pm 1.9$	$\mathbf{95.4\pm 1.7}$	$\underline{94.6\pm 1.8}$
Hopper-m-r	$90.4\pm 3.4$	$96.2\pm 2.8$	$\mathbf{98.3\pm 2.1}$	$\underline{96.9\pm 2.6}$
Hopper-m-e	$100.6\pm 2.0$	$104.4\pm 1.5$	$\mathbf{105.2\pm 1.3}$	$\underline{104.7\pm 1.4}$
Walker2d-m	$79.6\pm 3.6$	$84.4\pm 3.1$	$\underline{86.1\pm 2.8}$	$\mathbf{86.5\pm 2.8}$
Walker2d-m-r	$86.3\pm 4.2$	$92.1\pm 3.6$	$\underline{93.1\pm 2.9}$	$\mathbf{94.4\pm 2.3}$
Walker2d-m-e	$107.1\pm 2.2$	$111.3\pm 1.8$	$\underline{112.3\pm 1.5}$	$\mathbf{112.3\pm 1.2}$

Appendix

Table 5: Sensitivity analysis of the JEPA predictor depth

L_{P}

on D4RL locomotion and navigation tasks.

L_{P}

denotes the number of Transformer blocks in the latent-space JEPA predictor. Results are reported as normalized returns over five random seeds. The best result is shown in bold , while the second-best result is underlined .

Dataset	$K=5$	$K=10$	$K=20$	$K=30$	$K=40$
Hopper-m	$89.8\pm 2.8$	$93.7\pm 2.1$	$\mathbf{95.4\pm 1.7}$	$\underline{94.9\pm 1.9}$	$94.0\pm 2.2$
Hopper-m-r	$89.2\pm 3.9$	$94.7\pm 3.0$	$\mathbf{98.3\pm 2.1}$	$\underline{97.5\pm 2.4}$	$96.5\pm 2.7$
Hopper-m-e	$101.4\pm 2.2$	$103.8\pm 1.8$	$\mathbf{105.2\pm 1.3}$	$\underline{104.8\pm 1.5}$	$104.3\pm 1.7$
Walker2d-m	$80.2\pm 4.0$	$84.6\pm 3.3$	$\mathbf{86.1\pm 2.8}$	$\underline{85.8\pm 3.0}$	$85.2\pm 3.3$
Walker2d-m-r	$86.8\pm 3.9$	$91.8\pm 3.1$	$\mathbf{94.4\pm 2.3}$	$\underline{93.6\pm 2.6}$	$92.7\pm 3.0$
Walker2d-m-e	$108.2\pm 2.3$	$110.9\pm 1.9$	$\underline{112.3\pm 1.5}$	$\mathbf{112.5\pm 1.6}$	$111.6\pm 1.9$

Appendix

Table 6: Sensitivity analysis of the sequence length

K

on D4RL locomotion and navigation tasks. Results are reported as normalized returns over five random seeds. The best result is shown in bold , while the second-best result is underlined .

Method	Hopper	Walker2d	HalfCheetah	Ant	Avg.
DT-RTG-AdaLN	71.0	67.8	71.9	64.6	68.8
PDT	77.2	73.7	80.5	72.4	76.0
PDT-AdaLN	78.4	75.0	81.3	73.6	77.1
PDT-CG	80.2	77.2	82.8	75.2	78.9
PDT-AdaLN-CG	81.5	78.1	83.9	75.5	79.5
DT-JEPA	89.4	85.6	91.2	83.1	87.3

Appendix

Table 7: Additional ablation experiment. PDT-AdaLN applies AdaLN-based conditioning to PDT, PDT-CG augments PDT with the same critic-guided adaptation used in DT-JEPA, and PDT-AdaLN-CG combines both modifications. All methods are evaluated on the same unseen reward objectives across Hopper, Walker2d, HalfCheetah, and Ant as in Table 2 .

Transformer-Based RL

Transformer-based reinforcement learning uses attention-based sequence models to learn, plan, and control from experience across tasks.

See all 53 papers

Sep 21, 2026

D-JEPA: A Decision-Aligned Latent World Model

Latent world models predict the consequences of actions, but accurate prediction does not guarantee that latent distance reflects which candidate will execute successfully. We identify a decision-local prediction gap: among the few futures competing for execution, a candidate predicted closer to the goal can produce a worse realized outcome than an available alternative. We introduce D-JEPA, a decision-aligned latent world model that learns decision-relevant relations among candidate futures from executed outcomes. A bounded, permutation-equivariant operator jointly reasons over goal-relative predictive features and ordinal evidence, refining pretrained predictive geometry where action choices are most consequential. Restricted predictor adaptation and a shared ordinal interface extend this alignment across complementary predictive geometries. D-JEPA further realizes the learned decision structure in JEPA-compatible future representations, enabling deployment through native latent-distance planning. Evaluations across latent control, manipulation, pretrained action-producing models, physical robots and autonomous driving demonstrate improved action selection, including 87.89% success on PushT, a 15.04-point average gain on RoboTwin, and a 17-point gain on physical robot tasks. These results establish decision-relevant relational structure as a direct bridge between predictive world modeling and effective control.

Shuaijun Liu, Chengyu Wu, Qifu Wen +5

The Hong Kong University of Science and Technology (Guangzhou) · Boston University · Shanghai Jiao Tong University

World Models World Model-Based Planning Joint-Embedding Predictive Architecture Model-Based Planning Latent World Models

Ask PDF

Jan 24, 2025

ACT-JEPA: Novel Joint-Embedding Predictive Architecture for Efficient Policy Representation Learning

Learning efficient representations for decision-making policies is a challenge in imitation learning (IL). Current IL methods require expert demonstrations, which are expensive to collect. Additionally, they are not explicitly trained to understand the environment. Consequently, they have underdeveloped world models. Self-supervised learning (SSL) offers an alternative, as it can learn a world model from diverse, unlabeled data. However, most SSL methods are inefficient because they operate in raw input space. In this work, we propose ACT-JEPA, a novel architecture that unifies IL and SSL to enhance policy representations. It is trained end-to-end to jointly predict 1) action sequences and 2) latent observation sequences. To learn in latent space, we utilize Joint-Embedding Predictive Architecture, which allows the model to filter out irrelevant details and learn a robust world model. We evaluate ACT-JEPA in different environments and across multiple tasks. Our results show that it outperforms the strongest baseline in all environments. ACT-JEPA achieves up to 40% improvement in world model understanding and up to 10% higher task success rate. Finally, we show that predicting latent observation sequences effectively generalizes to predicting action sequences. This work demonstrates how integrating IL and SSL leads to efficient policy representation learning, an improved world model, and a higher task success rate.

Aleksandar Vujinovic, Aleksandar Kovacevic

Faculty of Technical Sciences, University of Novi Sad

Representation Learning World Model Learning Self-Supervised Learning Joint-Embedding Predictive Architecture Imitation Learning Policy Learning

Ask PDF

Sep 30, 2026

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

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.

Zheyuan Zhang, Suyu Ye, Nakul Agarwal +5

Honda Research Institute USA · Johns Hopkins University

Test-Time Adaptation Latent Dynamics Modeling World Model-Based Planning Joint-Embedding Predictive Architecture Action-Conditioned World Models Model-Based Planning

Ask PDF

Abstract

Transformer-Based RL

Explore similar work

D-JEPA: A Decision-Aligned Latent World Model

ACT-JEPA: Novel Joint-Embedding Predictive Architecture for Efficient Policy Representation Learning

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