Source-linked AI summary
Generalized Decision Transformer for Offline Hindsight Information Matching
Hiroki Furuta, Yutaka Matsuo, Shixiang Shane Gu
TL;DR
RL’s sample inefficiency motivates extracting more learning signal from each trajectory, and prior hindsight methods use future information for conditional-policy learning. The paper unifies these methods as HIM and proposes GDT, including CDT and BDT variants. CDT supports offline multi-task state-marginal matching, while BDT performs competitively in offline multi-task imitation learning.
Problem
RL sample inefficiency makes extracting more learning signal from each trajectory a central problem.
Source: §1 · p. 1Method
The paper frames future-conditioned learning as hindsight information matching and generalizes DT with feature-function and aggregator choices that yield CDT and BDT.
Source: §1 · p. 2Results
CDT is reported as the first effective offline multi-task state-marginal matching method, while BDT significantly outperforms simple DT-based imitation variants.
Source: §5.2 · p. 6, §7 · p. 9Takeaways & Limitations
HIM and GDT provide a unified formulation that includes DT, CDT, and BDT and applies to HIM problems with appropriate feature and aggregator choices.
Source: §7 · p. 9The preliminaries use Markovian policy notation for simplicity, though non-Markov policies can be represented with an augmented MDP containing past states.
Source: §3 · p. 3
Abstract
from arXiv · showhide
How to extract as much learning signal from each trajectory data has been a key problem in reinforcement learning (RL), where sample inefficiency has posed serious challenges for practical applications. Recent works have shown that using expressive policy function approximators and conditioning on future trajectory information -- such as future states in hindsight experience replay or returns-to-go in Decision Transformer (DT) -- enables efficient learning of multi-task policies, where at times online RL is fully replaced by offline behavioral cloning, e.g. sequence modeling. We demonstrate that all these approaches are doing hindsight information matching (HIM) -- training policies that can output the rest of trajectory that matches some statistics of future state information. We present Generalized Decision Transformer (GDT) for solving any HIM problem, and show how different choices for the feature function and the anti-causal aggregator not only recover DT as a special case, but also lead to novel Categorical DT (CDT) and Bi-directional DT (BDT) for matching different statistics of the future. For evaluating CDT and BDT, we define offline multi-task state-marginal matching (SMM) and imitation learning (IL) as two generic HIM problems, propose a Wasserstein distance loss as a metric for both, and empirically study them on MuJoCo continuous control benchmarks. CDT, which simply replaces anti-causal summation with anti-causal binning in DT, enables the first effective offline multi-task SMM algorithm that generalizes well to unseen and even synthetic multi-modal state-feature distributions. BDT, which uses an anti-causal second transformer as the aggregator, can learn to model any statistics of the future and outperforms DT variants in offline multi-task IL. Our generalized formulations from HIM and GDT greatly expand the role of powerful sequence modeling architectures in modern RL.
1 INTRODUCTION
RL’s sample inefficiency motivates extracting more supervision from each trajectory, and hindsight methods use future information to accelerate contextual-policy learning. The paper unifies these methods as hindsight information matching and proposes GDT, which yields new algorithms through small changes to DT.
1 INTRODUCTION
- Hindsight information matching (HIM) unifies future-conditioned methods that train policies to generate trajectories satisfying specified properties, including distributional properties.Behavioral cloning can serve as the learning objective.Source: §1 · pp. 1–2
- Generalized Decision Transformer (GDT) applies transformer-based hindsight behavioral cloning to future-information matching with only small architectural changes to DT.The feature function and anti-causal aggregator determine which information statistic is matched.Source: §1 · p. 2, Figure 1 · p. 2
- The paper introduces offline multi-task state-marginal matching and imitation learning, with a Wasserstein-distance evaluation metric for both problems.It develops CDT for state-marginal matching and BDT for imitation learning.Source: §1 · p. 2
- CDT replaces DT’s anti-causal summation with binning, while BDT uses a second transformer to model full-trajectory imitation.The resulting choices target offline multi-task state-marginal matching and imitation learning, respectively.Source: Figure 1 · p. 2, §1 · p. 2
2 RELATED WORK
Prior work uses hindsight information in off-policy RL, behavioral cloning, and sequence modeling, while state-marginal matching offers a distinct distribution-matching objective. The paper situates its methods at the intersection of these lines of work.
Hindsight Reinforcement Learning and Behavior Cloning
- Hindsight methods relabel data for parameterized rewards, and rich future-state examples have supported generalizable policies learned with behavioral cloning.Transformers have also produced strong offline RL and goal-based RL results [Chen et al., 2021a; Janner et al., 2021].Source: §2 · pp. 2–3
Distributional Reinforcement Learning and State-Marginal Matching
- State-marginal matching minimizes divergence from a target state distribution, unlike categorical distributional RL, which optimizes a non-conditional reward-maximizing policy.Many state-marginal matching methods require on-policy samples, while the paper presents CDT as an offline multi-task benchmarked method.Source: §2 · p. 3
- State-marginal matching is also connected to inverse RL and feature-expectation matching, although many such approaches are online methods.The cited work notes an exception in Kostrikov et al. (2020).Source: §2 · p. 3
RL and Imitation Learning as Sequence Modeling
- Transformer adoption in RL has lagged other fields partly because of challenges applying transformers over time in online RL.Recent work has applied sequence models to one-shot imitation and offline learning.Source: §2 · p. 3
3 PRELIMINARIES
The preliminaries define the RL objective, state-marginal matching, and parameterized rewards as foundations for the paper’s offline multi-task problems. They also note how the policy notation extends to non-Markov policies.
3 PRELIMINARIES
- The preliminaries define an MDP through its state and action spaces, transition function, initial-state distribution, reward, and discount factor.The policy notation is Markovian for simplicity, but can represent non-Markov policies by augmenting the MDP with a context window of past states.Source: §3 · p. 3
3.1 STATE MARGINAL MATCHING
- State-marginal matching seeks a policy whose state marginal minimizes divergence from a target distribution, rather than maximizing task reward.The target can represent uniform exploration, scripted behaviors, or expert data; the objective depends on dynamics and the target distribution.Source: §3.1 · p. 4
3.2 PARAMETERIZED RL OBJECTIVES
- Parameterized RL learns a policy conditioned on reward parameters, optimizing expected reward across parameter values weighted by their distribution.These conditioned policies are multi-task RL problems within contextual MDPs.Source: §3.2 · p. 4
4 HINDSIGHT INFORMATION MATCHING
Hindsight information matching treats future-derived trajectory statistics as contexts for learning policies that reproduce those statistics. Its choices of features and information summaries organize existing approaches and support new variants.
Hindsight algorithms
- Hindsight algorithms use future trajectory information to mine positive examples for contextual policies, which can accelerate RL or support behavioral cloning.Examples include HER, TDM, LfP, hindsight multi-task RL, and return-conditioned RL.Source: §4 · p. 4
Information statistics
- Information statistics can capture properties of a partial trajectory, including distributional summaries such as means, variances, and higher-order moments.A feature function Φ maps state-action pairs into features used to define those statistics.Source: §4 · p. 4
HIM problem types
- Existing HIM methods are grouped into four problem types according to IΦ(τ), with CDT adding a distribution-based category and BDT adapting DT to full-trajectory imitation.The table describes a coarse classification; CDT uses a histogram statistic, while BDT targets trajectory imitation.Source: Table 1 · p. 5, §4 · p. 5
Hindsight information matching
- Hindsight information matching uses trajectory-derived information statistics as policy contexts, making the corresponding trajectory examples optimal for those contexts.Behavioral cloning can train on these context-trajectory pairs.Source: §4 · p. 5
HIM algorithm choices
- Changing IΦ(τ) can yield new HIM algorithms, while behavioral cloning applies across the listed problem types because it avoids optimizing the policy against every information-matching objective.The paper contrasts this broad applicability with RL optimization, which is non-trivial for some return-, trajectory-, and distribution-based objectives.Source: §4 · p. 5
5 GENERALIZED DECISION TRANSFORMER
GDT generalizes Decision Transformer by varying the feature function and anti-causal aggregator for hindsight information matching. The section defines offline multi-task SMM and IL, a Wasserstein evaluation metric, and CDT and BDT variants.
5 GENERALIZED DECISION TRANSFORMER
- GDT generalizes DT by varying the feature function and anti-causal aggregator, with behavioral cloning enabling training across HIM problems.The paper uses transformers as its architecture and studies CDT and BDT variants.Source: §5 · pp. 5–6
5.1 TASK DEFINITIONS AND METRICS
- Offline multi-task SMM uses a specified feature function, while IL uses identity or learned features and is treated as SMM on full state.Both tasks are evaluated with a Wasserstein loss between state-marginal and target distributions, estimated by feature-space binning in practice.Source: §5.1 · p. 6
5.2 CATEGORICAL DECISION TRANSFORMER FOR DISTRIBUTION MATCHING
- CDT discretizes low-dimensional feature spaces into categorical approximations and uses recursive Bellman-like computation to form future-information targets.The paper presents it as an effective algorithm for offline multi-task SMM.Source: §5.2 · p. 6
- When Φ is unavailable, DT variants learn it through auto-encoding, contrastive, or end-to-end action-prediction losses for full-state SMM or offline multi-task IL.The authors argue that matching learned feature means can suffice to match distributions to arbitrary precision when Φ learns sufficient state features.Source: §5.3 · p. 6
5.4 BI-DIRECTIONAL DECISION TRANSFORMER FOR ONE-SHOT IMITATION LEARNING
- BDT uses an anti-causal transformer over reverse-order states and achieves substantially better offline multi-task IL results than DT-X variants.It assumes identity Φ and learns representations within the aggregator.Source: §5.4 · p. 7
6 EXPERIMENTS
Across MuJoCo offline experiments, CDT matches state-feature targets, including unseen and synthesized distributions, while BDT performs strongly on full-state imitation using learned features. The experiments evaluate these methods across reward and state-feature matching, diverse unseen targets, and full-state distribution matching.
6 EXPERIMENTS
- Experiments use D4RL medium-expert data across four MuJoCo tasks, compare against DT, Meta-BC, and FOCAL, and measure Wasserstein-1 feature-distribution distance.The evaluation holds out ten trajectories for testing and averages results over 20 rollouts per four random seeds.Source: §6 · p. 7
6.1 REWARD AND STATE-FEATURE MATCHING
- CDT matches unseen x-velocity target distributions better than baselines, and also outperforms them on 2D Ant xy-velocity matching.The study additionally tests shifted and synthesized novel target distributions.Source: Table 2 · p. 7, §6.1 · p. 7
6.2 GENERALIZATION TO UNSEEN TARGET DISTRIBUTION
- Single-task offline datasets may not reveal generalization to diverse behaviors, so the study tests synthetic bimodal z-velocity and meta-learning cheetah-velocity targets.The authors note that these datasets are collected by single-task reward-maximization policies.Source: §6.2 · p. 8
6.2.1 SYNTHESIZING UNSEEN BI-MODAL DISTRIBUTION
- CDT matches uni-modal and synthesized bi-modal z-velocity distributions better than DT and FOCAL, and is comparable to Meta-BC.The Figure 2(a) target includes running forward followed by backflipping in one rollout.Source: Figure 2 · p. 8, Table 3 · p. 8, §6.2.1 · p. 8
6.2.2 DIVERSE UNSEEN DISTRIBUTION FROM META LEARNING TASK
- For held-out cheetah velocities, CDT outperforms DT and FOCAL and is slightly better than Meta-BC on offline multi-task SMM.The held-out target velocities are 0.5, 1.5, and 2.5 within the 0.0–3.0 training range.Source: Table 4 · p. 9, §6.2.2 · p. 9
6.3 ONE-SHOT DISTRIBUTION MATCHING IN FULL STATE
- BDT outperforms learned-feature variants and Meta-BC on x-velocity matching, and is comparable to CDT or DT with longer input contexts.The reported Table 5 results use m = 16; the full results also vary m across 1, 4, and 16.Source: Table 5 · p. 9, §6.3 · p. 9
7 CONCLUSION
The conclusion presents HIM as a unifying formulation and GDT as a generalization of DT that includes CDT and BDT. The accompanying implementation and baseline details describe their hyperparameter reference and compare Meta-BC with FOCAL.
7 CONCLUSION
- GDT unifies HIM algorithms and includes DT, CDT, and BDT as special cases; CDT enables offline state-marginal matching, while BDT outperforms simple DT-based imitation variants.The paper proposes HIM and GDT as frameworks for applying sequence modeling to a broader range of RL problems.Source: §7 · p. 9
A HYPER-PARAMETER OF GENERALIZED DECISION TRANSFORMERS
- The implementation builds CDT and BDT on Chen et al. (2021a)'s official codebase and follows most of its hyperparameters.Table 6 lists hyperparameters for DT, CDT, and BDT.Source: A HYPER-PARAMETER OF GENERALIZED… · p. 17, Table 6 · p. 17
B DETAILS OF BASELINES
- Meta-BC performs better than FOCAL in the experiments, which the authors attribute to RL methods struggling with the distribution-matching objective and FOCAL optimizing task reward instead.The authors also note that BC often converges faster than offline RL methods.Source: B DETAILS OF BASELINES · p. 17
C DETAILS OF EVALUATION
The evaluation section motivates Wasserstein-1 distance as a quantitative way to assess distribution matching when only samples are available, then describes its empirical estimation. It also distinguishes this no-adaptation, distribution-distance evaluation from offline meta-RL and meta-imitation settings.
C DETAILS OF EVALUATION
- Prior SMM evaluations relied on qualitative rollout-density visualizations, while task-reward evaluation is unsuitable for the paper’s suboptimal or multi-task trajectories.The authors note that reward-based evaluation can be used when expert trajectories are available, but that condition does not fit these experiments.Source: C DETAILS OF EVALUATION · p. 18
- Wasserstein-1 distance between empirical categorical feature distributions provides a shared evaluation metric for offline multi-task SMM and IL when their densities are unavailable.Features are discretized into bins, and evaluation averages over test trajectories; task reward is unsuitable for the paper’s suboptimal or multi-task trajectories.Source: C DETAILS OF EVALUATION · p. 18
- The empirical metric discretizes feature space into B bins per dimension, with range predefined from offline datasets, and compares categorical distributions from target trajectories and policy rollouts.Policy rollouts use 4 seeds × 20 rollouts; the metric is averaged over test-set trajectories.Source: C DETAILS OF EVALUATION · p. 18
D CONNECTION TO META LEARNING
- Unlike offline meta-RL and meta-imitation settings that use test-time demonstrations for adaptation, this work assumes no adaptation at test time and evaluates feature-distribution distance rather than task reward.The comparison is summarized across the problem settings in Table 8.Source: D CONNECTION TO META LEARNING · p. 19, Table 8 · p. 19
E DETAILS OF EXPERIMENTS
The supplementary experiments describe the feature distributions used and show that trajectory quality corresponds to different distribution shapes.
E DETAILS OF EXPERIMENTS
- Sorting trajectories by cumulative reward yields distinct feature-distribution shapes for best, middle, and worst quality, while the full dataset appears to be a weighted mixture.Figure 3 visualizes reward and x-velocity distributions in the D4RL medium-expert datasets.Source: E DETAILS OF EXPERIMENTS · p. 20, Figure 3 · p. 20
E.1 QUANTITATIVE AND QUALITATIVE RESULTS OF REWARD AND STATE-FEATURE MATCHING
In reward-distribution matching, Categorical DT (CDT) generally outperforms DT on offline multi-task SMM. CDT rollout distributions also closely match held-out target distributions across the illustrated environments and trajectory qualities.
E.1 QUANTITATIVE AND QUALITATIVE RESULTS OF REWARD AND STATE-FEATURE MATCHING
- CDT achieves better reward-distribution matching than DT on offline multi-task SMM, measured by Wasserstein-1 distance.The comparison evaluates Wasserstein-1 distance between discretized rollout and target feature distributions; Figure 4 shows reward and x-velocity matching for best and middle held-out trajectories.Source: Table 9 · p. 20, §E.1 · p. 20, Figure 4 · p. 20
E.2 EVALUATION ON TASK PERFORMANCE
The supplementary experiments report that DT performs better than CDT on task rewards, whereas CDT handles two-dimensional feature-distribution matching and expert-to-medium distribution shifts more successfully. They also detail the construction of synthetic bimodal and diverse unseen-distribution evaluations.
E.2 EVALUATION ON TASK PERFORMANCE
- DT is consistently better than CDT on task-reward evaluations and achieves performance similar to held-out trajectories.Table 10 evaluates task rewards while conditioning on held-out trajectories.Source: Table 10 · p. 21, §E.2 · p. 21
E.3 2D STATE-FEATURE DISTRIBUTION MATCHING
- In Ant-v3 xy-velocity matching, CDT performs better on the joint Wasserstein-1 metric and handles expert-to-medium target shifts, while DT fits the medium distribution even for expert targets.CDT bins each feature dimension separately; DT instead sums each velocity over trajectories and normalizes by maximum horizon. Figure 5 illustrates expert and medium target cases.Source: §E.3 · p. 21, Table 11 · p. 21, Figure 5 · p. 21
E.4 DETAILS OF SYNTHESIZING UNSEEN BI-MODAL DISTRIBUTION
- The synthetic bimodal evaluation combines backflipping and forward-running trajectories, then creates patchworked test trajectories by joining 500-step segments from each behavior.The backflip behavior is induced by adding an absolute z-velocity term to the HalfCheetah reward; the dataset combines 500 collected trajectories with 500 D4RL expert trajectories.Source: §E.4 · p. 22
E.5 DETAILS OF DIVERSE UNSEEN DISTRIBUTION FROM META LEARNING TASK
- The evaluation prepared 31 target x-velocities from 0.0 to 3.0 at 0.1 intervals; 10 trajectories at 0.5, 1.5, and 2.5 were held out, with the rest used for training.The cheetah reward specified target velocity, and the horizon was 200 steps.Source: §E.5 · p. 22
E.6 QUANTITATIVE AND QUALITATIVE RESULTS OF ONE-SHOT DISTRIBUTION MATCHING IN FULL STATE
The supplemental results compare BDT on one-shot distribution matching, test CDT on shifted and synthetic targets, and describe categorical conditioning and learned information-statistics encoders. BDT performs consistently well on offline multi-task matching, while CDT handles supported shifts and tends to match even synthetic bimodal targets.
E.6 QUANTITATIVE AND QUALITATIVE RESULTS OF ONE-SHOT DISTRIBUTION MATCHING IN FULL STATE
- BDT is consistently better than other strategies and Meta-BC for offline multi-task SMM, and is comparable to CDT or DT with context length N=50.The experiments vary embedding size m and, for BDT, context windows N=20, 50, and 100; BDT improves with larger embeddings.Source: §E.6 · p. 22, Table 12 · p. 23, Table 13 · p. 23
E.7 SHIFTING THE TARGET DISTRIBUTION
- CDT handles shifted reward and state-feature targets better than DT, especially in hopper, when targets remain within the dataset distribution's support.The evaluation shifts reward and velocity values by bin_size × {−3.0, −2.0, −1.0, 0.0, +1.0, +2.0, +3.0}.Source: Table 14 · p. 24, Table 15 · p. 24, §E.7 · p. 24
E.8 SYNTHESIZING UNREALISTIC TARGET DISTRIBUTION
- CDT tends to match manually synthesized Gaussian state-feature targets, including bimodal cases, although results are worse than for realistic test targets.The synthetic targets are not constrained to be physically realizable.Source: Table 16 · p. 26, §E.8 · p. 26, Figure 9 · p. 26
F DETAILS FOR CATEGORICAL DECISION TRANSFORMER
- CDT discretizes features into B bins and recursively computes categorical future-feature distributions backward before providing the resulting statistics to the transformer.For multidimensional features, the procedure is applied per dimension, approximating the joint with a product of independent marginals.Source: F DETAILS FOR CATEGORICAL DECISION… · p. 27
G DETAILS OF DECISION TRANSFORMER WITH LEARNED Φ
- Decision Transformer variants can use auto-encoder or contrastive objectives to learn information statistics, with encoder training performed unsupervised, jointly, or by pretraining and freezing.The contrastive setup uses a query encoder and momentum encoder; the paper notes that more advanced temporally extended objectives remain future work.Source: G DETAILS OF DECISION TRANSFORMER WITH… · p. 27