Rethinking Optimal Transport in Offline Reinforcement Learning

TL;DR

This paper introduces a novel offline reinforcement learning algorithm based on optimal transport, which stitches together the best behaviors from diverse datasets to improve policy performance in continuous control tasks.

Contribution

It redefines offline RL as an optimal transport problem and proposes a new algorithm to extract policies that map states to distributions of optimal expert actions.

Findings

Demonstrates improved performance over existing methods on D4RL benchmarks.

Effectively extracts policies that focus on the best behaviors in sub-optimal datasets.

Validates the approach on continuous control problems with empirical results.

Abstract

We propose a novel algorithm for offline reinforcement learning using optimal transport. Typically, in offline reinforcement learning, the data is provided by various experts and some of them can be sub-optimal. To extract an efficient policy, it is necessary to \emph{stitch} the best behaviors from the dataset. To address this problem, we rethink offline reinforcement learning as an optimal transportation problem. And based on this, we present an algorithm that aims to find a policy that maps states to a \emph{partial} distribution of the best expert actions for each given state. We evaluate the performance of our algorithm on continuous control problems from the D4RL suite and demonstrate improvements over existing methods.