The Surprising Ineffectiveness of Pre-Trained Visual Representations for Model-Based Reinforcement Learning

Summary

This paper presents a study on pre-trained visual representations (PVRs) to improve sample efficiency of model-based reinforcement learning methods (MBRL). It studies a range of PVRs with different architectures, pre-training objectives and data modalities on two recent model-based algorithms: DreamerV3 and TD-MPC2. The authors propose to replace image observations with frozen PVRs features, replacing the original CNN encoder/decoder with a trainable linear layer. The paper focuses on 9 diverse continuous control tasks from 2 different domains: The DeepMind Control Suite (DMC) and ManiSkill2. The results show that the use of PVRs surprisingly does not help to improve performance and has diminishing returns compared to training encoder representations from scratch.

Strengths

The paper proposes a very interesting study on the application of PVRs for two recent model-based approaches. The use of PVRs for model-based RL is a promising area of research. The performance and learning speed of model-based agents like Dreamer is very linked to the quality of representations learned by the world model. Benefiting from already pre-trained representations to reduce training time or increase final performance would be very appreciated since the training time of model-based agents can take multiple days for some advanced tasks. The paper evaluates a diverse range of pre-trained models, showing the normalized returns of each pre-trained model. The authors find that learning the encoder representation from scratch leads to better performance for DreamerV3 and TD-MPC2, contradicting expectations following previous works. The paper also experiments with out-of-distribution (OOD) versions of the environments (colors, object sizes), finding that PVRs do not help either to improve performance in this setting. The paper is clear and well written.

Weaknesses

The study on the potential effect of using PVRs on MBRL methods is limited. The paper only studies the use of frozen PVRs, where the world model gradients are not back-propagated to the pre-trained weights during training. Solving RL tasks often requires the agent to pay attention to small details and learn very detailed representations, especially for DMC tasks. It would have been interesting to study the use of PVRs without frozen weights during training. As mentioned in the conclusion limitations, the authors only study the use of PVRs for continuous control tasks. It would have been interesting to study the application of PVRs to another domain based on discrete actions (Atari games, DeepMind Lab 3D environments, Minecraft...) or more aligned with data seen during PVR pre-training.

Questions

Did the authors experiment with fine-tuning the PVRs during training instead of using frozen representations ? It would have been interesting to experiment with optimizable PVRs (without reconstruction loss in the case of DreamerV3). Maybe the use of PVRs could exempt DreamerV3 from using a reconstruction loss for learning encoder representations. The use of frozen representations may also be the cause of the observed results, all information details required to achieve good performance may not be recoverable from the compressed pre-trained representations. Did the authors experiment with reconstructing the image observations from frozen PVRs in order to study the possibility of recovering image information from representations ? Also to assess the quality of trajectories imagined by the world model. Did the authors experiment with using representations of other layers than the final layer? The usefulness of representations of pre-trained models can vary greatly depending on the chosen layer. Are the results different when using the features of inner layers or a collection of layers features ? For the in-distribution autoencoder PVR, what policy is used to generate the data used for pre-training ? Is it a random policy, a collection of policies to ensure diverse behaviors, or a pre-trained policy ? I ask the question because during training, the PVR should provide representations for the widest distribution of agent trajectories or positions. DreamerV3 has sometimes troubles to reconstruct agent situations that are seen often when training on DMC tasks. I would be ready to increase my score if additional experiments and analysis are provided by the authors.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors adequately addressed the limitations in the conclusion.

Summary

This paper contains an extensive study on the potential benefits of pre-trained visual representations (PVRs) for model-based reinforcement learning (MBRL). Given the success of PVRs for model-free RL (MFRL), there are reasons to believe that PVRs are equally beneficial for MBRL, something that has not yet been explored to the same degree. Surprisingly, tests on two different RL architectures in combination with at least eight different PVR variants indicate that PVR-based solutions rarely perform better than representations learned from scratch when applied for MBRL. In most cases, it seems preferable not to use PVRs for MBRL, which should come as a surprise to many, given recent advances in applications based on PVRs.

Strengths

Given the recent development of PVRs and the increasing number of successful applications of PVRs, it is interesting to know what the opportunities and limiting factors are for PVRs when adopted for training agents using reinforcement learning. A study with this as a focus, such as the one presented in this paper, should interest many. Besides the most important conclusion drawn in the presented study, i.e. the fact that the benefits of PVRs for MBRL can be questioned, it contains several experiments to paint a more complete picture. It is concluded that the sample efficiency of MBRL agents rarely improves when supported by PVRs compared to those with representations learned from scratch. PVRs do not seem to make agents more robust of out-of-distribution (OOD) conditions either. It is also shown that the dynamics and reward prediction errors of RL agents are similar regardless of whether representations are learned from scratch or with PVRs. Also interesting is a study on the benefits of different properties of PVRs for generalization, where the most important property was shown to be data diversity during training, which might not be that surprising. The fact that language conditioning and the choice of transformer architecture have such a little impact might be more unexpected though.

Weaknesses

From the experiments, it can be concluded that PVRs hardly benefit MBRL, but the paper does not really try to answer why there is a difference between MBRL and MFRL in this regard. Could it be that the way DreamerV3 and TD-MPC2 are modified in this study, the PVR modules $g_*$ are trained to preserve semantic information, while in effect also suppressing spatial information? Earlier studies on PVRs for MFRL, cited in the paper, do not seem to do this but instead keep a path that allows a policy to be trained with some spatial information left intact. The authors seem surprised that language conditioning is not good for generalization, since it should provide semantically relevant features. Such a claim completely misses the fact that for e.g. manipulation it is often more important for a system to know where an object is located, how it is oriented and how large it is, than what class the object belongs to or any other semantic information. Most of our visual representations are explicitly trained to be invariant to transformations that might be particularly important for the task given to the agent. It is worth noting that the experiments were conducted in simulation, not on real robots. The paper does not draw conclusions that go beyond what can be observed in the experiments, but the differences between simulated and real-world conditions are not really regarded. There are other potential benefits of MBRL when applied to real robots. With MBRL any action executed on the system can be exploited for pre-training, while MFRL is typically focused on only one particular type of task. It might be that this would not really affect the conclusions as expressed in the paper, but the experimental conditions between simulations and real-world experiments can be very large, The same is true regarding conclusions drawn from the OOD experiments. In the paper, the nature of randomizations is the same for both the ID and OOD sets, even if the two sets don’t overlap. If you were to do the experiments in the real world, the conditions would be very different.

Questions

* What is believed to be the reason why the benefits of PVRs for MBRL seem so limited, in particular, given earlier success for MFRL? * To what extent do the tested PVRs preserve spatial information relevant for tasks such as manipulation?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The fact that more experiments are needed to draw final conclusions is highlighted as a limitation, which is true indeed given the diversity of domains for which MBRL can be used and the computational demands for experiments in each such domain.

Summary

This paper conducts a thorough set of experiments to evaluate the performance of pretrained visual representations (PVR) in model-based reinforcement learning. Empirical results show that PVR performs worse than learning from scratch, which is possibly due to the large reward prediction error.

Strengths

- This paper shows how pre-trained visual representations perform in MBRL, which has been relatively unexplored in previous research. - A thorough set of experiments reveals that pre-trained representations perform worse than learning from scratch, which is surprising and could inspire future research to analyze this phenomenon further.

Weaknesses

- This paper gives some possible reasons for the performance degradation of pre-trained representations. However, it lacks sufficient evidence to further support the claims. For example, Why does the reward prediction accuracy outweigh the dynamic prediction accuracy? Why do pre-trained representations lead to information loss about reward prediction? I encourage the author to dig deeper into the specific properties of pre-trained representations used in MBRL.

Questions

- I think the performance degradation of PVR is mainly due to the fact that they are not trained on the experimental data used in this paper, as Autoencoder performs similarly compared to learning from scratch. Therefore, could you fine-tune the pre-trained visual encoder and then evaluate the performance? I guess slight fine-tuning can boost the performance a lot. - Why does VC-1 perform almost the best in maniskill (see figure 3 and figure 4)? Is VC-1 specifically pre-trained on these manipulation data? - Another possible reason is that models pre-trained on large datasets are usually hard to transfer to low-data regimes. See UniSim [1] for details. [1] Yang et al., Learning Interactive Real-World Simulators, ICLR 2024

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

Limitations have been discussed in this paper.

Summary

Pre-trained Visual Representations (PVRs) has been widely applied to many domains to improve OOD generalization and sample efficiency, including model-free reinforcement learning (RL). This paper explores the application of PVRs in model-based RL, which has not been done. Experiments on two suite of simulated control environments (Deepmind control and Maniskill2) show that current PVRs are outperformed by representations learned from scratch in terms of sample efficiency and OOD generalization. The paper further analyzes the reward error and dynamics prediction error of various PVRs and provides meaningful insights.

Strengths

- The paper is well written and easy to follow. The research questions that the paper investigate are clearly stated. The proposed benchmark is well-motivated. - The proposed benchmark has great potential in advancing research in developing better PVRs for MBRL. - The paper provides sufficient detail for the experiments allowing easier reproduction of the results. - The experiments includes a large suite of popular PVRs, improving the credibility of the conclusions.

Weaknesses

- The number of algorithms in each property category is small, which means that the performance of one algorithm can easily affect the average performance of the category, making the comparison of different categories (Figure 5) less significant. For example, the "sequential data" category and the "temporal loss" category only differ in one method (VC-1), and attributing the difference caused by this one method to the sequential data property might not be correct, since there are many other differences between VC-1 and R3M/VIP.

Questions

- Different tasks can have very different reward functions even for the same environment, and we have no clue on what the reward function is when pre-training self-supervised visual representations. Then, why should we expect a single PVR to perform well on learning good reward models for all kinds of tasks? In other words, is the task of building PVRs that can learn good reward models a somewhat intractable one?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The proposed benchmark only evaluated on two suite of simulated control environments, which, as the authors also pointed out in the paper, is (1) somewhat insufficient, and (2) very challenging for the PVRs concerned because many of them are never trained on similar data distributions (especially for DMC). Therefore, saying that PVRs are "ineffective" or MBRL seems like an overstatement to me. Strong claims need strong evidence and the evidence provided in the paper is in my opinion insufficient. Of course I still appreciate the value of the proposed benchmark in that it brings the problem of applying PVRs to MBRL to the community's attention, but I do not agree with putting such an assertive statement as the title of the paper.

Thanks for your response

After reading your rebuttal, I've decided to maintain my rating. It is still hard to determine the accurate answer to why PVR is useless for MBRL. I believe that future research should focus on how to properly incorporate PVR for better performance, like in many other machine learning scenarios. However, this paper is still limited to providing some possible approaches to addressing this challenge. I encourage the authors to study this problem further and dig deeper. That will be a high-quality paper.