Decision-Focused Learning in MDPs: An Occupancy Measure Approach

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.