Zero Shot Generalization of Vision-Based RL Without Data Augmentation

Thank you for your in-depth feedback and for providing references to work on bi-simulation for visual generalization. We will add bi-simulation to the related works section and fix the reward / RL objective notation issues for the final iteration of the manuscript. We respond to individual comments and questions below.

Bi-simulation may be worth looking into.

Thank you for bringing this to our attention and it is indeed relevant. It aligns with data augmentation (DA) and latent regularization methods like our RePo baseline. From Appendix A.1 Equation 8, DA is defined as attempting to learn an “optimality invariant” Q-function such that $Q^*(**z**, a) = Q^*(**z**', a)$ for two states $**z**$ and $**z'**$ that are semantically the same but perhaps differ visually. Similarly, bi-simulation defines an optimal representation where states that are associated with similar return are closer together in some metric space, and one where task-irrelevant features are removed from the representation. However, our core claim is that for generalist agents, what is “task-irrelevant” can vary depending on the task, so removing such information may be detrimental. For example, a flying bird is irrelevant to autonomous driving until it becomes a collision risk. Rather than removing “task-irrelevant” features, we argue that generalist agents should retain all information about the environment, where the onus is on the policy network to decide what variables are relevant for taking optimal actions.

Since our method uses a reconstruction loss, we retain all information about the environment but specifically in a factorized representation, making our method equivariant to visual distribution shifts. Thus, our method makes no distinction between task-relevant and task-irrelevant features – it simply attempts to factorize the latent structure into separate dimensions without bias. The policy and critic networks decide what is task-relevant and what isn’t. Though this can make our method sensitive to “task-irrelevant” features for any given task, the associative memory component mitigates much of the problem by mapping OOD values (e.g., a new robot color at test time) to known values seen during training before passing the representation to the actor/critic networks.

A figure showing the architecture and...

Thank you for your recommendation. We will update Figure 1 to include a diagram showing the inference procedure of the associative latent model.

A figure illustrating the learned representation...

Here is the first figure requested by the reviewer: https://drive.google.com/file/d/1P0zkA4s4bRJoNWCNhLWmqKLDOUI-GmbW/view?usp=sharing. We take our model ($z_d=12$) trained on the standard, unmodified DMControl environment and perform rollouts on two instances of the color-hard evaluation environment with the same initial conditions. The two plots in the top row show the values that each of the 12 latent dimensions takes on over the course of the trajectory for both environments. The bottom row of images correspond to the initial state of the latent trajectory plots in the top row. We find that latent dimensions that wildly vary are the same in the left and right plot, specifically the ones colored red, pink, green, and brown. As the reviewer predicted, some of the dimensions vary significantly while others are more similar throughout the rollout, likely representing the divide between task-relevant and task-irrelevant features.

Additionally, show if you can use this approach to modify the images...

Regarding the reviewer’s second point, we do precisely this in Figure 6 (bottom) and Figure 8 in the appendix (A.3). In this experiment, we take our model trained on the color-hard environment, sample a batch of images with randomized colors, and encode them into our disentangled latent representation. We then hold all but one latent dimension fixed, interpolate that latent dimension from min to max value (x-axis), and visualize the resulting image using the decoder. The y-axis is several of the latent dimensions we interpolated. From Figure 6, we can see that interpolating the first two latent dimensions (first two rows) corresponds to changing the color of the scene, while the latent dimension corresponding to the bottom row changes the joint angle of the left knee.

What happens if you use to large of a latent space |z_d|? Does this break things?

This doesn't break the algorithm, but it can cause performance degradation if $|z_d|$ strays too far away from the true number of sources of variation. Several of the latent dimensions will end up corresponding to the same physical attributes. We provide a study on the effects of $|z_d|$ in the Appendix Section A.8.

Since each dimension is treated independently...

The number of codebooks 1:1 corresponds to the number of dimensions in the latent space, and so the scaling is linear i.e. if $|z_d|$ is 20, then we have 20 codebooks.

Thank you for your in-depth feedback and for providing the additional reference. We will add it to the related works and discussion sections for the camera-ready version of our manuscript. We respond to individual comments and concerns below.

However, I have some concerns in its novelty...

As the reviewer correctly pointed out, self-supervised disentanglement methods currently do not provide a way to extract or disentangle temporal information. To address this, we take a fixed window of consecutive image observations, encode each observation separately into disentangled latent embeddings, and then extract temporal information via a 1D convolutional network (See figure 2), which is one of our novel contributions.

Our other novel contributions are as follows:

• We show that the latent model in QLAE is equivalent to a Hopfield network with fixed, predetermined memories when the quantization loss is removed. We replace this with the attention mechanism used in modern Hopfield networks, boosting performance over standard QLAE.

• Prior RL disentanglement methods, such as DARLA, use a two-stage approach—disentangling latents with random policy data and then training an optimal policy on the fixed representation. This approach is suboptimal because (a) the random agent may not explore the full state space, and (b) critic gradients cannot backpropagate to the latent space or encoder, which is critical for good performance. Our framework jointly disentangles the latent space and trains the policy while allowing critic gradients to update the latent model, leading to significantly better performance than DARLA.

I am not sure whether these experiment setups are common knowledge in this field, but I would encourage the authors to explain some settings at least in the appendix...

We will add a summary of the DMControl Generalization Benchmark (DMCGB) to the appendix in the final, camera-ready version. For the reviewer’s reference, DMCGB is a wrapper over the standard DMControl (DMC) benchmark, which focuses on optimal control. DMCGB introduces visual distribution shifts to assess zero-shot generalization of agents trained on high-dimensional image observations. The “in-distribution” (or “training” environment as we refer to it in the paper) is the standard DMC benchmark, which emits unmodified image observations. The “color-hard” environment is an OOD setting that perturbs colors randomly on reset, while the DistractingCS environment applies camera jitter and plays random videos from a pre-recorded dataset. Both environments modify only visuals, leaving task dynamics unchanged. Our method trains solely on the unmodified DMC environment and is periodically evaluated on DMCGB environments, testing extrapolative generalization. In contrast, methods like SVEA apply random transformations such as overlaying images from the Places Dataset containing 10 million real world images, or applying random convolutions which change the colors in the scene, making their generalization results more representative of interpolative generalization.

Besides, it would also be helpful to provide some failure cases to show the limitation of the OOD generalization.

The results on the DistractingCS environment highlight the limitations of our method’s OOD generalization. While our performance is on par with SVEA, an ideal model would distinguish between the distracting background video and the agent in the foreground, resulting in minimal performance loss. The observed performance drop suggests room for improvement, making this a promising direction for future research.

Since disentanglement is an emphasis of this paper, it is great to have some analysis...

The representation is disentangled such that each latent dimension corresponds to one unique aspect of the image. As an example, in the Walker2D task, a given dimension could be the color of the robot, one of the robot’s joint angles, the floor, the sky, etc. We qualitatively show the physical meanings of some of the latent dimensions in the “latent traversal plots” in Figure 6 in the main paper, and in the appendix section A.3. In these experiments, we sample a batch of images, set all but one latent dimension static, and then interpolate one latent dimension from min to max value (x-axis) and generate the resulting images using the decoder. The y-axis (rows) are interpolations of different dimensions of the latent space. This allows us to see the physical meaning of each latent dimension, and indeed we find that the latent variables correspond to unique attributes such as a joint angle or the color of the robot/background as with Figure 6 (bottom) where the colors are randomized.

Thank you for your feedback and suggestions on how we could improve the paper's clarity. We respond to individual comments and questions below.

Some theoretical discussions, particularly regarding disentanglement and associative memory mechanisms, could be more clearly explained for a broader audience. -- While ablations exist, a more controlled comparison of ALDA with and without associative memory would further clarify its unique benefits.

Because association is implicit in the dynamics of the latent model that QLAE uses to perform disentanglement, there is no way to remove just the associative part of ALDA and perform an experiment. The closest comparison we can do is the comparison between QLAE and BioAE presented in Figure 3. Like BioAE, QLAE uses activation energy minimization as an auxiliary objective in order to disentangle the latent representation. The only other differences between the two are that BioAE does not have an associative latent model and that BioAE also enforces nonnegative activations, so it essentially functions as ALDA without associative memory. As per reviewer N37F's suggestion, we plan to update Figure 1 to include a diagram of the inference procedure of the associative latent model so that the role of association in our method is more clearly conveyed.

How does ALDA perform if the associative memory component is removed or replaced with a simpler alternative? ALDA without associative memory essentially functions as DARLA (see Figure 5 for comparison) or BioAE (see Figure 3 for comparison), which are purely disentanglement methods, both of which perform worse at visual generalization compared to ALDA. For more details, please see our response above.

Have you tested ALDA on more complex or real-world RL tasks beyond DeepMind Control Suite? How well does it scale? Please see our response to Reviewer Tie6 under "The environments the authors use are toy control environments...".

Thank you for your comments and feedback on the paper. We will fix the grammar errors for the camera-ready version of the manuscript. Regarding your questions and concerns, we respond to each one individually below.

I think it would be nice to discuss this for SVEA vs ALDA since that would be an interesting statistic to see.

The type of data augmentation that’s applied can affect the performance, stability, and sample complexity of the underlying RL algorithm. As noted in the SVEA paper, “More recently, extensive studies on data augmentation have been conducted with RL, and conclude that, while small random crops and translations can improve sample efficiency, most data augmentations decrease sample efficiency and cause divergence”. To explore this, we trained SVEA directly on the color-easy evaluation environment in DMControl. The color-easy environment randomizes the agent, sky, and background colors on reset, but not to the extreme RGB values that the color-hard environment does. This form of augmentation, often used with lighting randomizations for Sim2Real RL deployment [1], also helps assess generalization to OOD visual shifts, since in this experiment the DistractingCS and to an extent color-hard environments will be OOD with respect to the training data.

We compared ALDA, standard SVEA, and SVEA (color-easy) on the "Walker Walk" task here: https://drive.google.com/file/d/1hwFX6glI8-IW6i4vGqDSqo-MuCzDo5F2/view?usp=sharing. SVEA (color-easy) underperformed compared to both vanilla SVEA and ALDA, particularly in the training and color-hard environments. We suspect that the diversity of the 10 million real-world images from the Places Dataset is crucial for SVEA’s generalization and training stability, ensuring that evaluation environments remain in-distribution, or at least well within the support of the training distribution.

[1] "Maniskill3: Gpu parallelized robotics simulation and rendering for generalizable embodied ai."

Why did the authors not try SVEA + ALDA?...

ALDA aims to factorize the image distribution into latent variables, and random overlays could disrupt this by obfuscating the underlying structure. What we wish to convey is, rather than solely relying on data augmentations or brute forcing the generalization problem by collecting massive datasets, models that learn the underlying structure from fewer examples can allocate remaining compute/data budgets to other tasks. For instance, if a robot agent can learn SO(2) invariance from a single object rotated in various ways and can tease out the notion of rotational invariance, then it should be able to generalize rotational invariance to other objects, removing the need for exhaustive data collection on all possible orientations of all possible objects. Unlike computer vision or language data, robot data is more difficult to collect and not as widely available. As of now, the field is allocating a significant amount of data collection efforts to viewpoint, color, lighting, background, [...] randomizations, but if we can alleviate this on the model/architecture side, then those efforts can be spent elsewhere.

In instances where data is cheap or large datasets are readily available, we completely agree with the reviewer that the data should be leveraged. However, solely relying on data may not be sufficient to achieve true generalization, and may instead be obfuscating deeper issues within current robot learning methods. In the SO(2) invariance example, random rotations in CV or viewpoint variations in robotics may not truly be capturing the SO(n) group given current architectures [1, 2], indicating that data alone is insufficient for solving generalization.

[1] "Progress and limitations of deep networks to recognize objects in unusual poses."

[2] "On the Ability of Deep Networks to Learn Symmetries from Data: A Neural Kernel Theory."

The environments the authors use are toy control environments...

We agree that exploring this method in more complex environments would be beneficial and are working on extending our approach to benchmarks like Sim2Real transfer of manipulation policies trained via behavior cloning as part of a separate investigation. However, the focus of this work is to explore whether combining association with latent disentanglement enables zero-shot generalization—an idea that has not been studied before. To that end, we chose RL as the driving optimizer and the DMControl benchmark so that we can study visual generalization in isolation without worrying about the complexities of Sim2Real, harder tasks, real hardware, etc. We also maintain that visual distribution shifts remain a challenging problem to solve, regardless of the difficulty of the underlying task. Finally, recent works addressing visual generalization in RL, many of which we included as baselines, primarily evaluate on DMControl. Extending our method and all baselines to other benchmarks would be beyond the scope of this study.