Composite Flow Matching for Reinforcement Learning with Shifted-Dynamics Data

TL;DR

CompFlow tackles reinforcement learning with shifted-dynamics offline data by introducing a composite flow that reuses knowledge from the offline transition model to better approximate the online dynamics. The method leverages a Wasserstein-distance based estimation of the dynamics gap via optimal-transport flow matching, enabling principled gap measurement and an optimistic data-collection strategy that prioritizes high-gap regions. The authors provide theoretical guarantees showing reduced generalization error and improved performance bounds, and demonstrate strong empirical gains on Gym-MuJoCo benchmarks with various dynamics shifts and in wildlife-conservation simulations. This approach enhances sample efficiency and robustness when offline data comes from a different dynamical regime, with practical implications for real-world RL where online interaction is costly or risky.

Abstract

Incorporating pre-collected offline data from a source environment can significantly improve the sample efficiency of reinforcement learning (RL), but this benefit is often challenged by discrepancies between the transition dynamics of the source and target environments. Existing methods typically address this issue by penalizing or filtering out source transitions in high dynamics-gap regions. However, their estimation of the dynamics gap often relies on KL divergence or mutual information, which can be ill-defined when the source and target dynamics have disjoint support. To overcome these limitations, we propose CompFlow, a method grounded in the theoretical connection between flow matching and optimal transport. Specifically, we model the target dynamics as a conditional flow built upon the output distribution of the source-domain flow, rather than learning it directly from a Gaussian prior. This composite structure offers two key advantages: (1) improved generalization for learning target dynamics, and (2) a principled estimation of the dynamics gap via the Wasserstein distance between source and target transitions. Leveraging our principled estimation of the dynamics gap, we further introduce an optimistic active data collection strategy that prioritizes exploration in regions of high dynamics gap, and theoretically prove that it reduces the performance disparity with the optimal policy. Empirically, CompFlow outperforms strong baselines across several RL benchmarks with shifted dynamics.

Figures (6)

• Figure 1: Comparison between direct and composite flow matching. Composite flow first transports from a Gaussian latent variable to the offline transition distribution, then adapts to the online distribution via optimal transport flow matching.
• Figure 2: Overall Framework of CompFlow. To estimate the dynamics gap, we propose composite flow matching, which computes the Wasserstein distance between offline and online transition dynamics. Guided by the estimated dynamics gap, we augment policy training with offline transitions that exhibit low discrepancy from the online dynamics, and incorporate a behavior cloning objective to stabilize learning. To enhance data diversity and facilitate adaptation, we further encourage exploration in regions with high dynamics gap.
• Figure 3: Comparison of MSE between direct flow and composite flow.
• Figure 4: Comparison of return under different data selection ratios across tasks.
• Figure 5: Comparison of return under different exploration strengths across tasks.

Theorems & Definitions (22)

• Definition 2.1: Online Policy Learning with Shifted-Dynamics Offline Data
• Lemma 2.1: Return Bound between Two Environments lyu2024cross
• Theorem 3.1: Conditions for Composite Flow Yielding Smaller Errors
• Remark 3.2
• Proposition 3.3: Informal; shared-latent coupling is $W_2$-optimal
• Theorem 3.4: Large Dynamics Gap Exploration Reduces Performance Gap
• Lemma F.0: Return Bound between Two Environments lyu2024cross
• proof
• Lemma F.1: Extended telescoping lemma
• Lemma F.2