Score-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport

TL;DR

Diffusion policies excel at modeling multi-modal actions but struggle with distribution shifts when used in imitation learning alone. The paper introduces OTPR, an Optimal Transport-guided Score-based diffusion Policy for Reinforcement Learning fine-tuning, which uses the $Q$-function as a transport cost and treats the policy as an OT map to integrate demonstrations with online RL. It adds Masked Optimal Transport guided by expert keypoints and a compatibility-based resampling scheme to stabilize training, and derives a dual OT formulation to obtain a soft state-action coupling. Empirical results on three simulation tasks show OTPR achieving superior accuracy and robustness, especially in sparse-reward or long-horizon settings, with code released at the provided URL.

Abstract

Diffusion policies have shown promise in learning complex behaviors from demonstrations, particularly for tasks requiring precise control and long-term planning. However, they face challenges in robustness when encountering distribution shifts. This paper explores improving diffusion-based imitation learning models through online interactions with the environment. We propose OTPR (Optimal Transport-guided score-based diffusion Policy for Reinforcement learning fine-tuning), a novel method that integrates diffusion policies with RL using optimal transport theory. OTPR leverages the Q-function as a transport cost and views the policy as an optimal transport map, enabling efficient and stable fine-tuning. Moreover, we introduce masked optimal transport to guide state-action matching using expert keypoints and a compatibility-based resampling strategy to enhance training stability. Experiments on three simulation tasks demonstrate OTPR's superior performance and robustness compared to existing methods, especially in complex and sparse-reward environments. In sum, OTPR provides an effective framework for combining IL and RL, achieving versatile and reliable policy learning. The code will be released at https://github.com/Sunmmyy/OTPR.git.

Figures (4)

• Figure 1: Overview of OTPR. (a) Estimation Optimal Transport Plan: A stochastic dual approach with two parametrized dual variables is introduced to estimate the optimal transport plan with $Q$-cost from state distribution and action distribution. (b) Training: OTPR pre-trains a diffusion model from the expert’s data. It then iteratively performs RL to optimize a $Q$-function and trains diffusion models by score matching. (c) Online Interaction: In the inference step, the policy makes action inference by iteratively denoising a random noise, conditioned on current state. OTPR then employs the compatibility function $H$ to reweight each action before resampling one.
• Figure 2: Learning curves of online fine-tuning with various methods. Observe that OTPR largely always dominates or attains similar performance to the next best method. Other methods for fine-tuning diffusion policies (IDQL, DQL, DPPO) are a bit unstable, and perform substantially worse.
• Figure 3: (left) Comparison between OTPR with different guidance (H, Q and A). (right) Comparison between OTPR with (OTPR-M) and without (OTPR-U) the expert data mask.
• Figure 4: This example demonstrates a clear and concise visualization of a Q-value matrix alongside its corresponding estimated optimal transport plan.

Theorems & Definitions (9)

• Proposition 4.1
• Proposition 4.2
• Theorem 4.3
• Proposition 2.1
• proof
• Proposition 2.2
• proof
• Theorem 2.3
• proof