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 describes Imagining with Derived Memory (IDM), a model-based reinforcement learning method based on a world model. The proposed method is inspired by the memory prosthesis proposed by neuroscientists. The method introduces an additional disturbance into the latent variable of the recurrent state-space model (RSSM) to facilitate the robustness of the model. The experiment shows the IDM outperformed a baseline method, Dreamer, in several tasks.