ReFORM: Reflected Flows for On-support Offline RL via Noise Manipulation

Songyuan Zhang, Oswin So, H M Sabbir Ahmad, Eric Yang Yu, Matthew Cleaveland, Mitchell Black, Chuchu Fan

Keywords: Offline reinforcement learning, support constraint, flow model

TL;DR: Considering offline RL, we propose ReFORM, a two-stage flow policy that realizes the support constraint by construction and avoids the OOD issue without constraining the policy improvement.

Abstract:

Offline reinforcement learning (RL) aims to learn the optimal policy from a fixed
dataset generated by behavior policies without additional environment interactions. One common challenge that arises in this setting is the out-of-distribution
(OOD) error, which occurs when the policy leaves the training distribution. Prior
methods penalize a statistical distance term to keep the policy close to the behavior policy, but this constrains policy improvement and may not completely
prevent OOD actions. Another challenge is that the optimal policy distribution
can be multimodal and difficult to represent. Recent works apply diffusion or
flow policies to address this problem, but it is unclear how to avoid OOD errors
while retaining policy expressiveness. We propose ReFORM, an offline RL method
based on flow policies that enforces the less restrictive support constraint by construction. ReFORM learns a behavior cloning (BC) flow policy with a bounded
source distribution to capture the support of the action distribution, then optimizes
a reflected flow that generates bounded noise for the BC flow while keeping the
support, to maximize the performance. Across 40 challenging tasks from the OGBench benchmark with datasets of varying quality and using a constant set of
hyperparameters for all tasks, ReFORM dominates all baselines with hand-tuned
hyperparameters on the performance profile curves.

Primary Area: reinforcement learning

Submission Number: 8077