Distributionally-Constrained Policy Optimization via Unbalanced Optimal Transport

TL;DR

This paper introduces a novel constrained policy optimization method in reinforcement learning using unbalanced optimal transport, enabling effective handling of distribution constraints with an actor-critic algorithm.

Contribution

It formulates constrained policy optimization as unbalanced optimal transport and develops a general RL objective optimized via Dykstra's algorithm, including an actor-critic implementation.

Findings

Effective handling of distribution constraints in RL

Demonstrated success on various applications

Robust performance with sampling-based implementation

Abstract

We consider constrained policy optimization in Reinforcement Learning, where the constraints are in form of marginals on state visitations and global action executions. Given these distributions, we formulate policy optimization as unbalanced optimal transport over the space of occupancy measures. We propose a general purpose RL objective based on Bregman divergence and optimize it using Dykstra's algorithm. The approach admits an actor-critic algorithm for when the state or action space is large, and only samples from the marginals are available. We discuss applications of our approach and provide demonstrations to show the effectiveness of our algorithm.