Decision-Focused Learning in MDPs: An Occupancy Measure Approach

2026-10-06

AI Summary

This study addresses the poor scalability of decision-focused learning in Markov decision processes caused by exhaustive traversal of the entire state space. To overcome this limitation, it proposes an occupancy measure-based linear programming reformulation. The core innovations include introducing an augmented Lagrangian surrogate combined with an occupancy measure LP layer to enable efficient gradient computation, employing randomized row sketching to smooth gradient discontinuities, and designing learnable soft state aggregation alongside neural network function approximation to handle continuous state spaces. Experimental results demonstrate that the proposed method significantly reduces computational costs in multi-task scenarios while achieving lower regret compared to KKT-based baselines and two-stage approaches.

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.

Problem

• Decision-Focused Learning
• Markov Decision Process
• Scalability
• Occupancy Measure
• Linear Programming

Innovation

• Decision-Focused Learning
• Occupancy Measure
• Augmented Lagrangian Surrogate
• Soft State Aggregation
• Markov Decision Process