Model-Based Reinforcement Learning via Imagination with Derived Memory

1 Introduction . Model-based Reinforcement Learning ( MBRL ) approaches benefit from the knowledge of a model , allowing them to summarize the agent ’ s past experiences and foresight the future outcomes in a compact latent space [ 1 , 2 , 3 , 4 ] . Previous approaches can be generally categorized into four branches , including 1 ) Dyna-style algorithms , 2 ) sampling-based methods , 3 ) value expansion , and 4 ) backpropagation through paths . First , the Dyna-style algorithms [ 5 ] generate virtual interaction ∗Yuzheng Zhuang is the corresponding author . Yao Mu conducted this work during the internship in Huawei Noah ’ s Ark Lab . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . data ( e.g . ME-TRPO [ 6 ] , MBPO [ 7 ] ) to improve policy optimization . Second , the sampling-based methods , such as [ 8 , 9 , 10 ] , choose action by planning or shooting ( e.g . PlaNet [ 8 ] , PETS [ 9 ] ) . Third , the value expansion approaches , such as [ 11 , 12 ] , use model-augmented rollouts to improve the estimation accuracy of cumulative returns . Finally , the methods of backpropagation through paths [ 13 , 14 ] employ the analytic gradient of state value through the dynamic transition to learn the policy . For the tasks with high-dimensional sensory inputs , backpropagation through paths is more computationally efficient compared with conventional planning algorithms , which generated numerous rollouts for the best action sequence selection procedure . Dreamer [ 4 ] , as a landmark of such methods , achieves state-of-the-art ( SOTA ) performance on visual control tasks . It summarizes the agent ’ s high-dimensional experiences into a compact latent space by the dynamic model , which is analogized to the way human memories are stored , and benefits behaviours learning by latent imaginations . However , recent breakthroughs based on such a method focus on optimizing the policy on imaginary trajectories with only original memory on the agent ’ s real experiences and leave the diversity of imagination largely unstudied , resulting in noise-sensitive policy and inefficient learning . Recently , neuroscientists demonstrated that `` memory prosthesis '' , which transforms the brain ’ s activity patterns to the electrode signals to stimulate the human ’ s hippocampus [ 15 , 16 ] , could facilitate humans memory augmentation via an external implant . Inspired by such insight , our idea is to improve behaviour learning by constructing a “ memory prosthesis ” with externally derived memory [ 17 ] that each signal corresponds to a state transition sequence , which can enrich the diversity of the imagination without any interaction with the real environments . We propose a novel model-based reinforcement learning algorithm , called Imagining from Derived Memory ( IDM ) , which aims to improve policy robustness and sample efficiency by the enriched and reliable imagination . In order to generate diverse derived memory , transformations are performed on states of the latent trajectories , which are sampled from the experiences . The agent learns to obtain future rewards by the model-based imagination from both the derived and the original memories . The optimal policy is optimized by a novel actor-critic algorithm , where the state values optimize Bellman consistency for imagined rewards , and the policy maximizes the state values by propagating their analytic gradients back through the dynamics . In this paper , we aim to improve the diversity of imagination for model-based policy optimization with the derived memory . Our main contributions are listed below . • We design an effective memory generator that simultaneously achieves the imagination diversity as well as the relevancy between the derived memory and the real physical states , which avoids the distortion after decoding . • We propose a model-based policy optimization framework , unifying the analytical gradients of the imagination from both derived and original memories with learnable confidence . We also provide the upper bound of the value estimation error in the frame of IDM . This sheds light on further improvement of model-based reinforcement learning in the future . • Experiments on DMControl [ 18 ] tasks demonstrate that IDM outperforms existing SOTA approaches in terms of robustness to uncertainty and further improves the sample efficiency of the model-based method . Ablation experiment results verify the superiority of IDM . 2 Preliminaries . 2.1 Model-based Reinforcement Learning . Model-based reinforcement learning aims at optimizing a policy to maximize the cumulative rewards , by accessing to a ( known or learned ) model of the environment . We denote a time step as t , a state at t as st ∈ S , action at t as at ∈ A , reward function r ( st , at ) , policy πφ ( st ) and a world model pθ to characterize the process of interacting with the environment . The goal is to find a policy parameter φ∗ , which maximizes the long-horizon summed discounted rewards represented by a parameterized value function , vψ ( st ) . = E ( ∑∞ i=t γ i−tri ) with parameter ψ and the discounted factor γ . 2.2 World Model . The world model characterizes the process of interaction with the environment . The recurrent stochastic state model ( RSSM ) , which is used in PlaNet and Dreamer , is the SOTA solution for building latent dynamics [ 8 ] . RSSM can facilitate long-term predictions and enable the agent to imagine thousands of trajectories in parallel . It can represent both the historical memory and the dynamic uncertainty by using a recurrent neural network . Let us denote the sequence of observations as õ = { o0 , o1 , . . . , oT } , action sequence as ã = { a0 , a1 , . . . , aT } , rewards as r̃ = { r0 , r1 , . . . , rT } , and the corresponding latent states z̃ = { z0 , z1 , . . . , zT } . RSSM assumes that the latent state zt consists of zt = ( st , ht ) , where st , ht are the probabilistic and deterministic variables respectively . RSSM consists of four key components , including ( 1 ) the representation model p ( zt|zt−1 , at−1 , ot ) , which encodes observations and actions to create continuous vector-valued latent states with Markovian transitions , ( 2 ) the transition model q ( zt|zt−1 , at−1 ) , which is built by a recurrent neural network , and is to predict the future latent states with historical memory and actions , ( 3 ) the observation model q ( ot|zt ) , which is used to provide a reconstruction learning signal , and ( 4 ) The reward model q ( rt|zt ) predicts the rewards given to the latent states . We use p to denote distributions that generate samples in the real environment and q to denote their approximations that enable latent imagination . Specifically , the transition model allows us to predict ahead in the compact latent space without any interaction in the real environment . 3 Robust World Model for Derived Memory Generation . Based on RSSM , diverse memory vectors could be generated intuitively by transforming the original memories while putting forward higher requirements on the world model ’ s robustness to the latent space noise . Because if the derived memory can not correspond to the real physical state , i.e. , the reconstructed image is distorted , it could not effectively enrich the imagination but bring disturbance for policy learning . As shown in the second row in Figure 1a and Figure 1b , if we add an random noise to the latent state zt = ( ht , st ) , the derived latent state z′t = ( h ′ t , s ′ t ) in RSSM results in low-fidelity reconstructed images , which could not be mapped to the real physical states . To alleviate this issue , we propose a Robust Recurrent Stochastic State ( R2S2 ) model to improve the latent model ’ s robustness . We add a decoding constraint to the RSSM , which improves relevancy between the derived latent states and the real physical states , thus avoiding the distortion after decoding . Denote ξh is a noise that obeys N ( 0 , σh ) , h′t is the derived hidden state which adds the noise ξh to the hidden state ht , ξs is a noise that obeys N ( 0 , σs′ ) , s′t is the derive stochastic state which adds the noise ξs to st , ot is the reconstructed image from zt = ( ht , st ) , and o′t is the reconstructed image from z′t = ( h ′ t , s ′ t ) . The components of the R2S2 are optimized jointly to increase the variational lower bound [ 19 ] with a decoding constraint . The variational lower bound JM includes reconstruction terms for observations J tO and rewards J t R and a KL regularizer J t D [ 4 ] . The expectation is taken under the dataset and representation model , minJM = Ep ( ∑ t ( J tO + J t R + J t D ) ) s.t.DKL { p ( o′t|z′t ) , p ( ot|zt ) } ≤ δ ( 1 ) J tO = log q ( ot|zt ) J tR = log q ( rt|zt ) J tD = −βDKL { p ( zt|zt−1 , at−1 , ot ) , q ( zt|zt−1 , at−1 ) } . As demonstrated in Figure 1 , our proposed R2S2 model generates high-fidelity reconstructions , which suggests the enriched dynamic representation is capable for facilitate behaviour learning . The decoding constraint leaves a margin δ between the derived branch ’ s output and the original branch ’ s output . The constrained problem is solved by the fixed penalty method [ 20 ] where the penalty coefficient is given in the appendix A.4 . As shown in the third row of Figure 1a and Figure 1b , after adding noise on the original latent state zt = ( ht , st ) in R2S2 model , the reconstructed images with derived latent state z′t = ( h ′ t , s ′ t ) are sharper than RSSM , and are demonstrated different states of motion compared to the the reconstructed images with original latent state zt = ( ht , st ) . In summary , the decoding constraint used in the R2S2 model aims to guide the derived memory in latent space to meet a basic condition that the neighbouring latent states from real data should also follow the real physical states . 4 Policy Optimization with Derived Memory . By leveraging the R2S2 model , we propose an Imagining with Derived Memory ( IDM ) algorithm that aims to enable the agent to learn the policy with diverse imagination . As shown in Figure 2 , the derived memory is generated by the transformation from the original memory . The agent imagines the future states and rewards with both derived and original memories . The policy is improved by the enriched imagination with prediction reliability weight under the actor-critic framework . 4.1 Imagining with Derived Memory . We add several zero-mean noises ξh ∼ N ( 0 , σ2h ) on the hidden state ht and ξs ∼ N ( 0 , σ2s′ ) on the stochastic state st ∼ N ( µs , σ2s ) in the original memory to extend the agent ’ s memory which includes diverse embedded historical information . Thus the agent could learn behaviours with the trajectories in an enlarged latent space , rather than just the agent ’ s real experiences . The Gaussian noise in latent space could be theoretically mapped to any type of uncertainty in image level , and this principle is extensively used in VAE and GAN . Adding Gaussian perturbations should be an intuitive way to simulate the influences in latent space caused by the uncertainty of input image , since the stochastic part of latent state st in the R2S2 model obeys Gaussian distribution , any image will be encoded to a Gaussian variable in latent space with R2S2 model . The latent trajectories are imagined from both original memory and the derived memory by the learned transition model . The method to generate derived memory can be summarized as Equation ( 2 ) , for the i-th derived trajectory , the deriving process contains two key steps : derived memory generation and trajectory imagination . Derived memory generation : hit ∼ N ( ht , σ2h ) , sit ∼ N ( µs , ( σs + σs′ ) 2 ) Trajectory imagination : hit , s i t ∼ p ( hit , sit|hit−1 , sit−1 , ait−1 ) , ait ∼ π ( ait|sit , hit ) ( 2 ) In order to ensure the reliability of the imagination with derived memory and further mitigate the model error ’ s impact , we consider the reliability of the imagined trajectories by adding a trajectory evaluator . Every imagined trajectory is used with a reliability multiplier given by the evaluator . In the region where the model has reliable predict accuracy , the derived memory should be able to improve the policy robustness and sample efficiency significantly . We penalize the weight of the trajectories in the region with poor model prediction . The trajectory evaluator E aims to discriminate between real and overshoot sequences predicted by the model . In every iteration , we sample a number of state-action sequences generated by real interactions as real sequences Seqreal = [ zτ , aτ , zτ+1 ] t+T τ=t , use the latent model to predict the overshoot trajectories from the same start points as the fake sequences Seqfake = [ zt , at , ẑt+1 ] ∪ [ ẑτ , aτ , ẑτ+1 ] t+Tτ=t+1 . Train the evaluator to output the probability of whether a real trajectory , which is equal to maximize the following objective function Jeva = E ( zτ , aτ ) ∼D { t+N∑ τ=t [ log ( E ( zτ , aτ , zτ+1 ) ) + log ( 1− E ( ẑτ , aτ , ẑτ+1 ) ) ] } ( 3 ) where E ( · ) is approximated by the combination of the MLP and the outermost Softmax and zt = ẑt since zt is the start point of the overshoot trajectory .

Summary: This paper presents a model-based reinforcement learning agent that uses imagined trajectories, obtained via rollouts from noised initial states observed by the agent, to improve data efficiency and robustness. The approach builds on top of Dreamer, a popular MBRL approach that trains a latent dynamics model and trains a policy via value gradients computed on imagined rollouts from the latent model in an actor-critic setup. The presented approach, called Imagining with Derived Memory (IDM), improves upon Dreamer by generating diverse imagined rollouts; this is done by adding random perturbations to the latent state of the initial observation from which an imagined rollout can be derived. By augmenting the training data of standard imagined rollouts (whose initial states were observed in the environment) with derived rollouts from these perturbed initial states the approach generates diverse additional data to effectively double the batch size.  To ensure that the derived rollouts are physically plausible two additional changes are made to Dreamer. First, a decoding constraint is added to the model’s training loss; this loss constrains the reconstructions from the original and perturbed latents to remain close to each other, thereby encouraging the perturbed latents to also encode plausible latent states. Second, to ensure that only physically plausible rollouts are used for training, the approach additionally trains a discriminator that scores each trajectory; this discriminator is trained to distinguish between real and imagined trajectories and in practice scores “real-looking”, and therefore more plausible trajectories higher than implausible ones. This score is used to weight the loss for training the policy and value function, thereby mitigating the effect of implausible trajectories affecting learning. Additionally, a theoretical analysis of the IDM framework is also provided which upper bounds the error between true returns and returns achieved in the model.  The approach is tested on six tasks in the control suite, both with and without observation uncertainty (random image rotations within +- 5 degrees) and compared with Dreamer and DrQ, a state of the art model-free approach. IDM performs better than both baselines on 3/6 tasks without observation uncertainty, and significantly outperforms both baselines on all tasks with observation uncertainty added. A further study that ablates different parts of the algorithm including the use of derived rollouts, the use of the discriminator for scoring rollouts and the decoding constraint is provided; in general the combination of all three is needed for improving data efficiency in comparison to baselines.