NeurIPS 2026 Papers

Conference on Neural Information Processing Systems

Decision breakdown

6,437papers

Counted from the official paper list; may differ slightly from the number the conference announced.

  • Oral
    991.5%
  • Spotlight
    2403.7%
  • Poster
    6,09894.7%

Browse all papers

Topic: Reinforcement Learning · 760 papers · Reset filters

Learning to Surpass: Training Tool-Using Agents with Anchored Feedback

John Gkountouras, Fengjun Wang, Angelantonio Castelli, Satendra Kumar

Expand abstract

Training tool-using agents for planning is challenging when high-quality demonstrations exist only as final outcomes without the intermediate tool traces that produced them. A common response is to impute traces and apply imitation learning, but this approach is indirect and structurally limited to matching the reference rather than improving upon it. We propose anchored comparative feedback: rather than imitating a reference plan, we train agents to surpass it under a comparative LLM judge while satisfying hard constraints. The anchor provides a stable optimization target, and comparing two identically rendered plans neutralizes surface-level judge biases. On TripTailor, a tool-augmented travel planning benchmark, anchored GRPO achieves 80.1% final success rate compared to 43.7% for imitation learning on imputed traces, with gains on both constraint satisfaction and LLM-judged preference that persist across judge families, rubrics, and a held-out reward model. Gains generalize beyond TripTailor, with a fivefold improvement on the held-out TravelPlanner benchmark. These results suggest that anchored comparative feedback offers an effective approach for learning transferable tool-using planners from outcome-only supervision.

Poster
Learning to Persuade Exposes How Easily LLMs Abandon Correct Beliefs

Nimet Beyza Bozdag, Emre Can Acikgoz, Gokhan Tur, Dilek Hakkani-Tur

Expand abstract

Persuasion is a core dynamic of natural language communication, shaping how large language models (LLMs) update beliefs, resolve disagreements, and reach decisions. As LLMs increasingly debate, advise, and think collaboratively with humans and each other, resistance to harmful persuasion becomes a core requirement for reliable behavior. Yet we show that this requirement is far from met: a single targeted persuasive argument is enough to collapse model accuracy to near zero, even when the argument is factually false. We formalize this threat as adversarial persuasion and introduce an adversarial reinforcement learning framework that trains persuader agents to change a target model's answer in a single interaction. First, we show that optimizing persuasion strategies through trial and error exposes vulnerabilities that static prompting misses: RL-trained persuaders raise persuasion success from approximately 24\% to over 93\% against the training-time persuadee. Second, we find that these learned strategies transfer to unseen models, achieving 83\% attack success on Qwen-14B, 79\% on Llama-3.1-8B, and 25\% on GPT-4o-mini. Third, we demonstrate that a curriculum that bootstraps on more persuadable open-weight models before targeting harder models further increases GPT-4o-mini attack success from 25\% to 38\%. Moreover, our results reveal that optimized persuaders increasingly rely on credibility-based tactics, including fabricated citations and false authoritative evidence. Together, these findings expose a critical weakness in current LLM agents: even when they initially reason correctly, they can be steered toward false conclusions by optimized natural language influence. This positions persuasion robustness as a necessary safety criterion for multi-agent and human-AI decision-making systems.

Poster
Learning to Learn from Multimodal Experience

Xingyu Sui, Weixiang Zhao, Yongxin Tang, Yanyan Zhao, Yang Wu, Dandan Tu, Bing Qin

Expand abstract

Experience-driven learning has emerged as a promising paradigm for enabling agents to improve from interaction trajectories by accumulating and reusing past experience. However, existing approaches are predominantly developed in textual settings and rely on manually designed memory schemas, limiting their applicability to multimodal environments. In real-world scenarios, experience is inherently multimodal, involving heterogeneous signals across perception, reasoning, and action, which makes effective memory design significantly more challenging. In particular, the optimal way to structure and utilize multimodal experience is highly task-dependent and evolves over time, rendering fixed memory designs insufficient. In this work, we propose a new paradigm, learning to learn from multimodal experience, which shifts memory design from a predefined component to an adaptive and learnable process. Our framework enables agents to dynamically construct, organize, and utilize memory based on task requirements and interaction history, effectively learning how to structure experience for improved performance. Experiments demonstrate that adaptive memory design substantially enhances agent performance and generalization across multimodal tasks, highlighting the critical role of learning memory mechanisms in experience-driven learning.

Poster
Learning to Group and Order: Cross-Instance Self-Supervised RL for Vision-Centric MLLMs

Songming Yang, Hong-Tao Yu, Xiu-Shen Wei

Expand abstract

Reinforcement learning (RL) post-training has become an important paradigm for improving multimodal large language models (MLLMs), especially with automatically verifiable rewards. For vision-centric MLLMs, self-supervised visual pretext tasks provide such rewards without human annotations. However, existing self-supervised RL methods mostly exploit within-instance structure, such as spatial or temporal ordering within a single image or video, offering limited supervision for cross-instance discrimination and fine-grained comparison. We introduce Group-and-Order Self-Supervised Reinforcement Learning (GO-SSL), a cross-instance jigsaw framework for vision-centric MLLM post-training. Given two visual instances, GO-SSL mixes their local elements and trains the model to group them by instance of origin while recovering the order within each group, coupling instance-level comparison with spatial, temporal, or geometric structure recovery. This paired formulation turns the same unlabeled data into richer verifiable supervision through diverse pairings, shuffles, and hard-pair curricula. We instantiate GO-SSL mainly on cross-image jigsaw tasks and further extend it to 3D depth and video temporal ordering. Across diverse benchmarks, GO-SSL consistently improves over the base MLLM and prior baseline methods, with clear gains in fine-grained perception and spatial understanding. These results suggest a data-efficient direction for self-supervised RL in MLLMs: constructing comparative cross-instance contexts can provide richer transferable supervision than merely scaling data volume or adding isolated pretext tasks. Code will be released later.

Poster
Learning to Follow In-Context Watermark Instructions via Self-Distillation

Yepeng Liu, Tianyi Chen, Xuandong Zhao, Dawn Song, Yuheng Bu

Expand abstract

In-context watermarking (ICW) prepends an instruction to a query asking the model to embed a statistically detectable signal in its response. It thus equips LLMs with a watermarking interface that third parties can invoke without access to model internals. Its reliability hinges on the LLM following the instruction without degrading answer quality, yet how well current LLMs do so has not been measured. We introduce ICWBench, a benchmark of three verifiable ICW instruction families, each scored on both detectability and answer quality. Evaluating 14 frontier proprietary and open-source LLMs, we find that none of the evaluated LLMs achieves both objectives across all three families. To address this, we propose a self-contained two-stage training method, requiring no distillation from a stronger model, no manual annotation, and no pre-existing ICW IF ability. The first stage, self-distillation with logits perturbation (SDLP), uses the same base LLM as both teacher and student: an instruction-equivalent decoding-time logits perturbation makes the teacher follow the ICW instruction, and the student is trained to match the teacher's output distribution. The second stage applies reinforcement learning with the automatic verifier as the reward. Applied to Qwen3-14B, the weakest of the 14 evaluated LLMs in ICW IF, our method raises average TPR@$1\%$FPR across three ICW instructions from $0.100$ to $0.974$, achieving a more favorable trade-off than the frontier proprietary LLMs we evaluate.

Poster
Learning to Explore with Parameter-Space Noise: A Deep Dive into Parameter-Space Noise for Reinforcement Learning with Verifiable Rewards

Bizhe Bai, Xinyue Wang, Peng Ye, Tao Chen

Expand abstract

Reinforcement Learning with Verifiable Rewards (RLVR) enhances large language model (LLM) reasoning, yet growing evidence indicates an \emph{exploration ceiling}: they tend to reweight existing solution traces and produce low-diversity rollouts. This restricts the discovery of novel solution perspective and weakens out-of-distribution generalization. We address this limitation by directly perturbing the rollout policy's parameters while updating the clean policy model as usual, an approach that achieves better temporal consistency and yields more complex behavior patterns. To stabilize learning from these perturbed rollouts, which are inherently off-policy, we incorporate truncated importance sampling. Furthermore, we introduce a lightweight, dynamic noise scheduler based on semantic diversity and model self-certainty to optimally adjust the noise scale. Instantiated on GRPO, our method (PSN-GRPO) demonstrates its efficacy through four key contributions: (1) it improves high-budget pass@$k$, increases semantic and operational diversity, increases CoT consistency, and remains complementary to exploration-oriented RLVR techniques; (2) it validates efficacy across multiple models, including Qwen2.5, Qwen3, and Llama-3.1; (3) it shows robust performance on out-of-distribution science and coding benchmarks; and (4) it uncovers novel solution perspectives for problems the original model could not solve, as shown by qualitative analysis.

Poster
Learning to Cut: Reinforcement Learning for Benders Decomposition

Haochen Cai, Xian Yu

Expand abstract

Benders decomposition (BD) is a widely used solution approach for solving two-stage stochastic programs arising in real-world decision-making under uncertainty. However, it often suffers from slow convergence as the master problem grows with an increasing number of cuts. In this paper, we propose Reinforcement Learning for BD (RLBD), a framework that adaptively selects cuts using a neural network-based stochastic policy. The policy is trained using a policy gradient method via the REINFORCE algorithm. We evaluate the proposed approach on a two-stage stochastic electric vehicle charging station location problem and compare it with vanilla BD and LearnBD, a supervised learning approach that classifies cuts using a support vector machine. Numerical results demonstrate that RLBD achieves substantial improvements in computational efficiency and exhibits strong generalization to problems with similar structures but varying data inputs and decision variable dimensions.

Poster
Learning to Bid in Repeated Second-Price Auctions with Dynamic Values and Aggregated Feedback

Benjamin Heymann, Otmane Sakhi

Expand abstract

We study the problem of learning to bid when the bidder’s *value is dynamic*, i.e., when the current value depends on past outcomes. Specifically, we consider a bidder participating in repeated second-price auctions whose value depends on the time elapsed since their last successful bid, with auctions arriving in continuous time and only aggregated feedback revealed at the end of the horizon. Such a bidder must **(1)** balance the immediate benefit of winning the current auction against its impact on future values and **(2)** learn unknown environmental parameters. We derive regret bounds for a class of learning methods that combine plug-in estimators with a differential-equation characterization of the optimal policy, and show that a specific confidence bound algorithm learns the optimal policy with a near optimal regret of $\tilde{\mathcal{O}}(\log N)$ for piecewise linear primitives, and $\tilde{\mathcal{O}}(N^{1/3})$ for general, smooth primitives, achieving these regrets without explicit randomization. These theoretical results are supported by numerical experiments.

Poster
Learning to Audit ML Models with Theory of Mind

Yao Rong, Vaibhav Unhelkar

Expand abstract

When auditing machine learning models, human experts sequentially inspect model predictions and explanations, gradually building a mental model of the model’s behavior to uncover errors. This manual process is critical for safe deployment in high-stakes domains but fundamentally limited in scalability. In this work, we propose a novel method to simulate this human auditing behavior for automating the auditing process. We formalize auditing as a sequential decision-making problem, modeled as a Markov decision process. Our state representation summarizes the auditor’s accumulated knowledge of model predictions and feature usage. Building on this formalization, we solve the MDP using reinforcement learning, yielding RLAuditor, which learns to select informative test samples to efficiently uncover model errors. We further connect our formalization to Theory of Mind, showing that our state mirrors how human auditors build mental models: its update rule is equivalent to Bayesian belief revision and it converges to the model’s behavioral pattern at an interpretable rate. We validate our approach on multiple ML models and datasets across different data modalities. In a human subject study, we demonstrate that RLAuditor helps human auditors produce more accurate audit reports in less time.

Oral
Learning in Context, Guided by Choice: A Reward-Free Paradigm for Reinforcement Learning with Transformers

Juncheng Dong, Moyang Guo, Bowen He, Ethan Fang, Zhuoran Yang, Vahid Tarokh

Expand abstract

In-context reinforcement learning (ICRL) leverages the in-context learning capabilities of transformer models (TMs) to efficiently generalize to unseen sequential decision-making tasks without parameter updates. However, existing ICRL methods rely on explicit reward signals during pretraining, which limits their applicability when rewards are ambiguous, hard to specify, or costly to obtain. To overcome this limitation, we propose a new learning paradigm, \emph{In-Context Preference-based Reinforcement Learning} (ICPRL), in which both pretraining and deployment rely solely on preference feedback, eliminating the need for reward supervision. We study two variants that differ in the granularity of feedback: \emph{Immediate Preference-based RL} (I-PRL) with per-step preferences that generalize dueling bandits, and \emph{Trajectory Preference-based RL} (T-PRL) with trajectory-level comparisons. We first show that supervised pretraining, a standard approach in ICRL, remains effective under preference-only context datasets, demonstrating the feasibility of in-context solving unseen tasks using only preference signals. To further improve data efficiency, we introduce alternative preference-native frameworks for I-PRL and T-PRL that directly optimize TM policies from preference data without requiring reward signals nor optimal action labels. Experiments on dueling bandits, navigation, and continuous control tasks demonstrate that ICPRL enables strong in-context generalization to unseen tasks, achieving performance comparable to ICRL methods trained with full reward supervision.

Poster
Learning in Causal Markov Games

Aurghya Maiti, Elias Bareinboim

Expand abstract

Markov Games are the standard formal model for multi-agent reinforcement learning, capturing agents that act in a shared state and optimize rewards over time. However, in real-world settings, agents’ decisions are often influenced by unobserved factors, such as cognitive biases, behavioral tendencies, or intuitive signals, that also affect rewards and future states. Ignoring these unobserved confounders can make standard learning methods converge to suboptimal policies. In such settings, optimal play may require policies that condition on counterfactual signals, requiring reasoning at the counterfactual layer of the Pearl Causal Hierarchy. In this paper, we introduce Causal Markov Games (CMGs), a framework for modeling sequential multi-agent decision making in the presence of unobserved confounding. We show that CMGs strictly generalize Markov Games, with arbitrarily large gaps between classical interventional equilibria and causal counterparts. We then develop two learning algorithms under different observability assumptions. The first, CNash-VI-FO, is a model-based learner with finite-sample guarantees when agents’ natural actions or intuitions are revealed post hoc. The second, CNash-VI-NO, is an explore-then-exploit procedure with asymptotic guarantees for the setting in which opponents’ natural actions are never observed. To address scalability, we further provide a drop-in counterfactual augmentation of deep MARL. Empirically, on a confounded windy variant of the Multi-Particle Environment and the Iterated Causal Prisoner’s Dilemma, counterfactual agents strictly dominate their non-causal counterparts.

Poster
Learning from Language Feedback via Variational Policy Distillation

Yang Li, Erik Nijkamp, Semih Yavuz, Shafiq Joty

Expand abstract

Reinforcement learning from verifiable rewards (RLVR) suffers from sparse outcome signals, creating severe exploration bottlenecks on complex reasoning tasks. Recent on-policy self-distillation methods attempt to address this by utilizing language feedback to generate dense, token-level supervision. However, these approaches rely on a fixed, passive teacher to interpret the feedback. As the student policy improves, the teacher's zero-shot assessment capabilities plateau, ultimately halting further learning. To overcome this, we propose Variational Policy Distillation (VPD), a framework that formalizes learning from language feedback as a Variational Expectation-Maximization (EM) problem. VPD co-evolves both policies: in the E-step, the teacher is actively refined on trajectory outcomes via an adaptive trust-region update, translating textual feedback into a dynamically improved target token distribution. In the M-step, the student internalizes this dense distributional guidance on its own on-policy rollouts. By continuously improving the teacher's ability to extract actionable signals from textual critique, VPD overcomes the limitations of passive distillation. Evaluated across diverse sources of diagnostic feedback on scientific reasoning and code generation tasks, VPD consistently outperforms both standard RLVR and existing self-distillation baselines. Finally, by stress-testing our framework on rigid mathematical reasoning and cold-start regimes, we illuminate the fundamental bounds of feedback-driven self-distillation compared to pure environment-driven RL.

Poster
Learning When to Think: Adaptive Internal Computation for Reinforcement Learning

Rohit Kumar Salla, Simon Stepputtis

Expand abstract

As agents have to solve diverse tasks from grid-worlds to robotics and language modeling, the difficulty of each step varies, yet policies are rigid and spend a fixed amount of compute on each step. A natural approach is to spend compute on the harder states while saving it on easy ones, allowing agents to retain performance while saving valuable compute budgets. To this end, we propose Adaptive Internal Computation (AIC), a lightweight policy architecture that iteratively refines latent states and learns when to halt using a value-of-computation signal.Concretely, AIC repeatedly refines a latent state representation and learns a halting head from a value-of-computation signal, the per-step gain in expected return, allowing the policy to stop refining as soon as further computation no longer improves the action it would take. A central challenge, however, is that matched-budget gains alone cannot establish that compute is being directed to the states that benefit the most from them. We therefore introduce a six-component causal evaluation protocol leveraging a targeted-versus-random ablation that tests whether removing compute from AIC's predicted high-compute states hurts performance more than removing the same amount of compute from randomly chosen states, thereby validating that AIC can successfully identify high-compute states. Across twelve diverse tasks, including grid worlds, simulated robotics experiments, and language modeling, we demonstrate that AIC outperforms fixed-compute policies across a set of five underlying model architectures. In particular, architectures leveraging AIC retain or improve performance while spending about 52\% less compute compared to their respective fixed counterparts.

Poster
Learning When to Stop: Selective Imitation Learning Under Arbitrary Dynamics Shift

Surbhi Goel, Jonathan Pei, James Wang

Expand abstract

Behavior cloning provides strong imitation learning guarantees when training and test environments share the same dynamics. However, in many deployment settings the test environment's transitions differ from training, and classical offline IL offers no recourse: the learner must commit to an action at every state, even when its demonstrations are uninformative and could lead to arbitrary degradation of performance. This motivates the study of _selective_ imitation, where the learner may choose to _stop_ when it cannot act reliably. We introduce a model for selective imitation under arbitrary dynamics shift: given labeled expert demonstrations from a training environment and unlabeled state trajectories from the same expert in a test environment, the learner outputs a _selective policy_ that is _complete_ (rarely stops in training) and _sound_ (incurs low regret before stopping in test). Our algorithm, $\operatorname{SeqRejectron}$, constructs a stopping rule using a small set of _validator policies_ whose size is independent of the horizon or policy class. For deterministic policies, this yields horizon-free $\tilde{O}(\log|\Pi|/\epsilon^2)$ sample complexity, assuming sparse costs. For stochastic policies, we obtain analogous horizon-free guarantees using a cumulative Hellinger stopping time. We extend the framework to misspecified experts and different expert policies across train and test and obtain results that gracefully degrade with the amount of misspecification.

Poster
Learning Robust Reasoning through Guided Adversarial Self-Play

Shuozhe Li, Vaishnav Tadiparthi, Kwonjoon Lee, Nakul Agarwal, Hossein Nourkhiz Mahjoub, Ehsan Moradi Pari, Lizhang Chen, Du Cheng, Amy Zhang, Liu Leqi

Expand abstract

Reinforcement learning from verifiable rewards (RLVR) produces strong reasoning models, yet these models can fail catastrophically when the conditioning context is fallible (e.g., corrupted chain-of-thought, misleading partial solutions, or mild input perturbations), because standard RLVR optimizes final-answer correctness only under clean conditioning. We introduce GASP (Guided Adversarial Self-Play), a robustification method that explicitly trains error detection and repair capabilities using only outcome verification. Without human labels or external teachers, GASP instantiates an adversarial self-play game within a single model: a polluter learns to induce failure through locally coherent corruptions, while an agent learns to diagnose and recover under the same corrupted conditioning. To address the scarcity of successful recoveries early in training, we propose in-distribution repair guidance, an auxiliary imitation objective on self-generated repairs that increases recovery probability while preserving previously acquired capabilities. Across four open-weight models (1.5B–8B), GASP converts strong-but-brittle reasoners into substantially more robust ones that better withstand misleading and perturbed context, while often improving clean accuracy. Further analysis shows that adversarial corruptions create an effective curriculum, and that in-distribution guidance enables rapid recovery learning with minimal representational drift.

Poster
Learning Reusable Motor Motifs for Continuous Animal Behavior Modeling

Jiyi Wang, Jingyang Ke, Bo Dai, Anqi Wu

Expand abstract

Animals generate complex behaviors by flexibly recombining a finite set of motor primitives, but existing behavior segmentation methods oversimplify this process by imposing discrete syllables under restrictive generative assumptions. Here, we introduce Motif-based Continuous Dynamics (MCD), a framework that models behavior as a continuous, compositional process driven by reusable motor motifs. MCD leverages reinforcement learning (RL) to (1) discover interpretable motif representations via transition-based representation learning, and (2) model behavior as a time-varying mixture of these motifs through motif-based policies. This formulation avoids restrictive dynamics assumptions while capturing continuity, compositionality, and long-term dependencies in behavior. Across simulated gridworld tasks, maze navigation, and animal behavior datasets, MCD identifies reusable motifs and interprets the trajectories with them. These results provide a generative account of behavior as combinations of fundamental motifs, offering a flexible and interpretable framework for studying natural behavior.

Poster
Learning Process Rewards via Visitation Matching for Efficient RL

Raymond Tsao, Andrew Wagenmaker, Sergey Levine

Expand abstract

In many modern applications of reinforcement learning (RL), the natural reward for a task of interest is inherently sparse: a reward of 0 is given everywhere except when the task is completed, when a reward of 1 is given. Training a policy to maximize such a sparse reward requires solving a challenging credit assignment problem, leading to slow or ineffective RL improvement. We propose a simple approach to transform a sparse outcome reward into a dense process reward. Our approach relies on training a discriminator to distinguish between previous successful and unsuccessful trajectories, and using this discriminator to incentivize the RL-learned policy to match the state-action visitations of successful trajectories, while avoiding those of unsuccessful trajectories. By incentivizing the policy to match the visitations over all states, not just those that correspond to task success, this reward provides dense feedback on whether progress is being made towards task completion, and, we show, provably achieves this without changing the optimal policy. Focusing on finetuning of robotic control policies, we demonstrate that our approach leads to significantly faster RL finetuning performance on both simulated and real-world manipulation tasks, as compared to simply maximizing the sparse outcome reward.

Poster
Learning Planning Budgets in Real-Time RL

Aneesh Muppidi, Firas Darwish, Dylan Cope, Joao Henriques, Jakob Foerster

Expand abstract

Deliberating takes time. In real-time settings, that time is not free. Standard reinforcement learning (RL) sidesteps this as the environment waits indefinitely for the agent's decision. Instead, we study real-time RL environments where the environment progresses while waiting for the agent's action. Building on prior real-time formalizations, we introduce variable-delay real-time RL, where the agent chooses how long to deliberate at each decision point since the environment progresses. For the planning agents we use, the right delay is state-dependent, and naively planning how long to plan can paralyze the agent. We instead approach this setting by training a lightweight gating policy on top of a planner to select state-dependent planning budgets. Across real-time Pac-Man, Tetris, Snake, Speed Hex, and Speed Go, our gating policy outperforms fixed-budget and heuristic baselines, and transfers to a real-time setup where the environment and agent run on two different GPUs.

Poster
Learning Neuronal Wiring Rules from Morphological Token Sequences

Sidharth Goel, Erdem Varol

Expand abstract

Does a neuron’s shape predict whom it connects to? Peters’ rule, a canonical geo- metric principle in connectomics, predicts connectivity from spatial arbor overlap, but true wiring also depends on molecular and morphological specificity. We ask how far neuron skeletons, with identity and graph features withheld, predict directed connections in the FlyWire Drosophila connectome. Because most non- connected pairs are far apart, headline metrics are dominated by easy negatives that geometry can reject. We therefore evaluate controlled contrasts: plausible non-connections whose arbors are nearby, overlapping, or cell-type matched, but that still do not synapse. Across a controlled model ladder, from arbor overlap and hand-crafted morphology baselines to Connectoformer (a bidirectional Trans- former with pooled and pairwise cross-attention scoring), performance improves when the model gains the corresponding signal: global shape, local branch com- patibility, and tree context. Morphology and arbor co-occupancy are complemen- tary: morphology discriminates similar-looking unconnected pairs but can score pairs whose arbors never touch; arbor co-occupancy does the inverse. By suppress- ing essentially-zero-overlap pairs without boosting high-overlap pairs, a fixed one- sided arbor prior performs well in both regimes. The encoder also transfers across Drosophila individuals: zero-shot evaluation on the male CNS connectome beats arbor overlap on hard-negative discrimination. Stratified evaluation reveals a clean division of labor: geometry rules out pairs that cannot physically connect, and morphology picks the actual partners from what remains.

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

Jaehun Shon, Jinha Choi, Jongwook Jeon, Jongmin Lee

Expand abstract

Offline reinforcement learning aims to learn a policy solely from fixed datasets, which often contain multimodal action distributions. Flow policies can naturally represent such multimodal behaviors, but learning an efficient one-step flow policy remains challenging: standard value guidance often leads to mode collapse or exploits overestimation bias in out-of-distribution regions. To address this, we introduce One-step Flow policy via Optimal Transport (OptiFlow), a framework for one-step flow policy learning as a structured sample-allocation problem. OptiFlow jointly trains a value-aware reference flow policy and an efficient one-step policy, coupling their action samples through state-wise entropic optimal transport. For each state, critic-estimated values define the priority of distillation target actions, while the action-distance cost ensures geometrically compatible pairings. By avoiding direct critic maximization, our transport-guided approach enables in-distribution exploitation by anchoring the one-step policy to high-value, dataset-supported modes without the risk of out-of-distribution divergence. Experimental results demonstrate that OptiFlow effectively captures optimal multimodal behaviors and achieves strong performance across diverse offline RL benchmarks.

Poster