Introduction
Complex decision-making domains such as robotic control, autonomous driving, and embodied intelligence often involve large-scale offline datasets collected across diverse tasks and policies. Learning general decision knowledge from such heterogeneous data and efficiently adapting it to new downstream tasks is becoming increasingly important [22 ] . Motivated by the success of pretraining–finetuning paradigms in computer vision [7 , 16 ] and natural language processing [8 , 5 ] , reinforcement learning (RL) research has recently focused on leveraging offline trajectories to learn reusable policies and representations that can generalize across tasks. Offline RL provides a natural foundation for this direction, as it enables policy learning entirely from previously collected experience without requiring additional interaction with the environment. More recently, inspired by the strong sequence modeling capability of Transformers Vaswani et al., [28 ] , subsequent studies have reformulated offline policy learning as a conditional sequence modeling problem Chen et al., [6 ] , offering a new perspective for offline RL.
Within the CSM formulation, policy inference is conditioned on a specific target together with the historical context. Following this paradigm, the representative Decision Transformer (DT) conditions action generation on desired future returns and has demonstrated strong performance across a wide range of offline RL tasks. Despite its effectiveness, return-oriented CSM methods still suffer from several limitations. On the one hand, expected future returns provide an ambiguous description of future behaviors, since distinct future trajectories with different state transitions and action patterns may share similar cumulative returns. On the other hand, return conditioning inherently relies on task-specific reward functions, restricting policy pretraining to reward-annotated trajectories and thereby limiting scalability and transferability. To alleviate these limitations, Generalized DT (GDT) extends return conditioning to generalized hindsight statistics extracted from future trajectories Furuta et al., [11 ] , while Pretrained DT (PDT) encodes future trajectory segments into embeddings for future-conditioned action generation Xie et al., [31 ] . However, these approaches still rely on realized future trajectory information during pretraining and require appropriate future conditions to be specified or selected at deployment. This raises an important question: How can we explicitly learn the distribution of future dynamics in the latent space solely from observable history, while ensuring the predicted representations preserve sufficient information about the future to guide DT’s action generation?
Recent advances in offline RL have begun to explore policy learning under reward-free and unsupervised pretraining paradigms [26 , 23 , 2 ] . Since the transition dynamics p ( s ′ ∣ s , a ) p(s' \mid s, a) p ( s ′ ∣ s , a ) [15 , 14 ] . However, conventional generative world-modeling approaches often rely on observation reconstruction or generation objectives, which may require modeling fine-grained details that are not directly relevant to downstream decision making. Recently, Joint-Embedding Predictive Architectures (JEPA) have emerged as a promising alternative for predictive representation learning [21 , 1 , 3 ] . Rather than reconstructing future observations in the original input space, JEPA predicts their latent representations from observed context, encouraging the model to capture abstract information that is predictive of future dynamics.
Inspired by this paradigm, we propose DT-JEPA, a generalizable offline pretraining and online fine-tuning framework that learns transferable decision knowledge from reward-free offline trajectories and efficiently adapts to unseen downstream tasks. First, to capture the predictive information and uncertainty about future dynamics, we design a novel probabilistic predictor to model the conditional distribution of future representations in the latent space. Second, we incorporate the predicted future representation into DT through adaptive layer normalization (AdaLN) Peebles & Xie, [24 ] , dynamically modulating historical representations based on the future-predictive information to guide DT’s action generation. Finally, to adapt the pretrained policy to a new task with an unseen reward function, we develop a critic-guided online fine-tuning mechanism. The DT-JEPA policy and critic are jointly optimized on newly collected trajectories, where the critic estimates action values under the target reward and provides Q-value guidance for adapting the policy toward the new task objective.
Theoretically, we show that 1) JEPA objective function maximize a variational lower bound on 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 ) \mathcal{O}(N^{-1/4}) O ( N − 1/4 )
Preliminaries
Offline Reinforcement Learning and Decision Transformer
We consider a Markov decision process (MDP) M = ( S , A , P , r , γ , ρ 0 ) \mathcal{M} = (\mathcal{S}, \mathcal{A}, P, r, \gamma, \rho_0) M = ( S , A , P , r , γ , ρ 0 ) S \mathcal{S} S A \mathcal{A} A P ( s ′ ∣ s , a ) P(s' \mid s, a) P ( s ′ ∣ s , a ) r : S × A → R r : \mathcal{S} \times\mathcal{A} \to\mathbb{R} r : S × A → R γ ∈ [ 0 , 1 ) \gamma\in[0, 1) γ ∈ [ 0 , 1 ) ρ 0 \rho_0 ρ 0 [19 ] . A policy π ( a ∣ s ) \pi(a \mid s) π ( a ∣ s ) J ( π ) = E π [ ∑ t γ t r ( s t , a t ) ] J(\pi) = \mathbb{E}_{\pi} \left[\sum_{t} \gamma^t r(s_t, a_t)\right] J ( π ) = E π [ ∑ t γ t r ( s t , a t ) ] Q π ( s t , a t ) = E π [ ∑ k γ k r ( s t + k , a t + k ) ∣ s t , a t ] Q^\pi(s_t, a_t) = \mathbb{E}_{\pi} \left[\sum_{k} \gamma^k r(s_{t+k}, a_{t+k}) \mid s_t, a_t\right] Q π ( s t , a t ) = E π [ ∑ k γ k r ( s t + k , a t + k ) ∣ s t , a t ]
In offline reinforcement learning, the agent cannot interact with the environment during training. Instead, learning is performed exclusively from a fixed dataset D \mathcal{D} D [22 ] . DT reformulates offline RL as a conditional sequence modeling problem. For each trajectory, DT augments the state–action sequence with the return-to-go (RTG) R t = ∑ t ′ = t T r t ′ R_t = \sum_{t'=t}^{T} r_{t'} R t = ∑ t ′ = t T r t ′ τ ^ = ( R 1 , s 1 , a 1 , R 2 , s 2 , a 2 , … , R T , s T , a T ) \hat{\tau} = (R_1, s_1, a_1, R_2, s_2, a_2, \ldots, R_T, s_T, a_T) τ ^ = ( R 1 , s 1 , a 1 , R 2 , s 2 , a 2 , … , R T , s T , a T ) π θ ( a t ∣ τ ^ 1 : t − 1 , s t , R t ) \pi_\theta(a_t \mid\hat{\tau}_{1:t-1}, s_t, R_t) π θ ( a t ∣ τ ^ 1 : t − 1 , s t , R t )
L D T = E τ ^ ∼ D ^ [ ∑ t = 1 T − log π θ ( a t ∣ τ ^ 1 : t − 1 , s t , R t ) ] (1) \mathcal{L}_{\mathrm{DT}} = \mathbb{E}_{\hat{\tau} \sim\hat{\mathcal{D}}} \left[\sum_{t=1}^{T} -\log\pi_\theta(a_t \mid\hat{\tau}_{1:t-1}, s_t, R_t)\right] \tag*{(1)} L DT = E τ ^ ∼ D ^ [ t = 1 ∑ T − log π θ ( a t ∣ τ ^ 1 : t − 1 , s t , R t ) ] (1) where D ^ \hat{\mathcal{D}} D ^
Joint-Embedding Predictive Architecture
The Joint Embedding Prediction Architecture (JEPA) is a self-supervised representation learning framework based on latent-space prediction [21 , 1 ] . Its core idea is not to directly reconstruct the original observed data, but rather to predict the latent-space embeddings of missing target data. Given an input sample, JEPA constructs a context view x C x_C x C x T x_T x T f θ f_\theta f θ f θ ′ f_{\theta'} f θ ′ z C = f θ ( x C ) z_C = f_\theta(x_C) z C = f θ ( x C ) z T = f θ ′ ( x T ) z_T = f_{\theta'}(x_T) z T = f θ ′ ( x T ) f θ ′ f_{\theta'} f θ ′ θ ′ ← α θ ′ + ( 1 − α ) θ \theta' \leftarrow\alpha\theta' + (1-\alpha)\theta θ ′ ← α θ ′ + ( 1 − α ) θ α ∈ [ 0 , 1 ) \alpha\in[0, 1) α ∈ [ 0 , 1 ) [12 ] . Building on this, JEPA uses a predictor g ϕ g_\phi g ϕ z C z_C z C ξ T \xi_T ξ T z ^ T = g ϕ ( z C , ξ T ) \hat{z}_T = g_\phi(z_C, \xi_T) z ^ T = g ϕ ( z C , ξ T )
L J E P A = ∥ g ϕ ( f θ ( x C ) , ξ ^ T ) − f θ ′ ( x T ) ∥ 2 2 . (2) \mathcal{L}_{\mathrm{JEPA}} = \left\|g_\phi\left(f_\theta(x_C), \hat{\xi}_T\right) - f_{\theta'}(x_T)\right\|_2^2. \tag*{(2)} L JEPA = g ϕ ( f θ ( x C ) , ξ ^ T ) − f θ ′ ( x T ) 2 2 . (2) Thus, through the above method, JEPA ensures that the latent representation z C z_C z C x T x_T x T
Adaptive Layer Normalization
AdaLN is an adaptive normalization mechanism that injects conditional information into the feature representations of neural networks [24 ] . AdaLN dynamically generates corresponding affine transformation parameters based on an external conditional variable z z z γ ( z ) = MLP γ ( z ) \gamma(z) = \operatorname{MLP}_{\gamma}(z) γ ( z ) = MLP γ ( z ) β ( z ) = MLP β ( z ) \beta(z) = \operatorname{MLP}_{\beta}(z) β ( z ) = MLP β ( z ) x x x
AdaLN ( x , z ) = γ ( z ) ⊙ x − μ ( x ) σ 2 ( x ) + ϵ + β ( z ) , (3) \operatorname{AdaLN}(x, z) = \gamma(z) \odot\frac{x - \mu(x)}{\sqrt{\sigma^2(x) + \epsilon}} + \beta(z), \tag*{(3)} AdaLN ( x , z ) = γ ( z ) ⊙ σ 2 ( x ) + ϵ x − μ ( x ) + β ( z ) , (3) where μ ( x ) \mu(x) μ ( x ) σ 2 ( x ) \sigma^2(x) σ 2 ( x ) ϵ \epsilon ϵ ⊙ \odot ⊙ z z z
Method
In this section, we present DT-JEPA, an offline pretraining and online finetuning-based predictive sequence-modeling framework for generalizable decision making. During the pretraining phase, a variational JEPA model is designed to infer future representations from multi-task offline data without reward supervision. Subsequently, the predicted future latent guides DT’s action generation by effectively modulating transformer’s hidden representations. During the finetuning phase, task-specific critics is designed to efficiently adapt the pretrained policy toward unseen reward objectives with limited environment interaction. Theoretically, we analyze the effectiveness of the predictive representation learning objective and derive a finite-sample performance bound for the future-conditioned policy learning.
Motivation
The DT uses RTG, i.e. R t = ∑ i = t T r t R_t = \sum_{i=t}^{T} r_t R t = ∑ i = t T r t T ≥ t \mathcal{T}_{\geq t} T ≥ t T ≥ t A \mathcal{T}_{\geq t}^{A} T ≥ t A T ≥ t B \mathcal{T}_{\geq t}^{B} T ≥ t B R t A R_t^{A} R t A R t B R_t^{B} R t B C t C_t C t I ( C t ; T ≥ t ) I(C_t; \mathcal{T}_{\geq t}) I ( C t ; T ≥ t )
I ( C t ; T ≥ t ) = H ( T ≥ t ) − H ( T ≥ t ∣ C t ) I(C_t; \mathcal{T}_{\geq t}) = H(\mathcal{T}_{\geq t}) - H(\mathcal{T}_{\geq t} \mid C_t) I ( C t ; T ≥ t ) = H ( T ≥ t ) − H ( T ≥ t ∣ C t ) Here, H ( T ≥ t ) H(\mathcal{T}_{\geq t}) H ( T ≥ t ) T ≥ t \mathcal{T}_{\geq t} T ≥ t H ( T ≥ t ∣ C t ) H(\mathcal{T}_{\geq t} \mid C_t) H ( T ≥ t ∣ C t ) C t C_t C t Figure 1 , the smaller H ( T ≥ t ∣ C t ) H(\mathcal{T}_{\geq t} \mid C_t) H ( T ≥ t ∣ C t ) C t C_t C t I ( C t ; T ≥ t ) I(C_t; \mathcal{T}_{\geq t}) I ( C t ; T ≥ t ) H ( T ≥ t ∣ R t ) H(\mathcal{T}_{\geq t} \mid R_t) H ( T ≥ t ∣ R t ) I ( R t ; T ≥ t ) I(R_t; \mathcal{T}_{\geq t}) I ( R t ; T ≥ t ) F ^ t \hat{F}_t F ^ t
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.
Variational Joint-Embedding Future Conditioning
To predict a more effective future representation as the conditional variable for the DT’s action generation, we constructing a JEPA-based future representation learning mechanism. Specifically, the observable history of DT at each time step is defined as τ t − K + 1 : t = ( s t − K + 1 : t , a t − K + 1 : t − 1 ) \tau_{t-K+1:t} = (s_{t-K+1:t}, a_{t-K+1:t-1}) τ t − K + 1 : t = ( s t − K + 1 : t , a t − K + 1 : t − 1 ) K K K τ t − K + 1 : t \tau_{t-K+1:t} τ t − K + 1 : t f θ f_{\theta} f θ τ t − K + 1 : t \tau_{t-K+1:t} τ t − K + 1 : t H t \mathbf{H}_{t} H t
H t = f θ ( τ t − K + 1 : t ) . (4) \mathbf{H}_{t} = f_{\theta}\left(\tau_{t-K+1:t}\right). \tag*{(4)} H t = f θ ( τ t − K + 1 : t ) . (4) Moreover, based on the future trajectory τ t : t + K − 1 \tau_{t:t+K-1} τ t : t + K − 1 D \mathcal{D} D f θ ′ f_{\theta'} f θ ′ F t \mathbf{F}_{t} F t
F t = f θ ′ ( τ t : t + K − 1 ) . (5) \mathbf{F}_{t} = f_{\theta'}\left(\tau_{t:t+K-1}\right). \tag*{(5)} F t = f θ ′ ( τ t : t + K − 1 ) . (5) where θ ′ \theta' θ ′ θ \theta θ H t \mathbf{H}_{t} H t q η ( ⋅ ∣ H t ) = N ( μ ^ t P , diag ( ( σ ^ t P ) 2 ) ) q_{\eta}\left(\cdot\mid\mathbf{H}_{t}\right) = \mathcal{N}\left(\hat{\boldsymbol{\mu}}_{t}^{P}, \operatorname{diag}\left((\hat{\boldsymbol{\sigma}}_{t}^{P})^{2}\right)\right) q η ( ⋅ ∣ H t ) = N ( μ ^ t P , diag ( ( σ ^ t P ) 2 ) ) μ ^ t P \hat{\boldsymbol{\mu}}_{t}^{P} μ ^ t P σ ^ t P \hat{\boldsymbol{\sigma}}_{t}^{P} σ ^ t P H t \mathbf{H}_{t} H t F ^ t ∼ q η ( ⋅ ∣ H t ) \hat{\mathbf{F}}_{t} \sim q_{\eta}\left(\cdot\mid\mathbf{H}_{t}\right) F ^ t ∼ q η ( ⋅ ∣ H t )
To encourage the predicted future representation F ^ t \hat{\mathbf{F}}_{t} F ^ t H t \mathbf{H}_{t} H t τ ≥ t \tau_{\geq t} τ ≥ t I ( τ ≥ t ; H t ) I(\tau_{\geq t}; \mathbf{H}_{t}) I ( τ ≥ t ; H t )
Theorem 1. The vaiational lower bound of mutual information I ( τ ≥ t ; H t ) I(\tau_{\geq t}; \mathbf{H}_{t}) I ( τ ≥ t ; H t )
I ( τ ≥ t ; H t ) ≥ E p ( F t , H t ) [ log q η ( F t ∣ H t ) ] + H ( F t ) , (6) I(\tau_{\geq t}; \mathbf{H}_{t}) \geq\mathbb{E}_{p(\mathbf{F}_{t},\mathbf{H}_{t})}\left[\log q_{\eta}(\mathbf{F}_{t} \mid\mathbf{H}_{t})\right] + H(\mathbf{F}_{t}), \tag*{(6)} I ( τ ≥ t ; H t ) ≥ E p ( F t , H t ) [ log q η ( F t ∣ H t ) ] + H ( F t ) , (6) where p ( F t , H t ) p(\mathbf{F}_{t}, \mathbf{H}_{t}) p ( F t , H t ) H ( F t ) H(\mathbf{F}_{t}) H ( F t ) F t \mathbf{F}_{t} F t
The proof of the Theorem 1 are provided in Appendix § 8 .§ 1 . Since H ( F t ) H(\mathbf{F}_{t}) H ( F t )
L J E P A = E ( τ t − K + 1 : t , τ t : t + K − 1 ) ∼ D [ − log q η ( F t ∣ H t ) ] (7) \mathcal{L}_{\mathrm{JEPA}} = \mathbb{E}_{(\tau_{t-K+1:t},\tau_{t:t+K-1})\sim\mathcal{D}}\left[-\log q_{\eta}(\mathbf{F}_{t} \mid\mathbf{H}_{t})\right] \tag*{(7)} L JEPA = E ( τ t − K + 1 : t , τ t : t + K − 1 ) ∼ D [ − log q η ( F t ∣ H t ) ] (7) where D \mathcal{D} D θ ′ \theta' θ ′
AdaLN-based Latent Conditioned Action Generation
To effectively utilize the predicted future representation F ^ t \hat{\mathbf{F}}_{t} F ^ t τ t − K + 1 : t \tau_{t-K+1:t} τ t − K + 1 : t h t \mathbf{h}_{t} h t F ^ t \hat{\mathbf{F}}_{t} F ^ t h t \mathbf{h}_{t} h t F ^ t \hat{\mathbf{F}}_{t} F ^ t
γ t = MLP γ ( F ^ t ) , β t = MLP β ( F ^ t ) . (8) \gamma_{t} = \operatorname{MLP}_{\gamma}\left(\hat{\mathbf{F}}_{t}\right), \quad\beta_{t} = \operatorname{MLP}_{\beta}\left(\hat{\mathbf{F}}_{t}\right). \tag*{(8)} γ t = MLP γ ( F ^ t ) , β t = MLP β ( F ^ t ) . (8) The historical representation is normalized and subsequently modulated as
h ~ t = ( 1 + γ t ) ⊙ LayerNorm ( h t ) + β t , (9) \tilde{\mathbf{h}}_{t} = (1+\gamma_{t}) \odot\operatorname{LayerNorm}(\mathbf{h}_{t}) + \beta_{t}, \tag*{(9)} h ~ t = ( 1 + γ t ) ⊙ LayerNorm ( h t ) + β t , (9) where ⊙ \odot ⊙ h ~ t \tilde{\mathbf{h}}_{t} h ~ t a ^ t \hat{\mathbf{a}}_{t} a ^ t
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.
Reward-Free Pretraining and Reward-Oriented Finetuning
During the pretraining stage, we first train the policy using a standard behavior cloning (BC) objective on large-scale trajectory data of multiple tasks [6 ] . Specifically, the BC loss is defined as
L B C = E τ t − K + 1 : t ∼ D , F ^ t ∼ q η [ ∑ m = t − K + 1 t ( a m − a ^ m ) 2 ] (10) \mathcal{L}_{\mathrm{BC}} = \mathbb{E}_{\tau_{t-K+1:t}\sim\mathcal{D},\,\hat{F}_t\sim q_\eta}\left[\sum_{m=t-K+1}^{t}(a_m-\hat{a}_m)^2\right] \tag*{(10)} L BC = E τ t − K + 1 : t ∼ D , F ^ t ∼ q η [ m = t − K + 1 ∑ t ( a m − a ^ m ) 2 ] (10) where a ^ m = ϕ ( τ t − K + 1 : t , F ^ t ) \hat{a}_m=\phi(\tau_{t-K+1:t},\hat{F}_t) a ^ m = ϕ ( τ t − K + 1 : t , F ^ t ) ϕ \phi ϕ
L ( ϕ , θ , η ) = L B C + λ J L J E P A , (11) \mathcal{L}(\phi,\theta,\eta)=\mathcal{L}_{\mathrm{BC}}+\lambda_J\mathcal{L}_{\mathrm{JEPA}}, \tag*{(11)} L ( ϕ , θ , η ) = L BC + λ J L JEPA , (11) where λ J \lambda_J λ J
To characterize whether the predictive-conditioned policy learned from a finite offline dataset can approximate its population counterpart, we further provide a finite-sample performance analysis.
Theorem 2. Consider a finite-horizon MDP with horizon H H H [ 0 , 1 ] [0,1] [ 0 , 1 ] q η ( F ∣ x ) q_\eta(F\mid x) q η ( F ∣ x ) π ˉ η P \bar{\pi}_\eta^{\mathrm{P}} π ˉ η P π ˉ η , ϕ ^ \bar{\pi}_{\eta,\hat{\phi}} π ˉ η , ϕ ^ C F C_F C F Π F \Pi_F Π F c F c_F c F ϵ a p p \epsilon_{\mathrm{app}} ϵ app 1 − δ 1-\delta 1 − δ
J ( π ˉ η P ) − J ( π ˉ η , ϕ ^ ) ≤ O ( C F H 2 [ c F ( log ( ∣ Π F ∣ / δ ) N ) 1 / 4 + ϵ a p p ] ) . (12) J(\bar{\pi}_\eta^{\mathrm{P}})-J(\bar{\pi}_{\eta,\hat{\phi}})\leq O\left(C_FH^2\left[\sqrt{c_F}\left(\frac{\log(|\Pi_F|/\delta)}{N}\right)^{1/4}+\sqrt{\epsilon_{\mathrm{app}}}\right]\right). \tag*{(12)} J ( π ˉ η P ) − J ( π ˉ η , ϕ ^ ) ≤ O ( C F H 2 [ c F ( N log ( ∣ Π F ∣/ δ ) ) 1/4 + ϵ app ] ) . (12) Proof: See the Appendix B.2. Theorem 2 shows that, under a given future predictor, the learned predictive-conditioned policy approaches its population counterpart as the amount of offline data increases, up to the approximation error. This result supports learning a transferable predictive-conditioned policy from large-scale reward-free trajectories. Building on the pretrained policy, we further adapt the policy to a new task j j j [35 ] .
Specifically, the pretrained policy first interacts with the environment of task j j j τ j = ( s 1 j , a 1 j , … , s T j , a T j ) \tau^j = (s_1^j, a_1^j, \ldots, s_T^j, a_T^j) τ j = ( s 1 j , a 1 j , … , s T j , a T j ) D j \mathcal{D}^j D j D j \mathcal{D}^j D j D ^ j \widehat{\mathcal{D}}^j D j Q ω 1 j j ( s t j , a t j ) Q_{\omega_1^j}^j(s_t^j, a_t^j) Q ω 1 j j ( s t j , a t j ) Q ω 2 j j ( s t j , a t j ) Q_{\omega_2^j}^j(s_t^j, a_t^j) Q ω 2 j j ( s t j , a t j ) a t j a_t^j a t j s t j s_t^j s t j j j j [10 ] . For the two critic networks and the DT policy network, we respectively maintain EMA target networks parameterized by ω ˉ 1 j \bar{\omega}_1^j ω ˉ 1 j ω ˉ 2 j \bar{\omega}_2^j ω ˉ 2 j ϕ ˉ \bar{\phi} ϕ ˉ
Q ^ m j = ∑ k = m t − 1 γ k − m r k j + γ t − m min d ∈ { 1 , 2 } Q ω ˉ d j j ( s ^ t j , a ^ t j ) , (13) \widehat{Q}_m^j = \sum_{k=m}^{t-1} \gamma^{k-m}r_k^j + \gamma^{t-m}\min_{d\in\{1,2\}} Q_{\bar{\omega}_d^j}^j(\hat{s}_t^j,\hat{a}_t^j), \tag*{(13)} Q m j = k = m ∑ t − 1 γ k − m r k j + γ t − m d ∈ { 1 , 2 } min Q ω ˉ d j j ( s ^ t j , a ^ t j ) , (13) where γ \gamma γ r k j r_k^j r k j k k k j j j ω d j \omega_d^j ω d j
L ( ω d j ) = E τ t − K + 1 : t j ∼ D ^ j , F t j ∼ f θ , a ^ t j ∼ ϕ ˉ [ ∑ m = t − K + 1 t − 1 ( Q ω d j j ( s m j , a m j ) − Q ^ m j ) 2 ] . (14) \mathcal{L}(\omega_d^j) = \mathbb{E}_{\tau_{t-K+1:t}^j\sim\widehat{\mathcal{D}}^j,\mathcal{F}_t^j\sim f_\theta,\hat{a}_t^j\sim\bar{\phi}}\left[\sum_{m=t-K+1}^{t-1}\left(Q_{\omega_d^j}^j(s_m^j,a_m^j)-\widehat{Q}_m^j\right)^2\right]. \tag*{(14)} L ( ω d j ) = E τ t − K + 1 : t j ∼ D j , F t j ∼ f θ , a ^ t j ∼ ϕ ˉ [ m = t − K + 1 ∑ t − 1 ( Q ω d j j ( s m j , a m j ) − Q m j ) 2 ] . (14) Consequently, based on the actions a ^ t j \hat{a}_t^j a ^ t j j j j
L Q j = − E τ t − K + 1 : t j ∼ D ^ j , F t j ∼ f θ [ ∑ m = t − K + 1 t − 1 Q ω 1 j j ( s m j , a ^ m j ) ] . (15) \mathcal{L}_{Q}^j = -\mathbb{E}_{\tau_{t-K+1:t}^j\sim\widehat{\mathcal{D}}^j,\mathcal{F}_t^j\sim f_\theta}\left[\sum_{m=t-K+1}^{t-1}Q_{\omega_1^j}^j(s_m^j,\hat{a}_m^j)\right]. \tag*{(15)} L Q j = − E τ t − K + 1 : t j ∼ D j , F t j ∼ f θ [ m = t − K + 1 ∑ t − 1 Q ω 1 j j ( s m j , a ^ m j ) ] . (15) Accordingly, the loss function for fine-tuning the DT on task j j j
L j ( ϕ , θ , η ) = L B C j + λ J j L J E P A j + λ Q L Q j , (16) \mathcal{L}^j(\phi,\theta,\eta)=\mathcal{L}_{\mathrm{BC}}^j+\lambda_J^j\mathcal{L}_{\mathrm{JEPA}}^j+\lambda_Q\mathcal{L}_Q^j, \tag*{(16)} L j ( ϕ , θ , η ) = L BC j + λ J j L JEPA j + λ Q L Q j , (16) where λ J j \lambda_J^j λ J j λ Q \lambda_Q λ Q
Experiments
Our experiments systematically evaluate the proposed DT-JEPA framework centered on the following three questions:
RQ1: Can DT-JEPA learn transferable knowledge from reward-free datasets to support downstream policy adaptation?
RQ2: Can DT-JEPA effectively transfer the knowledge acquired during the pretraining phase to unseen new tasks?
RQ3: How do the latent future representation F ^ t \widehat{F}_t F t
Baselines
We compare DT-JEPA with a variety of offline RL and DT-based methods. TD3-BC [9 ] , a widely used offline RL method that combines value-based policy optimization with behavior cloning regularization; TD-JEPA [2 ] , which learns long-term latent-predictive representations from reward-free offline transitions for zero-shot downstream RL; QT [17 ] , which augments Transformer-based conditional sequence modeling with Q-value regularization; ODT [35 ] , an online DT that adapts the pretrained sequence policy through limited environment interaction; Prompt DT [32 ] , which conditions DT on trajectory demonstrations to support few-shot policy generalization across tasks; VDT [34 ] , which incorporates value guidance into DT to improve policy optimization; LSDT [29 ] , which jointly models long-term and local trajectory dependencies to enhance sequence-based decision making; and PDT [31 ] , a method that conditions policy learning on true future trajectory information and subsequently adapts the pretrained policy to downstream tasks.
Benchmark Datasets and Evaluation Protocol
We evaluate DT-JEPA on diverse tasks from the D4RL benchmark. For MuJoCo locomotion tasks, the experiments include four sequential control environments: Hopper , Walker2d , HalfCheetah , and Ant . For each environment, we consider three offline datasets: medium , medium-replay , and medium-expert , which differ in trajectory quality and diversity. Moreover, we evaluate on long-horizon navigation benchmarks AntMaze and Maze2D . For AntMaze , we consider four dataset configurations, including umaze , umaze-diverse , medium-play , and medium-diverse . For Maze2D , we evaluate on the umaze , medium , and large environments with progressively challenging navigation. For compact presentation, we use the following abbreviations: m/r/m-e for medium/medium-replay/medium-expert , u/u-d/m-p/m-d for the four AntMaze variants, and u/m/l for umaze/medium/large in Maze2D. Moreover, we use the D4RL normalized return as the primary evaluation metric. For each environment and dataset, we perform independent training with multiple random seeds and report the mean normalized return together with its standard deviation.
Experimental Results
RQ1: As shown in Table 1 , DT-JEPA achieves the best performance on 10 out of 12 Gym MuJoCo datasets and remains competitive across the navigation tasks. In particular, DT-JEPA ranks first on three of the four medium-replay datasets, indicating that predictive representations are effective at extracting useful decision information from heterogeneous offline behaviors. More importantly, DT-JEPA outperforms PDT on 18 out of 19 datasets, suggesting that JEPA-based predictive representation learning provides a stronger decision prior for policy learning. Meanwhile, DT-JEPA consistently surpasses TD-JEPA, showing the benefit of integrating latent-predictive representation learning with sequence-based policy optimization.
Table 1. Normalized returns on D4RL locomotion and navigation tasks. We report the mean performance together with the standard deviation over 5 random seeds.
RQ2: As shown in Figure 3 , we first evaluate the convergence of DT-JEPA on one unseen reward function in each of the four MuJoCo environments. DT-JEPA reaches high-return regions with fewer interaction episodes and consistently converges faster than Prompt DT, ODT, and PDT. This advantage indicates that critic-guided adaptation can effectively steer the pretrained policy toward new reward objectives, thereby improving downstream adaptation efficiency. Figure 4 further plots the averaged normalized return achieved on 3 unseen reward objectives of the four MuJoCo tasks. DT-JEPA consistently outperforms the competing methods across the four MuJoCo tasks, indicating that the predictive future representations learned during reward-free pretraining capture reusable dynamics information for the adaptation of diverse downstream tasks.
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.
RQ3: To examine the contribution of each component to downstream generalization, we conduct ablation studies on unseen reward objectives across Hopper, Walker2d, HalfCheetah, and Ant. As shown in Table 2 , future-conditioned variants consistently outperform scalar return-based conditioning, while probabilistic JEPA further improves over deterministic future prediction. This indicates that the learned latent future representation F ^ t \hat{F}_{t} F ^ t
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.
Conclusion
In this paper, we introduced DT-JEPA, a predictive representation learning framework that integrates JEPA with DT for transferable decision making. DT-JEPA learns probabilistic future representations from offline trajectories and incorporates them into action generation through AdaLN. We further developed a critic-guided fine-tuning mechanism to adapt the pretrained policy to unseen reward objectives with limited interaction. Theoretically, we analyze the predictive learning objective and derive a finite-sample performance bound for future-conditioned policy learning. Experiments show that DT-JEPA achieves strong performance and efficient adaptation, validating predictive future representations as an effective condition information for generalizable sequence decision making.
AI Use Statement
Generative AI tools were used solely for language editing and presentation, including improvements to grammar, clarity, concision, and readability. They were not used to develop the research methodology, formulate mathematical claims, implement the proposed method, generate or process experimental data, conduct experiments, or interpret the results. All AI-assisted revisions were carefully reviewed and verified by the authors, who take full responsibility for the final content of the paper.
Related Work
Reinforcement learning algorithms face the challenge of long-term credit assignment, which is typically addressed through temporal-difference (TD) learning [19 ] . Transformers [28 ] , originally developed for sequence modeling, provide a powerful mechanism for capturing long-range dependencies across trajectories. Decision Transformer (DT) [6 ] builds on this perspective and formulates reinforcement learning as a conditional sequence modeling problem, representing trajectories as sequences of returns, states, and actions and predicting actions conditioned on desired returns. Prompt-DT [32 ] further introduces trajectory demonstrations as prompts to encode task-specific information for few-shot policy generalization. Generalized DT (GDT) [11 ] interprets conditional policy learning through hindsight information matching and extends the conditioning signal beyond returns to future states and other trajectory statistics. Meta-DT [30 ] further applies conditional sequence modeling to offline meta-RL by learning task representations through a disentangled world model for few-shot and zero-shot adaptation. These studies suggest that the conditioning information is central to the generalization ability of sequence-based policies.
JEPA-based Representation Learning for RL
Future latent prediction has been widely used in reinforcement learning to learn decision-relevant representations and improve sample efficiency without reconstructing raw observations [25 , 13 ] . More recently, JEPA provide a general framework for predictive representation learning by predicting the latent representation of a target from an observed context rather than reconstructing the target in the original input space [21 ] . This work has been later extended to video representation learning through V-JEPA, where models learn by predicting latent representations of masked spatiotemporal regions [3 ] . In reinforcement learning, several representative methods have explored latent prediction from different perspectives. Recent methods such as TD-JEPA [2 ] further study reward-free representation learning through policy-conditioned temporal prediction, emphasizing long-term dynamics that can support zero-shot control. Kenneweg et al. Kenneweg et al., [20 ] further combine JEPA with PPO, jointly optimizing action-conditioned latent prediction and actor–critic objectives for reinforcement learning from visual observations.
Reward-free Pretraining and Zero-shot RL
Reward-free RL focuses on policy learning from unlabeled datasets without relying on task-specific reward supervision. Forward-Backward (FB) representations learn state-action occupancy representations that support zero-shot policy construction for downstream rewards [27 ] (Toutiati & Ollivier 2021). HILP instead learns a temporally structured latent space for discovering and composing goal-directed behaviors from reward-free data [23 ] . RLDP explores reusable policy representations learned from offline trajectories to facilitate generalization to unseen reward specifications [18 ] . These approaches highlight different ways of extracting task-agnostic structure from reward-free experience for downstream control. These methods demonstrate that reward-free RL methods can acquire reusable behavioral knowledge without specific reward supervision.
Proofs
Proof of Theorem 1
According to the chain rule of mutual information, we have
I ( H t ; T ≥ t , F t ) = I ( H t ; T ≥ t ) + I ( H t ; F t ∣ T ≥ t ) = I ( H t ; F t ) + I ( H t ; T ≥ t ∣ F t ) . (17) I(\mathbf{H}_t;\mathcal{T}_{\geq t},\mathbf{F}_t)=I(\mathbf{H}_t;\mathcal{T}_{\geq t})+I(\mathbf{H}_t;\mathbf{F}_t\mid\mathcal{T}_{\geq t})
=I(\mathbf{H}_t;\mathbf{F}_t)+I(\mathbf{H}_t;\mathcal{T}_{\geq t}\mid\mathbf{F}_t). \tag*{(17)} I ( H t ; T ≥ t , F t ) = I ( H t ; T ≥ t ) + I ( H t ; F t ∣ T ≥ t ) = I ( H t ; F t ) + I ( H t ; T ≥ t ∣ F t ) . (17) Since F t F_t F t τ ≥ t \tau_{\geq t} τ ≥ t τ ≥ t \tau_{\geq t} τ ≥ t F t F_t F t H t H_t H t
I ( H t ; F t ∣ τ ≥ t ) = 0. (18) I(H_t;F_t \mid\tau_{\geq t}) = 0. \tag*{(18)} I ( H t ; F t ∣ τ ≥ t ) = 0. (18) Substituting equation 18 into equation 17 yields
I ( H t ; τ ≥ t ) = I ( H t ; F t ) + I ( H t ; τ ≥ t ∣ F t ) . (19) I(H_t;\tau_{\geq t}) = I(H_t;F_t) + I(H_t;\tau_{\geq t} \mid F_t). \tag*{(19)} I ( H t ; τ ≥ t ) = I ( H t ; F t ) + I ( H t ; τ ≥ t ∣ F t ) . (19) Since mutual information is non-negative,
I ( H t ; τ ≥ t ∣ F t ) ≥ 0 , (20) I(H_t;\tau_{\geq t} \mid F_t) \geq0, \tag*{(20)} I ( H t ; τ ≥ t ∣ F t ) ≥ 0 , (20) we obtain
I ( H t ; τ ≥ t ) ≥ I ( F t ; H t ) . (21) I(H_t;\tau_{\geq t}) \geq I(F_t;H_t). \tag*{(21)} I ( H t ; τ ≥ t ) ≥ I ( F t ; H t ) . (21) According to the definition of mutual information, the mutual information between F t F_t F t H t H_t H t
I ( F t ; H t ) = E p ( F t , H t ) [ log p ( F t , H t ) p ( F t ) p ( H t ) ] . (22) I(F_t;H_t) = \mathbb{E}_{p(F_t,H_t)}\left[\log\frac{p(F_t,H_t)}{p(F_t)p(H_t)}\right]. \tag*{(22)} I ( F t ; H t ) = E p ( F t , H t ) [ log p ( F t ) p ( H t ) p ( F t , H t ) ] . (22) The joint distribution can be factorized as
p ( F t , H t ) = p ( F t ∣ H t ) p ( H t ) . (23) p(F_t,H_t) = p(F_t \mid H_t)p(H_t). \tag*{(23)} p ( F t , H t ) = p ( F t ∣ H t ) p ( H t ) . (23) Substituting equation 23 into equation 22, we obtain
I ( F t ; H t ) = E p ( F t , H t ) [ log p ( F t ∣ H t ) p ( H t ) p ( F t ) p ( H t ) ] = E p ( F t , H t ) [ log p ( F t ∣ H t ) p ( F t ) ] . (24) \begin{aligned}
I(F_t;H_t) &= \mathbb{E}_{p(F_t,H_t)}\left[\log\frac{p(F_t \mid H_t)p(H_t)}{p(F_t)p(H_t)}\right] \\
&= \mathbb{E}_{p(F_t,H_t)}\left[\log\frac{p(F_t \mid H_t)}{p(F_t)}\right]. \tag*{(24)}
\end{aligned} I ( F t ; H t ) = E p ( F t , H t ) [ log p ( F t ) p ( H t ) p ( F t ∣ H t ) p ( H t ) ] = E p ( F t , H t ) [ log p ( F t ) p ( F t ∣ H t ) ] . (24) Further expanding the logarithmic term gives
I ( F t ; H t ) = E p ( F t , H t ) [ log p ( F t ∣ H t ) − log p ( F t ) ] = E p ( F t , H t ) [ log p ( F t ∣ H t ) ] − E p ( F t , H t ) [ log p ( F t ) ] . (25) \begin{aligned}
I(F_t;H_t) &= \mathbb{E}_{p(F_t,H_t)}[\log p(F_t \mid H_t)-\log p(F_t)] \\
&= \mathbb{E}_{p(F_t,H_t)}[\log p(F_t \mid H_t)]-\mathbb{E}_{p(F_t,H_t)}[\log p(F_t)]. \tag*{(25)}
\end{aligned} I ( F t ; H t ) = E p ( F t , H t ) [ log p ( F t ∣ H t ) − log p ( F t )] = E p ( F t , H t ) [ log p ( F t ∣ H t )] − E p ( F t , H t ) [ log p ( F t )] . (25) Since log p ( F t ) \log p(F_t) log p ( F t ) H t H_t H t H t H_t H t
E p ( F t , H t ) [ log p ( F t ) ] = ∫ p ( F t , H t ) log p ( F t ) d F t d H t = ∫ [ ∫ p ( F t , H t ) d H t ] log p ( F t ) d F t = ∫ p ( F t ) log p ( F t ) d F t = − H ( F t ) . (26) \begin{aligned}
\mathbb{E}_{p(F_t,H_t)}[\log p(F_t)]
&= \int p(F_t,H_t)\log p(F_t)\,\mathrm{d}F_t\,\mathrm{d}H_t \\
&= \int\left[\int p(F_t,H_t)\,\mathrm{d}H_t\right]\log p(F_t)\,\mathrm{d}F_t \\
&= \int p(F_t)\log p(F_t)\,\mathrm{d}F_t \\
&= -H(F_t). \tag*{(26)}
\end{aligned} E p ( F t , H t ) [ log p ( F t )] = ∫ p ( F t , H t ) log p ( F t ) d F t d H t = ∫ [ ∫ p ( F t , H t ) d H t ] log p ( F t ) d F t = ∫ p ( F t ) log p ( F t ) d F t = − H ( F t ) . (26) Therefore, equation 25 can be rewritten as
I ( F t ; H t ) = H ( F t ) + E p ( F t , H t ) [ log p ( F t ∣ H t ) ] . (27) I(F_t;H_t) = H(F_t) + \mathbb{E}_{p(F_t,H_t)}[\log p(F_t \mid H_t)]. \tag*{(27)} I ( F t ; H t ) = H ( F t ) + E p ( F t , H t ) [ log p ( F t ∣ H t )] . (27) Equivalently,
H ( F t ∣ H t ) = − E p ( F t , H t ) [ log p ( F t ∣ H t ) ] , (28) H(F_t \mid H_t) = -\mathbb{E}_{p(F_t,H_t)}[\log p(F_t \mid H_t)], \tag*{(28)} H ( F t ∣ H t ) = − E p ( F t , H t ) [ log p ( F t ∣ H t )] , (28) and thus
I ( F t ; H t ) = H ( F t ) − H ( F t ∣ H t ) . (29) I(F_t;H_t) = H(F_t) - H(F_t \mid H_t). \tag*{(29)} I ( F t ; H t ) = H ( F t ) − H ( F t ∣ H t ) . (29) However, the true distribution p ( F t ∣ H t ) p(F_t \mid H_t) p ( F t ∣ H t ) q η ( F t ∣ H t ) q_{\eta}(F_t \mid H_t) q η ( F t ∣ H t ) η \eta η
For any given H t \mathbf{H}_t H t
D K L ( p ( F t ∣ H t ) ∥ q η ( F t ∣ H t ) ) ≥ 0. (30) D_{\mathrm{KL}}\left(p(\mathbf{F}_t \mid\mathbf{H}_t) \mathbin{\|} q_\eta(\mathbf{F}_t \mid\mathbf{H}_t)\right) \ge0. \tag*{(30)} D KL ( p ( F t ∣ H t ) ∥ q η ( F t ∣ H t ) ) ≥ 0. (30) Expanding the KL divergence gives
D K L ( p ( F t ∣ H t ) ∥ q η ( F t ∣ H t ) ) = ∫ p ( F t ∣ H t ) log p ( F t ∣ H t ) q η ( F t ∣ H t ) d F t = E p ( F t ∣ H t ) [ log p ( F t ∣ H t ) ] − E p ( F t ∣ H t ) [ log q η ( F t ∣ H t ) ] ≥ 0. (31) \begin{aligned}
D_{\mathrm{KL}}\left(p(\mathbf{F}_t \mid\mathbf{H}_t) \mathbin{\|} q_\eta(\mathbf{F}_t \mid\mathbf{H}_t)\right)
&= \int p(\mathbf{F}_t \mid\mathbf{H}_t)\log\frac{p(\mathbf{F}_t \mid\mathbf{H}_t)}{q_\eta(\mathbf{F}_t \mid\mathbf{H}_t)}\,\mathrm{d}\mathbf{F}_t \\
&= \mathbb{E}_{p(\mathbf{F}_t\mid\mathbf{H}_t)}[\log p(\mathbf{F}_t \mid\mathbf{H}_t)]-\mathbb{E}_{p(\mathbf{F}_t\mid\mathbf{H}_t)}[\log q_\eta(\mathbf{F}_t \mid\mathbf{H}_t)]\ge0. \tag*{(31)}
\end{aligned} D KL ( p ( F t ∣ H t ) ∥ q η ( F t ∣ H t ) ) = ∫ p ( F t ∣ H t ) log q η ( F t ∣ H t ) p ( F t ∣ H t ) d F t = E p ( F t ∣ H t ) [ log p ( F t ∣ H t )] − E p ( F t ∣ H t ) [ log q η ( F t ∣ H t )] ≥ 0. (31) It follows that
E p ( F t ∣ H t ) [ log p ( F t ∣ H t ) ] ≥ E p ( F t ∣ H t ) [ log q η ( F t ∣ H t ) ] . (32) \mathbb{E}_{p(\mathbf{F}_t\mid\mathbf{H}_t)}[\log p(\mathbf{F}_t \mid\mathbf{H}_t)]\ge\mathbb{E}_{p(\mathbf{F}_t\mid\mathbf{H}_t)}[\log q_\eta(\mathbf{F}_t \mid\mathbf{H}_t)]. \tag*{(32)} E p ( F t ∣ H t ) [ log p ( F t ∣ H t )] ≥ E p ( F t ∣ H t ) [ log q η ( F t ∣ H t )] . (32) Further taking the expectation over p ( H t ) p(\mathbf{H}_t) p ( H t )
E p ( H t ) E p ( F t ∣ H t ) [ log p ( F t ∣ H t ) ] ≥ E p ( H t ) E p ( F t ∣ H t ) [ log q η ( F t ∣ H t ) ] . (33) \mathbb{E}_{p(\mathbf{H}_t)}\mathbb{E}_{p(\mathbf{F}_t\mid\mathbf{H}_t)}[\log p(\mathbf{F}_t \mid\mathbf{H}_t)]\ge\mathbb{E}_{p(\mathbf{H}_t)}\mathbb{E}_{p(\mathbf{F}_t\mid\mathbf{H}_t)}[\log q_\eta(\mathbf{F}_t \mid\mathbf{H}_t)]. \tag*{(33)} E p ( H t ) E p ( F t ∣ H t ) [ log p ( F t ∣ H t )] ≥ E p ( H t ) E p ( F t ∣ H t ) [ log q η ( F t ∣ H t )] . (33) Since
p ( F t , H t ) = p ( H t ) p ( F t ∣ H t ) , (34) p(\mathbf{F}_t,\mathbf{H}_t)=p(\mathbf{H}_t)p(\mathbf{F}_t \mid\mathbf{H}_t), \tag*{(34)} p ( F t , H t ) = p ( H t ) p ( F t ∣ H t ) , (34) equation 33 is equivalent to
E p ( F t , H t ) [ log p ( F t ∣ H t ) ] ≥ E p ( F t , H t ) [ log q η ( F t ∣ H t ) ] . (35) \mathbb{E}_{p(\mathbf{F}_t,\mathbf{H}_t)}[\log p(\mathbf{F}_t \mid\mathbf{H}_t)]\ge\mathbb{E}_{p(\mathbf{F}_t,\mathbf{H}_t)}[\log q_\eta(\mathbf{F}_t \mid\mathbf{H}_t)]. \tag*{(35)} E p ( F t , H t ) [ log p ( F t ∣ H t )] ≥ E p ( F t , H t ) [ log q η ( F t ∣ H t )] . (35) Substituting equation 35 into equation 27, we obtain
I ( F t ; H t ) = H ( F t ) + E p ( F t , H t ) [ log p ( F t ∣ H t ) ] ≥ H ( F t ) + E p ( F t , H t ) [ log q η ( F t ∣ H t ) ] . (36) \begin{aligned}
I(\mathbf{F}_t;\mathbf{H}_t)&=H(\mathbf{F}_t)+\mathbb{E}_{p(\mathbf{F}_t,\mathbf{H}_t)}[\log p(\mathbf{F}_t \mid\mathbf{H}_t)]\\
&\ge H(\mathbf{F}_t)+\mathbb{E}_{p(\mathbf{F}_t,\mathbf{H}_t)}[\log q_\eta(\mathbf{F}_t \mid\mathbf{H}_t)]. \tag*{(36)}
\end{aligned} I ( F t ; H t ) = H ( F t ) + E p ( F t , H t ) [ log p ( F t ∣ H t )] ≥ H ( F t ) + E p ( F t , H t ) [ log q η ( F t ∣ H t )] . (36) Combining equation 21 and equation 36, we finally obtain
I ( H t ; τ ≥ t ) ≥ I ( F t ; H t ) ≥ E p ( F t , H t ) [ log q η ( F t ∣ H t ) ] + H ( F t ) . (37) \begin{aligned}
I(\mathbf{H}_t;\tau_{\ge t})&\ge I(\mathbf{F}_t;\mathbf{H}_t)\\
&\ge\mathbb{E}_{p(\mathbf{F}_t,\mathbf{H}_t)}[\log q_\eta(\mathbf{F}_t \mid\mathbf{H}_t)]+H(\mathbf{F}_t). \tag*{(37)}
\end{aligned} I ( H t ; τ ≥ t ) ≥ I ( F t ; H t ) ≥ E p ( F t , H t ) [ log q η ( F t ∣ H t )] + H ( F t ) . (37) Proof of Theorem 2
Theorem 2. Consider a finite-horizon MDP with horizon H H H [ 0 , 1 ] [0,1] [ 0 , 1 ] q η ( F ∣ x ) q_\eta(F\mid x) q η ( F ∣ x ) π ˉ η \bar{\pi}_\eta π ˉ η π ˉ η , ϕ ^ \bar{\pi}_{\eta,\hat{\phi}} π ˉ η , ϕ ^ [4 ] . Under bounded occupancy mismatch C F C_F C F Π F \Pi_F Π F c F c_F c F ϵ a p p \epsilon_{\mathrm{app}} ϵ app 1 − δ 1-\delta 1 − δ
J ( π ˉ η ) − J ( π ˉ η , ϕ ^ ) ≤ O ( C F H 2 [ c F ( log ( ∣ Π F ∣ / δ ) N ) 1 / 4 + ϵ a p p ] ) . (38) J(\bar{\pi}_\eta)-J(\bar{\pi}_{\eta,\hat{\phi}})\le O\left(C_FH^2\left[\sqrt{c_F}\left(\frac{\log(|\Pi_F|/\delta)}{N}\right)^{1/4}+\sqrt{\epsilon_{\mathrm{app}}}\right]\right). \tag*{(38)} J ( π ˉ η ) − J ( π ˉ η , ϕ ^ ) ≤ O ( C F H 2 [ c F ( N log ( ∣ Π F ∣/ δ ) ) 1/4 + ϵ app ] ) . (38) Proof. Let β \beta β D \mathcal{D} D q η ( F ^ t ∣ H t ) q_\eta(\hat{\mathbf{F}}_t\mid\mathbf{H}_t) q η ( F ^ t ∣ H t ) H t = f θ ( τ t − K + 1 : t ) \mathbf{H}_t=f_\theta(\tau_{t-K+1:t}) H t = f θ ( τ t − K + 1 : t ) F ^ t ∼ q η ( ⋅ ∣ H t ) \hat{\mathbf{F}}_t\sim q_\eta(\cdot\mid\mathbf{H}_t) F ^ t ∼ q η ( ⋅ ∣ H t ) π ϕ ( a t ∣ τ t − K + 1 : t , F ^ t ) \pi_\phi(a_t\mid\tau_{t-K+1:t},\hat{\mathbf{F}}_t) π ϕ ( a t ∣ τ t − K + 1 : t , F ^ t ) ϕ \phi ϕ
π ϕ ( a t ∣ τ t − K + 1 : t , F ^ t ) = N ( a t ; a ^ t , σ a 2 I ) . (39) \pi_\phi\left(a_t\mid\tau_{t-K+1:t},\hat{\mathbf{F}}_t\right)=\mathcal{N}\left(a_t;\hat{a}_t,\sigma_a^2 I\right). \tag*{(39)} π ϕ ( a t ∣ τ t − K + 1 : t , F ^ t ) = N ( a t ; a ^ t , σ a 2 I ) . (39) Thus, the negative log-likelihood satisfies
− log π ϕ ( a t ∣ τ t − K + 1 : t , F ^ t ) = 1 2 σ a 2 ∥ a t − a ^ t ∥ 2 2 + C , (40) -\log\pi_\phi\left(a_t\mid\tau_{t-K+1:t},\hat{\mathbf{F}}_t\right)=\frac{1}{2\sigma_a^2}\lVert a_t-\hat{a}_t\rVert_2^2+C, \tag*{(40)} − log π ϕ ( a t ∣ τ t − K + 1 : t , F ^ t ) = 2 σ a 2 1 ∥ a t − a ^ t ∥ 2 2 + C , (40) where C C C ϕ \phi ϕ (13) . Let d β ( τ t − K + 1 : t ) d_\beta(\tau_{t-K+1:t}) d β ( τ t − K + 1 : t ) β \beta β F ^ t \hat{F}_t F ^ t H t H_t H t
P η ( τ t − K + 1 : t , a t , F ^ t ) = d β ( τ t − K + 1 : t ) β ( a t ∣ τ t − K + 1 : t ) q η ( F ^ t ∣ H t ) . (41) P_\eta\left(\tau_{t-K+1:t}, a_t, \hat{F}_t\right) = d_\beta(\tau_{t-K+1:t})\beta(a_t \mid\tau_{t-K+1:t})q_\eta(\hat{F}_t \mid H_t). \tag*{(41)} P η ( τ t − K + 1 : t , a t , F ^ t ) = d β ( τ t − K + 1 : t ) β ( a t ∣ τ t − K + 1 : t ) q η ( F ^ t ∣ H t ) . (41) Consequently, conditioned on the observable historical trajectory and the predicted future representation, we have
P η ( a t ∣ τ t − K + 1 : t , F ^ t ) = d β ( τ t − K + 1 : t ) β ( a t ∣ τ t − K + 1 : t ) q η ( F ^ t ∣ H t ) d β ( τ t − K + 1 : t ) q η ( F ^ t ∣ H t ) = β ( a t ∣ τ t − K + 1 : t ) . (42) P_\eta\left(a_t \mid\tau_{t-K+1:t}, \hat{F}_t\right) = \frac{d_\beta(\tau_{t-K+1:t})\beta(a_t \mid\tau_{t-K+1:t})q_\eta(\hat{F}_t \mid H_t)}{d_\beta(\tau_{t-K+1:t})q_\eta(\hat{F}_t \mid H_t)} = \beta(a_t \mid\tau_{t-K+1:t}). \tag*{(42)} P η ( a t ∣ τ t − K + 1 : t , F ^ t ) = d β ( τ t − K + 1 : t ) q η ( F ^ t ∣ H t ) d β ( τ t − K + 1 : t ) β ( a t ∣ τ t − K + 1 : t ) q η ( F ^ t ∣ H t ) = β ( a t ∣ τ t − K + 1 : t ) . (42) For any predictive future-conditioned policy π ϕ ∈ Π F \pi_\phi\in\Pi_F π ϕ ∈ Π F
R η ( π ϕ ) = − E τ t − K + 1 : t ∼ d β , F ^ t ∼ q η ( ⋅ ∣ H t ) , a t ∼ β ( ⋅ ∣ τ t − K + 1 : t ) [ log π ϕ ( a t ∣ τ t − K + 1 : t , F ^ t ) ] . (43) \mathcal{R}_\eta(\pi_\phi) = -\mathbb{E}_{\substack{\tau_{t-K+1:t}\sim d_\beta,\\ \hat{F}_t\sim q_\eta(\cdot\mid H_t),\\ a_t\sim\beta(\cdot\mid\tau_{t-K+1:t})}}\left[\log\pi_\phi\left(a_t \mid\tau_{t-K+1:t}, \hat{F}_t\right)\right]. \tag*{(43)} R η ( π ϕ ) = − E τ t − K + 1 : t ∼ d β , F ^ t ∼ q η ( ⋅ ∣ H t ) , a t ∼ β ( ⋅ ∣ τ t − K + 1 : t ) [ log π ϕ ( a t ∣ τ t − K + 1 : t , F ^ t ) ] . (43) For each τ t − K + 1 : t \tau_{t-K+1:t} τ t − K + 1 : t F ^ t \hat{F}_t F ^ t
− E a t ∼ β ( ⋅ ∣ τ t − K + 1 : t ) [ log π ϕ ( a t ∣ τ t − K + 1 : t , F ^ t ) ] = − E a t ∼ β ( ⋅ ∣ τ t − K + 1 : t ) [ log β ( a t ∣ τ t − K + 1 : t ) ] + E a t ∼ β ( ⋅ ∣ τ t − K + 1 : t ) [ log β ( a t ∣ τ t − K + 1 : t ) π ϕ ( a t ∣ τ t − K + 1 : t , F ^ t ) ] = H ( β ( ⋅ ∣ τ t − K + 1 : t ) ) + D K L ( β ( ⋅ ∣ τ t − K + 1 : t ) ∥ π ϕ ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) ) . (44) \begin{aligned}
-\mathbb{E}_{a_t\sim\beta(\cdot\mid\tau_{t-K+1:t})}\left[\log\pi_\phi\left(a_t \mid\tau_{t-K+1:t}, \hat{F}_t\right)\right]
&= -\mathbb{E}_{a_t\sim\beta(\cdot\mid\tau_{t-K+1:t})}\left[\log\beta\left(a_t \mid\tau_{t-K+1:t}\right)\right] \\
&\quad+ \mathbb{E}_{a_t\sim\beta(\cdot\mid\tau_{t-K+1:t})}\left[\log\frac{\beta\left(a_t \mid\tau_{t-K+1:t}\right)}{\pi_\phi\left(a_t \mid\tau_{t-K+1:t}, \hat{F}_t\right)}\right] \\
&= H\left(\beta(\cdot\mid\tau_{t-K+1:t})\right) + D_{\mathrm{KL}}\left(\beta(\cdot\mid\tau_{t-K+1:t})\Vert\pi_\phi(\cdot\mid\tau_{t-K+1:t}, \hat{F}_t)\right). \tag*{(44)}
\end{aligned} − E a t ∼ β ( ⋅ ∣ τ t − K + 1 : t ) [ log π ϕ ( a t ∣ τ t − K + 1 : t , F ^ t ) ] = − E a t ∼ β ( ⋅ ∣ τ t − K + 1 : t ) [ log β ( a t ∣ τ t − K + 1 : t ) ] + E a t ∼ β ( ⋅ ∣ τ t − K + 1 : t ) log π ϕ ( a t ∣ τ t − K + 1 : t , F ^ t ) β ( a t ∣ τ t − K + 1 : t ) = H ( β ( ⋅ ∣ τ t − K + 1 : t ) ) + D KL ( β ( ⋅ ∣ τ t − K + 1 : t ) ∥ π ϕ ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) ) . (44) Taking expectation over the historical trajectory and the predicted future representation gives
R η ( π ϕ ) = H β + E η ( π ϕ ) , (45) \mathcal{R}_\eta(\pi_\phi) = \mathcal{H}_\beta+ \mathcal{E}_\eta(\pi_\phi), \tag*{(45)} R η ( π ϕ ) = H β + E η ( π ϕ ) , (45) where
H β = E τ t − K + 1 : t ∼ d β [ H ( β ( ⋅ ∣ τ t − K + 1 : t ) ) ] , (46) \mathcal{H}_\beta= \mathbb{E}_{\tau_{t-K+1:t}\sim d_\beta}\left[H\left(\beta(\cdot\mid\tau_{t-K+1:t})\right)\right], \tag*{(46)} H β = E τ t − K + 1 : t ∼ d β [ H ( β ( ⋅ ∣ τ t − K + 1 : t ) ) ] , (46) E η ( π ϕ ) = E τ t − K + 1 : t ∼ d β , F ^ t ∼ q η ( ⋅ ∣ H t ) [ D K L ( β ( ⋅ ∣ τ t − K + 1 : t ) ∥ π ϕ ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) ) ] . (47) \mathcal{E}_\eta(\pi_\phi) = \mathbb{E}_{\tau_{t-K+1:t}\sim d_\beta,\hat{F}_t\sim q_\eta(\cdot\mid H_t)}\left[D_{\mathrm{KL}}\left(\beta(\cdot\mid\tau_{t-K+1:t})\Vert\pi_\phi(\cdot\mid\tau_{t-K+1:t}, \hat{F}_t)\right)\right]. \tag*{(47)} E η ( π ϕ ) = E τ t − K + 1 : t ∼ d β , F ^ t ∼ q η ( ⋅ ∣ H t ) [ D KL ( β ( ⋅ ∣ τ t − K + 1 : t ) ∥ π ϕ ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) ) ] . (47) Since H β \mathcal{H}_\beta H β
π η P ∈ arg min π ϕ ∈ Π F R η ( π ϕ ) = arg min π ϕ ∈ Π F E η ( π ϕ ) , \pi_\eta^{\mathrm{P}} \in\underset{\pi_\phi\in\Pi_F}{\arg\min}\ \mathcal{R}_\eta(\pi_\phi) = \underset{\pi_\phi\in\Pi_F}{\arg\min}\ \mathcal{E}_\eta(\pi_\phi), π η P ∈ π ϕ ∈ Π F arg min R η ( π ϕ ) = π ϕ ∈ Π F arg min E η ( π ϕ ) , E η ( π η P ) = inf π ϕ ∈ Π F E η ( π ϕ ) ≤ ϵ a p p F . (48) \mathcal{E}_\eta(\pi_\eta^{\mathrm{P}}) = \underset{\pi_\phi\in\Pi_F}{\inf}\ \mathcal{E}_\eta(\pi_\phi) \leq\epsilon_{\mathrm{app}}^{F}. \tag*{(48)} E η ( π η P ) = π ϕ ∈ Π F inf E η ( π ϕ ) ≤ ϵ app F . (48) Since the predicted future representations are sampled from the same predictor during policy learning and policy execution, we have
π ˉ η P ( a t ∣ τ t − K + 1 : t ) = ∫ q η ( F ^ t ∣ H t ) π η P ( a t ∣ τ t − K + 1 : t , F ^ t ) d F ^ t , \bar{\pi}_\eta^{\mathrm{P}}(a_t \mid\tau_{t-K+1:t}) = \int q_\eta(\hat{F}_t \mid H_t)\pi_\eta^{\mathrm{P}}(a_t \mid\tau_{t-K+1:t}, \hat{F}_t)\,d\hat{F}_t, π ˉ η P ( a t ∣ τ t − K + 1 : t ) = ∫ q η ( F ^ t ∣ H t ) π η P ( a t ∣ τ t − K + 1 : t , F ^ t ) d F ^ t , π ˉ η , ϕ ^ ( a t ∣ τ t − K + 1 : t ) = ∫ q η ( F ^ t ∣ H t ) π ϕ ^ ( a t ∣ τ t − K + 1 : t , F ^ t ) d F ^ t . (49) \bar{\pi}_{\eta,\hat{\phi}}(a_t \mid\tau_{t-K+1:t}) = \int q_\eta(\hat{F}_t \mid H_t)\pi_{\hat{\phi}}(a_t \mid\tau_{t-K+1:t}, \hat{F}_t)\,d\hat{F}_t. \tag*{(49)} π ˉ η , ϕ ^ ( a t ∣ τ t − K + 1 : t ) = ∫ q η ( F ^ t ∣ H t ) π ϕ ^ ( a t ∣ τ t − K + 1 : t , F ^ t ) d F ^ t . (49) Next, we relate the discrepancy between these two policies. For timestep h h h d π ˉ η P ( τ h − K + 1 : h ) = P π ˉ η P ( τ h − K + 1 : h ) d_{\bar{\pi}_\eta^{\mathrm{P}}}(\tau_{h-K+1:h}) = P_{\bar{\pi}_\eta^{\mathrm{P}}}(\tau_{h-K+1:h}) d π ˉ η P ( τ h − K + 1 : h ) = P π ˉ η P ( τ h − K + 1 : h ) d π ˉ η , ϕ ^ h ( τ h − K + 1 : h ) = P π ˉ η , ϕ ^ ( τ h − K + 1 : h ) d^h_{\bar{\pi}_{\eta,\hat{\phi}}}(\tau_{h-K+1:h}) = P_{\bar{\pi}_{\eta,\hat{\phi}}}(\tau_{h-K+1:h}) d π ˉ η , ϕ ^ h ( τ h − K + 1 : h ) = P π ˉ η , ϕ ^ ( τ h − K + 1 : h ) Δ h = ∥ d π ˉ η P h − d π ˉ η , ϕ ^ h ∥ 1 \Delta_h = \left\|d^h_{\bar{\pi}_\eta^{\mathrm{P}}} - d^h_{\bar{\pi}_{\eta,\hat{\phi}}}\right\|_1 Δ h = d π ˉ η P h − d π ˉ η , ϕ ^ h 1 δ h = 2 E τ h − K + 1 : h ∼ d π ˉ η P h [ D T V ( π ˉ η P ( ⋅ ∣ τ h − K + 1 : h ) , π ˉ η , ϕ ^ ( ⋅ ∣ τ h − K + 1 : h ) ) ] \delta_h = 2\mathbb{E}_{\tau_{h-K+1:h}\sim d_{\bar{\pi}^{P}_{\eta}}^h}\left[D_{\mathrm{TV}}\left(\bar{\pi}^{P}_{\eta}(\cdot\mid\tau_{h-K+1:h}),\bar{\pi}_{\eta,\hat{\phi}}(\cdot\mid\tau_{h-K+1:h})\right)\right] δ h = 2 E τ h − K + 1 : h ∼ d π ˉ η P h [ D TV ( π ˉ η P ( ⋅ ∣ τ h − K + 1 : h ) , π ˉ η , ϕ ^ ( ⋅ ∣ τ h − K + 1 : h ) ) ] Δ 1 = 0 \Delta_1=0 Δ 1 = 0 P K ( τ h − K + 2 : h + 1 ∣ τ h − K + 1 : h , a h ) P_K(\tau_{h-K+2:h+1}\mid\tau_{h-K+1:h},a_h) P K ( τ h − K + 2 : h + 1 ∣ τ h − K + 1 : h , a h ) P ( s h + 1 ∣ s h , a h ) P(s_{h+1}\mid s_h,a_h) P ( s h + 1 ∣ s h , a h ) h > 1 h>1 h > 1
d π ˉ η P h ( τ h − K + 1 : h ) = ∫ d π ˉ η P h − 1 ( τ h − K : h − 1 ) π ˉ η P ( a h − 1 ∣ τ h − K : h − 1 ) × P K ( τ h − K + 1 : h ∣ τ h − K : h − 1 , a h − 1 ) d a h − 1 d τ h − K : h − 1 , d_{\bar{\pi}^{P}_{\eta}}^h(\tau_{h-K+1:h})
=
\int d_{\bar{\pi}^{P}_{\eta}}^{h-1}(\tau_{h-K:h-1})\bar{\pi}^{P}_{\eta}(a_{h-1}\mid\tau_{h-K:h-1})
\times P_K(\tau_{h-K+1:h}\mid\tau_{h-K:h-1},a_{h-1})\,da_{h-1}\,d\tau_{h-K:h-1}, d π ˉ η P h ( τ h − K + 1 : h ) = ∫ d π ˉ η P h − 1 ( τ h − K : h − 1 ) π ˉ η P ( a h − 1 ∣ τ h − K : h − 1 ) × P K ( τ h − K + 1 : h ∣ τ h − K : h − 1 , a h − 1 ) d a h − 1 d τ h − K : h − 1 , d π ˉ η , ϕ ^ h ( τ h − K + 1 : h ) = ∫ d π ˉ η , ϕ ^ h − 1 ( τ h − K : h − 1 ) π ˉ η , ϕ ^ ( a h − 1 ∣ τ h − K : h − 1 ) × P K ( τ h − K + 1 : h ∣ τ h − K : h − 1 , a h − 1 ) d a h − 1 d τ h − K : h − 1 . (50) d_{\bar{\pi}_{\eta,\hat{\phi}}}^h(\tau_{h-K+1:h})
=
\int d_{\bar{\pi}_{\eta,\hat{\phi}}}^{h-1}(\tau_{h-K:h-1})\bar{\pi}_{\eta,\hat{\phi}}(a_{h-1}\mid\tau_{h-K:h-1})
\times P_K(\tau_{h-K+1:h}\mid\tau_{h-K:h-1},a_{h-1})\,da_{h-1}\,d\tau_{h-K:h-1}. \tag*{(50)} d π ˉ η , ϕ ^ h ( τ h − K + 1 : h ) = ∫ d π ˉ η , ϕ ^ h − 1 ( τ h − K : h − 1 ) π ˉ η , ϕ ^ ( a h − 1 ∣ τ h − K : h − 1 ) × P K ( τ h − K + 1 : h ∣ τ h − K : h − 1 , a h − 1 ) d a h − 1 d τ h − K : h − 1 . (50) Expanding the definitions and adding and subtracting the cross term d π ˉ η P h − 1 ( τ h − K : h − 1 ) π ˉ η , ϕ ^ ( a h − 1 ∣ τ h − K : h − 1 ) d_{\bar{\pi}^{P}_{\eta}}^{h-1}(\tau_{h-K:h-1})\bar{\pi}_{\eta,\hat{\phi}}(a_{h-1}\mid\tau_{h-K:h-1}) d π ˉ η P h − 1 ( τ h − K : h − 1 ) π ˉ η , ϕ ^ ( a h − 1 ∣ τ h − K : h − 1 )
Δ h = ∫ ∣ d π ˉ η P h ( τ h − K + 1 : h ) − d π ˉ η , ϕ ^ h ( τ h − K + 1 : h ) ∣ d τ h − K + 1 : h ≤ ∫ ∣ ∫ d π ˉ η P h − 1 ( τ h − K : h − 1 ) [ π ˉ η P ( a h − 1 ∣ τ h − K : h − 1 ) − π ˉ η , ϕ ^ ( a h − 1 ∣ τ h − K : h − 1 ) ] × P K ( τ h − K + 1 : h ∣ τ h − K : h − 1 , a h − 1 ) d a h − 1 d τ h − K : h − 1 ∣ d τ h − K + 1 : h + ∫ ∣ ∫ [ d π ˉ η P h − 1 ( τ h − K : h − 1 ) − d π ˉ η , ϕ ^ h − 1 ( τ h − K : h − 1 ) ] × π ˉ η , ϕ ^ ( a h − 1 ∣ τ h − K : h − 1 ) P K ( τ h − K + 1 : h ∣ τ h − K : h − 1 , a h − 1 ) d a h − 1 d τ h − K : h − 1 ∣ d τ h − K + 1 : h ≤ ∫ d π ˉ η P h − 1 ( τ h − K : h − 1 ) ∫ ∣ π ˉ η P ( a h − 1 ∣ τ h − K : h − 1 ) − π ˉ η , ϕ ^ ( a h − 1 ∣ τ h − K : h − 1 ) ∣ d a h − 1 d τ h − K : h − 1 + Δ h − 1 = 2 E τ h − K : h − 1 ∼ d π ˉ η P h − 1 [ D T V ( π ˉ η P ( ⋅ ∣ τ h − K : h − 1 ) , π ˉ η , ϕ ^ ( ⋅ ∣ τ h − K : h − 1 ) ) ] + Δ h − 1 = δ h − 1 + Δ h − 1 . (51) \begin{aligned}
\Delta_h
&= \int\left|d_{\bar{\pi}^{P}_{\eta}}^h(\tau_{h-K+1:h})-d_{\bar{\pi}_{\eta,\hat{\phi}}}^h(\tau_{h-K+1:h})\right|d\tau_{h-K+1:h}\\
&\le\int\left|\int d_{\bar{\pi}^{P}_{\eta}}^{h-1}(\tau_{h-K:h-1})\left[\bar{\pi}^{P}_{\eta}(a_{h-1}\mid\tau_{h-K:h-1})-\bar{\pi}_{\eta,\hat{\phi}}(a_{h-1}\mid\tau_{h-K:h-1})\right]\right.\\
&\qquad\left.\times P_K(\tau_{h-K+1:h}\mid\tau_{h-K:h-1},a_{h-1})\,da_{h-1}\,d\tau_{h-K:h-1}\right|d\tau_{h-K+1:h}\\
&\quad+\int\left|\int\left[d_{\bar{\pi}^{P}_{\eta}}^{h-1}(\tau_{h-K:h-1})-d_{\bar{\pi}_{\eta,\hat{\phi}}}^{h-1}(\tau_{h-K:h-1})\right]\right.\\
&\qquad\left.\times\bar{\pi}_{\eta,\hat{\phi}}(a_{h-1}\mid\tau_{h-K:h-1})P_K(\tau_{h-K+1:h}\mid\tau_{h-K:h-1},a_{h-1})\,da_{h-1}\,d\tau_{h-K:h-1}\right|d\tau_{h-K+1:h}\\
&\le\int d_{\bar{\pi}^{P}_{\eta}}^{h-1}(\tau_{h-K:h-1})\int\left|\bar{\pi}^{P}_{\eta}(a_{h-1}\mid\tau_{h-K:h-1})-\bar{\pi}_{\eta,\hat{\phi}}(a_{h-1}\mid\tau_{h-K:h-1})\right|\,da_{h-1}\,d\tau_{h-K:h-1}+\Delta_{h-1}\\
&=2\mathbb{E}_{\tau_{h-K:h-1}\sim d_{\bar{\pi}^{P}_{\eta}}^{h-1}}\left[D_{\mathrm{TV}}\left(\bar{\pi}^{P}_{\eta}(\cdot\mid\tau_{h-K:h-1}),\bar{\pi}_{\eta,\hat{\phi}}(\cdot\mid\tau_{h-K:h-1})\right)\right]+\Delta_{h-1}\\
&=\delta_{h-1}+\Delta_{h-1}. \tag*{(51)}
\end{aligned} Δ h = ∫ d π ˉ η P h ( τ h − K + 1 : h ) − d π ˉ η , ϕ ^ h ( τ h − K + 1 : h ) d τ h − K + 1 : h ≤ ∫ ∫ d π ˉ η P h − 1 ( τ h − K : h − 1 ) [ π ˉ η P ( a h − 1 ∣ τ h − K : h − 1 ) − π ˉ η , ϕ ^ ( a h − 1 ∣ τ h − K : h − 1 ) ] × P K ( τ h − K + 1 : h ∣ τ h − K : h − 1 , a h − 1 ) d a h − 1 d τ h − K : h − 1 ∣ d τ h − K + 1 : h + ∫ ∫ [ d π ˉ η P h − 1 ( τ h − K : h − 1 ) − d π ˉ η , ϕ ^ h − 1 ( τ h − K : h − 1 ) ] × π ˉ η , ϕ ^ ( a h − 1 ∣ τ h − K : h − 1 ) P K ( τ h − K + 1 : h ∣ τ h − K : h − 1 , a h − 1 ) d a h − 1 d τ h − K : h − 1 d τ h − K + 1 : h ≤ ∫ d π ˉ η P h − 1 ( τ h − K : h − 1 ) ∫ π ˉ η P ( a h − 1 ∣ τ h − K : h − 1 ) − π ˉ η , ϕ ^ ( a h − 1 ∣ τ h − K : h − 1 ) d a h − 1 d τ h − K : h − 1 + Δ h − 1 = 2 E τ h − K : h − 1 ∼ d π ˉ η P h − 1 [ D TV ( π ˉ η P ( ⋅ ∣ τ h − K : h − 1 ) , π ˉ η , ϕ ^ ( ⋅ ∣ τ h − K : h − 1 ) ) ] + Δ h − 1 = δ h − 1 + Δ h − 1 . (51) Repeatedly applying the recursion and using Δ 1 = 0 \Delta_1=0 Δ 1 = 0
Δ h ≤ δ h − 1 + δ h − 2 + ⋯ + δ 1 = ∑ j = 1 h − 1 δ j . (52) \Delta_h\le\delta_{h-1}+\delta_{h-2}+\cdots+\delta_1=\sum_{j=1}^{h-1}\delta_j. \tag*{(52)} Δ h ≤ δ h − 1 + δ h − 2 + ⋯ + δ 1 = j = 1 ∑ h − 1 δ j . (52) Let d π ˉ η P ( τ t − K + 1 : t ) d_{\bar{\pi}^{P}_{\eta}}(\tau_{t-K+1:t}) d π ˉ η P ( τ t − K + 1 : t ) d π ˉ η , ϕ ^ ( τ t − K + 1 : t ) d_{\bar{\pi}_{\eta,\hat{\phi}}}(\tau_{t-K+1:t}) d π ˉ η , ϕ ^ ( τ t − K + 1 : t ) δ h \delta_h δ h
∥ d π ˉ η P − d π ˉ η , ϕ ^ ∥ 1 ≤ 1 T ∑ h = 1 T Δ h ≤ 2 T E τ t − K + 1 : t ∼ d π ˉ η P [ D T V ( π ˉ η P ( ⋅ ∣ τ t − K + 1 : t ) , π ˉ η , ϕ ^ ( ⋅ ∣ τ t − K + 1 : t ) ) ] . (53) \left\|d_{\bar{\pi}^{P}_{\eta}}-d_{\bar{\pi}_{\eta,\hat{\phi}}}\right\|_1
\le\frac{1}{T}\sum_{h=1}^{T}\Delta_h
\le2T\mathbb{E}_{\tau_{t-K+1:t}\sim d_{\bar{\pi}^{P}_{\eta}}}\left[D_{\mathrm{TV}}\left(\bar{\pi}^{P}_{\eta}(\cdot\mid\tau_{t-K+1:t}),\bar{\pi}_{\eta,\hat{\phi}}(\cdot\mid\tau_{t-K+1:t})\right)\right]. \tag*{(53)} d π ˉ η P − d π ˉ η , ϕ ^ 1 ≤ T 1 h = 1 ∑ T Δ h ≤ 2 T E τ t − K + 1 : t ∼ d π ˉ η P [ D TV ( π ˉ η P ( ⋅ ∣ τ t − K + 1 : t ) , π ˉ η , ϕ ^ ( ⋅ ∣ τ t − K + 1 : t ) ) ] . (53) We now relate the occupancy discrepancy to the return difference. Since rewards are bounded, the difference between the expected rewards at timestep h h h
J ( π ˉ η P ) − J ( π ˉ η , ϕ ^ ) ≤ ∑ h = 1 T ( Δ h + δ h ) ≤ ∑ h = 1 T δ h + ∑ h = 1 T ∑ j = 1 h − 1 δ j ≤ T ∑ h = 1 T δ h = 2 T 2 E τ t − K + 1 : t ∼ d π η P τ [ D T V ( π ˉ η P ( ⋅ ∣ τ t − K + 1 : t ) , π ˉ η , ϕ ^ ( ⋅ ∣ τ t − K + 1 : t ) ) ] . (54) \begin{aligned}
J(\bar{\pi}_{\eta}^{\mathrm{P}})-J(\bar{\pi}_{\eta,\hat{\phi}}) \\
&\leq\sum_{h=1}^{T}(\Delta_h+\delta_h)\leq\sum_{h=1}^{T}\delta_h+\sum_{h=1}^{T}\sum_{j=1}^{h-1}\delta_j\leq T\sum_{h=1}^{T}\delta_h \\
&=2T^{2}\mathbb{E}_{\tau_{t-K+1:t}\sim d_{\pi_{\eta}^{\mathrm{P}}}^{\tau}}\left[D_{\mathrm{TV}}\left(\bar{\pi}_{\eta}^{\mathrm{P}}(\cdot\mid\tau_{t-K+1:t}),\bar{\pi}_{\eta,\hat{\phi}}(\cdot\mid\tau_{t-K+1:t})\right)\right]. \tag*{(54)}
\end{aligned} J ( π ˉ η P ) − J ( π ˉ η , ϕ ^ ) ≤ h = 1 ∑ T ( Δ h + δ h ) ≤ h = 1 ∑ T δ h + h = 1 ∑ T j = 1 ∑ h − 1 δ j ≤ T h = 1 ∑ T δ h = 2 T 2 E τ t − K + 1 : t ∼ d π η P τ [ D TV ( π ˉ η P ( ⋅ ∣ τ t − K + 1 : t ) , π ˉ η , ϕ ^ ( ⋅ ∣ τ t − K + 1 : t ) ) ] . (54) To characterize the effect of distribution shift between the population policy and the offline behavior data, define C F : = ess sup τ t − K + 1 : t d π η P τ ( τ t − K + 1 : t ) d β ( τ t − K + 1 : t ) C_{F}:=\operatorname*{ess\,sup}_{\tau_{t-K+1:t}}\frac{d_{\pi_{\eta}^{\mathrm{P}}}^{\tau}(\tau_{t-K+1:t})}{d_{\beta}(\tau_{t-K+1:t})} C F := ess sup τ t − K + 1 : t d β ( τ t − K + 1 : t ) d π η P τ ( τ t − K + 1 : t ) d π η P τ ≪ d β d_{\pi_{\eta}^{\mathrm{P}}}^{\tau}\ll d_{\beta} d π η P τ ≪ d β C F C_{F} C F
J ( π ˉ η P ) − J ( π ˉ η , ϕ ^ ) J(\bar{\pi}_{\eta}^{\mathrm{P}})-J(\bar{\pi}_{\eta,\hat{\phi}}) J ( π ˉ η P ) − J ( π ˉ η , ϕ ^ ) ≤ 2 C F T 2 E τ t − K + 1 : t ∼ d β [ D T V ( π ˉ η P ( ⋅ ∣ τ t − K + 1 : t ) , π ˉ η , ϕ ^ ( ⋅ ∣ τ t − K + 1 : t ) ) ] . (55) \leq2C_{F}T^{2}\mathbb{E}_{\tau_{t-K+1:t}\sim d_{\beta}}\left[D_{\mathrm{TV}}\left(\bar{\pi}_{\eta}^{\mathrm{P}}(\cdot\mid\tau_{t-K+1:t}),\bar{\pi}_{\eta,\hat{\phi}}(\cdot\mid\tau_{t-K+1:t})\right)\right]. \tag*{(55)} ≤ 2 C F T 2 E τ t − K + 1 : t ∼ d β [ D TV ( π ˉ η P ( ⋅ ∣ τ t − K + 1 : t ) , π ˉ η , ϕ ^ ( ⋅ ∣ τ t − K + 1 : t ) ) ] . (55) We next expand the total variation discrepancy between the two marginal predictive-conditioned policies. From equation 49,
2 E τ t − K + 1 : t ∼ d β [ D T V ( π ˉ η P ( ⋅ ∣ τ t − K + 1 : t ) , π ˉ η , ϕ ^ ( ⋅ ∣ τ t − K + 1 : t ) ) ] = E τ t − K + 1 : t ∼ d β [ ∫ ∣ π ˉ η P ( a t ∣ τ t − K + 1 : t ) − π ˉ η , ϕ ^ ( a t ∣ τ t − K + 1 : t ) ∣ d a t ] = E τ t − K + 1 : t ∼ d β [ ∬ q η ( F ^ t ∣ H t ) ∣ π η P ( a t ∣ τ t − K + 1 : t , F ^ t ) − π ϕ ^ ( a t ∣ τ t − K + 1 : t , F ^ t ) ∣ d F ^ t d a t ] ≤ E τ t − K + 1 : t ∼ d β [ ∫ q η ( F ^ t ∣ H t ) ∫ ∣ π η P ( a t ∣ τ t − K + 1 : t , F ^ t ) − π ϕ ^ ( a t ∣ τ t − K + 1 : t , F ^ t ) ∣ d a t d F ^ t ] = 2 E τ t − K + 1 : t ∼ d β F ^ t ∼ q η ( ⋅ ∣ H t ) [ D T V ( π η P ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) , π ϕ ^ ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) ) ] . (56) \begin{aligned}
2\mathbb{E}_{\tau_{t-K+1:t}\sim d_{\beta}}\left[D_{\mathrm{TV}}\left(\bar{\pi}_{\eta}^{\mathrm{P}}(\cdot\mid\tau_{t-K+1:t}),\bar{\pi}_{\eta,\hat{\phi}}(\cdot\mid\tau_{t-K+1:t})\right)\right] \\
&=\mathbb{E}_{\tau_{t-K+1:t}\sim d_{\beta}}\left[\int\left|\bar{\pi}_{\eta}^{\mathrm{P}}(a_t\mid\tau_{t-K+1:t})-\bar{\pi}_{\eta,\hat{\phi}}(a_t\mid\tau_{t-K+1:t})\right|\,da_t\right] \\
&=\mathbb{E}_{\tau_{t-K+1:t}\sim d_{\beta}}\left[\iint q_{\eta}(\hat{F}_t\mid H_t)\left|\pi_{\eta}^{\mathrm{P}}(a_t\mid\tau_{t-K+1:t},\hat{F}_t)-\pi_{\hat{\phi}}(a_t\mid\tau_{t-K+1:t},\hat{F}_t)\right|\,d\hat{F}_t\,da_t\right] \\
&\leq\mathbb{E}_{\tau_{t-K+1:t}\sim d_{\beta}}\left[\int q_{\eta}(\hat{F}_t\mid H_t)\int\left|\pi_{\eta}^{\mathrm{P}}(a_t\mid\tau_{t-K+1:t},\hat{F}_t)-\pi_{\hat{\phi}}(a_t\mid\tau_{t-K+1:t},\hat{F}_t)\right|\,da_t\,d\hat{F}_t\right] \\
&=2\mathbb{E}_{\substack{\tau_{t-K+1:t}\sim d_{\beta}\\\hat{F}_t\sim q_{\eta}(\cdot\mid H_t)}}\left[D_{\mathrm{TV}}\left(\pi_{\eta}^{\mathrm{P}}(\cdot\mid\tau_{t-K+1:t},\hat{F}_t),\pi_{\hat{\phi}}(\cdot\mid\tau_{t-K+1:t},\hat{F}_t)\right)\right]. \tag*{(56)}
\end{aligned} 2 E τ t − K + 1 : t ∼ d β [ D TV ( π ˉ η P ( ⋅ ∣ τ t − K + 1 : t ) , π ˉ η , ϕ ^ ( ⋅ ∣ τ t − K + 1 : t ) ) ] = E τ t − K + 1 : t ∼ d β [ ∫ π ˉ η P ( a t ∣ τ t − K + 1 : t ) − π ˉ η , ϕ ^ ( a t ∣ τ t − K + 1 : t ) d a t ] = E τ t − K + 1 : t ∼ d β [ ∬ q η ( F ^ t ∣ H t ) π η P ( a t ∣ τ t − K + 1 : t , F ^ t ) − π ϕ ^ ( a t ∣ τ t − K + 1 : t , F ^ t ) d F ^ t d a t ] ≤ E τ t − K + 1 : t ∼ d β [ ∫ q η ( F ^ t ∣ H t ) ∫ π η P ( a t ∣ τ t − K + 1 : t , F ^ t ) − π ϕ ^ ( a t ∣ τ t − K + 1 : t , F ^ t ) d a t d F ^ t ] = 2 E τ t − K + 1 : t ∼ d β F ^ t ∼ q η ( ⋅ ∣ H t ) [ D TV ( π η P ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) , π ϕ ^ ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) ) ] . (56) Using the behavior policy as an intermediate distribution, the triangle inequality gives
D T V ( π η P ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) , π ϕ ^ ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) ) ≤ D T V ( π η P ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) , β ( ⋅ ∣ τ t − K + 1 : t ) ) + D T V ( β ( ⋅ ∣ τ t − K + 1 : t ) , π ϕ ^ ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) ) . (57) \begin{aligned}
D_{\mathrm{TV}}\left(\pi_{\eta}^{\mathrm{P}}(\cdot\mid\tau_{t-K+1:t},\hat{F}_t),\pi_{\hat{\phi}}(\cdot\mid\tau_{t-K+1:t},\hat{F}_t)\right) \\
&\leq D_{\mathrm{TV}}\left(\pi_{\eta}^{\mathrm{P}}(\cdot\mid\tau_{t-K+1:t},\hat{F}_t),\beta(\cdot\mid\tau_{t-K+1:t})\right) \\
&+D_{\mathrm{TV}}\left(\beta(\cdot\mid\tau_{t-K+1:t}),\pi_{\hat{\phi}}(\cdot\mid\tau_{t-K+1:t},\hat{F}_t)\right). \tag*{(57)}
\end{aligned} D TV ( π η P ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) , π ϕ ^ ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) ) ≤ D TV ( π η P ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) , β ( ⋅ ∣ τ t − K + 1 : t ) ) + D TV ( β ( ⋅ ∣ τ t − K + 1 : t ) , π ϕ ^ ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) ) . (57) Applying Pinsker’s inequality, D T V ( P , Q ) ≤ 1 2 D K L ( P ∥ Q ) D_{\mathrm{TV}}(P,Q)\leq\sqrt{\frac{1}{2}D_{\mathrm{KL}}(P\|Q)} D TV ( P , Q ) ≤ 2 1 D KL ( P ∥ Q )
2 E τ t − K + 1 : t ∼ d β [ D T V ( π ˉ η P ( ⋅ ∣ τ t − K + 1 : t ) , π ˉ η , ϕ ^ ( ⋅ ∣ τ t − K + 1 : t ) ) ] 2\mathbb{E}_{\tau_{t-K+1:t}\sim d_{\beta}}\left[D_{\mathrm{TV}}\left(\bar{\pi}_{\eta}^{\mathrm{P}}(\cdot\mid\tau_{t-K+1:t}),\bar{\pi}_{\eta,\hat{\phi}}(\cdot\mid\tau_{t-K+1:t})\right)\right] 2 E τ t − K + 1 : t ∼ d β [ D TV ( π ˉ η P ( ⋅ ∣ τ t − K + 1 : t ) , π ˉ η , ϕ ^ ( ⋅ ∣ τ t − K + 1 : t ) ) ] ≤ 2 E τ t − K + 1 : t , F ^ t [ 1 2 D K L ( β ( ⋅ ∣ τ t − K + 1 : t ) ∥ π η P ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) ) ] + 2 E τ t − K + 1 : t , F ^ t [ 1 2 D K L ( β ( ⋅ ∣ τ t − K + 1 : t ) ∥ π ϕ ^ ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) ) ] ≤ 2 E η ( π ˉ η P ) + 2 E η ( π ˉ ϕ ^ ) ≤ 2 ϵ a p p F + 2 E η ( π ˉ ϕ ^ ) . (58) \begin{aligned}
\leq2\mathbb{E}_{\tau_{t-K+1:t},\hat{F}_t}\left[\sqrt{\frac{1}{2}D_{\mathrm{KL}}\left(\beta(\cdot\mid\tau_{t-K+1:t})\|\pi_{\eta}^{\mathrm{P}}(\cdot\mid\tau_{t-K+1:t},\hat{F}_t)\right)}\right] \\
&+2\mathbb{E}_{\tau_{t-K+1:t},\hat{F}_t}\left[\sqrt{\frac{1}{2}D_{\mathrm{KL}}\left(\beta(\cdot\mid\tau_{t-K+1:t})\|\pi_{\hat{\phi}}(\cdot\mid\tau_{t-K+1:t},\hat{F}_t)\right)}\right] \\
&\leq\sqrt{2\mathcal{E}_{\eta}(\bar{\pi}_{\eta}^{\mathrm{P}})}+\sqrt{2\mathcal{E}_{\eta}(\bar{\pi}_{\hat{\phi}})} \\
&\leq\sqrt{2\epsilon_{\mathrm{app}}^{F}}+\sqrt{2\mathcal{E}_{\eta}(\bar{\pi}_{\hat{\phi}})}. \tag*{(58)}
\end{aligned} ≤ 2 E τ t − K + 1 : t , F ^ t [ 2 1 D KL ( β ( ⋅ ∣ τ t − K + 1 : t ) ∥ π η P ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) ) ] + 2 E τ t − K + 1 : t , F ^ t [ 2 1 D KL ( β ( ⋅ ∣ τ t − K + 1 : t ) ∥ π ϕ ^ ( ⋅ ∣ τ t − K + 1 : t , F ^ t ) ) ] ≤ 2 E η ( π ˉ η P ) + 2 E η ( π ˉ ϕ ^ ) ≤ 2 ϵ app F + 2 E η ( π ˉ ϕ ^ ) . (58) Combining equation (55) and equation (58) , we finally have
J ( π ˉ η P ) − J ( π ˉ η , ϕ ^ ) ≤ C F T 2 [ 2 ϵ a p p F + 2 E η ( π ^ ϕ ) ] . (59) J(\bar{\pi}_{\eta}^{\mathrm{P}})-J(\bar{\pi}_{\eta,\hat{\phi}})
\le C_F T^2\left[\sqrt{2\epsilon_{\mathrm{app}}^F}+\sqrt{2\mathcal{E}_{\eta}(\hat{\pi}_{\phi})}\right]. \tag*{(59)} J ( π ˉ η P ) − J ( π ˉ η , ϕ ^ ) ≤ C F T 2 [ 2 ϵ app F + 2 E η ( π ^ ϕ ) ] . (59) Thus, given N N N
R ^ η ( π ϕ ) = − 1 N ∑ i = 1 N log π ϕ ( a t i ∣ τ t i − K + 1 : t i , F ^ t i ) , \widehat{\mathcal{R}}_{\eta}(\pi_{\phi})=-\frac{1}{N}\sum_{i=1}^{N}\log\pi_{\phi}\left(a_{t_i}\mid\tau_{t_i-K+1:t_i},\hat{F}_{t_i}\right), R η ( π ϕ ) = − N 1 i = 1 ∑ N log π ϕ ( a t i ∣ τ t i − K + 1 : t i , F ^ t i ) , π ^ ϕ ^ ∈ arg min π ϕ ∈ Π F R ^ η ( π ϕ ) . (60) \hat{\pi}_{\hat{\phi}}\in\arg\min_{\pi_{\phi}\in\Pi_F}\widehat{\mathcal{R}}_{\eta}(\pi_{\phi}). \tag*{(60)} π ^ ϕ ^ ∈ arg π ϕ ∈ Π F min R η ( π ϕ ) . (60) Using R η ( π ϕ ) = H β + E η ( π ϕ ) \mathcal{R}_{\eta}(\pi_{\phi})=H_{\beta}+\mathcal{E}_{\eta}(\pi_{\phi}) R η ( π ϕ ) = H β + E η ( π ϕ )
E η ( π ^ ϕ ^ ) = E η ( π ^ ϕ ^ ) − E η ( π η P ) + E η ( π η P ) = R η ( π ^ ϕ ^ ) − R η ( π η P ) + E η ( π η P ) ≤ R η ( π ^ ϕ ^ ) − R η ( π η P ) + ϵ a p p F ≤ [ R η ( π ^ ϕ ^ ) − R ^ η ( π ^ ϕ ^ ) ] + [ R ^ η ( π ^ ϕ ^ ) − R ^ η ( π η P ) ] + [ R ^ η ( π η P ) − R η ( π η P ) ] + ϵ a p p F . (61) \begin{aligned}
\mathcal{E}_{\eta}(\hat{\pi}_{\hat{\phi}})&=\mathcal{E}_{\eta}(\hat{\pi}_{\hat{\phi}})-\mathcal{E}_{\eta}(\pi_{\eta}^{\mathrm{P}})+\mathcal{E}_{\eta}(\pi_{\eta}^{\mathrm{P}})\\
&=\mathcal{R}_{\eta}(\hat{\pi}_{\hat{\phi}})-\mathcal{R}_{\eta}(\pi_{\eta}^{\mathrm{P}})+\mathcal{E}_{\eta}(\pi_{\eta}^{\mathrm{P}})\\
&\le\mathcal{R}_{\eta}(\hat{\pi}_{\hat{\phi}})-\mathcal{R}_{\eta}(\pi_{\eta}^{\mathrm{P}})+\epsilon_{\mathrm{app}}^F\\
&\le\left[\mathcal{R}_{\eta}(\hat{\pi}_{\hat{\phi}})-\widehat{\mathcal{R}}_{\eta}(\hat{\pi}_{\hat{\phi}})\right]+\left[\widehat{\mathcal{R}}_{\eta}(\hat{\pi}_{\hat{\phi}})-\widehat{\mathcal{R}}_{\eta}(\pi_{\eta}^{\mathrm{P}})\right]\\
&\quad+\left[\widehat{\mathcal{R}}_{\eta}(\pi_{\eta}^{\mathrm{P}})-\mathcal{R}_{\eta}(\pi_{\eta}^{\mathrm{P}})\right]+\epsilon_{\mathrm{app}}^F. \tag*{(61)}
\end{aligned} E η ( π ^ ϕ ^ ) = E η ( π ^ ϕ ^ ) − E η ( π η P ) + E η ( π η P ) = R η ( π ^ ϕ ^ ) − R η ( π η P ) + E η ( π η P ) ≤ R η ( π ^ ϕ ^ ) − R η ( π η P ) + ϵ app F ≤ [ R η ( π ^ ϕ ^ ) − R η ( π ^ ϕ ^ ) ] + [ R η ( π ^ ϕ ^ ) − R η ( π η P ) ] + [ R η ( π η P ) − R η ( π η P ) ] + ϵ app F . (61) Since π ^ ϕ ^ \hat{\pi}_{\hat{\phi}} π ^ ϕ ^ Π F \Pi_F Π F
R ^ η ( π ^ ϕ ^ ) − R ^ η ( π η P ) ≤ 0 , \widehat{\mathcal{R}}_{\eta}(\hat{\pi}_{\hat{\phi}})-\widehat{\mathcal{R}}_{\eta}(\pi_{\eta}^{\mathrm{P}})\le0, R η ( π ^ ϕ ^ ) − R η ( π η P ) ≤ 0 , E η ( π ^ ϕ ^ ) ≤ 2 sup π ϕ ∈ Π F ∣ R η ( π ϕ ) − R ^ η ( π ϕ ) ∣ + ϵ a p p F . (62) \mathcal{E}_{\eta}(\hat{\pi}_{\hat{\phi}})\le2\sup_{\pi_{\phi}\in\Pi_F}\left|\mathcal{R}_{\eta}(\pi_{\phi})-\widehat{\mathcal{R}}_{\eta}(\pi_{\phi})\right|+\epsilon_{\mathrm{app}}^F. \tag*{(62)} E η ( π ^ ϕ ^ ) ≤ 2 π ϕ ∈ Π F sup R η ( π ϕ ) − R η ( π ϕ ) + ϵ app F . (62) Let c F c_F c F ℓ ϕ = − log π ϕ ( a t ∣ τ t − K + 1 : t , F ^ t ) \ell_{\phi}=-\log\pi_{\phi}\left(a_t\mid\tau_{t-K+1:t},\hat{F}_t\right) ℓ ϕ = − log π ϕ ( a t ∣ τ t − K + 1 : t , F ^ t ) ϕ ∈ Π F \phi\in\Pi_F ϕ ∈ Π F c F c_F c F 1 − δ 1-\delta 1 − δ
Pr ( sup π ϕ ∈ Π F ∣ R η ( π ϕ ) − R ^ η ( π ϕ ) ∣ > c F log ( 2 ∣ Π F ∣ / δ ) 2 N ) ≤ δ , \Pr\left(\sup_{\pi_{\phi}\in\Pi_F}\left|\mathcal{R}_{\eta}(\pi_{\phi})-\widehat{\mathcal{R}}_{\eta}(\pi_{\phi})\right|>c_F\sqrt{\frac{\log(2|\Pi_F|/\delta)}{2N}}\right)\le\delta, Pr ( π ϕ ∈ Π F sup R η ( π ϕ ) − R η ( π ϕ ) > c F 2 N log ( 2∣ Π F ∣/ δ ) ) ≤ δ , sup π ϕ ∈ Π F ∣ R η ( π ϕ ) − R ^ η ( π ϕ ) ∣ ≤ c F log ( 2 ∣ Π F ∣ / δ ) 2 N . (63) \sup_{\pi_{\phi}\in\Pi_F}\left|\mathcal{R}_{\eta}(\pi_{\phi})-\widehat{\mathcal{R}}_{\eta}(\pi_{\phi})\right|\le c_F\sqrt{\frac{\log(2|\Pi_F|/\delta)}{2N}}. \tag*{(63)} π ϕ ∈ Π F sup R η ( π ϕ ) − R η ( π ϕ ) ≤ c F 2 N log ( 2∣ Π F ∣/ δ ) . (63) Substituting this result into equation (62) , we obtain
E η ( π ^ ϕ ^ ) ≤ ϵ a p p F + 2 c F log ( 2 ∣ Π F ∣ / δ ) 2 N . (64) \mathcal{E}_{\eta}(\hat{\pi}_{\hat{\phi}})\le\epsilon_{\mathrm{app}}^F+2c_F\sqrt{\frac{\log(2|\Pi_F|/\delta)}{2N}}. \tag*{(64)} E η ( π ^ ϕ ^ ) ≤ ϵ app F + 2 c F 2 N log ( 2∣ Π F ∣/ δ ) . (64) Finally, substituting equation (64) into equation (59) , we have gives
J ( π η P ) − J ( π ˉ η , ϕ ^ ) ≤ C F T 2 [ 2 ϵ a p p F + 2 ϵ a p p F + 4 c F log ( 2 ∣ Π F ∣ / δ ) 2 N ] ≤ C F T 2 [ 2 2 ϵ a p p F + 2 c F ( log ( 2 ∣ Π F ∣ / δ ) 2 N ) 1 / 4 ] = O ( C F T 2 [ c F ( log ( ∣ Π F ∣ / δ ) N ) 1 / 4 + ϵ a p p F ] ) . (65) \begin{aligned}
J(\pi_{\eta}^{\mathrm{P}})-J(\bar{\pi}_{\eta,\hat{\phi}})
&\le C_F T^2\left[\sqrt{2\epsilon_{\mathrm{app}}^F}+\sqrt{2\epsilon_{\mathrm{app}}^F+4c_F\sqrt{\frac{\log(2|\Pi_F|/\delta)}{2N}}}\right]\\
&\le C_F T^2\left[2\sqrt{2\epsilon_{\mathrm{app}}^F}+2\sqrt{c_F}\left(\frac{\log(2|\Pi_F|/\delta)}{2N}\right)^{1/4}\right]\\
&=\mathcal{O}\left(C_FT^2\left[\sqrt{c_F}\left(\frac{\log(|\Pi_F|/\delta)}{N}\right)^{1/4}+\sqrt{\epsilon_{\mathrm{app}}^F}\right]\right). \tag*{(65)}
\end{aligned} J ( π η P ) − J ( π ˉ η , ϕ ^ ) ≤ C F T 2 2 ϵ app F + 2 ϵ app F + 4 c F 2 N log ( 2∣ Π F ∣/ δ ) ≤ C F T 2 [ 2 2 ϵ app F + 2 c F ( 2 N log ( 2∣ Π F ∣/ δ ) ) 1/4 ] = O ( C F T 2 [ c F ( N log ( ∣ Π F ∣/ δ ) ) 1/4 + ϵ app F ] ) . (65) Proof completed.
Experimental Details
Experimental Tasks and Datasets
Environments. We evaluate DT-JEPA on diverse D4RL benchmarks. For MuJoCo locomotion tasks, the experiments include four sequential control environments: Hopper, HalfCheetah, Walker2d, and Ant. These environments cover different levels of dynamical complexity, ranging from low-dimensional planar locomotion to high-dimensional three-dimensional quadrupedal control.
Figure 5. Illustration of four MuJoCo environments utilized in the paper.
Hopper. Hopper is a planar locomotion environment in which a single-legged robot must maintain balance while repeatedly hopping forward. The observation is 11-dimensional and contains the configurations and velocities of the robot body and articulated joints. The action space is 3-dimensional, with each action specifying the torque applied to one actuated joint. Due to its small action space but sensitive balance dynamics, Hopper provides a relatively compact benchmark for evaluating temporal decision making.
HalfCheetah. HalfCheetah considers a planar articulated robot whose objective is to generate coordinated motions for fast forward locomotion. Its observation space contains 17 dimensions describing joint configurations and velocities, while the 6-dimensional continuous action controls the torques applied to the actuated joints in the front and rear limbs. Compared with Hopper, HalfCheetah requires coordination across a larger set of coupled joints and admits a wider range of locomotion behaviors.
Walker2d. Walker2d is a two-dimensional bipedal locomotion task in which the agent controls two articulated legs to move forward while keeping the torso upright. The environment has a 17-dimensional observation space and a 6-dimensional action space. The actions correspond to torques applied to the thigh, leg, and foot joints of the two legs. Successful control requires both forward progression and stable coordination between the two legs, making the task particularly sensitive to errors in long-horizon action prediction.
Ant. Ant is a three-dimensional quadrupedal locomotion environment with four articulated legs and eight actuated joints. The agent must coordinate all four legs to produce stable forward motion. In the D4RL MuJoCo setting, the observation is 111-dimensional, incorporating joint configurations, velocities, and contact-related information, while the action space is 8-dimensional and specifies the torques applied to the eight controllable joints. Its substantially higher state and action complexity makes Ant a challenging benchmark for evaluating whether learned representations can capture long-horizon and multi-joint dynamics.
For Maze2D navigation tasks, we consider 3 environments with increasing difficulty: Umaze, Medium, and Large. In Maze2D, the agent controls a point mass to navigate through a two-dimensional maze toward a fixed goal. The environment has a 4-dimensional state space describing the position and velocity of the agent, and a 2-dimensional continuous action space specifying its acceleration. The reward is determined by the distance between the agent and the goal. The four environments mainly differ in maze structure and navigation horizon, providing progressively more challenging settings for evaluating long-horizon decision making.
Umaze. Umaze introduces a compact U-shaped maze that requires the agent to navigate around obstacles rather than moving directly toward the goal. Compared with Open, successful control requires the policy to account for the maze structure and coordinate a sequence of actions over a longer navigation horizon.
Medium. Medium contains a larger and more complex maze with longer feasible paths between the initial position and the goal. The agent must make appropriate route choices while maintaining stable motion through multiple corridors, increasing the demand for long-term trajectory reasoning.
Large. Large is the most challenging Maze2D configuration, featuring the largest maze and the longest navigation horizon. Reaching the target requires the agent to coordinate decisions over extended trajectories and avoid locally attractive but globally ineffective movements, making it a challenging benchmark for long-horizon navigation.
For AntMaze navigation tasks, we consider three maze configurations: Umaze, Medium, and Large. AntMaze extends Maze2D by replacing the point-mass agent with a quadrupedal Ant robot, resulting in a more challenging control problem that combines locomotion and long-horizon navigation. The environment has a 27-dimensional state space and an 8-dimensional continuous action space. The three configurations mainly differ in maze size and navigation complexity, providing progressively more challenging settings for evaluating long-horizon decision making.
Figure 6. (a) Umaze; (b) Medium; (c) Large. Illustration of four AntMaze environments utilized in the paper.
Umaze. Umaze is the smallest and simplest AntMaze configuration. The agent controls the Ant robot to navigate through a compact U-shaped maze toward the target location. Due to the relatively short navigation distance and simple maze structure, Umaze provides a basic setting for evaluating whether the learned policy can combine low-level locomotion control with goal-directed navigation.
Medium. Medium uses a larger maze with longer navigation paths and more complex route choices than Umaze. The agent must maintain stable quadrupedal locomotion while navigating through multiple maze corridors toward the target. The increased spatial scale requires the policy to capture longer-term dependencies between local movements and the overall navigation objective.
Large. Large is the most challenging of the three AntMaze configurations, featuring the largest maze and the longest navigation horizon. Successful control requires the agent to coordinate continuous locomotion over a long sequence of decisions while following an appropriate route through the maze. This setting therefore provides a demanding benchmark for evaluating long-horizon planning and the ability to exploit predictive representations over extended trajectories.
Offline datasets. For each locomotion domain, we consider three commonly used D4RL dataset qualities: medium , medium-replay , and medium-expert . These datasets differ not only in behavioral quality but also in the diversity and consistency of the trajectories they contain.
The medium dataset is collected using a partially trained policy with moderate task performance. It therefore mainly contains coherent but suboptimal behaviors. The medium-replay dataset consists of transitions stored in the replay buffer throughout the training process of the medium-level policy. As a result, it contains a substantially broader mixture of behaviors, including early-stage, low-quality, and partially improved trajectories. The medium-expert dataset combines trajectories from medium and expert policies, producing a heterogeneous dataset that contains both suboptimal and near-optimal behaviors.
For the AntMaze domain, we consider four commonly used D4RL dataset configurations: umaze , umaze-diverse , medium-play , and medium-diverse . These datasets differ in both maze complexity and the diversity of navigation behaviors contained in the offline data.
The umaze dataset uses the smallest maze layout with a fixed starting location and a fixed goal, and therefore contains relatively consistent navigation behaviors. The umaze-diverse dataset uses the same maze layout but introduces more diverse navigation trajectories through varying navigation goals and starting conditions. For the larger Medium maze, medium-play is collected using a set of hand-designed starting and goal locations, producing structured but heterogeneous navigation behaviors. In contrast, medium-diverse uses more broadly distributed navigation goals, resulting in a more diverse offline dataset with a wider range of trajectories. These configurations therefore progressively increase the difficulty of learning reusable long-horizon navigation behaviors from offline experience.
For the Maze2D domain, we evaluate on three D4RL environments: umaze , medium , and large . Unlike AntMaze, these configurations mainly differ in maze scale and navigation complexity rather than in the data-collection strategy. The umaze environment contains a compact U-shaped layout and relatively short navigation paths. The medium environment introduces a larger maze with longer routes and more complex spatial dependencies, while large further increases the maze size and navigation horizon. The corresponding offline datasets contain continuous navigation experience toward different goals, making the larger layouts increasingly challenging for learning long-horizon dependencies and stitching useful behaviors across trajectories.
Baselines
We provide a more detailed introduction to the baselines used in our experiments in this section.
TD3-BC [9 ] is a widely used offline reinforcement learning method built upon Twin Delayed Deep Deterministic Policy Gradient (TD3). It combines value-based policy optimization with a behavior cloning objective that regularizes the learned policy toward actions contained in the offline dataset. This simple modification reduces the tendency of the policy to exploit out-of-distribution actions while still allowing policy improvement through critic guidance. TD3-BC provides a representative non-sequence-modeling offline RL baseline for evaluating whether trajectory-level sequence modeling offers advantages over conventional actor–critic methods.
TD-JEPA [2 ] is a reward-free zero-shot RL method based on temporal-difference latent prediction. It learns state and task encoders together with a policy-conditioned multi-step predictor from offline reward-free transitions, encouraging the latent representations to capture long-term policy-dependent dynamics. The learned predictor approximates successor features in the latent space and is used to train a family of latent-conditioned policies. At test time, downstream reward functions are mapped into the learned task space to retrieve corresponding policies without further policy training. Since TD-JEPA also exploits JEPA-style latent prediction for reward-free representation learning, it provides a closely related baseline for evaluating the predictive representation learning mechanism of DT-JEPA.
ODT [35 ] extends DT to the offline-to-online reinforcement learning setting. It first learns a return-conditioned sequence policy from offline trajectories and subsequently adapts the pretrained policy through limited online environment interactions. In contrast to the deterministic action prediction commonly used in the original DT, ODT employs a stochastic policy and incorporates an entropy-regularized objective to encourage exploration during online fine-tuning. We include ODT as a representative sequence-
modeling baseline for evaluating the effectiveness of directly adapting a pretrained DT with online reinforcement learning.
QT [17 ] is a Transformer-based offline RL method that combines conditional sequence modeling with Q-value regularization. QT learns an action-value function using an n n n
Prompt DT [32 ] introduces trajectory prompts into DT to improve generalization across tasks. Instead of relying solely on return-to-go and historical state–action sequences, Prompt DT conditions the policy on short trajectory demonstrations collected from the target task. These demonstrations provide task-specific contextual information that allows the Transformer to infer the desired behavior through in-context sequence modeling. Prompt DT therefore serves as an important comparison for evaluating whether explicit task demonstrations are necessary for cross-task adaptation.
VDT [34 ] augments the DT framework with value-based guidance. While standard Decision Transformer mainly learns actions through supervised sequence prediction conditioned on return-to-go, VDT additionally exploits value information to provide a more informative signal for policy optimization. By incorporating value estimation into sequence-based decision making, VDT aims to alleviate the limitations of purely behavior-cloning-style training and improve the quality of generated actions. We use VDT to compare DT-JEPA with a value-guided variant of DT.
LSDT [29 ] enhances DT by jointly modeling long-term and local dependencies in offline trajectories. Rather than relying on a single temporal representation to capture all historical information, LSDT explicitly considers complementary dependencies at different temporal scales, enabling the policy to better exploit both long-horizon behavioral structure and short-term transition information. This design is particularly relevant for continuous-control problems where current actions may depend simultaneously on recent dynamics and longer-term trajectory context. LSDT is included to assess whether explicitly improving temporal dependency modeling is sufficient to achieve strong generalization.
PDT [31 ] is a closely related reward-free sequence-model pretraining method based on future conditioning. During pretraining, PDT encodes a succeeding trajectory segment into a latent future representation and conditions the policy on this latent variable for action prediction, allowing the model to learn diverse behaviors from reward-free and potentially suboptimal trajectories. A future prior is additionally learned to generate latent future representations from the current state. During downstream adaptation, PDT associates future representations with returns through a return prediction model and uses this information to select desirable futures for policy execution. Since PDT explicitly exploits future trajectory information without requiring reward labels during pretraining, it provides one of the most directly related baselines for evaluating DT-JEPA.
DT-JEPA Training and Deployment Algorithms
In this section, we provide the detailed training algorithms of DT-JEPA. As described in the main paper, the overall workflow of DT-JEPA mainly consists of two phases: reward-free offline pretraining and critic-guided online fine-tuning for unseen downstream tasks.
During the offline pretraining phase, DT-JEPA jointly learns predictive future representations and the future-conditioned decision policy from the offline multi-task trajectory dataset D \mathcal{D} D τ t − K + 1 : t \tau_{t-K+1:t} τ t − K + 1 : t τ t + 1 : t + K − 1 \tau_{t+1:t+K-1} τ t + 1 : t + K − 1 D \mathcal{D} D f θ f_{\theta} f θ H t \mathbf{H}_t H t f θ ′ f_{\theta'} f θ ′ F t \mathbf{F}_t F t H t \mathbf{H}_t H t q η q_{\eta} q η F ^ t \hat{\mathbf{F}}_t F ^ t Algorithm 1 Reward-Free Pretraining of DT-JEPA
Require: Offline multi-task dataset D \mathcal{D} D B B B K K K λ J \lambda_J λ J ρ \rho ρ f θ f_\theta f θ f θ ′ f_{\theta^{\prime}} f θ ′ q η q_\eta q η ϕ \phi ϕ
1: Initialize target encoder: θ ′ ← θ \theta^{\prime} \leftarrow\theta θ ′ ← θ
2: while not converged do
▷ Sample trajectory segments
3: Sample { ( τ t − K + 1 : t ( i ) , τ t : t + K − 1 ( i ) ) } i = 1 B ∼ D \left\{\left(\tau_{t-K+1:t}^{(i)},\tau_{t:t+K-1}^{(i)}\right)\right\}_{i=1}^{B} \sim\mathcal{D} { ( τ t − K + 1 : t ( i ) , τ t : t + K − 1 ( i ) ) } i = 1 B ∼ D
▷ Predict future latent representations
4: for i = 1 , … , B i=1,\ldots,B i = 1 , … , B do
5: H t ( i ) ← f θ ( τ t − K + 1 : t ( i ) ) \mathbf{H}_{t}^{(i)} \leftarrow f_\theta\left(\tau_{t-K+1:t}^{(i)}\right) H t ( i ) ← f θ ( τ t − K + 1 : t ( i ) )
6: F t ( i ) ← sg [ f θ ′ ( τ t : t + K − 1 ( i ) ) ] \mathbf{F}_{t}^{(i)} \leftarrow\operatorname{sg}\left[f_{\theta^{\prime}}\left(\tau_{t:t+K-1}^{(i)}\right)\right] F t ( i ) ← sg [ f θ ′ ( τ t : t + K − 1 ( i ) ) ]
7: q η ( ⋅ ∣ H t ( i ) ) = N ( μ ^ t ( i ) , diag ( ( σ ^ t ( i ) ) 2 ) ) q_\eta\left(\cdot\mid\mathbf{H}_{t}^{(i)}\right)=\mathcal{N}\left(\hat{\boldsymbol{\mu}}_{t}^{(i)},\operatorname{diag}\left((\hat{\boldsymbol{\sigma}}_{t}^{(i)})^2\right)\right) q η ( ⋅ ∣ H t ( i ) ) = N ( μ ^ t ( i ) , diag ( ( σ ^ t ( i ) ) 2 ) )
8: Sample F ^ t ( i ) ∼ q η ( ⋅ ∣ H t ( i ) ) \hat{\mathbf{F}}_{t}^{(i)} \sim q_\eta\left(\cdot\mid\mathbf{H}_{t}^{(i)}\right) F ^ t ( i ) ∼ q η ( ⋅ ∣ H t ( i ) )
9: end for
▷ Compute JEPA prediction loss
10: L J E P A = − 1 B ∑ i = 1 B log q η ( F t ( i ) ∣ H t ( i ) ) \mathcal{L}_{\mathrm{JEPA}}=-\frac{1}{B}\sum_{i=1}^{B}\log q_\eta\left(\mathbf{F}_{t}^{(i)}\mid\mathbf{H}_{t}^{(i)}\right) L JEPA = − B 1 ∑ i = 1 B log q η ( F t ( i ) ∣ H t ( i ) )
▷ Future-conditioned action generation
11: for i = 1 , … , B i=1,\ldots,B i = 1 , … , B do
12: Obtain DT hidden representation h t ( i ) ← D T ϕ ( τ t − K + 1 : t ( i ) ) \mathbf{h}_{t}^{(i)} \leftarrow\mathrm{DT}_{\phi}\left(\tau_{t-K+1:t}^{(i)}\right) h t ( i ) ← DT ϕ ( τ t − K + 1 : t ( i ) )
13: γ t ( i ) ← M L P γ ( F ^ t ( i ) ) , β t ( i ) ← M L P β ( F ^ t ( i ) ) \boldsymbol{\gamma}_{t}^{(i)}\leftarrow\mathrm{MLP}_{\gamma}\left(\hat{\mathbf{F}}_{t}^{(i)}\right),\quad\boldsymbol{\beta}_{t}^{(i)}\leftarrow\mathrm{MLP}_{\beta}\left(\hat{\mathbf{F}}_{t}^{(i)}\right) γ t ( i ) ← MLP γ ( F ^ t ( i ) ) , β t ( i ) ← MLP β ( F ^ t ( i ) )
14: h ~ t ( i ) ← ( 1 + γ t ( i ) ) ⊙ LayerNorm ( h t ( i ) ) + β t ( i ) \tilde{\mathbf{h}}_{t}^{(i)}\leftarrow\left(1+\boldsymbol{\gamma}_{t}^{(i)}\right)\odot\operatorname{LayerNorm}\left(\mathbf{h}_{t}^{(i)}\right)+\boldsymbol{\beta}_{t}^{(i)} h ~ t ( i ) ← ( 1 + γ t ( i ) ) ⊙ LayerNorm ( h t ( i ) ) + β t ( i )
15: Predict actions a ^ m ( i ) ← ϕ ( τ t − K + 1 : t ( i ) , F ^ t ( i ) ) \hat{\mathbf{a}}_{m}^{(i)}\leftarrow\phi\left(\tau_{t-K+1:t}^{(i)},\hat{\mathbf{F}}_{t}^{(i)}\right) a ^ m ( i ) ← ϕ ( τ t − K + 1 : t ( i ) , F ^ t ( i ) )
16: end for
▷ Compute behavior cloning loss
17: L B C = 1 B ∑ i = 1 B ∑ m = t − K + 1 t ∥ a m ( i ) − a ^ m ( i ) ∥ 2 2 \mathcal{L}_{\mathrm{BC}}=\frac{1}{B}\sum_{i=1}^{B}\sum_{m=t-K+1}^{t}\left\|\mathbf{a}_{m}^{(i)}-\hat{\mathbf{a}}_{m}^{(i)}\right\|_{2}^{2} L BC = B 1 ∑ i = 1 B ∑ m = t − K + 1 t a m ( i ) − a ^ m ( i ) 2 2
18: Update ϕ \phi ϕ θ \theta θ η \eta η L B C + λ J L J E P A \mathcal{L}_{\mathrm{BC}}+\lambda_J\mathcal{L}_{\mathrm{JEPA}} L BC + λ J L JEPA
19: Update target encoder via EMA: θ ′ ← ρ θ ′ + ( 1 − ρ ) θ \theta^{\prime}\leftarrow\rho\theta^{\prime}+(1-\rho)\theta θ ′ ← ρ θ ′ + ( 1 − ρ ) θ
20: end while
21: return pretrained DT-JEPA policy ϕ \phi ϕ f θ f_\theta f θ q η q_\eta q η
is subsequently transformed into feature-wise scaling and shifting parameters and injected into the DT through AdaLN, allowing future-predictive information to modulate the historical representation used for action generation. The JEPA module and DT policy are jointly optimized by minimizing the predictive loss L J E P A \mathcal{L}_{\mathrm{JEPA}} L JEPA L B C \mathcal{L}_{\mathrm{BC}} L BC
During the online fine-tuning phase, the pretrained DT-JEPA policy is adapted to an unseen task j j j D j \mathcal{D}^{j} D j D ^ j \hat{\mathcal{D}}^{j} D ^ j Q ω 1 j Q_{\omega_{1}}^{j} Q ω 1 j Q ω 2 j Q_{\omega_{2}}^{j} Q ω 2 j L Q j \mathcal{L}_{Q}^{j} L Q j L B C j \mathcal{L}_{\mathrm{BC}}^{j} L BC j L J E P A j \mathcal{L}_{\mathrm{JEPA}}^{j} L JEPA j L Q j \mathcal{L}_{Q}^{j} L Q j
Algorithm 2 Critic-Guided Online Fine-Tuning of DT-JEPA
Require: Pretrained DT-JEPA policy ϕ \phi ϕ q η q_{\eta} q η λ Q \lambda_Q λ Q λ J \lambda_J λ J γ \gamma γ ρ \rho ρ
Initialize task-specific dataset D j \mathcal{D}^{j} D j
Initialize critics Q ω 1 j Q_{\omega_1}^{j} Q ω 1 j Q ω 2 j Q_{\omega_2}^{j} Q ω 2 j
Initialize target networks: ω ˉ 1 j ← ω 1 j \bar{\omega}_1^{j} \leftarrow\omega_1^{j} ω ˉ 1 j ← ω 1 j ω ˉ 2 j ← ω 2 j \bar{\omega}_2^{j} \leftarrow\omega_2^{j} ω ˉ 2 j ← ω 2 j ϕ ˉ j ← ϕ j \bar{\phi}^{j} \leftarrow\phi^{j} ϕ ˉ j ← ϕ j
while not converged do ▷ Collect task-specific trajectories
Add newly collected trajectories to D j \mathcal{D}^{j} D j
Construct high-return subset D ^ j ⊆ D j \hat{\mathcal{D}}^{j} \subseteq\mathcal{D}^{j} D ^ j ⊆ D j
Sample { τ t − K + 1 : t j , ( i ) } i = 1 B ∼ D ^ j \{\tau_{t-K+1:t}^{j,(i)}\}_{i=1}^{B} \sim\hat{\mathcal{D}}^{j} { τ t − K + 1 : t j , ( i ) } i = 1 B ∼ D ^ j Update task-specific critics
for i = 1 , … , B i = 1,\ldots,B i = 1 , … , B do
Generate target action a ^ t j , ( i ) ∼ ϕ ˉ j ( τ t − K + 1 : t j , ( i ) ) \hat{a}_{t}^{j,(i)} \sim\bar{\phi}^{j}(\tau_{t-K+1:t}^{j,(i)}) a ^ t j , ( i ) ∼ ϕ ˉ j ( τ t − K + 1 : t j , ( i ) )
For m = t − K + 1 , … , t − 1 m = t-K+1,\ldots,t-1 m = t − K + 1 , … , t − 1
Q ^ m j , ( i ) = ∑ k = m t − 1 γ k − m r k j , ( i ) + γ t − m min d ∈ { 1 , 2 } Q ω ˉ d j ( s t j , ( i ) , a ^ t j , ( i ) ) \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}\left(s_{t}^{j,(i)},\hat{a}_{t}^{j,(i)}\right) Q ^ m j , ( i ) = k = m ∑ t − 1 γ k − m r k j , ( i ) + γ t − m d ∈ { 1 , 2 } min Q ω ˉ d j ( s t j , ( i ) , a ^ t j , ( i ) ) end for
for d ∈ { 1 , 2 } d \in\{1,2\} d ∈ { 1 , 2 } do 13.
L c r i t i c j ( ω d j ) = 1 B ∑ i = 1 B ∑ m = t − K + 1 t − 1 ( Q ω d j j ( s m j , ( i ) , a m j , ( i ) ) − Q ^ m j , ( i ) ) 2 \mathcal{L}_{\mathrm{critic}}^{j}(\omega_{d}^{j})=\frac{1}{B}\sum_{i=1}^{B}\sum_{m=t-K+1}^{t-1}\left(Q_{\omega_{d}^{j}}^{j}\left(s_{m}^{j,(i)},a_{m}^{j,(i)}\right)-\hat{Q}_{m}^{j,(i)}\right)^{2} L critic j ( ω d j ) = B 1 i = 1 ∑ B m = t − K + 1 ∑ t − 1 ( Q ω d j j ( s m j , ( i ) , a m j , ( i ) ) − Q ^ m j , ( i ) ) 2 Update ω d j \omega_{d}^{j} ω d j L c r i t i c j ( ω d j ) \mathcal{L}_{\mathrm{critic}}^{j}(\omega_{d}^{j}) L critic j ( ω d j )
end for ▷ Critic-guided policy optimization
Generate DT-JEPA actions a ^ m j , ( i ) \hat{a}_{m}^{j,(i)} a ^ m j , ( i ) ϕ j \phi^{j} ϕ j f θ j f_{\theta}^{j} f θ j p η j p_{\eta}^{j} p η j
L Q j = − 1 B ∑ i = 1 B ∑ m = t − K + 1 t − 1 Q ω 1 j j ( s m j , ( i ) , a ^ m j , ( i ) ) \mathcal{L}_{Q}^{j}=-\frac{1}{B}\sum_{i=1}^{B}\sum_{m=t-K+1}^{t-1}Q_{\omega_{1}^{j}}^{j}\left(s_{m}^{j,(i)},\hat{a}_{m}^{j,(i)}\right) L Q j = − B 1 i = 1 ∑ B m = t − K + 1 ∑ t − 1 Q ω 1 j j ( s m j , ( i ) , a ^ m j , ( i ) ) Update ϕ j \phi^{j} ϕ j θ j j \theta_{j}^{j} θ j j η j j \eta_{j}^{j} η j j
L B C j + λ J L J E P A j + λ Q L Q j \mathcal{L}_{\mathrm{BC}}^{j}+\lambda_{J}\mathcal{L}_{\mathrm{JEPA}}^{j}+\lambda_{Q}\mathcal{L}_{Q}^{j} L BC j + λ J L JEPA j + λ Q L Q j Update target critics via EMA:
ω ˉ d j ← ρ ω ˉ d j + ( 1 − ρ ) ω d j , d ∈ { 1 , 2 } \bar{\omega}_{d}^{j}\leftarrow\rho\bar{\omega}_{d}^{j}+(1-\rho)\omega_{d}^{j},\qquad d\in\{1,2\} ω ˉ d j ← ρ ω ˉ d j + ( 1 − ρ ) ω d j , d ∈ { 1 , 2 } Update target policy via EMA: ϕ ˉ j ← ρ ϕ ˉ j + ( 1 − ρ ) ϕ j \bar{\phi}^{j}\leftarrow\rho\bar{\phi}^{j}+(1-\rho)\phi^{j} ϕ ˉ j ← ρ ϕ ˉ j + ( 1 − ρ ) ϕ j
end while
return adapted DT-JEPA policy ϕ j \phi^{j} ϕ j f θ j f_{\theta}^{j} f θ j q η j q_{\eta}^{j} q η j
Hyperparameters
Table 3 summarizes the common hyperparameters used across all experiments. Unless otherwise specified, the same configuration is used for all environments.
Construction of Unseen Reward Tasks
To evaluate the ability of DT-JEPA to generalize beyond the reward objectives observed during pretraining, we construct a set of unseen downstream tasks by modifying the reward functions of the four MuJoCo locomotion environments: Hopper, HalfCheetah, Walker2d, and Ant. Across these tasks, the state space, action space, transition dynamics, and termination conditions of each
caption: tbl:3 : Common hyperparameters used for DT-JEPA.
environment remain unchanged; only the reward function is modified. Therefore, the resulting tasks share the same underlying dynamics while requiring different behavioral objectives.
For each environment, we consider three downstream objectives: run forward , run backward , and jump . These reward functions are not used during reward-free pretraining and are introduced only during downstream adaptation.
Let v x v_x v x x x x a a a z z z z i n i t z_{\mathrm{init}} z init c c o n t a c t c_{\mathrm{contact}} c contact
HalfCheetah. Following Yu et al. [33 ] , the forward, backward, and jump tasks are defined as
r H C f o r w a r d ( s , a ) = max { v x , 3 } − 0.1 ∥ a ∥ 2 2 , (66) r_{\mathrm{HC}}^{\mathrm{forward}}(s,a) = \max\{v_x,3\} - 0.1\lVert a\rVert_2^2, \tag*{(66)} r HC forward ( s , a ) = max { v x , 3 } − 0.1 ∥ a ∥ 2 2 , (66) r H C b a c k w a r d ( s , a ) = − max { v x , 3 } − 0.1 ∥ a ∥ 2 2 , (67) r_{\mathrm{HC}}^{\mathrm{backward}}(s,a) = -\max\{v_x,3\} - 0.1\lVert a\rVert_2^2, \tag*{(67)} r HC backward ( s , a ) = − max { v x , 3 } − 0.1 ∥ a ∥ 2 2 , (67) r H C j u m p ( s , a ) = − 0.1 ∥ a ∥ 2 2 + 15 ( z − z i n i t ) . (68) r_{\mathrm{HC}}^{\mathrm{jump}}(s,a) = -0.1\lVert a\rVert_2^2 + 15(z-z_{\mathrm{init}}). \tag*{(68)} r HC jump ( s , a ) = − 0.1 ∥ a ∥ 2 2 + 15 ( z − z init ) . (68) The first two objectives encourage locomotion along opposite horizontal directions, whereas the jump objective removes the horizontal velocity reward and instead encourages upward body displacement.
Walker2d. For Walker2d, the corresponding reward functions are
r W f o r w a r d ( s , a ) = v x − 0.001 ∥ a ∥ 2 2 , (69) r_{\mathrm{W}}^{\mathrm{forward}}(s,a) = v_x - 0.001\lVert a\rVert_2^2, \tag*{(69)} r W forward ( s , a ) = v x − 0.001 ∥ a ∥ 2 2 , (69) r W b a c k w a r d ( s , a ) = − v x − 0.001 ∥ a ∥ 2 2 , (70) r_{\mathrm{W}}^{\mathrm{backward}}(s,a) = -v_x - 0.001\lVert a\rVert_2^2, \tag*{(70)} r W backward ( s , a ) = − v x − 0.001 ∥ a ∥ 2 2 , (70) r W j u m p ( s , a ) = − ∣ v x ∣ − 0.001 ∥ a ∥ 2 2 + 10 ( z − z i n i t ) . (71) r_{\mathrm{W}}^{\mathrm{jump}}(s,a) = -\lvert v_x\rvert- 0.001\lVert a\rVert_2^2 + 10(z-z_{\mathrm{init}}). \tag*{(71)} r W jump ( s , a ) = − ∣ v x ∣ − 0.001 ∥ a ∥ 2 2 + 10 ( z − z init ) . (71) For the jump task, the horizontal-velocity penalty discourages the agent from obtaining high reward simply through horizontal locomotion, while the height term encourages vertical movement.
Ant. For the Ant environment, we additionally retain the contact penalty in all three objectives:
r A f o r w a r d ( s , a ) = v x − 0.5 ∥ a ∥ 2 2 − 0.005 c c o n t a c t , (72) r_{\mathrm{A}}^{\mathrm{forward}}(s,a) = v_x - 0.5\lVert a\rVert_2^2 - 0.005c_{\mathrm{contact}}, \tag*{(72)} r A forward ( s , a ) = v x − 0.5 ∥ a ∥ 2 2 − 0.005 c contact , (72) r A b a c k w a r d ( s , a ) = − v x − 0.5 ∥ a ∥ 2 2 − 0.005 c c o n t a c t , (73) r_{\mathrm{A}}^{\mathrm{backward}}(s,a) = -v_x - 0.5\lVert a\rVert_2^2 - 0.005c_{\mathrm{contact}}, \tag*{(73)} r A backward ( s , a ) = − v x − 0.5 ∥ a ∥ 2 2 − 0.005 c contact , (73) r A j u m p ( s , a ) = − ∣ v x ∣ − 0.5 ∥ a ∥ 2 2 − 0.005 c c o n t a c t + 10 ( z − z i n i t ) . (74) r_{\mathrm{A}}^{\mathrm{jump}}(s,a) = -\lvert v_x\rvert- 0.5\lVert a\rVert_2^2 - 0.005c_{\mathrm{contact}} + 10(z-z_{\mathrm{init}}). \tag*{(74)} r A jump ( s , a ) = − ∣ v x ∣ − 0.5 ∥ a ∥ 2 2 − 0.005 c contact + 10 ( z − z init ) . (74) The contact and control penalties regularize the behavior of the high-dimensional quadruped, while the velocity and height terms define the task-specific objective.
Hopper. Hopper tasks following the same design principle as Walker2d, which has a similar planar locomotion setting:
r H f o r w a r d ( s , a ) = v x − 0.001 ∥ a ∥ 2 2 , (75) r_{\mathrm{H}}^{\mathrm{forward}}(s,a) = v_x - 0.001\lVert a\rVert_2^2, \tag*{(75)} r H forward ( s , a ) = v x − 0.001 ∥ a ∥ 2 2 , (75) r H b a c k w a r d ( s , a ) = − v x − 0.001 ∥ a ∥ 2 2 , (76) r_{\mathrm{H}}^{\mathrm{backward}}(s,a) = -v_x - 0.001\lVert a\rVert_2^2, \tag*{(76)} r H backward ( s , a ) = − v x − 0.001 ∥ a ∥ 2 2 , (76) r H j u m p ( s , a ) = − ∣ v x ∣ − 0.001 ∥ a ∥ 2 2 + 10 ( z − z i n i t ) . (77) r_{\mathrm{H}}^{\mathrm{jump}}(s,a) = -\lvert v_x\rvert- 0.001\lVert a\rVert_2^2 + 10(z-z_{\mathrm{init}}). \tag*{(77)} r H jump ( s , a ) = − ∣ v x ∣ − 0.001 ∥ a ∥ 2 2 + 10 ( z − z init ) . (77) Addtional Results
This section reports additional experiments, as well as detailed numerical results for plots in the main part of the paper.
Effect of predicted future representations on the action distributions.
To examine whether the JEPA-predicted future representations effectively participate in decision making, we visualize the action distributions induced by different sampled future latents. We conduct this analysis on the Walker2d task using a DT-JEPA agent pretrained on the corresponding offline dataset, and visualize the action distributions under different future representations sampled from the probabilistic JEPA predictor. Given the same historical context, we sample three future representations from the probabilistic JEPA predictor and use them to condition the DT policy through AdaLN. As shown in Figure 7 , different JEPA latents induce clearly distinguishable action distributions across multiple action dimensions at both timesteps 1000 and 2000. Since the observable history is kept unchanged, the resulting variation can be attributed to the predicted future representation, confirming that F ^ t \hat{F}_t F ^ t
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 p ^ i = N i / ( N Δ a ) \hat{p}_i = N_i/(N\Delta a) p ^ i = N i / ( N Δ a ) N i N_i N i i i i N N N Δ a \Delta a Δ a
More importantly, the differences are structured across action dimensions rather than appearing as uniform perturbations. This suggests that different samples from the predictive distribution encode distinct possible future evolutions, which are translated by AdaLN into different control preferences. In this sense, the probabilistic JEPA predictor does not merely provide an additional stochastic input, but induces a family of future-conditioned action distributions under the same observed history. These observations provide evidence that the predicted latent future serves as a meaningful decision variable linking future-dynamics prediction to action generation in DT-JEPA.
Effect of predicted future representations on the attention score.
To further understand how different future-conditioning strategies affect historical information utilization, we visualize the attention scores of four variants in Figure 8 . Here, DT adopts the conventional RTG token as the conditioning signal; RTG-AdaLN retains RTG but injects it through AdaLN; JEPA-Token replaces RTG with the predicted future representation F ^ t \hat{F}_t F ^ t DT-JEPA combines predictive future representations with AdaLN modulation.
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.
A key observation is that the attention patterns become progressively more concentrated around the main diagonal from DT to DT-JEPA. This phenomenon is particularly meaningful for HalfCheetah, whose dynamics exhibit a strong local Markov structure: the state and action at the current step are primarily determined by temporally nearby transitions, while distant historical tokens are generally less informative for immediate action prediction. Therefore, an attention pattern concentrated around the diagonal indicates that the model assigns larger weights to temporally adjacent and dynamically relevant trajectory information rather than distributing attention broadly over the entire history.
The original DT exhibits a relatively diffuse attention pattern, suggesting that the RTG-conditioned Transformer does not sharply distinguish between locally relevant and less informative historical tokens. Introducing AdaLN with the same RTG condition already increases the locality of attention, as shown by RTG-AdaLN, indicating that feature-wise modulation provides a more direct mechanism for conditioning the self-attention computation. Replacing RTG with the predicted future representation further strengthens the diagonal structure in JEPA-Token. This suggests that F ^ t \hat{F}_t F ^ t
DT-JEPA exhibits the clearest and most concentrated diagonal attention pattern among all variants. This result indicates that predictive future conditioning and AdaLN modulation play complementary roles: the predicted future representation provides information about the likely subsequent dynamics, while AdaLN injects this information directly into the historical feature representations, allowing the Transformer to more selectively emphasize temporally local transitions that are most relevant under the underlying Markov dynamics.
Hyperparameter Sensitivity
Sensitivity to J J J
We further investigate the sensitivity of DT-JEPA to the training-stage weight λ J \lambda_J λ J { 0.001 , 0.01 , 0.1 , 1 , 10 } \{0.001, 0.01, 0.1, 1, 10\} { 0.001 , 0.01 , 0.1 , 1 , 10 } medium datasets. For navigation tasks, AntMaze is evaluated on umaze , while Maze2D includes the umaze , medium , and large configurations.
As shown in Figure 9 , the preferred value of λ J \lambda_J λ J λ J = 0.1 \lambda_J = 0.1 λ J = 0.1 λ J = 1 \lambda_J = 1 λ J = 1 λ J \lambda_J λ J λ J \lambda_J λ J
Figure 9. Sensitivity analysis of the training-stage JEPA loss weight λ J \lambda_J λ J λ J \lambda_J λ J { 0.001 , 0.01 , 0.1 , 1 , 10 } \{0.001, 0.01, 0.1, 1, 10\} { 0.001 , 0.01 , 0.1 , 1 , 10 }
The observed variation also reflects the different roles of predictive representation learning across tasks. A small λ J \lambda_J λ J λ J \lambda_J λ J λ J \lambda_J λ J
Sensitivity to J ′ J' J ′ Q Q Q
We further analyze the sensitivity of DT-JEPA to the two training-stage coefficients λ J \lambda_J λ J λ Q \lambda_Q λ Q Figure 10 , the preferred hyperparameter values vary across environments, indicating that the balance between predictive representation learning and value-guided policy optimization is task-dependent.
Figure 10. Sensitivity analysis of λ J \lambda_J λ J λ Q \lambda_Q λ Q The upper panel varies the weight of the JEPA predictive objective λ J \lambda_J λ J λ Q \lambda_Q λ Q
For λ J ′ \lambda'_J λ J ′ λ J = 0.1 \lambda_J = 0.1 λ J = 0.1 λ J \lambda_J λ J
For λ Q \lambda_Q λ Q λ Q \lambda_Q λ Q λ Q = 0.1 \lambda_Q = 0.1 λ Q = 0.1 λ Q = 0.01 \lambda_Q = 0.01 λ Q = 0.01
Overall, DT-JEPA shows stable performance over a moderate range of both hyperparameters, while the optimal trade-off depends on the environment.
Analysis on Trajectory Stitching Ability
As shown in Table 4 , removing the top high-return trajectories makes the offline datasets substantially more fragmented and reduces the availability of complete high-quality behaviors. Under this setting, DT-JEPA maintains strong performance across most locomotion and navigation tasks, indicating that its policy does not rely solely on imitating complete successful trajectories. Instead, the learned future representations appear to provide a useful predictive structure for connecting locally effective behavior segments that may originate from different trajectories.
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.
This advantage is particularly evident on the long-horizon navigation tasks. DT-JEPA performs consistently well across the AntMaze variants and the more challenging Maze2D configurations, where successful behavior often requires combining useful sub-trajectories over extended horizons. This suggests that predicting future latent dynamics helps the policy reason beyond the immediate local transition and identify behavior fragments that are compatible with a desirable future evolution. In contrast, methods that mainly rely on sequence imitation or fixed future conditioning are generally more affected when coherent high-return trajectories are removed, since the offline data provide weaker direct supervision for complete decision sequences. QT remains highly competitive on several datasets, which is consistent with the benefit of value-based trajectory stitching: its learned Q-function can favor locally high-value actions even when they are observed in different trajectories. DT-JEPA, however, shows a more consistent advantage across the degraded datasets, suggesting that predictive future representations provide a complementary way to recover long-horizon behavioral structure without explicitly relying on value-based action selection.
Sensitivity to the JEPA Predictor Depth
We further investigate the sensitivity of DT-JEPA to the capacity of the latent predictor by varying the number of Transformer blocks L P ∈ { 1 , 2 , 3 , 4 } L_P \in\{1, 2, 3, 4\} L P ∈ { 1 , 2 , 3 , 4 } Table 5 , the predictor depth has a clear influence on downstream performance. Increasing L P L_P L P
Table 5. Sensitivity analysis of the JEPA predictor depth L P L_P L P L P L_P L P bold , while the second-best result is <u>underlined</u>.
The effect of predictor depth also varies noticeably across task families. For the locomotion benchmarks, performance improves relatively smoothly from one to three blocks. Hopper and Ant consistently favor a three-block predictor, whereas HalfCheetah reaches comparable or slightly better performance with only two blocks. This behavior is consistent with the relatively smooth dynamics of HalfCheetah, for which additional predictor capacity brings limited benefit. Walker2d exhibits a slightly different pattern, where four blocks provide marginal improvements over three blocks. However, the differences between L P = 3 L_P = 3 L P = 3 L P = 4 L_P = 4 L P = 4
The dependence on predictor depth becomes more pronounced on the navigation benchmarks. Averaged over AntMaze and Maze2D, increasing the predictor depth from one to three blocks produces a substantially larger improvement than that observed on the locomotion tasks. In particular, AntMaze benefits consistently from increasing L P L_P L P
Sensitivity to the Sequence Length
We further study the sensitivity of DT-JEPA to the sequence length K K K Table 6 , increasing K K K K = 5 K = 5 K = 5 K = 10 K = 10 K = 10 K = 20 K = 20 K = 20
Table 6. Sensitivity analysis of the sequence length K K K
The benefit of increasing K K K K = 20 K = 20 K = 20 K = 30 K = 30 K = 30 K = 30 K = 30 K = 30 K = 20 K = 20 K = 20 K = 40 K = 40 K = 40
Additional Ablation Results
Table 7 further provides a controlled comparison on the same unseen downstream tasks as Table 2 , aiming to determine whether the generalization advantage of DT-JEPA can be mainly attributed to AdaLN-based conditioning or critic-guided adaptation. Compared with RTG-based conditioning, PDT already achieves substantially stronger performance, suggesting that incorporating future-oriented information provides a more informative basis for adapting the policy to reward objectives that are not observed during pretraining.
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 .
Applying AdaLN to PDT leads to a consistent but relatively modest improvement. This trend is consistent with the main ablation results, where replacing conventional RTG conditioning with AdaLN alone only provides limited gains. These observations indicate that AdaLN improves how the conditioning information interacts with the policy, but does not fundamentally determine the quality of the conditioning signal itself. In other words, a more flexible modulation mechanism is beneficial only when the underlying representation contains sufficiently informative task-relevant structure.
A more pronounced improvement is observed when critic-guided adaptation is introduced to PDT. This shows that downstream value information is important for aligning a pretrained future-conditioned policy with an unseen reward objective. Combining critic guidance with AdaLN further improves PDT across the four environments, indicating that the two mechanisms play complementary roles: AdaLN facilitates the utilization of conditioning information inside the policy, whereas critic guidance provides reward-oriented supervision that steers the policy toward desirable behaviors during adaptation.
Nevertheless, a clear performance gap remains between PDT-AdaLN-CG and the full DT-JEPA on all evaluated unseen tasks. More importantly, when considered together with Table 2 , PDT-AdaLN-CG still performs below DT-JEPA without critic guidance. This comparison is particularly informative because the strengthened PDT baseline has already been equipped with both AdaLN and the same critic-guided adaptation mechanism used by DT-JEPA. Therefore, the remaining advantage cannot be explained solely by the modulation architecture or downstream critic supervision. Instead, it points to the importance of the future conditioning representation itself.
The key distinction is that DT-JEPA does not rely only on a generic future embedding as the conditioning variable. It explicitly predicts future latent representations from the observed trajectory history and models their conditional distribution through JEPA learning. As a result, the conditioning signal is directly grounded in the current behavioral context and captures predictive information about how the trajectory may evolve. Such history-dependent future information can provide a more transferable decision prior when the downstream reward objective changes, which helps explain why DT-JEPA retains an advantage even after the alternative baseline is equipped with comparable conditioning and adaptation mechanisms.
