[Paper Review] TreeQN and ATreeC: Differentiable Tree-Structured Models for Deep Reinforcement Learning

Gregory Farquhar, Tim Rocktäschel+2 more|arXiv (Cornell University)|Oct 31, 2017

Reinforcement Learning in Robotics28 citations

TL;DR

This paper introduces TreeQN and ATreeC, differentiable, recursive tree-structured models that integrate end-to-end trained transition models into deep reinforcement learning for improved on-line planning. By formulating tree backups as differentiable operations, the models learn transition dynamics specifically for value estimation, outperforming n-step DQN, A2C, and value prediction networks on box-pushing and Atari games, with deeper trees often yielding better performance.

ABSTRACT

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.

Motivation & Objective

• To address the challenge of learning accurate transition models for on-line planning in complex, high-dimensional environments where model errors limit planning utility.
• To improve sample efficiency and planning accuracy in model-free deep RL by embedding a differentiable tree-structured value estimation process directly into the Q-function or policy network.
• To train the transition model end-to-end with the policy and value function, ensuring it is optimized for actual planning performance rather than observation reconstruction.
• To explore whether auxiliary losses can ground the transition model more strongly in the environment while preserving performance and enabling interpretable internal planning.

Proposed method

• TreeQN constructs a differentiable, recursive tree by applying a shared, learned transition model in an abstract state space, with Q-values computed via a tree backup that aggregates rewards and next-state values.
• The tree structure is differentiable through backpropagation, enabling end-to-end training of the transition model, reward head, and value head jointly.
• ATreeC extends TreeQN by adding a softmax layer on top of the tree output to form a stochastic policy network, enabling actor-critic training.
• The model uses a differentiable tree backup operation that computes Q-values as a recursive sum of immediate rewards and discounted next-state values, with shared parameters across tree nodes.
• Auxiliary losses are introduced to improve transition model fidelity, including reconstruction loss in the observation space and prediction of future states in the abstract space.
• The entire architecture is trained end-to-end using policy gradient or Q-learning objectives, with the transition model optimized for planning accuracy rather than generative reconstruction.

Experimental results

Research questions

• RQ1 Can a differentiable, recursive tree-structured model improve on-line planning in deep reinforcement learning when transition models are trained end-to-end?
• RQ2 Does training the transition model for planning performance rather than observation reconstruction lead to better sample efficiency and final performance?
• RQ3 Can deeper trees in TreeQN and ATreeC yield better performance than shallower trees or standard DQN architectures?
• RQ4 How do auxiliary losses for transition model supervision affect planning accuracy and model interpretability?
• RQ5 Can the integration of differentiable tree search into value functions or policies outperform existing model-based and model-free baselines on complex control tasks and Atari games?

Key findings

• TreeQN outperforms n-step DQN and value prediction networks (VPNs) on 18 out of 26 Atari games, with significant gains on games like Ms. Pac-Man and Q*bert.
• ATreeC matches or exceeds A2C performance across all Atari environments, with stronger performance on Q*bert and Krull, though it suffers from premature policy collapse on Seaquest.
• TreeQN-2 achieves a mean human-normalized score of 9302 on Atari, surpassing the best reported score of 7860 for n-step DQN and 8241 for A2C.
• In the box-pushing domain, TreeQN and ATreeC outperform n-step DQN and A2C, with TreeQN-2 achieving a final score of 15688 compared to 14468 for n-step DQN.
• Deeper trees (e.g., TreeQN-2) often outperform shallower ones, indicating that recursive planning improves value estimation.
• Ablation studies show that grounding the reward function improves performance, but learning strongly grounded transition models without performance degradation remains an open challenge.