OMPO: A Unified Framework for RL under Policy and Dynamics Shifts

TL;DR

OMPO introduces a unified approach for reinforcement learning under policy and dynamics shifts by matching transition occupancy, leading to improved performance across diverse environments and settings.

Contribution

The paper proposes a novel occupancy-matching framework with a tractable min-max formulation and an actor-critic architecture, advancing RL under policy and dynamics shifts.

Findings

OMPO outperforms existing baselines in various environments.

OMPO performs well with domain randomization in robotics.

The method is effective under both stationary and nonstationary dynamics.

Abstract

Training reinforcement learning policies using environment interaction data collected from varying policies or dynamics presents a fundamental challenge. Existing works often overlook the distribution discrepancies induced by policy or dynamics shifts, or rely on specialized algorithms with task priors, thus often resulting in suboptimal policy performances and high learning variances. In this paper, we identify a unified strategy for online RL policy learning under diverse settings of policy and dynamics shifts: transition occupancy matching. In light of this, we introduce a surrogate policy learning objective by considering the transition occupancy discrepancies and then cast it into a tractable min-max optimization problem through dual reformulation. Our method, dubbed Occupancy-Matching Policy Optimization (OMPO), features a specialized actor-critic structure equipped with a distribution discriminator and a small-size local buffer. We conduct extensive experiments based on the OpenAI Gym, Meta-World, and Panda Robots environments, encompassing policy shifts under stationary and nonstationary dynamics, as well as domain adaption. The results demonstrate that OMPO outperforms the specialized baselines from different categories in all settings. We also find that OMPO exhibits particularly strong performance when combined with domain randomization, highlighting its potential in RL-based robotics applications