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

D-JEPA: A Decision-Aligned Latent World Model
Sep 21, 2026Shuaijun Liu, Chengyu Wu, Qifu Wen +5World Models World Model-Based Planning
The Hong Kong University of Science and Technology (Guangzhou) · Boston University · Shanghai Jiao Tong University
ACT-JEPA: Novel Joint-Embedding Predictive Architecture for Efficient Policy Representation Learning
Jan 24, 2025Aleksandar Vujinovic, Aleksandar KovacevicRepresentation Learning World Model Learning
Faculty of Technical Sciences, University of Novi Sad
JEPA-TTT: Persistent Test-Time Training of Latent World Models for Planning under Dynamics Shifts
Sep 30, 2026Zheyuan Zhang, Suyu Ye, Nakul Agarwal +5Test-Time Adaptation Latent Dynamics Modeling
Honda Research Institute USA · Johns Hopkins University

Sep 21, 2026J/K move · Enter open · S save

D-JEPA: A Decision-Aligned Latent World Model

Shuaijun Liu, Chengyu Wu, Qifu Wen, Feiyang You+4

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

Open paper AskPDF arXiv

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.

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