Decision-Focused Learning in MDPs: An Occupancy Measure Approach

Zihao Zhao; Ashwath K. Karunakaram; Ali Eshragh; Yuexing Li; Kai Wang

cs.LGOct 6, 2026

Decision-Focused Learning in MDPs: An Occupancy Measure Approach

Authors: Zihao Zhao, Ashwath K. Karunakaram, Ali Eshragh, Yuexing Li, Kai Wang

Organizations: Georgia Institute of Technology, Atlanta, GA · Johns Hopkins Carey Business School, Washington, DC · International Computer Science Institute, Berkeley, CA

Abstract

In this work, we consider decision-focused learning (DFL) for a Markov decision process (MDP), where existing methods differentiate through the KKT conditions of the Bellman equation and require solving a linear system over all state-action pairs, limiting its scalability. We address this by reformulating the MDP as an occupancy measure-based linear program (LP), whose feasible region is induced by predicted dynamics, and we derive a closed-form gradient by identifying the active constraints in the feasible polyhedron via the pivoting algorithm. This occupancy measure-based LP layer raises two challenges: (1) LP's solution gradient is discontinuous when active constraints change, and (2) the LP backward cost still scales with the state size, which is costly for large or continuous state spaces. We address the challenges with an augmented Lagrangian surrogate and smooth the boundary jumps by random row sketching of the constraints, and a learnable soft state-aggregation layer and its function-approximation generalization that scales the LP to large finite and continuous-state MDPs. Across multiple tasks, our methods reach lower regret than KKT-based DFL and two-stage baselines with significantly lower computation cost. The source code for all experiments is available at https://github.com/A-Eshragh/State_Aggregation_Project.

Figures & tables

Metric	KKT	LP (ours)	LP + State Agg (ours)
Core Matrix to Invert	KKT System ( [ 48 ] , Eq.(7))	Optimal Basis $H_{\theta}^{B}$	Aggregated Basis $\hat{H}_{\theta}^{B}$
Matrix Dimension	$\|\mathcal{S}\|\|\mathcal{A}\|\times\|\mathcal{S}\|\|\mathcal{A}\|$	$\|\mathcal{S}\|\times\|\mathcal{S}\|$	$M\times M$
Inversion Complexity	$\mathcal{O}((\|\mathcal{S}\|\|\mathcal{A}\|)^{\omega})$	$\mathcal{O}(\|\mathcal{S}\|^{\omega})$	$\mathcal{O}(M^{\omega})$

Table 1: Comparison of computational complexity between the KKT approach, our LP approach, and our soft state-aggregation extension. Here

M

denotes the number of aggregated states with

M\ll|{\mathcal{S}}|

and

\omega<2.373

denotes the matrix multiplication exponent.

Figure 1: Loss (top) and gradient (bottom) of three DFL objectives as

\lambda_{\mathrm{pred}}

varies; red dashed indicate

\lambda_{\mathrm{true}}

and green dots indicate the loss minimum. (1) The occupancy proxy has a shifted minimum. (2) The augmented Lagrangian has a minimum at the true parameter. (3) Sketch averaging empirically reduces boundary sensitivity.

Figure 2: Training reward across epochs on three tasks. DFL-QP and DFL-SA achieve the highest training rewards, while DFL-Sketch and the other LP-based variants are generally competitive with the two-stage baseline.

Figure 3: Test regret on three tasks (lower is better). DFL-Sketch improves over DFL-LP on Inventory and CartPole, showing the benefit of smoothing across basis changes. DFL-SA further achieves regret close to DFL-QP at substantially lower cost.

Figure 4: Total training time for different tasks. DFL-QP is the most expensive method because it differentiates through the full KKT system. DFL-Sketch adds overhead from solving multiple sketched LP s, while DFL-SA reduces this cost through state aggregation.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

	Inventory	Cliff Walking	CartPole
MDP
State space size $\|{\mathcal{S}}\|$	$36$ ( $s_{\max}=30$ , $b_{\max}=5$ )	$48$ ( $4{\times}12$ grid)	continuous, $\mathbb{R}^{4}$
Action space size $\|{\mathcal{A}}\|$	$31$ ( $a_{\max}=30$ )	$4$	$2$
Discount $\beta$	$0.95$	$0.95$	$0.99$
Horizon $T$	N/A	N/A	$200$
Predictor / features

Appendix

Table 2: Key hyperparameters and model architectures for each task.

Method	Mean regret $\downarrow$	SE	Wins vs. TS
Two-stage	14.75	0.29	–
DFL-LP	14.14	0.18	7/10
DFL-Sketch	11.34	1.44	8/10
DFL-QP	9.53	0.84	10/10
DFL-SA ( $M=6$ )	7.21	0.15	10/10
DFL-SA ( $M=12$ )	7.76	0.15	10/10

Appendix

Table 3: Paired Cliff Walking test regret over 10 random seeds. Lower regret is better. Win counts compare each method against the two-stage baseline on the same seed. SE denotes the standard error of the mean.

Method	Best-so-far reward
DFL-LP	$170.5\pm 32.7$
DFL-QP	$192.4\pm 12.6$
Two-stage (BC + LP, no DFL gradient)	$168.1\pm 24.8$
MBRL (no BC anywhere)	$136.8\pm 66.6$
MBRL + BC reward shaping ( $\alpha=2$ )	$127.3\pm 61.3$

Appendix

Table 4: CartPole best-so-far reward (mean

\pm

standard deviation). The MBRL rows use

5

seeds and the additional configuration described above; the table does not establish a matched-data comparison.

Method	Test regret $\downarrow$	Total (s)
Two-stage	$24.537\pm 4.573$	$68.2\pm 11.2$
DFL-QP	$19.121\pm 2.156$	$3649.4\pm 165.5$
DFL-LP	$23.056\pm 3.294$	$972.0\pm 34.0$
DFL-Feas	$22.834\pm 3.414$	$871.0\pm 15.7$
DFL-Sketch (Ours)	$19.434\pm 2.363$	$897.7\pm 33.2$
DFL-SA ( $M=10$ ) (Ours)	$22.659\pm 2.589$	$226.5\pm 8.0$

Appendix

Table 5: Larger inventory MDP,

|{\mathcal{S}}|\times|{\mathcal{A}}|=51\times 31=1581

(mean

\pm

standard error over

10

seeds).

Method	Test regret $\downarrow$	Total (s)
Two-stage	$53.788\pm 7.710$	$227.3\pm 6.4$
DFL-QP	$44.545\pm 5.823$	$34791.0\pm 756.4$
DFL-LP	$40.690\pm 4.946$	$19915.3\pm 527.1$
DFL-Feas	$62.904\pm 8.230$	$19832.6\pm 618.7$
DFL-Sketch (Ours)	$48.591\pm 5.718$	$13790.6\pm 513.4$
DFL-SA ( $M=10$ ) (Ours)	$42.568\pm 5.814$	$540.6\pm 3.8$

Appendix

Table 6: Larger inventory MDP,

|{\mathcal{S}}|\times|{\mathcal{A}}|=101\times 31=3131

(mean

\pm

standard error over

10

seeds).

$\rho$	Test regret $\downarrow$	Flow-constraint MSE
$1$	$15.629\pm 2.296$	$23.98$
$10$	$14.845\pm 1.877$	$23.07$
$10^{2}$	$14.124\pm 1.808$	$21.26$
$10^{3}$	$12.819\pm 1.185$	$18.00$
$10^{4}$ (used in the paper)	$12.552\pm 0.991$	$17.82$
$10^{5}$	$13.305\pm 1.458$	$17.26$

Appendix

Table 7: Sensitivity of DFL-QP on the inventory MDP to the penalty weight

\rho

(mean

\pm

standard error over

10

seeds).

LP layer	Test regret $\downarrow$	Infeasible train solves
Hard	$63.024\pm 4.911$	$0.0\,\%$
Proj-LP (no diagonal step)	$24.147\pm 5.758$	$61.8\,\%$
Proj-MDP (DFL-SA)	$14.966\pm 2.319$	$0.0\,\%$

Appendix

Table 8: Ablation of the diagonal projection on the inventory MDP (

M=10

; mean

\pm

standard error over

10

seeds).

Explore similar work

Decision-Focused On-Policy Learning for Contextual Linear Optimization with Partial Feedback
May 31, 2026Wyame Benslimane, Tinghan Ye, Pascal Van Hentenryck +1Contextual Bandit Framework On-Policy
Department of Industrial Engineering and Operations Research University of California, Berkeley · H. Milton Stewart School of Industrial and Systems Engineering Georgia Institute of Technology
Decision-Focused Learning via Tangent-Space Projection of Prediction Error
May 2, 2026Junhyeong Lee, Sangjin Jin, Yongjae LeePlan Regret Prediction Error
Department of Industrial Engineering, Ulsan National Institute of Science and Technology, Ulsan, South Korea.
Adaptive Resolving Methods for Markov Decision Processes with Function Approximations
May 17, 2025Jiashuo Jiang, Yinyu Ye, Yiming ZongMarkov Decision Processes Optimal Policies
Department of Industrial Engineering & Decision Analytics, The Hong Kong University of Science and Technology · Department of Management Science & Engineering, Stanford University

May 31, 2026, cs.LGJ/K move · Enter open · S save

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

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

Contextual Bandit Framework On-Policy Feedback Latent Variable

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

Decision-Focused Learning via Tangent-Space Projection of Prediction Error

Adaptive Resolving Methods for Markov Decision Processes with Function Approximations