SiMPO: Measure Matching for Online Diffusion Reinforcement Learning

TL;DR

SiMPO introduces a unified measure matching framework for diffusion RL that incorporates negative reweighting, leading to improved policy optimization and performance.

Contribution

It generalizes diffusion RL with a measure matching approach that allows for signed measures and negative reweighting, providing theoretical insights and practical benefits.

Findings

SiMPO outperforms existing methods in empirical evaluations.

Negative reweighting helps avoid suboptimal actions.

The framework offers flexible reweighting schemes tailored to reward landscapes.

Abstract

A commonly used family of RL algorithms for diffusion policies conducts softmax reweighting over the behavior policy, which usually induces an over-greedy policy and fails to leverage feedback from negative samples. In this work, we introduce Signed Measure Policy Optimization (SiMPO), a simple and unified framework that generalizes reweighting scheme in diffusion RL with general monotonic functions. SiMPO revisits diffusion RL via a two-stage measure matching lens. First, we construct a virtual target policy by $f$-divergence regularized policy optimization, where we can relax the non-negativity constraint to allow for a signed target measure. Second, we use this signed measure to guide diffusion or flow models through reweighted matching. This formulation offers two key advantages: a) it generalizes to arbitrary monotonically increasing weighting functions; and b) it provides a principled justification and practical guidance for negative reweighting. Furthermore, we provide geometric interpretations to illustrate how negative reweighting actively repels the policy from suboptimal actions. Extensive empirical evaluations demonstrate that SiMPO achieves superior performance by leveraging these flexible weighting schemes, and we provide practical guidelines for selecting reweighting methods tailored to the reward landscape.