# Context-Former: Stitching via Latent Conditioned Sequence Modeling

arXiv: https://arxiv.org/abs/2401.16452 · HTML: https://mlsanity.com/paper/2401.16452/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Decision transformers and sequence modeling, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 34 of its 50 references that are in the map.
- First posted on arXiv 29 Jan 2024.

## Placed in

- Offline RL and imitation learning from data / [Decision transformers and sequence modeling](https://mlsanity.com/offlinerl/decision-transformers-and-sequence-modeling/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 30 of 34)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processing](https://mlsanity.com/paper/2107.13586/) (2021)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning](https://mlsanity.com/paper/1910.00177/) (2019)
- [Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble](https://mlsanity.com/paper/2110.01548/) (2021)
- [IQ-Learn: Inverse soft-Q Learning for Imitation](https://mlsanity.com/paper/2106.12142/) (2021)
- [Imitation Learning via Off-Policy Distribution Matching](https://mlsanity.com/paper/1912.05032/) (2019)
- [Imitation from Observation: Learning to Imitate Behaviors from Raw Video via Context Translation](https://mlsanity.com/paper/1707.03374/) (2017)
- [Prompting Decision Transformer for Few-Shot Policy Generalization](https://mlsanity.com/paper/2206.13499/) (2022)
- [Generalized Decision Transformer for Offline Hindsight Information Matching](https://mlsanity.com/paper/2111.10364/) (2021)
- [Q-learning Decision Transformer: Leveraging Dynamic Programming for Conditional Sequence Modelling in Offline RL](https://mlsanity.com/paper/2209.03993/) (2022)
- [Recent Advances in Imitation Learning from Observation](https://mlsanity.com/paper/1905.13566/) (2019)
- [Supported Policy Optimization for Offline Reinforcement Learning](https://mlsanity.com/paper/2202.06239/) (2022)
- [Diagnosing Bottlenecks in Deep Q-learning Algorithms](https://mlsanity.com/paper/1902.10250/) (2019)
- [Offline Learning from Demonstrations and Unlabeled Experience](https://mlsanity.com/paper/2011.13885/) (2020)
- [You Can't Count on Luck: Why Decision Transformers and RvS Fail in Stochastic Environments](https://mlsanity.com/paper/2205.15967/) (2022)
- [Scalar reward is not enough: A response to Silver, Singh, Precup and Sutton (2021)](https://mlsanity.com/paper/2112.15422/) (2021)
- [Versatile Offline Imitation from Observations and Examples via Regularized State-Occupancy Matching](https://mlsanity.com/paper/2202.02433/) (2022)
- [Pre-Training to Learn in Context](https://mlsanity.com/paper/2305.09137/) (2023)
- [Beyond Reward: Offline Preference-guided Policy Optimization](https://mlsanity.com/paper/2305.16217/) (2023)
- [DITTO: Offline Imitation Learning with World Models](https://mlsanity.com/paper/2302.03086/) (2023)
- [Prompt-Tuning Decision Transformer with Preference Ranking](https://mlsanity.com/paper/2305.09648/) (2023)
- [Discriminator-Guided Model-Based Offline Imitation Learning](https://mlsanity.com/paper/2207.00244/) (2022)

## Built on by (most cited first, 3 of 3)

- [Generative Trajectory Stitching through Diffusion Composition](https://mlsanity.com/paper/2503.05153/) (2025)
- [TDMPBC: Self-Imitative Reinforcement Learning for Humanoid Robot Control](https://mlsanity.com/paper/2502.17322/) (2025)
- [Imitating from auxiliary imperfect demonstrations via Adversarial Density Weighted Regression](https://mlsanity.com/paper/2405.20351/) (2024)

## Abstract (arXiv)

Offline reinforcement learning (RL) algorithms can learn better decision-making compared to behavior policies by stitching the suboptimal trajectories to derive more optimal ones. Meanwhile, Decision Transformer (DT) abstracts the RL as sequence modeling, showcasing competitive performance on offline RL benchmarks. However, recent studies demonstrate that DT lacks of stitching capacity, thus exploiting stitching capability for DT is vital to further improve its performance. In order to endow stitching capability to DT, we abstract trajectory stitching as expert matching and introduce our approach, ContextFormer, which integrates contextual information-based imitation learning (IL) and sequence modeling to stitch sub-optimal trajectory fragments by emulating the representations of a limited number of expert trajectories. To validate our approach, we conduct experiments from two perspectives: 1) We conduct extensive experiments on D4RL benchmarks under the settings of IL, and experimental results demonstrate ContextFormer can achieve competitive performance in multiple IL settings. 2) More importantly, we conduct a comparison of ContextFormer with various competitive DT variants using identical training datasets. The experimental results unveiled ContextFormer's superiority, as it outperformed all other variants, showcasing its remarkable performance.