Future-Conditioned Unsupervised Pretraining

• Future-conditioned unsupervised pretraining is a technique that leverages entire future sub-trajectories as a supervisory signal, enabling action prediction without explicit reward labels.
• It employs instantiations like the Pretrained Decision Transformer (PDT) and Self-Predictive Goal-Conditioned Pretraining (SGI) to use latent future embeddings for robust representation learning.
• Empirical results demonstrate that this approach outperforms traditional return-conditioned methods, yielding efficient fine-tuning and improved performance on suboptimal or unlabeled datasets.

Future-conditioned unsupervised pretraining is a methodology in reinforcement learning (RL) and sequence modeling that leverages the information contained in entire future sub-trajectories, rather than single scalar returns, to structure action prediction and representation learning in the absence of explicit reward labels. This framework is especially influential in settings characterized by reward-free offline data, where standard return-conditioned paradigms become brittle or inapplicable. The two leading instantiations are the Pretrained Decision Transformer (PDT) (Xie et al., 2023) and Self-Predictive Goal-Conditioned Pretraining (SGI) (Schwarzer et al., 2021), each demonstrating that exposures to future trajectory segments during pretraining provides a rich, task-agnostic supervisory signal that accelerates, and in many cases surpasses, supervised pretraining for RL.

1. Motivation and Core Concepts

Return-conditioned supervised learning—such as Decision Transformer (DT) or its stochastic variant ODT—frames offline RL as sequence modeling, feeding the history of states, current observation, and a scalar "return-to-go" into a transformer and maximizing the likelihood of observed actions. However, this approach demands fully reward-labeled datasets and struggles with unlabeled, suboptimal, or reward-free data sources, which are increasingly central in large-scale RL, imitation, or video domains (Xie et al., 2023).

Future conditioning addresses these limitations by embedding the entire future sub-trajectory following a given state into a latent representation, conditioning action prediction on this latent both during pretraining (unsupervised) and downstream finetuning (possibly with reward). This approach generalizes the conditioning context, enabling richer coverage of possible behaviors and enhancing generalization and controllability. In representation learning for RL, future prediction and goal-sampling similarly supply auxiliary signals that support learning versatile feature spaces (Schwarzer et al., 2021).

2. Formalism and Model Architectures

Pretrained Decision Transformer (PDT)

Consider an MDP $(\mathcal{S},\mathcal{A},P,r)$ and an offline dataset $\mathcal{D} = \{\tau^{(m)}\}_{m=1}^M$ of trajectories $\tau = (s_0, a_0, s_1, a_1, ..., s_T)$. During pretraining, no reward labels are assumed. For each time $t$, a history segment of length $K$, $\tau_{t:t+K-1}$, and its immediately following future $\tau_{t+K:t+2K-1}$, are extracted.

A future encoder $g_\theta$ embeds the future segment into a Gaussian latent: $$z \sim g_\theta(\cdot \mid \tau_{t+K:t+2K-1}), \quad z \in \mathbb{R}^d,$$ and a future prior $p_\theta(z \mid s_t)$ is trained to predict likely $\mathcal{D} = \{\tau^{(m)}\}_{m=1}^M$0's from current state alone. The GPT-style transformer receives the sequence

$\mathcal{D} = \{\tau^{(m)}\}_{m=1}^M$1

as input, with each state $\mathcal{D} = \{\tau^{(m)}\}_{m=1}^M$2 and action $\mathcal{D} = \{\tau^{(m)}\}_{m=1}^M$3 embedded and a single latent $\mathcal{D} = \{\tau^{(m)}\}_{m=1}^M$4 appended. All self-attention layers access the future embedding, ensuring joint conditioning on history and future. Action prediction is then

$\mathcal{D} = \{\tau^{(m)}\}_{m=1}^M$5

maximized via a behavior-cloning objective with entropy regularization (Xie et al., 2023).

Self-Predictive Goal-Conditioned Pretraining (SGI)

SGI operates by pretraining a convolutional state encoder $\mathcal{D} = \{\tau^{(m)}\}_{m=1}^M$6 to map states to latent vectors $\mathcal{D} = \{\tau^{(m)}\}_{m=1}^M$7 and a latent dynamics model $\mathcal{D} = \{\tau^{(m)}\}_{m=1}^M$8 to predict multi-step future latents. Synthetic goals—sampled from future states within or across trajectories, possibly with noise or permutation—provide additional context for unsupervised goal-conditioned Q-learning. The model ensemble leverages self-predictive representations (SPR) (Schwarzer et al., 2021).

3. Objective Functions, Training Procedure, and Regularization

PDT Joint Objective

During unsupervised pretraining, the total loss comprises a behavior-cloning (BC) term and a future-regularization term: $\mathcal{D} = \{\tau^{(m)}\}_{m=1}^M$9 where: $\tau = (s_0, a_0, s_1, a_1, ..., s_T)$0 and

$\tau = (s_0, a_0, s_1, a_1, ..., s_T)$1

The pretraining leverages the reparametrization trick for latent sampling. The regularization coefficient $\tau = (s_0, a_0, s_1, a_1, ..., s_T)$2 directly modulates the trade-off between behavior diversity (small $\tau = (s_0, a_0, s_1, a_1, ..., s_T)$3) and consistency (large $\tau = (s_0, a_0, s_1, a_1, ..., s_T)$4), where each dataset admits its own optimal regime (Xie et al., 2023).

SGI Pretraining Loss

SGI combines three unsupervised losses: $\tau = (s_0, a_0, s_1, a_1, ..., s_T)$5 where $\tau = (s_0, a_0, s_1, a_1, ..., s_T)$6 enforces multi-step future prediction in latent space, $\tau = (s_0, a_0, s_1, a_1, ..., s_T)$7 (inverse modeling) prevents representation collapse, and $\tau = (s_0, a_0, s_1, a_1, ..., s_T)$8 enables goal-conditioned Q-learning with synthetic goals (Schwarzer et al., 2021).

4. Fine-tuning and Downstream Control

After unsupervised pretraining, RL agents are typically fine-tuned on relatively small, on-task labeled datasets. In PDT, transitions are collected online, and a return predictor $\tau = (s_0, a_0, s_1, a_1, ..., s_T)$9 is trained to align latents $t$0 with high-return outcomes. Bayes’ rule is used to steer the latent prior towards rewarding futures: $t$1 Controllable sampling of $t$2’s based on predicted return enables the selection of high-return behavior modes during execution. Finetuning minimizes a composite loss including $t$3, $t$4, and $t$5 (Xie et al., 2023).

In SGI, the pretrained encoder and auxiliary models are incorporated into a Rainbow DQN agent, with learning rates for pretrained modules considerably reduced to preserve prior-learned features. Finetuning is conducted under strict data constraints, such as 100K environment steps on the Atari-100K benchmark (Schwarzer et al., 2021).

5. Empirical Evaluation and Comparative Performance

Empirical studies demonstrate that future-conditioned pretraining delivers strong or superior outcomes, particularly under challenging data regimes.

PDT Evaluation

Benchmarks on Gym MuJoCo tasks (D4RL, medium and medium-replay) show:

Environment	SAC	ACL	ODT-0	ODT Finetuned	PDT-0	PDT Finetuned
hopper-med-replay	24 ± 10	52 ± 49	74	75 ± 6	28	85 ± 5
walker2d-med-replay	–	–	–	70 ± 3	–	59 ± 15

On suboptimal datasets, PDT shows consistent gains over ODT after finetuning (hopper-medium-replay: 85 ± 5 for PDT vs. 75 ± 6 for ODT), with ablative experiments confirming the criticality of future conditioning and the regularization trade-off (Xie et al., 2023).

Generalization studies show that PDT's reward-agnostic latent space enables rapid adaptation to new reward functions, attaining, after 200K steps, total scores of ≈ 268 (PDT) vs. ≈ 101 (ODT) across several jump and forward-jump benchmarks.

SGI Evaluation

Ablations on 26 Atari games indicate that multi-step future prediction (SPR), inverse modeling, and goal RL each contribute to final performance, with the full SGI combination achieving the best median human-normalized score (0.679 versus 0.343 with no pretraining). Larger models pretrained with future-conditioned objectives significantly outperform smaller networks post-pretraining, highlighting the value of abundant and varied future state examples (Schwarzer et al., 2021).

6. Analysis, Behavioral Properties, and Limitations

Future-conditioned mandatory variables such as $t$6 in PDT can be manipulated to generate diverse behaviors. Sampling distinct latents at the initial state leads to widely varying action histograms, demonstrating substantial control over behavioral modes. The return predictor $t$7 provides effective ordinal ranking of latents; sampling from higher predicted-return percentiles correlates monotonically with realized return in downstream rollouts.

Reward-agnostic pretraining delivers strong transfer: policies pretrained with no knowledge of specific objectives rapidly align to new or altered reward specifications via minimal additional learning—new mapping $t$8. However, the additional computational requirements and dataset-specific tuning of future regularization, as well as the open problem of optimal future-latent coding (possible extensions to VQ-VAEs, normalizing flows, or diffusion-based priors), present ongoing challenges (Xie et al., 2023).

In SGI, ablation studies highlight the necessity of both future prediction and inverse modeling to avoid representation collapse and ensure data-efficient downstream learning (Schwarzer et al., 2021).

7. Relationship to Broader Research and Implications

The future-conditioned paradigm stands at the intersection of unsupervised representation learning, sequence modeling, and control. Distinct from standard return-conditioning, conditioning on rich, trajectory-scale futures exposes the learner to the full spectrum of possible behaviors, enforces diversity in learned features, and empowers sample-efficient fine-tuning or reward-agnostic transfer. As highlighted in both PDT and SGI, models explicitly trained to predict, encode, or reach flexible future contexts can outperform methods trained purely with direct or reward-provided signals when operating on real-world, noisy, or unlabeled data. This suggests direct future conditioning is a promising direction for scaling RL to unannotated, heterogeneous datasets and environments (Xie et al., 2023, Schwarzer et al., 2021).