Published as a conference paper at ICLR 2018 TREEQN AND ATREEC: DIFFERENTIABLE TREE-STRUCTURED MODELS FOR DEEP REINFORCEMENT LEARNING
Combining deep model-free reinforcement learning with on-line planning is a promising approach to building on the successes of deep RL. On-line planning with look-ahead trees has proven successful in environments where transition models are known a priori. However, in complex environments where transition models need to be learned from data, the deficiencies of learned models have limited their utility for planning. To address these challenges, we propose TreeQN, a differentiable, recursive, tree-structured model that serves as a drop-in replacement for any value function network in deep RL with discrete actions. TreeQN dynamically constructs a tree by recursively applying a transition model in a learned abstract state space and then aggregating predicted rewards and state-values using a tree backup to estimate Q-values. We also propose ATreeC, an actor-critic variant that augments TreeQN with a softmax layer to form a stochastic policy network. Both approaches are trained end-to-end, such that the learned model is optimised for its actual use in the tree. We show that TreeQN and ATreeC outperform n-step DQN and A2C on a box-pushing task, as well as n-step DQN and value prediction networks (Oh et al., 2017) on multiple Atari games. Furthermore, we present ablation studies that demonstrate the effect of different auxiliary losses on learning transition models.
# Update after the rebuttal
Thank you for the rebuttal.
The authors claim that the source of objective mismatch comes from n-step Q-learning, and their method is well-justified in 1-step Q-learning. However, there is still a mismatch even with 1-step Q-learning because the bootstrapped target is also computed from the TreeQN. More specifically, there can be a mismatch between the optimal action sequences computed from TreeQN at time t and t+1 if the depth of TreeQN is equal or greater than 2. Thus, the author's response is still not convincing to me.
I like the overall idea of using a tree-structured neural network which internally performs planning as an abstraction of Q-function, which makes implementation simpler compared to VPN. However, the particular method (TreeQN) proposed in this paper introduces a mismatch in the model learning as mentioned above. One could argue that TreeQN is learning an "abstract" planning rather than "grounded" planning. However, the fact that reward prediction loss is used to train TreeQN significantly weakens this claim, and there is no such an evidence in the paper. 
In conclusion, I think the research direction is worth pursuing, but the proposed modification from VPN is not well-justified.
# Summary
This paper proposes TreeQN and ATreeC which perform look-ahead planning using neural networks. TreeQN simulates the future by predicting rewards/values of the future states and performs tree backup to construct Q-values. ATreeC is an actor-critic architecture that uses a softmax over TreeQN. The architecture is trained through n-step Q-learning with reward prediction loss. The proposed methods outperform DQN baseline on 2D Box Pushing domain and outperforms VPN on Atari games.
[Pros]
- The paper is easy to follow.
- The application to actor-critic setting (ATreeC) is novel, though the underlying idea was proposed by [O'Donoghue et al., Schulman et al.].
[Cons]
- The proposed method has a technical issue.
- The proposed idea (TreeQN) and underlying motivation are almost same as those of VPN [Oh et al.], but there is no in-depth discussion that shows why TreeQN is potentially better than VPN. 
- Comparison to VPN on Atari is not much convincing. 
# Novelty and Significance
- The underlying motivation (planning without predicting observations), the architecture (transition/reward/value functions applied to the latent state space), and the algorithm (n-step Q-learning with reward prediction loss) are same as those of VPN. But, the paper does not provide in-depth discussion on this. The following is the differences that I found from this paper, so it would be important to discuss why such differences are important.
1) The paper emphasizes the "fully-differentiable tree planning" aspect in contrast to VPN that back-propagates only through "non-branching" trajectories during training. However, differentiating TreeQN also amounts to back-propagating through a "single" trajectory in the tree that gives the maximum Q-value. Thus, the only difference between TreeQN and VPN is that TreeQN follows the best (estimated) action sequence, while VPN follows the chosen action sequence in retrospect during back-propagation. Can you justify why following the best estimated action sequence is better than following the chosen action sequence during back-propagation (see Technical Soundness section for discussion)?
2) TreeQN only sets targets for the final Q-value after tree backup, whereas VPN sets targets for all intermediate value predictions in the tree. Why is TreeQN's approach better than VPN's approach? 
- The application to actor-critic setting (ATreeC) is novel, though the underlying idea of combining Q-learning with policy gradient was proposed by [O'Donoghue et al.] and [Schulman et al.].
# Technical Soundness
- The proposed idea of setting targets for the final Q-value after tree backup can potentially make the temporal credit assignment difficult, because the best estimated actions during tree planning does not necessarily match with the chosen actions. Suppose that TreeQN estimated "up-right-right" as the best future action sequence the during 3-step tree planning, while the agent actually ended up with choosing "up-left-left" (this is possible because the agent re-plans at every step and follows epsilon-greedy policy). Following n-step Q-learning procedure, we end up with setting target Q-value based on on-policy action sequence "up-left-left", while back-propagating through "up-right-right" action sequence in the TreeQN's plan. This causes a wrong temporal credit assignment, because TreeQN can potentially increase/decrease value estimates in the wrong direction due to the mismatch between the planned actions and chosen actions. So, it is unclear why the proposed algorithm is technically correct or better than VPN's approach (i.e., back-propagating through the chosen actions in the search tree).
 
# Quality
- Comparison to VPN on Atari is not convincing because TreeQN-1 is actually (almost) equivalent to VPN-1, but the results show that TreeQN-1 performs much better than VPN on many games. Since the authors took the numbers from [Oh et al.] rather than replicating VPN, it is possible that the gap comes from implementation details (e.g., hyperparameter). 
# Clarity
- The paper is overall easy to follow and the description of the proposed method is clear.