Model-Based Reinforcement Learning via Imagination with Derived Memory
Authors: Yao Mu
Yuzheng Zhuang
Bin Wang
Guangxiang Zhu
Wulong Liu
Jianyu Chen
Ping Luo
Shengbo Eben Li
Chongjie Zhang
Jianye HAO
Venue: NeurIPS 2021 Poster (2021)
Keywords: Model-based Reinforcement Learning
Policy Robustness
Visual Control Task
Derived Memory
World Models

Abstract: Model-based reinforcement learning aims to improve the sample efficiency of policy learning by modeling the dynamics of the environment. Recently, the latent dynamics model is further developed to enable fast planning in a compact space. It summarizes the high-dimensional experiences of an agent, which mimics the memory function of humans.  Learning policies via imagination with the latent model shows great potential for solving complex tasks. However, only considering memories from the true experiences in the process of imagination could limit its advantages. Inspired by the memory prosthesis proposed by neuroscientists, we present a novel model-based reinforcement learning framework called Imagining with Derived Memory (IDM). It enables the agent to learn policy from enriched diverse imagination with prediction-reliability weight, thus improving sample efficiency and policy robustness. Experiments on various high-dimensional visual control tasks in the DMControl benchmark demonstrate that IDM outperforms previous state-of-the-art methods in terms of policy robustness and further improves the sample efficiency of the model-based method. 

Model-Based Reinforcement Learning via Imagination with Derived Memory
Yao Mu The University of Hong Kong
Yuzheng Zhuang Huawei Noah's 's Ark Lab
Bin Wang
Huawei Noah's 's Ark Lab
muyao@connect. www.eeett.......... .ku.hh.
zhuangnuzhenghuuarei.com
wangbin1580huuweii. com
Guangxiang Zhu
Wulong Liu Huawei Noah's 's Ark Lab
Jianyu Chen
Tsinghua University guangkiangghu@@uttook..cc
Tsinghua University
Liwrulongohuavei.c com
j jianyuchenttsinghua edu. cn
Ping Luo
Shengbo Eben Li
Chongjie Zhang
The University of Hong Kong pluodcs. hku. hk
Tsinghua University
Tsinghua University
lishbo@tsinghua. edu. cn
chongjie@tsiinhuua edu.
Jianye Hao Huawei Noah's Ark Lab haojianye@huawei. com
Abstract
Model-based reinforcement learning aims to improve the sample efficiency of policy learning by modelling the dynamics of the environment. Recently, the latent dynamics model has been further developed to enable fast planning in compact space. It summarizes the high-dimensional experiences an agent, which mimics the memory function of humans. Learning policies via imagination with the latent model shows great potential for solving complex tasks. However, only considering memories from the true experiences in the process of imagination could limit its advantages. Inspired by the memory prosthesis proposed by neuroscientists, we present novel model-based reinforcement learning framework called Imagining with Derived Memory (IDM). It enables the agent to learn policy from enriched diverse imagination with prediction--eliabiiity weight, thus improving sample efficiency and policy robustness. Experiments on various high-dimensional visual control tasks in the DMControl benchmark demonstrate that IDM outperforms previous state-of-the-an art methods in terms of policy robustness and further improves the sample efficiency the model- method.
Introduction
Model-hased Reinforcement Learning (MBRL) approaches benefit from the knowledge of model, allowing them to summarize the agent's past experiences and foresight the future outcomes in compact latent space branches, including 1) Dyna-style algorithms, 2) sampling-basee based methods, 3) value expansion, and 4) backpropagation through paths. First, the Dyna-style algorithms generate virtual interaction
Previous approaches can be generally categorized into four
"Yuzheng Zhuang is the corresponding author. Yao Mu conducted this work during the internship in Huawei
Noah's 's Ark Lab.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

data (e.g. ME-TRPO 6 MBPO (7] to improve policy optimization. Second, the sampling-base based methods, such as 8 9 10, choose action by planning or shooting (e.g. PlaNet [8]. 9)TS Third, the value expansion approaches, such as [1] 12]. use model-augmenttdd rollouts to improve the estimation accuracy of cumulative returns. Finally, the methods of backpropagation through paths employ the analytic gradient of state value through the dynamic transition to learn the policy.
13 14
For the tasks with high-dimensional sensory inputs, backpropagation through paths more com- putationally efficient compared with conventional planning algorithms, which generated numerous rollouts for the best action sequence selection procedure. Dreamer 4 as landmark of such methods, achieves satte-of-the-eee (SOTA) performance on visual control tasks. It summarizes the agent's 's high- digensiiensi experiences into 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-sensitivy policy and inefficient learning. Recently, neuroscientists demonstrated that "memory prosthesis", which transforms the brain's 's activity patterns to the electrode signals to stimulate the human's 's hippocampus [ [6], could facilitate humans mem- ory 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 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 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.
states of the latent trajectories, which are sampled from
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 model-based policy optimization framework, umifying 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 the future.
Experiments on DMControl [8] 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 (known learned) model of the environment. We denote time step as a state t as s, E S, action at t as a, E A, reward function r(s;, a.), policy (st) and a world model pe to characterize the process of interacting with the environment. The goal is find a policy parameter pp" which maximizes the long-horizon summed discounted rewards represented by a parameterized value function, Uh (8x) E x ( (r:) with parameter de and the discounted factor 7.
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, 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 = {oo, 01,...,or}, action sequence as â = {ro, -,,-.., pr), and the corresponding latent states ğ RSSM assumes that the latent state consists of 2 (st, where St, he are the probabilistic and deterministic variables respectively. RSSM consists of four key components, including (1) the representation model p(z: 2 , ,t),, which encodes observations and actions create continuous vector-valus valued latent states with Markovian transitions, (2) the transition model ((2 1 which is built by recurrent neural network, and is to predict the future latent states with historical memory and actions, (3) the observation model g( (2). which is used to provide reconstruction learning signal, and (4) The reward model grralz,) predicts the rewards given to the latent states. We use p to denote distributions that generate samples in the real environment and g 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.
{ao; 01,0011....000. (T), rewards
{20, .1...., ,T
3 Robust World Model for Derived Memory Generation
Based on RSSM, diverse memory vectors could 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 La and Figure [bb we add an random noise to the latent state 24 (he, 9,) the derived latent state zf (hG, s,) in RSSM results in low-fidelity reconstructed images, which could not be mapped to the real physical states.
Original
RSSM
RSSM
Our
Ours
(a) Add random noise on RSSM and R2S2 (ours) in Walker-
Original
original
RSSM
RSSM
Ours
Ours
(b) Add random noise on RSSM and R2S2 (ours) model in Finger-Spin
Figure 1: Comparison of reconstruction results of the derived latent states with different world models. The first row the original image inputs, the second row the images decoded from the derived latent states by RSSM, the third row is the images decoded from the derived latent states by R2S2 model (ours). (Note that the shown results are randomly selected from test episodes, and the derived latent states are generated by adding Gaussian noise obeys (0, 0.5) to each dimension of (he st)).
To alleviate this issue, we propose 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 εn is a noise that obeys (000, on), , R is the derived hidden state which adds the noise &n to the hidden state hG, & is noise that obeys N (0, ox),, s is the derive stochastic

---

state which adds the noise &. to is the reconstructed image from 2 (he, st), and of is the reconstructed image from 2/ (hi, st).
The components of the R2S2 are optimized jointly to increase the variational lower bound 19 with a decoding constraint. The variational lower bound Jar includes reconstruction terms for observations Jó and rewards Jhz and KL regularizer JD representation model,
44. The expectation is taken under the dataset and
min TM- = Ep Σ (Jo + Jh + xb)) S.L.DLL p((o)12)). s
(I)
J/ = log (()())))
=og JD = 33DKL (2)1211,, ( ( (011))11))
As demonstrated in Figure I. 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 la and Figure [b] after adding noise on the original latent state 2, (h4, 8:) in R2S2 model, the reconstructed images with derived latent state of (K₁ s() are sharper than RSSM, and are demonstrated different states of motion compared to the the reconstructed images with original latent state 2n (he, ). In summary, the decoding constraint used in the R2S2 model aims to guide the derived memory in latent space to meet basic condition that the neighbouring latent states from real data should also follow the real physical states.
Policy Optimization with Derived Memory
By leveraging the R2S2 model, we propose 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 Ça ~ (0, o on the hidden state hú and E. ~ NN(, on the stochastic state s, ~ Ni,,,, of) in the original memory extend the agent's memory which includes diverse embedded historical information. Thus the agent could learn behaviours with the trajectories an enlarged latent space, rather than just the agent's real experiences. The Gaussian noise 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 latent state 8; 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 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:
h Nh,,, of)) Trajectory imagination:
st ~~>Ns (o. + +z,)
(2)
h st p(hi, si
de~ ((alls) k/)
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 trajectory evaluator. Every imagined trajectory used with a reliability multiplier given by the evaluator. the region where the model has reliable predict accuracy, the derived memory should be able to improve

---

03
Real image sequenoes embedding
Evaluator
d2
Imagined lakent sequences
12
ht
Reliable Score
hg
Actrr Critic
(b) Training Evaluator
p(o,
p(11)))
π
G
a2
Da (z)pp(olz))
Ze (p)
0g
Framework of IDM
(c) Decoding Constraint
Figure 2: The overall framework of the IDM. The original memory 21 = (h1,s)) $1) is embedded from high-dimensional state, and the derived memory (h1- s1) generated by the transformation from the original memory. The decoding constraint is used to avoid reconstructed distortion. Latent trajectories are imagined from both the original memory and the derived memory. Each trajectory is given reliability weight by the evaluator. The optimal policy π is learned by the enriched imagination with prediction reliability weight under the actor-critic framework.
the policy robustness and sample efficiency significantly. We penalize the weight of the trajectories in the region with poor model prediction. The trajectory evaluator aims 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 real sequences Seqeal [cy, a, 2r+1), 1 r=
use the latent model to predict the overshoot trajectories from the same start points as the fake sequences Segtake = [Zt, ae, 4 [2,, a, 2 2++1]=11+1 Train the evaluator to output the probability of whether real trajectory, which equal maximize the following objective function
C+N N
Jeva = ( D
Σ
[log (E (Z,, a,, 27+1)) + log (1 - E (λ,, ar, & 21-1))1
(3)
where E.) is approximated by the combination of the MLP and the outermost Softmax and 21 = & since 2, is the start point of the overshoot trajectory.
4.2 Policy Optimization
The policy (2+) is optimized by the imagined trajectories {2r, ar, rr rr 14+ We estimate every state's 's value in the imagined trajectory to trade off the difficulty of long-
back-propagation and efficiency of gradient utilization. That means a single trajectory {2r, ar ) +++ generates many sub-trajectories with different rollout steps. Every sub-trajectory is used update the actor and critic
with different reliability weight w(2{) given by the trajectory evaluator and is normalized in batch.
W = Elt of 2
- (Wori (Dari,derri Dderij
((+))
(4)
w((+))
i= + W(+))
N
where the batch size, Doi represents the trajectories imagined with original memory, and deeri represents the trajectories imagined with derived memory.
The value model Up (2+) estimates the expected imagined rewards that the action model achieves from each state In attempt trade off bias and variance, calculate the target state value by
5

---

TD-À method [2]] as in Dreamer, which is an exponentially-weie average of the estimates for different & to make the estimation robust to different prediction horizons.
( Σ
(z+)
VN (2+)
(5)
Vă (2+) = Egp 40
Σ
On - (zN)
(6)
h = R,
where the V estimates rewards beyond steps with the learned value model.
The critic is updated to regress the targets around which we stop the gradient [21] The reliability weight used in the critic updating, which is called weighted regression of target value. The objective for the critic og (zr), in turn, is to regress the target value
Σ w
min Σ
l0w (z+) - VA (2+) 11
2
Different from DDPG 22 and SAC 23 which only maximize immediate Q-values, we leverage gradients through transitions and update the policy by backpropagation through the value model.
The objective of the actor with policy The (2r) is to
N
+++ H
max
Σ
Σ
VZ (2+)
(8)
Since all steps implemented neural networks, analytically compute Jactorr by stochastic backpropagation. We use reparameterization for continuous actions and latent states,
ar tanh (Mip (s+) + o (8)) e e~ c(0,1.
4.3 Theoretical analysis of IDM Framework
We give the upper bound of the value estimation error for MBRL under the IDM framework based on (24). Denote the policy distribution shift between the current policy and the data-collecting policy TD as max, Drv (T [T] S ex by the maximum total-variation distance and the model gen- eralization error cn max Eg.:D [Drv (p(2'r 2, 2,a) Ipp' (z',r 2. where p p=2'rr 2, a) the real state transition distribution and p' (2',r 2, )) is the estimated state transition distri- bution. Under the IDM framework, if the model error under the updated policy is bounded by max Ezr(Ders,_)r _D=se) [Drv (p(2', r 2, (z'((2', 21 w(2) given by the evaluator is bounded by max { [[((2) - p(z', rr 12,0000 2.a)/po /po ((z', 2, a.))}} s Ew. then the returns estimation error upper bound is
I 6m and the reweighting coefficient
branch [x]]
n755] S 2rmax
+ (1 7)2 (1-7)
g+cr
+
k
2m)
(10)
dew! where ct Cu + CT de, The detailed proof can be seen in the Appedix well the IDM helps to approximate the imagined distribution p'(z', r 2, a) the real distribution p(z',r a). The trajectories imagined with the derived memory increase the sample diversity, while the evaluator plays the role of importance sampling by reweighting the trajectories by w(z). Coherent work of these two modules should effectively decrease the cu. thus reducing the estimation error of the value function. If the evaluator removed, the imagined trajectories, which are unreliable. will negatively impact the policy learning. If the derived memory module is removed, the effect of importance sampling will be weakened due the lack of sample diversity. In addition, the positive effect on sample diversity of derived memory needs to be guaranteed by the decoding constraint in the R2S2 model, which is used ensure the relevance between the derived memory and real physical states. Corresponding experiments are in Section 5.2
A.5 Ew indicates how

---

Table
Performance comparison with image uncertainty
IDM
Dreamer
DrQ
100
500 K
100
500 K
100
500
Walker- Walk
527.1 + 166
914.9 + 61
372.5 + 19
876.3 + 48
122.5 + 15
151.9 23
Walker-Run
159.2 = 32
529.3
116.1 +
464.2 - 65
58.3 + 14
105.2 25
Hopper-Stand
108.5 = 37
341.2
85.6 52
430.2 128
65.7 24
252.6 38
Cartpole-Swingun
329.6 = 58
659.4 136
303.3 + 43
423.2 66
157.3 112.6
Cheetah-Run
197.4 - 69
598.2 1:22:44 42
174.9 + 132
401.34 104
85.9 + 22
96.3 - 38
Finger-Spin
415.8 = 54
505.5 =1
283.5 + 42
374.9 + 88
145.9 + 33
436.3 + 27
Table 2: Performance comparison without external image uncertainty (under the same setting in Dreamer)
IDM
Dreamer
DrQ
100
500 K
100
500 K
100
500 K
Walker- Walk
630 + 113
959 = 22
277 = 12
897 + 49
612 = 164
921 + 45
Walker-Run
211 + 162
536 - 25
161:112
482 + 68
185 = 148
459 - 139
Hopper-Stand
216 + 73
819 + 27
44
734 + 136
106 +
789 + 38
Sartpole-Swinghn
594 + 63
824 + 63
326 + 27
762 + 27
759 + 92
868 + 10
Cheetah-Ran -un
328 + 48
631 + 103
235 137
570 + 253
344 + 67
660 + 96
Finger-Spin
623 + 76
875 +
341 - 70
769 + 183
901 + 104
938 + 103
Experiments
The IDM framework is implemented Algorithm I (see Appendix A.2 We compare with both the
SOTA of model-based methods and model-free methods. We mainly choose Dreamer as the model- based baseline to test the effectiveness of the IDM framework. As for model-free methods, DrQ [25 uses data regularized Q-learning method to learn the policy and shows significant advantage in sample efficiency compared to CURL [26). We choose the DrQ (55 as the model-free baseline.
5.1 Overall Performance
Robustness to Uncertainty: In attempt to test the agent's ability to handle environmental
uncertainty, the experiments are implemented in DMControl tasks [18] with observation uncertainty. Robotics may encounter observation uncertainties like the random swing of the input image caused by the robot's motion when the camera's 's fixed device is loose. We add a random rotation (-50, +5%) the imput image from the Mujoco environment (also used in RAD and DrQ paper as an augmented method) to simulate the possible random swing that may be experienced by robotics. Random rotation a representative uncertainty the agent may encounter in the real world. During the training process of RL, images are only collected horizontally from the simulator, but in the real world, the camera could be skewed in a different direction, which changes the angles that agents or objects appear in the input image. Under this setting, we test the performance of IDM together with the Dreamer and the DrQ. The image uncertainty is added only at test time, and the purpose of this experiment is to test the robustness of policies learned by different methods when encountering unknown image uncertainties. The effectiveness the IDM framework tackle the input uncertainty demonstrated in Table and as Figure
We also the robustness DrQ. Since DrQ's input is stack of three sequential images (3 continuous images), we implement 2 different experiments: 1) randomly rotate the 3 images by the same angle and 2) randomly rotate the images by different angles. All the rotate angles are sampled from (-50, 5"). The results show that under the first type of uncertainty DrQ's performance dropped significantly and under the second type of uncertainty DrQ almost failed. As shown in Figure the training curve (the orange one) of DrQ demonstrates that it learns behaviour successfully; however, its performance decreases significantly when it is tested with image uncertainty, as illustrated the testing curve (the blue one). The reason why DrQ performs poorly under external uncertainty is that DrQ has no guarantee to resist the uncertainty out of optimality invariant image transformations proposed in their paper, which especially constrains shifting under #4 pixels from the original image.
7

---

Walker walk
Cartpole swingup
Walker- run
1000
800
800
Average return
600
Average return
600
400
Average
400
200
Dreamer
200
Dreamer
Dieamer
IDM
IDM
IDM
0
2
2
4
Steps
105
Steps
.05
Steps
Les
Finger-spin
Cheetah-run
Hopper-stand
10000
400
Average return
Average riturn
Average
300
400
200
100
Dreamer
Dreamer
200
Dreamer
IDM
DMM
IDM
Stops
Steps
Steps
Figure 3: IDM exceeds Dreamer and DrQ at visual control tasks that testing with image uncertainty (the image uncertainty is added only at test time).
Our experiments show that although DrQ can effectively learn policy in training process with the given 5 random seeds, its performance overcoming uncertainty is much worse than IDM.
Walker-Walk
Cartpole-Swingup
Walker-Run
800
Avcrage return
Average return
E00
600
600
Drq w/ uncertainty Drq wio uncortainty
Drq w/ uncertainty Dra w/o uncortainty
Drq w/ uncertainty
Avcrage
400
400
400
Drq we uncertainty
200
200
200
0.0
0.2
0.4
0.6
1.0
1.0
1.0
Steps
Steps
Steps
Figure 4: DrQ's performance decreased due to the observation uncertainty (the uncertainty is performed by the same parameter for each frame in the stack).
Sample efficiency: order to show the sample efficiency of IDM, compared to Dreamer and DrQ [25] we test the performance of IDM in DMControl environments without image uncertainty with random seeds. illustrated in Table IDM outperforms Dreamer and is comparable with the DrQ in environments without external uncertainty.
Our selection of the environment Table 1 and Table 2 follows the below principle ensure fairness: tasks in which MBRL is more competitive than MFRL and 3 tasks in which MFRL is more competitive than MBRL. The measure competitiveness is the comparison between the final performance of Dreamer and D4PG (top model-free method) with 1e9 steps shown in Dreamer paper.
Overall, the experimental results suggest that IDM outperforms Dreamer and DrQ in terms of robustness to uncertainty and further improves the sample efficiency of the model-based method, which proves that our proposed IDM can facilitate learning better behaviour with limited experiences.
5.2 Ablation Studies
In IDM, there are three key components: decoding constraint, derived memory and trajectory evaluator. We implement 4 groups of ablation experiments which are tested in DMControl tasks with

---

observation uncertainty. Note that the derived memory is based on the decoding constraint. If we keep the derived memory module without adding decoding constraint, is theoretically invalid.
(1) IDM without the evaluator(w/o EVA): In this setting, latent trajectories are imagined with derived memory but without the trajectory evaluator. Although it performs better than Dreamer in Walker-Wail Walk, it has lower performance in Hopper-Stand due to lack of reliable guarantee for the derived data.
(2) IDM without derived memory (w/o DM): In this setting, the constraint for RSSM and the trajectory evaluator kept. As illustrated in Figure [5] it performs better than Dreamer but is insufficient compared with IDM. For instance, at the 100k steps of Walker walk, its performance is 26.6% better than Dreamer, but IDM is 41.6% better than Dreamer.
(3) IMM without decoding constraint and derived memory (w/o CST and DM): this setring, only the trajectory evaluator is kept, which is used to avoid the damage caused by the trajectory with low prediction accuracy. The experimental results show that its performance is basically similar to Dreamer, which means that if we only use to reweight the latent trajectories imagined with original memory, makes little difference.
(4) Dreamer with doubled batch size: We compare IDM with Dreamer with two times of batch size than IDM than Dreamer, but IDM also outperforms it.
as shown Figure Doubled-Dreamer shows less improvement
Overall, all the components in IDM critical, and they complement each other. The importance of different components depends on the characteristic of the environment. For example, in Walker- Walk, which is relatively stable bipedal walking environment compared to single-legged robot, the generated derived memory does not hugely diverge from the real states experienced by the agent. Therefore, shown Fīg the derived memory provides a major contribution. On the contrary, due to the single-legged robot's instability, the derived memory generated by Hopper-stand usually contains poor quality states. Thus, as shown in Figure c(c), the evaluator plays a significant role in this environment, but not in Walker. The ablation results is in some difficult tasks, such as Walker-run (the agent needs to stand up first and then accelerate to run), the generated derived memory is not as helpful to improve the performance as the other environments, so does the evaluator. As shown in Figure b)), there is little difference in the results of ablation. This result also enlightens us to find better way to construct derived memory in future work.
Walken Wa.k
Walken Ru
Hopper-tand
600
800
800
500
Auergge retum
600
600
400
400
w/o DM
200
IDN w/ CST and
200
IDM w/o EVA
1.00
200
IM
Doubled
Sceps (a)
Staps
Stops
(b)
Cartporle-Siinge.
800
700
5.0
mamm
Avarqqe roum
hvergg retum
400
400
500
300
300
20a
200
200
100
100
Shopp
1c5
Steps
1c5
Stops
1e5
Figure Contribution each component IDM to the performance improvement.

---

6 Related Works
(I) Model-based policy optimization: Prior works can be divided into the following categories. including Dyna-like algorithms,, value expansion, backpropagation through paths, and sampling-basee planning. Dyna-like algorithms alternate between model learning from environmental interaction, data generation, and policy improvement by model-free methods [5], e.g.. ME-TRPO [6]. Value expansion algorithm utilizes the dynamic model to improve the estimation accuracy of the cumulative return [.I. The backpropagation through path method is built upon using the analytic gradient of state value backpropagate through the dynamic transition. Typical algorithms in this category include algorithm planning is to choose the best action by large number of samples and regard it the objective for the policy network. Representative algorithms include the cross-entropy method (28) in continuous space, which is used in PlaNet [8] and PETS [9] and MCTS (10) in discrete space. MuZero further incorporates MCTS and achieves remarkable performance on Atari, and board games [29].
Dreamer 44) PILCO [27], and SVG [1]. Dreamer is the most representative model- for visual control tasks. The idea of sampling-basse
(2) Latent world model: In MBRL for high dimension input tasks, the latent world model offers a flexible way to represent key information of the observations. World Models [ learn latent dynamics in a two-stage process: representation learning and latent dynamics learning. However, lacks the coordination of the two processes. PlaNet 2 proposes recurrent stochastic state model (RSSM). which learns them jointly. Dreamer 4 utilizes the RSSM to make the long-term imagination. [30) represents the model uncertainty by neural network ensemble to approximately infer model posteriors and further improves the performance of PlaNet. Dreamer V2 [1] proposes discrete world models for discrete visual control tasks like Atari games. BIRD prediction accuracy by maximizing the mutual information between the trajectories imagined with memory and real trajectories but does not consider the diversity of memory. Recently, contrastive learningbbased approaches are gathering increasing attention. It has been applied to learn latent world models 4 32 33 34] motivated from different perspectives. Hafner et al. 4 try to use contrastive learning as alternative to image reconstruction, however, the contrastive learned agent gives worse performance compared with the one learned by image reconstruction. Dreaming (33] derives an InfoMax objective of contrastive learning from the evidence lower bound of Dreamer and removes the decoder to solve the object vanishing problem of the representation methods based on image reconstruction. CVRL 34 emphasizes the strength of contrastive learning handling complex visual observations and uses a hybrid training scheme that includes a model-based scheme (Dreamer) and model-free scheme (SAC). Since the key idea of IDM is not in conflict with contrastive learning. we can extend our framework by incorporating advanced contrastive representation learning to further improve the overall performance, which is left as future work.
B improves the
(3) Data augmentation for reinforcement learning: Data augmentation has been investigated in the context of RL to improve generalization and sample efficiency. CURL (26] minimize the contrastive loss between an image and its augmented version to learn a better latent representation. RAD (35) attempts to directly train the RL objective on multiple augmented views of the observations without any auxiliary loss. DrQ (25) aims to solve the distribution mismatch problem in off-policy RL, which utilized random cropping and regularized Q-functions in conjunction.
Conclusion
We present a novel model-based framework, called Imagining with Derived Memory (IDM), aiming to improve the policy robustness and the sample efficiency by effectively enhancing the agent's imagination diversity. IDM optimizes parameterized policy with the extended imagined trajectories from both original and derived memory by propagating analytic gradients of multi-step values back through learned latent dynamics. It outperforms previous methods in terms of the robustness to uncertainty and further improves the sample efficiency of the model-based method on a variety of challenging continuous control tasks with image inputs. Other methods to generate derived memory, such as employing different assumptions for latent noise distribution, will be investigated in the future. Furthermore, since method is compatible with representation learning techniques. extending our framework by incorporating advanced contrastive representation learning (e.g., RAD, CURL, and Dreaming) to further improve the sample efficiency could also be left future work.
10

---

Broader Impact
Sample efficiency and policy robustness are key factors for the development real- weal-world applications on automatic control, such as robotics and autonomous driving. The proposed IDM algorithm is situated in model-based RL and learns behaviours via imagination with derived memory, which enables the agent to tackle unseen situations by efficient enhancement on imagination. Thus it further improves policy robustness and sample efficiency over existing works, which provides a broader prospect for real-world applications. The technique to generate and utilize the derived memory is wide applicable computational efficient and could motivate both research and the industrial community to refine the presented method for addressing the uncertainty in different sources of collected data and making up the gap between the real world and training environments. In the long run, this paper contributes the foundation for realizing robust artificial intelligence, which will improve efficiency, reduce cost and risk in industrial operation, agriculture and also facilitate people's life by enhancing daily life intelligence levels.
Acknowledgments and Disclosure of Funding
The authors would like thank the anonymous reviewers for their valuable comments and helpful suggestions. The work is supported by Huawei Noah's 's Ark Lab; Ping Luo is supported by the General Research Fund of Hong Kong No.27208720.
References
[1] D. Ha and Schmidhuber. Recurrent world models facilitate policy evolution. 2018.
[2] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Leannng latent dynamics for planning from pixels. In International Conference on Machine Learning, pages 2555-2565, 2019.
[3] G. Zhu, M. Zhang, H. Lee, and C. Zhang. Bridging imagination and reality for model-based deep reinforcement learning. 2020.
[4] Danijar Hafner, Timothy Lillierap, Jimmy Ba, and Mohammad Norouzi. Dream control: Learning behaviors by latent imagination. arXiv preprini UrXiv:1912.001 2019.
[5] Richard S Sutton, Csaba Szepesvári, Alborz Geramifard, and Michael P Bowling. Dyna- style planning with linear function approximation and prioritized sweeping. arXiv preprini arXiv: 1206.32206.385 2012.
[6] Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Picter Abbecl. Model-ensemble trust-region policy optimization. arXiv preprint arXiv: 10592, 2018.
[7] M. Janner, J. Fu, M. Zhang, and S. Levine. When to trust your model: Model-based policy optimization. 2019.
[8] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. arXiv preprint arXiv:1811.0455 2018
[9] Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. In Proceedings of the 32nd International Conference Neural Information Processing Systems, pages 9759944777 2018.
[10] David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7666):5474333333) 359, 2017.
[11] Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael I Jordan, Joseph E Gonzalez, and Sergey Levine. Model-based value estimation for efficient model-free reinforcement learning. arXiv preprint
2018.
11

---

[12] Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample- efficient reinforeement learning with stochastic ensemble value expansion. In Advances in Neural Information Processing Systems, pages
2018.
[13] Michael Fairbank. Reinforcement learning by value gradients. arXiv preprint arXiv:0803.3339
2008.
[14] Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap. Tom Erez, and Yuval Tassa. Learning continuous control policies by stochastic value gradients. In Advances in Neural Information Processing Systems, pages 2944-2952, 2015.
[15] G De Lavilléon, M. M. Lacroix, L.. Rondi-Reig, and K. Benchenane. Explicit memory creation during sleep demonstrates a causal role of place cells in navigation. Nature Neuroxcience, 18(4):493, 2015.
[16] Robert E Hampson, Dong Song, Brian S Robinson, Dustin Fetterhoff, Alexander S Dakos, Brent M Roeder, Xiwei She, Robert T Wicks, Mark R Witcher, Daniel E Couture, et al. Developing a hippocampal neural prosthetic to facilitate human memory encoding and recall. Journal of neural engineering, 15(3):036014, 2018.
[17] Karaiskaki and Xenia Andrea, Anastassiou- pages 1-3. Springer International Publishing. Cham, 2019.
Anastassoo--Hadjichrllaaattt Hadjicharalambous. Inceptive and Derived Memory,
[18] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmalcki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv: 180100060, 2018
[19] Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An in- troduction to variational methods for graphical models. Machine learning, 37(2):183-233 1999.
[20] Bertsekas and DimitriP. Constrained optimization and Lagrange multiplier methods. Con- strained optimization and Lagrange multiplier methods, 1982.
[21] Aravind Srinivas, Allan Jabri, Pieter Abbeel, Sergey Levine, and Chelsea Finn. planning networks. arXiv preprint 1 1804.00645, 2018.
Universal
[22] Timothy P Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint 1509.0297/99 2015.
[23] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor, arXiv preprint 1rXiv:1801.01290, 2018.
[24] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine When to trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems, pages 12519-12530, 2019.
[25] Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint wrXiv:2004.13649 2020.
[26] Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised represen- tations for reinforcement learning. arXiv preprint arXiv:2003.06441 2020.
[27] Mare Deisenroth and Carl E Rasmussen. Pilco: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference machine learning (ICML-11), pages 465 4722, 2011.
[28] James Bekker. Applying the cross-entropy method in multi-objective optimisation of dynamic stochastic systems. PhD thesis, Stellenbosch: Stellenbosch University, 2012.
[29] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604-6666 2020.
12

---

[30] Masashi Okada, Norio Kosaka, and Tadahiro Taniguchi. Planet of the bayesians: Reconsidering and improving deep planning network by incorporating bayesian inference. arXiv preprint urXiv:2003.000 2020.
[31] D. Hafner, T. Lillicrap, M. Norouzi, and J. Ba. Mastering atari with discrete world models. 2020.
[32] Thomas Kipf, Elise van der Pol. and Max Welling. Contrastive learning of structured world models, arXiv preprint orXiv:/911.122 2019.
[33] Masashi Okada and Tadahiro Taniguhii. Dreaming: Model-based reinforcement learning by
latent imagination without reconstruction. arXiv preprint arXiv:2007.14533, 2020.
[34] Xiao Ma, Siwei Chen, David Hsu, and Wee Sun Lee. Contrastive variational reinforcement learning for complex observations. arXiv preprint 1rXiv:2008002430 2020.
[35] Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. arXiv preprint arXiv:2004. 14990, 2020.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section 5.2
(C) Did you discuss any potential negative societal impacts of your work? [Yes] See
Section 7
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See Section 4.3
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix A.3
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or as URL.? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Appendix A.2]
(c) Did you report error bars (c.g., with respect to the random seed after running experi- ments multiple times)? [Yes] See Section
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See Appendix [A.4
4. If you are using existing assets (e.g., code, data, models) or curating/relcasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any assets either in the supplemental material URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you's 're
using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "kvBdO2w7c69", "reviewer": "Reviewer_URjo", "summary": "Summary\n\nIn this paper, the authors consider the latent space in imaginary space. \nA GAN styled latent space noise is trained to add more diversity and robustness to imaginary planning.\n", "questions": "", "limitations": "", "rating": 6, "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": "\n\nMain Review\n\nPros:\n1. The algorithm is neat and well explained.\nGood related work is provided in the paper, and the algorithm is well explained.\n\n2. Theoretical analysis is provided in the paper, making the algorithm more convincing and providing good intuition and inspiration to future research.\n\n3. The problem studied is challenging and interesting.\nImage space continuous control problem is considered one the hardest problems in RL or robotics. This paper will contribute to this line of research.\n\nCons:\n1. The novelty is lacking.\nThe algorithm’s contribution is relatively incremental compared to the existing research, \ngiven that it is built upon Dreamer and the elements that are novel is the added noise.\nBut still it remains unclear how robust or scalable the algorithm is (from the experiment section),\nas shown in my later comments, \nthe experiments have a major flaw of not able to reproduce the performance of the original paper which open-sourced their code-base.\nThat being sad, it is hard to draw any conclusions of how good the proposed algorithm is in terms of practical performance. (updated according to AC's request)\n\n\n2. The experiment section is not convincing.\nMostly I believe the only dreamer is used as the baseline. \nAnd the performance of dreamer on different environments tested is having a large gap between the performance reported in the original paper.\nFor example cartpole-swingup is around 800 in the original dreamer paper, but was shown around 400 in the given paper. The case is consistent in Walker-Run, Finger-Spin, Cheetah-Run, Hopper-Stand.\nConsidering Dreamer is open-sourced, I wonder if there's a good explanation for the performance gap.\n\n\nUpdated after the rebuttal.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "limitation addressed and no obvious negative societal impact found in the paper.", "review_text": "", "extra_scores": {"ethical_concerns": "no ethical concerns.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2 hours"}, "extra_text": {}}, {"review_id": "dum_KxAPzP4", "reviewer": "Reviewer_S88K", "summary": "The authors introduce a variant of the model-based Dreamer algorithm with an added constraint at training time: that perturbed latents (with Gaussian noise) should lead to similar observation reconstructions. This is leveraged to augment the real trajectories with imagined trajectories that start from nearby latent states. Finally, a learned discriminator of real vs imagined latent trajectory is employed to reweight the trajectories in the policy optimization objective. The approach is tested in continuous control domains. \n", "questions": "", "limitations": "", "rating": 6, "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": "The approach proposed by the authors is reasonable and novel but the writing is unfortunately a bit hard to follow at times and the language is imprecise/inconsistent in places.\n\nFor example, it would make sense to be more precise about what constitutes real and derived memories (cf L121) given this is a crucial aspect of the paper. Sometimes this seems to be interchanged with real and “overshoot” (or fake) sequences or trajectories. The notation around these concepts doesn’t help:\n- L140: Seq_fake is defined using $\\hat\\tau_{t+1}$, but it’s not clear what this is.\n- On page 5, $z$ sometimes refers to a single latent, and sometimes seems to refer to a full trajectory (e.g. Eq 4).\n- $\\mathcal{E}$ takes a sequence as input in Eq3, and then is given a transition tuple as input in Eq 4.\n\nThe exposition in Alg 1 in the appendix does clarify some things, but the notation should be made consistent.\n\nThe experimental work compares the proposed approach (IDM) to Dreamer and DrQ on some Mujoco domains. In the normal conditions (Eq 2), IDM performs similarly to DrQ (if anything DrQ performs better since the difference between IDM and DrQ in the Walker tasks does not seem to be statistically significant). Nevertheless, the fact that it consistently outperforms the baseline (Dreamer) suggests that the mechanisms of the approach are beneficial in these domains. The provided ablations reinforce that point. \n\nThe experiments with the randomized visual inputs are a bit odd. IDM performs a lot better than Dreamer+DrQ there.\n- Is the randomization done during the training as well or only at test time?\n- It’s not clear that perturbing latents would match the image rotation perturbation? So why is IDM performing better here? \n\nOverall, the approach may be promising but I don’t think the paper is polished enough in its current form to be accepted.\n\nMore comments / questions:\n- The transition for the deterministic latent variables h is not explained in the paper. \n- Fig 1: left/right should be the same top trajectories, otherwise it’s not really a fair comparison. For example the trajectory in a) shows a falling walker vs a standing walking in b), small perturbations in these regions might have completely different effects. \n- Many parts of the paper repeat the equations from the Dreamer paper. In this case please say “as in Dreamer” rather than “similar to Dreamer” if the approach is unchanged (Eq 5,6). Similar comment for L160-161, this is not specific to your approach.\n- The relation of the analysis in 4.3 to [31] is not mentioned. \n- How was the noise variance for the latent perturbation chosen? Was the regular transition noise variance similarly tuned?\n\nMinor:\n\n- L162 object -> objective\n- Sentence on L33 should be rephrased.\n- Figure 3 caption doesn’t make sense.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "ZvE2qpcE4D", "reviewer": "Reviewer_EMUa", "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.\n\nTo 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.\n\nThe 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.", "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": "Originality: The approach builds upon prior work, primarily Dreamer, with the key addition being the use of imagined rollouts from noisy latent states to augment the training data. Other additions such as the decoding constraint and discriminator for scoring trajectories are needed to stabilise learning when using these derived rollouts. The paper does a good job of establishing related work and connections to prior work. I find the general idea of augmenting training data with diverse imagined rollouts quite interesting and useful,  but there are several open questions that could have been addressed. \n\n1. The approach only tries a simple way of doing this which is to add random Gaussian perturbations to the initial state latents. Are there other ways in which diverse imagined trajectories could be generated? Are there different perturbations that could be applied (e.g. correlated perturbations in different directions in the latent space)? It would have been nice to see some comparisons or a discussion on this topic.\n2. Similarly, could the decoding constraint be made proportional to the amount of noise added? i.e. latents that are close in magnitude need to have a stronger KL constraint than those that are perturbed more heavily.  \n\nQuality & Clarity: IDM shows significant improvements in performance upon Dreamer and DrQ on the control suite tasks with observation uncertainty and similar or slightly lower performance on some tasks without uncertainty. The ablation study also provides some interesting insights; the fact that the scoring of trajectories is primarily needed for complex tasks such as the hopper where the model can struggle to produce plausible rollouts compared to walker-walk where the dynamics are more straightforward is a very useful insight. The theoretical analysis mostly mirrors that of Janner et al. with the addition of the weight from the discriminator providing a slightly different bound. While the results are reasonable the paper could be significantly strengthened further by running a few additional experiments and adding some discussion to provide further insight. The writing could also be further improved to improve consistency and clarity. A few questions/suggestions:\n1. IDM seems to use a very long imagination horizon (H=50) for the control tasks, which is also far larger than the sequence length (L=15) on which the model was trained. Dreamer uses H=15 (from the paper). Is there a reason why this horizon was needed? Would it not hurt to use a horizon much longer than the sequence length as it might result in spurious gradients (even with the lambda returns). It would also be useful to see an experiment where the horizon is varied; does the horizon matter for achieving good performance?\n2. What is the architecture of the MLP used as the discriminator? Does it concat all observations from all time steps to compute a score? Is there a different discriminator for each trajectory length (since Dreamer uses sub-trajectory returns)?\n3. Is there a specific reason why the random rotation perturbation was chosen? What about other forms of perturbations? Results on an additional form of perturbation would add more strength to the paper. \n4. What is the form of observation uncertainty used for DrQ in Figure 4? Is this the same random rotation for each frame in the stack or different random rotations? \n5. A more pertinent comparison w.r.t observation uncertainty would be RAD; RAD is trained with different image augmentations (including rotations). It would be interesting to see how well that performs to get a baseline number.\n6. A key advantage of IDM is that it can be used to generate significant additional diverse imaginations. The paper only explores generating a single additional imagined rollout per state, i.e. the batch size is increased by a factor of 2. Is there any reason not to generate significantly more data? It would be great to see an experiment varying the effective batch size (e.g. 1 (no additional imagined data), 2 (current IDM), 4 (3x derived trajectories, 1x standard imagined trajs), 8 (7x/1x derived/standard imagined trajs) and so on. One would expect to see some gains for further increase in imagined samples up to a point.\n7. In terms of clarity, the writing could be further improved (e.g. Eqn 2 can be better phrased to show that a_t is first sampled and used to generate h_t+1, s_t+1, and a for loop can be added to show the full imagined rollout). The notation can be made more consistent (e.g. p_\\theta is used in the main text interchangeably with p). Both these would improve readability.\n\nSignificance: The approach presented in this paper provides an interesting way to augment training data with diverse imagined rollouts to improve learning efficiency and robustness to perturbations. While the initial results are promising, more experiments, ablations and discussions are needed to further strengthen the paper before the results can be attractive to the broader research community. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper does discuss some limitations of the proposed approach but a better discussion of the strengths and weaknesses are needed and can add value to the paper; adding some directions for future work would also be helpful. Societal impact is not addressed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "8MpOaF4TW-i", "reviewer": "Reviewer_8jCq", "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.", "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": "This paper describes Imagining with Derived Memory (IDM), a model-based reinforcement learning method based on a world model.\nGenerally, the paper is written properly and intuitively. The approach is understandable.\nThe theoretical analyses also add value to the paper.\nAs a whole, the paper has good originality, quality, and clarity.\nHowever, the significance is not so high for two reasons.\n1) The theoretical explanation about why and how the proposed method, IDM, contributes to the improvement of the performance is missing. Therefore, it is hard to recognize the theoretical contribution.   \n2)  The experiments are limited in tasks and analytical discussion is missing (i.e., they only show RL performance). The experiments is not sufficient enough to support all the authors' claims.\nAlso, the following points written in limitations should be considered.\n\n \n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "To improve the paper further, please consider the following comments.\n\n(1) Contrastive learning \nIn world-model learning, contrastive learning-based approaches are gathering attention. They perform SOTA in many tasks, i.r., Dreamer has not already been SOTA. The paper refers to the topic just in conclusion and implicitly in 6.3. It's better to mention the approach in 6.3 explicitly.\nAlso, CVRL [1] is worth mentioning. \n\n[1] Xiao Ma and Siwei Chen and David Hsu and Wee Sun Lee, Contrastive Variational Model-Based Reinforcement Learning for Complex Observations, Proceedings of the 4th Conference on Robot Learning (CoRL), 2020\n\nThe reference to Dreaming is missing.\n\n[2] Masashi Okada, Tadahiro Taniguchi, Dreaming: Model-based Reinforcement Learning by Latent Imagination without Reconstruction, IEEE International Conference on Robotics and Automation (ICRA) 2021\n\n(2) Relevancy between the latent states and the real physical states\n\nIn Section 3, the authors describe\n\"We add a decoding constraint to the RSSM, which improves relevancy\n between the derived latent states and the real physical states, thus avoiding the distortion after decoding.\"\n\nHowever, the reason why the approach can improve the relevancy between the latent states and the \"physical\" states is not explained appropriately. The metric on images is different from the metric on a physical system in robot control.  I think there are no clear reasons and evidence that the proposed method improves the relevancy.\n\n(3) Evaluation and baseline methods\nThe method is compared with Dreamer on six tasks alone.  Further evaluation employing other baseline methods, which the authors mention in Section 6, for example, is expected to show the effectiveness of the method.\nAlso, reference [2] is pointing out that Dreamer has an \"object vanishing problem.\"  In the experiments in this paper, the tasks involving the \"object vanishing problem\" are omitted. The reason why they picked up the six tasks is described in 5.1. However, it will be fair to show the results of other tasks. To my understanding, the selection of tasks seems to be a bit biased. \n\n(4)  Related work\nThe proposed method extended Dreamer by introducing diversity, i.e., sampling-based approximate distribution, at a latent pace to improve the robustness of the world model. From this viewpoint, the following paper is related.  I suggest the authors mention the study.\n\n[3] Masashi Okada, Norio Kosaka, Tadahiro Taniguchi, PlaNet of the Bayesians: Reconsidering and Improving Deep Planning Network by Incorporating Bayesian Inference, IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS) 2020\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The submission introduces Imagining with Derived Memory (IDM), a novel extension of the Dreamer agent, a model based RL methodology using rollouts in latent space to improve policy training. 
The proposed method wants to improve both training efficiency and robustness, and at its core involves regularising the dynamic model by explicitly constraining the imagined trajectories to be smoother wrt to perturbation of the latents, i.e. perturbed latents should map to similar reconstructions.  The smooth model is used to generate a richer set of trajectories, compared to Dreamer, by perturbing initial states collected interacting with the environment, and reweighting them in the final training loss using a learned discriminator function, ensuring plausibility of the sampled trajectories.  

I expect the new elements presented in this paper to be of broad interest to the community; in particular, the way training data is augmented is widely applicable, and could motivate more groups to refine the presented technique, addressing different sources of uncertainty in the trajectory data.
One drawback of the paper it that the experimental section only presents limited (although quite promising!) results on continuous control tasks; while more experiments and ablations would make the paper more immediately impactful, reviewers agreed that the current experimental setup is sufficient to support the claims made in the submission.

Finally, the reviewers made numerous suggestions on how to improve the presentation of both the methods and results, and authors have already incorporated feedback, or committed to include it in the next iteration of the paper.
Decision: Accept (Poster)