Source-linked AI summary
From Suboptimal to Optimal: One-Step Latent Flow Policies for Offline Reinforcement Learning
Shenghua Lu, Zixuan Yang, Rongxu Zhao, Bohao Qu, Jifeng Hu, Hechang Chen
TL;DR
Offline RL must learn effective policies from fixed datasets, while existing generative policies can require costly multi-step sampling and offer no mechanism to lift suboptimal behavior. MF2B learns shared, value-guided mean-flow mappings in a high-level latent decision space and completes refinement in one step. Across nine D4RL Gym-MuJoCo tasks, it achieves an average score of 90.91, exceeding the cited flow-policy baselines.
Problem
Offline RL must learn policies from fixed datasets, where suboptimal behavior is present, while generative policies can require costly multi-step sampling and lack a mechanism to lift suboptimal behaviors toward high-value regions.
Source: §3.1 · p. 3Method
MF2B encodes trajectory chunks as high-level latent decisions and uses a shared conditional mean velocity field for one-step flows from Gaussian-induced sources to positive behavior and from suboptimal decisions to high-value regions.
Source: §1 · p. 2Results
On nine D4RL Gym-MuJoCo tasks, MF2B achieves an average score of 90.91, surpassing FlowQ (85.6), SSCP (87.9), and FQL (79.2).
Source: §8 · p. 11Takeaways & Limitations
The high-level representation compresses sequence length from T to K ≈ T/H, reducing computational overhead, while suboptimal-to-high-value transport uses low-quality trajectories for policy improvement.
Source: §8 · p. 11
Abstract
from arXiv · showhide
Offline reinforcement learning aims to learn effective policies from fixed datasets without environment interaction. Although generative models such as diffusion and flow matching enhance policy expressiveness, they typically start from uninformative Gaussian noise and generate actions through multi-step iteration, leading to high computational costs and failing to effectively leverage suboptimal trajectories in datasets. We propose MeanFlow2Better (MF2B), which reformulates offline RL as generative decision chunk optimization from current performance status to better performance status in a high-level latent decision space. MF2B maps trajectory chunks to high-level latent decision representations through an encoder, placing generation starting points within a structured decision space sculpted by behavioral data. The model learns a shared conditional mean velocity field in this space, constructing two flow mappings: from Gaussian-Induced Decision Sources toward positive behavioral decision regions, and from suboptimal decisions toward high-value decision regions. Both flow mappings complete decision refinement through one-step ODE integration over the unit time interval $\Delta t=1$, requiring neither multi-step iterative sampling nor backpropagation through time. Furthermore, we propose hierarchical chunk-level advantage to guide the flow direction toward high-return regions, where the high-level decisions cover multiple low-level actions, leading to sequence length compression from $T$ to $K \approx T/H$ and significantly lower computational overhead. On the D4RL Gym-MuJoCo benchmark, MF2B achieves an average score of 90.91, matching or surpassing strong baselines.
1 Introduction
Offline RL must learn from fixed datasets despite complex behavior distributions and abundant suboptimal trajectories. MF2B uses one-step, high-level latent decision flows to refine decisions, and reports an average score of 90.91 across nine D4RL Gym-MuJoCo tasks.
1 Introduction
- Offline RL faces complex, multimodal behavior distributions and underuses suboptimal trajectories, which make up most practical datasets.Online interaction may be expensive, unsafe, or infeasible in settings such as robotics, healthcare, and autonomous driving.Source: §1 · p. 1
- MF2B treats suboptimal trajectories as starting points for improvement, using Gaussian-Induced Decision Sources as a semantic and value null-space origin.The source representations are sculpted into task-relevant decision representations.Source: §1 · p. 1
- MF2B reformulates offline RL as generative decision-chunk optimization in a high-level latent space, using a shared conditional mean velocity field and hierarchical chunk-level advantage guidance.Its two branches transport Gaussian-Induced Decision Sources toward positive behavioral regions and suboptimal decisions toward high-value regions.Source: §1 · p. 2
- MF2B completes decision refinement with one Euler step over the unit interval, rather than multi-step ODE integration.The method predicts conditional mean velocity for Δt = 1.Source: §1 · p. 2
- 90.91 average score across nine D4RL Gym-MuJoCo tasks establishes state-of-the-art performance among one-step flow policies.The method description also reports that this surpasses all one-step flow policy baselines.Source: §1 · p. 2
2 Related Work
Prior offline RL work addresses complex action distributions with diffusion and flow policies, but commonly relies on multi-step generation or simple Gaussian priors. MF2B distinguishes itself through latent decision representations and shared, value-guided dual-branch mean-flow mappings.
2 Related Work
- Earlier offline RL methods often used unimodal Gaussian policies, while diffusion methods such as DQL, IDQL, and EDP model multimodal actions but retain costly multi-step denoising.The cited conservative methods include CQL, TD3+BC, and IQL [Kostrikov et al., 2021].Source: §2 · p. 2
- Flow policies such as FQL and FlowQ use flow matching for offline RL, but typically rely on simple Gaussian priors and multi-step sampling.MeanFlowQL and SSCP are described as approaches to one-shot generation.Source: §2 · p. 2
- MF2B encodes trajectory chunks as high-level latent decisions, creating short-range, approximately linear flow mappings instead of modeling policies in raw action space.The paper contrasts this with raw-action-space flow policies [Park et al., 2025, Alles et al., 2025].Source: §2 · p. 2
- MF2B shares a conditional mean velocity field across a source-to-positive baseline branch and a value-guided suboptimal-to-high-value branch.The latter uses hierarchical chunk-level advantage weighting; the paper describes the combined approach as mean flow matching with value-guided dual-branch mapping in latent decision space.Source: §2 · p. 3
3 Preliminaries
Offline RL learns from fixed datasets without further environment interaction, often containing suboptimal behavior. The preliminaries introduce IQL value estimation and flow matching as foundations for the method.
3.1 Offline Reinforcement Learning
- Offline RL learns an optimal policy solely from a fixed transition dataset, which may contain suboptimal behavior and provides no further exploration.The setting excludes environment interaction during learning.Source: §3.1 · p. 3
3.2 Implicit Q-Learning
- IQL uses expectile regression to estimate state values as a high expectile of behavior values, providing value guidance without explicit policy improvement steps.Its asymmetric loss distinguishes it from fitting Q-value means.Source: §3.2 · p. 3
3.3 Flow Matching
- Flow matching learns a time-varying velocity field whose ODE transports samples from a simple prior distribution to a target data distribution.The preliminaries describe training by regressing target vector fields that define desired probability paths.Source: §3.3 · p. 4
- Mean velocity fields model average velocity between two times, enabling one-step generation from a prior without iterative ODE integration.As the interval between times approaches zero, mean velocity converges to marginal velocity.Source: §3.3 · p. 4
4 MeanFlow2Better Method
MF2B improves offline-RL decisions by flowing trajectory chunks through a compact latent decision space. A shared conditional mean-flow field maps value-neutral or suboptimal sources toward positive or higher-value decisions, which are decoded into action sequences.
4 MeanFlow2Better Method
- MF2B encodes trajectory chunks as high-level decisions, learns conditional mean velocities across two branches, and decodes refined decisions into action sequences.A chunk-level advantage critic supplies value guidance during training only; the flow interval is fixed to [0, 1].Source: §4 · p. 4
4.1 Chunk-Level High-Level Latent Decision Representation
- Chunk encoding compresses a T-step trajectory into K ≈ T/H decisions, separating high-level behavioral intent from low-level action execution.The decoder conditions action reconstruction on states and the latent decision for each chunk.Source: §4.1 · p. 5
4.2.1 Gaussian-Induced Decision Source
- Gaussian-Induced Decision Sources encode Gaussian trajectories after their reward, return-to-go, chunk return-to-go, and total return are set to zero.This places the sources at the zero point of the value dimension.Source: §4.2.1 · p. 5
4.3.1 Source-to-Positive Alignment Mean-Flow Branch
- The Source-to-Positive Alignment branch transports Gaussian-Induced Decision Sources toward positive behavioral decision regions.Its source is value-neutral and is not constrained to a specific return interval in the dataset.Source: §4.2.1 · p. 5
4.2.2 Suboptimal-to-High-Value Trajectory Pairing
- The Quality-Uplift branch learns to transport encoded suboptimal decisions toward encoded high-value decisions drawn from return-partitioned trajectory pools.The high-value timestep coverage ratio controls the timestep budget of the high-value pool.Source: §4.2.2 · p. 5
4.3 Dual-Branch Conditional Latent Decision Mean Flow Mapping
- Both branches share a conditional latent mean-flow field but differ in their source, target, and conditioning.The local condition uses the target trajectory’s return-to-go, while the global condition is a return gap.Source: §4.3 · p. 5
4.3.2 Quality-Uplift Mean-Flow Branch
- The Quality-Uplift branch transports suboptimal latent decisions toward high-value decisions, then decodes the refined decisions into action sequences.This branch instantiates the hypothesis that learnable improvement directions exist between suboptimal and high-value behaviors in latent decision space.Source: §4.3.2 · p. 6
4.3.3 High-Level Decision Transformer Architecture
- A bidirectional Transformer attends across the full sequence of K ≈ T/H chunk tokens to predict globally coherent mean velocities.The flow field takes concatenated latent decision representations as input tokens.Source: §4.3.3 · p. 6
4.4 Hierarchical Chunk-Level Latent Decision Value Estimation and Advantage Guidance
A hierarchical chunk-level advantage critic provides value guidance while limiting indiscriminate fitting to dataset behavior frequency. Its weighting preserves quality preference without allowing extreme chunks to dominate flow supervision.
4.4 Hierarchical Chunk-Level Latent Decision Value Estimation and Advantage Guidance
- The critic evaluates chunk-level transitions rather than primitive actions or full trajectories and is used only during training.It follows IQL and estimates value with an expectile baseline of behavioral decision values.Source: §4.4 · p. 6
- The Quality-Uplift branch interpolates advantage weights with uniform weights to preserve quality preference while preventing extreme chunks from dominating supervision.The Source-to-Positive branch uses uniform weighting.Source: §4.4 · p. 6
- Low-value chunks remain source endpoints for Quality-Uplift flow, with the critic adjusting their training contribution rather than discarding them.Source: §4.4 · p. 6
4.5 Joint Training Objective
The joint objective trains action reconstruction and flow prediction while preventing flow-target construction from sending gradients into the encoder. Inference refines a source decision in one step before decoding a chunk of actions.
4.5 Joint Training Objective
- Stop-gradient on both endpoints prevents flow-supervision targets from drifting jointly with encoder parameters.The encoder instead receives end-to-end gradients through action reconstruction and the decoding path from the refined decision.Source: §4.5 · p. 6
- With t0=0 and t1=1, the algorithm partitions offline data by return into Itop and Ilow, samples low-return, top-return, Gaussian-noise, and dataset trajectories, encodes them, and computes hierarchical chunk-level IQL advantages.Source: §4.5 · p. 7
- At inference, the policy encodes a source chunk, refines its latent decision with one mean-flow step, decodes H actions, and executes them.The environment buffer is updated after each chunk of actions.Source: §4.5 · p. 7
4.6 One-Step High-Level Decision Refinement and Inference
At test time, refinement starts from a source trajectory and uses a single forward pass to predict the mean velocity across the full interval, avoiding iterative ODE integration.
4.6 One-Step High-Level Decision Refinement and Inference
- A single forward pass predicts the full mean velocity over [t0, t1], transporting an existing suboptimal behavior toward high-value regions without iterative ODE integration.The method learns directional transport in structured high-level decision space rather than generating trajectories from scratch.Source: §4.6 · p. 7
5 Experiments
The experiments evaluate MF2B on nine D4RL Gym-MuJoCo locomotion tasks, using normalized returns averaged across seeds and comparing against regression, diffusion, and flow policies.
5.1 Environmental Setup
- MF2B is evaluated on nine Gym-MuJoCo locomotion tasks spanning HalfCheetah, Hopper, and Walker2d under medium, medium-replay, and medium-expert datasets.The comparison is presented in Table 1.Source: §5.1 · p. 7, Table 1 · p. 8
5.2 Evaluation Metrics
- Evaluation reports D4RL normalized scores averaged over five training seeds, with each model tested for 50 episodes per task.Scores map random policy performance to 0 and domain-specific expert performance to 100.Source: §5.2 · p. 8
5.3 Baseline Models
- MF2B is compared with regression-based policies, diffusion policies, and one-step flow policies across representative offline RL baselines.The listed methods include BC, IQL, CQL, Diffuser, DQL, FQL, and FlowQ.Source: §5.3 · p. 8
6 Experimental Results
On Gym9, MF2B scores 90.91 on average, surpassing all listed flow policies and most listed diffusion policies.
6.1.1 Gym-MuJoCo Locomotion Task Comprehensive Comparison
- 90.91 average score on Gym9 surpasses all listed flow policies and most listed diffusion policies.Flow-policy scores are FQL 79.2, FlowQ 85.6, and SSCP 87.9; diffusion-policy scores include DQL 87.9 and TCD 85.7.Source: §6.1.1 · p. 8
7 Ablation Studies
The ablations show that each tested mechanism contributes to MF2B’s Gym-MuJoCo performance, while latent-space analyses and sensitivity experiments examine the flow’s geometry and configuration.
7.1 Module Ablation Experiments
- The ablations separately test mean-velocity transport, chunk-level value guidance, directional pool pairing, and source-to-positive alignment.Each variant changes one mechanism relative to the complete model.Source: §7.1 · pp. 8–9
7.1.1 Module Ablation Experiment Results
- MF2B averages 90.91 and outperforms all four variants: DLM −4.47, NVG −4.07, UP −5.71, and NAF −5.28.The variants remove or replace one mechanism while keeping other settings identical.Source: §7.1 · p. 8, Table 2 · p. 9, §7.1.1 · p. 9
7.1.2 High-Level Latent Decision Space Distribution Analysis
- Zpos aligns with Ztop, Zneg separates from behavioral representations, and Zlow overlaps Ztop, supporting complementary structure formation and within-space quality uplift.The analyses concern high-level latent representations, not the original action space.Source: §7.1.2 · pp. 9–10, Figure 1 · p. 10
7.2.1 Chunk Length and High-Level Context
- On three medium-replay tasks, H=10 and Kmax=100 perform best among tested chunk configurations at approximately equal primitive-step coverage.Both finer H=5 and coarser H=20 granularities slightly reduce performance.Source: §7.2.1 · p. 10, Table 3 · p. 10, §7.2 · p. 10
- With H fixed at 10, performance degrades sharply at Kmax=25 and Kmax=50; Kmax=100 covers the full 1000-step horizon.The results indicate a need for sufficient high-level context to capture cross-stage task structure.Source: §7.2.1 · p. 10, Table 4 · p. 10
7.2.2 Trajectory Coverage Ratio
- Coverage ratios 0.1 and 0.5 perform close to the default 0.2, while 0.8 causes sharp drops on Hopper and Walker2d.The paper attributes the drops to reduced directional contrast from an excessively large high-value pool.Source: §7.2.2 · p. 11, Table 5 · p. 11
8 Conclusion
MF2B refines offline-RL decisions through one-step flow mappings in a chunk-level latent space, and averages 90.91 on nine D4RL Gym-MuJoCo tasks. Its learned decision space remains constrained by the diversity of the offline dataset.
8 Conclusion
- A shared conditional mean velocity field maps Gaussian-Induced Decision Sources toward positive behavioral regions and suboptimal decisions toward high-value regions, guided by hierarchical chunk-level advantage.Source: §8 · p. 11
- One-step Euler integration over ∆t = 1 avoids iterative sampling and backpropagation through time, while chunk-level latents compress sequence length from T to K ≈ T/H.Source: §8 · p. 11
- 90.91 average score on 9 D4RL Gym-MuJoCo tasks surpasses FlowQ (85.6) and SSCP (87.9), and significantly outperforms FQL (79.2).Source: §8 · p. 11
- The latent decision space is constrained by offline-dataset diversity, so excessively narrow behavioral support may limit its quality.Source: §8 · p. 11