# BATS: Best Action Trajectory Stitching

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

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Placed in Offline RL algorithms, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 21 of its 30 references that are in the map.
- First posted on arXiv 26 Apr 2022.

## Placed in

- Offline RL and imitation learning from data / [Offline RL algorithms](https://mlsanity.com/offlinerl/offline-rl-algorithms/)

## Builds on (most cited first, 21 of 21)

- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [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)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction](https://mlsanity.com/paper/1906.00949/) (2019)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [AWAC: Accelerating Online Reinforcement Learning with Offline Datasets](https://mlsanity.com/paper/2006.09359/) (2020)
- [Behavior Regularized Offline Reinforcement Learning](https://mlsanity.com/paper/1911.11361/) (2019)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [MOReL : Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2005.05951/) (2020)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [COMBO: Conservative Offline Model-Based Policy Optimization](https://mlsanity.com/paper/2102.08363/) (2021)
- [Critic Regularized Regression](https://mlsanity.com/paper/2006.15134/) (2020)
- [Neural Episodic Control](https://mlsanity.com/paper/1703.01988/) (2017)
- [Scalable methods for computing state similarity in deterministic Markov Decision Processes](https://mlsanity.com/paper/1911.09291/) (2019)
- [Model-Free Episodic Control](https://mlsanity.com/paper/1606.04460/) (2016)
- [BAIL: Best-Action Imitation Learning for Batch Deep Reinforcement Learning](https://mlsanity.com/paper/1910.12179/) (2019)
- [Generalizable Episodic Memory for Deep Reinforcement Learning](https://mlsanity.com/paper/2103.06469/) (2021)
- [DeepAveragers: Offline Reinforcement Learning by Solving Derived Non-Parametric MDPs](https://mlsanity.com/paper/2010.08891/) (2020)

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

- [Elastic Decision Transformer](https://mlsanity.com/paper/2307.02484/) (2023)
- [Generative Trajectory Stitching through Diffusion Composition](https://mlsanity.com/paper/2503.05153/) (2025)
- [Model-based Trajectory Stitching for Improved Offline Reinforcement Learning](https://mlsanity.com/paper/2211.11603/) (2022)
- [Free from Bellman Completeness: Trajectory Stitching via Model-based Return-conditioned Supervised Learning](https://mlsanity.com/paper/2310.19308/) (2023)
- [Value Memory Graph: A Graph-Structured World Model for Offline Reinforcement Learning](https://mlsanity.com/paper/2206.04384/) (2022)
- [Model-based trajectory stitching for improved behavioural cloning and its applications](https://mlsanity.com/paper/2212.04280/) (2022)
- [Refining Compositional Diffusion for Reliable Long-Horizon Planning](https://mlsanity.com/paper/2605.03075/) (2026)
- [MGDA: Model-based Goal Data Augmentation for Offline Goal-conditioned Weighted Supervised Learning](https://mlsanity.com/paper/2412.11410/) (2024)
- [What Do You Need for Compositional Generalization in Diffusion Planning?](https://mlsanity.com/paper/2505.18083/) (2025)
- [NEEDLEWORK: Offline Rewriting of Robot Data with Verified Local Stitches](https://mlsanity.com/paper/2610.02339/) (2026)

## Abstract (arXiv)

The problem of offline reinforcement learning focuses on learning a good policy from a log of environment interactions. Past efforts for developing algorithms in this area have revolved around introducing constraints to online reinforcement learning algorithms to ensure the actions of the learned policy are constrained to the logged data. In this work, we explore an alternative approach by planning on the fixed dataset directly. Specifically, we introduce an algorithm which forms a tabular Markov Decision Process (MDP) over the logged data by adding new transitions to the dataset. We do this by using learned dynamics models to plan short trajectories between states. Since exact value iteration can be performed on this constructed MDP, it becomes easy to identify which trajectories are advantageous to add to the MDP. Crucially, since most transitions in this MDP come from the logged data, trajectories from the MDP can be rolled out for long periods with confidence. We prove that this property allows one to make upper and lower bounds on the value function up to appropriate distance metrics. Finally, we demonstrate empirically how algorithms that uniformly constrain the learned policy to the entire dataset can result in unwanted behavior, and we show an example in which simply behavior cloning the optimal policy of the MDP created by our algorithm avoids this problem.