Search papers, labs, and topics across Lattice.
This paper reformulates decision-focused learning (DFL) for MDPs by casting the problem as an occupancy measure-based linear program, overcoming the scalability limitations of prior KKT-based methods. To address gradient discontinuities and state-space scaling, they introduce an augmented Lagrangian surrogate with random row sketching and a learnable soft state-aggregation layer. Experimental results demonstrate lower regret and significantly reduced computational cost compared to existing DFL and two-stage baselines across various tasks.
Occupancy measure formulation enables decision-focused learning in MDPs with closed-form gradients, overcoming KKT-based scalability bottlenecks and achieving state-of-the-art regret.
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.
