Decision-Focused Learning in MDPs: An Occupancy Measure Approach

1 Introduction

• KKT-based MDP differentiation solves a linear system over |S||A| state-action pairs, limiting scalability and potentially suffering poor conditioning as the discount factor approaches one.The conditioning issue arises from the Bellman operator's nonlinearity.Source: §1 · p. 1
• The occupancy-measure LP has a non-degenerate optimal basis under a strictly positive initial distribution and requires only a |S| × |S| basis-matrix inversion for its gradient.This gives a closed-form, well-conditioned gradient and lower cost, especially when |A| is large.Source: §1 · pp. 1–2
• Predicted dynamics make the occupancy feasible only for the predicted MDP, so the paper uses an augmented Lagrangian surrogate to correct loss bias and random row sketching to reduce abrupt gradient changes.With an optimal dual solution of the true LP, the surrogate has the true parameter as a global minimizer; averaging solutions from different sketches empirically reduces gradient changes near basis boundaries.Source: §1 · p. 2
• Learned soft state aggregation compresses the basis system from |S| × |S| to M × M, with M ≪ |S|, and function approximation extends the construction to continuous-state MDPs.The introduction compares computational complexity with the KKT approach, but the supplied table passage provides no cell values.Source: §1 · p. 2, Table 1 · p. 2
• Across three tasks, the proposed methods target competitive decision quality with lower training cost than multiple baselines.The introduction describes learnable soft state aggregation and function approximation as extensions to larger finite and continuous-state MDPs.Source: §1 · p. 2

2 Related Work

• Prior DFL approaches address differentiation through an argmin with KKT implicit differentiation, solver relaxations, surrogate losses, and learned task-aligned losses.DFL trains predictors end-to-end against decision regret rather than parameter accuracy.Source: §2 · p. 2
• Prior constrained-DFL work addresses infeasibility from mispredicted constraints, but often studies mixed-integer programs or contextual stochastic LPs without sequential structure.In this MDP setting, a single prediction determines both objective and constraints through predicted dynamics, creating coupled sensitivity.Source: §2 · pp. 2–3
• Unlike prior state-aggregation methods that operate on a fixed MDP, this work differentiates a learnable aggregation layer end-to-end through a predicted MDP.Prior approaches include homomorphism-based, projection-based, representation-learning, and probabilistic-membership methods.Source: §2 · p. 3

3 Problem Statement

• An occupancy measure records the expected discounted frequency of visiting each state and taking each action, and MDP optimization can be written as a linear program over it.The LP uses flow constraints determined by the transition model.Source: §3 · p. 3
• Decision-focused learning trains the predictor to maximize the value of the resulting policy under true parameters, whereas a two-stage pipeline minimizes parameter prediction error before solving the MDP.The two objectives can diverge because lower prediction error does not necessarily improve downstream decision quality.Source: §3 · p. 3
• In the occupancy LP, predicted dynamics determine the feasible region, so a predicted occupancy can violate true flow constraints and shift the training-loss minimum away from the true parameter.The predicted occupancy may score well under the true reward without corresponding to a feasible policy under true dynamics.Source: §3 · p. 4
• The paper targets these issues with an LP-layer gradient, an augmented Lagrangian surrogate and row sketching, and soft state aggregation that compresses the gradient system.The proposed aggregation further reduces the system from |S| × |S| to M × M, where M ≪ |S|.Source: §3 · p. 4

4.1 Closed-form Gradient of the LP Layer

• The LP-layer gradient requires an |S| × |S| basis solve rather than the |S||A| × |S||A| system used by KKT-based DFL, avoiding its conditioning issues.Under a strictly positive initial distribution, the occupancy LP is non-degenerate and has a well-defined optimal basis.Source: §4 · p. 4, §4.1 · p. 4

4.2 DFL Training Loss

• The exact LP gradient applies within a locally stable basis region; changes in the active basis can still produce discontinuous loss or gradient jumps.The paper uses sketch-and-solve averaging to make gradients less sensitive across these transitions.Source: §4.1 · p. 4, §4.2 · p. 5, Figure 1 · p. 5
• Because predicted parameters enter both the reward objective and transition-dependent flow constraints, the predicted occupancy can violate true flow constraints and bias the loss minimum.The occupancy is feasible under predicted dynamics, not necessarily under the true dynamics.Source: §4.2 · p. 5
• The augmented Lagrangian combines decision quality under the true reward with a linear dual correction and a weighted penalty for violating true flow constraints.The quadratic penalty strengthens the correction as the predicted occupancy moves farther from the true-feasible set.Source: §4.2 · p. 5
• For any penalty ρ > 0, the surrogate is optimal when the predicted occupancy is optimal for the true LP, so the true parameter is a global minimizer.The dual solution is available from the true LP already solved during training.Source: §4.2 · p. 5
• Random row sketching averages solutions from LPs with randomized constraints, moving the average off a single vertex and empirically reducing gradient sensitivity to basis changes.The average solution need not satisfy the original flow constraints, and finite averaging does not guarantee smoothness or preserve the true-parameter minimum guarantee.Source: §4.2 · pp. 5–6

5 State Aggregation

• A learned soft state-aggregation layer compresses the |S| × |S| gradient system to an M × M system, with M ≪ |S|, and extends to general-state MDPs through function approximation.This addresses the remaining state-size bottleneck after occupancy-LP differentiation has reduced the dependence on action-space size.Source: §5 · p. 6

5.1 Soft State Aggregation for Finite State MDPs

• For finite MDPs, learnable softmax memberships form an aggregated MDP whose flow constraints use an M × M basis, and the aggregated occupancy can be lifted back to the original states.Soft assignments remain differentiable and are trained jointly with the predictor.Source: §5.1 · pp. 6–7

5.2 Function Approximation for General State MDPs

• For continuous-state MDPs, neural networks represent occupancy and test functions to form a finite-dimensional projected LP, whose coefficients are estimated from offline transitions.The transition estimate uses current state-action features and observed next-state test functions, avoiding the need for a learned transition model.Source: §5.2 · p. 7

6 Experiments

• The evaluation covers three tasks spanning finite-state and general-state MDPs, with comparisons against two-stage and several DFL baselines.The baselines include KKT-based DFL-QP, the occupancy-LP gradient, feasibility-aware DFL, sketching, and state aggregation.Source: §6 · pp. 7–8
• Reported numbers are averaged over 10 random seeds, with standard deviation shown by shaded regions and error bars.Source: §6 · p. 8

6.1 Inventory Problem

• The inventory task is a finite-state ordering problem with Poisson demand of unknown rate, immediate deterministic lead time, and expected-profit maximization.The state tracks net inventory, including backlog, and the known selling price and unit ordering cost determine stage reward.Source: §6.1 · p. 8

6.2 Cliff Walking

• Cliff Walking uses a 4×12 grid with four actions and latent weather regimes that alter slip probabilities and can make the short path near the cliff risky.The regimes have overlapping weather-feature distributions, requiring inference of the hazardous region from noisy features.Source: §6.2 · p. 8

6.3 Continuous State MDP

• CartPole is evaluated as a continuous-state MDP with a four-dimensional state, two actions, a 200-step horizon, and unknown transitions observed through an offline buffer.Its transitions use standard rigid-body dynamics integrated with one Euler step of size 0.02.Source: §6.3 · pp. 8–9
• CartPole uses jointly trained neural basis and test functions, with behavior-cloning reward shaping on the LP forward solve while regret is evaluated on raw reward.The two-stage baseline shares the pretrained features and policy extraction but omits the DFL gradient phase.Source: §6.3 · p. 9

Decision quality

• On CartPole, DFL-SA (d=90) reaches mean regret of approximately 3, versus 10 for DFL-QP and 45 for two-stage; on shifted ICE-MID, DFL-SA also lowers regret on Cliff Walking.On Inventory, gains over two-stage are modest because that baseline is already competitive.Source: §7 · p. 9, Figure 3 · p. 10

Training reward

• DFL-QP and DFL-SA achieve the highest training rewards across three tasks, while DFL-Sketch and other LP-based variants are generally competitive with the two-stage baseline..Source: Figure 2 · p. 9

Computation cost

• On Inventory with |S|=101, DFL-SA (M=30) achieves the lowest mean regret while training approximately 40× faster than DFL-QP.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.Source: §7 · p. 9, Figure 4 · p. 10

Effect of the state aggregation size

• Increasing aggregation size reduces regret across evaluated configurations, including CartPole’s decrease from 15 to 3 as d rises from 30 to 90.Training time generally increases with aggregation size, though CartPole timings are not strictly monotone because solver behavior also affects wall-clock cost.Source: §7 · p. 10

Other ablation studies

• Additional ablations examine scaling to larger Inventory state spaces, sensitivity to feasibility penalty ρ, and the diagonal projection used in soft aggregation.Source: §7 · p. 10

• The differentiable LP layer uses a closed-form gradient requiring only a |S| × |S| basis solve.Source: §8 · p. 10
• An augmented Lagrangian surrogate addresses predicted-dynamics infeasibility, random row sketching smooths gradients across basis boundaries, and learned aggregation reduces basis size to M ≪|S|.Function approximation extends the aggregation approach to continuous-state MDPs.Source: §8 · p. 10
• Across Inventory, Cliff Walking, and CartPole, the methods achieve lower regret than two-stage learning and decision quality comparable to or better than KKT-based DFL at substantially lower computational cost.Source: §8 · p. 10

A.1 Sample-based estimators for the projected LP

• Under the behavior-weighted convention, occupancy is expressed relative to the data distribution to avoid dividing by small behavior probabilities, at the cost of solving a behavior-weighted rather than Lebesgue-weighted projected LP.The estimators for Ψ⊤PΦ, b, and the predicted reward are plain empirical averages.Source: §A.1 · p. 15
• The CartPole implementation weights the B block by the state marginal rather than the state-action marginal, making its LP a heuristic surrogate rather than an exact estimate of the behavior-weighted LP.The consistent estimator removes this mismatch and yields a policy proportional to πb(a | s) Φ(s, a)⊤w∗.Source: §A.1 · p. 15
• Backpropagation through Eq. (4) reduces to differentiating per-sample features ψi and ϕj when the projected-LP coefficients depend on predicted dynamics through empirical sums.Source: §A.1 · p. 15

A.2 Derive discrete-state soft aggregation from continuous-state function approximation

• The soft aggregation construction forms state features from expected next-state embeddings across actions, then projects them to cluster logits for a temperature-scaled softmax membership assignment.The reduced MDP is constructed from weighted averages of the original MDP quantities.Source: §A.2 · p. 16
• For soft memberships, replacing W⊤W with diag(µ) yields valid aggregated-MDP flow constraints but is an approximation; the replacement is exact for hard memberships.The aggregated occupancy is defined by ˆy(m, a) := µm w(m,a).Source: §A.2 · pp. 16–17

A.3 Proof for Proposition 4.1

• The proof invokes the KKT conditions of the true LP as sufficient for optimality.Source: §A.3 · p. 17

B. Details of Experiments

• Experiments ran on one shared workstation with two Intel Xeon Gold 6548Y+ processors, 128 logical cores, and 1.5 TB of memory.Source: §B.1 · p. 17

B.1 Compute resources, hyperparameters and model architectures

• Within each task, methods share the predictor architecture, optimizer, and data; results are averaged over 10 random seeds.Source: §B.1 · p. 17

B.2 Details for Cliff Walking MDP Task

• A compact slip tensor has one ambient slot and six volatile-cell slots that override it near the cliff, with categorical movement probabilities for each action.The structured slip model is expanded into transitions by applying the five outcomes and enforcing boundaries, cliff resets, and goal termination.Source: §B.2 · p. 17
• The shifted test regime uses only ICE-MID, while training follows a skewed 15:1:1:1 prior and makes CLEAR and ICE-MID features indistinguishable.This creates a setting where a predictor may fit the training distribution yet select a poor policy under test dynamics.Source: §B.2 · p. 18
• Regret evaluates each predicted-policy choice in the sample’s true MDP, against an oracle value computed using the true transition kernel.Thus regret measures downstream cost under misspecified dynamics rather than prediction error.Source: §B.2 · p. 18
• DFL-SA aggregates 48 states to 12, 18, or 24 before the LP solve, while DFL-Sketch averages gradients from eight randomized LP sketches retaining 20% of rows.The aggregate occupancy is disaggregated for evaluation; sketching uses no aggregation.Source: §B.2 · p. 18
• Table 3 uses an earlier regret scale and aggregation grid than Figure 3, so its win counts do not establish per-seed wins or significance for the main-figure configuration.Its paired win counts compare methods against the two-stage baseline on the same seed.Source: §B.2 · p. 18, Table 3 · p. 19

B.3 Details for Continuous MDP Task

• Because CartPole reward is action-independent, the LP forward solve adds a behavior-cloning bonus with α=2, while the DFL regret loss uses raw reward.A behavior-cloning head predicts actions from state features.Source: §B.3 · p. 19
• The CartPole two-stage baseline shares DFL’s BC-pretrained features, shaping, and policy extraction, but omits the DFL gradient phase and holds features fixed.This controls for BC pretraining and reward shaping when comparing decision-aware feature learning.Source: §B.3 · p. 19
• CartPole LP-2 coefficients are estimated from offline transitions and initial states, with the transition kernel entering through observed next states rather than a learned model.The two-stage baseline holds BC-pretrained features fixed throughout.Source: §B.3 · p. 19
• The additional CartPole MBRL rows have lower observed means and larger across-seed standard deviations, but use a different configuration and do not establish a matched-data effect.The additional runs use five seeds; the main LP experiments use the Table 2 configuration.Source: §B.3 · p. 19, Table 4 · p. 20

B.4 Scaling to larger inventory MDPs

• At |S|=101, DFL-SA (M=30) achieves the lowest mean regret and trains 40.3× faster than DFL-QP and 23.1× faster than DFL-LP.Table 6 covers the 101 × 31 occupancy-variable setting.Source: §B.4 · p. 20, Table 6 · p. 20
• At |S|=51, DFL-SA (M=30) has mean regret 19.249 versus DFL-QP’s 19.121 and trains 4.3× faster.Table 5 covers the 51 × 31 occupancy-variable setting.Source: §B.4 · p. 20, Table 5 · p. 20

B.5 Sensitivity to the feasibility penalty ρ

• ρ = 104 is selected by validation regret; flow-constraint MSE decreases monotonically with increasing ρ, while regret slightly worsens at ρ = 105.Regret is relatively stable for ρ ∈ [103, 105].Source: §B.5 · p. 20, Table 7 · p. 21