Learning MDPs from Features: Predict-Then-Optimize for Sequential Decision Making by Reinforcement Learning
Authors: Kai Wang
Sanket Shah
Haipeng Chen
Andrew Perrault
Finale Doshi-Velez
Milind Tambe
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: Optimization
predict-then-optimize
decision-focused learning
differentiable optimization
reinforcement learning
MDP
low-rank approximation
Woodbury matrix identity
KKT conditions
optimality conditions
sequential decision problems

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.

Learning MDPs from Features: Predict-Thee Then-Optimize for Sequential Decision Problems by Reinforcement Learning
Kai Wang Harvard University Cambridge, MA kaiwangeg. harvard.
Sanket Shah Harvard University Cambridge, MA sanketshahOg. harvard. edu
Haipeng Chen Harvard University Cambridge, MA hpchen@seas hww.aaaa.aaaaaadd.dd... edu
Andrew Perrault*
Finale Doshi-Velex
Milind Tambe
The Ohio State University Columbus, OH perrault. 178.su. edu
Harvard University Cambridge, MA finale@seas harvard.
Harvard University Cambridge, MA
edu
Abstract
In the predict-then-opimiize 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-focusee 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 context sequential decision problems (formulated as MDPs) that are solved via reinforcement learning. In particular, trajectories from training MDPs, which we use to train a predictive model that generalizes to unseen test MDPs without trajectories. Two significant computa- tional challenges arise in applying decision-focuset learning to MDPs: (i) large state and action spaces make infeasible for existing techniques to differentiate through MDP problems, and (ii) the high-dimensional policy space, as parameter- ized by neural network, makes differentiating through a policy expensive. We resolve the first challenge by sampling provably unbiased derivatives to approxi- mate and differentiate through optimality conditions, and the second challenge by using low-rank approximation to the high- 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 unseen tasks.
framework in the
are given environment features and set
Introduction
Predict-thed seaetttthe---otiiii 5 is framework for solving optimization problems with unknown parame- ters. Given such a problem, we first train predictive model to predict the missing parameters from problem features. Our objective is maximize the resulting decision quality when the optimization problem is subsequently sol ved with the predicted parameters [25 28 Recent work on the decision- focused learning approach 7 36) embeds the optimization problem 4 into the training pipeline and trains the predictive model end-to-end to optimize the final decision quality. Compared with
* This work was done while author was Harvard University.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

more traditional "two-stage" stage" approach which maximizes the predictive accuracy of the model (rather than the final decision quality), the rrcision-focusee learning approach can achieve a higher solution quality and generalize better to unseen tasks.
This paper studies the treooii-ten-optimiize 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 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 poolyy evaluation (OPE) as shown in Figure[[ ! 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 Ransfer-RR.Aa unpatrolled areas. These settings differ from those considered in transfer- RL. 221 24 29 3T and
meta-RL S 10 33
35
40
because we generalize across different MDPs by explicitly predicting
the mapping function from features to missing MDPs parameters, while transfer/metaa RL achieve generalization by learning hidden representation of different MDPs implicitly with trajectories.
The main contribution of this paper 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: Bellman- based approach where mean-squared Bellman error is minimized, and policy gradient-based approach where the expected cumulative reward is maximized. We convert these optimality conditions into their corresponding Karush-Kuhn- Kuhn-Tuccer (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- 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 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 deriva- tives. 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 low-rank approximation to further approximate the sample-baeed second-order derivatives. This low-rank approximation reduces both the computation cost and the memory usage of differentiating through KKT conditions.
learning to
We empirically test our decision-focusee algorithms on three settings: grid world with unknown rewards, and mare--inding and Tuberculosis treatment problems where transition probabilities are unknown. Decision-focused learning achieves better OPE performance in unseen test MDPs than two- stage approach, and our low-rank approximations significantly scale-up decision-focused learning.
Related Work
Differentiable optimization Amos et al. (2] propose using quadratic program as a differentiable
layer and embedding into deep learning pipeline, and Agrawal al. [ extend their work convex programs. recision-focused learning 7 36 focuses on the
[5] 9 framework by embedding an optimization layer into training pipeline, where the optimization layers can be convex 77.. linear [22) 36]. and non-convex (2 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-fort actions, reducing the dimensionality. Karkus et al. [16] differentiate through an algorithm by unrolling and relaxing

---

Backpropagation (decision-focused)
Predictive model ())
vector
Prediction MDP parameter mw (x)
Planning RL solver max (x)
Optimal policy (0)
Evaluation
(to be learned)
Evaly (xr (0)
Backpropagation (two-stage)
Trajectories T' = [(s, a.r); [(s,a,))
Predictive loss
(available in the training MDP problems only)
Offline off-policy evaluation
Figure I: We consider learning a predictive model to map from features unknown MDP parameters and obtaining a policy by solving the predicted MDP with RL. Two-stage learning learns the predictive model by minimizing a predictive loss function, whereas decision- learning is trained end-to-end to maximize the final off pollcy evaluation performance.
all the strict operators by soft operators. Futoma et al. 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 low-rank approximation to form an unbiased estimate of the optimality conditions differentiate through, and show that the approach of Futoma et al. 11 is included in ours as special case.
Predict-then-optimize and offline reinforcement learning The idea of planning under predicted MDP arises in model-based RI. as certainty equivalence [19. It has been extended to offline settings [17 39 who learn pessimistic MDP before solving for the policy. Our setting differs because of the presence of features and train-test split- test MDPs are completely fresh without any trajectories. Our setting also resembles meta RL (c.g.. 88. 10
33 35 40) and transfer RL (e.g.. 24 29 31I.) 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.
21
2 Problem Statement
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 (8, so A. T, R) with an initial state sp. 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 0" Additionally, we assume there are problem features x associated with each MDP, where (0", 2)) is correlated and drawn from the same unknown, but fixed, distributior
We are given a set of training MDPs and a set of test MDPs, each with missing parameters / and features 2. Each training MDP is accompanied by set of trajectories T performed by behavior policy Thch 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 m., map from features 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 predictive model m.. predict the missing parameters 0 = (z)) The predicted parameters are used to solve the test MDPs, yielding the policy π (m)(()) where our offline off off-policy evaluation (OPE) metric is maximized. This framework is illustrated in Figure |
'Examples of the missing parameters 0" include the poaching risk of different locations in wildlife conser- vation and the transition probability patients' healthiness in healtheare 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.
3

---

Offline off-policy evaluation We evaluate policy x in fully offline setting with trajectories T { T; (8t1, aas T. Sih ash Ti generated from the MDP using behavior policy Theh-
We use an OPE metric used by Futoma et al. we evaluate policy and trajectories 7 as:
(x)
VCWPDIS AESS/ EES((T)
(1)
where VCWPDIS (T) = Ex=
rupu(*] and ESS Ei= E Piee e.e) (t)
E. and P(((T) is the ratio of (la,)))))))
the proposed policy and the behavior policy likelihoods up to time pa( = Ie mailaa Par)"
Optimization formulation Given set of training features and trajectories Dimin {1T,, Ti) our goal is to learn a predictive model m. to optimize the training performance:
max D (= (x)())
The testing performance is evaluated on the unseen test set Diett {ax Ti) with trajectories hidden from training, and only used for evaluation: D(,, [Evaly ( (T* (m)))
3 Two-stage Learning
To learn the predictive model m.r from trajectories, the standard approach is to minimize an expected predictive loss, which is defined by comparing the prediction (m) with the trajectories TT:
min (ir, C(O,T)
where
E.TT Lo(T), = m
(3)
For example, when the rewards 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 the negative log-likelihood of seeing the trajectories in the training set.
In the first stage, to train the predictive model, we gradient descent minimize the loss function defined in Equation (3) and make prediction on the model parameter = (,) of each problem. In the second stage, solve each MDP problem with the predicted parameter 0 using an RL algorithm to generate the optimal policy x² (0). 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.
4 Decision-focused Learning
In this section, we present our main contribution, decision-focuseee learning in sequential decision problems, as illustrated in Figure [ 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 m,, we can directly optimize the objective in Equation (2) by running end-to-end gradient descent to update the predictive model mw:
d Eval(TT) Eval(TT) dx d8 dw
(4)
dr d0 dw
We assume the policy T (0) is either stochastic and smooth with respect to the change in the parameter 9 which is common in settings with continuous state action spaces, or that appropriate regularization term [2 13 is used to improve the smoothness of the policy. More discussions about the smoothness can be found in Appendix B.I
This gradient computation requires us to back-propagate from the final evaluation through the MDP layer to the predictive model m. that we want to update. The major challenge in Equation (4) is to compute da" 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 de computational challenges associated with the derivative computation.
We then discuss two

---

4.1 Optimality Conditions in MDPs
When the predicted model parameter e (=) is given, the MDP can be solved by any RL algorithm to get an optimal policy x². Here we discuss two common optimality conditions in MDPs, differing by the use of policy gradient or Bellman equation:
Definition 1 (Policy gradient-base based optimality condition). Defining Jo(T) to be the expected camu- lative reward under policy n, the optimality condition of the optimal policy IT" is:
arg max, (
where JØ(N) == ET Go(T)
(5)
where (T) is the discounted value trajectory given parameter and the expectation is taken over the trajectories following the optimal policy and transition probability (as part of 0).
Definition (Bellman-base based optimality condition). Defining J be the mean geea--suuared Bellman erro under policy æ, the optimality condition of the optimal policy HT" (valuation function) is:
1
x arg migminJ... Jo(
where Jo(a) = E- (T, mx)
(6)
x
where óØ(T, T.) E (s....")er Qx(s,a) - Rø(s, - -ai~xx (n','aa') is the total Bellman error of trajectory T, and each individual term dG(T, depend on the parameter 0 because the Bellman error depends on the immediate reward Rg. which can be a part of the MDP parameter 0. The expectation in Equation (6) is taken over all the trajectories generated from policy and transition probability (as part of 8).
4.2
Backpropagating Through Optimality and KKT Conditions
To compute the derivative of the optimal policy x (0) in an MDP with respect to the MDP parameter 8, we differentiate through the KKT conditions of the corresponding optimality conditions:
Definition 3 (KKT Conditions). Given objective (T( in an MDP problem, since the policy param- eters are unconstrained, the necessary KKT conditions can be written as: VoJoo(T)) O.
In particular, computing the total derivative of KKT conditions gives:
d V JoJ" V
Jo(T") dm
d0 = VorJo( ( di
a0 vooo(==)) +
d0 dT
d0
dB = -V
da We can use Equation (7) replace the term de in Equation 44 to compute the full gradient to back-propagate from the final evaluation to the predictive model parameterized by w:
d Eval(=*) Eval( dw
d0
Br Jo Jo(TT)
(8)
dx
dw
4.3 Computational Challenges in Backward Pass
Unfortunately, although we can write down and differentiate through the KKT conditions analytically, we cannot explicitly compute the second-order derivatives JeJe(("" and Jo(T") in Equation due to the following two challenges:
Large state and action spaces involved in optimality conditions The objectives (x="") in Defi- Litioon[ame Definition Deffinttionniivolll involve 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 derivatives V 22 and V2 Jo
High--dimensional policy space parameterized by neural networks In MDPs solved by model- free deep RL algorithms, the policy space E II 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 Jo("") E (Rdi (edlim[=)xxdmm dx) intractable to compute, store, invert.
We use the same notation J to denote both the expected cumulative reward and the expected Bellman error simplify the later analysis of decision- focused learning.
5

---

Sampling Unbiased Derivative Estimates
In both policy gradient-based and Bellman- 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 linear program there are infinitely many constraints, making it intractable to differentiate through.
Inspired by the policy gradient theorem, although cannot compute the exact gradient of the objective, we can sample a set of trajectories т = {si, @1, ..,.... 8h: Oh. Th from policy and model parameter 0 with finite time horizon h. Denoting po(T; æ)) to be the likelihood of seeing trajectory т, we can compute an unbiased derivative estimate for both optimality conditions:
Theorem (Policy gradient-basee based unbiased derivative estimate). We follow the notation of Defini- tion 1 and define Do (T( = Σm Ein 7 Ro(sj, az) log o (a:)s:). We have:
(OT = Er~r, V Va log + 22
V Jo( Bron 0 V LO
(9)
Jo(T)) = Er~,, V Vo log + 2.
Theorem 2 (Bellman-based unbiased derivative estimate). We follow the notation in Definition
define (T( = 2 En~e, [T, m)]. have:
V Jo(T) = Er
oV 5 + 22 x log
= VoJo(T) = Eran, V VV + +(()]
Jo = E 8
8VO VV log
V
+V 0(5²)
(10)
For the latter, we apply the fact that at the near-optimal policy, the Bellman error close to 0 and thus each individual component 0(T, v) is small to simplify the analysis. Refer to the appendix for the full derivations of Equations (9) and (TO).
Equations 99) and (10) provide 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 0 to run multiple simulations to collect set of trajectories. These trajectories from predicted parameters can be used to compute each individual derivative in Equations (9) and (TO).
6 Resolving High-dimensional Derivatives by Low-rank Approximation
(ection[5]prvvio 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 J=="((")) to compute (TT)) Jo(T") which is an inevitable step toward computing the full gradient of decision-focusee learning in Equation (8). In this section, discuss various ways to approximate (TT)((T") 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(m) (T) and < n to be the number of trajectories we sample to compute the derivatives.
6.1 Full Hessian Computation
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 π I7 small because the full expressions in Equations (9) and (IO) involve computing second-order derivatives c.g., 2 in Equation (IO), which is still challenging to compute and store when the matrix size the Hessian matrices and the matrix inversion with S w < 2.373 the complexity order of matrix inversion.
large. The computation cost is O(n k) + O(n") dominated by computing all
6.2 Approximating Hessian by Constant Identity Matrix
One naive way to approximate the Hessian (or") is to simply use a constant identity matrix el. We choose < 0 for the policy gradient-based optimality in Definition I because the optimization

---

Algorithm 1: Decision-focused Learning for MDP Problems with Missing Parameters
Parameter: Training set Drnin {T.,, Ta) Ie learning rate a, regularization À 0.1
2 Initialization; Initialize predictive model myw : X - θ parameterized by w
3 for epoch E {1. ........, }, each training instance (x, T) E Diroin do
Forward: Make prediction 0 = (x) Compute two-stage loss C(0, T). Run model-free
RL to get an optimal policy x (0) on MDP problem using parameter
Backward: Sample set of trajectories under 8, n² to compute V²Jo(="), Jo(""
Gradient step: Set Aw Evaly
_ by Equation ( with predictive loss L as regularization. Run gradient ascent to update model: w f- w + aAw
der
dw
7 Return:
Predictive model myw-
problem is maximization problem and thus is locally concave the optimal solution, whose
Hessian is negative semi-definite. Similarly, we choose e > 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 al. [TT 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 V (o(T)) which requires ((n)).
6.3 Low-rank Hessian Approximation and Application of Woodbury Matrix Identity
A compromise between the full Hessian and using a constant matrix is approximating the second- order derivative terms in Equations (9) and (10) by constant identity matrices, while computing the first-order derivative terms with auto autoddiffernntiation. Specifically, given a set of l sampled trajectories T2, Tk Equations and I0) can be written and approximated in the following form:
K Σ (ueu, + +cl) = UV +el
Jo 20
(neuy + Hi)
(11)
where U [u1, u2.... uel/V/ E B V 11. U2 .... uyl/ VV E Tgnn and u, v2 E TR. p and (T)), and H; corresponds to the remaining
correspond to the first-order derivatives in Equations terms that involve second-order derivatives. We use a constant identity matrix to approximate H,, while explicitly computing the remaining parts to increase accuracy.
However, we still cannot explicitly expand UVT R_x= since the dimensionality is too large. Therefore, we apply Woodbury matrix identity
337 to invert Equation [I]
(V (UVT + +cl)-1 1 - 1 -CCl---
(12)
where V E R can be efficiently computed with much smaller & < n. When we compute the full gradient for decision-focuseee 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(kk ke)) that is much smaller than the full computation cost O(n + n²).
The full algorithm for decision-focused learning MDPs is illustrated Algorithm I
Example MDP Problems with Missing Parameters
Gridworld with unknown cliffs We consider Gridworld environment with 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 (5, 1) located at the top right corner. Each intermediate state has reward associated with it, where most of them give the agent a reward drawn from A A(O, 1) but 20% of the them are cliffs and give (-10 1) penalty to the agent. The agent has prior information
*The gradient of Bellman error can be written as Jo where the policy л is the parameters the value function approximator and / is defined as the expected Bellman error. The derivative of the final gradient can be written as Je ") = Je de by chain rule, which matches the last three terms in Equation when the Hessian is approximated by identity matrix.
7

---

Table I: OPE performances of different methods on the test MDPs averaged over 30 independent runs. Decision-focure learning methods consistently outperform two-stage approach, with some exception using identity matrix based Hessian approximation which may lead to high gradient variance.
Gridworld
Snare
Tuberculosis
Trajectories
Random
Near-optimal
Random
Near-optimal
Random
Near-optimal
TS
-12.0+1.3 1.3
0.2
0.8+
3.7.+..3
35.8=1.5
3.6
PG-Id
1.2
0.7
-0.1
3.6
38.4=1.5
40.7
Bellman-Id
-9.64
4.6400.7
0.7
3.6--0.3
39.1=1.7
1.7
PG-W
112
5.5
1.2+0.4
4.8
38.1=1.5
10.8
Bellman-V W
111.3+1.44
0.8+0.8
1.5+:04
0.3
88.6=1.6
11.1
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 set of historical trajectories from the training MDPs. The agent learns predictive model to map from the features of 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 snare at each time step. Motivated by (38), we assume ranger who can visit one site per time step and observes whether a snare is present. If a smare 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 sequential problem rather than 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 '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 discrete, given predictive model of the missing transition probability, the agent uses double deep Q-learning (DDON) 32 to solve the predicted MDPs.
Partially observable patient treatment problems with missing transition probability We con- sider a version of the Tuberculosis Adherence Problem explored in (33). 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 take medication for 30 days. Each day, healthcare worker chooses one patient to intervene on. They observe whether that patient is adhering or not, and improve the patient's 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 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 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 its continuous belief state equivalence and solve using DDQN.
Please refer to Appendix C for more details about problem setup in all three domains.
8 Experimental Results and Discussion
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 equation- based (Bellman), and two different Hessian approximations (Identity, Woodbury) defined in Sec- tion 6 Computing the full Hessian (as in Section () is computationally intractable. Across all three examples, we use 7 training MDPs, 1 validation MDP, and 2 test MDPs, each with 100 trajectories.
8

---

DP-rc-teenig
by
del aeldddaan-- manddayy
DF-IG Iddentiy =: DF-PG Woodbury tiDF-PG- Full
Training evaluation
runtime per step [rac]
Training predictivn
60
20- 2..
10.0
7.5
25
30
73
100
125
8
grid cells in gridworld
Training Epochs
Traininng Epochs
(c) Backpropagation runtime per gra- dient step per instance
(a) Predictive loss in training MDPs (b) Performance in training MDPs
Figure 2: Learning curves of Gridworld problem with near-optimal trajectories. Two-stage minimizes the predictive loss in Figure[((a) but this does not lead to good training performance Figure 2(b) Figure 3(c) shows the backpropagation runtime per gradient step per instance of three Hessian approximations, which becomes intractable when trained for multiple instances and multiple epochs.
gP-10-teo gy-eniyyen
2 3--Belman--ae ty
Woodbury
del Bicoaa-
E Dr-PG-identity C DDF-G-Woodbury DD--CG-ull
1.
Training evaluation
runtmme per stop (zac)
Training predictive
1000
100
10
15
20
25
sites conservation parks
Tranning Epochs
Training Epochs
(c) Backpropagation runtime per gra- dient step per instance
(a) Predictive loss in training MDPs (b) Performance in training MDPs
Figure 3: Learning curves of snare finding problems with random trajectories. Two-stage achieves both low predictive loss Figure 3(a) and high training OPE in Figure but the test performance
poor in Table 1 Figure 3(C) plots the backpropagation runtime per gradient step per instance.
The predictive model is trained on the training MDP trajectories for 100 epochs. Performance is eval- uated under the Off-Policy Evaluation (OPE) metric of Equation ( 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 Table T
Decision-focused learning with the Woodbury matrix identity outperforms two-stage learning Table [ 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- IId in the snare problem). In contrast, decision-focus focused methods using a more accurate low-rank approximation and Woodbury matrix identity (i.e., PG-W W and Bellman-We W), discussed in Section 6.3 dominate two-stage performance in the test MDPs across all settings.
we plot the predictive loss curve the training MDPs over different training epochs of Gridworld and snare problems. In particular, two-stage approach is trained minimize such loss, but fails to win in Table [ 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 consistent with other studies in mismatch of predictive loss and evaluation metric
Low predictive loss does not imply a winning policy In Figures 2(a)
((a)
14 15 20
Comparison between different Hessian approximations rate Hessian approximation (identity) does not always lead to poorer performance. We hypothesize that this due to the non-convex off- 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. Table T the performance of Dellman-Identin and PG- PG-Identity sometimes lead to wins over two stage and sometimes losses.
Table [ we notice that more inaccu-

---

Comparison between policy gradient and Bellman-basee decision-focused learning We ob- serve that the dellman-based recision-focussee approach consistently outperforms the policy gradient- based 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 hy- pothesize 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 [ show the backpropagation runtime of the policy- gradient based optimality condition per gradient step per training instance across different Hessian approximations and different problem sizes 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 day complete.
Analytically, let n be the dimensionality of the policy space and A < n be the number of sampled trajectories used approximate the derivatives. As shown in Section the computation cost of full Hessian On + n²) is quadratic in and strictly dominates all the others. In contrast, the costs of the identity matrix approximation (nn)) and the Woodbury approximation O(nk + A²) are both linear in n. The Woodbury method offers an option to get more accurate Hessian at low additional cost.
9
Conclusion
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 recision-focusee 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 reserve future direction.
Acknowledgement
This research was supported by MURI Grant Number W911NF-17-1-00700 and 1991111888--0200 Chen and Perrault acknowledge support from the Center for Research on Computation and Society. Doshi-Velez acknowledges support from NSF CAREER II5-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-Veler 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.
References
[1] A. Agrawal, B. Amos, S. Barratt, S. Boyd, S. Diamond, and Z. Kolter. Differentiable convex optimization layers. arXiv preprint arXiv: 1240,/233 2019.
[2] B. Amos and J. Z. Kolter. Optnet: Differentiable optimization International Conference Machine Learning, pages 136-145. PMLR, 2017.
layer in neural networks, In
[3] B. Amos, D. Rodriguez, Sacks, B. Boots, and J. 7. Kolter. Differentiable mpe for end-to-end planning and control. arXiv preprint arXiv: 1810.131000 2018.
10

---

[4] S. Bai, J. Z. Kolter, and V. Koltun. Deep equilibrium models. arXiv preprint arXiv:1909.01377, 2019.
[5] O. E. Balghiti, A. N. Elmachtoub, P. Grigas, and A. Tewari. Generalization bounds in the predict-then-optimizze framework. arXiv preprint arXiv: 1488, 2019.
[6] P. Branco, L. Torgo, and R. P. Ribeiro. A survey of predictive modeling imbalanced domains. ACM Computing Surveys (CSUR), 49(2):1-50, 2016.
[7] P. L.. Donti, B. Amos, and J. Z. Kolter. base- end-to-end model learning in stochastic optimization. NeurlPS, 2017.
[8] Y. Duan, Schulman, X. Chen, P. Bartlett, I. Sutskever, and P. Abbeel. RI 2: Fast reinforce- ment learning via slow reinforcement learning. arXiv preprint arXiv:1611.0277) 2016.
[9] A. N. Elmachtoub and P. Grigas. Smart "predict, then optimize". Management Science, 2021.
[10] C. Finn, P. Abbeel, and S. Levine. Model-agnostic meta-learning for fast adaptation of deep networks. International Conference on Machine Learning, pages 1126-1122 135. PMLR, 2017.
[11] J. Futoma, M. C. Hughes, and F. Doshi-Velez. Popcorn: Partially observed prediction con- strained reinforcement learning. arXiv preprint arXiv:2001.04032, 2020.
[12] T. Haarnoja, H. Tang, P. Abbeel, and S. Levine. Reinforcement learning with deep energy- basedd policies. In International Conference on Machine Learning, pages 1352-1361. PMLR, 2017.
[13] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International Conference on Machine Learning, pages 1861-1870. PMLR, 2018.
[14) C. Huang. S. Zhai, W. Talbott, M. B. Martin, Sun, C. Guestrin, and J. Susskind. Addressing the loss-metric mismatch with adaptive loss alignment. In International Conference on Machine Learning, pages 2891-2900. PMLR, 2019.
[15] J. M. Johnson and T. M. Khoshgoftaar. Survey deep learning with class imbalance. Journal of Big Data, 6(1):1-54, 2019.
[16] Karkus, X. Ma, D. Hsu, L. P. Kaelbling, W. S. Lee, and T. Lozano-Pérez, Differentiable algorithm networks for composable robot learning. arXiv preprint arXiv:1905.11602, 2019.
[17] R. Kidambi, A. Rajeswaran, P. Netrapalli, and T. Joachims. Morel: Model-based offiine reinforcement learning, arXiv preprint arXiv:2005.00 05951, 2020.
[18] J. A. Killian, B. Wilder, A. Sharma, V. Choudhary, B. Dilkina, and M. Tambe. Learning to prescribe interventions for tuberculosis patients using digital adherence data. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 2430 1430-2438, 2019.
[19] R. Kumar and P. Varaiya. Stochastic systems: Estimation, identification, and adaptive control. SIAM, 2015.
[20] N. Lambert, B. Amos, O. Yadan, and R. Calandra. Objective mismatch in model-based reinforcement learning. arXiv preprint arXiv:2002.04523 2020.
[21] A. Lazaric. Transfer in reinforcement learning: a framework and a survey. In Reinforcement Learning, pages 143-173. Springer, 2012.
[22] Mandi and T. Guns. Interior point solving for Ip-based prediction+ optimisation. arXiv preprint arXiv:2010. 139443, 2020.
[23] A. Mate, J. A. Killian, H. Xu, A. Perrault, and M. Tambe. Collapsing bandits and their application to public health interventions. arXiv preprint vrXiv:2007.04432, 2020.
[24] A. Y. Ng, D. Harada, and S. Russell. Policy invariance under reward transformations: Theory and application to reward shaping. In Icml, volume 99, pages 278-287 1999.

---

[25] C. Ning and F. You. Optimization under uncertainty in the era of big data and deep learning: When machine learning meets mathematical programming. Computers & Chemical Engineering. I 125:434 448, 2019.
[26] A. Perrault, B. Wilder, E. Ewing, A. Mate, B. Dilkina, and M. Tambe. End-to-end game-focused learning of adversary behavior in security games. In Proceedings of the AAAI Conference on 2020.
Artificial Intelligence, volume 34, pages
[27] A. Raffin, A. Hill, M. Ermestus, A. Gleave, A. Kanervisto, and N. Dormann. Stable baselines3. https://gt thub. combuu..coom//aa com RM/ RtttaallM//ttbeeeeee -ataaa//eeaababbbbaaelleeeee 2019.
[28] N. V. Sahinidis. Optimization under uncertainty: state-of-the-art and opportunities. Computers Chemical Engineering, 28(6-7):971-995 2004.
[29] Schaul, D. Horgan, K. Gregor, and D. Silver. Universal value function approximators. In
International conference machine learning, pages 1312-1320. PMLR, 2015.
[30] R. S. Sutton. Integrated architectures for learning, planning, and reacting based on approximating dynamic programming. In Machine learning proceedings 1990, pages 216-224. Elsevier, 1990.
[31] M. E. Taylor and P. Stone. Transfer learning for reinforcement learning domains: survey. Journal of Machine Learning Research, 10(7), 2009.
[32] H. Van Hasselt, A. Guez, and D. Silver. Deep reinforcement learning with double q-learning. In AAAI, 2016.
[33] J. X. Wang. 7. Kurth-Nelson, D. Tirumala, H. Soyer, J. 7. Leibo, R. Munos, C. Blundell, D. Ku- maran, and M. Botvinick. Learning reinforcement learn. arXiv preprint arXiv: 161 1.05763, 2016.
[34] K. Wang, B. Wilder, A. Perrault, and M. Tambe. Automatically learning compact quality- surrogates for optimization problems. arXiv preprint arXiv:2006. 2020.
[35] Y. Wang, Y. Ge, L. Li, R. Chen, and T. Xu. Offline meta-level model-based reinforcement learning approach for cold-start recommendation. arXiv preprint arXiv:2012.02476, 2020.
[36] B. Wilder, B. Dilkina, and M. Tambe. Melding the data-decisions pipeline: Decision-focused learning for combinatorial optimization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 1658-1665. 2019.
[37] M. Woodbury. Inverting modified matrices. Statistical Research Group, 1950.
[38] L. Xu, S. Gholami, S. Mc Carthy, B. Dilkina, A. Plumptre, M. Tambe, R. Singh, M. Nsubuga, J. Mabonga. M. Driciru, et al. Stay ahead of poachers: Illegal wildlife poaching prediction and patrol planning under uncertainty with field test evaluations (short version). In 2020 IEEE 36th International Conference on Data Engineering (ICDE), pages 1898-1901. IEEE, 2020.
[39] T. Yu, A. Kumar, R. Rafailov, A. Rajeswaran, S. Levine, and C. Finn. Combo: Conservative offline model-based policy optimization. arXiv preprint arXiv:2 2021.
[40] L. Zintgraf, K. Shiarlis, M. Igl, S. Schulze, Y. Gal, K. Hofmann, and S. Whiteson. Varibad: A very good method for bayes-adaptes deep rl via meta-learning In International Conference on Learning Representations, 2019.
12


Reviews:
[{"review_id": "pEBXCPp7_a", "reviewer": "Reviewer_R1HM", "summary": "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.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "[Post-rebuttal comments]\nThank you for your response. I have no further questions. \n\n[Significance]\nThe paper offers clearly significant contributions, particularly with respect to the approach to differentiating the MDP layer. However, the significance of the application (learning missing parameter values) is not entirely clear to this reviewer. Overall, the paper does not provide much motivation for the significance of this problem and the need for advanced solution methods. The paper could benefit from a discussion along these lines.\n\n[Novelty]\nThe paper uses argmin differentiation techniques applied to a novel setting and proposes clever approximations to handle the very large resulting computational problem.\n\n[Readability & Technical Quality]\nThe paper is very well written, easy to follow, makes economic use of symbolism, and introduces concepts in a reader-friendly order. It is easy to understand despite several interacting concepts in a complex model. The experimental results and caveats are described very nicely. However, the description of the case studies in the experimental portion could be more detailed. It is somewhat unclear the exact specifications of the reinforcement learning tasks; even a small image, if available, would be helpful. \n\nA few specific comments on aspects that are unclear:\n\n - Are there only a portion of missing parameters in this setting, or are all parameters missing? If only a portion is assumed to be missing, shouldn't each case study come with a description of which parameters are missing, and in what proportion? E.g., 70,80% or reward values missing. If this is the case, then the significance of the problem being solved is in more need of explanation and motivation.\n\n - The specific form of the underlying models that produce trajectories in the case studies -- are these specified? I.e., the architecture of the policy network.\n\n - Perhaps something should be said about the initialization of the learned parameters in the case studies\n\n - It seems that it is not mentioned until the appendix, that the underlying deep Q learning models in the case studies require their policy functions to be replaced with a smoothed approximation in order for the methodology to be applied. The appearance of Q-learning models in the main text results is confusing since the methods described prior assume a continuous policy. If not already, this should be clarified in the main text and the smoothed relaxations should be specified with some mathematical detail, possibly in the appendix.\n\n - For clarity, the nature of the loss functions l targeting the trajectories in Eqn (3) should be described.\n\nThe methods introduced are very well motivated on a technical level; the difficulties are thoroughly explained and novel approaches are used to handle them. Thorough proofs are provided, and the technical approach is largely convincing.\n\n [Minor Comments]\n\n - In Eqn (5), if appropriate - the equation seems like an abuse of notation, slightly confusing. Perhaps putting an argmax on the RHS or splitting the expression into 2 lines would be more clear.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations of the methods are discussed in the paper. ", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "~4"}, "extra_text": {}}, {"review_id": "FpkYRy1nIV", "reviewer": "Reviewer_DDs9", "summary": "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.\n\nThe 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.", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "To the best of my knowledge, this is the first paper that extends the predict-then-optimize framework to sequential decision-making processes. The paper is well-written and easy to follow. The material in the main paper is mathematically sound. A few comments - \n\n1. L227 - It is not clear to me why Eq(8) is \"equivalent to the idea of differentiating through the final gradient of Bellman error..\". Please elaborate.\n2. L299-300 - While the experimental analysis as presented clearly shows that the proposed PG-W and Bellman-W outperform TS, it cannot be conclusively stated that they outperform the low accuracy approximation of the Hessian through a constant identity matrix. It is true in the case of the Snare dataset, but not the other two. Do the authors have a prescriptive guidance as to when the constant identity matrix based approximation may be preferable over the low-rank approximation. Additionally, since the latter is more expensive, is there a bound on how much worse will the former be than the latter?\n3. Is it possible to extend Figures 2(c) and 3(c) to include backpropagation runtime for Bellman optimality condition?\n4. The experimental results section could be further strengthened by benchmarking the proposed approach against comparable approaches - perhaps, Karkus et al. [15] and Futoma et al. [10].", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Not applicable", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5 hours"}, "extra_text": {}}, {"review_id": "E4s8ku0ao7a", "reviewer": "Reviewer_HiZT", "summary": "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Summary\n\nThe 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.\n\n# Detailed Summary\n\n## The Good\n\nI think that the proposed method is quite novel. The paper addresses the computational issues of inferring MDP parameters when differentiating through KKT conditions. The experiment section showcases the benefits of applying the method proposed in the paper. Each component of the final algorithm is relatively well-argued (except the metric chosen).\n\n## The Bad\n\nI think that the assumptions on the MDP parameters $\\theta$ and their restrictions have not been well discussed, the actual motivation for the problem statement is a bit lacking, since this is a relatively untypical problem setting that doesn't appear in RL that much, I think that the authors should motivate it better (with where this is applicable and what is the benefit as opposed to vanilla RL). \n\nFurthermore, I think that the experiment section is weak, it mostly showcases the method in toy tasks, but as the authors note, and I understand,  learning MDPs from data is hard, so I wouldn't dwell on this that much.\n\n## Detailed Comments\n\n\nl.105: would be helpful here to give an example of what kind of features x are to be expected. Are they images, observations, sequences of observations etc. This description should argue why is this method useful.\n\n\nl.109: do you mean rather that optimal trajectories are hidden, the trajectories from $\\pi^*$?\nl.109: please state explicitly that the predictive model is parametrized by parameters $\\omega$, eg neural network.\n\n\nl.115: I think that the choice of evaluation metric warrants a discussion of why this is a good metric, because people (for example myself) are perhaps not familiar with the work of Futoma et al.\n\nl.141: by intuition, I would say that the assumption that the policy is smooth in $\\theta$ is very restrictive in the kind of problems that you can solve, eg. just turning transitions on or off by setting transition probabilities  to 0. Maybe you can comment a bit more on the restrictions imposed by this assumption.\n\nDef. 1: I have a slight issue with notation here, oftentimes in RL papers $J$ is used to denote the policy return as far as I know, so perhaps another notation here would be more suitable to ddenote the Bellman error under policy.\n\nl.161: you are saying that the reward depends on parameter $\\theta$, this is already restricting the kind of problems that you can solve, since you are talking about very specific MDP parameters. Otherwise, if the parameters affect the transition probabilities only, or the discount factor, the one-step reward might not be affected, but the solution of the MDP is.\n\nSec. 7: the set of experiments is rather small, and they are mostly toy experiments.\n\nl.258:  I am not familiar with this snare-finding  MDP setting, how is this not a bandit problem where each arm is the site to be visited and the reward is sampled with some probability? Here the reward doesn't depend on the set of features?\n\nl.277:  in this example, the reward doesn't depend on the set of features? Please state explicitly what is the set of features here.\n\nFig. 2: the color for the DF-PG-Identity setting needs to be changed, currently it's barely visible in the performance plots.\n\n## Related Work\n\nAdditional related work that concerns itself with inferring parameters of MDP in constrained setting that you should consider adding to related work:\n\nhttp://proceedings.mlr.press/v139/vlastelica21a.html\n\n\n\n## Conclusion\n\nAll in all, I think it's a good paper that could use a few tweaks. \n\n\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"ethical_concerns": "There are no ethical concerns.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "CNvWmwVFGtL", "reviewer": "Reviewer_gBYr", "summary": "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. ", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Post-rebuttal comment: I thank the authors for answering my questions and providing additional ablations. I will keep my score.\n\nOriginality: The problem setting is not entirely original but is interesting as it has connections to different areas such as meta-RL. I enjoyed reading the end-to-end approach which tackles a core problem of differentiating a policy’s offline off-policy evaluation with respect to its parameters. Differentiating through an optimality condition is a popular method in related works on embedding an optimization layer in a neural network model, for example [1]. The extension to MDPs is definitely interesting and I think the authors did a good job addressing the technical challenges from such an extension.\n\nQuality: The method is technically sound. I have two questions about the details in algorithm 1.\n1. The forward step (line 4) requires obtaining an optimal policy on the predicted MDP problem. Since the KKT condition also applies to the optimal policy, what is the impact if running a model-free RL algorithm only produces a sub-optimal policy? \n\n2. The gradient step (line 6) mentioned using the predictive loss as a regularizer. Did you do an ablation study on how important this regularization is? What are some high-level explanations on why this regularization is necessary?\n\nClarity: This paper is well-written and easy to follow. Good job!\n\nSignificance: I think the idea presented in this work is of interest to the broad NeurIPS community. One weakness is that the three empirical settings are synthetic which may hide some potential issues of the algorithm. For example, related to the first question in the quality section, the robustness of the algorithm should be tested in more complex settings where obtaining an optimal policy is difficult.\n\n[1]: End to end learning and optimization on graphs, Bryan Wilder, Eric Ewing, Bistra Dilkina, Milind Tambe, NeurIPS 2019\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors raised some interesting hypotheses related to the design choices in section 8. It would be great if they can also provide concrete analyses on those hypotheses. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors propose a "predict then optimize" approach for sequential decision problems. In other words, they consider a family of problems where features of the problem are given that are indicative of the MDP definition (reward and/or transition function). The proposed solution strategy is then a pipeline, where first missing MDP parameters are predicted and subsequently the MDP is solved using policy-based or value-based methods. Importantly, unlike e.g. many model-based RL approaches, this is not done in a 2-stage approach (where the dynamics model or reward function are trained on a separate supervised prediction loss) but end-to-end (the prediction of MDP parameters is learned to optimize the final RL loss). The main challenge in applying this framework is the large state-action spaces and large policy spaces concerned. So, the authors focus their attention on approximating key steps of the algorithm to obtain practicable algorithms in this setting. 

The reviewers are unanimously positive about this paper. 
- All reviewers consider the paper novel (though one reviewer considers it not entirely original)
- The reviewers consider the paper mathematically sound and the design of algorithmic components to be well justified
- The reviewers considered the paper to be well-written
- While the experiments to showcase clear benefits, they are mostly limited to toy tasks and the results could benefit from additional comparison
- Three reviewers did mention that the problem statement is not well-motivated. What in what 'real' settings can we apply this method?

Considering the unanimous positivity and agreement on novelty, soundness, and clarity, I recommend this paper to be accepted. 
Decision: Accept (Spotlight)