Decision-Focused Learning in MDPs: An Occupancy Measure Approach

AI summary

Overview

Research area: Machine learning, specifically decision-focused learning (DFL), differentiable optimization, and Markov decision processes (MDPs), with connections to linear programming, occupancy measures, and state aggregation.

Technical level: Advanced. The paper works with LP duality, KKT conditions, basis matrices, matrix inversion complexity, and differentiable optimization layers.

Scope: The paper replaces the KKT-based differentiation used in prior decision-focused learning for MDPs with a differentiable occupancy-measure linear program, adding an augmented Lagrangian training loss, random row sketching, and learnable soft state aggregation to make the approach faster and more scalable.

What This Paper Is About

In many sequential decision problems, a model must predict unknown parameters (such as demand or transition dynamics) that are then fed into an optimizer. Standard "two-stage" pipelines train the predictor to minimize prediction error, but accurate predictions do not necessarily produce good decisions. Decision-focused learning instead trains the predictor end-to-end on the quality of the resulting decision. When the downstream problem is an MDP, existing DFL methods differentiate through the KKT conditions of the Bellman equation, which requires solving a linear system of dimension |S||A| × |S||A| and scales poorly. This paper reformulates the MDP as a linear program over occupancy measures whose feasible region is induced by the predicted dynamics, and derives a closed-form gradient for that LP layer.

Key Contributions

1. A differentiable occupancy-measure LP layer for MDPs. The paper proposes an LP layer with a closed-form, well-conditioned gradient whose main computational cost is inverting a single |S| × |S| basis matrix, rather than the |S||A| × |S||A| system required by the KKT approach.

2. An augmented Lagrangian DFL surrogate with a global optimum guarantee. The training loss combines the downstream true reward with a linear dual correction and a quadratic penalty on true flow-constraint violation. Proposition 4.1 states that the true parameter is a global minimizer. Random row sketching of the flow constraints is used to empirically reduce sensitivity to basis changes.

3. Scaling to large finite and continuous-state MDPs via soft state aggregation and function approximation. A learnable softmax membership matrix compresses the |S| × |S| system to an M × M system with M ≪ |S|, and the same construction is extended to continuous states through occupancy features Φ and test functions Ψ trained jointly with the predictor. The paper reports competitive decision quality with lower training cost than multiple baselines across three tasks.

4. Empirical comparison across three tasks. Finite-state inventory, stochastic 4×12 Cliff Walking, and continuous-state CartPole, compared against two-stage, KKT-based DFL (DFL-QP), a feasibility-aware constrained-DFL baseline (DFL-Feas), DFL-LP, DFL-Sketch, and DFL-SA, averaged over 10 random seeds.

Main Findings

• Lower regret than KKT-based and two-stage baselines: On all three tasks, the best-performing DFL variants achieve lower mean test regret than the two-stage baseline. DFL-Sketch improves upon DFL-LP on Inventory and CartPole, reducing mean regret on CartPole from approximately 24 to 9.

• Closed-form gradient via the optimal basis: Because the predicted parameter θ enters both the LP objective (through rewards) and the flow constraints (through transitions), the constraint Jacobian ∂H^B_θ/∂θ is generically non-zero, giving a non-zero closed-form gradient. The paper notes that when θ enters only the LP objective over a fixed feasible region, dy*/dθ ≡ 0 and KKT differentiation provides no learning signal.

• Non-degeneracy under a strictly positive initial distribution: When γ(s) > 0 for all states s, every basis feasible solution is non-degenerate, so the positive entries of the selected solution uniquely determine an invertible basis matrix H^B_θ ∈ R^{|S|×|S|}.

• Bias of the naive occupancy proxy: The predicted occupancy is feasible only under the predicted dynamics H_θ and generically violates the true flow constraints H_true, so it can score higher on r_true^T y than any truly feasible occupancy. Minimizing −r_true^T y*_θ need not recover θ_true. The augmented Lagrangian surrogate (with linear dual term and quadratic penalty) has a minimum at the true parameter instead.

• Complexity reduction: Table 1 reports that the KKT approach inverts a |S||A| × |S||A| matrix at cost O((|S||A|)^ω), the LP approach inverts a |S| × |S| basis at cost O(|S|^ω), and the LP plus state aggregation inverts an M × M basis at cost O(M^ω), where ω < 2.373 is the matrix multiplication exponent.

• Problem sizes in the experiments: The occupancy LP has |S||A| variables and |S| flow constraints, giving 36 × 31 = 1116 variables for Inventory and 48 × 4 = 192 for Cliff Walking.

• DFL-SA performance (partial): With sufficiently large aggregation sizes, DFL-SA achieves regret comparable to DFL-QP on Inventory and Cliff Walking and lower mean regret on CartPole. All experiments average over 10 random seeds with shaded regions and error bars indicating standard deviation.

• Cost reporting: The paper states that methods reach lower regret with significantly lower computation cost, and measures total wall-clock training time including forward optimization and backpropagation. The specific wall-clock numbers are not included in the available text.

Methodology in Plain English

The starting point is a classical fact: solving an MDP can be written as a linear program over the discounted state-action occupancy measure y, where the objective is the reward dotted with y and the constraint H_θ y = γ encodes flow conservation under the predicted dynamics. The paper differentiates through this LP.

Step 1 — Get a usable gradient. Two-stage learning has no decision signal. Differentiating through the KKT conditions of Bellman optimality requires inverting a matrix of size |S||A| × |S||A|. In the occupancy LP, the solution sits on a vertex, and at that vertex only |S| variables (those in the "basis") are nonzero. The paper identifies this basis by pivoting, assumes it is locally stable, and then uses the constraint equation to derive an exact closed-form derivative of the basic variables with respect to θ. This requires inverting only the |S| × |S| basis matrix.

Step 2 — Fix the training objective. Using the raw predicted occupancy as the training signal is biased because that occupancy satisfies the predicted dynamics, not the true ones. The authors add two correction terms to the loss: a linear term using the optimal dual solution of the true LP, and a quadratic penalty on violation of the true flow constraints. They prove the true parameter is a global minimizer of this surrogate. No extra optimization is needed because LP solvers return primal and dual solutions together.

Step 3 — Smooth the jumps. When the optimal basis changes, the loss and its gradient jump. Standard smoothing adds noise to the objective, but here the constraints also depend on the prediction. So the authors instead draw K Gaussian sketches of the flow-constraint rows, keeping only a fraction k = α|S| of them, solve each sketched LP, and average the solutions. A mass-conservation row 1^T y = 1/(1−β) is kept as an extra constraint to prevent the smaller sketched feasible set from becoming unbounded. Sketching is used only during training; inference solves the unsketched problem.

Step 4 — Shrink the state space. To avoid inverting an |S| × |S| matrix when states are numerous or continuous, the paper introduces a soft state-aggregation layer: a learnable membership matrix W with simplex rows collapses states into M clusters (M ≪ |S|), producing an aggregated reward, aggregated transition kernel, and aggregated flow constraint. The aggregated LP has the same form as the original, so the same closed-form gradient applies to an M × M basis. For continuous states, occupancy features Φ and test functions Ψ are represented by neural networks and trained jointly with the predictor; the state-space integrals are estimated from offline transitions (s_t, a_t, r_t, s_{t+1}) by sample averaging rather than a learned transition model.

Step 5 — Lift back. After solving in the compressed space, the solution is lifted back to the original state-action space for evaluation and action normalization.

Why This Matters

Impact on research. Prior DFL for MDPs treated the policy as the solution of a nonlinear system and differentiated through KKT conditions, which is expensive and can be poorly conditioned as the discount factor approaches one. This paper shows that moving to the occupancy-measure LP gives a non-degenerate, well-defined basis, a closed-form gradient, and a matrix-inversion cost that no longer grows with the action space. It also connects learnable state aggregation and representation learning to decision-focused learning, where prior aggregation work operated on a fixed MDP rather than differentiating end-to-end through a predicted one. The setting also differs from prior constrained-DFL work in that a single prediction determines both the objective and the constraints simultaneously, creating a coupled sensitivity that earlier constrained-DFL methods for mixed-integer programs or contextual stochastic LPs do not have.

Real-world applications (potential, motivated by the paper's own examples and framing):

• Inventory control under uncertain demand, the paper's motivating example and first experimental task, where order quantities are chosen from noisy demand signals.
• Stochastic navigation and path planning under changing environmental conditions, as in the Cliff Walking task with latent weather regimes determining slip probabilities.
• Control of continuous physical systems from offline data, as in the CartPole task with unknown dynamics observed only through a buffer of transitions.
• Domains where the paper's framing applies generally: sequential decisions whose parameters must be predicted and then optimized over, which the authors motivate with inventory-style problems.

Industry relevance. The paper's central selling point is computational cost: inverting an |S| × |S| basis instead of an |S||A| × |S||A| system, and further compressing to M × M with M ≪ |S|. For practitioners applying decision-focused learning to real sequential decision problems with large action spaces, this addresses the main obstacle to scaling DFL beyond small MDPs. The sketching option is described as requiring multiple LP solves that can be parallelized, each using only a small fraction of constraints (e.g., 10%).

Future Directions

• Provable smoothness for the sketched loss. The paper explicitly states that the true-parameter minimum guarantee does not automatically extend to the sketched loss, because the averaged solution need not satisfy the original flow constraints and finite averaging does not guarantee smoothness. Closing this gap is an open question.

• Characterizing approximation error from state aggregation. The soft aggregation replaces W^T W with diag(μ) to obtain a valid aggregated MDP and lifts the aggregated solution back via the membership matrix. How the aggregation size M and feature dimensions d trade off against decision quality is the natural next question; the paper reports d = 1024 for DFL-LP and DFL-QP and d ∈ {30, 60, 90} for DFL-SA on CartPole.

• Extending the bias guarantee beyond the unsketched loss. Proposition 4.1 covers the augmented Lagrangian surrogate under the exact LP solution. Extending the guarantee to the training procedure actually used (sketch averaging plus aggregation) is not established.

• Broader task coverage. The evaluation uses three tasks (inventory, Cliff Walking, CartPole). Whether the occupancy-LP layer with state aggregation holds up on larger finite MDPs, higher-dimensional action spaces, or domains where the occupancy features are harder to learn is not reported in the available content.

Target Audience

Researchers and graduate students working on decision-focused learning, differentiable optimization layers, and learning for sequential decision-making, particularly those already familiar with MDPs, LP duality, and implicit differentiation. Practitioners who need to train predictive models whose outputs feed a downstream planner or controller and who care about training cost at scale would also benefit, though the mathematical development (basis matrices, augmented Lagrangians, sketching) makes this an advanced read. Readers new to DFL or LP duality would need substantial background preparation.

Authors’ 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.