D IFFERENTIABLE T RUST R EGION L AYERS FOR D EEP R EINFORCEMENT L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Trust region methods are a popular tool in reinforcement learning as they yield robust policy updates in continuous and discrete action spaces.
However, enforcing such trust regions in deep reinforcement learning is difﬁcult.
Hence, many approaches, such as Trust Region Policy Optimization (TRPO) and Proximal Policy
Optimization (PPO), are based on approximations.
Due to those approximations, they violate the constraints or fail to ﬁnd the optimal solution within the trust region.
Moreover, they are difﬁcult to implement, lack sufﬁcient exploration, and have been shown to depend on seemingly unrelated implementation choices.
In this work, we propose differentiable neural network layers to enforce trust regions for deep Gaussian policies via closed-form projections.
Unlike existing methods, those layers formalize trust regions for each state individually and can complement existing reinforcement learning algorithms.
We derive trust region projections based on the Kullback-Leibler divergence, the Wasserstein L2 distance, and the Frobenius norm for Gaussian distributions. We empirically demonstrate that those projection layers achieve similar or better results than existing methods while being almost agnostic to speciﬁc implementation choices.
I NTRODUCTION
Deep reinforcement learning has shown considerable advances in recent years with prominent application areas such as games (Mnih et al., 2015; Silver et al., 2017), robotics (Levine et al., 2015), and control (Duan et al., 2016).
In policy search, policy gradient (PG) methods have been highly successful and have gained, among others, great popularity (Peters & Schaal, 2008).
However, often it is difﬁcult to tune learning rates for vanilla PG methods, because they tend to reduce the entropy of the policy too quickly.
This results in a lack of exploration and, as a consequence, in premature or slow convergence.
A common practice to mitigate these limitations is to impose a constraint on the allowed change between two successive policies.
Kakade & Langford (2002) provided a theoretical justiﬁcation for this in the approximate policy iteration setting.
Two of the arguably most favored policy search algorithms, Trust Region Policy Optimization (TRPO) (Schulman et al., 2015a)
and
Proximal Policy Optimization (PPO) (Schulman et al., 2017), follow this idea using the KullbackLeibler divergence (KL) between successive policies as a constraint.
We propose closed-form projections for Gaussian policies, realized as differentiable neural network layers. These layers constrain the change in successive policies by projecting the updated policy onto trust regions.
First, this approach is more stable with respect to what Engstrom et al.
(2020) refer to as code-level optimizations than other approaches. Second, it comes with the beneﬁt of imposing constraints for individual states, allowing for the possibility of state-dependent trust regions. This allows us to constrain the state-wise maximum change of successive policies. In this we differ from previous works, that constrain only the expected change and thus cannot rely on exact guarantees of monotonic improvement.
Furthermore, we propose three different similarity measures, the KL divergence, the Wasserstein L2 distance, and the Frobenius norm, to base our trust region approach on.
The last layer of the projected policy is now the the trust region layer which relies on the old policy as input. This would result in a ever-growing stack of policies which is clearly infeasible. To circumvent this issue we add an additional loss to the reinforcement learning objective which ensures the input and output of the projection stay close together.
While this still results in an approximation of the trust region update, we show that the trust regions are properly enforced.
We also extend our approach to allow for a controlled evolution of the entropy of the policy, which has been shown to increase performance in difﬁcult exploration problems (Pajarinen et al., 2019; Akrour et al., 2019).
We compare and discuss the effect of the different similarity measures and the entropy control on the optimization process.
Additionally, we benchmark our algorithm against existing methods and demonstrate that we achieve similar or better performance.
R ELATED W ORK
Approximate Trust Regions.
Bounding the size of the policy update in policy search is a common approach.
While, originally, Kakade & Langford (2002) focused on a method based on mixing policies, nowadays most approaches use KL trust regions to bound the updates.
Peters et al.
(2010) proposed a ﬁrst approach to such trust regions by formulating the problem as a constraint optimization and provided a solution based on the dual of that optimization problem.
Yet, this approach is not straightforwardly extendable to highly non-linear policies, such as neural networks.
In an attempt to transfer those ideas to deep learning, TRPO (Schulman et al., 2015a) approximates the KL constraint using the Fisher information matrix and natural policy gradient updates (Peters & Schaal, 2008; Kakade, 2001), along with a backtracking line search to enforce a hard KL constraint.
Yet, the resulting algorithm scales poorly.
Thus, Schulman et al.
(2017) introduced PPO, which does not directly enforce the KL trust region, but clips the the probability ratio in the importance sampling objective.
This allows using efﬁcient ﬁrst-order optimization methods while maintaining robust training.
However, Engstrom et al.
(2020) and Andrychowicz et al.
(2020) recently showed that implementation choices are essential for achieving state-of-the-art results with PPO. Code-level optimizations , such as reward scaling as well as observation, reward, gradient and value function clipping, can even compensate for removing core parts of the algorithm, e. g, the clipping of the probability ratio.
Additionally, PPO heavily relies on its exploration behavior and might get stuck in local optima (Wang et al., 2019).
Tangkaratt et al.
(2018) use a closed-form solution for the constraint optimization based on the method of Lagrangian multipliers.
Yet, they require a quadratic parametrization of the Q-Function, which can limit the performance.
Pajarinen et al.
(2019) introduced an approach based on compatible value function approximations to realize KL trust regions.
Based on the reinforcement learning as inference paradigm (Levine, 2018), Abdolmaleki et al.
(2018) introduced an actor-critic approach using an Expectation-Maximization based optimization with KL trust regions in both the E-step and M-step.
Song et al.
(2020) proposed an on-policy version of this approach using a similar optimization scheme and constraints.
Projections for Trust Regions.
Akrour et al.
(2019) proposed Projected Approximate Policy Iteration (PAPI), a projection-based solution to realize KL trust regions.
Their method projects an intermediate policy, which already satisﬁes the trust region constraint, onto the constraint bounds.
This maximizes the size of the update step.
However, PAPI relies on other trust region methods to generate this intermediary policy and cannot operate in a stand-alone setting.
Additionally, the projection is not directly part of the policy optimization but applied afterwards, which can result in sub-optimal policies.
Further, in view of high computational complexity, TRPO and PAPI have to simplify the problem of constraining the change by working with the expected KL divergence.
Opposed to that, we realize the projections as fully differentiable network layers which are directly included in the optimization process.
Additionally, our projections enforce the constraints per state. This allows for better control of the change between subsequent policies and for state-dependent trust regions.
For the KL-based projection layer we need to resort to numerical optimization and implicit gradients for convex optimizations (Amos & Kolter, 2017; Agrawal et al., 2019).
Thus, we investigate two alternative projections based on the Wasserstein L2 (W2) and Frobenius norm, which allow for closed form solutions.
Both, Wasserstein and Frobenius norm have found only limited application in reinforcement learning.
Pacchiano et al.
(2020) use the Wasserstein distance to score behaviors of agents.
Richemond & Maginnis (2017) proposed an alternative algorithm for bandits with
Wasserstein based trust regions.
Song & Zhao (2020) focus on solving the trust region problem for distributional policies using both KL and Wasserstein based trust regions for discrete action spaces.
Our projections are applicable independently of the underlying algorithm and only assume a Gaussian policy, a common assumption for continuous action spaces.
Several authors (Dalal et al., 2018; Chow et al., 2019; Yang et al., 2020) used projections as network layers to enforce limitations in the action or state space given environmental restrictions, such as robotic joint limits.
Entropy Control.
Abdolmaleki et al.
(2015) introduced the idea of explicitly controlling the decrease in entropy during the optimization process, which later was extended to deep reinforcement learning by Pajarinen et al.
(2019) and Akrour et al.
They use either an exponential or linear decay of the entropy during policy optimization to control the exploration process and escape local optima.
To leverage those beneﬁts, we embed this entropy control mechanism in our differentiable trust region layers.
P RELIMINARIES AND P ROBLEM S TATEMENT
We consider the general problem of a policy search in a Markov Decision Process (MDP) deﬁned by the tuple ( S , A , T , R , P 0 , γ ) .
We assume the state space S and action space A are continuous and the transition probabilities
T : S ×
A × S → [0 , 1] describe the probability transitioning to state s t + ∈ S given the current state s t ∈ S and action a t ∈ A . We denote the initial state distributions as P 0 : S → [0 , 1] .
The reward returned by the environment is given by a function
R : S ×
A → R and γ ∈ [0 , 1) describes the discount factor.
Our goal is to maximize the expected accumulated discounted reward R γ =
E
[Equation]
R [Equation] s t , [Equation]
t )] . To ﬁnd the optimal policy, traditional PG methods often make use of the likelihood ratio gradient and an importance sampling estimator.
Moreover, instead of directly optimizing the returns, it has been shown to be more effective to optimize the advantage function as this results in an unbiased estimator of the gradient with less variance
[Equation]
where A π ( s, a ) =
E [ R γ | s 0 = s, a 0 = a ; π ] − E
[ R γ | s {+
+}
= s ; π ] describes the advantage function, and the expectation is w.r.t
π θ old , i.e. s (cid:48) ∼ T ( ·| s, a ) , a ∼ π θ old ( ·| s ) , s 0 ∼ P 0 ( s 0 ) , s ∼ ρ π θ old where ρ π θ old is a stationary distribution of policy π θ old .
The advantage function can be estimated by generalized advantage estimation (GAE) (Schulman et al., 2015b).
This objective is used by most TRPO and PPO implementations, even though GAE was not used in the original TRPO paper.
Trust region methods use additional constraints for the given objective.
Using a constraint on the maximum KL over the states has been shown to guarantee monotonic improvement of the policy (Schulman et al., 2015a).
However, as all current approaches do not use a maximum KL constraint but an expected KL constraint, thus the monotonic improvement guarantee also does not hold exactly.
We are not aware of such results for the W2 distance or the Frobenius norm.
For our projections we assume Gaussian policies
π θ old ( a t | s t )
=
N ( a t | µ old ( s t ) , Σ old ( s t )) and π θ ( a t | s t )
=
N ( a t | µ ( s t ) , Σ( s t )) represent the old as well as the current policy, respectively.
We explore different trust regions on top of Equation 1 that employ different similarity measures between old and new distributions, i. e. the frequently used reverse KL, the Wasserstein L2 distance, and the Frobenius norm.
Reverse KL Divergence.
The KL divergence between two Gaussian distributions with means µand µ 2 and covariances Σ 1 and Σ 2 can generally be written as
[Equation]
where dim is the dimensionality of µ 1 , µ 2 .
The KL uses the Mahalanobis distance to measure the similarity between the two mean vectors.
The difference of the covariances is measured by the difference in shape, i.e., the difference in scale, given by the log ratio of the determinants, plus the difference in rotation, given by the trace term.
Given the KL is non-symmetric, it is clearly not a distance, yet still a frequently used divergence between distributions.
We will use the more common reverse KL for our trust region, where the ﬁrst argument is the new policy and the second is the old policy.
Wasserstein Distance.
The Wasserstein distance is a distance measure based on an optimal transport formulation, for more details see Villani (2008).
The Wasserstein-2 distance for two Gaussian distributions can generally be written as
[Equation]
A key difference to the KL divergence is that the Wasserstein distance is a symmetric distance measure, i. e., W 2 ( q, p ) =
W 2 ( p, q ) .
Our experiments also revealed that it is beneﬁcial to measure the W2 distance in a metric space deﬁned by the covariance of the old policy distribution, denoted here as Σ 2 , as the distance measure is then more sensitive to the data-generating distribution.
The
W2 distance in this metric space reads
[Equation]
Frobenius Norm.
The Frobenius norm is a matrix norm and can be directly applied to the difference of the covariance matrices of the Gaussian distributions.
To measure the distance of the mean vectors, we will, similar to the KL divergence, employ the Mahalanobis distance as this empirically lead to an improved performance in comparison to just taking the squared distance.
Hence, we will denote the following metric as Frobenius norm between two Gaussian distributions
[Equation]
The Frobenius norm also constitutes a symmetric distance measure. [Equation]
IFFERENTIABLE T RUST -R EGION L AYERS FOR G AUSSIAN P OLICIES
We present projections based on the three similarity measures, i. e., Frobenius norm, Wasserstein L2 distance, and KL divergence.
These projections realize state-wise trust regions and can directly be integrated in the optimization process as differentiable neural network layers.
Additionally, we extend the trust region layers to include an entropy constraint to gain control over the evolution of the policy entropy during optimization.
The trust regions are deﬁned by a distance or divergence d ( π ( ·| s ) , π old ( ·| s )) between probability distributions.
Complementing Equation 1 with the trust re- gion constraint leads to
[Equation]
While, in principle, we want to enforce the constraint for every possible state, in practice, we can only enforce them for states sampled from rollouts of the current policy.
To solve the problem in Equation 2, a standard neural network will output the parameters µ, Σ of a Gaussian distribution π θ , ignoring the trust region bounds.
These parameters are provided to the trust region layers, together with the mean and covariance of the old policy and a parameter specifying the size of the trust region (cid:15) .
The new policy is then given by the output of the trust region layer.
Since the old policy distribution is ﬁxed, all distances or divergences used in this paper can be decomposed into a mean and a covariance dependent part, we also use separate trust regions and bounds for mean and covariance, as this gives the algorithm more ﬂexibility.
The trust region layers aim to project π θ into the trust region by ﬁnding parameters ˜ µ and ˜Σ that are closest to the original parameters µ and Σ while satisfying the trust region constraints.
The projection is based on the same distance or divergence which was used to deﬁne the respective trust region.
Formally, this corresponds to the following optimization problems for each s
[Equation]
and
[Equation]
where ˜ µ s and ˜Σ s are the optimization variables for state s . Here, d mean is the mean dependent part and d cov is the covariance dependent part of the employed distance or divergence.
For brevity of notation we will neglect all dependencies on the state in the following.
We denote the projected policy as ˜ π
( a | s ) =
N ( a | ˜ µ, ˜Σ) .
4. P ROJECTION OF THE M EAN
For all three trust region objectives we make use of the same distance measure for the mean, i. e., the Mahalanobis distance.
Hence, the optimization problem for the mean is given by
[Equation]
By making use of the method of Lagrangian multipliers (see Appendix B.2)
[Equation]
we can formulate the dual and solve it for the projected mean ˜ µ as
[Equation]
This equation can directly be used as mean for the Gaussian policy, while it easily allows to compute gradients.
Note, that for the mean part of the KL we would need to use the Σ − 1 instead of Σ − 1 old in the objective of Equation 5.
Yet, this objective still results in a valid trust region problem which is much easier to optimize.
4. P ROJECTION OF THE C OVARIANCE
Frobenius Projection.
The Frobenius projection formalizes the trust region for the covariance with the squared Frobenius norm of the matrix difference, which yields
[Equation]
We again use the method of Lagrangian multipliers (see Appendix B.3) and obtain the covariance ˜Σ as
[Equation]
) where η is the corresponding Lagrangian multiplier.
Wasserstein Projection.
Deriving the Wasserstein projection follows the same procedure.
We obtain the following optimization problem
[Equation]
where I is the identity matrix.
A closed form solution to this optimization problem can be found by using the methods outlined in Takatsu (2011). However, we found the resulting solution for the projected covariance matrices to be numerically unstable.
Therefore, we made the simplifying assumption that both the current Σ and the old covariance Σ old commute with ˜ Σ .
Further, under the common assumption of diagonal covariances, the commutativity assumption holds.
For the more general case of arbitrary covariance matrices, we would need to ensure the matrices are sufﬁciently close together, which is effectively ensured by Equation 8.
Again, we introduce Lagrange multipliers and solve the dual problem to obtain the optimal primal and dual variables (see Appendix B.4).
Note however, that here we chose the square root of the covariance matrix 1 as primal variable.
The corresponding projection for the square root covariance ˜Σ 1 / 2 is then
[Equation]
where η is the corresponding Lagrangian multiplier.
We see the same pattern emerging as for the
Frobenius projection.
The chosen similarity measure directly reappears in the expression for the Lagrangian multipliers and the primal variables are weighted averages of the corresponding parameters of the old policy and the predicted one. [Figure] n t r [Figure]
FrobeniusWassersteinKL
KL Projection.
Identically to the previous two projections, we reformulate Equation 4 as
[Equation]
where dim ( a ) is the dimensionality of the action space.
It is impossible to obtain a fully closed form solution for this problem.
However, following Abdolmaleki et al.
, we can obtain the precision ˜Λ = ˜Σ − 1 by interpolation between the precision matrices of the old policy π old and the current policy π
[Equation]
where η is the corresponding Lagrangian multiplier and g ( η ) the dual function.
While this dual cannot be solved in closed form, an efﬁcient solution using a standard numerical optimizer, such as BFGS, is still possible, as it is a 1-D convex optimization.
Yet, we want to propagate gradients through our projection and thus also need to propagate them back through the numerical optimization.
To this end, we follow Amos & Kolter (2017) and compute those gradients by taking the differentials of the KKT conditions of the dual.
We refer to Appendix B.5 for more details and derivations.
Entropy Control.
Previous works Akrour et al.
(2019) and Abdolmaleki et al.
(2015) have shown the beneﬁts of introducing an entropy constraint H (˜ π ) = β in addition to the trust region constraints. Such a constraint allows for more control over the exploration behavior of the policy.
In order to endow our algorithm with this improved exploration behavior, we make use of the results from Akrour et al.
and scale the standard deviation of the Gaussian distribution with a scalar factor exp { ( β − H ( π θ )) / dim ( a ) } , which can also be individually computed per state.
4. A NALYSIS OF THE P ROJECTIONS
It is instructive to compare the three projections.
The covariance update is an interpolation for all three projections, yet the quantities that are interpolated differ.
For the Frobenius projection we directly interpolate between the old and current covariances (Equation 7), for the W2 projection between their respective matrix square roots (Equation 9), and for the KL projection between their inverses (Equation 11).
In other words, each projection suggests which parametrization to use for the covariance matrix.
The different interpolations also have an interesting effect on the entropy of the resulting covariances which can be observed in Figure 1.
Further, we can prove the following theorem about the entropy of the projected distributions
Theorem 1 Let π θ and π θ old be Gaussian and η ≥ 0 , then for the entropy of the projected distribution H (˜ π ) it holds that H (˜ π ) ≥ minimum ( H ( π θ ) , H ( π θ old )) for the Frobenius (Equation 7) and the Wasserstein projection (Equation 9), as well as, H (˜ π ) ≤ maximum ( H ( π θ ) , H ( π θ old )) for the KL projection (Equation 11).
The proof is based on the multiplicative version of the Brunn-Minkowski inequality and can be found in Appendix B.1.
Intuitively, this implies that the Frobenius and Wasserstein projections [Table] act more aggressively , i. e., they rather yield a higher entropy, while the KL projection acts more conservatively , i. e., it rather yields a smaller entropy.
This could also explain why many KL based trust region methods lose entropy too quickly and converge prematurely.
Yet, by introducing an explicit entropy control, those effects can be mitigated.
4. S UCCESSIVE P OLICY U PDATES
The above projections can directly be implemented for training the current policy.
Note, however, that at each epoch i the policy π i predicted by the network before the projection layer does not respect the constraints and thus has to rely on calling this layer. Yet, the policy of the projection layer ˜ π i depends on the parameters of π i and the old policy network π i, old = ˜ π i − 1 .
This would result in an ever-growing stack of policy networks which would become increasingly costly to evaluate. In other words, ˜ π i is computed using all stored networks of π i , π i − 1 , . . .
, π 0 . We now discuss the parametrization of ˜ π via amortized optimization. We need to encode the information of the projection layer into the parameters θ of the next policy, i.e. ˜ π ( a | s ; θ ) =
p ◦ π θ ( a | s ) is a composition function in which p denotes the projection layer.
Basically, the output of π θ is ( µ, Σ) , and p computes (˜ µ, ˜Σ) according Equations 6, 7, 9, or 11.
Formally, we aim to ﬁnd a set of parameters θ ∗ =
arg min θ E s ∼ ρ π old [ d (˜ π ( ·| s ) , π θ ( ·| s ))] , where ρ π old is the state distribution of the old policy and d is the similarity measure used for the projection, such that we minimize the expected distance or divergence between the projection and the current policy prediction.
The most intuitive way to solve this problem is to use the existing samples for additional regression steps after the policy optimization.
Yet, this adds a computational overhead.
Therefore, we propose to concurrently optimize both objectives during training by penalizing the main objective, i. e.,
[Equation]
Note that the importance sampling ratio is computed based on Gaussian distribution generated by the trust region layer and not directly from the network output. Furthermore, the gradient for the regression penalty does not ﬂow through the projection, it is solely acting as supervised learning signal.
As appropriate similarity measures d for the penalty, we resort to the measures used in each projection.
For a detailed algorithmic view see Appendix A.
E XPERIMENTS
Mujoco Benchmarks We evaluate the performance of our trust region layers regarding sample complexity and ﬁnal reward in comparison to PAPI and PPO on the OpenAI gym benchmark suite (Brockman et al., 2016).
We explicitly did not include TRPO in the evaluation, as Engstrom et al.
(2020) showed that it can can achieve similar performance to PPO.
For our experiments, the PAPI projection and its conservative PPO version are executed in the setting sent to us by the author.
The hyperparameters for all three projections and PPO have been selected with Optuna (Akiba et al., 2019).
See Appendix D for a full listing of all hyperparameters.
We use a common set of hyperparameters for all environments but the Humanoid, for which we optimized separately.
Next to the standard PPO implementation with all code-level optimizations we further evaluate PPO-M, which only leverages the core PPO algorithm.
Our projections and PPO-M solely use the observation normalization, network architecture, and initialization from the original PPO implementation. All [Figure] [Figure]
[Figure]
FROBW2KLPAPIPPO -
[Figure] Samples d m ea n ( π , π o l d ) Ant-v2-W α =
α = 1 α =α = 8 α =α = 25 α =α = 75 α =
[Figure] [Figure] [Figure] algorithms parametrize the covariance as a non-contextual diagonal matrix.
We refer to the Frobenius projection as FROB , the Wasserstein projection as W2 , and the KL-projection as KL . Table 1 gives an overview of the ﬁnal performance and convergence speed on the Mujoco benchmarks, Figure 4 in the appendix displays the full learning curves.
After every epoch, we evaluate ﬁve episodes without applying exploration noise to obtain the return values.
However, note that we did not include the entropy projection for PAPI to provide a fair comparison to PPO.
The results show that our trust region layers are able to perform similarly or better than PPO and PAPI across all tasks.
While the performance on the Hopper-v2 is comparable, the projections signiﬁcantly outperform all baselines on the HalfCheetah-v2.
The KL projection even demonstrates the best performance on the remaining three environments. Besides that, the experiments present a relatively balanced performance between projections, PPO, and PAPI. The differences are more apparent when comparing the projections to PPO-M, which uses the same implementation details as our projections.
The asymptotic performance of PPO-M is on par for the Humanoid-v2, but it convergences much slower and is, moreover, noticeably weaker on the remaining tasks.
Consequently, the approximate trust region of
PPO alone is not sufﬁcient for good performance, only paired with certain implementation choices. Still, the original PPO cannot fully replace a mathematically sound trust region as ours, although it does not exhibit a strong performance difference. For this, Figure 2 visualizes the mean KL divergence at the end of each epoch for all methods.
Despite the fact that neither W2 nor Frobenius projection use the KL, we leverage it here as a unifying measure to compare the change in the policy distribution.
All projections are characterized by an almost constant change between each epoch, whereas for PPO-M the changes are highly inconsistent.
The code-level optimizations of PPO can mitigate this to some extend but cannot properly enforce the desired constant change in the policy distribution.
We found especially the learning rate decay as major factor that contributes well to the rather good behavior of PPO.
Albeit, PAPI provides a similar principled trust region projection as we do, it still has some inconsistency by approaching the bound iteratively.
Entropy Control To demonstrate the effect of combining our projections with entropy control, as described in Section 4.2, we evaluate all Mujoco tasks again for this extended setting.
The target entropy in each iteration i is computed by exponentially decaying the initial entropy H 0 to κ with temperature τ as κ + ( H 0 − κ ) τ 10 iN , where N is the total number of training steps.
The bottom of Table 1 shows the results for our projections with entropy control.
Especially on the more complex tasks, all three projections signiﬁcantly beneﬁt from the entropy control.
Their asymptotic performance for the HalfCheetah-v2, Ant-v2, and Humanoid-v2 increases and yields a much faster convergence in the latter.
For the other Mujoco tasks the performance remains largely constant as the complexity of these tasks is insufﬁcient to beneﬁt from the explicit entropy control, as also noted by Pajarinen et al. ; Abdolmaleki et al. Entropy control is particularly helpful in environments which require more exploration. Contextual Covariances.
To emphasize the advantage of state-wise trust regions we consider the case of policies with state-dependent covariances.
Existing methods, such as PPO and TRPO, are rarely used in this setting. In addition, PAPI cannot project the covariance in the contextual case. Further, Andrychowicz et al. (2020) demonstrated that for the standard Mujoco benchmarks, contextual policies are not beneﬁcial in an on-policy setting.
Therefore, we evaluate on a task motivated from optimal control which beneﬁts from a contextual covariance.
We extend the Mujoco Reacherv2 to a 5-link planar robot, the distance penalty to the goal is now only provided in the last time step, t = 200 , and the observation space now contains the current time step t .
This semi-sparse reward speciﬁcation imposes a signiﬁcantly harder exploration problem as the agent is only provided with feedback at the last time step.
We again tuned hyperparameters using Optuna Akiba et al.
and did not use entropy control.
The results are presented in Figure 2 (right), comparing the approaches with and without contextual covariances, except for PAPI which is not applicable in the contextual case.
All three projections signiﬁcantly outperform the baseline methods with the non-contextual covariance.
Additionally, both the W2 and KL projection improve their results in the contextual case.
In contrast, all baseline methods decrease in performance and are not able to leverage the advantage of contextual information.
The poor performance for the baselines mainly originates from incorrect exploitation.
PPO reduces the covariance too quickly, whereas PAPI reduces it too slowly, which leads for both to a suboptimal performance.
However, the Frobenius projection also does not beneﬁt from the contextual policy because of numerical instabilities, arising from small covariance values towards convergence.
Those issues can be mitigated using a smaller covariance bound, but they cannot be fully avoided.
The KL projection, while yielding the best results throughout all experiments, is the computationally most expensive.
Hence, if contextual covariances are not required, the KL is the best choice for most problems.
As a bonus, it has all properties of existing KL-based trust region methods that have monotonic improvement guarantees.
Nevertheless, for quick benchmarks or contextual covariances, the W2 is preferred, given it is slightly less prone to hyperparameter choices and does not add any computational overhead as the KL does.
Trust Region Regression Loss.
Lastly, we investigate the main approximation of our approach, the trust region regression loss (Equation 12).
In the following ablation, we evaluate how different choices of the regression weight α affect the constraint satisfaction.
Figure 2 (center) shows the Mahalanobis distance between the current policy mean at the end of each epoch without the projection and the old policy mean for different α values.
In addition, for one run we remove the regression penalty from the loss but instead run several iterations of regressing the current policy towards the projection after each policy update.
One key observation is that decreasing the penalty up to a certain threshold leads to larger changes in the policy and pushes the mean closer to its maximum bound.
Intuitively, this can be explained by the construction of the bound.
As the penalty is only added to the loss, when the bound is violated, large changes in the policy are punished while making small steps might slightly lower the achieved reward, but does not incur any other negative effect.
By selecting a larger α , this behavior is ampliﬁed.
Furthermore, we can see that small values of α yield a behavior which is similar to the full regression setting.
Consequently, it is justiﬁed to use a computationally simpler penalty, instead of performing a full regression after each policy update.
D ISCUSSION AND F UTURE W ORK
In this work we proposed differentiable projection layers to enforce trust region constraints for
Gaussian policies in deep reinforcement learning.
While being more stable than existing methods, they also offer the beneﬁt of imposing the constraints on a state level.
Unlike previous approaches that only constrain the expected change between successive policies and for which monotonic improvement guarantees thus only hold approximately, we can constrain the maximum change.
Our results illustrate that trust regions are an effective tool in policy search for a wide range of different similarity measures.
Apart from the commonly used reverse KL, we also leverage the Wasserstein distance and Frobenius norm.
We demonstrated the subtle but important differences between those three different types of trust regions and showed our benchmark performance is on par or better than existing methods that use more code-level optimizations.
For future work we plan to continue our work with complex environments, in particular with contextual covariances.
Additionally, more sophisticated heuristics or learning methods could be used to adapt the trust region bounds for better performance.
Lastly, we are interested in using our trust region layers for other deep reinforcement learning approaches, such as actor-critic methods.
B D ERIVATIONS B. P ROOF OF T HEOREM
This section provides a proof for Theorem 1.
We mainly used the multiplicative version of the
Brunn-Minkowski inequality
[Equation] where Σ 1 , Σ 2 are p.s.d, α, β are positive, and α + β = 1 .
Frobenius Projection
[Equation]
Wasserstein Projection Let k denote the dimensionality of the distributions under consideration
[Equation]
≥ 0 .
[Equation] [Equation] + log (cid:12)(cid:12)
[Equation]
0 .(cid:12)(cid:12) 1 η + (cid:12)(cid:12) Σ
[Table]
old (cid:12)(cid:12)
ηη +
[Table]
+ log (cid:12)(cid:12) Σ
[Table]
(cid:12)(cid:12) 1 η + (cid:12)(cid:12)
[Equation]
[Table]
[Equation] (cid:12)(cid:12) [Equation]
KL Projection
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
(cid:16)(cid:12)(cid:12) (2 πe Σ) − [Equation]
[Equation]
the fact that: det
[Equation]
B. M EAN P ROJECTION
First, we consider only the mean objective
[Equation]
which give us the following dual
[Equation]
Differentiating w.r
[Equation]
t. ˜ µ yields
[Equation]
Setting the derivative to 0 and solving for ˜ µ gives
[Equation]
Inserting the optimal mean ˜ µ ∗ in Equation 13 results in
[Equation]
Thus, differentiating w.r.t ω yields
[Equation]
Now solving ∂ L ( ω ) ∂ω !
=
for ω , we arrive at
[Equation]
B. F ROBENIUS C OVARIANCE P ROJECTION
We consider the following objective for the covariance part
[Equation]
with the corresponding Lagrangian
[Equation]
[Equation]
Differentiating w.r.t.
˜Σ yields
[Equation]
We can again solve for ˜Σ by setting the derivative to 0 ,
i
[Equation]
e.,
[Equation]
Inserting ˜Σ ∗ into Equation 14 yields the dual function
[Equation]
Differentiating w.r.t.
η results in
[Equation]
Hence, ∂ L ( η ) ∂η !
=
yields
[Equation]
B. W ASSERSTEIN C OVARIANCE P ROJECTION
As described in the main text, the Gaussian distributions to have been rescaled by Σ − 1 old to measure the distance in the metric space that is deﬁned by the variance of the data.
For notational simplicity, we show the derivation of the covariance projection only for the unscaled scenario.
The scaled version can be obtained by a simple redeﬁnition of the covariance matrices.
For our covariance projection we are interested in solving the following optimization problem
[Equation]
which leads to the following Lagrangian function
[Equation]
Assuming that ˜Σ commutes with Σ as well as Σ old
[Equation]
Equation 15 simpliﬁes to
[Equation]
where S is the unique positive semi-deﬁnite root of the positive semi-deﬁnite matrix ˜Σ , i.e.
S = ˜Σ 1 / 2 .
Instead of optimizing the objective w.r.t ˜Σ , we optimize w.r.t S in order, which greatly simpliﬁes the calculation.
That is, we solve
[Equation]
for S , which leads us to
[Equation]
Inserting this into Equation 16 yields the dual function
[Equation]
The derivative of the dual w.r
[Equation]
t. η is given by
[Equation]
Now solving ∂ L ( η ) ∂η !
=
for η [Equation] we arrive at
[Equation]
B. KL-D IVERGENCE P ROJECTION
We derive the KL-Divergence projection in its general form, i [Equation] e., simultaneous projection of mean and covariance under an additional entropy constraint
[Equation]
Instead of working with this minimization problem we consider the equivalent maximization problem
[Equation]
which is similar to the one considered in Model Based Relative Entropy Strochastic Search (MORE) (Abdolmaleki et al., 2015), with a few distinctions.
To see those distinctions let η and ω denote
the
Lagrangian multipliers corresponding to the KL and entropy constraint respectively and consider the Lagrangian corresponding to the optimization problem in Equation
[Equation]
Opposed to Abdolmaleki et al.
we are not working with an unknown reward but using the log density of the target distribution π instead.
Thus we do not need to ﬁt a surrogate and can directly read off the parameters of the squared reward.
They are given by the natural parameters of π , i.e, Λ = Σ − 1 and q = Σ − 1 µ .
Additionally, we need to add a constant 1 to ω to account for the additional entropy term in the original objective, similar to (Arenz et al., 2018).
Following the derivations from Abdolmaleki et al.
and
Arenz et al.
we can obtain a closed form solution for the natural parameters of ˜ π , given the Lagrangian multipliers η and ω
[Equation]
and
[Equation]
To obtain the optimal Lagrangian multipliers we can solve the following convex dual function using gradient descent
[Equation]
Given the optimal Lagrangian multipliers, η ∗ and ω ∗ we obtain the parameters of the optimal distribution ˜ π ∗ using Equation 18.
Forward Pass For the forward pass we compute the natural parameters of π , solve the optimization problem and compute mean and covariance of ˜ π ∗ from the optimal natural parameters.
The corresponding compute graph is given in Figure 3.
Backward Pass Given the computational graph in Figure 3 gradients can be propagated back though the layer using standard back-propagation.
All gradients for the analytical computations (black arrows in Figure 3) are straight forward and can be found in (Petersen & Pedersen, 2012).
For the gradients of the numerical optimization of the dual (red arrows in Figure 3) we follow Amos & Kolter (2017) and differentiate the KKT conditions around the optimal Lagrangian multipliers computed during the forward pass.
The KKT Conditions of the dual are given by
[Equation]
[-η ∗ ) = 
 and m 2 ( − ω ∗ ) = 
-][Figure]
here m = ( m 1 , m 2 ) T denotes the Lagrangian multipliers for the box constraints of the dual ( η and ω need to be non-negative).
Taking the differentials of those conditions yields the equation system
[Equation]
which is (analytically) solved to obtain the desired partial derivatives
[Equation]
Implementation We implemented the whole layer using C++, Armadillo, and OpenMP for parallelization.
The implementation saves all necessary quantities for the backward pass and thus a numerical optimization is only necessary during the forward pass.
Before we perform a numerical optimization we check whether it is actually necessary.
If the target distribution ˜ π is within the trust region we immediately can set ˜ π ∗ =
π θ , i.e., the forward and backward pass become the identity mapping.
This check yield signiﬁcant speed-ups, especially in early iterations, if the target is still close to the old distribution.
If the projection is necessary we use the L-BFGS to optimize the 2D convex dual, which is still fast.
For example, for a 17 -dimensional action space and a batch size of 512 , such as in the Humanoid-v2 experiments, the layer takes roughly 170 ms for the forward pass and 3 .
5 ms for the backward pass if the all 512 Gaussians are actually projected 2 .
If none of the
Gaussians needs to be projected its less than 1 ms for forward and backward pass.
Simpliﬁcations If only diagonal covariances are considered the implementation simpliﬁes signiﬁcantly, as computationally heavy operations (matrix inversions and cholesky decompositions) simplify to pointwise operations (divisions and square roots).
If only the covariance part of the KL is projected, we set µ old = µ = ˜ µ ∗ and dµ = 0 which is again a simpliﬁcation for both the derivations and implementation.
If an entropy equality constraint, instead of an inequality constraint, it is sufﬁcient to remove the ω > 0 constraint in the dual optimization.
C A DDITIONAL R ESULTS
Figure 4 shows the training curves for all Mujoco environments with a 95% conﬁdence interval.
Besides the projections we also show the performance for PAPI and PPO.
In Figure 5 the projections also leverage the Entropy control based on the results from from Akrour et al.
[Figure]
[Figure]
D H YPERPARAMETERS
Tables 2 and 3 show the hyperparameters used for the experiments in Table 1.
Target entropy, temperature, and entropy equality are only required when the entropy projection is included in the layer, otherwise those values are ignored
[Table] [Table]
[Table] [Table]
[Table] [Table]
penalty
[Table]
n
[Table]
normalized observations
[Table]
normalized rewards
[Table]
True
[Table]
observation clip n.a. n.a. reward clip n.a. n.a. vf clip n.a. n.a.
[Table]
clip n
[Table]
contextual
[Table]
hidden layers [64, 64] hidden layers vf [64, 64] hidden activation tanh [Table] [Table] context)
n
[Table]
a.
[Table]
n
[Table]
n
[Table]
[Table]
n
[Table]
n
[Table]
a.
[Table]
n
[Table]
a
[Table]
False n.a.
[Table]
adam
[Table]
[Table]
1e
[Table]
[Table]
[Table]
1e
[Table]
minibatch
[Table]
entropy
[Table]
penalty max entropy weight regression weight 8.
n.a. normalized observations
True normalized rewards
False
True
False observation clip n.a. n.a. reward clip n.a. n.a. vf clip n.a. n.a. importance ratio clip n.a. contextual std
False hidden layers [64, 64] hidden layers vf [64, 64] hidden activation tanh