Decision-Focused Learning in MDPs: An Occupancy Measure Approach

arXiv:2610.08384 · cs.LG · Submitted 2026-10-06 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Decision-Focused Learning in MDPs".

Jane: The gist The authors propose an occupancy-measure-based linear program layer for decision-focused learning in Markov decision processes that offers a closed-form gradient and scales to large state spaces through learnable…

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So we’re looking at this paper called "Decision-Focused Learning in MDPs: An Occupancy Measure Approach". Jane, what’s the big idea here, in plain English?

Jane: It takes a problem where you're trying to make decisions in a Markov decision process, or MDP, and it says instead of using the standard way—which is differentiating through those complex KKT conditions of the Bellman equation—we can rephrase it as an occupancy measure-based linear program, or LP.

Tom: An LP? That sounds like we’re turning a complex decision problem into a standard linear math problem. What does that actually mean for scaling this thing up?

Lu: It means the feasible region of this new LP is defined by the predicted dynamics, and they claim that because of how it’s set up, we get two big wins. First, under a strictly positive initial distribution, the LP is non-degenerate. That means we get a well-defined optimal basis and we avoid those KKT conditioning issues you usually run into with standard methods >#pg2.

Meng: Avoiding KKT conditioning sounds good for theory, but what about the practical cost when we have a huge state space? If the LP still has to deal with every state-action pair, that’s not scalable for what we actually run in production >#pg1.

Lalam: The core idea is that instead of solving a massive system over all states and actions, you only need to invert a single S by S basis matrix to get the gradient. That’s much cheaper when the number of actions, or A, is large >#pg2.

Tom: So we replace solving the whole KKT system with just inverting one big basis matrix? That sounds like a huge computational win for training speed. What about those practical hurdles they mention?

Jane: The authors admit that this LP reformulation has two real challenges they have to fix in practice, and those are biased loss minimum and gradient discontinuity >#pg2.

The paper's summary: Tom: Let’s talk about what those challenges are, because if you can’t train it reliably, the math doesn't matter for us in production. What’s the first issue they highlight?

Jane: They point out that standard constrained decision-focused learning often predicts only objective coefficients, but here in MDPs, the predicted parameters show up both in the reward part of the objective and inside those constraints through transitions >#pg2.

Lu: So the predicted occupancy doesn't actually satisfy the true flow constraints, which means if you train using that prediction directly in your loss function, you end up with a biased loss minimum >#pg2.

Meng: That sounds like we’re training the model toward a wrong target because its predictions don't match reality on the true dynamics >#pg2. How do they fix that bias?

Jane: They introduce something called an augmented Lagrangian surrogate, which adds two things to the loss: a linear dual correction and a quadratic penalty on how much those flow constraints are violated >#pg3. They prove that this surrogate means the true parameter is a global minimizer of it >#pg6.

Tom: An augmented Lagrangian with penalties sounds like adding complexity to the training objective, but they claim it keeps the solution globally optimal for the true parameters. That’s a strong claim if they can actually manage it without making training too unstable >#pg6.

Lalam: I see how that helps stabilize things; it’s essentially a way of telling the model, "Hey, if you violate the true dynamics constraints, pay a penalty."

The paper's improvements: Tom: Okay, so we have this core idea and these fixes for bias. What about scaling? Because if this only works for small state spaces—say fewer than fifty states—it’s not useful for inventory or complex robotics >#pg1.

Jane: That’s where they introduce state aggregation, which is a classic technique in reinforcement learning. They suggest using a learnable soft state-aggregation layer to compress the massive S by S system down into something much smaller, say M by M, where M is much less than S >#pg3.

Lu: This layer uses a learned softmax membership matrix to map the original states to these abstract representations, and they claim this reduces the per-step gradient cost from O(S omega) down to O(M omega), which is a significant speedup >#pg8.

Meng: If we compress the system size by using M instead of S, that means less computation on every single training step, which is exactly what I need to see when I’m trying to ship something at scale >#pg8.

Tom: And for continuous state spaces? How do they handle those where there are infinitely many states? They can’t just use a fixed matrix anymore, right?

Jane: For continuous states, they use neural networks to create occupancy features and test functions, which lets them estimate those state-space integrals directly from offline transitions rather than needing a learned transition model >#pg6.

Lalam: So for continuous spaces, the neural network estimates the math for us so we can still train that projected LP using the basis-based gradient they derived >#pg6.

Conclusion: Tom: We’ve covered the occupancy measure approach, from how it reformulates MDPs into LPs to how they fix those bias and discontinuity problems with surrogates and sketching. What’s the final word on what this means for decision-focused learning in general?

Jane: This paper presents a way to get a closed-form gradient that only requires inverting that single S by S basis matrix, which is a big step because it makes the gradient calculation much more efficient than solving the full KKT system >#pg6.

Lu: And with the scaling mechanisms they introduced—the soft aggregation layer and function approximation for continuous states—they’ve shown how to project this LP approach onto large finite and continuous state MDPs, which is where most of the practical work needs to happen >#pg8.

Meng: From a deployment standpoint, the speedup from reducing the cost from O(S omega) to O(M omega) is what actually matters for us when we think about running this kind of system in production >#pg8.

Lalam: I'm pretty sure this paper establishes that you can get decision-focused learning in MDPs with a tractable, well-conditioned layer that scales through these aggregation and approximation techniques >#pg6.

Tom: Exactly. So, the occupancy measure approach provides a way to tackle the state space curse by replacing complex nonlinear equations with linear constraints and then providing an efficient path for optimization >#pg1. That’s what we’ve got from "Decision-Focused Learning in MDPs: An Occupancy Measure Approach" >#pg6.

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

Georgia Institute of Technology · Johns Hopkins Carey Business School · International Computer Science Institute

cs.LG

Submitted: 2026-10-06

Updated: 2026-10-06

Comments: Accepted at NeurIPS 2026

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 86/100

The gist: The gist The authors propose an occupancy-measure-based linear program layer for decision-focused learning in Markov decision processes that offers a closed-form gradient and scales to large state

Key concepts

Occupancy Measure LP
The MDP is treated as an LP where the feasible region is defined by predicted dynamics. This formulation avoids KKT conditioning issues under positive initial distributions and allows differentiation through a single basis matrix, leading to a well-conditioned gradient.
Augmented Lagrangian Surrogate
This technique adds a linear dual correction and a quadratic penalty term to the loss function. It ensures that the true parameter is a global minimizer of this surrogate, effectively correcting for biases where predicted occupancy might violate true flow constraints.
Soft State-Aggregation Layer
This layer uses a learned softmax membership matrix to compress the large state space into a smaller one. This reduces the size of the basis system from |S| x |S| to M x M, significantly lowering per-step gradient costs for large MDPs.
Random Row Sketching
This method is used on flow constraints to empirically smooth out abrupt gradient changes that occur when the optimal basis shifts. It helps stabilize the training process by averaging solutions near basis boundaries.

Summary

The gist The authors propose an occupancy-measure-based linear program layer for decision-focused learning in Markov decision processes that offers a closed-form gradient and scales to large state spaces through learnable aggregation.

How it works

The core idea is to reformulate the MDP as an occupancy measure-based linear program (LP), whose feasible region is induced by predicted dynamics, which has two key advantages: first, under a strictly positive initial distribution, the occupancy LP is non-degenerate, yielding a well-defined optimal basis and avoiding KKT conditioning issues; second, differentiating through the LP layer requires inverting only a single S × S basis matrix, giving a closed-form, well-conditioned gradient and substantially lower cost when A is large<ref:2610.08384#pg6>. This approach replaces the nonlinear Bellman equation with linear constraints, and the resulting gradient depends only on a small basis matrix<ref:2610.08384#pg5>.

The paper addresses two practical challenges of directly using this LP reformulation: (a) Biased loss minimum, where predicted occupancy may violate the true flow constraints, and (b) Gradient discontinuity, where changes in the optimal basis can cause sudden jumps in the loss and its gradients<ref:2610.08384#pg6>. To tackle these, they introduce an augmented Lagrangian surrogate that adds both a linear dual correction and a quadratic penalty on the true-dynamics constraint to the training loss, proving that the true parameter is a global minimizer of this surrogate<ref:2610.08384#pg6>. Furthermore, they apply random row sketching to the flow constraints and average the resulting LP solutions to empirically reduce abrupt gradient changes near basis boundaries<ref:2610.08384#pg6>.

Scaling the approach

To scale up to large finite and continuous-state MDPs, they introduce two primary scaling mechanisms: (1) a learnable soft state-aggregation layer that uses a learned softmax membership matrix to compress the S × S system to M × M with M ≪ S; and (2) function approximation for continuous-state MDPs<ref:2610.08384#pg6>. For finite state MDPs, soft aggregation corresponds to a particular choice of (Φ, Ψ) derived from a learnable membership matrix W ∈ RS×M with simplex rows, which collapses the S × S basis matrix to an M × M matrix with M ≪ S, reducing the per-step gradient cost from O(Sω) to O(Mω)<ref:2610.08384#pg8>.

For continuous-state MDPs, they use neural networks to represent occupancy features Φ and test functions Ψ, estimating state-space integrals from offline transitions rather than requiring a learned transition model<ref:2610.08384#pg6>. This allows them to train the projected LP by backpropagating through these estimates and the LP using their basis-based gradient<ref:2610.08384#pg6>.

Theoretical Guarantees and Results

The closed-form gradient derived in Eq. (4) is exact within each basis region under the assumption that the optimal basis is locally stable, which means small perturbations in θ do not change the selected basis<ref:2610.08384#pg6>. They introduce an augmented Lagrangian surrogate, L(y∗θ):= −r⊤truey∗θ + (Htruey∗θ − γ)⊤ν∗ + ρ∥Htruey∗θ − γ∥2, which satisfies L(y∗θ) − L(y∗true) ≥ 0 for every predicted parameter θ, with equality iff y∗θ is optimal for the true LP, and thus θ = θ∗ is a global minimizer<ref:2610.08384#pg6>.

Experimental Performance

Experiments on Inventory, Cliff Walking, and CartPole show lower regret than the two-stage baseline (Figures 2 and 3), with DFL-Sketch improving upon DFL-LP on Inventory and CartPole, showing the benefit of smoothing across basis changes<ref:2610.08384#pg6>. For CartPole, DFL-SA (d=90) reduces mean regret to approximately 3, compared to 10 for DFL-QP and 45 for two-stage<ref:2610.08384#pg8>. The gains over two-stage are modest on Inventory, where the baseline is already competitive<ref:2610.08384#pg6>.

Computational Efficiency

The computational cost comparison shows that DFL-LP replaces differentiation through the full KKT system with a smaller basis solve, reducing training time by about 8× relative to DFL-QP on Inventory<ref:2610.08384#pg8>. DFL-SA further reduces the size of the basis system through learned aggregation, achieving strong decision quality at substantially lower cost than DFL-QP<ref:2610.08384#pg8>. On CartPole, DFL-SA (d=90) trains about 4× faster than DFL-LP and 7.5× faster than DFL-QP while achieving the lowest mean test regret<ref:2610.08384#pg8>.

Conclusion

The authors present a differentiable LP layer for decision-focused learning in MDPs, with a closed-form gradient requiring only a S × S basis solve<ref:2610.08384#pg6>. An augmented Lagrangian surrogate corrects the bias from predicted-dynamics infeasibility, while random row sketching smooths gradients across basis boundaries<ref:2610.08384#pg6>. Learned state aggregation reduces the basis size to M ≪ S, and function approximation extends the approach to continuous states<ref:2610.08384#pg6>. Experiments on Inventory, Cliff Walking, and CartPole show lower regret than two-stage and decision quality comparable to or better than KKT-based DFL at substantially lower computational cost<ref:2610.08384#pg6>.

How it works

The authors propose a novel occupancy-measure LP layer as a differentiable MDP layer with a closedform, well-conditioned gradient requiring inversion of a S × S basis matrix<ref:2610.08384#pg6>. They introduce an augmented Lagrangian DFL surrogate that augments the regret objective, with a global optimum guarantee, and employ random row sketching to empirically reduce sensitivity to basis changes<ref:2610.08384#pg6>.

State Aggregation in Markov Decision Processes

The paper addresses the curse of dimensionality by introducing a soft state-aggregation layer that is trained jointly with the predictor and extended to the general state MDPs through linear function approximation<ref:2610.08384#pg8>. This layer maps original states to abstract representations, collapsing the S × S basis matrix to an M × M matrix with M ≪ S and reducing the per-step gradient cost to O(Mω)<ref:2610.08384#pg8>.

Function Approximation for General State MDPs

For continuous state spaces, they use neural networks to represent occupancy features Φ and test functions Ψ, yielding a finite-dimensional LP<ref:2610.08384#pg6>. They estimate the projected LP coefficients from offline transitions rather than requiring a learned transition model<ref:2610.08384#pg6>.

How it works

The core idea is to reformulate the MDP as an occupancy measure-based linear program (LP), whose feasible region is induced by predicted dynamics, which has two key advantages: first, under a strictly positive initial distribution, the occupancy LP is non-degenerate, yielding a well-defined optimal basis and avoiding KKT conditioning issues; second, differentiating through the LP layer requires inverting only a single S × S basis matrix, giving a closed-form, well-conditioned gradient and substantially lower cost when A is large<ref:2610.08384#pg6>.

State Aggregation in Markov Decision Processes

The paper addresses the curse of dimensionality by introducing a soft state-aggregation layer that is trained jointly with the predictor and extended to the general state MDPs through linear function approximation<ref:2610.08384#pg8>.

Improvements for AI systems

  1. The system can perform end-to-end decision-focused learning (DFL) in Markov Decision Processes (MDPs) by training a predictive model that maximizes decision quality rather than just minimizing prediction error, as described by the objective: min w L(w):= r⊤ truey∗true − r⊤ truey∗θ (Eq. 1).

  2. The system can handle large or continuous state spaces efficiently by employing a learnable soft state-aggregation layer and its function-approximation generalization that scales the LP to large finite and continuous-state MDPs. This reduces the computational cost from O(Sω) to O(Mω), where M ≪ S.

  3. The system can be trained robustly against prediction errors by using an augmented Lagrangian surrogate that adds both a linear dual correction and a quadratic penalty on the true-dynamics constraint to address the biased loss minimum challenge (C1).

  4. The system can maintain training stability near optimal policy switches by applying random row sketching to the flow constraints and averaging the resulting LP solutions, which empirically reduces abrupt gradient changes near basis boundaries.

  5. The system can provide a closed-form, well-conditioned gradient for the MDP layer by deriving a gradient that depends only on inverting a single S × S basis matrix (Eq. 3), instead of solving the full SA × SA KKT system.

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.

Sources

Related papers