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

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