Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 8: Top 50% of accepted NeurIPS papers, clear accept
review_content: 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.

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.
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 - 

1. 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.
2. 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?
3. Is it possible to extend Figures 2(c) and 3(c) to include backpropagation runtime for Bellman optimality condition?
4. 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].
limitations_and_societal_impact: Not applicable
review_title: null

final_score: rating: 7: Good paper, accept
review_content: 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.
main_review: # 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.

# Detailed Summary

## The Good

I 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).

## The Bad

I 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). 

Furthermore, 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.

## Detailed Comments


l.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.


l.109: do you mean rather that optimal trajectories are hidden, the trajectories from $\pi^*$?
l.109: please state explicitly that the predictive model is parametrized by parameters $\omega$, eg neural network.


l.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.

l.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.

Def. 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.

l.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.

Sec. 7: the set of experiments is rather small, and they are mostly toy experiments.

l.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?

l.277:  in this example, the reward doesn't depend on the set of features? Please state explicitly what is the set of features here.

Fig. 2: the color for the DF-PG-Identity setting needs to be changed, currently it's barely visible in the performance plots.

## Related Work

Additional related work that concerns itself with inferring parameters of MDP in constrained setting that you should consider adding to related work:

http://proceedings.mlr.press/v139/vlastelica21a.html



## Conclusion

All in all, I think it's a good paper that could use a few tweaks. 





limitations_and_societal_impact: yes
review_title: null

final_score: rating: 7: Good paper, accept
review_content: 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.

main_review: [Post-rebuttal comments]
Thank you for your response. I have no further questions. 

[Significance]
The 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.

[Novelty]
The paper uses argmin differentiation techniques applied to a novel setting and proposes clever approximations to handle the very large resulting computational problem.

[Readability & Technical Quality]
The 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. 

A few specific comments on aspects that are unclear:

 - 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.

 - 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.

 - Perhaps something should be said about the initialization of the learned parameters in the case studies

 - 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.

 - For clarity, the nature of the loss functions l targeting the trajectories in Eqn (3) should be described.

The 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.

 [Minor Comments]

 - 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.
limitations_and_societal_impact: Limitations of the methods are discussed in the paper.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: 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. 
main_review: Post-rebuttal comment: I thank the authors for answering my questions and providing additional ablations. I will keep my score.

Originality: 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.

Quality: The method is technically sound. I have two questions about the details in algorithm 1.
1. 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? 

2. 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?

Clarity: This paper is well-written and easy to follow. Good job!

Significance: 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.

[1]: End to end learning and optimization on graphs, Bryan Wilder, Eric Ewing, Bistra Dilkina, Milind Tambe, NeurIPS 2019

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_title: null

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)