Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 6: Marginally above acceptance threshold
review_content: Summary:
The paper proposed a framework that combines TransE-style world model learning and value iteration networks. Section 3.2 gives a nice and clear overview of the algorithm: generating trees by calling the learned world model, run value iteration networks on the tree, and use the value iteration algorithms to compute the value and the policy at each state.
Since the output of the algorithm is a policy, the model was trained using PPO.

Comments:
The overall presentation of the algorithm is quite clear, and the proposed model is a nice combination of value iteration networks with learned world models. My comments will thus focus on the claims made by the authors and the completeness of the experiment section.

- Claims
1. The authors have made very strong arguments about the generality of the proposed method: "As a result, we are able to seamlessly run XLVINs with minimal configuration changes on environments from MDPs with known structure (such as grid-worlds), through pixel-based ones (such as Atari), all the way towards fully continuous-state environments, consistently outperforming or matching baseline models which lack XLVIN’s inductive biases." The authors have made several arguments about non-discrete and non-deterministic environments in the paper, however, they haven't shown any results on challenging continuous and stochastic environments. For example, how will the model perform in environments with high-dimensional action spaces: such as torques for robots? How will the model perform in partially-observable environments? If the proposed algorithm does not work for continuous action spaces (Page 4, footnote), the authors should make their claims weaker: finite and discrete action spaces + finite planning horizon makes the number of states visited also finite.

- Experiments
1. Overall the experiments are relatively weak. Adding more results on high-dimensional continuous environments (such as humanoids) or more difficult atari games could possibly help.
2. The contribution of the proposed algorithm comes from two parts: the world-model part and the VIN+Policy Network part. However, there are only few fine-grained ablation studies. For example, there have been a great number of models that learn the forward dynamics of the environment. How will other forward dynamics models perform in the same framework (e.g., Ha & Schmidhuber 2018a)? Does the TransE-style low-dimensional prior help?
3. Related to my previous comment: is there any way that the authors can make a quantitative evaluation of the learned world models?
4. The authors should make comparisons with other models that do explicit/implicit planning based on the same learned world models, e.g., Value Prediction Networks (Oh et al., 2017) and TreeQN (Farquhar et al., 2018).
5. There naturally exists a chicken-and-egg problem between learning world models and learning task-specific policies: learning good world models in complex environments requires good policies to collect useful data, while a better world model can improve task-specific policies. I think the authors should be more clear about how they are collecting data for training their TransE models and discuss the trade-off between exploration and exploitation in their model.
review_title: Interesting paper. Claims should be refined. More ablations should be included.

final_score: rating: 6: Marginally above acceptance threshold
review_content: **Summary**

The authors propose a generalization of Value Iteration Networks to unknown, potentially continuous state spaces. They describe a framework for leveraging a learned graph embedding model (TransE) in combination with a deep RL model and an execution model based on graphical message passing to perform a VI-like operation. The authors show improved performance compared to baselines on a grid-world task with a known MDP, as well as several simple continuous control environments and the Atari game Freeway.

**General**

Overall the approach is interesting, but the current results are somewhat limited and the model, in particular the executor, seems to depend critically on the availability of a suitable graph for pre-training. The authors leverage a "CartPole-inspired" graph for the majority of the tasks, and the tasks seem to have been chosen to be consistent with the statistics of this synthetic graph (from the text, Freeway's "'up-and-down' structure aligns it somewhat to environments like CartPole"). While the authors illustrate a limited ability to generalize from graph/task misalignment on the maze task (8x8 to 16x16), it's still not clear that the model will generalize to more complex tasks where it's not straightforward to pre-train on a representative graph. For instance, all of the non-maze tasks appear to use simple ternary actions (e.g. left, right and stay for CartPole; up, down and stay for Freeway). As well, working from a task-relevant pre-training graph appears to give the agent foreknowledge of the task itself, which makes comparisons to baselines like PPO unfair.

The main question here for me is whether the paper reaches the threshold on originality and significance. The paper builds on the VIN/GVIN work and particularly depends on Deac et al. 2020 for illustrating VI with GNNs. It provides the first extension of these works to more complex tasks with unknown MDPs. However, the limited results and dependence on the availability of a suitable graph makes the significance of these extensions unclear. Nonetheless, it could lay the framework for future work that pushes the VI framework into new domains.

Section 4.3 provides a nice clarification on what the model is doing under the hood.

**Specific**

- I would like to see more exploration of the executor depth, i.e. do we see gradual improvement as the depth is increased?

- Are the Freeway with respect to the first 1 million transitions (text) or 100000 (figure)?

- \alpha is used for both the hinge loss in Eq. 2 and 3.2 3.(c) which is somewhat confusing at first glance. 

**Pros:**
- Extends VI framework to more complex, unknown tasks.
- Well-written and clear, excluding the specific issues noted above.
- Nice qualitative results

**Cons:**
- Results are provided on a limited number of simple tasks, primarily using PPO as a baseline.
- Tasks seem to have been chosen such that they have a representative graph available.

review_title: Extends VI to new domains, but seems to depend critically on an available representative graph

final_score: rating: 6: Marginally above acceptance threshold
review_content: This paper proposed a novel policy prediction model that combines self-supervised contrastive learning, graph representation learning and neural algorithm execution to generalize the Value Iteration Networks to MDPs. The method described in the paper is a combination of existing works in the literature but seems to work well in practice. The experiments evaluate multiple aspects of the proposed model (E.g. number of executor layers, etc.) and show significant performance improvement over the existing approaches.

The latent representation used for policy prediction implicitly incorporates two-step aggregation 1) from initial representation extraction through encoding network and 2) from message passing, which seems helpful for the generalization.

As I am not actively tracking RL literature, I am not sure if there is a similar approach has already been proposed or not. My comments are based on the assumption that no existing work using GNN to do further representation learning on top of the individual encoded information. It would be better if the authors clearly state this novelty at the end of related works.

Pros
===
1. The paper is very well written and provided sufficient background knowledge to let the reader follow the description.
2. While it appears to be a simple combination of existing techniques, the proposed model shows the benefit of obtaining better latent state representations for the policy prediction task. 

Cons
===
1. The novelty of the proposed model is a bit weak in terms of a lack of specialization for this particular task. 
2. Figure 1 in the paper is not quite meaningful. A better demonstrative figure would help improve this paper.
review_title: Well written paper, approach proposed is simple and easy to understand

final_score: rating: 7: Good paper, accept
review_content: The paper tackles an open problem of the value-iteration-network-paradigm. The proposed method (XLVIN) has a conceptual edge over traditional value iteration networks in that it can be applied to continuous problems and problems where the state space is either too big or not fully known in advance. The experiments mostly succeed in making the case that XLVINs:

* are a drop-in replacement for traditional VINs in discrete settings
* are able to improve on models without explicit planning models in general

One issue with the first point is that the experiments on the $16\times16$ environments are in favor of VINs against XLVINs. The explanation provided in the paper makes sense but on the whole I get the impression that more experiments are necessary to asses whether XLVINs can truly replace VINs in discrete environments. 

There is also quite a bit of algorithmic complexity that results from using the graph neural network paradigm. This is justified by the fact that XLVINs are also applicable to a much wider range of problems. However, as a reader, I'd appreciate it a lot if the paper would spend time on discussing some of that complexity in greater detail (I go into specifics under "Cons"). Regardless, I recommend accepting the paper as the ideas presented here are relevant to the VIN-community.

### Pros

* The method is applicable to continuous problems and large or not fully-known state-spaces, where VINs are not.
* The experiments on continuous problems show an edge over methods without a planning component.
* The method has an edge in the low-data regime.

### Cons

* Some important questions are not addressed:
  * How does the method scale with the planning horizon $K$?
  * How feasible is it to produce synthetic MDPs on real-world problems?
  * How feasible is it to use generic graphs instead?

  It's understandable that some content needs to be left out due to space constraints, but at the same time, these questions would be at the heart of future work on this topic and the paper would benefit from addressing them.

* The experiments on the $16\times16$ mazes no longer support the claim that XLVINs can replace VINs gracefully. An explanation is provided for this, though it's still unclear how the two models compare in discrete environments other than the $8\times 8$ mazes. I feel like more experiments are needed here, perhaps in scenarios that don't give an unfair advantage to VINs. The mean- or max-pooled VINs are not convincing for that because they fail to solve even the simplest tasks.
review_title: This paper addresses the limitations of value iteration networks with regards to continuous settings and partially-known state spaces. This is achieved by approximating the full MDP with a subset of states and using graph neural networks to mimic dynamic programming. Experimental results show improvement over a model with no explicit planning module (PPO).

Meta-review: The work extends the line of work based on value iteration networks. The main goal is to extend VINS to continuous and partially observable state spaces. The approach combines self-supervised contrastive learning and graph representation learning with VINs to address these issues. Reviewers liked the premise of the paper and had several follow-up clarifications.  The authors provided the rebuttal and addressed some of the concerns. However, upon post rebuttal discussion, the reviewers decided to maintain their score. While everyone recommended weak acceptance, no one championed the paper. This was primarily due to the concerns in the empirical analysis. It is not clear that XLVINs are clearly outperforming VINs and Graph-VINs in all settings. All baselines are not present in all the environments, so it is difficult to draw a consistent conclusion. The paper is in a good state but not fully polished to infer clear conclusions about the effectiveness of the proposed approach. Please refer to the feedback below for more details. We believe strengthening the experimental results section will turn this paper into a very strong submission.
Decision: Reject