Zero Shot Generalization of Vision-Based RL Without Data Augmentation

According to the paper, the strong disentanglement requires that each dimension in the latent representation corresponds to a unique source of the variation in the image. As for the weak disentanglement, a more clear and precise definition (line 216 to line 219) might be helpful for the reader to understand. For example, one natural question to ask is if several dimensions in the latent representation correspond to one unique source, is this a weak disentanglement? I would appreciate it if you could provide concrete examples for these two types of disentanglement. This would help readers better understand the key concepts and their implications for the method.

Thank you for bringing this to our attention. We have updated the manuscript (Section 3) to include a more formal definition of weak disentanglement. To address your question, if several latent dimensions each encode only one source, but the source is the same (i.e. the latents are redundant), we would not consider this weak disentanglement. By defining weak and strong disentanglement, we wish to make a distinction between latent spaces where one or more but not all dimensions encode multiple sources i.e. weak disentanglement and latent spaces where each dimension encodes at most one source i.e. strong disentanglement. Consider a 5-dimensional latent space for a dataset containing five sources. In this example, $z=<(\hat{s_0}, \hat{s_1}), (\hat{s_2}), (\hat{s_3}), (\hat{s_3}), (\hat{s_2}, \hat{s_4})>$ would be weakly disentangled because dimensions 1 and 5 encode two sources. $z = <(\hat{s_0}), (\hat{s_2}), (\hat{s_1}), (\hat{s_3}), (\hat{s_4})>$ would be strongly disentangled, because each dimension encodes for only one latent source. We claim that, assuming linear independence between the sources, the latter is ideal for generalization because it achieves the greatest representative power of the data distribution.

And the relationship between disentanglement and generalization ability is not so clear to me. According to the line 233 to line 235, if I understand correctly, weak disentanglement can also lead to generalization. Is strong disentanglement a necessary condition for generalization or is weak disentanglement sufficient? If not, please explain more about the benefits and drawbacks of using strong disentanglement vs. weak disentanglement?

In some cases, weak disentanglement is sufficient. Data augmentation methods achieve a type of weak disentanglement that is sufficient for generalization on those specific benchmarks, as demonstrated by the SVEA baseline. These methods remove distractor information and only model task-relevant variables in the latent space, and so there is at least a partial factorization between task-relevant and task-irrelevant sources.. However, there is no guarantee that task-relevant sources i.e. agent dynamics are disentangled. In some cases this is an ok assumption to have, particularly when the task-relevant sources are not linearly independent. However the main drawback as we point out in the paper is that this approach requires an excessive amount of additional data and finetuning for training stability just to achieve weak disentanglement. We simply wish to show that there exists a body of work in the disentangled representation learning literature that achieves strong disentanglement in a much more data efficient way by exploring algorithmic and model changes rather than the brute force approach of collecting more data. Strong disentanglement, definitionally, will achieve the factorization of task-relevant and task-irrelevant variables just by trying to factorize all sources of the data distribution. There are, of course, downsides to this approach, in that it requires more inductive biases in the model architecture, it traditionally requires one to know the sources a priori (although we found a way around this issue through empirical analysis), and it struggles to capture temporal information.

Beyond being robust to distribution shifts, strong disentanglement is thought to be one of the components of compositional generalization. For example, if a human is presented with a red apple, and is able to factorize “color” from “object”, it is nontrivial to imagine a purple or a black apple despite not being presented with this data (and assuming those colors have been seen in other contexts). This would be extrapolative generalization, whereas traditionally when the deep learning literature refers to generalization, it refers to interpolative generalization. This may be a promising future research direction in other ML fields, such as data efficient generative modeling, world modeling for agents, etc. capable of extrapolative generalization.

What if you vary distracting intensity, how would that change the relative results, especially compared to SVEA.

Increasing the distracting intensity increases the degree of camera perturbations. We expect both methods to perform proportionally worse. We are running this experiment and will notify the reviewers when it has been added to the manuscript.

I also notice it perform worse than SVEA in the training environment, can you add interpretations on that?

We have several hypotheses for why this might be the case. ALDA’s latent space is discretized to a finite number of bins, whereas we are working with high-dimensional continuous control tasks. This discretization, while necessary for achieving disentangled representations and performing associations at test time, may cause approximation errors. In addition, as mentioned in the discussion section, the disentangled representation literature has primarily focused on using reconstructions in (V)AE’s with other soft constraints to disentangle the latent space, and thus struggles to encode temporal information across images. While we provide an initial solution in the form of a 1D CNN layer to capture temporal information, it is perhaps not as efficient as a method that can both capture and disentangle the temporal information. However, we have yet to design or come across such a method.

In figure 6, top half, the robot starts from different pose when traversing the latent code. what if it starts from the same as in the bottom half

The bottom and top figures are different training runs. The initial configuration of the robot does not affect our ability to perform latent traversals.

What if we compare the proposed method to data augmentation method that uses simpler augmentation, such as color jitter or random crop? For example

We are running SVEA with the “random conv” augmentation and will provide the results to the reviewer once they are available.

Can you combine the proposed approach with data augmentation? How would you do that?

Yes. Rather than applying color distortions to the data, we train directly on the “color hard” environment and show qualitative results in Figures 6 (bottom) and 8. This effectively simulates training on additional data the same way applying color distortions would. To clarify, we mean data augmentation in the sense that our algorithm will scale with more varied data. However, it is unclear whether certain specific data augmentation techniques, such as randomly overlaying images from a different data distribution onto the training distribution, which results in an unnatural image, is something that ALDA can effectively model. We have updated the manuscript to clarify this point.

Can you combine disentanglement / association with dynamics? How would do you that?

We assume the reviewer is referring to the agent dynamics. In this case, ALDA is likely doing this with the agent dynamics and not just task-irrelevant variables. The sources of the data distribution for DMControl include agent-dynamics such as rigid-body orientations. We see empirical evidence that ALDA learns to disentangle and gain fine-grained control over individual rigid bodies in Figure 6, Appendix A.2, and Appendix A.3.

More proof of ALDA's generalization ability would be obtained by extending the studies to more varied settings, such as Procgen. Procgen's randomly produced levels and diverse game environments present distinct challenges, unlike those found in the DMControl suite. These settings may demonstrate the efficacy of ALDA's generalization across diverse activities, offering a comprehensive assessment of its adaptation to novel and intricate situations.

We agree that extending the studies to more settings such as Procgen would help validate the capabilities of ALDA. However, the recent literature on zero-shot generalization in vision-based RL, which has formed the majority of our baselines, don’t test on these environments, instead opting for DistractingCS and / or the color shift environments on DMControl only. Finetuning each baseline and producing results on other benchmarks perhaps warrants a broader investigation on the state of zero-shot generalization in RL given current methods, but is out of scope for this paper.

More illustrations or pseudocode could help to clarify the function of associative memory in processing OOD samples. Particularly, pseudocode illustrating ALDA's processing of an OOD sample—from initial encoding to final output—would effectively elucidate the function of associative memory in the generalization process.

Thank you for bringing this to our attention. The processing of an OOD sample is exactly the same as the processing of a training sample. The only difference is that the continuous output of the encoder will be different with an OOD sample. When the associative latent dynamics maps the encoder output to the latent codes, the average “distance traversed” for each latent code for OOD samples will be higher than in-distribution samples. We have added pseudocode for both ALDA’s forward pass and the associative latent dynamics in the appendix (A.7).

An assessment of ALDA's effectiveness in real-world settings would demonstrate its applicability and scalability. For example, testing on a robotic manipulation job requiring the robot to generalize to objects of diverse forms, sizes, or textures not seen during training, could effectively showcase its generalization capabilities in practical applications.

We agree with this completely! However, testing our method in real environments requires other considerations, such as finetuning our method to work with RGB(D) data only and not proprioceptive state (whereas most, if not all sim2real papers use both), since including proprioceptive state as input somewhat defeats the purpose of our method, and including baselines focused more on sim2real rather than simulated RL environments. In addition, we are considering implementing our approach within other RL frameworks such as on-policy RL to test the generalizability of our approach from an algorithmic perspective. We have plans to carry out this work, but we believe this warrants a separate paper.

In order to show that ALDA is generalizable to a wider variety of OOD contexts, could you test it on other benchmarks, such as Procgen?

Please see our response to the reviewer’s comment above.

Could you elaborate on the associative memory dynamics in ALDA, perhaps using a pseudocode or flow diagram, to show how it specifically supports zero-shot generalization?

We have uploaded pseudocode showing ALDA’s forward pass and the memory retrieval mechanism accompanied by a short explanation in the appendix, section A.7.

Is the method applicable to an offline reinforcement learning setting, given that online reinforcement learning can be costly and risky in real-world environments? It would be beneficial to address the adjustments necessary for adapting ALDA for offline reinforcement learning, along with the potential obstacles the authors expect in this adaptation.

We believe ALDA can be adapted to the offline RL setting. Note that the disentangled representation learning literature typically trains models on “offline” pre-curated datasets, which is a very similar setting to the offline RL problem. One potential challenge is estimating the number of sources (i.e. the size of the latent representation) on large offline-RL datasets that are sufficiently diverse. Ideally ALDA would continuously disentangle new data into encodings of new sources as new data appeared, in which case a fixed vectorized representation would not be sufficient. Other self-supervised disentanglement techniques that bias the model to learn generic object-centric representations without labels or masks such as slot attention [1], combined with some form of associative memory, may be more suitable for the offline-RL setting.

[1] Locatello, Francesco, et al. "Object-centric learning with slot attention." Advances in neural information processing systems 33 (2020): 11525-11538.

The method was tested in a narrow set of test environments (color hard and distracting cs). The paper could be improved if the authors would show the benefit of the proposed approach on standard (and more challenging) procedurally generated vision-based datasets for zero-shot generalization in RL, such as Procgen [1] or Crafter [2].

We agree that it would be interesting to see how our method performs on other vision-based datasets curated for zero-shot generalization. However, much of the recent influential work on addressing zero-shot generalization, which we included as baselines, do not report results on these environments, and have focused solely on DistractingCS and/or color shift environments in DMControl. While we would eventually like to expand this work to other domains such as Procgen and improve on the existing algorithm, finetuning our method and all baselines to other domains is an extensive undertaking and out of scope for this paper.

In lines 180-184, the authors mention that one drawback of previous approaches is that the latent representation excludes all information not relevant to the training tasks, which can hinder adaptation to new tasks that involve information that was previously considered irrelevant. The paper would benefit from a concrete experiment that showcases this failure and demonstrates how the proposed approach handles it.

It is difficult to show this effect on the DMControl tasks given the lack of scene complexity. Other reviewers have pointed out that showing the performance of this method on real environments with varying objects, colors, lighting conditions, etc. might help showcase ALDA’s capabilities, which we agree with and have planned as future work. We believe real environments would be the more appropriate choice to showcase this effect, for example training to pick and place a single object with various other objects in the background, but then evaluating the model’s ability to perform pick and place on the objects that were in the background.

Some baselines are missing from the experiment section. For example, a comparison to methods that use the information bottleneck to exclude irrelevant information from the state representation in vision-based RL, such as [3].

The “RePo” baseline reported in the paper uses an information bottleneck in the form of a KL penalty to exclude irrelevant information from the latent space. We will add a short explanation of the RePo baseline to make this clear and refer the reviewer to [1] for more details on the implementation.

An ablation study of $|z_d|$ (the dimension of the latent space) is missing.

Thank you for bringing this to our attention. We will add this ablation to the paper and notify the reviewers when the manuscript has been updated with the additional experiments

Is it possible to utilize an RNN unit instead of frame stacking to incorporate temporal information into the latent code?

It should be possible to replace our 1D CNN layer with an RNN, but we didn’t try this since ALDA and several baselines only use 3 consecutive frames, thus not requiring learning of any long-horizon dependencies.

Section A.2 and Figure 7 are not clear to me. Would you please explain the experimental setting there and the results?

DMControl provides access to low-dimensional proprioceptive state information as the input representation as opposed to RGB images. Since ALDA attempts to learn a strongly disentangled latent representation, we were curious if any of the dimensions of ALDA’s learned representation from RGB data corresponded to any of the proprioceptive state variables. After training the agent, we performed a single rollout and logged the values of several select latent dimensions over time (top row) and compared them to several select rigid-body orientations over time (bottom row) that came from the ground truth proprioceptive state vector. While we can’t make any strong conclusions given that the encoder can map one-to-many, we do notice oscillatory behavior in these latent dimensions, same as the rigid-body orientations over time. This suggests, at least in a qualitative way, that the latent space is not only disentangling but possibly modeling individual rigid bodies of the agent as if we had used proprioceptive state information to begin with. We think this might be a very fruitful future line of research attempting to recover what has typically been hand-engineered proprioceptive state information (which is inherently disentangled) directly from high-dimensional data.

How was the hyperparameter search done for all the baselines?

We used the default hyperparameters provided by the codebases of each respective baseline.

[1] Zhu, Chuning, et al. "Repo: Resilient model-based reinforcement learning by regularizing posterior predictability." Advances in Neural Information Processing Systems 36 (2023): 32445-32467.