Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
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. 

1. 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.
2. 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.  

Quality & 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:
1. 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?
2. 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)?
3. 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. 
4. 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? 
5. 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.
6. 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.
7. 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.

Significance: 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. 
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_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
main_review: This paper describes Imagining with Derived Memory (IDM), a model-based reinforcement learning method based on a world model.
Generally, the paper is written properly and intuitively. The approach is understandable.
The theoretical analyses also add value to the paper.
As a whole, the paper has good originality, quality, and clarity.
However, the significance is not so high for two reasons.
1) 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.   
2)  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.
Also, the following points written in limitations should be considered.

 



limitations_and_societal_impact: To improve the paper further, please consider the following comments.

(1) Contrastive learning 
In 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.
Also, CVRL [1] is worth mentioning. 

[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

The reference to Dreaming is missing.

[2] Masashi Okada, Tadahiro Taniguchi, Dreaming: Model-based Reinforcement Learning by Latent Imagination without Reconstruction, IEEE International Conference on Robotics and Automation (ICRA) 2021

(2) Relevancy between the latent states and the real physical states

In Section 3, the authors describe
"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."

However, 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.

(3) Evaluation and baseline methods
The 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.
Also, 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. 

(4)  Related work
The 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.

[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
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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. 

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.

For 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:
- L140: Seq_fake is defined using $\hat\tau_{t+1}$, but it’s not clear what this is.
- On page 5, $z$ sometimes refers to a single latent, and sometimes seems to refer to a full trajectory (e.g. Eq 4).
- $\mathcal{E}$ takes a sequence as input in Eq3, and then is given a transition tuple as input in Eq 4.

The exposition in Alg 1 in the appendix does clarify some things, but the notation should be made consistent.

The 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. 

The experiments with the randomized visual inputs are a bit odd. IDM performs a lot better than Dreamer+DrQ there.
- Is the randomization done during the training as well or only at test time?
- It’s not clear that perturbing latents would match the image rotation perturbation? So why is IDM performing better here? 

Overall, the approach may be promising but I don’t think the paper is polished enough in its current form to be accepted.

More comments / questions:
- The transition for the deterministic latent variables h is not explained in the paper. 
- 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. 
- 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.
- The relation of the analysis in 4.3 to [31] is not mentioned. 
- How was the noise variance for the latent perturbation chosen? Was the regular transition noise variance similarly tuned?

Minor:

- L162 object -> objective
- Sentence on L33 should be rephrased.
- Figure 3 caption doesn’t make sense.


limitations_and_societal_impact: Yes.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: Summary

In this paper, the authors consider the latent space in imaginary space. 
A GAN styled latent space noise is trained to add more diversity and robustness to imaginary planning.

main_review: 

Main Review

Pros:
1. The algorithm is neat and well explained.
Good related work is provided in the paper, and the algorithm is well explained.

2. Theoretical analysis is provided in the paper, making the algorithm more convincing and providing good intuition and inspiration to future research.

3. The problem studied is challenging and interesting.
Image space continuous control problem is considered one the hardest problems in RL or robotics. This paper will contribute to this line of research.

Cons:
1. The novelty is lacking.
The algorithm’s contribution is relatively incremental compared to the existing research, 
given that it is built upon Dreamer and the elements that are novel is the added noise.
But still it remains unclear how robust or scalable the algorithm is (from the experiment section),
as shown in my later comments, 
the experiments have a major flaw of not able to reproduce the performance of the original paper which open-sourced their code-base.
That 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)


2. The experiment section is not convincing.
Mostly I believe the only dreamer is used as the baseline. 
And the performance of dreamer on different environments tested is having a large gap between the performance reported in the original paper.
For 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.
Considering Dreamer is open-sourced, I wonder if there's a good explanation for the performance gap.


Updated after the rebuttal.
limitations_and_societal_impact: limitation addressed and no obvious negative societal impact found in the paper.
review_title: null

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)