BATS: Best Action Trajectory Stitching

TL;DR

This paper introduces BATS, an offline reinforcement learning algorithm that constructs an augmented dataset with new trajectories using learned dynamics, enabling effective planning and policy improvement without constraints on logged data.

Contribution

BATS proposes a novel dataset augmentation method for offline RL by adding trajectories through learned models, facilitating planning and value estimation directly on the constructed MDP.

Findings

BATS can effectively plan trajectories using learned dynamics.

The method provides bounds on value functions based on dataset properties.

Behavior cloning the optimal policy of the constructed MDP avoids unwanted behaviors.

Abstract

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.