Title: No title found

Abstract: Decision-time planning policies with implicit dynamics models have been shown to work in discrete action spaces with Q learning. However, decision-time planning with implicit dynamics models in continuous action space has proven to be a difficult problem. Recent work in Reinforcement Learning has allowed for implicit model based approaches to be extended to Policy Gradient methods. In this work we propose Policy Tree Network (PTN). Policy Tree Network lies at the intersection of Model-Based Reinforcement Learning and Model-Free Reinforcement Learning. Policy Tree Network is a novel approach which, for the first time, demonstrates how to leverage an implicit model to perform decision-time planning with Policy Gradient methods in continuous action spaces. This work is empirically justified on 8 standard MuJoCo environments so that it can easily be compared with similar work done in this area. Additionally, we offer a lower bound on the worst case change in the mean of the policy when tree planning is used and theoretically justify our design choices.

1. INTRODUCTION:

Reinforcement Learning (RL), the study of learning what to do. Learning what to do in a given situation so as to maximize reward signals. Generally speaking, Reinforcement Learning problems are approached from either a Model-Free or Model-Based perspective. A Model-Free approach builds a policy through interacting with the environment and uses the acquired experience to improve the policy. A Model-Based approach builds a dynamics model of the environment and uses this improve a policy in a number of ways. Such as planning, exploration, and even training on imagined data (Sutton, 1990; Ha & Schmidhuber, 2018) . Model-Free Reinforcement Learning generally offers superior final performance. However, in exchange for strong final performance a large number of samples are required. Model-Based Reinforcement Learning offers better sample complexity but comes at the cost of a weaker final policy. In light of this, recent research (Wellmer & Kwok, 2019; Silver et al., 2016; Farquhar et al., 2017; Oh et al., 2017) has been focused on mixing both Model-Free and Model-Based Reinforcement Learning.

Model-Based Reinforcement Learning can be broken down further into two categories (Wellmer & Kwok, 2019) . The first is explicit dynamics models. Explicit dynamics models are when the observations are directly being reconstructed. The second is implicit dynamics models. Implicit dynamics models are when the dynamics model is learned indirectly. Some examples of this could be through predicting future rewards, values, policies or other auxiliary signals.

The work we present here, Policy Tree Network, lies at the intersection of Model-Free and Model-Based Reinforcement Learning. Our work builds off the work done by Wellmer & Kwok (2019) , where they showed how to build an implicit dynamics model for Policy Gradient methods in continuous action space. Policy Tree Network takes this a few steps further by showing how to leverage the implicit dynamics model for decision-time planning.

Our empirical results are on eight MuJoCo (Todorov et al., 2012) environments. In our experiments we validate the decision-time planning scheme that we introduce. Lastly we show Policy Tree Network outperforms the model-free baseline (Schulman et al., 2017) and the mixed model-free model-based baseline (Wellmer & Kwok, 2019) .

ATreeC (Farquhar et al., 2017) takes a similar approach to TreeQN except instead of a Q-learning approach ATreeC uses a policy gradient approach. We note that even though ATreeC uses a policy gradient approach it still does not work with continuous action spaces. Crucially, ATreeC optimizes the backed-up "Q" estimates to function as logits for a multinomial distribution. Because of this, the backed-up "Q" values in this case can be thought of as pseudo Q values.

Policy Prediction Network (PPN) (Wellmer & Kwok, 2019 ) is an approach to building implicit dynamics models with policy gradient methods. This was the first work to show how to build an implicit dynamics model in continuous action spaces. This was done by introducing overshooting objectives and a clipping scheme designed for overshooting objectives. The training algorithm is located in Section A.2.

EQUATION

L i,π t = 1 2 max(-ratio i t A GAE t+i , -ratio i t,clip A GAE t+i ) -α h H i t (2) L i,v t = 1 2 max((v i t,θ -R t+i ) 2 , (v i t,clip -R t+i ) 2 ) (3) L i,r t = 1 2 max((r i t,θ -r t+i ) 2 , (r i t,clip -r t+i ) 2 ) (4)

Where L i,π t , L i,v t , and L i,r t are the policy, value, and reward losses grounded at time t predicting i steps into the future. Additionally, subscript θ refers to estimates from the latest model parameters, H is policy entropy, A GAE t+i is the generalized advantage estimate (Schulman et al., 2016) , r is a reward target, r is a predicted reward, v is a predicted value, and R is a boot-strapped n-step return target. The α coefficients are hyper-parameters used to trade off importance of objectives. The clipped estimates are used to stay near estimates from old parameters(θ') and are further defined in Section A.3 along with the definition of the importance sampling ratio.

The short coming of PPN is that it did not lend itself well to leveraging the implicit dynamics model to perform decision-time planning and thus used the model-free policy to sample actions. Additionally, it does not make use of the Q function, a separate measure of how "good" an action is at a given state, as discussed in Section 3.2.

3. POLICY TREE NETWORK:

Policy Tree Network uses a combination of model-free and model-based techniques. Policy Tree Network builds off the PPN (Wellmer & Kwok, 2019) approach to learning an implicit-model. Actions during behavior time are chosen by a model-free policy. Learning is done with a model-based approach that follows the behavior policy's rollout trajectory. However, the test policy follows a model-based approach. An implicit transition model is embedded into the architecture. Through overshooting objectives, a dynamics model is learned and the collection of forward predictions offer the benefit of additional signal in the gradient updates.

Our novel contribution in this works is a decision-time planning algorithm that allows us to leverage the implicit dynamics model. The decision-time planning algorithm is based on building a tree of possible future values, actions, states, and rewards and then performing a tree backup algorithm we introduce. In previous works it was not possible to directly leverage the implicit model for decision time planning with policy gradient methods in continuous action space. Our empirical results in Section 4 demonstrate the advantage of PTN over the model-free baseline (PPO) and the mixed model-free & model-based baseline (PPN).

3.1. LEARNING:

Policy Tree Network is trained over a collection of overshooting objectives. We follow an identical clipping approach (Section A.3) to the one found in PPN. Targets used for training are computed exactly the same way as in PPN. More explicitly, no aspect of training depends on the decision-time planning algorithm described later in Section 3.2. While this might be desirable, we describe why it is difficult in A.1. We additionally introduce a term β Hafner et al. (2018) used to trade off importance of short and long-term predictions. Lastly, for simplicity we drop the extra value loss term found in PPN (Equation 1), leaving us with the following:

L t = d-1 i=0 β i (L i,π t + α v L i,v t + α r L i,r t ). Where L i,π

t , L i,v t , and L i,r t (equations 2, 3, 4) are the policy, value, and reward losses grounded at time t predicting i steps into the future.

The implicit dynamics model is jointly learned from the objectives in L t whenever i > 0. When i = 0 estimates are not required to pass through the transition network. In Section A.2 we show the training algorithms for PPO (Schulman et al., 2017) , PPN (Wellmer & Kwok, 2019) , and PTN which are quite similar.

PPN (Wellmer & Kwok, 2019) was shown to have returns drastically differ between training depth values. The point of β is to stabilize returns over different choices of training depth (Section A.6).

3.2. DECISION-TIME PLANNING:

Including branching in PTN is appealing because it would allow for decision-time planning. Furthermore, this appears to be a simple task since a transition model is already being learned. However, because the decision-time planning policy is difficult to directly measure with a PDF (Section A.1), we resort to performing decision-time planning exclusively at evaluation time.

Traditionally, at every state, the next action in a state is either determined by a policy, π, learned usually based on a policy gradient method, or an action value function Q that induces a policy that takes actions with the largest Q value. However, in PTN, we have access to independently learned policy and Q functions.1 There are three ways of using π and Q to create the final policy, π F . 1. setting π F = π as done in PPN, (Wellmer & Kwok, 2019) , which ignores the Q function.

2.. setting π:

F = argmax Q, called Q-backup which ignores π 2 . 3. setting π F = f (π, Q), called π-Q-backup, for some function f .

As (1) was described in PPN we will further dig into the details of (2) in Section 3.2.1 and (3) in Section 3.2.2. Methods discussed in this section do not impact the learning procedure discussed in Section 3.1 and are only used for decision-time planning at evaluation time.

We use a tree planning method which expands a tree up to a depth d expanding b branches at each depth and collecting reward, policy, and value predictions along the way. Tree expansion is done by recursively calling the core module (f core ) with b action samples and the abstract state. The action being passed to the core module could be sampled from a uniform distribution or the policy (π(a|ŝ i,j t,θ )), where ŝi,j t,θ represents the predicted abstract-state i steps into the future with branch j according to parameters θ. We show that if the correlation between π and Q is positive, sampling π will be more suitable and Section 4.2 shows that they are in fact positively correlated in practice. Candidate actions represent branches in the tree. Then the predicted estimates are backed-up with a TD-λ scheme (Sutton, 1988; Sutton & Barto, 2018; Farquhar et al., 2017) to assign a Q value to each branch as described in Algorithm 1. When performing decision-time planning, the action associated with the max backed-up Q-value is taken. This policy is referred to as π θ ,F , where θ denotes which parameters were used in the tree.

3.2.1. Q BACKUP:

Given an abstract state s = f enc θ (x) and an action a, the backed-up Q-value calculated from d-step planning is defined as:

EQUATION

where

s = f tr θ (s, a), v = f v θ (s)

, and r = f r θ (s, a). In Q backup, the final policy, π F , aims at selecting an action that maximizes the action value function, Q 0 . However, Q 0 is a combination of neural networks and finding arg max a Q 0 is a difficult problem. We use our tree expansion to achieve this. With tree expansion, at each depth i, we sample b actions based on π and choose the one with maximum Q i+1 to estimate arg max a Q i+1 in equation 6.

Intuitively, if π and Q are positively correlated, sampling from π allows us to find actions with higher Q compared with sampling from a uniform distribution, because, for a state s and an action a, larger π(a|s) value would suggest larger Q(s, a). This is the subject of the following lemma. Note that, for a given state s, the covariance between π(X|s) and Q(X, s) over a bounded action space is defined as cov U (π(X|s), Q(X, s)) where X is a random variable with distribution U , the uniform distribution over the action space. π and Q are positively correlated if cov U (π(X|s), Q(X, s)) is positive. Proof of the lemma can be found in the appendix.

Lemma 1. E X∼π [Q(X, s)] > E X∼U [Q(X, s)] if cov U (π(X|s), Q(X, s)) > 0.

Section 4.2 shows that the correlation between π and Q 0 is in fact positive in practice. Therefore, sampling actions according to π is justified.

Since sampling is done based on π, our final policy π F now depends on π as well as Q 0 . Next we investigate the relationship between π F , π and Q 0 . Consider the abstract state s = ŝ0,0 t,θ . Recall that Q 0 (a, s) is an action value function calculated based on the tree expansion and backup. Q 0 uses sampling, is probabilistic, and is a function of f core θ , λ, b and d. Action selection is done by sampling X i ∼ π for 1 ≤ i ≤ b and selecting arg max i Q 0 (X i , s). Thus, for an action a, cumulative density function (cdf) of π F (a|s) is given by

a -∞ π F (z|s)dz = P r(arg max i Q 0 (X i , s) ≤ a).

Intuitively, the branching factor b can be thought of as interpolating how much confidence we have in π and reward, value and transition networks (which make up Q 0 ): low b signifies confidence in π and larger b signifies more confidence in reward, value and transition networks. As b increases π F becomes less dependant on π. When b goes to infinity, π F only depends on reward, value and transition networks and not π. To investigate the maximum possible difference between mean of π and π F over all Q 0 functions we provide the next theorem. The theorem is provided in a one dimensional action space for better illustration, but the argument can be extended to higher dimensions. Proof of the theorem can be found in the appendix. Theorem 1. Let µ F be the mean of actions selected based on π F and µ, σ the mean and standard deviation of actions selected based on π. In the worst-case over all possible The bound is a lower bound on the worst-case and it shows that µ F can move at least as far as σ √ π from µ. In other words, there exist Q 0 functions for which the difference between mean of π F and π is at least σ √ π . Observe that the bound decreases as σ decreases, which intuitively makes sense since with lower σ, it takes more number of samples to obtain an action further away from µ.

Q 0 functions, |µ-µ F | ≥ Figure 3: PTN Backup Algorithm 1 π Q expand(s, i) initialize â, r, s, v , Q for j in b do â[j] ∼ π(s) r[j] = f r θ (s, â[j]) ŝ [j] = f tr θ (s, â[j]) v [j] = f va θ (ŝ [j]) if i + 1 == d then Q[j] = r[j] + γ v [j] π(a = â[j]|s) else Q tmp = π Q expand(ŝ [j], i + 1) Q[j] = r[j] + γ((1 -λ)v [j] + λ max â Q tmp) π(a = â[j]|s) return Q

The bound above sheds a light on one of the downsides of this approach. The final policy π F can become significantly different from π, depending on Q 0 and σ, even when b = 2 and as b grows it becomes only reliant on Q 0 (Section A.5 shows the numerical value of the bound for several b values). Although Section 4.2 shows empirically that π and Q 0 are on average correlated, there can be states where they are not. In such a scenario, Q backup may rely too heavily on Q 0 .

We previously mentioned that b can be seen as a confidence parameter, adjusting how much we rely on π and the networks that make up Q 0 . Nevertheless, it does not remedy the above mentioned problem. This is for two reasons, one is that b only takes discrete values. We mentioned above that b = 2 may shift the policy by too much, but we cannot reduce b any more without exclusively relying on π. Secondly, choosing a b for a given confidence level is a difficult problem, even if a discrete confidence parameter is sufficient. This is because b has a complicated relationship with π F , π and Q involving integrals. Thus, it is not clear that, for instance, if we have equal confidence in Q and π, what the value of b should be.

3.2.2. π-Q BACKUP:

Motivated by the issues discussed in section 3.2.1, we introduce a π-Q backup. Given a policy π and a Q-function Q, we want our final policy π F to depend on Q and π based on how much confidence we have in each. If we have equal confidence in both Q and π, then we would like π F to be a function depending on π and Q equally. One way to achieve this is using geometric mean of π and Q as an indicator of how good an action is3 . That is, we would like to select an action such that √ π × Q is maximized (note that if we have reasons to believe π or Q provide better estimates, we can take the weighted geometric mean).

The backup procedure is defined as mentioned in Section 3.2.1 , however, given an abstract state s = f enc θ (x) and an action a, the backed-up Q-value calculated from d-step planning is defined as:

Q i (s, a) = (r + γV i (s ))π(a|s) (7) V i (s) = v if i = d (1 -λ)v + λ max a Q i+1 (s, a) if i < d (8)

Note that neither Q i or V i in Equations (5, 6, 7, 8) are directly being optimized, they are only used for decision-time planning. Notice how estimates are scaled by likelihood of the action being sampled from π. This removes the previously mentioned issue of selected actions no longer depending on π. It's interesting to note that in this case when b = ∞ the policy is deterministic.

This does not manifest itself as an issue because the decision-time planning policy isn't used as the behavior policy so we are less concerned about exploration.

Note that we still need to estimate max a Q i+1 in equation 8. We take a similar approach to Qbackup mentioned in Section 3.2.1, since π and √ πQ are expected to be correlated. Furthermore, regarding the use of geometric mean, we note that it is more meaningful than the arithmetic mean in this scenario as the values may be of different scales and that it preserves proportional change as opposed to absolute change.

In Section 4.3 we will empirically explore both approaches to backup. The π-Q backup is shown in Algorithm 4. Though we note that if the scaling of Q estimates by π and the square root are dropped it would reduce to the max Q-backup algorithm described in Section 3.2.1.

4. EXPERIMENTS:

Our experiments seek to answer the following questions: (1) Is correlation between π and Q positive? (2) What style of backup performs better: a standard Q-value backup or a policy weighted Q-value backup and how does branching effect the returns of the decision-time planning policy? (3) Does PTN outperform the baselines? 4.1 EXPERIMENTAL SETUP Preprocessing in our experiments was done similarly to that of PPO (Schulman et al., 2017) and identical to the preprocessing in PPN (Wellmer & Kwok, 2019) . All models, PPO2, PPN, and PTN were implemented in Pytorch (Paszke et al., 2017) . The parameters (θ) are updated with the Adam optimizer (Kingma & Ba, 2014) . All the experiments are run for 1 million time steps unless otherwise noted.

Our PPO2 implementation uses the same hyperparameters as the baselines implementation (Dhariwal et al., 2017): 3 fully connected layers with 128 hidden units and tanh activations for the policy. 3 fully connected layers 128 hidden units and tanh activations for the critic. The largest difference in our PPO2 implementation is that we do not perform orthogonal initialization.

Our PTN implementation uses similar hyperparameters (identical to PPN): 2 fully connected layers with 128 hidden units and tanh activations for the embedding. 2 fully connected residual layers with 128 hidden units, tanh activations, and unit length projections of the abstract-state (Farquhar et al., 2017) for the transition module. 1 fully connected layer with 128 hidden units for the policy mean. 1 fully connected layer with 128 hidden units for the value. 1 fully connected layer with 128 hidden units for the reward. In practice we use Huber losses instead of L2 losses, as was done in related implicit model based works (Oh et al., 2017) .

We set β 0 = d and β >0 = 1. Values used for branching and depth are explicitly stated in corresponding experiments. We test on the 8 standard OpenAI MuJoCo environments (Todorov et al., 2012; Brockman et al., 2016) In this experiment we empirically measure the correlation between π and Q 0 (with d = 1). Positive correlation empirically justifies the use of lemma 1 that sampling from π offers a better approach to maximizing the backup objective (for both Q backup and πQ backup) as opposed to uniform sampling. To measure correlation we first train a PPN policy for 200 thousand time-steps on the MuJoCo environments. Next we collect trajectories over 10 thousand timesteps. Then at each observation we uniformly sample 100 actions and compute the corresponding Q-values and corresponding PDF points from π. This gives us 10 thousand estimates of correlation. From here we fit a normal distribution to the correlation results and report them in Table 3 .

Positive correlation tells us that sampling according to π will better maximize π F than a uniform sample over actions. This empirically justifies our approach to performing tree expansion. In this section we explore the impact of using planning as the test policy and which backup scheme captures the highest returns. We consider a PTN trained with a depth equal to 2 and we evaluate both backup procedures found in Section 3.2.

4.3. BRANCHING & BACKUP:

PTN π-Q backup PTN Q Backup b=1 b=2

As we can see in Table 2 , nearly all environment suffer from an increased branching factor when using Q backup. None of the environments show a clear trend that increasing the branching benefits returns. In fact, both Hopper and Walker2d environments show the opposite, as b increases returns decrease. Intuitively this makes sense, as b increases to infinity the action taken no longer relies on the policy π but instead relies entirely on the Q backup from reward and value predictions.

When the backup scheme takes into account the likelihood of an action coming from the policy, it is no longer true that as b increases to infinity the action selected by decision-time planning is independent from the policy. As we can see in Table 2 , every environment benefits from increasing branching used in decision-time planning when the backup scheme is π-Q backup discussed in Section 3.2.2. Comparing Q-backup and π-Q-backup in Table 2 it's obvious to see that the policy weighted Q backup offers much more robust returns. Now we clearly see that decision-time planning is indeed helpful. Furthermore, we now know that a π-Q backup offers a reliable approach. Going forward we will ask if it's possible to learn from this decision-time planning policy.

4.4. BASELINE COMPARISON:

To test our model, we chose to benchmark against PPO2 and PPN on eight MuJoCo (?) environments. We include our d = 2 b = 5 model in our baseline comparison. However, we note that it is possible that other configurations perform better on some environments.

As can be seen in Figure 5 , we find that PTN finds a better if not comparable policy in all of the environments. PTN has large performance gains over PPO's model-free approach, even though PTN's implicit planning module only looks a short distance into the future.

Environments with larger meaningful observation spaces, such as Humanoid, we suspect could benefit from a bigger latent space. In Humanoid, the latent space ( 128) is far smaller than the observation dimensions (376) and makes learning a useful implicit transition model infeasible. Given more compute we would explore increasing the latent space. 

5. CONCLUSION:

In this work we present for the first time an approach to decision-time planning for implicit dynamics models in continuous action space. We provide a theoretical justifications for our design choices and show strong empirical results which are improvements over previous related work. In future work we would like to investigate distilling the tree policy into the model-free policy, find a better sampling procedure for expansion, and incorporate research from recurrent networks to avoid issues with vanishing and exploding gradients in the implicit dynamics model.

A APPENDIX:

A.1 DECISION-TIME PLANNING AS THE BEHAVIOR POLICY Including branching in Policy Tree Network's behavior policy is appealing because it would allow for directly optimizing the test policy we intend to use and appears to be simple since an implicit transition model is already learned. However, as we will soon see this is not the case.

A natural first thought is just to recursively sample π b times up to a depth d, perform a backup on the expanded tree, and then take the action associated with the maximal base branch. This is the approach described in Section 3.2. However, this can not be used as the behavior policy because you are changing the distribution of how rollout actions are chosen in a way that is difficult to directly measure. The rollout policy π θ ,F is not equal to π θ used to do importance sampling in the policy gradient loss. When b = 1 then π θ ,F is the same as π θ . However as soon as b > 1 we begin to rely on the reward and value network to help decide which actions to take.

The previously mentioned issues can be avoided if π θ ,F is not used as the behavior policy. Instead with decision-time planning we can only use π θ ,F as our test policy. Interestingly, there is no guarantee that this should work because the policy we are optimizing for (π θ ) is not the same as the test policy. We note that a theoretical issue does arise when using π θ ,F as the test policy. The f v θ network parameters are trained to predict the value for policy π θ not π θ ,F . However, in practice this turns out not to be a large issue. 

A.3 CLIPPING ratio:

i t = π θ (a = a t+i |s = ŝi t,θ ) π θ (a = a t+i |s = ŝ0 t+i,θ ) ratio i t,clip = clip(ratio i t , 1 -, 1 + ), vi t,clip = clip(v i t,θ -v 0 t+i,θ , -, ) + v0 t+i,θ , ri t,clip = clip(r i t,θ -r0 t+i,θ , -, ) + r0 t+i,θ .

During training we follow the grounded clipping approach shown in PPN (Wellmer & Kwok, 2019) . The clipped estimates are defined above. Where is a hyperparameter that defines the size of the clipping region. Clipping all the network heads turns out to be imperative to the learning process (Wellmer & Kwok, 2019) 

. A.4 PROOFS Proof of Lemma 1. Note that cov U (π(X|s), Q(X, s)) = E U [Q(X, s)π(X|s)] - E U [Q(X, s)]E U [π(X|s)].

Let R be a bounded action space. By definition, the covariance equals

1 area(R) x∈R (Q(x, s)π(x|s)) -E U [Q(X, s)] 1 area(R) x∈R π(x|s), where x∈R π(x|s) = 1 and x∈R (Q(x, s)π(x|s)) = E π [Q(X, s)].

The result follows because the covariance is positive.

Proof of Theorem 1. We prove this by showing the existence of Q 0 functions for which the inequalities holds. We drop s from our notation for convenience and write Q(a, s) as Q(a) and π(a|s) as π(a).

Let X F ∼ π F , i.e., X F = arg max i Q 0 (X i ), where X i ∼ π. Consider a Q 0 (a) that is strictly decreasing as a increases. Then A similar argument but setting Q 0 to be a strictly increasing function in a gives the inequality for when µ F is larger than µ.

A.5 NUMERICAL EVALUATION OF BOUND IN THEOREM 1:

Fig. 6 shows the numerical values for the bound in Theorem 1 for several branching values. Note that the bound becomes larger than 1 × σ when b = 4. This means that when b = 4, the difference between the mean of π F and π can become larger than one standard deviation of π, which is a significant difference.

A.6 LEARNING MODIFICATIONS:

Notice that these returns are worse than PTN(with decision-time planning) shown in Figure 5 . This shows performance benefits from the modifications shown in Section 3.1(ex: β). The point of Wellmer & Kwok (2019) showed that optimal depth is highly dependent on the environment and returns can drastically differ. While the modifications certainly do not entirely fix this, we find that it mitigates large differences in returns from different training depths.

:

by independent we mean π is not just a policy induced by selecting max of Q.

when b = ∞

When a negative or zero valued Q exists in the tree backup, modify the geometric mean by adding | min(Q)| + δ to the Q estimates. δ is a positive small non-zero term which can be optimizedHabib (2012)

Figures and Tables:

Figure 1: Predict policies, rewards, abstract states, and the value of the abstract states (Wellmer & Kwok, 2019).

Figure 2: PTN Expansion, a, r, v outputs from f core are omitted for simplicity

Figure 4: pseudo-code for π-Q expansion and backup algorithm. Where the Q returned is a b dimensional vector. b|σ ∞ -∞ zφ(z)( erfc(z) 2 ( z √ 2 )) b-1 dz|, where erfc is the complementary error function and φ is the p.d.f of standard normal distribution. Setting b = 2, we get the looser bound |µ -µ F | ≥ σ √ π in the worstcase.

Figure 5: Results on 1 million step MuJoCo benchmark. Dark lines represent the mean return and the shaded region is a standard deviation above and below the mean.

PPN TRAINING ALGORITHM Algorithm 2 Policy Prediction Network(PPN) (Wellmer & Kwok, 2019) Initialize parameters θ θ = θ for iteration=1, 2, . . . do Run policy π θ in environment for n time steps Compute advantage estimates A GAE 1 , . . . , A GAE n for epoch= 1, . . . , K do Shuffle n samples into mini-batches of size M ≤ n for each mini-batch do T is the set of samples selected for the mini-batch L mb = 1 M t∈T L t Optimize L mb w.r.t. θ θ = θ PPO, PPN, and PTN use similar training algorithms, the main difference stems from how L t is defined. Where in Algorithm 2, n time steps are used to collect trajectories, K is the number of epochs, M is mini-batch size, and T represents the randomly sampled time steps in a specific minibatch.

-∞ π F (x)dx = P r(X F ≤ c) = P r(∃i, X i ≤ c)=1 -P r(∀ i , X i > c). Since X i 's are i.i.d, P r(∀ i , X i > c) = P r(X > c) b = erfc( π(c)dc. Substitute c = σz + µ, we have µ F = bσ φ(z)dz = 1, which gives the first bound. Setting b = 2, we have that µ F -µ = -σ √ π + µ -µ = -σ √ π .

Figure 6: Numerical evaluation of bound in Theorem 1 for various branching values train depth = 5 train depth = 10 PTN PPN PTN PPN Hopper-v2 1944.0 ± 258.5 2191.6 ± 183.3 1672.7 ± 457.5 1752.3 ± 625.4 Walker2d-v2 2936.9 ± 893.2 2808.4 ± 647.3 3054.8 ± 434.0 2565.9 ± 327.9 Swimmer-v2 89.4 ± 35.3 73.3 ± 34.5 80.3 ± 35.7 57.5 ± 30.3 HalfCheetah-v2 3439.8 ± 698.6 3410.5 ± 840.9 3638.9 ± 545.2 3632.8 ± 819.0 InvertedPendulum-v2 971.7 ± 56.5 979.6 ± 25.7 1000.0 ± 0.0 1000.0 ± 0.0 InvertedDoublePendulum-v2 4332.6 ± 192.9 3344.5 ± 234.4 4360.2 ± 145.1 3186.8 ± 87.7

Summary of correlation between Q and π.

Table content: <table><tr><td colspan="3">: Hopper, Walker2d, Swimmer, HalfCheetah, InvertedPendulum, In-</td></tr><tr><td>vertedDoublePendulum, Humanoid, and Ant.</td><td/><td/></tr><tr><td>4.2 CORRELATION</td><td/><td/></tr><tr><td colspan="3">sample correlation mean sample correlation variance</td></tr><tr><td>Hopper-v2</td><td>0.937</td><td>0.018</td></tr><tr><td>Walker2d-v2</td><td>0.971</td><td>0.017</td></tr><tr><td>Swimmer-v2</td><td>0.945</td><td>0.016</td></tr><tr><td>HalfCheetah-v2</td><td>0.968</td><td>0.018</td></tr><tr><td>InvertedPendulum-v2</td><td>0.969</td><td>0.008</td></tr><tr><td>InvertedDoublePendulum-v2</td><td>0.971</td><td>0.005</td></tr><tr><td>Humanoid-v2</td><td>0.934</td><td>0.022</td></tr><tr><td>Ant-v2</td><td>0.972</td><td>0.019</td></tr></table>

branching and backup experiment where d = 2 and λ = 0.95

Table content: <table><tr><td/><td/><td/><td/><td/><td>b=3</td><td/><td>b=5</td><td/><td>b=1</td><td/><td>b=2</td><td/><td>b=3</td><td/><td>b=5</td><td/></tr><tr><td>Hopper</td><td colspan="2">2296.8 ±</td><td colspan="2">2784.4 ±</td><td colspan="2">2981.3 ±</td><td colspan="2">3083.5 ±</td><td colspan="2">2155.9 ±</td><td colspan="2">1763.0 ±</td><td colspan="2">1400.8 ±</td><td colspan="2">1103.4 ±</td></tr><tr><td/><td>403.6</td><td/><td>357.8</td><td/><td>306.7</td><td/><td>244.1</td><td/><td>343.9</td><td/><td>431.9</td><td/><td>423.5</td><td/><td>398.1</td><td/></tr><tr><td>Walker2d</td><td colspan="2">3442.7 ±</td><td colspan="2">4080.3 ±</td><td colspan="2">4189.1 ±</td><td colspan="2">4253.8 ±</td><td colspan="2">3381.5 ±</td><td colspan="2">3539.7 ±</td><td colspan="2">3160.7 ±</td><td colspan="2">2722.8 ±</td></tr><tr><td/><td>182.0</td><td/><td>295.5</td><td/><td>355.9</td><td/><td>334.3</td><td/><td>378.3</td><td/><td>517.6</td><td/><td>743.5</td><td/><td>1033.0</td><td/></tr><tr><td>Swimmer</td><td>102.1</td><td>±</td><td>105.6</td><td>±</td><td>106.6</td><td>±</td><td>107.5</td><td>±</td><td>89.5</td><td>±</td><td>88.6</td><td>±</td><td>84.5</td><td>±</td><td>75.5</td><td>±</td></tr><tr><td/><td>28.4</td><td/><td>29.2</td><td/><td>29.2</td><td/><td>29.2</td><td/><td>35.1</td><td/><td>34.4</td><td/><td>32.0</td><td/><td>30.9</td><td/></tr><tr><td>HalfCheetah</td><td colspan="2">4082.1 ±</td><td colspan="2">4346.3 ±</td><td colspan="2">4459.8 ±</td><td colspan="2">4555.0 ±</td><td colspan="2">3435.8 ±</td><td colspan="2">3501.7 ±</td><td colspan="2">3526.3 ±</td><td colspan="2">3467.7 ±</td></tr><tr><td/><td>659.9</td><td/><td>649.1</td><td/><td>638.8</td><td/><td>643.6</td><td/><td>742.4</td><td/><td>748.2</td><td/><td>734.5</td><td/><td>741.1</td><td/></tr><tr><td>InvertedPendulum</td><td colspan="2">1000.0 ±</td><td colspan="2">1000.0 ±</td><td colspan="2">1000.0 ±</td><td colspan="2">1000.0 ±</td><td>995.2</td><td>±</td><td>998.7</td><td>±</td><td>998.6</td><td>±</td><td>998.2</td><td>±</td></tr><tr><td/><td>0.0</td><td/><td>0.0</td><td/><td>0.0</td><td/><td>0.0</td><td/><td>9.6</td><td/><td>2.6</td><td/><td>2.8</td><td/><td>3.6</td><td/></tr><tr><td>InvertedDoublePendulum</td><td colspan="2">6301.4 ±</td><td colspan="2">9347.6 ±</td><td colspan="2">9353.2 ±</td><td colspan="2">9356.9 ±</td><td colspan="2">6286.6 ±</td><td colspan="2">9244.3 ±</td><td colspan="2">9159.7 ±</td><td colspan="2">7929.4 ±</td></tr><tr><td/><td>299.8</td><td/><td>0.2</td><td/><td>0.2</td><td/><td>0.1</td><td/><td>308.4</td><td/><td>64.5</td><td/><td>225.1</td><td/><td>1498.2</td><td/></tr><tr><td>Ant</td><td colspan="2">1730.3 ±</td><td colspan="2">2109.3 ±</td><td colspan="2">2183.2 ±</td><td colspan="2">2370.0 ±</td><td colspan="2">1946.4 ±</td><td colspan="2">2047.1 ±</td><td colspan="2">2078.5 ±</td><td colspan="2">2013.3 ±</td></tr><tr><td/><td>225.7</td><td/><td>329.7</td><td/><td>324.8</td><td/><td>371.3</td><td/><td>266.1</td><td/><td>202.7</td><td/><td>197.5</td><td/><td>279.1</td><td/></tr><tr><td>Humanoid</td><td>747.8</td><td>±</td><td>788.0</td><td>±</td><td>817.8</td><td>±</td><td>895.5</td><td>±</td><td>933.1</td><td>±</td><td>957.3</td><td>±</td><td>919.2</td><td>±</td><td>859.8</td><td>±</td></tr><tr><td/><td>280.7</td><td/><td>355.6</td><td/><td>410.5</td><td/><td>566.3</td><td/><td>432.9</td><td/><td>449.2</td><td/><td>345.9</td><td/><td>272.9</td><td/></tr></table>

6 ± 183.3 1672.7 ± 457.5 1752.3 ± 625.4 Walker2d-v2 2936.9 ± 893.2 2808.4 ± 647.3 3054.8 ± 434.0 2565.9 ± 327.9 Swimmer-v2 89.4 ± 35.3 73.3 ± 34.5 80.3 ± 35.7 57.5 ± 30.3 HalfCheetah-v2 3439.8 ± 698.6 3410.5 ± 840.9 3638.9 ± 545.2 3632.8 ± 819.0 InvertedPendulum-v2 971.7 ± 56.5 979.6 ± 25.7 1000.0 ± 0.0 1000.0 ± 0.0 InvertedDoublePendulum-v2 4332.6 ± 192.9 3344.5 ± 234.4 4360.2 ± 145.1 3186.8 ± 87.7 Returns from PTN(no π F /decision-time planning) and PPN using only the model-free policies β(found in PTN but not PPN) is to stabilize returns over different values of training depth. In PPN,

Table content: <table/>

Peer review (iclr2020):
The paper proposes a modification to Policy Prediction Networks (PPN) in which the learned transition-, reward- and value function models are used at test-time in a planning procedure. A second contribution is the *pi-Q-backup* which uses the geometric mean of both the policy and the value function as maximisation target of the planning step. Overall, I find the idea interesting and the experimental evaluations promising. However, I am voting for *weak reject* for the reasons outlined below. If some (or all) of them are address, I*d be happy to raise my score. - I found the paper hard to understand. In particular, the algorithm PPN on which this work is build is not explained at all, requiring the reader to read the original PPN paper. Including a description of PPN, including it*s main features, would greatly help the paper. Second, I am still not sure I correctly undestand when each component is used. As I currently understand it, the main usage of the model during training time is to compute the Advantages in equation (2)? Or are those computed based on rollouts? If so, where is the model actually being used during training? - Figure 1 is taken directly from the PPN paper without any reference or citation (as far as I can tell). - For the comparison in Figure 5, it would be great if PPN could also be tested with the newly introduced parameter _x0008_eta_i. At the moment, it is hard to tell whether the performance gains are due to _x0008_eta or due to the proposed planning scheme. - I*m confused about Theorem 1: Wouldn*t we want an upper bound on the difference of means? Also, what does *worst-case* mean? Is that for the *worst* Q-function we could choose? Edit: Thank you for your comments and updated manuscript. I think the writing has improved significantly, but could still be further improved and clarified. In particular, the question of how the model and other components at various points in time could be made more obvious. I found the authors* response to R3 here helpful as well. At least for me some of the confusion arises not due to the complexity of the proposed approach, but just because combining real and *simulated* transitions can be used and mixed in so many different ways that it*s important to be clear about it. Also, at least personally, I found the explanation *Learning is done with a model-based approach that follows the behaviour policy’s rollout trajectory. However, the test policy follows a model-based approach* still not very helpful. Overall, I think the presentation is on a good way but needs some more work. With that being said and now having a better understanding of the algorithm I think this is very interesting work. However, I share R1*s concerns about the computation of the pi-Q backup, in particular that it seems arbitrary and doesn*t handle negative values. I*m also not convinced that adding |min(Q)| is a good solutions as a) we don*t always have access to that value and b) If I*m not wrong, than pi_F is not invariant under a shift of Q. I*m wondering why the authors decided to take the geometric mean instead of following the more typically used approach of using exp(Q/Temperature)*pi to combine Q-function and a policy distribution (see e.g. the *Control as Inference* literature, the *Maximum a Posteriori Policy Optimisation* or *Soft Actor Critic* algorithms, in particular the *Soft Value functions*. I think this should at the very least be an ablation study, and could be even performing better and at the very least be robust to negative values. Overall, I think this is very interesting work and could become a very strong paper, but I will remain to recommend a *weak reject* because I think it needs some more work to get there.