Learning Multimodal One-step Flow Policy via Value-weighted Optimal Transport

• Distribution shift makes value estimates unreliable when a policy selects actions poorly supported by fixed offline data, motivating support-aware learning.Online exploration may also be costly or unsafe in robotics and control.Source: §1 · p. 1
• Multimodal datasets can defeat unimodal policies, which may collapse distinct valid actions into a suboptimal or physically invalid mean.Flow policies can model multimodal actions, and deterministic probability paths make one-step learning a natural efficiency target.Source: §1 · p. 1
• One-step flow learning is challenging because pointwise distillation can collapse behaviors, while direct critic maximization can exploit local optima or overestimated out-of-distribution actions.Source: §1 · pp. 1–2
• OptiFlow couples value-aware reference-flow and one-step-policy samples through entropic optimal transport, using critic values for target priority and action distance for compatible pairings.Source: §1 · p. 2
• OptiFlow avoids direct critic maximization and anchors one-step learning to high-value reference actions while allowing flexible noise-to-action mappings that preserve multimodal structure.The paper reports multimodal capture in controlled diagnostics and strong offline RL benchmark performance.Source: §1 · p. 2

• Diffusion- and flow-based policies model multimodal offline actions [15]; FQL uses critic-guided improvement, while GFP combines value-aware behavior cloning with one-step distillation.Source: §2 · p. 2
• Offline RL balances policy improvement with support preservation through conservative values [3, 4], behavior regularization [5], or value-weighted extraction [22]; OptiFlow instead shapes transport supervision with critic values.Source: §2 · p. 2
• Retrieval constraints and optimal transport methods [25] [26] keep policies near data or align distributions, whereas OptiFlow uses state-wise action-sample coupling to produce value-weighted, geometry-aware regression targets.Source: §2 · p. 2

3.1 Offline Reinforcement Learning

• 3.1 Offline Reinforcement Learning: The setting uses a discounted MDP and a fixed dataset from potentially diverse, unknown behavior policies, without additional environment interaction.OptiFlow learns a critic from the dataset and uses its values in a transport problem for updating the one-step policy.Source: §3.1 · pp. 2–3

3.2 Flow Policies

• 3.2 Flow Policies: Flow matching learns a state-conditioned vector field from noise toward dataset actions; inference integrates the field, while a one-step policy directly maps noise to actions.The multi-step reference models multimodal behavior, and transport-guided updates train the efficient one-step policy.Source: §3.2 · p. 3

3.3 Entropic Optimal Transport

• 3.3 Entropic Optimal Transport: Entropy regularization smooths the coupling between discrete distributions and enables efficient computation with Sinkhorn matrix scaling.The regularized transport framework is later applied state-wise to align one-step and multimodal reference policies.Source: §3.3 · p. 3

• OptiFlow separates policy improvement into a dataset-anchored reference flow, a value-aware geometric coupling, and transport-guided target selection for the one-step policy.This structure avoids direct critic maximization, which can pull the policy toward local optima or out-of-distribution regions.Source: §4 · pp. 4, 6
• The critic supplies reference-policy training priorities and the reference-side transport marginal; value-aware behavior cloning makes the multi-step reference a data-supported proposal concentrated on high-return regions.The proposal restricts sample allocation to candidate actions described as both in-distribution and high-performing.Source: §4 · pp. 4–5
• The state-wise coupling assigns sampled one-step actions to reference actions, and maximum assignment extracts explicit regression targets from the soft transport plan.Source: §4 · p. 5
• OptiFlow uses a uniform marginal for one-step samples and a critic-value-weighted marginal for reference actions, while the transport cost encodes only geometric action proximity.This decouples value guidance from geometry, assigning more mass to high-return actions without putting value terms in the cost.Source: §4 · p. 5
• Transport-guided distillation routes one-step samples to geometrically aligned, distinct high-value modes, avoiding pointwise matching that can collapse multimodal behavior into a mean.The one-step policy is updated against fixed reference targets with the transport weights treated as constants.Source: §4 · p. 6

5.1 Controlled Multimodal Experiments

• The bandit diagnostic shows OptiFlow preserving multiple high-value modes within behavior support, unlike value maximization that can drift out of support.FQL-D under-prioritizes high-value modes, while FQL-Q can deviate from behavior support under value overestimation.Source: §5.1 · pp. 6–7, Figure 2 · p. 7
• Under asymmetric and bimodal value landscapes, value-weighted distillation collapses toward a dominant mode, while OptiFlow’s value-weighted reference marginal preserves multimodal structure.A uniform reference marginal with value placed in the transport cost misallocates mass because low-value actions receive equal total mass.Source: §5.1 · p. 7, Figure 3 · p. 7
• In a four-goal MDP with uneven dataset coverage, OptiFlow-BC maintains balanced coverage of the equally optimal modes and attains the highest return, while FQL trades diversity against quality.OptiFlow-BC retains the value-weighted OT coupling and transport-guided extraction but uses a standard BC reference policy.Source: §5.1 · p. 8, Figure 4 · p. 8

5.2 Benchmark Experimental Setup

• The evaluation covers reward-based OGBench navigation and manipulation, D4RL AntMaze and Adroit, and conventional and expressive generative-policy baselines.State-based OGBench tasks use 1M gradient steps and other listed task groups use 500K; the method samples 16 one-step and 64 reference actions per state and uses 30 Sinkhorn iterations. Table 1 aggregates OGBench families and D4RL domains.Source: §5.2 · pp. 8–9, Table 1 · p. 9

5.3 Main Results

• 48.3±0.4 is OptiFlow’s state-based OGBench average, versus 45.6 ± 0.4 for GFP and 39.0 ± 0.4 for FQL; its manipulation average is 44.5 ± 0.5.The manipulation average compares with 40.8 ± 0.5 for GFP and 36.3 ± 0.4 for FQL.Source: §5.3 · p. 9, Table 1 · p. 9
• Five of six OGBench manipulation families favor OptiFlow, while it improves over FQL on AntMaze-Giant, Humanoid-Medium, and Humanoid-Large.The five manipulation families are Cube-Single, Cube-Double, Cube-Triple, Puzzle-4×4, and Scene-Play.Source: §5.3 · p. 9
• On D4RL, OptiFlow’s AntMaze average exceeds GFP by 4.0 points and FQL by 4.3 points, while it remains competitive with flow-based baselines on Adroit.Source: §5.3 · p. 10

6 Conclusion

• OptiFlow couples one-step-policy and reference-policy samples through entropic optimal transport, using critic values to weight supervision and action distance to assign geometrically compatible pairs.This provides a structured alternative to independent noise-wise distillation and direct critic maximization.Source: §6 · p. 10
• OptiFlow preserves multimodal action structure in controlled diagnostics and achieves strong offline-RL benchmark performance, with especially clear gains on OGBench manipulation tasks.Source: §6 · p. 10
• The paper proposes using critic information to shape structured policy supervision rather than relying only on direct critic maximization.It frames this as a path toward value-guided generative policies that preserve multimodal behavior while retaining efficient one-step deployment.Source: §6 · p. 10

A Discussion

• The discussion examines why value belongs in the reference marginal, how global transport matching relates to multimodality, and how Sinkhorn iterations and policy samples affect computational cost.OptiFlow adds training-time computation from repeated sampling and Sinkhorn-based transport updates, though these matrix computations can be parallelized on modern GPUs.Source: §6 · p. 10, A Discussion · p. 14

A.1 Limitations

• Training requires state-wise entropic optimal transport over sampled actions and Sinkhorn iterations, making it more expensive than pointwise distillation or direct actor updates.After training, inference uses only the one-step flow policy and remains as efficient as other one-step policies.Source: §A.1 · p. 14
• OptiFlow may offer smaller advantages when the reference-action distribution is simple or a single local target already suffices for policy extraction.Its strongest benefit is when the reference flow policy supplies multiple useful candidate actions and supervision mass can be allocated advantageously.Source: §A.1 · p. 14

A.2 Why use a value-aware reference marginal in optimal transport?

• In a simple 2 × 2 example with a uniform reference marginal, incorporating value information solely through transport cost cannot guarantee policy improvement, regardless of cost design.The uniform-marginal value is 5.5.Source: §A.2 · p. 15
• For any feasible coupling, expected reference-action value is determined by the reference-side marginal, not the transport cost; with a uniform marginal, it is the average sampled value.This result extends beyond the 2 × 2 example to an arbitrary number of sampled actions.Source: §A.2 · p. 16
• As τ decreases, the value-weighted marginal’s expected assigned value rises from the candidate average toward the best candidate, while softmax avoids concentrating all supervision on one action.The authors describe τ as trading assigned value against mode coverage.Source: §A.2 · p. 16
• Replacing OptiFlow’s value-weighted marginal with a uniform marginal reduces average performance from 57.3 to 9.7, versus 57.3 to 48.6 when VaBC is replaced with standard BC.Table 2 compares full OptiFlow with each variant under the same evaluation protocol as Table 9.Source: Table 2 · p. 16, §A.2 · p. 16

A.3 Alternative distillation methods and why OptiFlow?

• In the equal-value bimodal example, OptiFlow’s global marginal preserves both modes, while left-initialized VDWM and Top-2 nearest-neighbor matching collapse to the left mode.The global coupling allocates mass across reference actions rather than selecting each target independently.Source: Figure 5 · p. 17, §A.3 · p. 18
• Value-weighted distillation weights reference actions by value, but naive ℓ2 distillation can collapse multimodal flow outputs to an intermediate mean action.The bandit example shows VWD collapsing to the middle point.Source: §A.3 · p. 17
• Top-K matching assigns each student action to its nearest selected high-value reference action, but K = 1 collapses to one mode and K = 3 can retain intermediate-region mass.Hard selection does not specify how much mass each selected reference action should receive, making the method sensitive to K.Source: §A.3 · p. 17
• Independent value-distance matching can send many samples to one mode when comparable-value targets are broken symmetrically by initialization, distance, or estimated value.With all particles initialized in the left half of the equal-value bimodal bandit, VDWM assigns them to the closer left mode.Source: Figure 5 · p. 17, §A.3 · p. 18

A.4 Computational cost

• Sinkhorn updates are not a dominant training-time bottleneck: reducing iterations from 30 to 20 or 10 yields only marginal runtime improvements.The default is 30 iterations; Table 3 reports average training-step time over 15 tasks on one NVIDIA H100 GPU.Source: §A.4 · p. 18, Table 3 · p. 18
• Runtime is more sensitive to reference-policy samples M than to one-step-policy samples N.Increasing M raises reference-action generation and matching costs substantially; increasing N adds comparatively moderate overhead.Source: §A.4 · p. 18, Table 4 · p. 19
• Sample generation and transport matching dominate training cost, while inference requires only one forward pass.Reducing Sinkhorn iterations from 30 to 20 or 10 yields only marginal runtime improvements.Source: §A.4 · pp. 18–19, Table 3 · p. 18, Figure 6 · p. 19

B Broader Impacts

• OptiFlow’s simulated-control results may reduce reliance on unsafe or costly online exploration, and its one-step deployment may suit latency-sensitive control.The broader-impact discussion still identifies out-of-dataset behavior and inaccurate critic values as deployment risks.Source: B Broader Impacts · p. 19

C.1 Algorithm

• The Sinkhorn procedure alternates row and column scaling updates to compute a transport plan from the cost matrix, marginals, and regularization parameter.The experiments use 30 iterations.Source: §C.1 · p. 20

Architectures

• The critic, reference flow policy, and one-step policy use four-layer 512-unit MLPs with GeLU activations; only the critic uses LayerNorm.The critic is a two-member ensemble, and policy outputs are clipped to [-1, 1] at use time.Source: D Implementation Details · p. 20

Flow teacher and image processing

• The reference flow uses linear-interpolant flow matching with uniform time sampling and 10 Euler steps; visual tasks use an IMPALA encoder and random-shift augmentation.Visual inputs stack three frames; the augmentation probability is 0.5.Source: D Implementation Details · p. 20

Transport and distillation

• Per state, the method transports 16 one-step actions against 64 reference actions, using squared-distance costs and a value-softmax reference marginal.Balanced Sinkhorn runs for 30 iterations with ε = 0.05; the actor combines value-aware BC and distillation losses equally.Source: D Implementation Details · p. 20

Training and evaluation

• Offline training runs for 1M steps on state-based OGBench and 500K on D4RL, with evaluation averaged over eight seeds unless otherwise noted.Online fine-tuning adds 1M steps using deterministic one-step-policy data and retains the offline dataset in the replay buffer.Source: D Implementation Details · pp. 20–21

Multimodal diagnostics

• The four-goal diagnostic tests whether transport-guided extraction retains four equally optimal modes despite asymmetric dataset coverage.The diagnostic removes 80% of transitions involving states or next states below y = x, leaving 55,044 transitions; assignment stability is also tested by resampling 64 reference actions 100 times.Source: D Implementation Details · p. 21

Hyperparameters

• Tables 5–8 provide shared and task-specific hyperparameters, including OptiFlow’s τ and η settings for OGBench and D4RL.The per-task overrides apply to OptiFlow and GFP; other methods use shared defaults unless specified.Source: Table 5 · p. 22, Table 6 · p. 22, Table 7 · p. 23

E Complete Experimental Results

• The appendix reports OGBench and D4RL benchmark results, averaging eight seeds except for visual tasks, which use four.The OGBench and D4RL tables specify their respective evaluation protocols and reported metrics.Source: E Complete Experimental Results · p. 24, Table 9 · p. 24, Table 11 · p. 25

F Ablation Studies

• The ablations examine temperature and Sinkhorn settings, action-sample counts, reference-policy quality, weighting alternatives, direct Q-maximization, and offline-to-online fine-tuning.They also test whether larger sampling budgets alone account for OptiFlow’s gains.Source: F Ablation Studies · p. 26

F.1 Sensitivity to value-weighting temperatures

• Performance varies smoothly around tuned τ and η settings, with no single temperature dominating the result.The authors report that appropriate task-specific settings can be found with minor tuning.Source: §F.1 · p. 26, Figure 7 · p. 26

F.2 Sensitivity to Sinkhorn hyperparameters

• With ε = 0.05, 30 Sinkhorn iterations perform best on all three swept tasks and add little overhead.The sweeps show a mild task-dependent trade-off; the cost matrix is normalized by its mean.Source: §F.2 · pp. 26–27, Table 12 · p. 26

F.3 Reference and one-step action sampling

• Performance depends mainly on reference sample count M and is largely insensitive to N; the default (N, M) = (16, 64) matches or achieves the best result on each task.Smaller M degrades performance sharply, while M = 256 gives consistent gains; the default balances performance and computation.Source: §F.3 · p. 27, Table 13 · p. 27

F.4 Controlling for the action-sampling budget

• FQL and GFP remain below OptiFlow’s 57.3 average even when using 64 one-step samples, versus OptiFlow’s 16.FQL averages 37.6, 38.5, and 38.1 at N = 1, 16, and 64; GFP averages 52.5, 51.3, and 51.2, respectively.Source: Table 14 · p. 27, §F.4 · p. 28

F.5 Disentangling reference policy quality and distillation

• With comparable reference performance, OptiFlow-BC beats FQL on 8 of 10 tasks and raises the one-step average from 37.6 to 48.6; full OptiFlow reaches 57.3.The authors describe value-aware BC and OT distillation as complementary; their reference policy averages 19.8 versus GFP’s 28.4, while its one-step policy averages 57.3 versus 52.5.Source: §F.5 · p. 28, Table 15 · p. 28

F.6 AWR-style weighting for value-aware behavior cloning

• AWR-style weighting improves some tasks but not consistently, averaging 52.5 ± 1.4 versus 57.3 ± 1.0 for the default weighting.The authors retain the default two-class softmax weighting.Source: §F.6 · p. 29, Table 16 · p. 29

F.7 Direct Q-maximization

• Direct Q-maximization has mixed, task-dependent effects and larger weights can collapse performance; the main results therefore use λq = 0.The authors report strong benchmark performance with transport-based distillation alone in their setting.Source: §F.7 · pp. 29–30, Table 17 · p. 29

G Offline-to-Online Fine-tuning

• Online fine-tuning improves over the offline checkpoint across 15 evaluated OGBench and D4RL tasks, without modifying the OptiFlow training procedure.The 15-task set matches FQL.Source: Table 19 · p. 30, G Offline-to-Online Fine-tuning · p. 31
• Newly collected transitions enter critic updates, VaBC reference-policy training, and transport-based target construction while retaining the offline checkpoint’s reference policy and one-step distillation structure.Online interaction continues from an offline OptiFlow checkpoint at 1M training steps.Source: G Offline-to-Online Fine-tuning · pp. 30–31
• Puzzle-4×4-play-singletask-task4-v0 has high seed variance, with five of eight seeds reaching 100 and three remaining at 0; adroit tasks require tuning.Relocate-cloned can improve with further tuning, but results are reported without task-specific online tuning unless otherwise stated in Table 18.Source: Table 18 · p. 30, G Offline-to-Online Fine-tuning · p. 31

10. Broader impacts

• The work involved neither crowdsourcing nor research with human subjects.Source: §10 · p. 38
• LLMs were not used as an important, original, or non-standard component of the core methods.Source: §10 · p. 38