T REE QN AND AT REE C: D IFFERENTIABLE T REE P LANNING FOR D EEP R EINFORCEMENT L
EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
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 deﬁciencies 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 planner.
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, with deeper trees often outperforming shallower ones.
We also present a qualitative analysis that sheds light on the trees learned by TreeQN.
I NTRODUCTION
A promising approach to improving model-free deep reinforcement learning (RL) is to combine it with on-line planning.
The model-free value function can be viewed as a rough global estimate which is then locally reﬁned on the ﬂy for the current state by the on-line planner.
Crucially, this does not require new samples but only additional computation, which is often available.
One strategy for on-line planning is to use look-ahead tree search (Knuth & Moore, 1975;
Browne et al., 2012).
Traditionally, such methods have been limited to domains where perfect environment simulators are available, such as board or card games (Coulom, 2006; Sturtevant, 2008).
However, in general, models for complex environments with high-dimensional observation spaces and complex dynamics must be learned from agent experience.
Unfortunately, to date, it has proven difﬁcult to learn models for such domains with sufﬁcient ﬁdelity to realise the beneﬁts of look-ahead planning (Oh et al., 2015; Talvitie, 2017).
A simple approach to learning environment models is to maximise a similarity metric between model predictions and ground truth in the observation space.
This approach has been applied with some success in cases where model ﬁdelity is less important, e.g. , for improving exploration (Chiappa et al., 2017;
Oh et al., 2015).
However, this objective causes signiﬁcant model capacity to be devoted to predicting irrelevant aspects of the environment dynamics, such as noisy backgrounds, at the expense of value-critical features that may occupy only a small part of the observation space (Pathak et al., 2017).
Consequently, current state-of-the-art models still accumulate errors too rapidly to be used for look-ahead planning in complex environments.
Another strategy is to train a model such that, when it is used to predict a value function, the error in those predictions is minimised.
Doing so can encourage the model to focus on relevant features of the observations.
An example is the predictron (Silver et al., 2017b), where the model is used to aid policy evaluation without addressing control.
Value prediction networks (VPNs) (Oh et al., 2017) take a similar approach but use the model to construct a look-ahead tree only when constructing bootstrap targets and selecting actions, similarly to TD-search (Silver et al., 2012).
Crucially, the model is not embedded in a planning algorithm during optimisation.
We propose a new approach to address the aforementioned problems: by formulating the tree search in a differentiable way and integrating it directly into the Q -function or policy, we train the entire agent, including its learned model, end-to-end, thus ensuring that the model is optimised for the correct goal and is suitable for on-line planning during execution of the policy.
Our approach can alternatively be viewed as a model-free method in which the fully connected layers of DQN are replaced by a recursive network that applies transition functions with shared parameters at each tree node expansion.
The resulting architecture, which we call TreeQN, encodes an inductive bias based on the prior knowledge that the environment is a stationary Markov process, which facilitates faster learning of better policies.
We also present an actor-critic variant, ATreeC, in which the tree planner is augmented with a softmax layer and used as a policy network.
The main contributions of this paper are: (i) we introduce TreeQN and ATreeC, end-to-end differentiable tree planners that serve as drop-in replacements of value and policy networks in deep RL with discrete actions; (ii) we show that TreeQN and ATreeC outperform their DQN-based counterparts in a box-pushing domain and a suite of Atari games, with deeper trees often outperforming shallower trees, and that TreeQN outperforms VPN (Oh et al., 2017) on most Atari games; (iii) we present a qualitative analysis of the learned trees in the box-pushing domain.
B ACKGROUND
We consider an agent learning to act in a Markov Decision Process (MDP), with the goal of maximising its expected discounted sum of rewards R t = (cid:80) ∞ t =
γ t r t , by learning a policy π ( s ) that maps states s ∈ S to actions a ∈ A .
The state-action value function ( Q -function) is deﬁned as Q π ( s , a ) =
E π
[ R t
[Equation]
] ; the optimal Q -function is
[Equation]
The Bellman optimality equation writes Q ∗ recursively as
[Equation]
where P is the MDP state transition function and r is a reward function, which for simplicity we assume to be deterministic.
Q -learning (Watkins & Dayan, 1992) uses a single-sample approximation of the contraction operator T to iteratively improve an estimate of Q ∗ .
In deep Q -learning (Mnih et al., 2015), Q is represented by a deep neural network with parameters θ , and is improved by regressing Q ( s , a ) to a target r +
γ max a (cid:48) Q ( s (cid:48) , a (cid:48) ; θ − ) , where θ − are the parameters of a target network periodically copied from θ .
We use a version of n -step Q -learning (Mnih et al., 2016) with synchronous environment threads.
In particular, starting at a timestep t , we roll forward n env = 16 threads for n = 5 timesteps each.
We then bootstrap off the ﬁnal states only and gather all n env × n = 80 transitions in a single batch for the backward pass, minimising the loss:
[Equation]
If the episode terminates, we use the remaining episode return as the target, without bootstrapping.
This algorithm’s actor-critic counterpart is A2C, a synchronous variant of A3C (Mnih et al., 2016) in which a policy π and state-value function V ( s ) are trained using the gradient: ∆ θ = (cid:88)
[Equation]
where A j is an advantage estimate given by (cid:80) jk =
1 γ j − k r t + n − k + γ j V
( s t + n ) − V ( s t + n − j ) , H is the policy entropy, β is a hyperparameter tuning the degree of entropy regularisation, and α is a hyperparameter controlling the relative learning rates of actor and critic.
These algorithms were chosen for their simplicity and reasonable wallclock speeds, but TreeQN can also be used in other algorithms, as described in Section 3.2.
Our implementations are based on
OpenAI Baselines (Hesse et al., 2017).
[Figure]
The canonical neural network architecture in deep RL with visual observations has a series of convolutional layers followed by two fully connected layers, where the ﬁnal layer produces one output for each action-value.
We can think of this network as ﬁrst calculating an encoding z t of the state s t which is then evaluated by the ﬁnal layer to estimate Q ∗ ( s t , a ) (see Fig. 1).
In tree-search on-line planning, a look-ahead tree of possible future states is constructed by recursively applying an environment model.
These states are typically evaluated by a heuristic, a learned value function, or Monte-Carlo rollouts.
Backups through the tree aggregate these values along with the immediate rewards accumulated along each path to estimate the value of taking an action in the current state.
This paper focuses on a simple tree-search with a deterministic transition function and no value uncertainty estimates, but our approach can be extended to tree-search variants like UCT (Kocsis & Szepesv´ari, 2006; Silver et al., 2016) if the components remain differentiable.
T REE QN
In this section, we propose TreeQN, a novel end-to-end differentiable architecture for tree planning.
We ﬁrst give an overview of the architecture, followed by details of each model component and the training procedure.
TreeQN uses a recursive tree-structured model between the encoded state z t and the predicted stateaction values Q ( s t , a ) .
Instead of directly estimating the state-action value from the current encoded state z t using fully connected layers as in DQN (Mnih et al., 2015)
(Fig. 1), TreeQN uses a recursive model to reﬁne its estimate of Q ( s t , a ) via learned transition, reward, and value functions (see Fig. 2).
Because these learned components are shared throughout the tree, TreeQN implements an inductive bias, missing from DQN, that reﬂects the prior knowledge that the Q -values are properties of a stationary Markov process.
Speciﬁcally, TreeQN learns an action-dependent transition function that, given a state representation z t , predicts the next state representation
z a i t +1 for action a i ∈ A , and the corresponding reward ˆ r a i t .
TreeQN applies this transition function recursively to construct a tree containing the state representations and rewards received for all possible sequences of actions up to some predeﬁned depth d (“Tree Transitioning” in Fig. 2).
The value of each predicted state V ( z ) is estimated with a value function module.
Using these values and the predicted rewards, TreeQN then performs a tree backup, mixing the k -step returns along each path in the tree using TD( λ ) (Sutton, 1988; Sutton & Barto, 1998).
This corresponds to “Value
Prediction” in Fig.
[Equation]
For 0 < λ < 1 , value estimates of the intermediate states are mixed into the ﬁnal Q -estimate, which encourages the intermediate nodes of the tree to correspond to meaningful states, and reduces the impact of outlier values.
When λ = 1 , Eq. 3 simpliﬁes to a backup through the tree using the Bellman equation:
[Equation]
We note that even for a tree depth of only one, TreeQN imposes a signiﬁcant structure on the value function by decomposing it as a sum of action-conditional reward and next-state value, and using a shared value function to evaluate each next-state representation.
Tree Transitioning
Value Prediction
[Figure]
Crucially, during training we backpropagate all the way from the ﬁnal Q -estimate, through the value prediction, tree transitioning, and encoding layers of the tree, i.e. , the entire network shown in Fig. 2.
Learning these components jointly ensures that they are useful for planning on-line.
3. M ODEL C OMPONENTS
In this section, we describe each of TreeQN’s components in more detail.
Encoder function.
As in DQN, a series of convolutional layers produces an embedding of the observed state, z t = encode ( s t ) .
Transition function.
Every action is parameterised by a fully connected layer.
We use residual connections (He et al., 2016) to calculate a next-state representation that carries information about the effect of taking action a i ∈ A :
[Equation]
where
W a i ∈ R k
×
k is a learnable action-dependent transition matrix.
The next-state representation is calculated for every action independently, but the transition function is shared for the same action throughout the tree.
A caveat is that the model can still learn to use different parts of the latent state space in different parts of the tree, which could undermine the intended parameter sharing in the model structure.
To help TreeQN learn useful transition functions that maintain quality and diversity in their plans, we introduce a unit-length projection of the state representations by simply dividing a state’s vector representation by its L2 norm before each application of the transition function, z t :=
z t / (cid:107) z t (cid:107) .
This prevents the magnitude of the representation from growing or shrinking, which encourages the behaviour of the transition function to be more consistent through the tree.
Reward function.
In addition to predicting the next state, we also predict the immediate reward for every action a i ∈ A in state z t using
[Equation]
where W r 1 ∈ R k × k , W r 2 ∈ R |A|× k and ReLU is the rectiﬁed linear unit (Nair & Hinton, 2010), and the predicted reward for a particular action ˆ r a i t is the i -th element of the vector ˆ r
( z t ) .
Value function.
The value of a state representation z is estimated with
[Equation]
where w ∈ R k .
3. T RAINING
The TreeQN architecture is fully differentiable, so we can directly use it in the place of a Q -function in any deep RL algorithm with discrete actions.
Differentiating through the entire tree ensures that the learned components are useful for planning on-line, as long as that planning is performed in the same way as during training.
While training, we also use an auxiliary reward-prediction loss, in the form of an L2 loss regressing ˆ r a 1 ( z t + j ) , the predicted reward at the ﬁrst level of the tree corresponding to the selected action a , to the true observed reward, for each of the n timesteps of n -step Q-learning:
[Equation]
where η is a hyperparameter weighting the loss.
In our experiments we simply use η = 1 , and we generally found the performance to be robust with respect to η , though this may be a consequence of the well behaved reward scaling in our environments.
The predicted rewards that this objective encourages the model to learn appear both in its own Q -value prediction and in the target for n -step Q -learning.
Consequently, we expect this auxiliary loss to be well aligned with the true objective.
By contrast, other potential auxiliary losses, e.g. , an observation-prediction loss, might help representation learning but would not explicitly learn any part of the desired target.
AT REE C
[Figure]
The intuitions guiding the design of TreeQN are as applicable to policy search as to valuebased RL, in that a policy can use a tree planner to improve its estimates of the optimal action probabilities (Gelly & Silver, 2007; Silver et al., 2017a).
As our proposed architecture is trained end-to-end, it can be easily adapted for use as a policy network.
In particular, we propose ATreeC, an actor-critic extension of TreeQN.
In this architecture, the policy network is identical to TreeQN, with an additional softmax layer that converts the Q estimates into the probabilities of a stochastic policy.
The critic shares the encoder parameters, and predicts a scalar state value with a single fully connected layer: V cr ( s ) = w (cid:62)
cr z + b cr .
The entire setup, shown in Fig. 3, is trained with A2C as described in
Section 2, with the addition of the same reward-prediction loss used for TreeQN.
TreeQN could also be used in the critic, but we leave this possibility to future work.
R ELATED W ORK
There is a long history of work combining model-based and model-free RL.
An early example is
Dyna-Q (Sutton, 1990) which trains a model-free algorithm with samples drawn from a learned model.
Similarly, van Seijen et al.
(2011) train a sparse model with some environment samples that can be used to reﬁne a model-free Q -function.
Gu et al.
(2016) use local linear models to generate additional samples for their model-free algorithm.
However, these approaches do not attempt to use the model on-line to improve value estimates.
In deep RL, value iteration networks (Tamar et al., 2016) use a learned differentiable model to plan on the ﬂy, but require planning over the full state space, which must also possess a spatial structure with local dynamics such that convolution operations can execute the planning algorithm.
The predictron (Silver et al., 2017b) instead learns abstract-state transition functions in order to predict values.
However, it is restricted to policy evaluation without control.
Value prediction networks (VPNs) (Oh et al., 2017) take a similar approach but are more closely related to our work because the learned model components are used in a tree for planning.
However, this tree is only used to construct targets and choose actions, and not to compute the value estimates during training.
Such estimates are instead produced from non-branching trajectories following on-policy action sequences.
By contrast, TreeQN is a uniﬁed architecture that constructs the tree dynamically at every timestep and differentiates through it, eliminating any mismatch between the model at training and test time.
Furthermore, we do not use convolutional transition functions, and hence do not impose spatial structure on the latent state representations.
These differences simplify training, allow our model to be used more ﬂexibly in other training regimes, and explain in part our substantially improved performance on the Atari benchmark.
Donti et al.
(2017) propose differentiating through a stochastic programming optimisation using a probabilistic model, to learn model parameters with respect to their true objective rather than a maximum-likelihood surrogate.
However, they do not tackle the full RL setting, and do not use the model to repeatedly or recursively reﬁne predictions.
Imagination-augmented agents (Weber et al., 2017) learn to improve policies by aggregating rollouts predicted by a model.
However, they rely on pretraining an observation-space model, which we argue will scale poorly to more complex environments.
Further, their aggregation of rollout trajectories takes the form of a generic RNN rather than a value function and tree backup, so the inductive bias based on the structure of the MDP is not explicitly present.
A class of value gradient methods (Deisenroth & Rasmussen, 2011; Fairbank & Alonso, 2012; Heess et al., 2015) also differentiates through models to train a policy.
However, this approach does not use the model during execution to reﬁne the policy, and requires continuous action spaces.
Oh et al.
(2015) and Chiappa et al.
(2017) propose methods for learning observation-prediction models in the Atari domain, but apply these models only to improve exploration.
Variants of scheduled sampling (Bengio et al., 2015) may be used to improve robustness of these models, but scaling to complex domains has proven challenging (Talvitie, 2014).
E XPERIMENTS
We evaluate TreeQN and ATreeC in a simple box-pushing environment, as well as on the subset of nine Atari environments that Oh et al.
(2017) use to evaluate VPN.
The experiments are designed to determine whether or not TreeQN and ATreeC outperform DQN, A2C, and VPN, and whether they can scale to complex domains.
We also investigate the impact of tree depth, and the extent to which the trees represent meaningful plans.
Full details of the experimental setup, as well as architecture and training hyperparameters, are given in the appendix.
Box Pushing.
First, we experiment in a simple box-pushing domain.
We randomly place an agent, 12 boxes, 5 goals and 6 obstacles on the center 6 × 6 tiles of an 8 × 8 grid.
Figure 5a shows a typical state.
The agent’s goal is to push boxes into goals in as few steps as possible while avoiding obstacles.
Boxes may not be pushed into each other.
The obstacles, however, are ‘soft’ in that they are do not block movement, but generate a negative reward if the agent or a box moves onto an obstacle.
This rewards better planning without causing excessive gridlock.
This environment is inspired by Sokoban, as used by Weber et al.
(2017), in that poor actions can generate irreversibly bad conﬁgurations.
However, the level generation process for Sokoban is challenging to reproduce exactly and has not been open-sourced.
More details of the environment and rewards are given in Appendix A.2.
Atari.
To demonstrate the general applicability of TreeQN and ATreeC to complex environments, we evaluate them on the Atari 2600 suite (Bellemare et al., 2013).
Following Oh et al.
(2017), we use their set of nine environments and a frameskip of 10 to facilitate planning over reasonable timescales.
R ESULTS & D ISCUSSION
In this section, we present our experimental results for TreeQN and ATreeC, as well as a qualitative analysis of learned trees for the box-pushing environment.
[Figure]
[Figure] [Figure]
[Figure] [Figure]
DQN
[Figure]
Steps
A v e r age
R
e w a r d o v e r
E p i s ode s
[Figure] [Figure]
7. B OX P USHING
Fig.
4a shows the results of TreeQN with tree depths 1, 2, and 3, compared to a DQN baseline.
In this domain, there is a clear advantage for the TreeQN architecture over DQN.
TreeQN learns policies that are substantially better at avoiding obstacles and lining boxes up with goals so they can be easily pushed in later.
TreeQN also substantially speeds up learning.
We hypothesise that the greater structure brought by our architecture regularises the model, encouraging appropriate state representations to be learned quickly.
Even a depth-1 tree improves performance signiﬁcantly, as disentangling the esimation of rewards and next-state values makes them easier to learn.
This is further facilitated by the sharing of value-function parameters across branches.
When trained with n -step Q-learning, there is a clear difference in the quality of TreeQN policies of different depths: deeper trees learn faster and plateau higher.
In the box-pushing domain, useful transition functions are relatively easy to learn, and the extra computation time with those transition modules can help reﬁne value estimates, yielding advantages for additional depth.
Fig.
4b shows the results of ATreeC with tree depths 1, 2, and 3, compared to an A2C baseline.
As with
TreeQN, ATreeC substantially outperforms the baseline.
Furthermore, thanks to its stochastic policy, it substantially outperforms TreeQN.
Whereas TreeQN and DQN sometimes indecisively bounce back and forth between adjacent states, ATreeC captures this uncertainty in its policy probabilities and thus acts more decisively.
However, unlike TreeQN, ATreeC shows no pronounced differences for different tree depths.
We hypothesise that this is due to a ceiling effect in this domain.
Nonetheless, this result demonstrates the ease with which TreeQN can be used as a drop-in replacement for any deep RL algorithm that learns policies or value functions for discrete actions.
To shed further light on how TreeQN works, we can also inspect the learned policies and trees, whose values often correspond to intuitive reasoning about sensible policies.
Fig.
5 shows one example.
In Fig.
5a, the agent can push a box onto a goal by moving either left or right.
Fig.
5c shows the corresponding tree values: each node contains the λ -mixed value of the path ending at that node, so the action chosen at each timestep is the ﬁrst action in the highest value path.
The two highest value paths are right-left, and left-right.
Thus, TreeQN has learned that both outcomes are positive and that returning for the other box is the correct next move.
Interestingly, the evaluation after the ﬁrst level prefers to move left ﬁrst, while the evaluation after the second level, after reﬁning its values, prefers to move right ﬁrst.
This superior path allows the agent to quickly reach another box adjacent to a goal in the top left.
Fig.
5b shows the environment after the agent takes its action (move right) and Fig.
5d shows the corresponding tree.
The agent now prefers to move left twice, scoring for pushing the second box.
+ + # − − # + #
X #
+ # # − + # # # − − # (
can score by
+ + # − − # + # −
X + # # − + # # # − − # z
t 1.
down 1.75 down 2.42 left 1.87 up 2.32 right 2.
left 2.61 down 2.88 left 2.39 up 2.93 right 1.
up 1.87 down 2.28 left 1.74 up 2.30 right 2.
right 2.37 down 2.98 left 2.48 up 2.73 right z t 1.
down 1.77 down 1.92 left 1.80 up 1.80 right 2.
left 1.97 down 2.13 left 1.95 up 2.03 right 1.
up 1.83 down 1.95 left 1.79 up 1.85 right 1.
right 1.80 down 1.96 left 1.83 up 1.85 right
[Figure]
Although walking over the red obstacle incurs a small penalty (accounted for by the lower absolute values in the tree), the correct action sequence is planned and executed.
However, many of the planned action sequences do not correspond to the selected future actions.
For example, the highest scoring path often consists of a repeated action, as in Fig.
5d, even if that is not the optimal sequence.
We hypothesise that the agents use these paths to conduct extra action-conditional computations, which makes sense given that the levels in the tree are only weakly tied to the timesteps of the underlying MDP.
Because TreeQN is ﬂexible and trained end-to-end, it can learn to use its transition function in different ways depending on the circumstance.
However, although this ﬂexibility can lead to improved performance, it can make the resulting plans less interpretable.
7. A TARI
Table 1 summarises all our Atari results, while Fig.
6 shows learning curves in depth.
TreeQN shows substantial beneﬁts in many environments compared to our DQN baseline, which itself often outperforms VPN (Oh et al., 2017).
ATreeC also often outperforms A2C. We present the mean performance of three random seeds, while the VPN results reported by Oh et al.
(2017), shown as dashed lines in Fig.
6, are the mean of the best ﬁve seeds of an unspeciﬁed number of trials.
[Figure]
A v e r age
R e w a r d o v e r
E p i s ode s
[Figure]
A v e r age
R e w a r d o v e r
E p i s ode s
[Figure]
[Table]
TreeQN.
In Alien, CrazyClimber [Table] and Seaquest, TreeQN signiﬁcantly outperforms the baselines, while we see smaller advantages for Enduro, MsPacman and Qbert.
Amidar and Krull show no signiﬁcant differences between methods, and Frostbite has degenerately poor performance for all methods.
Unsurprisingly, TreeQN performs well in domains that are suited to look-ahead planning, with simple dynamics that generalise well and tradeoffs between actions that become apparent only after several timesteps.
For example, an incorrect action in Alien can trap the agent down a corridor with an alien.
In Seaquest, planning ahead can help determine whether it is better to go deeper to collect more points or to surface for oxygen.
By contrast, Enduro is a racing game with mostly reactive decisions that can be easily adjusted in subsequent timesteps.
ATreeC. ATreeC, like TreeQN, outperforms its baseline strongly in Alien and Crazy Climber.
Compared to TreeQN, ATreeC’s performance is better across most environments, particularly on Qbert, reﬂecting an overall advantage for actor-critic also found by Mnih et al.
Performance is worse on Seaquest, revealing a deﬁciency in exploration as policy entropy collapses too rapidly.
In environments like Krull and Frostbite, all agents are clearly gated by their ability to explore.
Both of these games require the completion of sub-levels in order to accumulate large scores, and none of our agents reliably explore beyond the initial stage of the game.
Combining TreeQN and ATreeC with smart exploration mechanisms is an interesting direction for future work.
Compared to the box-pushing domain, there is less of a clear performance difference between trees of different depths.
In some environments (Alien, Seaquest), greater depth does appear to be employed usefully by TreeQN.
However, for the Atari domain the embedding size for the transition function we use is much larger (512 compared to 128), and the dynamics are much more complex.
Consequently, we expect that optimisation difﬁculties, and the challenge of learning abstract-state transition functions, impede the utility of deeper trees in some cases.
We look to future work to further reﬁne methods for learning to plan abstractly in complex domains.
An interesting behaviour in the Atari domain arises due to the over-speciﬁcation of the action space, which for instance contains a ‘no-op’ action, or a ‘ﬁre’ action in combination with hybrid move-and-ﬁre actions.
The agents can avoid taking these redundant actions, while still placing high value on these actions after the ﬁrst expansion in the tree to effectively carry out computations for action-unconditional rollouts.
The ﬂexibility of TreeQN and ATreeC enable our agents to ﬁnd this solution as a useful way to reﬁne value or policy estimates with further computation, but it can limit the interpretability of the trees in comparison to the box-pushing domain where good policies must use all actions.
C ONCLUSIONS & F UTURE W ORK
We presented TreeQN and ATreeC, new architectures for deep reinforcement learning in discreteaction domains that integrate differentiable on-line tree planning into the action-value function or policy.
Experiments on a box-pushing domain and a set of Atari games show the beneﬁt of these architectures over their model-free counterparts, as well as over VPN.
In future work we intend to investigate enabling more efﬁcient optimisation of deeper trees, encouraging the transition functions to produce interpretable plans, and integrating smart exploration.
Stochastic transition functions for non-deterministic environments are also an interesting avenue for research.
A A PPENDIX A. E FFECT OF INCREASED PARAMETERS ON N - STEP DQN
Naively adding parameters or depth to a DQN architecture is not guaranteed to improve performance, as shown in Fig.
Like our TreeQN implementation, the extra layers used in Fig.
7a have shared parameters, residual connections and state normalisation.
[Figure] s [Figure] [Figure]
[Figure]
4.0 1e
[Figure]
[Figure]
A. B OX P USHING
Environment.
For each episode a new level is generated by placing an agent, 12 boxes, 5 goals and 6 obstacles in the center 6 × 6 tiles of an 8 × 8 grid, sampling locations uniformly.
The outer tiles are left empty to prevent situations where boxes cannot be recovered.
The agent may move in the four cardinal directions.
If the agent steps off the grid the episode ends and it receives a penalty of -1.
If the agent moves into a box, it is pushed in the direction of movement.
Moving a box out of the grid generates a penalty of -0.1.
Moving a box into another box is not allowed and trying to do so generates a penalty of -0.1 while leaving all positions unchanged.
When a box is pushed into a goal, it is removed and the agent receives a reward of +1.
Obstacles generate a penalty of -0.2 when the agent or a box is moved onto them.
Moving the agent over goals incurs no penalty.
Lastly, at each timestep the agent receives a penalty of -0.01.
Episodes terminate when 75 timesteps have elapsed, the agent has left the grid, or no boxes remain.
The observation is given to the model as a tensor of size 5 × 8 × 8 .
The ﬁrst four channels are binary encodings of the position of the agent, goals, boxes, and obstacles respectively.
The ﬁnal channel is ﬁlled with the number of timesteps remaining (normalised by the total number of timesteps allowed).
Architecture.
The encoder consists of (conv-3x3-1-32, conv-3x3-1-32, conv-4x4-1-64, fc-128), where conv-wxh-s-n denotes a convolution with n ﬁlters of size w × h and stride s and fc-h denotes a fully connected layer with h hidden units.
All layers are separated with ReLU nonlinearities.
The hidden layer of the reward function MLP has 64 hidden units.
A. A TARI
Preprocessing of inputs follows the procedure of Mnih et al.
(2015), including concatenation of the last four frames as input, although we use a frameskip of 10.
Architecture.
The Atari experiments have the same architecture as for box-pushing, except for the encoder architecture which is as follows: (conv-8x8-4-32, conv-4x4-2-64, conv-3x3-1-64, fc-512).
A. O THER HYPERPARAMETERS
All experiments use RMSProp (Tieleman & Hinton, 2012) with a learning rate of 1e-4, a decay of α =0.99, and (cid:15)
=
1e-5.
The learning rate was tuned coarsely by running DQN on the Seaquest environment, and kept the same for all subsequent experiments (box-pushing and Atari).
For DQN and TreeQN, (cid:15) for (cid:15) -greedy exploration was decayed linearly from 1 to 0.05 over the ﬁrstmillion environment transitions observed.
For A2C and ATreeC, we use a value-function loss coefﬁcient α = 0 .
5 and an entropy regularisation β = 0 .
The reward prediction loss was scaled by η = 1 .
We use n steps = 5 and n envs = 16 , for a total batch size of 80.
For box-pushing, we add an L2 regularisation (weight 0.01) on the parameters of the transition function, which helped slightly.
We believe this is because the large matrix overparameterises the simple dynamics of
this environment