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

2026-05-31

AI Summary

This work addresses decision-focused learning in sequential contextual linear optimization under partial feedback. The authors propose an online policy learning method that jointly optimizes a predictive model and the downstream linear decision task through a stochastic predict-and-optimize framework. The key innovation lies in extending decision-focused learning to the online, partial-feedback setting for the first time and introducing a hybrid gradient estimator that integrates a scoring function with a plug-in architecture, effectively leveraging structural information from the downstream optimization problem. Empirical results demonstrate that the proposed approach significantly outperforms contextual bandit baselines across multiple tasks—including top-k selection, shortest path, combinatorial pricing, and real-world energy dispatch—achieving substantially lower cumulative regret.

Abstract

Decision-focused learning (DFL) trains predictive models by optimizing downstream decision quality rather than standalone prediction accuracy. For contextual linear optimization, most existing DFL methods assume offline data and full observations of the objective cost vector. We develop an on-policy learning method for sequential contextual linear optimization under partial feedback, generalizing the standard bandit feedback setting. 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. To update this distributional model, we 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 decision-focused plug-in component that uses an auxiliary nuisance estimate of the latent cost vector to exploit the downstream optimization structure, becoming more informative as the estimate improves. We prove an $\mathcal{O}(T^{-1/2})$ bound on the average squared policy-gradient norm, matching the standard non-convex SGD rate. Experiments on top-$k$ selection, shortest path, combinatorial pricing, and a real-data energy-scheduling benchmark show that the hybrid gradient approach achieves lower cumulative regret than contextual-bandit-style baselines across all benchmarks, using both Gaussian and richer conditional generative models. Code is available at https://github.com/Joeyetinghan/on-policy-bandit-dfl.

Problem

• decision-focused learning
• contextual linear optimization
• partial feedback
• on-policy learning
• bandit feedback

Innovation

• decision-focused learning
• on-policy learning
• partial feedback
• hybrid gradient estimator
• contextual linear optimization