← Search

Kaiwen Wang

13 accepted papers

2025

$Q\sharp$: Provably Optimal Distributional RL for LLM Post-Training

NeurIPS 2025poster

Reinforcement learning (RL) post-training is crucial for LLM alignment and reasoning, but existing policy-based methods, such as PPO and DPO, can fall short of fixing shortcuts inherited from pre-training. In this work, we introduce $Q\sharp$, a value-based algorithm for KL-regularized RL that guide…

Cited by 0SourcecodeScholar
2025

A Reductions Approach to Risk-Sensitive Reinforcement Learning with Optimized Certainty Equivalents

ICML 2025poster

We study risk-sensitive RL where the goal is learn a history-dependent policy that optimizes some risk measure of cumulative rewards. We consider a family of risks called the optimized certainty equivalents (OCE), which captures important risk measures such as conditional value-at-risk (CVaR), entro…

Cited by 0SourcePDFScholar
2025

Value-Guided Search for Efficient Chain-of-Thought Reasoning

NeurIPS 2025poster

In this paper, we propose a simple and efficient method for value model training on long-context reasoning traces. Compared to existing process reward models (PRMs), our method does not require a fine-grained notion of ``step,'' which is difficult to define for long-context reasoning models. By coll…

Cited by 0SourcecodeScholar
2024

Avoiding Object Damage in Robotic Manipulation

IROS 2024poster

The large-scale deployment of robotic manipulation systems in warehouses has highlighted the rare but costly problem of robot-induced object damage. We present a system that uses a classification model to predict whether an object will get damaged during robotic manipulation. The model uses object a…

Cited by 0SourceScholar
2024

Conditional Language Policy: A General Framework For Steerable Multi-Objective Finetuning

EMNLP 2024finding

Reward-based finetuning is crucial for aligning language policies with intended behaviors (*e.g.*, creativity and safety). A key challenge is to develop steerable language models that trade-off multiple (conflicting) objectives in a flexible and efficient manner. This paper presents Conditional Lang…

Cited by 17SourcePDFScholar
2024

Efficient and Sharp Off-Policy Evaluation in Robust Markov Decision Processes

NeurIPS 2024poster

We study the evaluation of a policy under best- and worst-case perturbations to a Markov decision process (MDP), using transition observations from the original MDP, whether they are generated under the same or a different policy. This is an important problem when there is the possibility of a shift…

2024

More Benefits of Being Distributional: Second-Order Bounds for Reinforcement Learning

ICML 2024poster

In this paper, we prove that Distributional Reinforcement Learning (DistRL), which learns the return distribution, can obtain second-order bounds in both online and offline RL in general settings with function approximation. Second-order bounds are instance-dependent bounds that scale with the varia…

Cited by 15SourcePDFScholar
2024

Switching the Loss Reduces the Cost in Batch Reinforcement Learning

ICML 2024poster

We propose training fitted Q-iteration with log-loss (FQI-LOG) for batch reinforcement learning (RL). We show that the number of samples needed to learn a near-optimal policy with FQI-LOG scales with the accumulated cost of the optimal policy, which is zero in problems where acting optimally achieve…

Cited by 6SourcePDFScholar
2023

The Benefits of Being Distributional: Small-Loss Bounds for Reinforcement Learning

NeurIPS 2023poster

While distributional reinforcement learning (DistRL) has been empirically effective, the question of when and why it is better than vanilla, non-distributional RL has remained unanswered. This paper explains the benefits of DistRL through the lens of small-loss bounds, which are instance-dependent b…

2022

Deep Multi-Modal Structural Equations For Causal Effect Estimation With Unstructured Proxies

NeurIPS 2022accept

Estimating the effect of intervention from observational data while accounting for confounding variables is a key task in causal inference. Oftentimes, the confounders are unobserved, but we have access to large amounts of additional unstructured data (images, text) that contain valuable proxy signa…

Cited by 15SourcePDFScholar
2022

Doubly Robust Distributionally Robust Off-Policy Evaluation and Learning

ICML 2022spotlight

Off-policy evaluation and learning (OPE/L) use offline observational data to make better decisions, which is crucial in applications where online experimentation is limited. However, depending entirely on logged data, OPE/L is sensitive to environment distribution shifts — discrepancies between the…

2022

Learning Bellman Complete Representations for Offline Policy Evaluation

ICML 2022oral

We study representation learning for Offline Reinforcement Learning (RL), focusing on the important task of Offline Policy Evaluation (OPE). Recent work shows that, in contrast to supervised learning, realizability of the Q-function is not enough for learning it. Two sufficient conditions for sample…