Learning MDPs from Features: Predict-Then-Optimize for Sequential Decision Problems by Reinforcement Learning
Year: 2021

Abstract:
In the predict-then-optimize framework, the objective is to train a predictive model, mapping from environment features to parameters of an optimization problem, which maximizes decision quality when the optimization is subsequently solved. Recent work on decision-focused learning shows that embedding the optimization problem in the training pipeline can improve decision quality and help generalize better to unseen tasks compared to relying on an intermediate loss function for evaluating prediction quality. We study the predict-then-optimize framework in the context of sequential decision problems (formulated as MDPs) that are solved via reinforcement learning. In particular, we are given environment features and a set of trajectories from training MDPs, which we use to train a predictive model that generalizes to unseen test MDPs without trajectories. Two significant computational challenges arise in applying decision-focused learning to MDPs: (i) large state and action spaces make it infeasible for existing techniques to differentiate through MDP problems, and (ii) the high-dimensional policy space, as parameterized by a neural network, makes differentiating through a policy expensive. We resolve the first challenge by sampling provably unbiased derivatives to approximate and differentiate through optimality conditions, and the second challenge by using a low-rank approximation to the high-dimensional sample-based derivatives. We implement both Bellman–based and policy gradient–based decision-focused learning on three different MDP problems with missing parameters, and show that decision-focused learning performs better in generalization to unseen tasks.

heading: 1 Introduction
text: Predict-then-optimize [5, 9] is a framework for solving optimization problems with unknown parameters. Given such a problem, we first train a predictive model to predict the missing parameters from problem features. Our objective is to maximize the resulting decision quality when the optimization problem is subsequently solved with the predicted parameters [25, 28]. Recent work on the decisionfocused learning approach [7, 36] embeds the optimization problem [1, 3, 4] into the training pipeline and trains the predictive model end-to-end to optimize the final decision quality. Compared with a
∗This work was done while the author was at Harvard University.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
more traditional “two-stage” approach which maximizes the predictive accuracy of the model (rather than the final decision quality), the decision-focused learning approach can achieve a higher solution quality and generalize better to unseen tasks.
This paper studies the predict-then-optimize framework in sequential decision problems, formulated as Markov decision processes (MDPs), with unknown parameters. In particular, at training time, we are given trajectories and environment features from “training MDPs.” Our goal is to learn a predictive model which maps from environment features to missing parameters based on these trajectories that generalizes to unseen test MDPs that have features, but not trajectories. The resulting “predicted” training and test MDPs are solved using deep reinforcement learning (RL) algorithms, yielding policies that are then evaluated by offline off-policy evaluation (OPE) as shown in Figure 1. This fully offline setting is motivated by real-world applications such as wildlife conservation and tuberculosis treatment where no simulator is available. However, such domains offer past ranger patrol trajectories and environmental features of individual locations from conservation parks for generalization to other unpatrolled areas. These settings differ from those considered in transfer-RL [21, 24, 29, 31] and meta-RL [8, 10, 33, 35, 40] because we generalize across different MDPs by explicitly predicting the mapping function from features to missing MDPs parameters, while transfer/meta RL achieve generalization by learning hidden representation of different MDPs implicitly with trajectories.
The main contribution of this paper is to extend the decision-focused learning approach to MDPs with unknown parameters, embedding the MDP problems in the predictive model training pipeline. To perform this embedding, we study two common types of optimality conditions in MDPs: a Bellmanbased approach where mean-squared Bellman error is minimized, and a policy gradient-based approach where the expected cumulative reward is maximized. We convert these optimality conditions into their corresponding Karush–Kuhn–Tucker (KKT) conditions, where we can backpropagate through the embedding by differentiating through the KKT conditions. However, existing techniques from decision-focused learning and differentiating through KKT conditions do not directly apply as the size of the KKT conditions of sequential decision problems grow linearly in the number of states and actions, which are often combinatorial or continuous and thus become intractable.
We identify and resolve two computational challenges in applying decision-focused learning to MDPs, that come up in both optimality conditions: (i) the large state and action spaces involved in the optimization reformulation make differentiating through the optimality conditions intractable and (ii) the high-dimensional policy space in MDPs, as parameterized by a neural network, makes differentiating through a policy expensive. To resolve the first challenge, we propose to sample an estimate of the first-order and second-order derivatives to approximate the optimality and KKT conditions. We prove such a sampling approach is unbiased for both types of optimality conditions. Thus, we can differentiate through the approximate KKT conditions formed by sample-based derivatives. Nonetheless, the second challenge still applies—the sampled KKT conditions are expensive to differentiate through due to the dimensionality of the policy space when model-free deep RL methods are used. Therefore, we propose to use a low-rank approximation to further approximate the sample-based second-order derivatives. This low-rank approximation reduces both the computation cost and the memory usage of differentiating through KKT conditions.
We empirically test our decision-focused algorithms on three settings: a grid world with unknown rewards, and snare-finding and Tuberculosis treatment problems where transition probabilities are unknown. Decision-focused learning achieves better OPE performance in unseen test MDPs than twostage approach, and our low-rank approximations significantly scale-up decision-focused learning.
Related Work
Differentiable optimization Amos et al. [2] propose using a quadratic program as a differentiable layer and embedding it into deep learning pipeline, and Agrawal et al. [1] extend their work to convex programs. Decision-focused learning [7, 36] focuses on the predict-then-optimize [5, 9] framework by embedding an optimization layer into training pipeline, where the optimization layers can be convex [7], linear [22, 36], and non-convex [26, 34]. Unfortunately, these techniques are of limited utility for sequential decision problems because their formulations grow linearly in the number of states and actions and thus differentiating through them quickly becomes infeasible. Amos et al. [3] avoid this issue by studying model-predictive control but limited to quadratic-form actions, reducing the dimensionality. Karkus et al. [16] differentiate through an algorithm by unrolling and relaxing
all the strict operators by soft operators. Futoma et al. [11] deal with large optimality conditions by differentiating through the last step of the value-iteration algorithm only. Instead, our approach does not rely on any MDP solver structure. We combine sampling and a low-rank approximation to form an unbiased estimate of the optimality conditions to differentiate through, and show that the approach of Futoma et al. [11] is included in ours as a special case.
Predict-then-optimize and offline reinforcement learning The idea of planning under a predicted MDP arises in model-based RL as certainty equivalence [19]. It has been extended to offline settings [17, 39], who learn a pessimistic MDP before solving for the policy. Our setting differs because of the presence of features and train-test split—our test MDPs are completely fresh without any trajectories. Our setting also resembles meta RL (e.g., [8, 10, 33, 35, 40]) and transfer RL (e.g., [21, 24, 29, 31].) Meta RL focuses on training a “meta policy” for a distribution of tasks (MDPs), leveraging trajectories for each. Transfer RL works by extracting transferable knowledge from source MDPs to target MDPs using trajectories. In contrast to these two paradigms, ours explicitly trains a predictive model (which maps problem features to missing MDP parameters) to generalize knowledge learned from the training set to the testing set using problem features, not trajectories.

heading: 2 Problem Statement
text: In this paper, we consider learning a predictive model to infer the missing parameters in a sequential decision-making task (formulated as MDPs) using the predict-then-optimize framework. Each MDP is defined by a tuple (S, s0,A, T,R) with an initial state s0, a possibly infinite state space S and possibly infinite action space A. We assume some parameters are missing in each MDP, which could be any portion of the transition function T and the reward function R. We denote the missing parameters vector by θ∗. Additionally, we assume there are problem features x associated with each MDP, where (θ∗, x) is correlated and drawn from the same unknown, but fixed, distribution2.
We are given a set of training MDPs and a set of test MDPs, each with missing parameters θ∗ and features x. Each training MDP is accompanied by a set of trajectories T performed by a behavior policy πbeh, consisting of a sequence of states, actions and rewards. In the test MDPs, trajectories from the behavior policy are hidden at test time. These testing MDPs are considered fresh instances that we have to generate a policy without using any trajectories. Thus, at training time, we learn a predictive model mw to map from features to missing parameters; at test time, we apply the same model to make predictions and plan accordingly without using trajectories.
Formally, our goal is to learn a predictive model mw to predict the missing parameters θ = mw(x). The predicted parameters are used to solve the test MDPs, yielding the policy π∗(mw(x)) where our offline off-policy evaluation (OPE) metric is maximized. This framework is illustrated in Figure 1.
2Examples of the missing parameters θ∗ include the poaching risk of different locations in wildlife conservation and the transition probability of patients’ healthiness in healthcare problems, where the corresponding problem features are terrain features of different locations and the characteristics of different patients that are correlated to the missing parameters, respectively. These correlated features allow us to predict the missing parameters even if we do not have any trajectories of the MDP.
Offline off-policy evaluation We evaluate a policy π in a fully offline setting with trajectories T = {τi}, τi = (si1, ai1, ri1, · · · , sih, aih, rih) generated from the MDP using behavior policy πbeh. We use an OPE metric used by Futoma et al. [11] — we evaluate a policy π and trajectories T as:
EvalT (π) := V CWPDIS(π)− λESS/ √ ESS(π) (1)
where V CWPDIS(π) := ∑h t=1 γ t ∑ i ritρit(π)∑ i ρit(π) and ESS(π) := ∑h t=1 ( ∑ i ρit) 2∑ i ρ 2 it , and ρit(π) is the ratio of
the proposed policy and the behavior policy likelihoods up to time t: ρit(π) := ∏t t′=1 π(ait′ |sit′ ) πbeh(ait′ |sit′ ) .
Optimization formulation Given a set of training features and trajectories Dtrain = {(xi, Ti)}i∈Itrain , our goal is to learn a predictive model mw to optimize the training performance:
maxw E(x,T )∈Dtrain [EvalT (π ∗(mw(x)))] (2)
The testing performance is evaluated on the unseen test set Dtest = {(xi, Ti)}i∈Itest with trajectories hidden from training, and only used for evaluation: E(x,T )∈Dtest [EvalT (π∗(mw(x)))].

heading: 3 Two-stage Learning
text: To learn the predictive model mw from trajectories, the standard approach is to minimize an expected predictive loss, which is defined by comparing the prediction θ = mw(x) with the trajectories T :
minw E(x,T )∼Dtrain L(θ, T ) where L(θ, T ) = Eτ∼T `θ(τ), θ = mw(x) (3)
For example, when the rewards are missing, the loss could be the squared error between the predicted rewards and the actual rewards we see in the trajectories for each individual step. When the transition probabilities are missing, the loss could be defined as the negative log-likelihood of seeing the trajectories in the training set.
In the first stage, to train the predictive model, we run gradient descent to minimize the loss function defined in Equation (3) and make prediction on the model parameter θ = mw(x) of each problem. In the second stage, we solve each MDP problem with the predicted parameter θ using an RL algorithm to generate the optimal policy π∗(θ). However, predictive loss and the final evaluation metric are commonly misaligned especially in deep learning problems with imbalanced data [6, 14, 15, 20]. This motivates us to learn the predictive model end-to-end and therefore avoid the misalignment.

heading: 4 Decision-focused Learning
text: In this section, we present our main contribution, decision-focused learning in sequential decision problems, as illustrated in Figure 1. Decision-focused learning integrates an MDP problem into the training pipeline to directly optimize the final performance. Instead of relying on a predictive loss to train the predictive model mw, we can directly optimize the objective in Equation (2) by running end-to-end gradient descent to update the predictive model mw:
d Eval(π∗) dw = d Eval(π∗) dπ∗ dπ∗ dθ dθ dw (4)
We assume the policy π∗(θ) is either stochastic and smooth with respect to the change in the parameter θ, which is common in settings with continuous state or action spaces, or that an appropriate regularization term [12, 13] is used to improve the smoothness of the policy. More discussions about the smoothness can be found in Appendix B.1.
This gradient computation requires us to back-propagate from the final evaluation through the MDP layer to the predictive model mw that we want to update. The major challenge in Equation (4) is to compute dπ ∗
dθ , which involves differentiating through an MDP layer solved by an RL algorithm. In the following section, we first discuss two different optimality conditions in MDPs, which are later used to convert into KKT conditions and differentiate through to compute dπ ∗
dθ . We then discuss two computational challenges associated with the derivative computation.

heading: 4.1 Optimality Conditions in MDPs
text: When the predicted model parameter θ = mw(x) is given, the MDP can be solved by any RL algorithm to get an optimal policy π∗. Here we discuss two common optimality conditions in MDPs, differing by the use of policy gradient or Bellman equation: Definition 1 (Policy gradient-based optimality condition). Defining Jθ(π) to be the expected cumulative reward under policy π, the optimality condition of the optimal policy π∗ is:
π∗ = argmaxπ Jθ(π) where Jθ(π) := Eτ∼π,θ Gθ(τ) (5) where Gθ(τ) is the discounted value of trajectory τ given parameter θ, and the expectation is taken over the trajectories following the optimal policy and transition probability (as part of θ). Definition 2 (Bellman-based optimality condition). Defining Jθ(π) to be the mean-squared Bellman error3 under policy π, the optimality condition of the optimal policy π∗ (valuation function) is:
π∗ = argmin π Jθ(π) where Jθ(π) := Eτ∼π,θ
1 2 δ2θ(τ, π) (6)
where δθ(τ, π) = ∑
(s,a,s′)∈τ Qπ(s, a)− Rθ(s, a)− γ Ea′∼π Qπ(s′, a′) is the total Bellman error of a trajectory τ , and each individual term δθ(τ, π) can depend on the parameter θ because the Bellman error depends on the immediate reward Rθ, which can be a part of the MDP parameter θ. The expectation in Equation (6) is taken over all the trajectories generated from policy π and transition probability (as part of θ).

heading: 4.2 Backpropagating Through Optimality and KKT Conditions
text: To compute the derivative of the optimal policy π∗(θ) in an MDP with respect to the MDP parameter θ, we differentiate through the KKT conditions of the corresponding optimality conditions: Definition 3 (KKT Conditions). Given objective Jθ(π) in an MDP problem, since the policy parameters are unconstrained, the necessary KKT conditions can be written as: ∇πJθ(π∗) = 0.
In particular, computing the total derivative of KKT conditions gives:
0 = d
dθ ∇πJθ(π∗)=
∂
∂θ ∇πJθ(π∗) +
∂
∂π ∇πJθ(π∗)
dπ∗
dθ =∇θπJθ(π∗)+∇2πJθ(π∗)
dπ∗
dθ
=⇒ dπ ∗
dθ = −(∇2πJθ(π∗))−1∇2θπJθ(π∗) (7)
We can use Equation (7) to replace the term dπ ∗
dθ in Equation (4) to compute the full gradient to back-propagate from the final evaluation to the predictive model parameterized by w:
d Eval(π∗) dw = −d Eval(π ∗) dπ∗ (∇2πJθ(π∗))−1∇2θπJθ(π∗) dθ dw (8)

heading: 4.3 Computational Challenges in Backward Pass
text: Unfortunately, although we can write down and differentiate through the KKT conditions analytically, we cannot explicitly compute the second-order derivatives∇2πJθ(π∗) and∇2θπJθ(π∗) in Equation (8) due to the following two challenges:
Large state and action spaces involved in optimality conditions The objectives Jθ(π∗) in Definition 1 and Definition 2 involve an expectation over all possible trajectories, which is essentially an integral and is intractable when either the state or action space is continuous. This prevents us from explicitly verifying optimality and writing down the two derivatives∇2πJθ(π∗) and∇2θπJθ(π∗).
High-dimensional policy space parameterized by neural networks In MDPs solved by modelfree deep RL algorithms, the policy space π ∈ Π is often parameterized by a neural network, which has a significantly larger number of variables than standard optimization problems. This large dimensionality makes the second-order derivative∇2πJθ(π∗) ∈ Rdim(π)×dim(π) intractable to compute, store, or invert.
3We use the same notation J to denote both the expected cumulative reward and the expected Bellman error to simplify the later analysis of decision-focused learning.

heading: 5 Sampling Unbiased Derivative Estimates
text: In both policy gradient–based and Bellman–based optimality conditions, the objective is implicitly given by an expectation over all possible trajectories, which could be infinitely large when either state or action space is continuous. This same issue arises when expressing such an MDP as a linear program — there are infinitely many constraints, making it intractable to differentiate through.
Inspired by the policy gradient theorem, although we cannot compute the exact gradient of the objective, we can sample a set of trajectories τ = {s1, a1, r1, . . . , sh, ah, rh} from policy π and model parameter θ with finite time horizon h. Denoting pθ(τ, π) to be the likelihood of seeing trajectory τ , we can compute an unbiased derivative estimate for both optimality conditions: Theorem 1 (Policy gradient-based unbiased derivative estimate). We follow the notation of Definition 1 and define Φθ(τ, π) = ∑h i=1 ∑h j=i γ jRθ(sj , aj) log π(ai|si). We have:
∇πJθ(π) = Eτ∼π,θ [∇πΦθ(τ, π)] =⇒ ∇2πJθ(π) = Eτ∼π,θ
[ ∇πΦθ · ∇π log p>θ +∇2πΦθ ] ∇2θπJθ(π) = Eτ∼π,θ [ ∇πΦθ · ∇θ log p>θ +∇2θπΦθ
] (9) Theorem 2 (Bellman-based unbiased derivative estimate). We follow the notation in Definition 2 to define Jθ(π) = 12 Eτ∼π,θ [ δ2θ(τ, π) ] . We have:
∇πJθ(π) = Eτ∼π,θ [ δ∇πδ + 1
2 δ2∇π log pθ
] =⇒ ∇2πJθ(π) = Eτ∼π,θ [ ∇πδ∇πδ> +O(δ) ] ∇2θπJθ(π)=Eτ∼π,θ [ ∇πδ∇θδ>+ ( ∇πδ∇θ log p>θ +∇π log pθ∇θδ>+∇2θπδ ) δ +O(δ2) ] (10)
For the latter, we apply the fact that at the near-optimal policy, the Bellman error is close to 0 and thus each individual component δ(τ, π) is small to simplify the analysis. Refer to the appendix for the full derivations of Equations (9) and (10).
Equations (9) and (10) provide a sampling approach to compute the second-order derivatives, avoiding computing an expectation over the large trajectory space. We can use the optimal policy derived in the forward pass and the predicted parameters θ to run multiple simulations to collect a set of trajectories. These trajectories from predicted parameters can be used to compute each individual derivative in Equations (9) and (10).

heading: 6 Resolving High-dimensional Derivatives by Low-rank Approximation
text: Section 5 provides sampling approaches to compute an unbiased estimate of second-order derivatives. However, since the dimensionality of the policy space dim(π) is often large, we cannot explicitly expand and invert∇2πJθ(π∗) to compute∇2πJθ(π∗)−1∇2θπJθ(π∗), which is an inevitable step toward computing the full gradient of decision-focused learning in Equation (8). In this section, we discuss various ways to approximate ∇2πJθ(π∗) and how we use low-rank approximation and Woodbury matrix identity [37] to efficiently invert the sampled Hessian without expanding the matrices. Let n := dim(π) and k n to be the number of trajectories we sample to compute the derivatives.

heading: 6.1 Full Hessian Computation
text: In Equations (9) and (10), we can apply auto-differentiation tools to compute all individual derivatives in the expectation. However, this works only when the dimensionality of the policy space π ∈ Π is small because the full expressions in Equations (9) and (10) involve computing second-order derivatives , e.g.,∇2πΦθ in Equation (10), which is still challenging to compute and store when the matrix size n× n is large. The computation cost is O(n2k) +O(nω) dominated by computing all the Hessian matrices and the matrix inversion with 2 ≤ ω ≤ 2.373 the complexity order of matrix inversion.

heading: 6.2 Approximating Hessian by Constant Identity Matrix
text: One naive way to approximate the Hessian∇2πJθ(π∗) is to simply use a constant identity matrix cI . We choose c < 0 for the policy gradient–based optimality in Definition 1 because the optimization
Algorithm 1: Decision-focused Learning for MDP Problems with Missing Parameters 1 Parameter: Training set Dtrain = {(xi, Ti)}i∈Itrain , learning rate α, regularization λ = 0.1 2 Initialization: Initialize predictive model mw : X → Θ parameterized by w 3 for epoch ∈ {1, 2, . . . }, each training instance (x, T ) ∈ Dtrain do 4 Forward: Make prediction θ = mw(x). Compute two-stage loss L(θ, T ). Run model-free RL to get an optimal policy π∗(θ) on MDP problem using parameter θ. 5 Backward: Sample a set of trajectories under θ, π∗ to compute ∇2πJθ(π∗),∇2θπJθ(π∗) 6 Gradient step: Set ∆w = d EvalT (π
∗) dw − λ dL(θ,T ) dw by Equation (8) with predictive loss L as
regularization. Run gradient ascent to update model: w ← w + α∆w 7 Return: Predictive model mw.
problem is a maximization problem and thus is locally concave at the optimal solution, whose Hessian is negative semi-definite. Similarly, we choose c > 0 for the Bellman–based optimality in Definition 2. This approach is fast, easily invertible. Moreover, in the Bellman version, Equation (8) is equivalent to the idea of differentiating through the final gradient of Bellman error as proposed by Futoma et al. [11]4. However, this approach ignores the information provided by the Hessian term, which can often lead to instability as we later show in the experiments. In this case, the computation complexity is dominated by computing∇2θπJθ(π∗), which requires O(nk).

heading: 6.3 Low-rank Hessian Approximation and Application of Woodbury Matrix Identity
text: A compromise between the full Hessian and using a constant matrix is approximating the secondorder derivative terms in Equations (9) and (10) by constant identity matrices, while computing the first-order derivative terms with auto-differentiation. Specifically, given a set of k sampled trajectories {τ1, τ2, · · · , τk}, Equations (9) and (10) can be written and approximated in the following form:
∇2πJθ(π) ≈ 1
k ∑k i=1 ( uiv > i +Hi ) ≈ 1 k ∑k i=1 ( uiv > i + cI ) = UV > + cI (11)
where U = [u1, u2, · · · , uk]/ √ k ∈ Rn×k, V = [v1, v2, · · · , vk]/ √ k ∈ Rn×k and ui, vi ∈ Rn correspond to the first-order derivatives in Equations (9) and (10), andHi corresponds to the remaining terms that involve second-order derivatives. We use a constant identity matrix to approximate Hi, while explicitly computing the remaining parts to increase accuracy.
However, we still cannot explicitly expand UV > ∈ Rn×n since the dimensionality is too large. Therefore, we apply Woodbury matrix identity [37] to invert Equation (11):
(∇2πJθ(π))−1 ≈ (UV > + cI)−1 = 1 c I − 1 c U(cI − V >U)−1V > (12)
where V >U ∈ Rk×k can be efficiently computed with much smaller k n. When we compute the full gradient for decision-focused learning in Equation (8), we can then apply matrix-vector multiplication without expanding the full high-dimensional matrix, which results in a computation cost of O(nk + kω) that is much smaller than the full computation cost O(n2k + nω).
The full algorithm for decision-focused learning in MDPs is illustrated in Algorithm 1.

heading: 7 Example MDP Problems with Missing Parameters
text: Gridworld with unknown cliffs We consider a Gridworld environment with a set of training and test MDPs. Each MDP is a 5×5 grid with a start state located at the bottom left corner and a safe state with reward drawn from N (5, 1) located at the top right corner. Each intermediate state has a reward associated with it, where most of them give the agent a reward drawn from N (0, 1) but 20% of the them are cliffs and give N (−10, 1) penalty to the agent. The agent has no prior information
4The gradient of Bellman error can be written as∇πJθ(π∗) where the policy π is the parameters of the value function approximator and J is defined as the expected Bellman error. The derivative of the final gradient can be written as∇w(∇πJθ(π∗)) = ∇2θπJθ(π∗) dθdw by chain rule, which matches the last three terms in Equation 8 when the Hessian is approximated by an identity matrix.
about the reward of each grid cell (i.e., the reward functions of the MDPs are unknown), but has a feature vector per grid cell correlated to the reward, and a set of historical trajectories from the training MDPs. The agent learns a predictive model to map from the features of a grid cell to its missing reward information, and the resulting MDP is used to plan. Since the state and action spaces are both finite, we use tabular value-iteration [30] to solve the MDPs.
Partially observable snare-finding problems with missing transition function We consider a set of synthetic conservation parks, each with 20 sites, that are vulnerable to poaching activities. Each site in a conservation park starts from a safe state and has an unknown associated probability that a poacher places a snare at each time step. Motivated by [38], we assume a ranger who can visit one site per time step and observes whether a snare is present. If a snare is present, the ranger removes it and receives reward 1. Otherwise, the ranger receives reward of −1. The snare can stay in the site if the ranger does not remove it, which makes the snare-finding problem a sequential problem rather than a multi-armed bandit problem. As the ranger receives no information about the sites that they do not visit, the MDP belief state is the ranger’s belief about whether a snare is present. The ranger uses the features of each site and historical trajectories to learn a predictive model of the missing transition probability of a snare being placed. Since the belief state is continuous and the action space is discrete, given a predictive model of the missing transition probability, the agent uses double deep Q-learning (DDQN) [32] to solve the predicted MDPs.
Partially observable patient treatment problems with missing transition probability We consider a version of the Tuberculosis Adherence Problem explored in [23]. Given that the treatment for tuberculosis requires patients to take medications for an extended period of time, one way to improve patient adherence is Directly Observed Therapy, in which a healthcare worker routinely checks in on the patient to ensure that they are taking their medications. In our problem, we consider 5 synthetic patients who have to take medication for 30 days. Each day, a healthcare worker chooses one patient to intervene on. They observe whether that patient is adhering or not, and improve the patient’s likelihood of adhering on that day, where we use the number of adherent patients as the reward to the healthcare worker. Whether a patient actually adheres or not is determined by a transition matrix that is randomly drawn from a fixed distribution inspired by [18]. The aim of the prediction stage is to use the features associated with each patient, e.g., patient characteristics, to predict the missing transition matrices. The aim of the RL stage is then to create an intervention strategy for the healthcare worker such that the sum of patient adherence over the 30-day period is maximised. Due to partial observability, we convert the problem to its continuous belief state equivalence and solve it using DDQN.
Please refer to Appendix C for more details about problem setup in all three domains.

heading: 8 Experimental Results and Discussion
text: In our experiments, we compare two-stage learning (TS) with different versions of decision-focused learning (DF) using two different optimality conditions, policy gradient (PG) and Bellman equationbased (Bellman), and two different Hessian approximations (Identity, Woodbury) defined in Section 6. Computing the full Hessian (as in Section 6.1) is computationally intractable. Across all three examples, we use 7 training MDPs, 1 validation MDP, and 2 test MDPs, each with 100 trajectories.
The predictive model is trained on the training MDP trajectories for 100 epochs. Performance is evaluated under the Off-Policy Evaluation (OPE) metric of Equation (1) with respect to the withheld test trajectories. In the following, we will discuss how DF variants work compared with TS methods, and explore why some methods are better. We use two different trajectories, random and near-optimal, in the training MDP to model different imbalanced information given to train the predictive model. The results are shown in Table 1.
Decision-focused learning with the Woodbury matrix identity outperforms two-stage learning Table 1 summarizes the average OPE performance on the test MDPs. We can see that in all of the three problem settings, the best performances are all achieved by decision-focused learning. However, when Hessian approximation is not sufficiently accurate, decision-focused learning can sometimes perform even worse than two-stage (e.g., PG-Id and Bellman-Id in the snare problem). In contrast, decision-focused methods using a more accurate low-rank approximation and Woodbury matrix identity (i.e., PG-W and Bellman-W), as discussed in Section 6.3, dominate two-stage performance in the test MDPs across all settings.
Low predictive loss does not imply a winning policy In Figures 2(a), 3(a), we plot the predictive loss curve in the training MDPs over different training epochs of Gridworld and snare problems. In particular, two-stage approach is trained to minimize such loss, but fails to win in Table 1. Indeed, low predictive loss on the training MDPs does not always imply a high off-policy evaluation on the training MDPs in Figure 2(b) due to the misalignment of predictive accuracy and decision quality, which is consistent with other studies in mismatch of predictive loss and evaluation metric [6, 14, 15, 20].
Comparison between different Hessian approximations In Table 1, we notice that more inaccurate Hessian approximation (identity) does not always lead to poorer performance. We hypothesize that this is due to the non-convex off-policy evaluation objective that we are optimizing, where higher variance might sometimes help escape local optimum more easily. The identity approximation is more unstable across different tasks and different trajectories given. In Table 1, the performance of Bellman-Identity and PG-Identity sometimes lead to wins over two stage and sometimes losses.
Comparison between policy gradient and Bellman-based decision-focused learning We observe that the Bellman-based decision-focused approach consistently outperforms the policy gradientbased approach when the trajectories are random, while the policy gradient-based decision-focused approach mostly achieves better performance when near-optimal trajectories are present. We hypothesize that this is due to the different objectives considered by different optimality conditions. The Bellman error aims to accurately cover all the value functions, which works better on random trajectories; the policy gradient aims to maximize the expected cumulative reward along the optimal policy only, which works better with near-optimal trajectories that have better coverage in the optimal regions.
Computation cost Lastly, Figures 2(c) and 3(c) show the backpropagation runtime of the policygradient based optimality condition per gradient step per training instance across different Hessian approximations and different problem sizes in the gridworld and snare finding problems. To train the model, we run 100 epochs for every MDP in the training set, which immediately makes the full Hessian computation intractable as it would take more than a day to complete.
Analytically, let n be the dimensionality of the policy space and k n be the number of sampled trajectories used to approximate the derivatives. As shown in Section 6, the computation cost of full Hessian O(n2 + nω) is quadratic in n and strictly dominates all the others. In contrast, the costs of the identity matrix approximation O(nk) and the Woodbury approximation O(nk + kω) are both linear in n. The Woodbury method offers an option to get a more accurate Hessian at low additional cost.

heading: 9 Conclusion
text: This paper considers learning a predictive model to address the missing parameters in sequential decision problems. We successfully extend decision-focused learning from optimization problems to MDP problems solved by deep reinforcement learning algorithms, where we apply sampling and low-rank approximation to Hessian matrix computation to address the associated computational challenges. All our results suggest that decision-focused learning can outperform two-stage approach by directly optimizing the final evaluation metric. The idea of considering sequential decision problems as differentiable layers also suggests a different way to solve online reinforcement learning problems, which we reserve as a future direction.

heading: Acknowledgement
text: This research was supported by MURI Grant Number W911NF-17-1-0370 and W911NF-18-1-0208. Chen and Perrault acknowledge support from the Center for Research on Computation and Society. Doshi-Velez acknowledges support from NSF CAREER IIS-1750358. The computations in this paper were run on the FASRC Cannon cluster supported by the FAS Division of Science Research Computing Group at Harvard University.
Declaration of Interests
Doshi-Velez reports grants from the National Institutes of Health, and personal fees from Davita Kidney Care and Google Health via Adecco, outside of the submitted work. Tambe is jointly affiliated with Google Research India, outside of the submitted work. Wang, Shah, Chen, and Perrault declare no competing interests.

Review summaries:
In this paper, the authors consider the problem of learning a predictive model for MDPs with missing parameters. Instead of a standard two-stage approach, an end-to-end approach based on differentiating through an optimality condition is proposed. To address the associated computation issues, a sampling method and a low-rank approximation method are introduced. Empirical evaluations on three synthetic settings are provided to show that the end-to-end approach outperforms the two-stage approach. 

The paper proposes a learning problem in which a partially defined MDP is provided with missing parameter values, particularly a subset of the transition probabilities and reward values. Input features defining each MDP are provided, along with sample trajectories from its stochastic policy. The goal is to learn to predict the missing parameter values, by targeting the sample trajectories given the MDP features. The prediction model is trained end-to-end across several stages: Prediction of missing MDP parameters from feature values, followed by reinforcement learning to solve the MDP for an optimal policy, and finally evaluation of the policy as an expectation of some loss over its trajectories. The primary challenge is differentiation through the reinforcement learning layer, which is achieved by differentiating through the optimality conditions of the decision problem. Computational intractability of differentiating with respect to large state and action spaces is handled with one of multiple proposed gradient approximations. 

The authors extend the "predict-then-optimize" framework to sequential decision problems. The framework which has been studied previously focuses on improving the quality of the decision-making process compared to the traditional two-stage prediction-optimization approach. In this paper, the sequential decision-making problem is formulated as a MDP, with the goal being to learn a predictive model that can predict missing parameters. At training time, a predictive model to map features to missing parameters is learned, and at test time, the learned model is used to make predictions without using trajectories. The predicted parameters are used to solve the test MDPs, yielding an optimal policy which is evaluated using an offline off-policy evaluation.  The authors study two types of optimality conditions: (1) a policy gradient based approach where the expected cumulative reward is maximized, (2) a Bellman based optimality condition where mean squared Bellman error is minimized. These optimality conditions are expressed using KKT conditions. However, due to the large state and action spaces involved in the optimization reformulation, and the high-dimensionality of the policy space, the authors propose a technique to sample an estimate of the first and second order derivatives to approximate the optimality and KKT conditions. To handle the second challenge, two options are proposed - approximating the Hessian by constant identity matrix, and using a low-rank Hessian approximation and application of Woodbury matrix identity.

The authors use the predict-then-optimize framework to tackle prediction of MDP parameters based on input features. They achieve this by formulating MDP optimality KKT conditions and differentiating through them. In addition, the paper provides an ubiased sampling method for estimation of the KKT derivatives and a low-rank approximation to the high-dimensional sample-based derivatives. The authors name their approach decision-focused learning.