Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This submission proposes eXecuted Latent Value Iteration Networks (XLVINs) to relax the assumptions of Value Iteration Networks. The method infers the dynamics of the MDP using known representation learning techniques and uses the inferred dynamics to construct a graph representation which is used for planning. The method allows continuous state spaces and unlike VINs, not restricted to grid-world.
main_review: 
Note:
I reviewed this paper for ICML 2021. The method is identical as far as I can tell. The writing quality has been improved at some places. In terms of experiments, the authors have removed known MDP maze experiments and have added one continuous space environment, LunarLander-v2. 


Strengths:
- The paper tackles an important problem of planning under uncertain maps which is of interest to the community.
- The method incorporates several prior techniques to tackle planning from pixel space which makes it complicated, but it is described in sufficient detail at most places.


Weakness:
- The paper aims to improve over VINs by tackling continuous space and pixel space environments. However, while doing so it seems the proposed method isn't as effective as VINs at discrete state planning under known MDPS. Instead of adding larger maze size experiments, the authors chose to remove all maze experiments. If the method isn't effective at known MDP discrete space planning, this should be acknowledged in the paper.
- For continuous state space tasks, the authors only conduct experiments in only 4 simple gym games and for pixel space observations, only 3 Atari games. The motivation of using these benchmarks for testing planning algorithms is unclear. It would make more sense to test the proposed method for planning in high dimensional continuous state space 3D manipulation or navigation tasks. In any case, the experimental evaluation needs to be more comprehensive in my opinion to establish the effectiveness of XLVINs.
- The paper misses some relevant related work, for example Universal planning networks (Srinivas et al. ICML 2018), Differentiable MPC for End-to-End Planning and Control (Amos et al., Neurips 2018). It would be useful to discuss how the proposed method differs from these methods.
- The baselines used are weak in my opinion. The proposed model learns a transition model between latent states. It seems more similar to model-based RL methods than to VINs. I would have liked to see empirical comparisons to Universal planning networks (Srinivas et al. ICML 2018) or PlaNet (Hafner et al. ICML 2019).


Update after author response:
The author response has addressed my major concerns. I am increasing my rating accordingly. I still believe larger mazes and robotics applications are better environments to test planning methods rather than a few Atari games.

limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: This work extends the value iteration network family of reinforcement learning agents to a much less constrained form to avoid the "algorithmic bottleneck". The proposed model XLVIN creates a plannable representation of the state space by 1) first learning an embedding space for the state transitions via a TransE (or contrastive) approach and 2) then aggregating several layers of look-ahead state embeddings via a pretrained GNN. The policy and value functions are parameterized on top of the plannable representation. Empirically XLVINs provide a significant boost on the stochastic control tasks and Atari games in the low-data regime.
main_review: Pros:
- This paper provides a solution to the issues presented in previous work along this line:
  - Not all high-dimensional state spaces can be mapped properly to a lower-dimensional space for VIN-like updates.
  - The implicit planning in previous work is prone to error, especially from the incorrect estimation of the reward and state values.
- The improvements compared with existing baselines in low-data regimes are significant.
- The paper is well written, and the experiments are solid with several extra studies beyond standard ones.

Cons:
- Since the complexity of the proposed method scales with O(|A|^K), the model can only look ahead for a very small number of steps.
  - Why would a small K help build a better state representation? For atari games, it seems that looking 1-2 steps ahead into the future won't provide a decent gain.
- The discussion of how the quality of the learned TransE state embeddings will affect the final performance is not included in the evaluation section. Such sensitivity analysis will be helpful in seeing how XLVINs can generalize to more complicated state spaces (e.g., molecule space) where TransE style embeddings are hard to learn.
limitations_and_societal_impact: Yes.
review_title: null

final_score: rating: 8: Top 50% of accepted NeurIPS papers, clear accept
review_content: summary: This paper proposes a new implicit planner module, XLVINs, for end-to-end policy learning, which mimics the value iteration planning procedure with a pre-trained graph convolution network and overcomes the existing limitations of previous VI-style planning modules. This paper carefully investigates the existing literature with a thorough analysis of the algorithmic limitations of previous VI modules. The experiments also provide a complete understanding of the proposed method. 
main_review: The paper is extremely well written with lots of insightful discussions and explanations. It is really a great pleasure to read this paper and everything is so easy to follow under its current form. The discussion of algorithmic bottleneck is insightful. 

From a technical perspective, the paper proposes to use a TransE style latent model and adopt a GCN-based VI module. These two techniques directly address the existing limitations of previous works, which look clearly novel to me. I like the idea of pre-training a VI-module using random graphs and was surprised to see it is really working ---- particularly the experiment where the ground truth information is recovered. 

In general, I like this paper and feel happy to see it gets accepted. 
limitations_and_societal_impact: There are some minor comments.

Regarding the TransE, it seems a bit unjustified since it seems that even if we simply train a neural network to predict the next state (e.g., a transition model, or a world model), things would still work without any changes? Is there any specific reason for using TransE other than its simplicity?

Finally, I definitely feel the experiment scenarios are sufficient in the current paper. If possible, the paper can be even stronger if results on more Atari games can be presented.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: In the context of implicit planning algorithms, the paper proposes to replace exact value iteration (VI) computations with a learned graph neural network that is trained to perform VI computations in an abstract latent space.

The paper shows competitive performance on several benchmark datasets using this approach, against a PPO baseline, and an implicit planning algorithm using exact VI.

main_review: The paper argues that performing exact VI computations in an implicit planning setup suffers from an "algorithmic bottleneck". They argue that this algorithmic bottleneck occurs when the exact VI computations are performed on poor initial estimates of the value function. 

To remedy this "algorithmic bottleneck" the paper proposes to instead perform (approximate) VI computations in a latent space using a pre-trained Graph Neural Network (GNN). The GNN consists of 3 parts, an encoder, an update function, and a decoder, and is pre-trained to perform approximate VI on a dataset of randomly sampled Markov Decision Processes, after which only the update function is kept and frozen.

Given this GNN, the "executor", that can perform approximate VI in a latent space, the paper uses a fairly standard implicit planning setup. A learned encoder maps the environment state into a latent space, and a learned transition function maps latent states and actions to the next (expected) latent state. Given this tree of latent states, the executor performs its approximate VI computations and finally, a learned policy and value network outputs the policy and value function.

The paper is exceptionally well written. Especially section 2 gives the reader an excellent overview of the relevant material needed to understand the contributions of the paper. The paper shows compelling results demonstrating that it's better than a solid baseline (PPO) and more importantly, ATreeC, which is the method they directly built on by replacing the exact VI computations.

I have a few comments that I think could make the paper even stronger. 

1) The paper argues that the algorithmic bottleneck is a problem, but doesn't really offer any evidence of this, except in the full experiments where they show their proposed method is better than the one using the exact VI computations. Further, the paper doesn't give any intuition as to why performing the VI computations in a high dimensional latent space might avoid the algorithmic bottleneck. I would greatly appreciate a small-scale experiment that directly shows the effect of the "algorithmic bottleneck" and how doing (approximate) VI in a higher dim latent space avoids it. The authors could use the simple MDPs they use to train the executor, and inject noise in the value functions and show how their learned executor is more robust to noise than the direct VI equations. This along with a few sentences on the intuition why doing approximate VI in a higher dim space is more robust would greatly improve the paper in my opinion. 

2) It was not entirely clear to me how the whole system was trained. In particular, I find it hard to understand how the encoding function and transition functions are trained. The loss in eq. 4 seems to indicate that they're trained entirely offline on (S,A,S') triplets, but that can't be true since then the learned embeddings wouldn't align with the one used by the executor. It would be very helpful if the authors included the joint loss that is optimized with PPO, and/or an algorithm describing the training procedure (training the executor first, then freezing, etc.).

A few minor things:
 - I'd have liked a small discussion of how well (or not) the approach handles multi-modal MDPs. Computing an expected state over a multi-modal outcome seems like a poor approach.
 - Line 156: "... and do not project our state to a low-dimensional space": well, you do, with your encoder function. I know what you mean, but this sentence is kind of confusing.
limitations_and_societal_impact: Yes
review_title: null

Meta-review: The paper proposes proposes eXecuted Latent Value Iteration Networks (XLVINs) to relax the assumptions of Value Iteration Networks.  While most neural architectures in deep RL are not well motivated, this paper builds on an important line of work that treats value representations as implicit planners.  This is very nice work that generalizes previous work.  Well done!
Decision: Accept (Spotlight)