Learning MDPs from Features: Predict-Then-Optimize for Sequential Decision Making by Reinforcement Learning

1 Introduction . 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 . 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 ( 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 ) ) ) ] . 3 Two-stage Learning . 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 .

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.