Research · Overview

Research

We build adaptive, person-centered methods for decision-making in messy, high-stakes settings, drawing on reinforcement learning, uncertainty quantification, causal inference, and language model reasoning.

Most reinforcement learning results are reported on problems that were built to be measurable: the state is observed, the reward arrives on the next step, the data was collected by the agent being trained, and a bad action costs you nothing but a reset. Real decisions are not like that. The settings we care about are partially observed and confounded, the consequence of a choice can arrive weeks after the choice, the data was collected by somebody else for reasons you cannot fully reconstruct, and some actions cannot be taken back.

Each area below is a distinct methodological thread, but they share a commitment: we are less interested in clean problems with established benchmarks than in the contextual, sequential, partially observed problems that demand new thinking, and in methods that researchers and practitioners can actually use.

Three square grid maps side by side, sharing a colour scale that runs from -1 in dark red through white at 0 to +1 in blue. The first, labelled LifeGate Env, shows the layout: a mostly black grid with a grey wall segment part-way down the left side, a yellow block filling the lower right, a blue strip along the top right edge, a red strip down the whole right edge, and a small white square marking the agent near the bottom. The second and third show the learned dead-end value and rescue value over those same cells. The dead-end map is pale in the upper left and deepens to dark red across the lower right; the rescue map is the near mirror image, deep blue through the upper left and fading to white toward the lower right.

In a small navigation task, separately learned dead-end and rescue values mark which regions of the grid still allow a good outcome and which no longer do.

  • Offline RL
  • Irreversibility & risk
  • Clinical & scientific

Medical Dead-ends and Learning to Identify High-Risk States and Treatments

A probability density curve, labelled Z-pi of s and a, drawn over return value on the horizontal axis and density on the vertical axis. The distribution has two humps: a small one over low returns near -0.5 and a much larger one near +0.2. The low-return tail beneath the small hump is shaded grey and cut by two vertical dashed lines — the value-at-risk threshold at roughly -0.5, and to its left the conditional value-at-risk, the average of that shaded tail, at roughly -0.65.

Conditional value at risk summarises the low-return tail of a return distribution instead of its average, which is what lets the method flag unrecoverable states at a chosen risk tolerance.

  • Offline RL
  • Irreversibility & risk
  • Clinical & scientific

Risk Sensitive Dead-end Identification in Safety-Critical Offline Reinforcement Learning

A left-to-right pipeline schematic. On the left, three dashed coloured traces of a multivariate time series X of t, with open circles marking observations at unevenly spaced times t0, t1, t2 through tT and the gap between two of them marked delta t. An arrow leads into a tapering block labelled Backbone NN, then a narrow block labelled Evidential Layer, which emits Normal Inverse-Wishart parameters phi. Two arrows branch to the right: the first to a shaded fan of predicted trajectories widening and narrowing around the observations, labelled calibrated, temporally correlated uncertainty; the second to a pair of small panels, a contour map of the distribution over the mean and covariance beside a scatter of orange and blue points split by a curved boundary, labelled uncertainty-guided, robust classification.

An irregularly sampled clinical time series is mapped to the parameters of an evidential distribution, so the model reports uncertainty that stays correlated across time rather than at isolated points.

  • Clinical & scientific

Continuous Time Evidential Distributions for Irregular Time Series

A tree diagram of a three-part action. From a filled root node, three branches labelled a1 = 0, a1 = 1 and a1 = 2 lead to boxed action vectors [0,0,0], [1,0,0] and [2,0,0], so each node holds a full action with the not-yet-chosen entries left at a default of zero. A branch labelled a2 = 1 descends from [0,0,0] to [0,1,0], which in turn splits along a3 = 1 and a3 = 2 to [0,1,1] and [0,1,2]. A smaller copy of the whole tree sits at lower right with one root-to-leaf path drawn in solid black and every other branch greyed out.

Laying a combinatorial action out as a tree lets the policy pick one sub-action per level and score a single path, rather than enumerating every combination of sub-actions.

  • Offline RL
  • Structured actions

BraVE: Offline Reinforcement Learning for Discrete Combinatorial Action Spaces

A left-to-right block diagram of a policy network. A circle holding the state s and a box labelled Sub-Action Embedding Table both feed a tall State Conditioning block, whose state-aware embeddings pass into a Transformer Block marked as repeated L times. Its output X enters a dashed group marked as repeated for each of the A sub-actions, containing a Decision MLP that emits distribution parameters and a Sample Sub-Action step that draws sub-action a-i. A Concat block gathers the sampled sub-actions into the final action vector a.

Self-attention over sub-action embeddings, conditioned on the state, lets each sub-action be chosen with the others in view and in no fixed order.

  • Structured actions

SAINT: Attention-Based Policies for Discrete Combinatorial Action Spaces

A line chart of learning progress against wall-clock time, with time in minutes from 0 to 300 on the horizontal axis and percentage of the reference FIQL asymptotic score on the vertical axis. A horizontal dashed line marks the 100 percent target. The thick dark navy SPIN curve sits flat at zero during pre-training, jumps past the target at a vertical dashed line annotated "ASM Learning Complete" around ten minutes, and stays above it. The three baseline curves — Factored IQL, SAINT and AR-IQL — climb gradually from the start and only approach the target somewhere between roughly 150 and 300 minutes.

On the humanoid-stand task, pre-training a model of which joint actions are valid lets policy learning reach the reference score within about ten minutes, where the baselines take hours.

  • Offline RL
  • Structured actions

Improving and Accelerating Offline RL in Large Discrete Action Spaces with Structured Policy Initialization

A three-part diagram summarising how to spend a fixed sampling budget, divided by dashed vertical rules and annotated with three numbered findings in salmon callout boxes. On the left, a set of three-dimensional axes — parallel rollouts per problem going up, problems per batch and sequential iterations along the base — carries a pale blue plane with dashed contours and a red line climbing through a row of yellow points, labelled: more compute, larger numbers of rollouts, on both easy and hard problem sets. In the middle, two small line charts plotted against sequential iterations: rollouts per problem rising in an S-shape, and problems per batch falling in the mirror image, labelled: with the total batch size held fixed, more iterations means more rollouts and fewer problems per batch. On the right, another set of the same three-dimensional axes with a flat orange surface low across the base, labelled: problems per batch has only a marginal effect within a moderate range.

A fixed sampling budget for reinforcement learning on language models is split three ways — rollouts per problem, problems per batch, and sequential iterations — and the number of rollouts per problem is the allocation that should grow as the budget grows.

  • LLM reasoning

IsoCompute Playbook: Optimally Scaling Sampling Compute for LLM RL

Three line charts compare summed terminal-token probability over 200 training steps with no EOS fix, in orange, and with a semantic EOS fix, in green. For Qwen3, both curves rise initially, but the unfixed probability later collapses to zero while the corrected probability remains near 0.8. For Llama 3.2, both begin near zero; the unfixed curve stays there while the corrected curve rises to roughly 0.3. For Gemma 3, the unfixed curve falls to zero early, while the corrected curve climbs sharply after about 60 steps and remains near 1.

Treating equivalent EOS tokens as one semantic stopping action prevents or reverses the collapse in termination probability across Qwen, Llama and Gemma.

  • LLM reasoning

When EOS Tokens Disagree: Understanding Length Inflation in On-Policy Distillation

A left-to-right process diagram with a purple Inference engine generating rollouts from a prompt, a green Training engine receiving synchronized parameters and producing token confidence values, and a blue GRPO token loss receiving a reward signal. The orange CIS path computes importance ratios for individual tokens, clips anomalous ratios with confidence-aware thresholds, and feeds the reweighted loss back into the update loop.

CIS corrects the mismatch between inference and training engines by reweighting token-level learning signals with confidence-aware caps.

  • LLM reasoning

Rethinking Training-Inference Mismatch in LLM Reinforcement Learning: Where It Arises and How to Correct It

A three-stage pipeline runs from left to right. In Candidate Pool, a question is sent to multiple teacher models to produce K candidate reasoning trajectories, from which B will be selected. In Rank by Score, the student model scores and sorts the trajectories; the top B sit inside a green dashed box and receive weights proportional to the gap between their scores and the first unselected score, while the remaining trajectories sit inside a red dashed box with zero weight. In Weighted SFT, the selected trajectories and their weights form a weighted loss that is used to fine-tune the student.

LARK scores teacher-generated reasoning trajectories by how efficiently the student can learn from them, then selects and weights the top candidates for supervised fine-tuning.

  • Data curation

LARK: Learnability-Grounded Trajectory Selection for Efficient Reasoning Distillation

Research · Areas

Data curation

Choosing the training examples that make a model learn the right things.

Data selection and trajectory curation for efficient mid-training, where coverage and learnability matter as much as example quality.

  • LLMs
  • Data Curation
  • Mid-training
Read more