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

Q-Probe: Scaling Image Quality Assessment to High Resolution via Context-Aware Agentic Probing

Xiang Li, Xueheng Li, Yu Wang, Xuanhua He, Zhangchi Hu, Chengjun Xie

Expand abstract

Reinforcement Learning (RL) has empowered Multimodal Large Language Models (MLLMs) to achieve superior human preference alignment in Image Quality Assessment (IQA). However, existing RL-based IQA models typically rely on coarse-grained global views, failing to capture subtle local degradations in high-resolution scenarios. While emerging "Thinking with Images" paradigms enable multi-scale visual perception via zoom-in mechanisms, their direct adaptation to IQA induces spurious "cropping-implies-degradation'' biases and misinterprets natural depth-of-field as artifacts. To address these challenges, we propose Q-Probe, the first agentic IQA framework designed to scale IQA to high resolution via context-aware probing. First, we construct Vista-Bench, a pioneering benchmark tailored for fine-grained local degradation analysis in high-resolution IQA settings. Furthermore, we propose a three-stage training paradigm that progressively aligns the model with human preferences, while simultaneously eliminating causal bias through a novel context-aware cropping strategy. Extensive experiments demonstrate that Q-Probe achieves state-of-the-art performance in high-resolution settings while maintaining superior efficacy across resolution scales.

Poster
Q-MMR: Off-Policy Evaluation via Recursive Reweighting and Moment Matching

Nan Jiang, Xiang Li

Expand abstract

We present a novel theoretical framework, Q-MMR, for off-policy evaluation in finite-horizon MDPs. Q-MMR learns a set of scalar weights, one for each data point, such that the reweighted rewards approximates the expected return under the target policy. The weights are learned inductively in a top-down manner via a moment matching objective against a value-function discriminator class. Notably, and perhaps surprisingly, a data-dependent finite-sample guarantee for general function approximation can be established under only the realizability of $Q^\pi$, with a **dimension-free** bound---that is, the error does **not** depend on the statistical complexity of the function class. We also establish connections to several existing methods, such as linear FQE. Further theoretical analyses shed new light on the nature of coverage, a concept of fundamental importance to offline RL.

Poster
Pushing Biomolecular Utility-Diversity Frontiers with Supergroup Relative Policy Optimization

Xinwu Ye, He CAO, Li Hao, Bin Feng, Zijing Liu, Robert Tang, Yu Li, Shenghua Gao

Expand abstract

Biomolecular generators are often adapted with reward feedback to improve task-specific utility, but pushing utility alone can concentrate generation on a narrow family of candidates. Maintaining diversity is difficult because sample diversity is a set-level property. We introduce Supergroup Relative Policy Optimization (SGRPO), a flexible GRPO-style framework that directly constructs rewards from set-level diversity. For each condition, SGRPO samples a supergroup of candidate sets, compares their diversity under the same condition, and redistributes the group diversity reward to individual rollouts through leave-one-out diversity contributions before combining it with rollout-level utility. This design decouples SGRPO from a particular generator, utility reward, or diversity metric, and allows instantiation with different GRPO-style approaches. We evaluate SGRPO on de novo small-molecule design, pocket-based small-molecule design, and de novo protein design, instantiating it with both GRPO and Coupled-GRPO across autoregressive and discrete diffusion generators. Across decoding sweeps, SGRPO expands the utility-diversity Pareto frontier and achieves the best frontier-level metrics relative to pretrained generators, GRPO, and memory-assisted GRPO when applicable. Our analyses further show that direct set-level diversity rewards remain effective with small groups and help preserve broader generation-distribution coverage during post-training. The code is available at https://anonymous.4open.science/r/SGRPO/README.md.

Poster
Provably Safe, Yet Performant Reinforcement Learning

Kai Yun, Zeyang Li, Navid Azizan

Expand abstract

Safe reinforcement learning (RL) aims to learn policies that optimize rewards while satisfying constraints. Predominant approaches rely on soft-constrained policy optimization, which has achieved empirical success but does not provide formal safety guarantees for the learned policy. In contrast, methods with strict guarantees typically rely on explicit certificate functions, whose construction requires the direct synthesis and verification of control-invariant sets, a process that scales poorly with state dimension and often yields overly conservative behavior. In this paper, we present the Provably Safe, yet Performant RL (PSP-RL) framework, a novel two-phase architecture for learning provably safe policies in a scalable manner, designed to overcome the key bottlenecks of prior methods. Rather than explicitly computing invariant sets, PSP-RL leverages a learned backup policy to forward-integrate the system dynamics, generating an implicit control-invariant set online. In the first phase, the backup policy is trained with our proposed safe-arrival value function, which characterizes the optimal backup policy for invariant-set construction. In the second phase, an RL policy is trained end-to-end through a differentiable projection layer that strictly enforces the safety guarantees induced by the learned backup policy. By maximizing the volume of the implicit control-invariant set in the first phase, the resulting PSP policy from the second phase is performant and scalable, while maintaining provable safety. Crucially, PSP-RL imposes no restrictions on the underlying RL algorithm and can be plugged into any existing training pipeline. We establish theoretical guarantees for the proposed framework and evaluate it on robotic control tasks with state dimensions up to 10, a regime in which prior provably safe RL methods struggle or become impractical.

Spotlight
Prospective Hindsight: Self-Calibrating Reinforcement Learning via Prediction–Reality Gaps

Jiaxin Zhang, Xiangyu Peng, Qinglin Chen, Yu Li, Hiroaki Hayashi, Chien-Sheng Wu

Expand abstract

Reinforcement learning for long-horizon agents relies on \emph{purely retrospective} training signals: credit is assigned only after observing environmental consequences, leaving the agent's belief at action time invisible to the gradient. We introduce \textbf{Prospective Hindsight (PH)}, a self-calibrating training principle that augments any retrospective base method with a signal derived from the gap between the agent's \emph{prospective prediction} (before feedback) and the \emph{retrospective evaluation} (after feedback). This per-rollout \emph{surprise} identifies samples where the agent's self-model is most inaccurate and amplifies their gradient contribution through a stop-gradient surprise-weighted advantage. Since the prospective predictor shares parameters with the policy, the two co-evolve, progressively shifting focus to the agent's remaining blind spots. We connect this principle to a privileged-information gap and show that minimizing the surprise residual provides a descent pathway on the agent's miscalibration rate; calibration thus emerges as a byproduct of optimization rather than from an added objective. On single-turn verifiable tasks and a multi-turn personal-agent task (under GRPO, on-policy distillation, and their combination), PH improves both task performance and calibration, with consistent gains across model scales. Notably, the dominant miscalibration mode shifts structurally between regimes, overconfident failures in single-turn, underconfident successes in multi-turn, yet the same training principle addresses both successfully.

Poster
Proportionality in Ranking Compression

Daniel Halpern, Xinyu Liu, Evi Micha, Yu Peng Ng, Warut Suksompong

Expand abstract

Motivated by applications in reinforcement learning from human feedback, we consider a setting where a given collection of rankings needs to be compressed into a smaller collection of rankings. This can also be seen as a combination of two fundamental social choice frameworks: ranking aggregation and multiwinner voting. We propose three proportionality notions that capture different types of representation in this setting: positional proportionality, pairwise proportionality, and proportionality for solid coalitions (PSC). On the one hand, we show that positional proportionality and PSC are always satisfiable for any input rankings and target compression size, and a desired output can be found in polynomial time. On the other hand, we prove that pairwise proportionality cannot be satisfied in general, but can nevertheless be attained when the input rankings are single-peaked or single-crossing.

Poster
Programmatic Reasoning with Structural Schema: A Unified Framework for Multi-Table Inference

Jialin Chen, Brandon Mayer, Michael Galkin, Sami Abu-El-Haija, Rex Ying, Bryan Perozzi

Expand abstract

Large language models (LLMs) have shown strong performance in single-table reasoning but struggle to generalize to multi-table scenarios due to schema heterogeneity, long input contexts, and the lack of structural awareness. We introduce StrucTab-R1, a schema-centric reasoning framework that decouples relational reasoning from raw data exposure. Our approach encodes relational databases as heterogeneous schema graphs, where tables, columns, and foreign-key constraints form distinct node and edge types, and employs a heterogeneous graph encoder with question-conditioned cross-attention pooling to distill compact, semantically grounded schema tokens. Rather than serializing table rows into the context, the LLM plans over the schema representation and emits a Chain-of-Execution: an executable reasoning trace that first identifies the relevant schema subgraph and then performs a sequence of composable tool-function calls (e.g., filter, join, aggregate). These calls are executed externally, with compact observations returned to support subsequent reasoning steps. This design confines the model's attention to relevant subgraphs, mitigating hallucinations in multi-hop join reasoning. To further improve reliability, we train the model with supervised traces followed by structure-aware reinforcement learning, where the reward jointly optimizes answer correctness, schema-region precision, and execution consistency. Experiments on single-table and multi-table benchmarks show that StrucTab-R1 improves execution accuracy over strong general-purpose and table-specialized LLMs and remains effective on large-schema settings where raw-context baselines fail.

Poster
ProgramBench: Can Language Models Rebuild Programs From Scratch?

John Yang, Kilian Lieret, Jeffrey Ma, Parth Thakkar, Dmitrii Pedchenko, Sten Sootla, Emily McMilin, Pengcheng Yin, Rui Hou, Gabriel Synnaeve, Diyi Yang, Ofir Press

Expand abstract

Turning ideas into full software projects from scratch has become a popular use case for language models. Agents are being deployed to seed, maintain, and grow codebases over extended periods with minimal human oversight. Such settings require models to make high-level software architecture decisions. However, existing benchmarks measure focused, limited tasks such as fixing a single bug or developing a single, specified feature. We therefore introduce ProgramBench to measure the ability of software engineering agents to develop software holisitically. In ProgramBench, given only a program and its documentation, agents must architect and implement a codebase that matches the reference executable’s behavior. End-to-end behavioral tests are generated via agent-driven fuzzing, enabling evaluation without prescribing implementation structure. Our 200 tasks range from compact CLI tools to widely used software such as FFmpeg, SQLite, and the PHP interpreter. We evaluate 9 LMs and find that none fully resolve any task, with the best model passing 95% of tests on only 3% of tasks. Models favor monolithic, single-file implementations that diverge sharply from human-written code.

Poster
Principled Policy Optimization for LLMs via Self-Normalized Importance Sampling

HuiYang Shao, Xintong Zhang, Junyi Liu

Expand abstract

Reinforcement Learning from Human Feedback (RLHF) is a key technique for aligning Large Language Models (LLMs) with human preferences. While Proximal Policy Optimization (PPO) is the standard algorithm, its reliance on a critic network incurs significant memory and computational costs. This has motivated the development of critic-free alternatives such as Group Relative Policy Optimization (GRPO) and Group Sequence Policy Optimization (GSPO). However, these methods suffer from a critical trade-off: they either employ theoretically unsound, high-variance estimators (GRPO) or introduce systematic bias to achieve stability, causing them to optimize a perturbed objective (GSPO). In this paper, we introduce SNIB (Self-Normalized Importance Sampling with a Baseline), a novel critic-free algorithm that addresses this dilemma by offering a method that is both stable and asymptotically correct. SNIB leverages principled self-normalized importance sampling to achieve the stability of modern methods without sacrificing asymptotic correctness. We provide a comprehensive theoretical analysis, proving that SNIB's gradient estimator is consistent with a finite-sample bias that decays as $O(1/G)$ under self-normalization. Furthermore, we demonstrate its superior robustness to reward model uncertainty and show that it preserves the principled trade-off between reward maximization and KL regularization, a property that is distorted by biased estimators. Our work establishes a theoretically-grounded foundation for building more stable and reliable critic-free RLHF algorithms.

Poster
PriFT: Prior-Support Guided Token Reweighting for Supervised Fine-Tuning

Ke Wang, Shuangqi Li, Mathieu Salzmann, Pascal Frossard

Expand abstract

Supervised fine-tuning (SFT) is computationally efficient and broadly applicable, but often shows weaker generalization than reinforcement learning (RL). A key limitation is its off-policy objective: SFT fits fixed demonstrations token by token, including targets that may be poorly aligned with the model's pretrained distribution. Recent token-reweighted SFT methods address this issue by assigning larger training weights to tokens that better align with the model's predictive distribution, using statistics such as target-token probability or entropy. However, computing these statistics from the online model being fine-tuned makes token weights trajectory-dependent, as the model's distribution rapidly departs from the pretrained model and induces self-reinforcing reweighting dynamics. We propose PriFT, Prior-support guided Fine-Tuning, which derives token weights from a frozen pretrained reference to obtain a stable reweighting signal unaffected by online fine-tuning dynamics. This signal estimates prior support: the extent to which each target token is supported by the pretrained model before task-specific adaptation. Across multiple existing weighting and selection rules, replacing online statistics with pretrained statistics consistently improves performance. We introduce two instantiations: PriFT-prob, which uses pretrained target-token probability, and PriFT-mass, which selects tokens by relative support under the pretrained distribution. Extensive experiments on mathematical reasoning, code generation, and medical question answering show that PriFT achieves state-of-the-art SFT results and provides a better initialization for subsequent RL training.

Poster
Prevailing Bisimulation Metric Learning Is Biased: Implicit Regularization and Its Remedy

Junqi Lu, Ruixiang Sun, Xin Li, Gaopeng Peng, Ao Zhang, Mingzhong Wang

Expand abstract

Deep bisimulation metric learning has emerged as a principled framework for learning robust state representations in reinforcement learning. However, prevailing methods rely on sample-based regression objectives that exhibit a critical stability flaw in stochastic environments. We identify this flaw as a manifestation of the **double sampling problem**: minimizing mean squared error against single-sample stochastic targets introduces an irreducible variance bias. Through the lens of stochastic differential equations (SDEs), we rigorously prove that this bias leads to a **Variance Trap**, wherein target variance scales with the encoder's sensitivity. This variance effectively acts as an implicit Jacobian regularizer, driving representation collapse in the presence of transition noise. To address this issue, we propose Bisimulation Saddle-Point Optimization (BSPO), a debiased framework that reformulates the bisimulation metric learning objective as a primal-dual saddle-point problem. By introducing an Auxiliary Network (AuxNet) to estimate the expected Bellman error, BSPO decouples the structural error from transition noise, eliminating the bias without requiring physically infeasible double sampling. To the best of our knowledge, this work is the first to identify that prevailing sample-based bisimulation metric learning objectives are systematically biased and to provide a practical debiasing framework. Our empirical evaluations on stochastic continuous/discrete control tasks and visually distracting environments demonstrate that BSPO significantly outperforms representative baselines, particularly in high-noise regimes where conventional methods fail to converge.

Poster
Preserving Exploration for LLM Reasoning via Mean Order-Statistic Alignment

Ruotian Peng, Yi Ren, Long-Fei Li, Yandong Wen

Expand abstract

Pre-trained large language models exhibit strong exploration ability, as evidenced by high pass@K scores, yet this ability degrades significantly during reinforcement learning (RL) fine-tuning due to mode collapse. A natural remedy is to impose token-wise KL divergence between the fine-tuned and base models; however, we argue that such objective introduces a fundamental mismatch, where token-wise KL operates locally at individual token positions, while exploration is a global behavior that emerges over sequences as a whole. Motivated by this insight, we propose Mean Order-Statistic Alignment (MOSA), which preserves exploration by aligning the fine-tuned model's global behavior to that of the base model. Specifically, MOSA captures global exploration through the mean order-statistic profile, which is obtained by computing the order statistics of each token's posterior over the vocabulary and averaging them across tokens at each rank position. The profile, designed a valid probability distribution, directly admits a principled KL-based objective. For efficiency, we further approximate the profile using only the top-K probabilities and a residual tail bucket, yielding a compact implementation with minimal computational overhead. Experiments on Countdown, 6 mathematical reasoning, and 2 coding benchmarks show that MOSA discovers more diverse plausible solutions, improves both pass@1 and pass@K, and scales effectively to models up to 32B parameters.

Poster
Preference Conditioned Multi-Objective Reinforcement Learning: Decomposed, Diversity-Driven Policy Optimization

Tanmay Ambadkar, Sourav Panda, Shreyash Kale, Jonathan Dodge, Abhinav Verma

Expand abstract

Multi-objective reinforcement learning (MORL) seeks to train agents capable of balancing conflicting objectives. While single preference-conditioned policies offer a highly scalable solution, existing approaches remain brittle in practice, frequently failing to recover dense Pareto fronts. We demonstrate that this failure stems from two structural pathologies: destructive advantage cancellation caused by premature Early Scalarization (ES), and representational mode collapse across the preference space. To overcome these bottlenecks, we introduce \toolname{}, a PPO-based framework that fundamentally reorganizes multi-objective optimization. By preserving per-objective learning signals through a decomposed pipeline and integrating preferences only after trust-region stabilization (Late-Stage Weighting), \toolname{} improves credit assignment under conflicting objectives. Concurrently, a scaled diversity regularizer encourages behavioral divergence proportional to preference distance. \toolname{} operates entirely within the efficient linear scalarization regime shared by standard deep MORL baselines. By reducing information loss caused due to linear scalarization rather than relying on expensive non-linear utility functions, it suggests that optimization bottlenecks play a significant role. Across available standard benchmarks, including high-dimensional and many-objective environments, \toolname{} consistently discovers broader, higher-quality Pareto fronts than prior methods, exceeding state-of-the-art hypervolume and expected utility using a single deployable policy.

Poster
PreFT: Prefill-only finetuning for inference efficiency

Andrew Lanpouthakoun, Aryaman Arora, Zhengxuan Wu, Dhruv Pai, Benjamin Keigwin, Dan Jurafsky, Chris Potts

Expand abstract

Large language models can now be personalised efficiently at scale using parameter efficient fine-tuning methods (PEFTs), but \emph{serving} user-specific PEFTs harms throughput, even with specialised kernels and memory management techniques. This is because, theoretically and empirically, a mismatch exists between prefill (processing a large number of tokens at once) and decode (generating a single token autoregressively): the latter has far lower throughput when serving multiple adapters. Rather than optimising performance relative to parameter count, for efficient multi-adapter serving, we instead ought to optimise performance relative to \textit{serving throughput}. We therefore propose \textbf{\texttt{PreFT}} (Prefill-only Finetuning), wherein we only apply the adapter to prefill tokens and discard it afterwards. \texttt{PreFT} significantly increases throughput with minimal effect on performance. We develop and release an efficient implementation of two prefill-only PEFTs, LoRA and ReFT, on the vLLM inference engine. We first show that serving multi-user \texttt{PreFT}s is vastly more efficient than traditional PEFTs ($1.90\times$ the throughput when serving $512$ adapters on Llama 3.1 70B). Then, we compare the performance of prefill-only vs.~all-token adapters on a variety of supervised finetuning and reinforcement learning tasks with LMs at varying scales. On SFT, we observe that the evaluation loss of \texttt{PreFT}s is higher, but can be compensated by increasing rank with nearly no reduction in throughput. On RL, we consistently find that \texttt{PreFT}s approach parity with standard PEFTs. Together, this work validates prefill-only adaptation of LLMs as a more favourable accuracy--throughput tradeoff than existing PEFTs for personalised serving.

Poster
Power Distribution Bridges Sampling, Self-Reward RL, and Self-Distillation

Akiyoshi Tomihari, Issei Sato

Expand abstract

Recent analyses question whether reinforcement learning (RL) is responsible for strong reasoning in large language models (LLMs). At the same time, distillation and inference-time sampling, including power sampling, have emerged as effective ways to improve LLM performance. However, the relationship among RL, distillation, and sampling remains unclear. In this study, we focus on the power distribution, the target distribution of power sampling, and show that the power distribution bridges sampling, self-reward KL-regularized RL, and self-distillation. From the sampling perspective, we show that inexpensive local approximations cannot reproduce sequence-level power without information about possible suffixes. From the RL perspective, the power distribution is the closed-form optimizer of KL-regularized RL when the model's sequence-level log-probabilities are used as the reward. This identification leads to \emph{power self-distillation}, an offline distillation surrogate that shares the same target distribution and amortizes the cost of power sampling into supervised training on teacher samples. We further show that power self-distillation can achieve self-reward sharpening, while improvement in a downstream true reward is governed by the covariance between true reward and self-reward under the power distribution. Experiments on reasoning tasks support our analysis: power sampling raises self-reward, true-reward gains depend on alignment with self-reward, and power self-distillation can match or exceed the performance of power sampling at much lower inference cost.

Poster
Posterior Sampling-based Online Learning for Episodic POMDPs

Dengwang Tang, Dongze Ye, Rahul Jain, Ashutosh Nayyar, Pierluigi Nuzzo

Expand abstract

Learning in POMDPs is known to be significantly harder than in MDPs. In this paper, we consider the online learning problem for episodic POMDPs with unknown transition and observation models. We propose a Posterior Sampling-based reinforcement learning algorithm for POMDPs (PS4POMDP), which is much simpler and more implementable compared to state-of-the-art optimism-based online learning algorithms for POMDPs. We show that the Bayesian regret of the proposed algorithm scales as the square root of the number of episodes and is polynomial in the other parameters. In a general setting, the regret scales exponentially in the horizon length $H$, and we show that this is inevitable by providing a lower bound. However, when the POMDP is undercomplete and weakly revealing (a common assumption in the recent literature), we establish a polynomial Bayesian regret bound. We finally propose a posterior sampling algorithm for multi-agent POMDPs, and show it too has sublinear regret.

Poster
Policy-Level Exploration for Coordinated Multi-Agent Reinforcement Learning

Yucong Zhang, Chao Yu

Expand abstract

In multi-agent reinforcement learning (MARL), agents must not only discover high-reward behaviours but also reliably reproduce coordinated behaviours over time. A fundamental challenge arises from a mismatch between action-level stochastic exploration and temporally extended coordination. To address this issue, we propose a policy-level exploration framework with Deterministic Action Execution (DAE) for MARL, which shifts locally random action exploration to temporally extended policy selection via an option-based exploration mechanism. Specifically, at each time step, each agent selects a policy from a shared policy set\textemdash comprising one exploitation policy and multiple exploration policies\textemdash using a learned option policy, and actions are executed deterministically conditioned on the selected policy. By maintaining deterministic execution under temporally extended policy selection, DAE facilitates temporally coordinated behaviours across multiple time steps. Comprehensive experiments on four multi-agent tasks\textemdash Predator Prey, StarCraft II micro management challenge (SMAC), SMACv2, Google Research Football, and Multi-Agent Coordination benchmark (MACO)\textemdash demonstrate that DAE outperforms mainstream baselines, achieving superior sample efficiency and final learning performance.

Poster
Policy-DRIFT: Dynamic Reward-Informed Flow Trajectory Steering

Atharva Mahajan, Abhijeet Yashwant Vishwasrao, Yuning Wang, Ricardo Vinuesa

Expand abstract

Skin-friction drag induced by wall-bounded turbulent flows accounts for a substantial fraction of energy consumption across commercial aerospace, wind energy, and marine transport. Its active reduction is one of the highest-value targets in engineering fluid dynamics. Deep reinforcement learning (DRL) has emerged as the leading approach for real-time flow control, yet its performance ceiling is set not by algorithmic capability but by reward structure, the naive scalar objective does not optimally reflect the underlying physics. Policy-DRIFT bypasses this ceiling by relocating reward information from policy gradients to generative model inference: a conditional flow matching model (CFM) constructs a physically-grounded manifold of realisable flow states spanning multiple control regimes, Terminal Reward Guidance (TRG) steers samples toward reward-maximising targets at inference, and a lightweight DRL policy, structurally decoupled from reward quality, tracks these full-field targets via root-mean-squared error (RMSE) minimisation. The test case is turbulent channel flow simulated using direct numerical simulation (DNS) at friction Reynolds number of $\mathrm{Re}_\\tau = 180$, which is the canonical benchmark for wall-bounded turbulence. Policy-DRIFT achieves $49\\%$ drag reduction approaching the theoretical upper bound, which is $\\approx 16\\%$ higher than the DRL benchmark, while consuming 37$\\times$ less actuation energy. Our approach combines generative methods with active flow control, marking a paradigm shift towards controlling complex physical systems efficiently.

Poster
Policy Gradient Algorithms in Average-Reward Multichain MDPs

Jongmin Lee, Ernest Ryu

Expand abstract

While there is an extensive body of research analyzing policy gradient methods for discounted cumulative-reward MDPs, prior work on policy gradient methods for average-reward MDPs has been limited, with most existing results restricted to ergodic or unichain settings. In this work, we first establish a policy gradient theorem for average-reward multichain MDPs based on the invariance of the classification of recurrent and transient states. Building on this foundation, we develop refined analyses and obtain a collection of convergence and sample-complexity results that advance the understanding of this setting. In particular, we show that the proposed $\alpha$-clipped policy mirror ascent algorithm attains an $\epsilon$-optimal policy with respect to positive policies.

Poster
PoEM: Predicting New RL Outcomes from Existing Policies

Kimia Hamidieh, Giannis Daras, Antonio Torralba

Expand abstract

Pre-trained models are routinely post-trained with reinforcement learning (RL) to maximize specific rewards, such as human alignment, correctness, or instruction following. This post-training process is computationally intensive, sometimes unstable, and has to be run from scratch every time the reward model changes or when we want to combine multiple rewards. We hence ask: given a new reward function, is it possible to predict the RL outcomes without actually running RL on it? We answer this in the affirmative by introducing PoEM, a framework to predict the outputs of RL on a new reward function using a set of models already post-trained on other rewards. First, we show that if the new reward function can be written as a linear combination of existing ones, then the new policy in log-space can be written as a linear combination of the existing log-policies. Surprisingly, even in cases where the rewards are not linearly connected, we observe that often log-policies from RL training span an approximately low-rank subspace across rewards. To our benefit, the weighting coefficients for this combination can be estimated using only the reward or basis policy outputs on the samples. We turn these observations into an algorithm that takes post-trained models and a new reward function, and approximates the target RL policy without actually running any additional RL training. We experimentally validate our approach across synthetic and real experiments, spanning both text and image modalities.

Poster