Decision-Focused On-Policy Learning for Contextual Linear Optimization with Partial Feedback

arXiv:2606.01081 · cs.LG · Submitted 2026-05-31 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: Today's paper: "Decision-Focused On-Policy Learning for Contextual Linear Optimization with Partial Feedback".

Jane: The gist:

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So we're looking at this paper now, "Decision-Focused On-Policy Learning for Contextual Linear Optimization with Partial Feedback". It tackles how to train predictive models by focusing on the quality of the decisions they actually make instead of just being right about a prediction.

Jane: It seems like the main gap they are filling is that most existing decision-focused learning methods assume you have all the data upfront and see every single cost vector, but this paper looks at sequential optimization problems where you only get partial feedback based on what you choose to do.

Lu: Exactly. They develop an on-policy learning method specifically for sequential contextual linear optimization under partial feedback, which is a setting they say isn't systematically captured by any existing framework yet.

Meng: So they propose a policy that learns to sample a cost-vector prediction from some conditional distribution and then immediately solves the resulting linear optimization problem based on that prediction. That sounds like it has real potential for online decision making where you can't wait for all the data.

Lalam: And to update this model of predictions, they introduce a two-component hybrid gradient estimator. The first part is a score function estimator, which gives an unbiased but potentially high-variance policy gradient estimate.

Tom: High variance is always a concern when you're dealing with online learning signals, but the second component seems like it's designed to help with that by using a decision-focused plug-in component. This part uses an auxiliary estimate of the latent cost vector to exploit the structure of the downstream optimization problem.

Jane: So they have this two-part update mechanism: one that gives you a clean but noisy signal, and another one that tries to make sense of the cost vector using some other guess, which gets better as that guess improves.

Lu: The analysis shows they achieve an O(T-one/two) bound on the average squared policy-gradient norm for this hybrid update, which matches the standard non-convex stochastic gradient descent rate. That's a solid theoretical baseline for convergence.

Meng: An O(T-one/two) rate is good because it means the learning process converges at a known speed, similar to what you see in standard SGD settings when you don't have all the data at once. But what about actual deployment?

Tom: The experiments on top-k selection, shortest path, combinatorial pricing, and even a real-data energy-scheduling benchmark show that this hybrid gradient approach under bandit feedback achieves lower mean cumulative regret than the contextual-bandit baselines across all four benchmarks.

Jane: And even compared to linear Thompson sampling on three of those benchmarks, it performs better when you're only seeing partial feedback. That suggests this method is robust and useful in real-world online scenarios where you don't have perfect information about the true cost vector.

Lu: Plus, they show that this framework works even with richer conditional generative cost models for the plug-in component, which expands its applicability beyond just simple linear models.

Meng: From a deployment standpoint, it sounds like this isn't just an academic exercise; it’s designed to handle the messiness of real online data where you only get partial feedback from your choices.

Lalam: The paper concludes that this hybrid update couples an unbiased score function update with a decision-focused plug-in update, and under the stated assumptions, its fixed-weight update reaches first-order stationarity at rate O(T-one/two).

Tom: So in simple terms, the authors are showing that by combining a standard gradient estimate with one tailored to exploit how decisions affect future costs, you get a solid convergence guarantee.

Jane: The real implication here is that this method provides more flexibility than some other approaches; it doesn't require you to have a closed-form posterior or an enumerable feasible set, just a distributional cost model and a linear optimization oracle.

Lu: It also suggests that the hybrid gradient is conditionally biased, but only through the nuisance error, which means they are exploring how that bias might improve performance if you adjust those learning parameters over time.

Meng: If you're deploying this in production, knowing it can handle partial feedback better than pure bandit baselines on all those benchmarks is a strong signal about its practical utility.

Lalam: The final results show that semi-bandit and full-information feedback lower regret relative to pure bandit feedback on all four problems, which shows the method's flexibility across different feedback settings.

Tom: It's clear that this paper is focused on building a tool that works in the messy, partial-feedback world of sequential optimization, moving beyond just having accurate predictions to actually making good decisions under uncertainty.

Conclusion: Jane: So, to wrap up, this paper is about training AI models to make better decisions in online optimization problems when you don't get all your data at once—partial feedback.

Tom: Right, so it’s about Decision-Focused On-Policy Learning for Contextual Linear Optimization with Partial Feedback. The authors are focusing on how to optimize the decision quality itself rather than just getting a prediction score right.

Lu: They introduce a method where you sample a cost vector prediction and then solve the optimization problem based on that, which is different from standard methods that wait for full information.

Jane: And they do this using this two-part gradient estimator—one part being the standard score function and another one being decision-focused plug-in to help exploit the structure of the linear optimization.

Tom: The key result is that under bandit feedback, this hybrid approach beats contextual bandit baselines on all four benchmarks, and even linear Thompson sampling on three of them. That's a solid comparison with concrete numbers there.

Meng: I’m interested in what they say about the regret bound, because for me at scale, we need to know the convergence speed. They show that the average squared policy-gradient norm is bounded by O(T negative one over two), which matches standard non-convex SGD rates.

Lalam: That rate means their fixed-weight update reaches first-order stationarity in that time frame, which is a strong theoretical guarantee for sequential learning under these conditions.

Jane: So, what does this mean practically? It suggests we can get better performance in real online scheduling or pathfinding when the feedback is incomplete, and this method offers flexibility with different types of cost models.

Tom: It means we don't necessarily need a perfect posterior distribution to do good work; we just need a distributional model that allows us to estimate those gradients effectively.

Lu: The implication is that this framework can be extended beyond linear feedback and risk-neutral objectives, which opens the door for more complex real-world problems.

Jane: That sounds like it could help us build more robust AI systems for environments where data collection is expensive or slow to get. What happens when we move from linear optimization to something non-linear?

Wyame Benslimane, Tinghan Ye, Pascal Van Hentenryck, Paul Grigas

Department of Industrial Engineering and Operations Research University of California, Berkeley · H. Milton Stewart School of Industrial and Systems Engineering Georgia Institute of Technology

cs.LG

Submitted: 2026-05-31

Updated: 2026-10-07

Code: https://github.com/Joeyetinghan/on-policy-bandit-dfl

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 83/100

The gist: The gist: This paper develops an on-policy learning method for sequential contextual linear optimization under partial feedback by introducing a hybrid gradient estimator that combines an unbiased

Key concepts

Decision-Focused Learning (DFL)
DFL trains predictive models by focusing on optimizing the quality of downstream decisions rather than just maximizing standalone prediction accuracy. In this context, it means learning a policy that directly improves the final optimization outcome.
On-Policy Learning
This refers to a learning method where the agent learns from data generated by its current policy. The proposed method is on-policy, meaning it updates its knowledge based on decisions made by the policy being learned, which is suitable for sequential tasks under partial feedback.
Hybrid Gradient Estimator
This estimator combines two parts: a score function term that provides an unbiased but potentially high-variance gradient estimate of the cost function, and a decision-focused plug-in term that uses an auxiliary estimate to better exploit the structure of the downstream optimization problem.

Terminology

Summary

The gist: This paper develops an on-policy learning method for sequential contextual linear optimization under partial feedback by introducing a hybrid gradient estimator that combines an unbiased score function term with a decision-focused plug-in term to exploit downstream optimization structure.

How it works

The proposed method learns a stochastic predict-then-optimize policy that samples a cost-vector prediction from a conditional distribution and solves the resulting downstream linear optimization problem. To update this distributional model, the authors introduce a two-component hybrid gradient estimator. The first component is a score function estimator, which provides an unbiased but potentially high-variance policy gradient estimate. The second is a decisionfocused plug-in component that uses an auxiliary nuisance estimate of the latent cost vector to exploit the downstream optimization structure.

Key Components and Estimators

The learning process involves several key mathematical constructs and estimators. The score function estimator is defined as gscore t:= yt ∇θ log pθt(ˆct xt), where cˆt ∼ pθt(· xt), wt = w⋆(ˆc), yt = c⊤t wt. This estimator satisfies E[gscore t] = ∇θVcost(θt). The decision-focused plug-in component uses a nuisance estimate fϕ to approximate the conditional mean cost vector f⋆(x) = E[c x].

Convergence and Stationarity

The convergence analysis targets first-order stationarity for non-convex objectives. The hybrid update is defined by a convex mixture of the two signals: gαt = α gscore t + (1 − α) gplug-in t. The analysis shows that in both cases, the optimization term is O(T−1/2), matching the standard non-convex SGD rate.

Experimental Results and Benchmarks

Experiments were conducted on four online contextual optimization benchmarks under bandit feedback. These include top-k selection, shortest path, combinatorial pricing, and a real-data energy-scheduling benchmark. The results show that under bandit feedback, DFHPG achieves lower mean cumulative regret than contextual-bandit baselines on all four benchmarks and than linear Thompson sampling on three.

Conclusion of Findings

DFHPG couples an unbiased score function update with a decision-focused plug-in update, and under the stated assumptions its fixed-weight update reaches first-order stationarity at rate O(T−1/2). Under bandit feedback, DFHPG and the plug-in-only DFHPG-0 reduce cumulative regret relative to the contextual-bandit baselines and the score-only ablation on all four benchmarks. DFHPG is also more flexible than linear bandit algorithms: it needs neither a closed-form posterior (LinTS) nor an enumerable feasible set (OFUL), only a distributional cost model with the required gradient estimators and a linear optimization oracle. Extending the framework beyond linear feedback and risk-neutral objectives is a natural next step. The paper demonstrates that the hybrid gradient is conditionally biased, but only through the nuisance error, as formalized by Lemma C.3. This suggests that adaptive schedules αt → 0 may improve empirical performance. The final results show that semi-bandit and full-information feedback lower regret relative to pure bandit feedback on all four problems. The paper concludes by showing that the surrogate plug-in gradient estimator can be used for any surrogate whose subgradient satisfies the regularity bounds of Assumption 9. The paper also details the implementation of various surrogate losses and generative cost models for the plug-in component. The paper provides extensive numerical evidence showing that DFHPG outperforms several baselines across various problem difficulties and feedback modes.

Improvements for AI systems

  1. Better handling of partial feedback in sequential decision-making by learning a stochastic predict-then-optimize policy. This improved system can operate effectively in real operational systems like vehicle routing or energy scheduling where only partial information about the true cost vector is revealed after a decision is made, as described by the paper's method: Our method learns a stochastic predict-then-optimize policy that samples a cost-vector prediction from a conditional distribution and solves the resulting downstream linear optimization problem.

  2. Incorporation of an adaptive mixing weight schedule for the hybrid gradient estimator. This allows the system to dynamically balance the unbiased but high-variance score function estimate against the lower-variance decision-focused plug-in term, as seen in Algorithm 1: "The schedule decays αt from αmax toward αmin, weighting the unbiased score function gradient more heavily early on (when the nuisance estimate is unreliable) and the lower-variance plug-in gradient later (as the nuisance estimate stabilizes)."

  3. Robustness to different cost model architectures via surrogate losses. The system can leverage diverse predictive models, such as conditional normalizing flows or diffusion models, by employing differentiable surrogates like SPO+, Perturbed Fenchel–Young loss (PFYL), or DPO, ensuring the learning process is not tied to a single prior distribution: We extend the decision-focused plug-in term to accommodate a broad class of DFL surrogates (e.g., the Smart Predict-then-Optimize surrogate (SPO+), Perturbation Gradient (PG), and Perturbed Fenchel–Young loss (PFYL)).

  4. Improved regret performance under various feedback regimes by leveraging the plug-in component for semi-bandit settings. The system can achieve lower cumulative regret when only partial information is revealed, as demonstrated by the result: Semi-bandit feedback lowers the cumulative regret of DFHPG by roughly 26–58% on the synthetic benchmarks and ∼39% on energy, and DFHPG-0 has lower mean regret than DFHPG in every column.

  5. Adaptation to real-world data characteristics using nuisance estimation. The system can learn the underlying conditional cost structure from observed feedback, which is crucial when full cost vectors are not available: "We introduce an auxiliary cost-vector regression model fϕ: X → R dw (the nuisance estimator), parameterized by ϕ ∈ Φ ⊂ R dϕ, which approximates the conditional cost mean f⋆(x) = E[c x] by minimizing the following least-squares objective based on the observed feedback v = H(w)c."

Sources

Related papers