# Decision-Focused Learning in MDPs: An Occupancy Measure Approach

arXiv: https://arxiv.org/abs/2610.08384 · HTML: https://mlsanity.com/paper/2610.08384/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 4 of its 53 references that are in the map.
- First posted on arXiv 6 Oct 2026.

## Builds on (most cited first, 4 of 4)

- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [Sketching as a Tool for Numerical Linear Algebra](https://mlsanity.com/paper/1411.4357/) (2014)
- [Average-reward model-free reinforcement learning: a systematic review and literature mapping](https://mlsanity.com/paper/2010.08920/) (2020)
- [Learning Decision-Focused Uncertainty Sets in Robust Optimization](https://mlsanity.com/paper/2305.19225/) (2023)

## Abstract (arXiv)

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.