Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Rethinking Optimal Transport in Offline Reinforcement Learning

Arip Asadulaev, Rostislav Korst, Alexander Korotin, Vage Egiazarian, Andrey Filchenkov, Evgeny Burnaev

TL;DR

This work proposes a novel algorithm that aims to find a policy that maps states to a distribution of the best expert actions for each given state using optimal transport.

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.

Rethinking Optimal Transport in Offline Reinforcement Learning

TL;DR

This work proposes a novel algorithm that aims to find a policy that maps states to a distribution of the best expert actions for each given state using optimal transport.

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.

Paper Structure

This paper contains 20 sections, 1 theorem, 17 equations, 4 figures, 5 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Optimal Transport
  4. Offline Reinforcement Learning
  5. Rethinking OT in RL
  6. Considering RL as OT problem
  7. Stitching with Partial OT
  8. Toy Experiments
  9. D4RL Experiments
  10. Benchmarks and Baselines
  11. Settings
  12. Results
  13. Parameter $w$
  14. Related work
  15. Conclusion and Discussion
  16. ...and 5 more sections

Key Result

Proposition 9.1

For any policy $\pi$, let's define its performance as $J(\pi) \stackrel{def}{=} {\mathbb{E}_\pi}[\sum_{t=0}^{\infty} \gamma^{t} r(s_{t}, a_{t})]$ over the trajectories obtained by following the policy $\pi$, $(s_{0} \sim S, a_{t} \sim \pi(s_{t}), s_{t+1} \sim P(\cdot \mid s_{t}, a_{t}))$. Let $\beta

Figures (4)

  • Figure 1: Toy experiments. (a) Left point $S_0$ denoting start and $S_T$ is the only rewarded, target location. Black curves visualize behavior trajectories $\beta$. (b) Best behavior policy $\beta^{*}$ according to the data, and the optimal policy $\pi^{*}$ that provides the optimal (shortest path) solution. (c) Comparison of the policy $\beta$ trained by minimizing the objective $-Q^\beta(s,\pi(s))$+BC and our policy $\pi$ trained by the PPL algorithm \ref{['algo:main']}.
  • Figure 2: Exponential moving average (coef. 0.3) curves of the normalized score curves for the Antmaze. Different colors of the curves represent results for $w=3$, $w=5$, $w=8$, $w=12$.
  • Figure 3: Normalized score curves on the Antmaze tasks, $\text{IPL}^{\text{CQL}}$ algorithm is blue, CQL is red
  • Figure :

Theorems & Definitions (1)

  • Proposition 9.1: Policy Improvement with Partial Policy Learning