Variational Reparametrized Policy Learning with Differentiable Physics
Authors: Zhiao Huang
Litian Liang
Zhan Ling
Xuanlin Li
Chuang Gan
Hao Su
Venue: Submitted to ICLR 2023 (2023)
Keywords: Differentiable Physics Reinforcement Learning

Abstract: We study the problem of policy parameterization for reinforcement learning (RL) with high-dimensional continuous action space. Our goal is to find a good way to parameterize the policy of continuous RL as a multi-modality distribution. To this end, we propose to treat the continuous RL policy as a generative model over the distribution of optimal trajectories. We use a diffusion process-like strategy to model the policy and derive a novel variational bound which is the optimization objective to learn the policy. To maximize the objective by gradient descent, we introduce the Reparameterized Policy Gradient Theorem. This theorem elegantly connects classical method REINFORCE and trajectory return optimization for computing the gradient of a policy. Moreover, our method enjoys strong exploration ability due to the multi-modality policy parameterization; notably, when a strong differentiable world model presents, our method also enjoys the fast convergence speed of trajectory optimization. We evaluate our method on numerical problems and manipulation tasks within a differentiable simulator. Qualitative results show its ability to capture the multi-modality distribution of optimal trajectories, and quantitative results show that it can avoid local optima and outperforms baseline approaches.

Under review as a conference paper at ICLR 2023
VARIATIONAL REPARAMETRIZED POLICY LEARNING WITH DIFFERENTIABLE PHYSICS
Anonymous authors Paper under double-blind review
ABSTRACT
We study the problem of policy parameterization for reinforcement learning (RL) with high-dimensional continuous action space. Our goal is to find a good way parameterize the policy of continuous RL multi-modality distribution. To this end, we propose to treat the continuous RL policy as a generative model over the distribution of optimal trajectories. We use diffusion process-like strategy model the policy and derive a novel variational bound which is the optimization objective to learn the policy. To maximize the objective by gradient descent, we introduce the Reparameterized Policy Gradient Theorem. This theorem elegantly connects classical method REINFORCE and trajectory return optimization for computing the gradient of policy. Moreover, method enjoys strong exploration ability due to the multi-modality policy parameterization; notably, when a strong differentiable world model presents, method also enjoys the fast convergence speed trajectory optimization. We evaluate our method on numerical problems and manipulation tasks within differentiable simulator. Qualitative results show its ability to capture the multi-modality distribution of optimal trajectories, and quantitative results show that can avoid local optima and outperforms bascline approaches.
INTRODUCTION
Reinforcement learning (RL) with high-ddiensional continuous action space notoriously hard despite its fundamental importance for many application problems such robotic manipulation (Ope- nAI al. 2019: Mu et al. (02T). Compared with the discrete action space counterpart setup, is much trickier to represent policies for continuous R the optimality theory of RL, the function class of discrete RL policies is simply categorical distributions (Sutton & Barto 2018), while the function class of continuous RL policies has include density functions of arbitrary probabilistic distributions.
In practice, popular frameworks (Silver al. 2014 Haarnoja et al. 2018; Schulman et al. 2017) of deep RL formulate the continuous policy as neural network that outputs a single-modality density function over the action space (e.g., a Gaussian distribution however, breaks the promise of RL being a global optimizer of the return function because the single- modality policy parameterization introduces local minima that are hard to escape using gradients w.r.t. distribution parameters. Besides, single-modality policy will significantly weaken the exploration ability of RL. algorithms because the sampled actions are usually concentrated around the modality. Our Bandit examples show how single-modality RL policy fails to solve simple continuous action bandit problem. Therefore, practice, continuous RL often requires meticulous reward design that
actions). This formulation,
takes considerable human effort (Mu et al. 2021 Savva et al. 2019 Yu et al. 2020 Makoviychuk
ct al. 2021 Zhu et al. 2020).
In this paper, propose principled framework to learn the continuous RL policy multi- modality density function. We provide a holistic solution to two closely-relatee questions: 1) how to parameterize the continuous RL policy? 2) how update the policy parameterized as in (1)?
Our parameterization the continuous RL policy is based on two ideas. First, we take a sequence modeling perspective (Chen al.) 202I)) and view the policy trajectory space (instead of the action space) (Ziebart (0100) Levine 2018). Under this sequence modeling perspective, we can sample population of trajectories that cover multiple modalities of
density function over the entire

---

Under review as a conference paper at ICLR 2023
trajectories, which allows us to coneurrently explore distant regions in the solution space. Second, we use a generative model to parameterize the multi-modalin policies, inspired by their success in modeling highly complicated distributions such natural images (Goodfellow et al. 2016: Zhu
let . 2017 RRombach et al. 2022 Ramesh et al. 202T). We introduce a sequence of latent variables 2, and learn a decoder that "reparameterizes" the random distribution into the multi-modality trajectory distribution (Kingma & Welling 2013). from which we can sample trajectories т. Our
policy network in a spirit, akin to diffusion models (Rombach et al.) 2022] Ramesh al. 202T] by learning to model the joint distribution p(2. r). As prototypical work, we prefer simple design and choose not to include the multi-step diffusion process as in image modeling.
Our choice of policy parameterization leads us to adopt the variational method (Kingma & Welling 2013; Haarnoja et al., 2018, Moon 1996) to derive an on-policy learning algorithm. Classical
on-policy learning leverages the policy gradient theorem (Sutton & Barto 2018). i.e., VJ((T) = E,, (R(T) V log p())] Because we model p(z, 7), that requires computing p(T) and its gradient with [ p(2, r) dz. However, marginalizing out is often intractable when is continuous, and it is well-known that optimizing the marginal distribution log p(T) by gradient descent suffers from local optimality issues (e.g., using gradient descent to optimize Gaussian mixture models which have latent variables is not effective, so EM is often used instead Ng (2000). To overcome these obstacles, we take a different route and adopt variational method (maximum entropy RL) that directly optimizes the joint distribution of the optimal policy without hassels of integrating over
We derive a novel variational bound which is the optimization objective for policy learning. To maximize this objective, introduce the Reparameterized Policy Gradient Theorem. The theorem states a principled way to compute the policy gradient by combining the eeward-weighte (as in the classical policy gradient theorem) and the path- wise gradient from a differentiable world model (as in classical trajectory optimization methods). The two sources of gradient complement each other the reward- weighted gradient improves the likelihood of selecting trajectories from regions containing high-reward ones, whereas the path-wise gradient suggests updates to make local improvements to trajectories. This combination allows us to enjoy the precise gradient computation from differentiable world models (Werling et al. [021 Hu et al. 2019 Huang et al. 2021 and also maintains the flexibility to sample and optimize the trajectory distribution globally. Note that this beautiful result is also consequence of introducing the latent variable 2.
gradient
The effectiveness of our method is rooted in our choice of generative policy modeling method. An ideal method needs to be powerful in modeling multi-modality distribution, and it needs to support sampling, density value computation, and stable gradient computation. Although there are other candidates for deep generative models, they often have limitations to being used for continuous policy modeling. For example, GAN-like generative models (Goodfellow et al. (2020) can only sample but not compute the density value. While normalizing flow (Rezende & Mohamed 2015) can compute the density value, they might not be as robust numerically due to the dependency on the determinant of the network Jacobian; moreover, normalizing flows must apply continuous transformations onto continuously connected distribution, making it difficult to model disconnected modes (Rasul et al. 2021
We apply our methods to several numerical problems and three trajectory optimization tasks to manip- ulate rigid-body objects and deformable objects supported by differentiable physics engines (Werling et al. 2021 Huang et al., 2020). These environments contain various local optima that challenge single-modan policies (either sewarddweighted policy gradients as in REINFORCE or path-wise gradients as in trajectory optimization). In contrast, our approach benefits from modeling the trajecto- ries as generative models and the ability to sample the trajectory space. In qualitative experiments, we observe a strong pattern of modai-ty distribution by visualizing the policy after the learning converges. By quantitative evaluation, the policy learned by our framework does not suffer from the local optimality issue and significantly outperforms baselines.
2 VARIATIONAL REPARAMETERIZED POLICY LEARNING
In this section, we introduce our basic framework. All notations follow the convention of the commu- nity. To be more clear, we leave a section that introduces background knowledge in Appendix AAaad mathematical table for notations in Appendix B
2

---

Under review as a conference paper at ICLR 2023
Decoder qu (z
Stari point
Encoder:pp (z]r)
Reward
Hidden space
Trajectory Distribution
Figure 1: a) the reward landscape where the agent needs to move from the red dots to the region containing a high reward: b) latent space modeled by random Gaussian; c) the state density of sampled trajectory from go (2,2 In (c), each red dot corresponds to a state in the trajectory. Our policy can be viewed as encoding the stochastic latent variable into the trajectory distribution through the decoder q=(z, r). We rely on encoder Pø(=/r) to ensure the cycle consistency between the latent z and the sampled trajectory T.
2.1 GENERATIVE MODELING OF OPTIMAL. TRAJECTORIES
Following (odorov 2006 2008 Toussaint 2009 Ziebart 2010 Kappen et al. 2012: Levine 2018
Haarnoja et al. 2018). Variational Reparameterized Policy (VRP) framework, shown in Fig.
views policy optimization as learning a generative model that generates optimal trajectories. To capture the multimodaliiies of the optimal trajectories, introduce latent distribution and construct a decoder-encoder structure similar to recent deep generative models (Kingma & Welling 2013 Ho
al. 2020 Rombach et al.) 2022; Ramesh al., 202T). The "encoder" maps trajectories to latent samples, while the decoder is neural network that "reparameterizes" random samples from the latent distribution into different trajectories.
Method Overview. As introduced in Sec. 2.2 the latent distribution and the reparameterizztion function together form policy that generate diverse trajectories. We use novel variational bound to approximate the posterior of optimal trajectories in Sec. 2.3 as the optimization objective. The variational bound naturally combines maximum entropy RL while also containing term to enforce cycle consistency (Zhu et al., 2017) between the encoder and the decoder, preventing the policy from mode collapse. We then introduce Reparameterized Policy Gradient Theorem in Sec. 2.4 to optimize our reparameterized policy to improve the variational bound. We find that the derived policy gradient includes two terms: one optimizes the latent distribution through reward-weiie weighted gradient as the classical policy gradient theorem, and another optimizes the reparameterization network, which can benefit from the path-wise gradient from differentiable world model. This helps our inference algorithm to enjoy the efficiency of differentiable physics while being able to sample globally when the reward landscape is discontinuous or non-convex.
2.2 TRAJECTORY GENERATION WITH REPARAMETERIZED POLICY
Let E 2 be latent variable, which can be either continuous categorical to model optimal trajectories. We design our "policy" as a joint distribution qo (z, r) of latent z and the trajectory т. To model sequential trajectory,
can consider the following two factorizations: 1. sample before т: 46 (2, 7) p()1)g((=(81) II p( a, (a. 8x): 2. sample with qo(2,
p(si (00 (81) I[= p( (0 11++1/90 18,, az) (x)ga(() 2, )oo (z) 12, 1) The latter allows us to model the hybrid policy in the option-critic (Bacon al. 2017). where
treat 2, the option per step.
Though sampling from may require us to sample т, 2 together, it is still convenient to view as the combination of two marginal distributions: (1) the policy for the latent representation po(z (iss (p(())())) ). and (2) a decoder pe Pe(T]a 22, 81) that reparameterizes z into real trajectory. Note that we will use pe refer to various marginal distribution qe (2, 7). ) For example, Po(T) Ja ge d2 will be the marginal distribution of trajectories sampled from (2(,,)).

---

Under review as a conference paper at ICLR 2023
The term "reparameterization" is inspired by the reparameterization trick in (Kingma & Welling 2013). When is Gaussian noise Ç. we can reparameterize the noise into actions by a = ((rr) and optimize it with gradients directly. The reparameterized policy go be viewed as instance of the stochastic computation graph (Schulman et al. 2015 Weber et al. 2019). In practice, go qola [0ls,22 2) are modeled by neural networks as in common policy learning frameworks (Haarnoja
ct al. 2018
2.3 THE VARIATIONAL LOWER BOUND FOR REPARAMETERIZED POLICIES
The Auxiliary Trajectory Encoder Intuitively, learning policy networks by inputting random distribution can help generate diverse examples, but it will also suffer from the issues of mode collapse Li et al. (2017). To build a connection between the trajectory and T and enforce a cycle consistency, we introduce auxiliary distribution (z(2))) which can look at the whole trajectories and map r back into the latent 2. We will illustrate how the auxiliary encoder naturally emerges from the variational lower bounds following (Rezende & Mohamed 2013 Levine 2018).
The Evidence Lower Bound With the help of the auxiliary encoder pa we can now define a joint distribution of optimality O, latent 2, and the trajectory T as (o 2, p(O where O and are independent conditioning on т. Treating (2, T) as the posterior approximation, we can write the Evidence Lower Bound (ELBO) for the optimality. Note that the following inequality holds for arbitrary distribution Pa and go.
log (O)) = E2, [log pa (0, log (2, +
ELBO
E r.r
[logps z) log [)())]
Eg Trags
[logp(O,7) + Po log qo (2,
Eg g
log log
log Poglo()))) log go T)
(I)
reward
prior
cross entragy
eatropy
The ELBO contains four parts that can all be computed directly given the sampled 2 and (the environment probability p(5t+1|s),0, az) is canceled as in (Levine) 2018)). The first two parts are R(T)/T++ c, where is the normalizing constant that can be ignored during optimization, and a prior distribution p(r), which is assumed to be known. The third part is the log-likelihood of based on encoder. It easy to see that if we fix qa, and maximize Ps alone will minimize the cross-entropy Ex,rogp log (2([T)), similarly to the supervised learning. 90(2.7) T. "d(
the predefined reward log p(O[r)=
It achieves optimality when Po(=/T) (=
modeling the posterior of = for sampled from qo. On the other hand, fixing Φ, the decoder qe is encouraged to generate trajectories that are easy to identify or classify; this helps to increase diversity and enforce a cycle consistent to avoid mode collapse. The fourth part is the policy entropy that enables maximum entropy exploration.
Maximizing all terms together for the parameters and 6 will minimize
Dkz (2, () r(O)) Dx (g0(2, ().(2.))))ll(()))) ())po)(()(((()(()1))))
where optimality can be achieved when p=(2)7) = (2/T) and pa (T))p((2))) po (2) + (T(
pendix (o( z)dz P(r(O). We discuss the method's connection with other methods Ap-
2.4 REPARAMETERIZED POLICY GRADIENT WITH DIFFERENTIABLE PHYSICS
Given 0, o we can treat the ELBO as our reward (T = log p(O]=) + p(T) + po log (2 r). The maximization w.r.t. ø is straightforward. As for qo. depending its structure, can optimize it with various on-policy or off- off-policy RL. methods as in (Weber et al., 2019).
Here study special case where (T is differentiable to 0 through sampled trajectory Thus, we can optimize A by first-order path-wise gradients directly. This case is very interesting as analytical gradients provide fast convergence speed but also suffer from the local optima issue (Li
4

---

Under review as a conference paper at ICLR 2023
al. 2017). while our reparameterization policy naturally includes latent distribution for sampling and a parameterization part for optimization, providing more chances to search globally. One can obtain a path-wise gradient from a differentiable simulator (Hu et al. 2019; Werling et al. 2021), learned neural model (Liang et al. 2022).
Formally, we hope to find to maxi- mize the expectation ee (T)] = J2. (0(2, 7) Rebt (T) The sampling procedure is usually non-differentiable. But if we let 0r fs 2, () and the dynamics S(+1 h((,, az) be differentiable functions, the trajectory becomes diffecrentaable function of 0 almost everywhere.
Reparameterized Policy in Differentiable Environments.
We further assume that 81 is fixed and the functions fo, h are deterministic to simplify the derivation. One can remove these assumptions easily through reparameterization as in (Silver et al. 2014 Heess et al. 2015) to generalize to stochastic functions. As a result, we can write qo(2, Tg) = qo (o (2, 81 $1)), where and the density qo is deterministic function of the sampled 2 and 81- In this case, E| (e)o((To) ]] = J- go (z, Ta) Relbo (To )dz. For each sampled trajectory To. we can compute its path- wise gradient w.r.t. to the policy parameter ! as (TRabo(Te). Then we have the following theorem to compute the gradient of the expected rewards:
Theorem (Reparameterized Policy Gradient Theorem). Given 81 and under regularity conditions in Appendix F.1 For almost every 0, the expected reward EeRetb (TO)] exists and differentiable, and its gradient can be computed by
(eB(RenO(TT)) [ go (2, Retbo(T8) V go To) +
V Reelbo (Tg)
dz.
Reward weighted Gradient (REINFORCE)
Parkwise Gradient (Trajectory Opimiiztiion
We provide a short proof in Appendix
F.2
Combine Sampling and Optimization. Our theorem shares the same spirit as in Schulman et al. (2015) and is related to the problem of exchange derivative and expectation [L'Ecuyer 1995). Here we want to emphasize its use case in trajectory optimization. Theorem I splits gradients into two parts. The first part is a reward- reward-weighted weighted gradient that aims to improve the likelihood of the distribution (0(2[T) go (z. TD) for high-reward ones. If the environment is differentiable, the second part may optimize the policy directly through a path-wise gradient to make local improvements to trajectories. The natural combination of the sampling and optimization may allow our approach to enjoy the benefits of both: it can search over the whole space but can also leverage local structures for fast convergence speed. This can not be achieved without the latent variable 2.
Our methods provide the potential to avoid discontinuities in the reward landscape. If ge 2, Ta) does not depend s>1, the reward-we weighteed gradient can optimize the sampling distribution directly without relying on the path-wise gradients, providing the policy a chance to ignore the discontinuous point and move to the high reward region directly. Another appealing point in Theorem | is that in Assumption 2 we only require Relbo to be Lipschitz continuous for (o(2, To) 2 O. Thus, our policy gradient estimate can make unbiased estimation so long as the probability density of sampling a discontinuous trajectory is negligible. Even if discontinuous point generates biased gradient estimation for sampled z, the zeroth-order part still has chance correct its density directly by observing the reward directly so long as the optimal trajectories have certain properties.
By defining latent distributions, building the reparameterized policy and the auxiliary encoder, then optimizing the ELBO with the reparameterized policy gradient through differentiable world model, we can combine search and optimization to enjoy the generative modeling of the optimal trajectories. We summarize the whole recipe in Algorithm of of Appendix D and describe implementation details in Appendix D
'eere assume contains only continuous variables to simplify the derivation. Adding discrete variables straightforward by summing all possibilities.
5

---

Under review as a conference paper at ICLR 2023
3 RELATED WORK
Differentiable Simulation and Trajectory Optimization. Trajectory optimization (Kelly 2017) aims to find trajectory which optimizes the target metries under given constraints over the tra- jectory. One branch of methods optimizes the trajectory using the gradient of the optimization objective (Ratliff et al. 2009] Kalakrishnan et al., 2011; Schulman et al., 2014). Recently, the
development of differentiable simulation technique (Hu et al. 2019; Werling et al. [021] QQao et al. 2021] Freeman et al., 202T Howell et al. 2022) enables compute the gradient from the reward function and optimize the trajectory in physical simulators. Such methods are efficient around the optimal solution: however, the optimization often gets stuck local optima elsewhere, especially when the reward function is not smooth. In our work, one term in the formula of our RPG theorem corresponds to the gradient from trajectory optimization.
Policy as Sequential Generative Model. Maximum entropy reinforcement learning (Todorov)
2006; 2008 Toussaint 2009; Ziebart 2010 Kappen al. 2012) can be viewed as variational
inference in probabilistic graphical models (Levine] 2018). which models optimality as observed variable and trajectory as a latent variable. When the demonstration or a fixed dataset is provided on the offtine RL setting (Chen et al. 2021 Reed et al. (2022), , policy learning is simplified as sequence modeling task (Chen et al. 2021 ZZeng et al,; Reed et al., 2022). They use autoregressive models to learn the distribution of the whole trajectory, including actions, states, and rewards, and use the action prediction as policy. In our work, we learn a sequential generative model of policy for online RL via the variational method. The policy can model any distributions in the trajectory space.
Variational Skill Discovery Our method is closely related to the work in unsupervised reinforce- ment learning (Eysenbach et al. 20188] Achiam et al., (218)) or diverse skill learning (Kumar et al., 2020 Osa 2022). These methods share the same technique of using neural networks to approxi- mate the posterior of the latent variable given either states or state-action pairs and encourage the policy to reach states consistent with the latent variables. However, these methods do not model the optimal trajectory distribution but only aim to generate a diverse set of solutions by adding the mutual information term as a reward bonus, resulting in different formulations and effects. For example, they all choose to fix initial latent distributions without optimizing them, limiting their ability to achieve optimality. Moreover, Eysenbach et al. (2018); Achiam et al. (2018) does not optimize the learned skill for the environment rewards; Osa et al. (2022) does not optimize the mutual information along trajectories; Kumar et al. (2020) needs to solve the optimization problem first before finding a diverse set of solutions. In contrast, we jointly optimize the latent representation and the policy with a single objective, providing a simple but unified perspective for previous approaches in optimization problems.
Hierarchical Methods As mentioned in Sec. 2.2 the hierarchical methods, e.g., option critic (Bacon) et al. 2017), can be seen as a special case of our method when we use sequence of latent variables 2 (Z1, ........ .T) to reparameterize the policy. Without optimizing the latent variable through the variational inference, most hierarchical RL methods often need special designs for the latent space, e.g., state-based subgoals (Kulkarni et al. 2016 Nachum et al. 2018b a) or predefined skills (Li le al. 2020) to avoid mode-collapse. Osa et al. (2019) regularized options to maximize the mutual information between the action and the options, which are very relevant to ours. However, it does not model temporal structures as ours ensure consistency along the trajectories. Hierarchical imitation learning Gupta et al. (2019); Pertsch et al. 2) Shankar & Gupta (2020); Jaang et al. (2022): Lynch al. (2020); Fang et al. (2020) extract temporal abstractions using generative models from demonstrations. InfoGAN (Li et al. 2017) and ASE (Peng et al.] 2022) uses adversarial training Goodfellow et al. p2 : Ho Ermon (2016) to imitate demonstrations. These works all rely on demonstrations rather than rewards to learn abstractions. For example, (Co-Reyes et al. 2018) learns representation on the collected dataset with variational inference and then utilizes the trained model for planning policy learning. The separation of the representation learning and reward maximization makes it differ from our methods: first, it requires a state reconstruction module to supervise the generative model, which is challenging for digh-dimensional observations; second, it optimizes neither the latent distribution nor the actions for the reward directly, thus requires additional planning procedure during the execution find suitable actions.

---

Under review as a conference paper at ICLR 2023
4 EXPERIMENT
4.1 NUMERICAL EXPERIMENTS
4.1.1
ENVIRONMENT AND EVALUATION
In this section, we investigate the following questions: 1) Whether our method is more efficient for trajectory optimization than gradient-free algorithms; 2) Whether our method can search over the large solution space and avoid local optima that might trap a single-modality policy.
We design the following environments to answer these questions. Bandit and 2: Our bandit problems in Figure (1) and Figure (2) have a 1d action space and a non-convex reward landscape. We initialize our policy as Gaussian centralized around 0 with its scale barely touching the right mode of reward. The reward function of Bandit 2 contains an additional discontinuous point. Move 1 and 2: An agent moves in 2D environment for fixed number of steps and receives reward according to its terminal position. The terminal reward landscape contains 4 Gaussian peaks (deeper color higher reward), shown in Figure (3) and (4). The agent is initialized at the center. Note that the reward of Move near the initialization point is constant, so the agent receives a non-zero gradient only if it arrives The terminal reward shown in Figure 2a (5). There is one obstacle consisting of three white circles. When an agent runs into an obstacle, it will bounce off (continuous collision detection (Hu al.) 2019)) due to an impulse normal to the contact surface. Note that there is a small dent between the three circles, creating local optimum.
position closer to one of the four Gaussian peaks. Move
We evaluate our method against the single modality policies learned from the following base- lines: 1) Reinforce (Williams) 1992), and 2) policies optimized through path-wise gradients with Adam (Kingma & Ba 2014). We compare their sample efficiency and final performance after training finished.
4.1.2 EXPERIMENT RESULTS
We plot the average episode return of each algorithm against the number of samples in Figure 2b Error bars show the standard deviations over five runs with different random seeds. The results suggest that our methods achieve better sample efficiency and find better solutions than baseline algorithms.
Bandi:
Bandi:
Move
1.00
1:00
2,5
2,5
0:
80
0.25
0:25
Artio
action
(a) Environments for numerical experiments. We plot the reward landscapes in the first two bandit environments and show the rewards of the terminating positions as contour plots in the right three 2d move environments. These environments all pose some level of challenge in optimizing reward/value landscapes that are discontinuous and/or locally optimal.
Bir
Bandlk 2
Mov 1
Move
Hove
:3.0
2.5
1.0
2.5
0.6
0.2
0.00
0.
05
1.0
0.0
0.5
-2.0
Deinforce
Adom
Our Methos
P70
(b) Training curve. X-axis: number of sampled states. Y-axis: average return.
Error bars show standard
deviations over runs. Hyperparameters are available in Table
We first take Bandit 2 in Figure 2a (2), which has two modes and regions containing zero gradients in its reward landscape, in Figure (2). In terms of optimality, our method converges to the globally optimal solution the end and achieves much higher return compared with other methods in Figure (2). Two reasons
example for comparison. We show the performance of each algorithm
7

---

Under review as a conference paper at ICLR 2023
prevent our method from being stuck at the local optima. First, the entropy term in the Rebo objective encourages the agent to perform a maximum entropy exploration over the whole reward landscape, while the cross entropy term drives the agent to sample distinguishable actions for different latent variables, causing different action modes determined by latent variables to move away from each other, resulting in better state coverage. Second, after discovering different local optima, our policy is able to maintain a multi-modality action distribution to exploit those local optima independently. By adjusting the probabilities of different modes through optimizing rewards with latent variables, can eventually converge to the global optimal solution. In experiments, we observed that although our method quickly found the right mode of the reward function, shown in Figure 2a (2), which is locally optimal, it would still increase the probabilities for the under-explored region due to the exploration
terms in our objective. After it found the left mode, it would maintain two action modes and optimize them together until the expected return of the left mode was larger than the right one and eventually converged to the left, finding the global optimum.
A notable feature of the multi-modality is that during the whole optimization process, the expected reward increases monotonically as only needs to increase the reward for modes with higher rewards. In contrast, when Adam and Reinforce converge to the right mode, even if we have introduced the maximum entropy reward to encourage exploration, different from our method, their single modality property prevents them from jumping from one mode to better mode as there no way go across the flat region in between without decreasing the reward.
In terms of sample efficiency, our method benefits from the pathwise analytical gradient, consistently outperforming Reinforce, which has noisy gradient estimate that may reduce convergence speed when the number of samples is limited. Though Adam increases rewards faster than our method in the beginning stage, it quickly gets stuck at the local optima. Current RL algorithms like PPO or SAC are also constrained by the single modality gaussian parameterization of the policy, these methods can only find suboptimal solutions, even when more than enough samples are given.
Remaining environments Experiments on the remaining four environments show similar behavior pattern that our method can converge to the global optima that the baselines fail to find. We observe that the curve of ADAM in (1gure(bb(()) even decreases after reaching local optimum. This is due to the discontinuity on the right of the reward landscape as shown in Figure 21)(1). When Gaussian policy reaches the right local optimum, the analytical gradients nearing that point still only contain a positive value, pushing the policy to go across the optimal point and move into the right flat region, causing the return to decrease to -0.5 in the end. Reinforce and PPO does not suffer from this issue, as it only uses the reward-wa weighted gradient estimate, and can correctly identify the left low-reward region and move away from it. SAC also does not suffer from this issue, as the policy optimizes the Q function, which smooth the discontinuity to allow the policy to converge at the local maximum stably. But it is still very challenging to find the global optima with similar reason as in the Bandit 2 environment.
In the case of sequential decision-making problems, our method consistently outperforms algorithms that maintain single modality policy in several move environments by large margin and can find global optima as shown in Figure 2b (3)-(5). We observed that Reinforce went to the left side mode in Move and Move as at the starting state there are no reward signals pointing to the global optimum in the upper side of the state space. Its variance is also high due to the increased problem dimensions. Adam even did not make any progress in Move as there is no gradient around the initial state. The Move 3 environment includes physics-based contacts in the dynamic system, while our method still can solve it with the analytical gradients, in presence of the discontinuities during contacts.
4.1.3 How LATENT DISTRIBUTIONS AFFECT THE TRAJECTORY MODELING?
Modeling the trajectories with a sequential { (2,} instead of single latent variable sampled before sampling the trajectory 7 can help factorize the trajectory space to build more compact representation. We construct an environment containing four obstacles to demonstrate its effects in Figure The agent on the bottom left must avoid obstacles and reach the goal at the top right corner. We study the case where = is a discrete variable. We first plot the state and the action distributions for policies learned by sampling a single categorical & before the trajectory T in Figure 3a Different colors represent the different values of 2. We plot the state and action distributions and use different colors to represent their corresponding 2 values. Though a categorical distribution can represent various paths toward the goal, the number of latent samples constrained its expressive power However, as shown in
8

---

Under review as a conference paper at ICLR 2023
Figure 3b model the distribution with sequence of 21, 22k+,22K+11 22K+1:.... :.N++ This brings the compositionality that allows the agents to select various & sequences to generate trajectories that better cover the state space. The action distribution symmetry of the optimal trajectories and the action pattern moving along different directions, which helps the poliey better cover the whole state space compared with single categorical distribution.
the right successfully captures the
state
action
state
action
(a) Sample z before
(b) Sample sequence (=.}
Figure 3: A navigation task to demonstrate the effects of learning sequence of An agent needs to move from the bottom left to the top right and avoid the green obstacles. Both the state and action space are 2D. Different colors represent different values of 2,.
4.2 EVALUATION ON MANIPULATION TASKS WITH DIFFERENTIABLE PHYSICS
We evaluate our approach in several physical environments demonstrate its potential in trajectory optimization. Figure + compares our approaches against the gradient-based trajectory optimizer with a single Gaussian head in three environments. The Grasp environment, based Physics (Werling et al., 2021), contains three balls. The agent needs to control a 3 DoF gripper to grasp the green ball and pick it up. However, the agent only knows that it needs grasp an object, and it does not receive any reward before lifting it. We use reward to minimize the distance between the gripper and the nearest balls, which creates local minimum that encourages the gripper touch the nearest ball in front of it, as shown in Figure 4 The first-order gradient-basse based trajectory optimization got stuck. contrast, our method can capture the multimodality of the reward landscape, approach different balls, and exploit different options (touching different balls) simultancously. In the end, method successfully finds the solution to pick up the green ball. The Rope and the Cut environments soft body environments built with PlasticineLab (Huang et al. 2020), each containing 3DoF rigid body manipulator. In the Rope environment, an agent needs to push the left side of a rope forward. Similarly, in the Cut environment, manipulators are initialized in the middle of the objects, which is far away from suitable contact points required finish the tasks, leading to local optima (Li al. 2022) shown in Fig. 4 In contrast, our method can explore different contact points and find correct solutions in the end. These experiments show the potential of our approach in learning the optimal policies with differentiable physics.
Nimble
agent needs cut off the left end of rope. However,
5 CONCLUSION AND FUTURE WORK
We derive a framework that models the policy of continuous RL by a multi-modality distribution in the variational inference framework. Under this framework, we also derive Reparameterized Policy Gradient Theorem which enjoys the advantage of both classical sample-based methods (as REINFORCE) and trajectory optimization methods (which require the support of differentiable world models).
Our framework opens a new venue of continuous RI. has deep connections with diffusion model, decision sequence modeling, and differentiable physics techniques. In the future, are interested in exploring how sequence modeling techniques, such as transformers, can be used to model the policy in our framework. We also interested in testing method for more complicated tasks, such as dexterous object manipulation. Finally, our method can be extended to offline RL and off ofllcy RL setups. We leave all these exploration efforts to future works.

---

Under review as a conference paper at ICLR 2023
Gaussian Policy
Our Method
Grasp
Rope
Cut
Figure 4: Visualization of policy distribution for manipulation with differentiable physics. We qualitatively compare the Gaussian policy learned by ADAM and the by our method. On the left, we overlay final states from three trajectories of the Gaussian policy. On the right, we also draw three samples (trajectories) from the policy, and we visualize the final state of each trajectory in different plot. We that the three final states from the Gaussian policy similar, while the final states are very different for trajectories from our policy. Additionally, one of the modalities of our method can solve the problem (the rightmost column), while the Gaussian policies get stuck in local minima.
policy learned
REPRODUCIBLITY STATEMENT
We will provide open-source implementation of our method GitHub. We will also update more illustrative examples and demos to our website The hyperparameter used in the experiment section be found in Table 2 in the appendix. The theoretical results also in multiple sections of our appendix.
REFERENCES
Joshua Achiam, Harrison Edwards, Dario Amodei, and Pieter Abbeel. Variational option discovery algorithms. arXiv preprint 1807 2018
Rika Antonova, Jingyun Yang, Krishna Murthy Jatavallabhula, and Jeannette Bohg. Rethinking opti- mization with differentiable simulation from global perspective. arXiv preprint arXiv:2207.0016 2022.
Pierre-Luc Bacon, Jean Harb, and Doina Precup. The option-critic architecture. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540 2016.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084-15097, 2021.
steve cheng. differentiation under the integral sign. Diiferrentiaiionnndeerbb www.tthteitnegtttnntggmmm
URL. https://planetmetth/ org
John Co-Reyes, YuXuan Liu, Abhishek Gupta, Benjamin Eysenbach, Pieter Abbeel, and Sergey
Levine. Self-consistent trajectory autoencoder: Hierarchical reinforcement learning with trajectory embeddings. International conference on machine learning, pp. 1009-1018. PMLR, 2018.
A. P. Dempster, N. M. Laird, and D. B. Rubin. Maximum likelihood from incom- plete data via the algorithm. Journal of the Royal Statistical Society: Series
10

---

Under review as a conference paper at ICLR 2023
B (Methodologicall) 39(1):122 1977. doi: http://oooaar//11111111111110109 tb01600.x. URL htttt:: wttp://essss....snssseeeeeeee www..eeioarr..wwe...c.....l... dowat//oo//aaae abs wttp:////ccoaaaa/...1111... 1977 tb01600.2 1517:600..9997..1110000...
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. arXiv preprini arXiv: tetr:00002000 2018.
Kuan Fang, Yuke Zhu, Animesh Garg, Silvio Savarese, and Li Fei-Fei. Dynamics learning with cascaded variational inference for multi-step manipulation. Conference on Robot Learning, pp. 42-52. PMLR, 2020.
Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem. Brax-a differentiable physics engine for large scale rigid body simulation. arXiv preprint arXiv:2/06. 13281, 2021.
Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1. MIT Press, 2016.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, (3(11))139-144 2020.
Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. arXiv preprint arXiv:1910.11956 2019.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pp. 1861-1870. PMLR, 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv: 2019.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. arXiv preprint arXiv:2203. 040.2295...... 2022.
Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap. Tom Erez, and Yuval Tassa. Learn- ing continuous control policies by stochastic value gradients. Advances in neural information processing systems, 28, 2015.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in neural information processing systems, 29, 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6640-6851, 2020.
Taylor A Howell, Simon Le Cleac' h, Zico Kolter, Mac Schwager, and Zachary Manchester. Dojo: A differentiable simulator for robotics. arXiv preprint arXiv:2203.0mm 00806, 2022.
Yuanming Hu, Luke Anderson, Tzu-Mao Li, Qi Sun, Nathan Carr, Jonathan Ragan-Kellee and Frédo Durand. Difftaich:: Differentiable programming for physical simulation. arXiv preprint arXiv: 110000093 2019.
Zhiao Huang, Yuanming Hu, Tao Du Siyuan Zhou, Hao Su, Joshua B Tenenbaum, and Chuang Gan. Plasticinelab: A soft-body manipulation benchmark with differentiable physics. In International Conference on Learning Representations, 2020.
Zhiao Huang, Yuanming Hu, Tao Du, Siyuan Zhou, Hao Su, Joshua B Tenenbaum, and Chuang Gan. Plasticinelab: soft-body manipulation benchmark with differentiable physics. arXiv preprint 1.Xi::2104.03311. 2021.
Zhengyao Jiang, Tianjun Zhang, Michael Janner, Yueying Li, Tim Rocktäschel, Edward Grefen- stette, and Yuandong Tian. Efficient planning in a compact latent action space. arXiv preprint arXiv:2208. 10291, 2022.
11

---

Under review as a conference paper at ICLR 2023
Mrinal Kalakrishnan, Sachin Chitta, Evangelos Theodorou, Peter Pastor, and Stefan Schaal. Stomp: Stochastic trajectory optimization for motion planning. In 2011 IEEE international conference on robotics and automation, pp. 3569945744 IEEE, 2011.
Hilbert J Kappen, Vicenç Gómez, and Manfred Opper. Optimal control graphical model inference problem. Machine learning, 87(2):159-182 2012.
Matthew Kelly. An introduction to trajectory optimization: How to do your own direct collocation.
SIAM Review, 59(4): 904, 2017.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations, 12 2014.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv:1312.6114 2013.
arXiv preprint
Tejas D Kulkarni, Karthik Narasimhan, Ardavan Sacedi, and Josh Tenenbaum. Hierarchical deep reinforcement learning: Integrating temporal abstraction and intrinsic motivation. Advances in neural information processing systems, 29, 2016.
Saurabh Kumar, Aviral Kumar, Sergey Levine. and Chelsea Finn. One solution is not all you need: Few-shot extrapolation via structured maxent rl. Advances in Neural Information Processing Systems, 13:8198-822 8210, 2020.
Pierre L'Ecuyer. Note:
On the interchange of derivative and expectation for likelihood ratio derivative estimators. Management Science, 41 41(4):738-74) 1995.
Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv: 00001100500000. 2018.
Chengshu Li, Fei Xia, Roberto
and Silvio Savarese. Hrl4in: Hicrarchical rein- forcement learning for interactive navigation with mobile manipulators. In Conference on Robot Learning, pp. 603 616. PMLR, 2020.
Sizhe Li, Zhiao Huang. Tao Du Hao Su, Joshua B Tenenbaum, and Chuang Gan. Contact points discovery for soft-body manipulations with diffcrentiable physics. In International Conference on Learning Representations, 2022.
Yunzhu Li, Jiaming Song, and Stefano Ermon. Infogail: Interpretable imitation learning from visual demonstrations. Advances in Neural Information Processing Systems, 30, 2017.
Jacky Liang, Mohit Sharma, Alex LaGrassa, Shivam Vats, Saumya Saxena, and Oliver Kroemer. Search-based task planning with learned skill effect models for lifelong robotic manipulation. In 2022 IEEE International Conference on Robotics and Automation (ICRA). IEEE, 2022.
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. In Conference on robot learning, pp. 1113-1132. PMLR, 2020.
Viktor Makoviychuk, Lukasz Wawrzyniak, Yunrong Guo, Michelle Lu, Kier Storey, Miles Macklin, David Hoeller, Nikita Rudin, Arthur Allshire, Ankur Handa, al. Isaac gym: High performance gpu-based physics simulation for robot learning. arXiv preprint arXiv:2 108. 10470, 2021.
Todd K Moon. The expectation-maxiiizzaa 47-60, 1996.
algorithm. IEEE Signal processing magazine, 13(6):
Tongzhou Mu, Zhan Ling, Fanbo Xiang, Derek Yang, Xuanlin Li, Stone Tao, Zhiao Huang, Zhi- wei Jia, and Hao Su. Maniskill: Generalizable manipulation skill benchmark with large-scale demonstrations. arXiv preprint arXiv:2107. 14483, 2021.
Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine. Near-optimal representation learning for hierarchical reinforcement learning. arXiv preprint arxiv:/810.01157, 2018a.
Ofir Nachum, Shixiang Shane Gu, Honglak Lee, and Sergey Levine. Data-efficient hierarchical reinforcement learning. Advances in neural information processing systems, 31, 2018b.
12

---

Under review as a conference paper at ICLR 2023
Andrew Ng. Cs229 lecture notes. CS229 Lecture notes, 1(1):133, 2000.
OpenAl, Ilge Akkaya, Marcin Andrychowicz, Maciek Chociej, Mateusz Litwin, Bob McGrew, Arthur Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas, Jonas Schneider, Nikolas Tezak, Jerry Tworek, Peter Welinder, Lilian Weng, Qiming Yuan, Wojciech Zaremba, and Lei Zhang. Solving rubik's cube with robot hand. CoRR, abs/1910.07113, 2019. URL
org/ /bgg/bb///// abs
.07113
advantage-weigheed information maximization. arXiv preprint 19XIv::901.01366 2019. Takayuki Osa, Voot Tangkaratt, and Masashi Sugiyama. Hierarchical reinforcement learning via
Takayuki Osa, Voot Tangkaratt, and Masashi Sugiyama. Discovering diverse solutions in deep reinforcement learning by maximizing state-action-based mutual information. Neural Networks, 152:90-104, 2022.
Xue Bin Peng, Yunrong Guo, Lina Halper, Sergey Levine, and Sanja Fidler. Ase: Large-scale reusable adversarial skill embeddings for physically simulated characters. arXiv preprint arXiv:2205.01 009900 2022.
Karl Pertsch, Youngwoon Lee, and Joseph Lim. Accelerating reinforcement learning with learned skill priors. In Conference on robot learning, pp. 188-204, PMLR, 2021.
Yiling Qiao, Junbang Liang, Vladlen Koltun, and Ming Lin. Differentiable simulation of soft multi-body systems. Advances in Neural Information Processing Systems, 34:17123-17135, 2021.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning, pp. 8821-8831. PMLR, 2021.
Kashif Rasul, Calvin Seward, Ingmar Schuster, and Roland Vollgraf. Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting. In International Conference Machine Learning, pp. 8857-8868. PMLR, 2021.
Nathan Ratliff, Matt Zucker, Andrew Bagnell, and Siddhartha Srinivasa. Chomp: Gradient optimization techniques for efficient motion planning. In 2009 IEEE International Conference on Robotics and Automation, pp. 489 494. IEEE, 2009.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.0611 06155 2022.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning, pp. 1530-1538. PMLR, 2015.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF Confer- ence on Computer Vision and Pattern Recognition, pp. 10684-10695, 2022.
Manolis Savva, Abhishek Kadian, Oleksandr Maksymets, Yili Zhao, Erik Wijmans, Bhavana Jain, Julian Straub, Jia Liu, Vladlen Koltun, Jitendra Malik, et al. Habitat: A platform for embodied ai research. ICCV, 2019.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839):604--99 2020.
John Schulman, Yan Duan, Jonathan Ho, Alex Lee, Ibrahim Awwal, Henry Bradlow, Jia Pan, Sachin Patil, Ken Goldberg, and Pieter Abbeel. Motion planning with sequential convex optimization and convex collision checking. The International Journal of Robotics Research, 33(9):1251-1270, 2014.
John Schulman, Nicolas Heess, Theophane Weber, and Pieter Abbeel. Gradient estimation using stochastic computation graphs. Advances in Neural Information Processing Systems, 28, 2015.
13

---

Under review as a conference paper at ICLR 2023
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Tanmay Shankar and Abhinav Gupta. Learning robot skills with temporal variational inference. In International Conference on Machine Learning, pp. 8624-8633, PMLR, 2020.
David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference on machine learning, Pp. 387-395. PMLR, 2014.
Hyung Ju Suh, Max Simchowitz, Kaiqing Zhang, and Russ Tedrake. Do diffcrentiable simulators give better policy gradients? In International Conference Machine Learning, pp. 20668-20696. PMLR, 2022.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Emanuel Todorov. Linearly-solvable markov decision problems. Advances in neural information processing systems, 19, 2006.
Emanuel Todorov. General duality between optimal control and estimation. 2008 47th IEEE Conference Decision and Control, pp. 42922. IEEE, 2008
Marc Toussaint. Robot trajectory optimization using approximate inference. In Proceedings of the 26th Annual International Conference on Machine Learning, ICML '09, pp. 1049-1056, New York. NY, USA, 2009, Association for Computing Machinery. ISBN 978 9781605585161 10.1 10.1145/1553374155522 1533500
doi:
ottp///..../.11...111111.1 11..1..1111155555 111.11111155553331153333 1553508
Théophane Weber, Nicolas Heess, Lars Buesing, and David Silver. Credit assignment techniques in stochastic computation graphs. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 2650-2660. PMLR, 2019.
Keenon Werling, Dalton Omens, Jeongseok Lee, loannis Exarchos, and C Karen Liu. Fast and feature-complete differentiable physies for articulated rigid bodies with contact. arXiv preprint arXiv:2103.16021, 2021.
Ronald J Williams.
Simple statistical gradient- algorithms for connectionist reinforcement learning. Machine learning, ((3):229-266, 1992.
Jie Xu, Viktor Makoviychuk, Yashraj Narang, Fabio Ramos, Wojeiech Matusik, Animesh Garg, and Miles Macklin. Accelerated policy learning with parallel differentiable simulation. arXiv preprint arXiv:2204.07137. 2022.
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in Neural Information Processing Systems, 34:25476-25488, 2021.
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world, benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pp. 1094-1100, PMLR, 2020.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. arXiv preprint arXiw:2202.05607.
Yun-Yan&o Zhu, Taesung Park, Phillip Isola, and Alexei A Efros. Unpaired image-to-image translation using cycle-consistent adversarial networks. In Proceedings of the IEEE international conference on computer vision, pp. 2223-2232, 2017.
Yuke Zhu, Josiah Wong, Ajay Mandlekar, and Roberto Martín-Martin, robosuite: A modular simulation framework and benchmark for robot learning. arXiv preprint arXiv:2009.122 2020.
Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. Camegie Mellon University, 2010.
14

---

Under review as a conference paper at ICLR 2023
A PRELIMINARY
Markov decision process A Markov decision process (MDP) is a tuple of (S, R.). where S is the state space and A is the action space. p(s'|s, [is' a) a) is the transition probability that transits state to another state s' after taking action a. The function R(s, a, a's" computes reward per transition. A policy <(]ss) selects action distribution according to the state s. Executing policy π starting from the initial state 81 with density p(8)) will result a trajectory т. which is a sequence of states and actions {s1, a1 ................................ at..... where a, T[als s st), $1+1 ~ p(s)s St: az). We also use the terminology environment to refer to an MDP in an RL problem. The discounted reward of trajectory is (T) Ei=1 7 R(s,, dt, $1+1) where 0 < 7 < is the discount factor to ensure the series converges. The goal of reinforcement learning (RL.) is to find a parameterized policy me that maximizes the expected reward Esy=p(e)) V (s1) E [R,(T) where we call V (s1) the value of the state $1 for the policy ne.
For simplicity, we focus on optimization problems in a finite horizon MDP in this paper. This means that we only consider the first T states and actions {s1, 01,........ ST, ar,ar, ar, $T+1). In this case, we can directly optimize for the total reward R(r) ΣT R(s,, ar $412). Defining the measure and ensuring the existence of policy gradients for infinite horizon MDP with differentiable physics requires more strict conditions, and we leave as future work.
Policy gradient and zeroth-order gradient estimate REINFORCE (Sutton & Barto 2018) approximates the expected reward with Monte-Carlo sampling and optimizes the policy param- eter 0 with a zeroth order gradient estimate, which can be written as Nampin Vo log (a') where Nsamples is the number of sampled trajectories, 7¹ and a {at) lathsesT represent the trajectory and action sequence per sample, and the log-likelihood can be computed by log (aaa = Eisier log No (aglss).
IV (+1)l x
Differentiable simulation In differentiable physics we consider as special MDP where the transition probability p(s'|s, a) can be represented as ö(s' - h(s, a) The transition function h S A S and reward function R(s, a, s') are deterministic and differentiable. If we further assume that the policy (a]s)) = 5(a fo())) is a differentiable deterministic function w.r.t. 0 and input s, then each element in a sampled trajectory T and the total reward R(T) are also differentiable w.r.t. each 8t, ag and 0. We can optimize the reward R(T) with gradient-based optimizers after computing the first-order path-wise gradient
@R(T) through back- back- mack-propogation.
89
It's argued that even if the environment is deterministic, it is still beneficial to optimize a stochastic policy (Suh et al., [022 Xu et al. (2022) to leverage the advantages of stochastic sampling for the non- smooth reward landscape (Hu et al., 2019 Xu et al. 2022; Antonova et al.) 2022). A typical choice to add Gaussian noise w to the action sequence a and optimize it with the reparameterization trick (Kingma & Welling 2013), which estimates the same policy gradient as REINFORCE by OE,, -V (a1) R Eiero
Nample VoR(r'(a + w!)), where a² and w are actions and Gaussian noises
00 to generate the i-th trajectory 7'
RL as Probabilistic Inference
The RL as inference framework (Todorov] 2006; 2008 Toussaint
2009 Ziebart 2010: Kappen et al. 2012 Levine 2018), which defines optimality pl p(O)T) & R(T)T
where T is defined temperature. It further defines prior distribution of the trajectory p(r) p(s1) IT' ((a: (p St. at). where p(a, (se) is known prior action distribution, e.g., Gaus- sian distribution. Thus, we can compute the possibility of the optimality p(O) p(7 (p))..
The goal of the framework is to approximate the posterior distribution of optimal trajectories P(T(O) p(O(r)e(r) )(T)) which can be considered as softmax weighting among all prior tra- jectories based on their rewards, sharing the same spirit to path integral and the SoftMax policy.
28 is the Dirac delta function.
15

---

Under review as a conference paper at ICLR 2023
B MATHEMATICAL NOTATIONS
Notation
Explanation
2
The latent variable
The sampled trajectory
p(O[T)
Optimality of a trajectory
P(T)
Prior distribution of trajectories (for example, random actions)
P(O)
I p(O)r)p(7)dd.
P(T(O)
The posterior policy want model
(2(
The introduced auxiliary encoder
Joint distribution of 7 by sampling from the reparameterized policy in the environment p($ oz).
q0 (2, T)
(a(ar (a 8t; 2) or qo (1t 14,
The policy to select actions at different steps
go (2(21) or qo (2+ S1,, 21-1)
The policy to sample latent
(=(=(T)
The probability sampling & or all 2 given T in qo
po (2)
The marginal distribution of in do conditioning
(T(T)
The marginal distribution of T in go
po(T(2)
The marginal distribution of r conditioning
Qo The variational lower bound defined in Eq.
Ralbo
to (s(,,,2 2,
The deterministic policy function.
h(st, az)
The deterministic state transition function.
20 (2/=)
Usually the (o(2, TD) in deterministic environments
Ty To (2)
In deterministic environment, the trajectory becomes function of 0 and The path-wise gradient obtained from the differentiable environment
V V PeRooooTT
C CONNECTION WITH GENERATIVE MODELS
Many generative models are based on ELBO
log p(x) Egwa(s) [logp(r, 2) log q[=]] + KT(g())lp((\lx))
One can compare our approaches with other generative models. VAE defines pa (z, 2) = po (x12))((2)
and ((2) (2(2x) and then optimize 6, o jointly to maximize the ELBO bound. By doing so, (=(x)) has to align with the true posterior of (=(z))) Thus
logp(z) 2 [og (x|²) + p(z) - log go (=(x))
The coneetation-maximizaan po(r, 2) = Pa (=))))))(()))))))))))))))))))))) with finite z, thus we can compute the expectation po(2])) explicitly and treat it as 2s. then E-step: compute pa p=(2(x: to solve max log Pa (a: D (1s (=|ax) (po)ppp((la (=(x))- M-step: fixing o, find maxę Egs (logp((I, 2 - Egg [log [o(2(x)] can maximize the ELBO bound.
algorithm (EM) (Dempster et al. 1977) assumes that we have
In Maximum Entropy RL, we have p(O,7) and we optimize (a(T(O) only. The ELBO bound becomes maximum entropy term Exar [oor((O()) + log 1r))
p(O)=)p(T) defined by the reward,
Latent
Encoder ((=|x)
Joint p(x, *)
MLE
VAE EM Diffusion
Pa #)
pe(al=)p(s)
p(a)
log (a()) (lll
(=())))pa
(wre he2)
IT: M(x:: / Pir- 8x1) (r(r)
p(xr) [T 1
p(mo)
pp(]) p(O <)) (=1r)p())
P(O)
RPG
x.
(=.
P(O)
Table 1: Comparison of different algorithms that optimize ELBO bounds for inference
D IMPLEMENTATION DETAILS
The algorithm is shown in Algorithm 1
16

---

Under review as a conference paper at ICLR 2023
Algorithm 1 Variational Reparameterized Policy
1: Input: Pa. (2,
2: Initialize Pé, y0-
3: while time remains do
4: Sample start state $1.
5: Sample from the latent code policy 2,
6:
Compute the variational lower bound Retoo based Eq.
7:
Update the auxiliary encoder Ps to maximize Relbo.
8: Update the parameterized policy gradient qo to maximize Relbo with Theorem
9: end while
Table Hyperparameters and rewards our algorithms.
Hyperparameter Env
Bandit
Bandit
Movel
Move2
Move3
environment steps
1
1
8
8
50
num epoch
50
50
50
50
50
batch size
1000
1000
256
256
500
samples per epoch batch size
50000
50000
12800
12800
25000
1000
1000
256
256
500
The learning rate (3 x 10-4) is the same across environments. Adam uses the exact same hyperparameter method except for the dimension of the latent variable is 1. REINFORCE uses the exact same hyperparameter as our method except for the dimension of the latent variable and the policy gradient is estimated using reinforce.
E BIAS OF THE FIRST-ORDER GRADIENT ESTIMATOR
In this section, analyze the difference between two types of policy gradient estimators, namely the Zeroth-Order (ZoBG) and First-Order (FoBG) gradient estimator (Suh et al. 2022). We establish necessary condition for two to be equivalent when both action and parameter are dimensional. From the ID case, we also show the intractability of constructing a gradient estimator similarly to these two in higher dimensions.
Lemma (Leibniz integral rule, general form). Let fx, t) be a function such that both f(z, and its partial derivative 0f(x.t) are continuous in and x in some region of the at plane, including dz S x1. Also suppose that the functions a(x) and b(x) are both continuous
a(x) st b(x), Zu I and both have continuous derivatives for Do S x S 21. Then, for zo S x x1
(()e) dx Ja(a)
pb(x) Jaco on t)dt + dx
d a(x) f(z, a(x))
fla,
b(x) f(x,b(x))
dx
Lemma 2 (Jump Discontinuity, ID version). Let f(x, t) be function such that both f(x, t) and its partial derivative 0f(x.t) are continuous in and in some region of the at plane. has jump discontinuity if and only if there exist b, = =q((x) such that
Uz.
f(z,t) + lim f(x,t) (+))(())-
lim
++)+(()++
We define i correspond the ith jump discontinuity of f. For convenience, also abbreviate the above as f(b,) * f(bt) in later parts of this section. Corollary (Leibniz integral rule, ID region form). Let f(x, t) be function such that both f(x, t) and its partial derivative 0((a)) are continuous in and in region of the at plane, including
17

---

Under review as a conference paper at ICLR 2023
E = R. Also suppose that the functions q(xx) for all are continuous and have contimuous derivatives for x RR. Then, by Lemma and for X
((2 dx J[1
a / Ox (2,
Σ
@q.(x)
((( S(x, (()))()()))
Ox
interior
boundary part
For convenience, we also refer to the first term as the interior part and the second term as the boundary part the later parts of this section.
Theorem 2 (Policy Gradient Theorem, zeroth order version). Let TT(a) be function such that TT(.) E C for all E R and (T)(a) (a) E C² for all a E A, = R. Let R(aa be a function with jump discontinuities and defined for all
d S To To(a)R(a)da A
fA
Omo(a) R(aadda
00
holds.
Proof. Define the integrand Re(a) f(0, a), by Corollary |
To T((a)R(a)da
d0
{ (a)(a) R(a)da + Σ
(q.(0)
00
00
Case 1: If (bbb,) + 0, by Lemma the definition of the ith jump discontinuity can be rewritten
(:(0) = b. 5(0,b,) + b(0,b;)) To(b; RR(b, ) + No ((b())(()))
R(b,) + R(b+))
To(-) (.) E C for all 0
which means (0) is constant with respect to 0. 89
04:(0)
0 reduces the boundary part to 0 and gives
d0 / To(a)R(o)da
d
[ (a)
R(a)da
00
Case 2: (b) = 0, by Lemma 2 then f(0,b) = To R(b) = 0, which means 5(0, f(0,b b-) and does not have jump discontinuity at b.
Lemma 3 (Law of the unconscious statistician, continuous random variable). Let x be random variable and let g(x)) be a function of this random variable.
Eig(z)) [ fx(2)((2)
Theorem 3 (Policy Gradient Theorem, first-order version). Let be a random variable such that a = h(0,, w) with probability distribution p(w) E C and @hilww) + O for all Let f(0, w) = p(w) R()((.W) with the ith jump discontinuity of f defined be (9). The first order gradient is
du
d8 / Ro(a)R(a)da p(w) ② 00 R(h(f,w))dw + Σ (q:(0)
d
dA - =(0.
Proof. Reparameterize a tt(0 w), substitute by definition
d0 / To(a) do
d
d p(w))((((, w)) d [ f(0,w)dw
18

---

Under review as a conference paper at ICLR 2023
By Lemma 2 the definition of the ith jump discontinuity can be rewritten (:(0) = b, ((0,b5) * f(0, );") ((((0, + ((b,))(((())))) bT))
R(h(0, (1((0.,,)) b, * R(h(0,b,))
p(w) E C
Suppose -N(0,)), a), h(0, w) 0 + ow and exist action = a0=000, (0,)))) such that R(a,) + R(a
the definition of qi (0) and its evaluation at local neighbour qi (0 + A0) can be rewritten as
4: (0) b: h(0, b.) dj
(: + A0) bg h( 20, b, a4
since function R is not parameterized, the location of its ith jump discontinuity is a, and does not change. Therefore, substitute in h(0, w) = 0 + aw and solve for b; and of gives 6; = (a, - 0)/a
and 6,- = (a: - (0 + A0))/a. Differentiate q: (0) gives
(0)
( + A0) (0)
lim 28-U
do
AO
b2² b;
lim
A8
AA
A0/a
lim
A8-10
A
#0
In general, ignoring the boundary term when estimating the left-hand side with first-order estimator leads to a bias due to the discontinuity of the reward/value function.
Therefore, have shown that gradient estimator for only the interior part of the first-order gradient is a biased estimator of the original policy gradient. The bias is the boundary term.
F PROOF OF THE REPARAMETERIZED POLICY GRADIENT THEOREM
F.1 REGULARITY CONDITIONS
Let 0 - E R without loss of generality and Ω = fz E 2|g0 (z, ra) > 0} be a measure space, then for = E 2, 0 E c). we make the following assumptions
Assumption 1 (Bounded ELBO). The reward R. prior density log)( log go (2 (T), log (2((T) are bounded for all т.
Assumption (Lipschitz Condition). qo(2[T). log log (=((),, log (o(227) prior distribution log P(T), action policy (st,, 2, f), reward R and dynamics hiaa, 8t) are Lipschitz continuous w.r.f. any 8, and continuous 2. In addition, qo is Lipschitz continuous w.r.f.
The probability density q0(2, To) and its derivative (z, is Lebesgue-in intesuueinttegable jointly over Q and their integration is bounded for all 0 E [-c,c.
Assumption 3
Assumption iis easy to guarantee when the state space and the action space are compact and the time horizon 7 is finite because continuous functions compact set are always bounded. Similarly, for bounded inputs, the common neural networks can also satisfy Assumption 2 Otherwise, we can simply clamp them before feeding them into neural network to make the input bounded, do not need the network and the environment to continuously differentiable as in Lemma [ but only absolutely continuous. We require the dynamics h(a,, sx) to be Lipschitz continuous, ruling out chaotic systems.
Assumption holds for common distributions. For example, given Gaussian distribution p(z) el ((22)) /VV22T))) the parameter o's partial derivative de (z) e(-=²/((o")) () (-o2 2²)/(v 2TO²) is simply a function of exponential family and can be verified to be absolutely inte- grable on IR and bounded for any domain E [a, b) a, O.
19

---

Under review as a conference paper at ICLR 2023
F.2 PROOF OF THEOREM 1
To differentiate under integration, we use the following Lemma (cheng) L L'Euyyer 1995).
Lemma (Leibniz integral rule). Let X be an open subset of IR, and St be a measure space. Suppose that a function f X x Ω + IK satisfies the following conditions:
1. f(x w) is measurable function of and w jointly and is integrable over w for almost all
x E X fixed.
2. For almost all w E Ω fxx,,)) is an absolutely continuous function of x (the derivative (f w)x exists almost everywhere).
3. 05/Or is "locally integrable", that is, for all compuct intevals (a, bj E X,
J." 1/
② [(x, Ox
dwdr < co.
Then for almost every x E X, its derivative exists and
d da J f(1, = [ Ox w)w
a
Then easy to prove Theorem
Proof. Define the measurable function f(0, 2) go (0 To) (elb)(TT))
Since Retbo (To) is bounded by a constant M by Assumption [ and qo To) is positive and Lebesgue- integrable by Assumption 3, their product's absolute value will be bounded by (o(2, To) M, (0,
becomes Lebesgue-integrable for every 0.
By Assumption Sa and assuming the horizon is finite, the trajectory will become Lipschitz function 0f(6.2)
with respect to 0 and z, so as (T)O)TT), (2(2, ra) and their products f(0, Thus almost everywhere for all = E Ω.
exists 5e
To bound its derivative, let Rebo(Te) I.-L.ischitz w.r.t. 0, notice that [oJ =) = Retoo (T)) Vogo(2, Tp) + (0(2, To) Va Relo S M [Vogo(2, TO)1 ) + (10 To)1
is dominated by the sum of two Lebesgue-integrnn
functions whose integration is bounded for all
0 E e by Assumption 3 Thus its gradient must locally integrable. Then applying Lemma we
can do differentiation under the integration sign and finish the proof as
V g E Reehoo
V (qa(2, To Retho (Ta)) dz
Rabo (To) Fo (Z, (=((8) g0(2, Vo (To)dz
[ go (z, To) Ratho (To) V (o Va (To) d2.
We want to emphasize the importance of the Lipschitz continuity that helps to bound the integration. Informally, the Lipschitz condition is directly related to the scale of the path-wis wise gradient. When the Lipschitz constant is too large, the environments will generate exploding gradients or contains sharp changes in the reward landscape, similarly to the discontinuity and causing empirical bias (Suh et al.) (2022). Besides, when the reward landscape has discontinuous points, there will be bias caused by the motion of the discontinuity, as suggested by the Leibniz integral rule. We refer interested readers to Appendix for detailed illustration of the difference between the zeroth-order gradient estimate and the first-order gradient estimate under the expectation we do not find in the literature. Note that we do not ask for the continuity of the gradients and support activation functions such as ReLU.
The finite horizon assumption is also necessary to make the theorem holds, otherwise the gradients Va Rabb((T) may be unbounded if we do not up a suitable discount factor.
20

---

Under review as a conference paper at ICLR 2023
G LEARNING A DIFFERENTIABLE WORLD MODEL
G.1 METHOD
To apply our method in non-differentiable environment, train differentiable world model joindy with the policy optimization as Model-based Reparameterized Policy Gradient (MBRPG). This method has been proven data-eflicient for policy optimization (Hatner et al. 2019 Schrittwieser
et al. 2020 Ye et al. 2021 Hansen et al., 2022
Specifically, besides of the encoder Pal po(21s, (2|s, a). the policy ITa (a|s, 2), (2(2)))) we additionally learn a deterministic dynamic network (s, a), a reward network (s, a), an observation encoder fe (o) and the Q 2-nntwork (w a, 2). In practice, they are all two-layer fully connected neural networks with hidden dimension 256. We also define o, 3 to weigh the entropy term and the cross entropy term the variational lower bound.
Given any 2 and any initial latent state s; = fa)) and an arbitrary action sequence, we can use the learned dynamic network to generate the trajectory by
a2 (a(as 2) or a, at
(policy) (reward)
r3 = (s at)
r/ a log No (urls)) 2) + 3 logpo(zlse,uu
(intrinsic reward) ( value) (dynamies)
Q: Qu(se,
81+1 (82 az)
for - i, 1, ,,. i K. Notice that during the model rollout, can either the action or sampled from the current policy Ta [t, 2) or action sequence {aî"} sampled from the replay buffer. When the action is sampled from the current policy To (a 18, 2). we obtain a Monte-Carlo estimate for the value of 8,, which can be used to optimize the policy Tz:
i
Vestimate (os. se 7 (Qitk- + x(4x)+ + Σ P(r+ + r()
(I)
To train the dynamics model, we sample trajectory segments of length K + 1 Tei+K {os, ag' ,r? 0 gt r self-supervise the dynamic network to ensure state consistency and avoid reconstruction as (Ye
from the replay buffer and select a latent We then " 0++x
al. 2021 Hansen et al. 2022) (in this case we let 0 = a. in the trajectory):
S+K-1
(T( =
Σ - wgloooo(/oo00000--0000
22) La ng(r?" + (oe+1, 2))) (Qn(:QQ:n"""
(2)
where ng(z) means stopping gradient and L1 1000, La L3 0.5 constant balance the loss. The whole process is illustrated in Algorithm 2
For all experiments, we take K = 6.
G.2 EXPERIMENT
We first compare our method with SAC Haarnjja al. (0118) on locomotion tasks (Cheetah-v3 and Humanoid-v3 -v3 in OpenAl Gym Brockman et al. (2016) in Figure ((a) and (b). We plot the learning curve of MBRPG together with the SAC's performance after training for 3 million steps. We can see that by learning the model, our method achieves on-par performance with fewer samples. In particular, our method only requires million samples to reach 15000 scores on Cheetah-v3 and only 0.5 million samples to reach 6000+ scores on Humanoid-v3 This attributes to the efficiency brought by learning a model for policy optimization. However, we observed that removing the latent variables and the variational bound, which is baseline model-based RL algorithm (MBRL), did not affect the performance of our method. We conjecture that these locomotion tasks do not require a multi-modal policy for exploration, and Gaussian policy is sufficient.
To illustrate the effectiveness of our approach, we have built two environments that need multi- modality explorations, illustrated in Figure In (a) AntMove environments, ant robot move
21

---

Under review as a conference paper at ICLR 2023
Algorithm 2 Model-based Variational Reparameterized Policy
1: Input: Pa, na, Nuy, Ry, Jav Q
2: while time remains do
3: Sample start state o1 and encode it f01) Select 2 from Te(2[s1).
4: Execute the policy (, 2) and store transitions into the replay buffer B.
5: Sample a batch of trajectory segment of length K {T 2} from the buffer B
6:
Optimize / using Equation 2
7:
Optimize [ 2) with gradient descent maxmize the value estimate in Equation 1 for
9, sampled from the buffer.
8:
Optimize (T(2|s)) (/s)) with policy gradient to maximize Vestimate (81, =) log (z(2())) for si
is sampled from the buffer.
9: Optimize o, a, and other auxiliary networks.
10: end while
(a) AntMvvv
(b) Cabinet
Bonus: 0. XX
X
Bonus: 1.0
Figure Non- cin-eifferentiabee environments
either upward or rightward to reach one of the two goal regions represented by the black circle. The agent always receives a penalty proportional go upward instead of moving right. However, if the agent chooses the longer right path, it will receive a higher bonus when it reaches the goal. Such a local optimum will trap a single modality agent, as shown in Figure (c). When remove the latent space and the latent policy p(2[s1), the baseline single-moda RL policy (MBRL) will get stuck in the local optima, while our approach (MBRPG), thanks to its ability maintain multi-modality trajectory distribution, will attempt go right when it is sub-optimal at the beginning and thus has a higher chance to find the global optima. We dof mobile robot for opening the cabinet shown in Sigure
its distance to the closest goal, which motivates it to
(c).
also test our method for casier for the robot to open the left door, but it will receive a higher reward when it opens the right one. Similarly, the normal RL agent modeled by Gaussian distribution (MBRL) fails explore way to open the right door. In contrast, our method explores the two directions simultaneously and can open the right one, in the end, resulting in higher rewards, as shown in Figure 6
is
[a] Cheetah-v3
(b) Humanaid.vv
(c) AntMove
(d) Cabinet
15000
SecE9-MM
Sac
150
mbrog mbri
200
mixl
ovg [rk return
2vg tral naturn
0030
200
2Mg trai rettum
100
10000
6000
5000
2000
0
0.0
0.5
1.
0.0
0.2
0.4
0.0
0.2
0.4
0.0
0.1
0.2
0.3
simulator steps M)
simulator steps (x
simulator steps M)
simulator steps (x M)
Figure 6: Experiments with learned world model.
22

---

Under review as a conference paper at ICLR 2023
std=3.00
std=1.00
stdw0.45
std=0.20
std=0.12
std=0.01
1.00
1.00
1.00
1.00-
1.00
1.00
0.75
0.75
0.75
0.75
0.75
0.75
20.50
0.50
0.50
0.25
0.25
0.25
0.25
0.25
0.25
0.00
0.00
0.00
0.00
0.00
0.00
2
action
action
action
action
action
action
Figure 7: The landscape of expected reward after a gaussian filter with gradually decreasing standard deviation.
H INSUFFICIENCY OF THE SINGLE MODALITY GAUSSIAN POLICY
Our method is motivated by the insufficiency of the Gaussian policy in solving non-convex problems. gaussian policy with large entropy has proven effective in many continuous control problems. However, even if the Gaussian distribution has very high initial variance, it may still fail to solve many tasks. The Gaussian parameterization can only model mngle-mydality action distribution, which limits the policy update in both REINFORCE and ADA. It would prevent them from converging to the global optima even though it can explore the whole reward landscape.
Without loss of generality, let take the 1d continuous bandit problem case, a Gaussian policy is fully determined by its mean ja and variance a². Both REINFORCE and ADAM maximize the expected reward Ea [R(aa] by computing its policy gradient VEn-W((..) [R(a)) (a)). We now assume R(a) is Lipschitz continuous and differentiable almost every- where. As we show in the paper, the gradients computed by the two methods are the same.
example. In this
If the standard deviation is not large enough, a Gaussian policy will suffer from the non- convexity issue like a deterministic optimizer. It is easy to see that V Fa-s((a)) [R(a)] =
2)²
VP f. 1/(oV/T) a deterministic function given by passing the original reward landscape through a Gaussian filter. As shown in Figure 7 optimizing R(a) with policy gradient will be equivalent to running gradient descent on F (μ) with a sufficient number of samples. So if non-convex reward R(a) is still non-convex after was smoothed by a Gaussian kernel, the random perturbation provided by the Gaussian noise will not help the policy jump out of local optima.
R(a)da VyF"(u), where (")))
1/((o/2T)
R(a)da
We now consider the method starting with a very large standard deviation a and gradually reducing to zero, shown in Figure We plot the F" for different std The leftmost shows F for large a and the rightmost close to the original reward function. We use the red dotted line to denote the optimum jr (a) = arg max,, F (J4) for each We can that for very large a, Fe is convex, and we can assume the policy gradient method finds the optimal f (a). However, /r" for very large is not necessarily close to the global optima μ (0). In the beginning, the right mode is optimal as it has high average reward (which can be visually measured by the area under the right mode). But when we reduce the gradually, the global maximum may drastically change from the right mode to the left mode as shown in the red block (middle two) of Figure because the left mode has higher extreme value that exceeds the right. However, at the moment that the optimum changes to the left mode, Na is already non-convex; the policy gradient method has little chance to discover the new global optima and will get stuck the right local optima. This behavior will devastate all bradie algorithms that optimize hingle-mdda policy (REINFORCE, ADAM, SAC, and PPO).
To illustrate the insufficiency of the gaussian policy with large standard deviation, we vary the initial standard deviation for REINFORCE and anneal to 0 in sufficiently large amount of training steps. Notice that the global optima have reward of 1.
23

---

Under review as a conference paper at ICLR 2023
Initial std
0.01
0.1
0.5
1
2
10
expected reward
0.50=0.00
0.50=0.00
0.50+0.00
0.50+0.00
0.49#0.00
0.47=0.00
0.42++0.00
Table Final performance of reinforce with linearly decreasing standard deviation.
ABLATION STUDIES
In this section, first study the importance of the trajectory encoder to understand what role it plays in helping maintain a multi-modality trajectory distribution. Then, we show how hyperparameter controlling the reward weight of entropy of action affects algorithms that have different policy paramterization.
Bandit 1
Bandit
Move 1
Move
Move
1.00
0.75
0.8
0.6
5.0
-3.5
0.25
2.5
0.2
0.0
-2.0
0
0.0
0.5
1.0
0.0
1.0
Reinforce
Adam
Ours
Ours (no info loss)
Figure 8: Ablating trajectory encoder
In Figure we show the importance of the trajectory encoder, We can see that when we do not optimize the mutual information between and r through log (2(r) (Ours (no info loss)), our method degrades to almost identical behavior Adam. We can that the best-performing method our method in its current form.
24


Reviews:
[{"review_id": "fou44K6gcGG", "reviewer": "Reviewer_5kNp", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper studies reinforcement learning (RL) in high-dimensional continuous action spaces. It introduces Variational Reparametrized Policy (VRP), which formulates the policy as a generative model for optimal trajectories, optimized using a variational approach. The paper claims that this allows modeling complex and multi-model trajectory distributions. It evaluates VRP on two differentiable environments (a bandit and a grasping task) and compares against standard REINFORCE and Adam.\n", "summary_of_the_review": "While generatively modeling multimodal trajectory distributions is an interesting idea, the submission lacks clarity and provides only very basic empirical evaluations with limited scope. I encourage the authors to improve the exposition of their work and extend the experimental evaluation. In its current form, I recommend to reject this submission.", "strength_and_weaknesses": "Strengths: the basic idea of generatively modeling multimodal trajectory distributions is interesting.\n\nWeaknesses:\n\n- The structure and writing of the paper is poor, which makes it hard to follow. \n- The paper only discusses the very limited setting of differentiable environments.\n- The empirical evaluation is of low quality.\n\nPlease find detailed comments below.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "LrrmtbLRihk", "reviewer": "Reviewer_CRPH", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The work proposes a method for exploration in reinforcement learning.\nThe idea combines the concepts of variational inference and control as inference. In the method,\na variational autoencoder is used to learn a latent representation of the optimal trajectory distributions, where\nthe probability of optimality of a trajectory is proportional to $\\exp(R(\\tau))$ as in the typical control as inference\nframework. The \"decoder\" in this framework includes the policy that takes the latent variable $z$ as an input and tries to maximize the rewards. The latent variable $z$ can either be sampled at the beginning of an episode and kept as a fixed input \nto the policy, or it can also be periodically resampled using a state-dependent sampler.\nIt seems the main benefit is that $z$ leads to temporally correlated exploration.\n\nOne of the most closely related works is SeCTAR (Co-Reyes et al., ICML 2018) which uses a similar trajectory modeling framework, but\nthey use planning instead of policy search for action selection.\nOther related works include skill or options learning frameworks that also use a latent variable input to the policy.\nAnother difference to prior works is that the current work uses differentiable simulation to aid in the optimization.\n\nExperiments were performed in small bandit and other tasks, as well as in manipulation tasks with differentiable physics.", "summary_of_the_review": "I recommend rejecting the paper.\n\nWhile the idea was theoretically founded, I believe the writing lacked focus making the paper not that clear. Moreover, the experimental work did not appear thorough, and I believe it requires more investigation. Including comparisons with existing benchmarks or implementations would be good. Ablation studies on the different components of the algorithm would also aid in convincing me.", "strength_and_weaknesses": "**Strengths**\n1. The method seemed like a technically sound way to learn a latent space for exploration.\n2. The idea is widely applicable in reinforcement learning.\n\n**Weaknesses**\n1. The experimental comparisons all appeared new without any comparison on previous benchmarks or against previous implementations. Including at least one comparison with previous work might be good.\n2. The theorems (primarily the gradient estimator and the discussion around discontinuities) directly follow from previous work, and I don't think there was anything to really derive. The proofs and regularity conditions appeared redundant to me as the continuity requirements are well-known.\n3. I think the work lacked focus. For example, I think the main points of interest are whether the encoder is useful, and whether the ELBO objective is a useful theory. I didn't see any ablation study removing the encoder, or experiments where one samples a random $z$ giving this as an input to the policy and optimizing the rewards as usual. I would have liked to see stronger justifications regarding why the proposed method is useful in exactly the current embodiment compared to slight modifications of it.\n4. The experiments appeared contrived to me. For example on the bandit 1 and 2 tasks, results for only one set of hyperparameters were shown. If the initial variance of the Gaussian distribution is set sufficiently high, I believe all of the methods should be able to solve these tasks. It would be good to test more hyperparameters and show the sensitivity of the methods.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "DNs20FOOc51", "reviewer": "Reviewer_LHSH", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposed a new RL algorithm for continuous control. A key idea is to represent the policy as a distribution of trajectories in the problem space to capture the multi-modality nature of the problem. A variational lower-bound is derived to create a practical optimization objective and with some assumptions on the smoothness of the problem a gradient calculation formulation is derived that can leverage differentiable physics simulation. The resulting method is applied to a few robotic control problems in simulation to show the effectiveness of the method.", "summary_of_the_review": "The presented method is interesting and novel as far as I can tell. \nAs mentioned in the sections above, it'll be helpful if more discussions and potentially comparisons can be provided.", "strength_and_weaknesses": "Strengths:\n1) The proposed method is solid \n2) The ability to leverage differentiable physics and mitigate local minima issue is desirable\n3) Experiments shows the effectiveness of the algorithm\n\nWeaknesses:\n1) It’s not clear how well the method can work if a differentiable physics is not available.\n2) It is mentioned that the method assumes lipschitz continuity in the problem while some of the presented tasks doesn’t seem to be so all the time (e.g. when the object make contact with the rope). How does the method handle this?\n3) Furthermore, I'm wondering how well would the model be able to handle non-stable dynamics such as balancing a pole, or bipedal locomotion, where a small perturbation would lead to drastically different states later on. \n4) The method seems to be also related to methods that uses mutual-information for encouraging diverse behaviors (e.g. Mutual Information State Intrinsic Control, Discovering Diverse Solutions in Deep Reinforcement Learning by Maximizing State-Action-Based Mutual Information, etc). Some discussions with these methods would be useful.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject