text: The authors propose a new network architecture for RL that contains some relevant inductive biases about planning. This fits into the recent line of work on implicit planning where forms of models are learned to be useful for a prediction/planning task. The proposed architecture performs something analogous to a full-width tree search using an abstract model (learned end-to-end). This is done by expanding all possible transitions to a fixed depth before performing a max backup on all expanded nodes. The final backup value is the Q-value prediction for a given state, or can represent a policy through a softmax. I thought the paper was clear and well-motivated. The architecture (and various associated tricks like state vector normalization) are well-described for reproducibility. Experimental results seem promising but I wasn’t fully convinced of its conclusions. In both domains, TreeQN and AtreeC are compared to a DQN architecture, but it wasn’t clear to me that this is the right baseline. Indeed TreeQN and AtreeC share the same conv stack in the encoder (I think?), but also have the extra capacity of the tree on top. Can the performance gain we see in the Push task as a function of tree depth be explained by the added network capacity? Same comment in Atari, but there it’s not really obvious that the proposed architecture is helping. Baselines could include unsharing the weights in the tree, removing the max backup, having a regular MLP with similar capacity, etc. Page 5, the auxiliary loss on reward prediction seems appropriate, but it’s not clear from the text and experiments whether it actually was necessary. Is it that makes interpretability of the model easier (like we see in Fig 5c)? Or does it actually lead to better performance? Despite some shortcomings in the result section, I believe this is good work and worth communicating as is.