Zero Shot Generalization of Vision-Based RL Without Data Augmentation
Authors: Sumeet Batra
Gaurav S. Sukhatme
Venue: ICML 2025 poster (2025)
Keywords: representation learning
reinforcement learning
disentangled representation learning
associative memory
zero-shot generalization
Primary area: reinforcement_learning->deep_rl
TL;DR: We demonstrate how disentangled representation learning and associative memory can be used to enable zero-shot generalization of vision-based RL agents without data augmentation or fine-tuning.

Abstract: Generalizing vision-based reinforcement learning (RL) agents to novel environments remains a difficult and open challenge. Current trends are to collect large-scale datasets or use data augmentation techniques to prevent overfitting and improve downstream generalization. However, the computational and data collection costs increase exponentially with the number of task variations and can destabilize the already difficult task of training RL agents. In this work, we take inspiration from recent advances in computational neuroscience and propose a model, Associative Latent DisentAnglement (ALDA), that builds on standard off-policy RL towards zero-shot generalization. Specifically, we revisit the role of latent disentanglement in RL and show how combining it with a model of associative memory achieves zero-shot generalization on difficult task variations *without* relying on data augmentation. Finally, we formally show that data augmentation techniques are a form of weak disentanglement and discuss the implications of this insight.

Zero shot Generalization of Vision-Based RL Without Data Augmentation
Sumeet Batra
Gaurav Sukhatme
Abstract
Generalizing vision-based reinforcement learn- ing (RI.) agents to novel environments remains difficult and open challenge. Current trends are to collect large-scale datasets or use data aug- mentation techniques to prevent overfitting and improve downstream generalization. However, the computational and data collection costs in- crease exponentially with the number of task vari- ations and can destabilize the already difficult task of training RL agents. In this work, we take in- spiration from recent advances in computational neuroscience and propose a model, Associative Latent DisentAnglement (ALDA), that builds on standard off-policy RL towards zero-shot general- ization. Specifically,
Associata
Figure I. Disentanglement + association. A disentangled repre- sentation is learned using the original training data (top). When encountering an OOD sample (bottom), individual latents can he compared and mapped back to known values (colored green). La- tent dimensions that are more OOD (colored red) can be mapped back without affecting other latent dimensions.
generalize. However, a significant drawback of these ap- proaches is, intuitively, that they require larger models, more training data, longer training times, and have greater train- ing instability that must be dealt with care when training RL agents.
revisit the role of latent disentanglement in RL and show how combining it with a model of associative memory achieves zero-shot generalization on difficult task varia- tions without relying on data augmentation. Fi- nally, we formally show that data augmentation techniques are form of weak disentanglement and discuss the implications of this insight.
Yet when we examine biological agents, we find that humans and, indeed, many other primates are able to quickly adapt to task variations and environment perturbations (DiCarlo et al., 2012: Friston, 2010). While all aspects of biologi- cal intelligence that contribute generalization have yet to be understood, there is understanding in the recent neuroscience literature of aspects related to representation learning that we look to for inspiration. Many parts of the brain in human and non-human primates contain neurons that represent single factors of variation within the environ- ment, such as grid cells (Hafting et al., 2005), object-vector cells (Høydal al., 2019), and border cells (Solstad et al., 2008) that represent euclidean spaces, distance objects, and distance to borders, respectively, Such disentangled representations have been theorized to facilitate composi- tional generalization (Higgins et al., 2018) and have been studied with curated datasets where the factors of varia- tion are known (Higgins al., 2017a; Whittington et al., 2023), and even within the context of RL (Higgins al., 2017b). It is then to our surprise, with limited exceptions (Dunion et al., 2023; Sax et al., 2018), that disentangled representation learning has not garnered much attention within robot learning or RL more generally. One potential reason for this is that learning disentangled representations while simultancously learning an RL policy is extremely difficult. Indeed, Higgins al. (2017b) required a two-stage
1. Introduction
Training generalist agents that can adapt novel environ- ments and unseen task variations is longstanding goal for vision-based RL. RL generalization benchmarks have focused on data augmentation to increase the amount of training data available to the agent while preventing model overfitting and increasing robustness environment per- turbations (Yarats al., 2021a; Almuzairee al., 2024; Hansen al., 2021). This follows the current trend in the broader robot learning community of training large models at scale on massive datasets (Kim et al., 2024; Hansen et al., 2024; Team et al., 2024) with the hope that the model will
'Equal contribution 'Department of Computer Science, Univer- sity of Southern California, Los Angeles, USA. Correspondence to: Sumeet Batra <ssbatra udc.edu>, Gaurav Sukhatme <gau- rav @dsc.edu>.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Zero shot Generalization of Vision-Based RL
approach where the disentangled representation was learned
2. Background
first, followed by policy learning. In addition, Yarats al. (2021b) found that using 8-VAE directly led to training instability and worse performance, instead opting to use deterministic autoencoder with softer constraints. Finally, there is counter-evidence by Schott et al. (2022) to sug- gest that, while disentangling representations may facilitate generalization, it alone cannot achieve out-of-distribution (OOD) generalization.
2.1. Reinforcement Learning
We wish to Iearn a policy π that maps states to opti- mal actions that maximize cumulative reward. The agent- enviromment interaction loop is typically formulated a Markov Decision Process (MDP) (S, A, R, R,P,)), where S and A are the state and action spaces, R(s, a) is the reward function, P(s'ls,a) a) is the probabilistic transition function, and 7 is the discount factor. The policy π learns a mapping of state to action with the objective of maxi- mizing cumulative discounted return G, E -E 7'r,] In vision-based RI., we do not assume access to the low dimensional state St E S. Instead, we must in- fer St given high-dimensional image observations o E O, making the problem a partially observable MDP, or POMDP A, P. P. 0),, where O is the space of high- dimensional observations.
We hypothesize that one of the potential key missing ingre- dients to OOD generalization is associative memory mech- anisms that use prior experiences to help inform decision- making in light of new data in a disentangled latent space. Intuitively, if the representation of a stored memory and new observation are disentangled, then the projection of the new observation onto the stored memories becomes factorized projection in which individual factors can be compared in- dependently of other factors (Figure 1). Indeed, many of the aforementioned single-factor neurons exist in the entorhinal cortex in the hippocampus (Hafting et al., 2005), respon- sible for, amongst other things, memory recollection and association. Recent literature suggests that the hippocampus learns flexible representations of memories by decompos- ing sensory information into reusable components and has been implicated in other cognitive tasks such as planning, decision-making, and imagining novel scenarios (Behrens et al., 2018; Rubin et al., 2014). It would then seem that the disentanglement of high-dimensional data into a modular, reusable representation is simply the first step in a multi-step process that enables generalization in biological agents.
Soft Actor-Critic (Haarnoja et al., 2018) is an off- off-policy actor-critic algorithm that jointly trains policy ι and state- action value function Q using the maximum entropy frame- work. The policy optimizes the maximum entropy objective argmax, x Σ; (s..,.)))ppp (r (=(ss))) The optimal Q function (s a) is estimated using temporal difference learning (Sutton, 1988) by minimizing the soft Bellman residual:
((Q) = -Ess,an [(((s),;) - -
Here, y is the soft Q-target, which is computed as = rt + y(mino, Qe; (8++1, a++1) - alog=(-18e+1)) We can then describe the policy's objective as:
Inspired by these two ingredients found in nature, we pro- pose a new method, ALDA (Associative Latent DisentAn- glement), that 1. learns a disentangled representation from the training data and 2. uses an associative memory model to recover data points in the original training distribution zero-shot given OOD data. We demonstrate how this ap- proach enables zero-shot generalization on common gen- eralization benchmark for vision-based RL without using data augmentation techniques or techniques that remove dis- tractor variables from the latent space. Finally, we provide a formal proof showing that data augmentation methods for vision-based RL. create what we refer to as weak disen- tanglement, where the latent space is perhaps partitioned into two or more categories but not perfectly factorized into individual subcomponents. We conclude by discussing the implications of this insight and future directions of this line of research.
J(T) = max a B. [min;=1,2 (Sc: (8).a)) - ologa (a. [s)]
A replay buffer DD is maintained that contains transition tuples (st, ae 81+1) collected from prior interactions of a potentially different behavior policy. Since the off-polios RL formulation does not require transitions to be from the current behavior policy, we can reuse prior experience to update the policy and the Q-function. We use SAC as a foundation for our method, and while we propose some architectural changes to improve the synergy between SAC and our method, the changes are generally applicable to most off- off-policy RI. algorithms.
2.2. Disentangled Representation Learning.
Nonlinear ICA:
The disentanglement problem is some- times formulated in the literature (Hsu et al., 2023) through nonlinear independent component analysis (ICA) due to their conceptual similarity. We follow suit since the no- tation will be useful in later sections. Suppose there are n, nonlinear independent variables S1, that are the
2

---

Zero shot Generalization of Vision-Based RL
sources of variation of the images in the data distribution. A data-generating model maps sources images:
(2017b), the best of our knowledge, is the only prior method that learns disentangled representation of the im- age inputs using a highly regularized 3-VAE (Higgins et al., 2017a) for zero-shot generalization in RL. DARLA's ap- proach is two-stage, where an initial dataset is collected by sampling random actions to first learn a disentangled latent representation, and then a policy is trained on this repre- sentation to maximize future return. However, a significant shortcoming is that random actions may not cover the full state distribution of the agent for more complicated tasks, whereas our method jointly learns the disentangled repre- sentation and the policy. SAC+AE (Yarats et al., 2021b) trains a decoder to reconstruct the images, resulting in a rich latent space that improves performance and sample effi- ciency on many vision-based tasks. Interestingly, SAC+AE mentions B-VAE's used by DARLA and proposes using deterministic variant with similar constraints that shows some zero-shot generalization capability, but the authors make no mention of disentanglement and instead conclude that the key ingredient was adding a reconstruction loss as an auxiliary objective.
The p(s) II p(s)). 0= g(s)
(L)
where g : S O is the non-linear data generating function. The nonlinear ICA problem is to recover the underlying sources given samples from this model. Similarly, the goal of disentanglement is to learn a latent representation z such that every variable 21, ...., 2n, E z corresponds to a dis- tinct source $1...., $n Unfortunately, nonlinear ICA is nonidentifiable that is, there are many decompositions of the data into sets of independent latents that fit the dataset, and so recovering the true sources reliably is impossible. Thus, the field of disentangled representation learning has focused more on empirical results and evaluation metrics on toy datasets where the true sources of variation are known. Given a dataset of paired source-data samples (s, o = g(s)), the goal is to learn an encoder f O + Z and de- coder 9 Z + O such that the disentanglement evaluation metrics are high while also maintaining acceptable recon- structions of the data. Disentanglement models are typically constructed as (variational) autoencoders (Whittington et al., 2023; Hsu et al., 2023; Higgins et al., 2017a) and are rarely applied outside of toy datasets.
Another promising approach to generalization is learning a task-centric or object-centric representation using auxiliary objectives. Yamada et al. (2022) learn task-centric rep- resentation by using expected discounted returns as labels, with the auxiliary task being to minimize the error between the predicted and true return values using the learned rep- resentation. Ferraro et al. (2023); Pore et al. (2024) use segmentation masks to learn object-centric representations that are robust background distractors. Zhang et al. (2021) utilizes bisimulation metrics to determine behavioral sim- ilarity between states, resulting in robust, invariant repre- sentations. One drawback of these approaches is that the latent representation overfits to the task by excluding all other information not relevant to the immediate task, usually citing that irrelevant information in the latent space hinders generalization performance. However, adapting to new task that involves information that was previously consid- ered irrelevant becomes a challenge for these methods. We hypothesize that the issue is not having irrelevant informa- tion in the latent space but rather that the latent variables are entangled without strong priors for disentanglement. A disentangled representation then paves the way for associa- tion, whereby individual dimensions of latent vectors from OOD images can be independently zero-shot mapped back to known values of those latent variables learned from the training data.
2.3. Generalization in Vision-Based RL
Image augmentation methods have shown success and have become the go-to method for generalizing iision-based RL algorithms such as Soft Actor-Critic (SAC) (Haarnoja et al., 2018) and TD3 (Fujimoto et al., 2018), generally using aug- mentations such as random crops, random distortions, and random image overlays to simulate distracting backgrounds. Methods such as DrQ (Yarats et al., 2021a), SADA (AI- muzairee et al., 2024), and SVEA (Hansen et al., 2021) regularize the Q function by providing the original and aug- mented images as inputs into the critic. In many cases, however, the image augmentations can put the training data within the support of the distributions of the evaluation environments. For example, the random conv image aug- mentation changes the color of the agent and/or background, and the policy is evaluated on an environment where the color of the agent is randomized. This brings into question whether these methods are truly capable of zero-shot extrap- olative generalization when the training data is made to be sufficiently similar to the test data.
Beyond image augmentation techniques are methods that perform self-supervisio using auxiliary objectives. Note that image augmentations for RL are also sometimes re- ferred to as self-supervised objectives, however we wish to make the distinction between methods that leverage data augmentation and those that do not, DARLA Higgins al.
2.4. Associative Memory
An associative memory (AM) network stores a set of pat- terns with the intent to retrieve the most similar stored pat- tern given an input. The best-known form is a Hopfield network, originally proposed in Hopfield (1982), which
3

---

Zero shot Generalization of Vision-Based RL
was inspired by how the brain is capable of recalling entire memories given partial or corrupted input (e.g., recalling food item given a particular smell). Classical Hopfield networks could only store and recall binarized memories, whereas modern (dense) Hopfield networks (Krotov & Hop- field, 2016) can work with continuous representations and are trainable as differentiable layers within existing Deep Learning frameworks (Ramsauer et al., 2021). The memory retrieval dynamics are typically formulated energy minimization. Let & E IRd be the input query pattern, and X = [_], E RdxM be memory patterns. In AM models, memories are stored on the local minima of the energy landscape, where the goal is to retrieve the closest stored pattern to & by minimizing energy. Modern Hopfield networks assume the following general form for the energy function:
Intuitively, if data augmentation enables learning latent representation such that the Qzz, a), function of is im- mune to distractor variables, then any dimension of the latent space that encodes for task-relevant variables cannot also encode for task-irrelevant variables. Otherwise, distri- bution shifts involving the task-irrelevant variables would affect the Q function and, thus, the performance of the agent. One of two conditions must be true: either is partitioned, where some variables approximate only sources from D, and others only sources from E, or z contains no informa- tion about sources in E. altogether, both of which are a form of weak disentanglement.
function of
We take a probabilistic perspective to see why this relation- ship is important. Suppose 81....k E D and s1_,..... E F. In order to learn a latent representation that contains no information about task-irrelevant sources $1+ data augmentation methods essentially estimate the marginal distribution over task-relevant sources:
M
E =
F(a E)
In particular, by setting F -Ise( (Ise = log-sum-exponent), the retrieval dynamics becomes Enew = Xsoftmax(8X7 E), which is the attention mecha- nism (Vaswani, 2017) and the backbone of modern Hopfield networks. Follow-up works such as Bietti et al. (2023) show a tight connection between the learning dynamics of Trans- formers and models of associative memory.
p(s1. .I..,... 9k) = -E Ep(s
Sky Sk, Sk+1, Sn.)
S611 5k12
The implication of this equation is that we must collect data for every possible variation of the task-irrelevant sources, which may be prohibitively expensive for real-world ap- plications. Instead, model with strong priors that leads to disentanglement without data augmentation essentially achieves the same result without the additional costs. Al- though the latent representation of such a model may still contain task-irrelevan features, the policy can learn to sim- ply ignore them or associate them with known values in the presence of OOD data, as is the case with ALDA. In addi- tion, these task insk-inrelevant features may become relevant if the task changes (e.g., if the current task is for a manipulator to stack a blue cube, and the next task is to stack a red cube), and so it may, in fact, be important to keep them.
3. On the Relationship Between Disentanglement and Data Augmentation
We begin by motivating the case for learning a disentangled representation for RL agents by showing connection be- tween data augmentation and disentangled representation learning. Specifically, we formally prove that data augmen- tation is a weak disentanglement of the latent space. We define weak disentanglement as partial factorization of some but perhaps not all latent dimensions of the latent space i.e. 3z E z, , 8y. 86 E s)cor(5), (x(=;) * O. Strong disentanglement, on the other hand, is a complete factor- ization where each latent dimension 2 E z encodes for unique source 8; E s and is thus linearly independent of other latent dimensions, which is the goal of disentangled representation learning. The full proof is provided in A.
4. Method
Experimental Setup. We first describe the generaliza- tion benchmark and our evaluation criteria to provide addi- tional context. We train on four challenging tasks from the DeepMind Control Suite (Tassa et al.. 2018). To evaluate zero-shot generalization capability, we periodically evalu- ate model performance under challenging distribution shifts from the DMControl Generalization Benchmark (Hansen & Wang, 2021) and the Distracting Control Suite (Stone et al., 2021) throughout training. Specifically, we have two evaluation environments: color hard, which randomizes the color of the agent and background to extreme RGB values, and distracting cs, which applies camera shaking and plays a random video in the background from the DAVIS 2017 dataset (Pont- Tooot-Tusttae et al., 2017).
Theorem 1: Suppose we are given a z = fg((s)),, where some latent dimension 2ę E z approximates one or more sources. We will denote the approximations as &;. We can categorize the sources into two categories, 1) and E, which correspond to task-relevant and task-irrelevan sources, re- spectively. For any such 2k. if Q" (z, a) is an optimality invariant optimal Q-function immune to distractor variables, then the following must be true of z:
cov(Si, 65lxx) = 0 Vs; E D, E 2k EZ.
4

---

Zero shot Generalization of Vision-Based RL
4.1. Disentanglement
We now describe our framework for jointly learning a dis- entangled representation and performing association. For latent disentanglement, we choose to use QLAE (Hsu et al., 2023), the current SOTA disentanglement method, which trains an encoder fo that maps to a continuous disentangled latent space, a discrete, parameterized latent model Lg. and a decoder Jo that reconstructs the observation. Similar to VQ-VAE (van den Oord al., 2017). QLAE learns a dis- crete latent representation, except that each dimension of the continuous representation is mapped to discrete scalar via a separate codebook, rather than having one codebook for the entire latent space. Concretely, Z is the set of latent codes defined by the Cartesian product of n, scalar code- books Z V1 X ... x Vn_ i.c., there are n2 latent variables and 1Vyl discrete categories per variable. The continuous outputs of the encoder are the latent variables, each of which is quantized to the nearest scalar value in their respective codebooks i.e. the number of codebooks equals the dimen- sionality of the latent representation.
A+1
(((
20
Leaes
AALON
Sapastoi4)
Skn.
Figure 2. Diagram of our method SAC+ALDA. Trainable com- ponents are colored yellow. A strongly regularized autoencoder and the quantized latent space enable latent disentanglement. The latent model is also responsible for association when encountering OOD inputs.
4.2. Association
The naive approach to performing association would be to feed the quantized latent representation through Hopfield network. However, upon closer inspection of QLAE's latent dynamics, we find that most of the components of generic associative memory model are already present, i.c., QLAE is implicitly also a Hopfield network. Figure 3 shows a com-
2d,- = rigminn,xey, [o()) - vaal.j = 1,..,,....
(3)
Since we cannot differentiate through argmin, as with VQ- VAE, the authors of QLAE use quantization and commit- ment losses and straight- gradient estimator (Ben-
gio al., 2013):
Walker Walk Color Hard
Walker Walk Distracting
1000
1000
Lquantize ||topGradiennf((())) - 2al13
(4)
Loommit []fo(x) StopGradient(za)) (za)1
Dwesard
600
600
The authors claim that while this is a failure mode for vector quantization, 2 is low-dimensional enough that, in practice, it does not meaningfully impact performance. While this may be true for standalone disentanglement benchmarks, we find that it causes training instability and performance degra- dation when jointly learning policy for high-dimensional continuous control problems. We propose a solution in section 4.2 from the viewpoint of associative memory.
400
400
200
200
0.4
1.0
0.4
2.0
steps
environment staps
1e6
BloAE
QLAE
Figure 3. Ablation comparing SAC+QLAE to SAC+BioAE on the two distribution shift evaluation environments for the walker walk task.
In vision-based RL, framestacking is commonly used to encode temporal information by stacking l RGB frames, resulting in encoder inputs and decoder outputs of shape R BB xW However, latent disentanglement models
parison of using SAC with QLAE vs with BioAE (Whitting- ton et al., 2023), another disentanglement method that uses biologically inspired constraints and achieves comparable results on disentanglement benchmarks (Hsu et al.. 2023). This comparison allows us to isolate and analyze how the associative latent dynamics contributes to generalization, since BioAE does not incorporate any associative mecha- nism while QLAE implicitly does. BioAE achieves strong initial performance the two evaluation environments, but slowly degrades over the course of training. We suspect both models overfit to the training environment, but QLAE is capable of zero-shot mapping OOD latent variables to in-distribution values. To understand how QLAE implicitly incorporates an associative mechanism, we first present the
perform poorly when given stacked images, struggling to separate independent sources. We observe this in Appendix Section A.5. To resolve this issue, we fold l: into the batch dimension and encode/decode batches of single images in R BB X resulting in a batch size Bk of disentangled latent vectors z E F Bkn., To incorporate temporal infor- mation, we reshape the batch of latent vectors into R Bx len,, and feed it into ID convolutional neural network (CNN), producing our final latent vector 2 E R Bxe, 2 is used as the state representation for the actor and critic networks, while the decoder network for the disentanglement model only ever receives the disentangled representation 24 as input.
5

---

Zero shot Generalization of Vision-Basce RL
general framework described in Millidge al. (2022) of universal Hopfield network that all feedforward associative memory networks in the literature can be factorized as:
1940
Waller Walk Toan
Pl kct Wolk Coror Hord
1600
Walker Walk Datrocting
100
600
00
600
600
300
= P sep((smm (X,
(5)
On ST 15
1i ano mest ctps (led)
0an
atalorries cape Eiato
ate roncent capa (290)
ALDA
QLAE
P is projection, sep is separation function, and sim is similarity function between the stored memories X and query Ç. In this case, equation 3 can be interpreted as follows: we compute similarity scores via L1 distance between fa( and StopGradient(za), and the separation function is the argmin operator. Since argmin retreives the desired value from the codebook, the projection function P in this case is the identity function. Through this view, we can rewrite the latent dynamics of QLAE in many ways, perhaps exchanging the similarity function for L2 distance or dot product, changing the separation function, etc., as long as it follows the framework of 5. Since Z is a product of scalar codebooks, L1 distance remains an appropriate choice for the similarity function. Instead, we augment the latent dynamics with a Softmax separation function follows:
Figure 4. Ablation comparing ALDA and QLAE. Average of 5 seeds, shaded region represents a 95% confidence interval.
Coommil + Creconnstruct
J(ALDA) = EornD [15o(o) - StopGradient(
[Softmax (- 3L,, (fo(o), V)) o 1)
(7)
+ +log go go (0rl2²) + Ao 1113 A
where the last two terms are weight-decay on the encoder and decoder parameters controlled by Xo and Xg. respec- tively. Observations collected by the policy are stored in a replay buffer D, from which batches are randomly sampled to train ALDA. We observe that this formulation consider- ably improves training and evaluation performance on the color hard environment and to a degree, on the Distract- ingCS environment, as shown in Figure 4. The obvious question is, how do set the dimensionality of za. which should 1:1 correspond to the number of sources n, if the number of sources is unknown? While there is no rigorous method to derive |2a| at this time, we empirically found that setting |aa] to within the ballpark of the size of the obser- vation spaces for the proprioceptive, state-based versions of the tasks seemed to work well. |zal is set to 12 for all reported tasks.
=d=- Softmax((3L vy)) O V
(6)
where 3 is scalar temperature parameter. Equation 6 can be interpreted in two ways. From an associative memory perspective, attention-basee Hopfield models apply Softmax to separate the local minima (stored memories) on the en- ergy landscape, where 3 controls the degree of separation, and so we've ve recovered the modern Hopfield memory re- trieval dynamics. From a purely mathematical perspective, we have what resembles the reparameterization Jang et al. (2017), although we do not perform any sampling in method. This lends a novel view on attention-base based Hopfield networks- models with high-temperaturg parameter can be interpreted as classifiers over X classes, where XX| is the number of stored memo- rics whose local minima are well separated on the energy landscape.
-oftmax- categorical
5. Experiments
We compare against several baselines that together represent the full range of different learning paradigms in the litera- ture that attempt to elicit zero-shot generalization. DARLA (Higgins et al., 2017b) is, the best of our knowledge, the only other algorithm that attempts to learn disentangled representation of the image distribution towards zero-shot generalization of aision-based RL. SAC+AE (Yarats et al., 2021b) uses a deterministic autoencoder with an auxiliary re- construction objective and strong regularization that demon- strates decent zero-shot generalization capability. RePo (Zhu et al., 2023) a model-based RL algorithm that learns task-centric latent representation immune to background distractors. Finally, SVEA (Hansen et al., 2021) is an off- policy RL algorithm that improves training stability and performance of off-policy RL under data augmentation. As in their paper, we use the random overlay augmentation for SVEA, where images sampled from the Places (Zhou et al.,
In the limit, as 3 goes to infinity, we achieve maximum separation between memories and recover equation 3. In practice, we choose a large value for 3 such that we retrieve one scalar from each codebook, as originally intended, al- though our method works well with smaller values of /3 (see appendix Section A.4 for additional results). Since large 3 values can cause downstream gradients to vanish, we find that keeping the commitment loss from equation 4 helps keep the outputs of the encoder close to the values of the latent model. However, we do not optimize the codebook towards the encoder outputs i.c., we omit Lquantiz. This can be interpreted as having a set of task-optimized memories that the encoder must learn to map to under the Hopfield interpretation. Our final objective for ALDA is then simply
6

---

Zero shot Generalization of Aision-Bascd RL
Cartpole
Walker
Cup
Finger Spin
3020
3000
1000
1000
Epiodd Hewwwer
100
140
(00)
040
430
200
0.3
02
Ep Rewardd Cotor 1ene
800
800
000
100
160
600
220
260
200
200
Co. Revaard Distreet
000
300
000
000
430
400
230
200
0.
03
0.3
02
03
04
0.
02 03
D4
codtror ment steps Ce5)
crvironment Steps (166)
oww stops 11c6)
eтntonmen: S1n05 (1e6)
ALDA
Figure 5. Top row: performance on the training environment. Middle and bottom rows are evaluation results on the "color hard" and DistractingCS evaluation environments, respectively, Average seeds, shaded region represents a 95% confidence interval.
Figure 6. Latent traversals. Each row corresponds to holding all but one latent variable fixed, interpolating it from min max value in its codebook, and decoding resulting vector an image. Top: Traversal of select latents for the standard training environment. Bottom: Traversal of select latents when training directly on the color hard environment. Latents that encode for distractor variables (e.g.- color) seemingly not simultancously encode task-relevant variables (e.g., agent dynamics). Visualizations latent traversals can be found Appendix section A.3.
2017) dataset of 10 million images overlayed during training. The training curves and evaluation "color hard" and DistractingCS are presented in Figure 5.
within the support of the data distributions of the evaluation environments. Nevertheless, ALDA performs comparably and, in some cases, is equal to SVEA despite only seeing images from the original task. Furthermore, we find that ALDA can even outperform SVEA SVEA is trained with different data augmentations not including image overlays from the Places dataset, shown in Appendix Section A. 1. Performance degrades severely for all algorithms on the DistractingCS environment. We suspect that, in addition to the already difficult task of ignoring the background video, camera shaking affects the implicitly learned dynamics, and thus, additional finetuning may be unavoidable for this task.
Excluding SVEA, ALDA outperforms all baselines on both distribution shift environments. ALDA also maintains sta- bility and high performance on the training environment, de- spite the disentanglement auxiliary objective and extremely strong weight decay (No, A 0.1) the encoder and decoder. We do not expect to outperform SVEA since it uses additional data sampled from dataset of 1.8 million diverse real-world scenes. likely putting the training data

---

Zero shot Generalization of Vision-Based RL
Still, ALDA performs better than all baselines excluding SVEA on Distracting CS well, even matching the perfor- mance of SVEA on cartpole balance and finger spin.
there remains portion of the distribution containing ev- ery possible edge case and unaccounted-for variation, com- monly referred to as "the long tail," that remains elusive because it is prohibitively expensive to account for every possible variation. Unfortunately, these uncaptured varia- tions are frequent enough due to the ever-changing dynami- cal nature and complexity of the real world that deploying agents in the real world remains challenging. Therefore, seems the case that data augmentation techniques, collecting massive datasets, and the like are not sufficient to develop generalist agents capable of adaptation the way humans and other animals are. That's 's not to say that data is not important or a fundamental ingredient to training machine learning models. In fact, the method proposed in this paper scales with more data as with prior works that leverage data aug- mentation techniques. Instead, our proposition is that if a data-driver model can generalize better with less data, then it will scale better with more data.
The disentangled representation learning field primarily uses toy datasets where the ground truth sources of the data dis- tribution are known. Therefore, all disentanglement metrics we are aware of require knowing the sources, making it difficult to quantitatively evaluate disentanglement perfor- mance DMControl. In the absence of any quantitative disentanglement metrics, we opt to show empirical evidence of disentanglement in our model, presented in Figure 6. In this experiment, we encode a random batch of observations o into the disentangled latent representation 2d = Ls (o)) We pick one latent dimension zd; of the latent vector and in- terpolate it from its minimum to maximum value in its code- book while holding all other latent variables fixed. The re- sulting perturbed latents are decoded using ALDA's learned decoder o' = Io (=a) and presented as row of images. We find that each latent tends to learn information about a single aspect of the robot, for example, the orientation of the torso (middle row of the top image) or left hip joint angle (top row of the top image). From the bottom image in Figure 6. we can see that ALDA encodes scene color information separately from robot dynamics when trained directly on the color hard environment. Unlike methods that learn invari- ant representations, ALDA does not discard aas--irelevvant information, but rather encodes all variables separately, re- sulting in an implicit separation between task-relevant and iieeeertll variables. Further empirical analysis on dis- entanglement of the image observations and what the latent dimensions learn to represent are presented in A.2.
In Section 3, we showed how data augmentation and dis- entangled representation learning aim to achieve the same result- a factorization of the latent space into separate components in order to improve downstream generalization performance. Given the additional computational and data collection costs and potential training instabilities that data augmentation methods may incur, it seems more fruitful to investigate models with inductive biases that elicit modular and generalizable representations without relying on data scaling laws. While presenting the model with sufficiently large and diverse datasets remains unquestionably important, we cannot rely solely upon the data in hopes that the model learns good representation. As with any other inductive biases, such as using CNNs for vision tasks or transformers for NLP tasks, inductive biases that elicit modular represen- tations while leveraging data are worth studying if we are to develop agents that can perform and adapt well in the real world.
6. Discussion
As stated previously, the disentanglement problem by exten- sion of nonlinear ICA underdetermined, so there are many ways the latent space may factorize, perhaps by represent- ing the sky and background with one latent or by separating them into two latents, etc. Given that both the task and re- construction gradients of the critic/decoder affect the latent model/encoder, an interesting scientific and philosophical implication is that the model is potentially biased towards a disentangled representation that is useful, although there is no way to quantitatively or qualitatively show such result at this time. Nevertheless, it remains an interesting line of further investigation from a scientific standpoint, and per- haps, philosophically, says something about whether the question "What is the ground truth factorized representation that best explains the data?" is even the right question to ask.
We hope that the work presented here inspires future re- search into novel models and architectures to learn represen- tations that enable the adaptability we see in our biological counterparts. We discuss some limitations of our method and promising directions for future research. One notable limitation that our disentangled latent representation zn does not explicitly account for temporal information since it primarily estimates the sources that produce the image distribution. Instead, we must capture temporal information in the downstream 1D-CNN layer as shown in Figure 2. How to learn a disentangled representation that contains sources of both the image data and temporal information for decision- makinn tasks remains an open question. Another limitation is that, while we introduce a simple Hopfield model as modification to QLAE, we do not take advantage of the more recent literature involving learnable attention- based or energy-based Hopfield networks (Ramsauer et al.,
RL agents deployed in the real world must constantly adapt to changing environmental conditions. Much of the variance
can be captured with a sufficiently large dataset. However,
8

---

Zero shot Generalization of Vision-Based RL
2021; Hoover al., 2024). Stronger Hopfield models that synergize well with disentangled representations is another potentially fruitful research direction.
Impact Statement
Agents are rapidly being integrated into our society across many different dimensions simultaneously, from web-based LLM agents driverless cars, to household robotics. gen- eral trend is that models become larger and trained vast datasets, the interpretability becomes diminished, which is problematic for Al safety research, especially in critical tasks as autonomous driving. It is also generally the case that methods which are more interpretable by de- sign are interpretable the expense of performance when compared to fully end-to-end systems without safety con- straints. We believe that encouraging learning disentangled representations in models not only enables greater in- terpretability and facilitates Al safety research, but does so while improving generalization rather than hindering it, and is thus a promising avenue of further research.
Given that we use very compact disentangled latent space with strong empirical evidence that individual latents cap- ture information about specific aspects of the agent, an inter- esting research direction is to investigate whether all or parts of the proprioceptive state representation can be recovered from image observations. We provide some preliminary evidence of this in the appendix (A.2). Beyond interpretabil- ity, such a model may yield better performance since state- based RL agents tend to perform better than baseed agents. Finally, while our work was inspired by the role of the hippocampus biological intelligence, the exact mechanisms of the machinery and how they interact with decision- decision-making, planning, and imagination components of biological brains are by no means precisely modeled in this paper, nor are all of the computations the hippocampus may be performing fully understood. Future collaborative research between the machine learning and neuroscience fields into data-driven computational models of these mech- anisms may yield even better-performing, adaptable agents.
Acknowledgements
We would like to thank Tommaso Salvatori for advice and insightful discussion on modern Hopfield networks. We also thank David Millard, Shashank Hegde, and Yutai Zhou for providing feedback on the proof.

---

Zero shot Generalization of Vision-Based RL
References
Machine Learning Research, pp. 1582-1591. PMLR, 2018. URL http::/prooceedinggg ml mlr.press// v80/ ht0/fujjmotolll html.
Almuzairce, A., Hansen, N., and Christensen, H. I. A recipe for unbounded data augmentation in visual reinforcement learning. CoRR, abs/2405, 17416, 2024. doi: 10.48550/ ARXIV 17416. URL https:: / http//di.og//////////// 48550/arXiv. 2405. 10.4451..114 17416.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In Dy, J. G. and Krause, A. (eds.), Proceedings of the 35th International Conference on Machine Learn- ing. ICML 2018, Stockholmsmässan, Stockholm, Swe- den, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 1856-1865. PMLR, 2018. URL http: pttpp//roceeeennceeeeepe/o....ae wwl. m.r..prss/ v80 / htt//aaazooaaaa.. html.
Behrens, T. E., Muller, T. H., Whittington, J. C., Mark, S., Baram, A. B.. Stachenfeld, K. L., and
Z. What is cognitive map? organizing knowledge for flexible behavior. Neuron, 100(2):990-50, 2018.
Bengio, Y., Léonard, N., and Courville, A. C. Estimating or propagating gradients through stochastic neurons for conditional computation. CoRR, abs/1308.3432, 2013. URL http::tt://aaa /ttp:v///aavvva// wtt.aaiwwi.....aaaarrrrrreaaeeae wttg/bs////// 1100..//133... 3432.
Hafting, T., Fyhn, M., Molden, S., Moser, M.-B., and Moser, E. I. Microstructure of spatial map in the entorhinal cortex. Nature, 436(7052):801-806 2005.
Bietti, A., Cabannes, V., Bouchacourt, D., Jégou, H., and Bottou, L. Birth of a transformer: A memory viewpoint. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, Decem- ber 10 16, 2023, 2023. URL http: /papers. 2023/ /hash/
Hansen, N. and Wang, X. Generalization in reinforcement
learning by soft data augmentation. In IEEE International Conference on Robotics and Automation, ICRA 2021, Xi'an ana China, May 30 June 5, 2021, pp- 1361 1-13617.
IEEE, 2021. doi: 10. 10.11009/CRA4444002220222 URL https: /tmp:://doii /tw..//iii////////////.... http://1...199///4448800000 4021.956111
103.
nips. cc/ httpe//aappppfeeeeeeeeeeee
eueseneNNedoe H., and Wang, X. Stabilizing deep learning with convnets and vision transformers under data augmentation. In Ranzato, M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, W. (eds.), Advances in Neural Information Processing Systems 34: Annual Conference Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 3680-3693, 2021. URL
html.
DiCarlo, J. J., Zoccolan, D., and Rust, N. C. How does the brain solve visual object recognition? Neuron, 73(3): 415-434, 2012.
Dunion, M., McInroe, T., Luck, K. S., Hanna, J. P., and Albrecht, S. V. Temporal disentanglement of repre- sentations for improved generalisation in reinforcement learning. In The Eleventh International Conference Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL https: // //openreviee ew net / forum? i PP ISLTT.
neur http:pps.cc//aper//222///////// 2o0ffee22aabfteeeeessdd000000000 2attt000000bbbaaacttt..ac Abstract. html.
Hansen, N., Su, H., and Wang, X. TD-MPC2: scalable, ro- bust world models for continuous control. The Twelfih International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenRe- view.net, 2024. URL https: wpppreveeeeeeeeeeeee nett/ forum? di0ouiistthhhht
Ferraro, S., Mazzaglia, P., Verbelen, T., and Dhoedt, B. FOCUS: object-centric world models for robotics manip- ulation. CORR, abs/2307.0244 2023. doi: 10.48550/ TRXIV.2307.02427, URL https: htt.dd..rgg///// 48550/arXiv. 02427.
Higgins, I., Matthey, L., Pal, A., Burgess, C. P., Glorot, X., Botvinick, M. M., Mohamed, S., and Lerchner, A. beta-vae: Learning basic visual concepts with a con- strained variational framework. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017a. URL https: / //opnnevveww.oo net wtw.:reeeeeoooiinieeeeeeeeee
Friston, K. The free-energy principle: a unified brain theory?
Nature reviews neuroscience, 11(2):177-138, 2010.
Fujimoto, S., Hoof, H., and Meger, D. Address- ing function approximation error in actor-critic meth- ods. In Dy, J. G. and Krause, A. (eds.), Proceedings of the 35th International Conference on Machine Learn- ing, ICML 2018, Stockholmsmässsnn Stockholm, Swe- den, July 10- 15, 2018, volume 80 of Proceedings of
Higgins, I., Pal, A., Rusu, A. A., Matthey, L., Burgess, P., Pritzel, A., Botvinick, M. M., Blundell, C., and
10

---

Zero shot Generalization of Vision-Based RL
Lerchner, A. DARLA: improving zero-shot transfer in reinforcement learning. In Precup, D. and Teh, Y. W. (eds.), Proceedings of the 34th International Conference on Machine Learning. ICML 2017, Sydney, NSW, Aus- tralia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pp. 1480-1490. PMLR, 2017b. URL http::tp//// pttpp/eeeceeensssceee mlr. prrss// v70/higgins17a, html.
10.46550AAXXV.2246009222 URL https: htops:///oo /doi. org/10. wtwp://.44855//aax..xxo. 09226. 09246.
Krotov, D. and Hopfield, J. Dense associative memory for pattern recognition. Lee, D. D., Sugiyama, M., von Luxburg. U., Guyon, I., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 29: Annual Conference Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pp. 1172- 1180, 2016.
Higgins, I., Sonnerat, N., Matthey, L., Pal, A., Burgess, C. P. Bosnjak, M., Shanahan, M., Botvinick, M. M., Has- sabis, D., and Lerchner, A. SCAN: learning hierarchical compositional visual concepts. In 6th International Con- ference Learning Representations, ICLR 2018, Van- couver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018. URL ht ttts::: // //apeneeview..o. net
URL httppss//proceeddngeee
eaae (a html.
Millidge, B., Salvatori, T., Song, Y., Lukasiewicz, T., and Bogacz, R. Universal hopfield networks: A general framework for single-shot associative memory models. In Chaudhuri, K., Jegelka, S., Song, L.,, Szepesvári, C., Niu, G., and Sabato, S. (eds.), International Conference on Machine Learning ICML 2022, 17-23 July 2022, Bal- timore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 15561-15583. PMLR, 2022. URL https: http:::://ppoceeeiinggee v162 mi (t55/mi1li0n http:/m111ddee222aaaaa2m html.
-RZ.
Hoover, B., Liang, Y., Pham, B., Panda, R., Strobelt, H., Chau, D. H., Zaki, M., and Krotov, D. Energy transformer. Advances in Neural Information Processing Systems, 36, 2024.
Hopfield, J. J. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 7955)22554-2588 1982.
malr.press.
Cont-Tuset, J., Perazzi, F,, Caelles, S., Arbeláez, P., Sorkine- Hornung, A., and Van Gool, L. The 2017 davis chal- lenge on video object segmentation. arXiv preprint 17X4:::70..0065, 2017.
Høydal, Ø. A., Skyteen, E. R., Andersson, S. O., Moser, M.- B., and Moser, E. I. Object-vector coding in the medial entorhinal cortex. Nature, 568(7752):400 404, 2019.
Pore, A., Muradore, R., and Dall' Alba, D. DEAR: disen- tangled environment and agent representations for re- inforcement learning without reconstruction. CoRR, 2as/2407.00633 2024. doi: 10.48550/ARXIV24400, 00633. URL https: http:://doi.oooo http:///doo..o/11..3//3. arXiv. 2011.240777 2407. 00633.
Hsu, K., Dorrell, W., Whittington, J. C. R., Wu, J., and Finn, C. Disentanglement via latent quantization. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information P Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, Decem- ber 10 16, 2023, 2023. URL http: //apers. nips.
Ramsauer, H., Schäfl, B., Lehner, J., Seidl, P., Widrich, M., Gruber, L., Holzleitner, M., Adler, T., Kreil, D. P., Kopp, M. K. Klambauer, G., Brandstetter, J., and Hochreiter, S. Hōрńicää
fiflookks & is all you need. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL ht https::///// / wpwprevieneeeeeeeeeeeee net forum?
html.
Jang, E., Gu, S., and Poole, B. Categorical reparame- terization with gumbel-softmax. In 5th International Conference Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https: //openreview.co//////// net formmm id-rkEEe
Rubin, R. D., Watson, D., Duff, M. C., and Cohen, N. J. The role of the hippocampus in flexible cognition and social behavior. Frontiers in human neuroscience, 8:742, 2014.
Kim, M. J., Pertsch, K., Karamcheti, S., Xiao, T., Bal- akrishna, A., Nair, S., Rafailov, R., Foster, E. P., Lam, G., Sanketi, P., Vuong, Q. Kollar, T., Burchfiel, B., Tedrake, R. Sadigh, D., Levine, S., Liang. P., and Finn, C. Openvla: An open-source vision-language action model. CoRR, abs/2406.09246,, 2024. doi:
Sax, A., Emi, B., Zamir, A. R., Guibas, L., Savarese. S., and Malik, J. Mid-level visual representations improve gener- alization and sample efficiency for learning visuomotor policies. arXiv preprint /81. 1971, 2018.

---

Zero shot Generalization of Vision-Based RL
Schott, L.,, von Ktigelgen, J., Träuble, F., Gehler, P. V., Rus- sell, C., Bethge, M., Schölkopf, B., Locatello, F., and Brendel, W. Visual representation learning does not gen- eralize strongly within the same domain. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenRe- view.net, 2022. URL https:///////m /ttp://ppeene //pp//ppeeeew..r//ee/e net/ Forum?ide Darumm?dd--UUPPIaaaS 9RUhPlaadhh
2023, Kigali, Rwanda, May 1-5, 2023. Open
2023. URL htt.t:::///pppeeviieeeee hetpiiewwweeeeeeeeeeeeeeeeeee forum?
Yamada, J., Pertsch, K., Gunjal, A., and Lim, J. Task- induced representation learning. In The Tenth Interna- tional Conference Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https: http///pppeeeee.eeeeeeeeeeeeeeep net.:oruntt 1d=02yxtIZA2P%
Solstad, T., Boccara, C. N., Kropff, E., Moser, M.-B., and Moser, E. I. Representation of geometric borders in the 2008, Stone, A., Ramirez, O., Konolige, K., and Jonschkowski, R. The distracting control suite- A challenging bench- mark for reinforcement learning from pixels. CoRR, abs/2101.02722, 2021. URL pettps://arxiv..eo 2obs/2101.0 02722.
Yarats, D. and Kostrikov, I. Soft actor-critic (sac) im- plementation in pytorch. deni hatissyaass/pytocchhccm sac,chh__aa,,000 2020.
entorhinal cortex. Science, 322(5909):
Yarats, D., Kostrikov, I., and Fergus, R. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021a. URL https: / openreview net forum? dd= 6 TvGa...
Sutton, R. S. Learning to predict by the methods of temporal differences. Mach. Learn,, 3:9-44, 1988. doi: 10. 1007/ 15001 URL hattps://doie org/10...10... 1007/ BF BF00115009.
Yarats, D., Zhang, A., Kostrikov, I., Amos, B., Pineau, J., and Fergus, R. Improving sample efficiency in model- free reinforcement learning from images. In Thirty-Fisth AAAI Conference Artificial Intelligence, AAAI 2021, Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Sympo- sium on Educational Advances in Artificial Intelligence, EAAI 2021, Virtual Event, February 2-9, 2021, pp. 10674- 1068 AAAI Press, 2021b. doi: 10. 10.1609/AAAL V35112. 17276. URL https:t/p///. http//di..ogg//////////oo /69/ aaai. v35i12.17277
Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., de Las Casas, D., Budden, D., Abdolmaleki, A., Merel, J., Lefrancq. A., Lillicrap. T. P., and Riedmiller, M. A. Deepmind control suite. CoRR, abs/1801 .00900, 2018. URL http: / http/arxivvv/////. org/ abs /101..1111.111111 00690.
Team, O. M., Ghosh, D., Walke, H., Pertsch, K., Black, K., Mees, O., Dasari, S., Hejna, J., Kreiman, T., Xu, C., Luo, J., Tan, Y. L., Chen, L. Y. Sanketi, P., Vuong, Q., Xiao, T., Sadigh, D., Finn, C., and Levine, S. Octo: An open- source generalist robot policy. CoRR, 10w/2405.12113
Zhang. A., McAllister, R. T., Calandra, R., Gal, Y., and Levine, S. Learning invariant representations for rein- forcement learning without reconstruction. In 9th Inter- national Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenRe- view.net, 2021. URL https: http:////eeeeeeeee..mmmm net/ forum (0d:===FFCDRRREEU
2024. doi: 10.48550/ARXII.2005.2222 12211. URL https: //doi. org http:r/g.........0 405.448550//r/// 2405 12213.
van den Oord, A., Vinyals, O., and Kavukcuoglu, K. Neural discrete representation learning. In Guyon, I., von Luxburg, U., Bengio, S., Wallach, H. M., Fergus, R. Vishwanathan, S. V. N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA. USA. pp. 6306-6315, 2017. URL https: httpps:/proceeedd hpppppeeeinngs.see neurips. wttppsppppppppppaaeaeeeeeeeeeeet ht1////aarr//11//aaa/ f (tteee0000cceeee0eee fbc rast Repo: Resilient model-based reinforcement learning html.
Zhou, B., Lapedriza, A., Khosla, A., Oliva, A., and Tor- ralba, A. Places: A 10 million image database for scene recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2017.
Zhu, C., Simchowitz, M., Gadipudi, S., and Gupta, A.
by regularizing posterior predictability. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Infor- mation Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023, 2023. URL http:/ /papers. nips. cos.cerrrpper.. fi (illes/paper/22/ //aash/
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Whittington, J. C. Dorrell, W., Ganguli, S., and Behrens, T. Disentanglement with biological constraints: A the- ory of functional cell types. In The Eleventh Interna- tional Conference on Learning Representations, ICLR
12

---

Zero shot Generalization of Vision-Based RL
html.
13

---

Zero shot Generalization of Vision-Based RL
A. Appendix
A.I. Proof
Preliminaries. We reintroduce the nonlinear ICA problem formulation here for the reader's reference. There are na nonlinear independent variables $1,..., Sn, that are considered the "ground truth" sources of variation of the data distribution. We assume there exists data-generating model that maps sources images:
((s) II (())), g(s)
where S - O the non-linear data generating function. For the purpose of this proof, we restrict S and O to be the space of sources and images within our dataset. The goal of nonlinear ICA, and by extension, disentangled representation learning, to recover the underlying sources given samples from this model. We claim that most, if not all, data augmentation techniques in Q-learning
form of weak disentanglement, where either the method factorizes the latent space into task-relevant fask-irrelevant variables removes task iaskiitrlevvant variables from the latent space entirely.
For given task, we can split the sources into two categories: sources $1..k, k < that are task-relevant, which we will call D, and sources $k.1....,, that are not, whose category we will refer to as E. The encoder maps observations to a latent space f: O - Z, and so z is function the sources z = fo(g(s)) We refer to & as an approximation to the true variable &t that exists one or more dimensions of z. We make no assumptions on whether the sources are entangled or disentangled z.
Optimality Invariant Image Transformations Described in Yarats et al. (2021a), data augmentation applied Q-learning can be formulated using the following general framework. An optimality state transformation h: O x T - O is mapping that preserves Q values.
-Ua(())))))))))))))))))))))) C,,
(8)
v are the parameters of ()...) drawn from the set all possible parameters T. In other words, defines the space of all possible data augmentations that should not affect the output of the Q-function.
Proposition 1: Let Φ - S be a function that perturbs sources 8j E. Then
h(o, g()(s)).
(9)
This follows from the definition of E in that any optimality -nveaiiat transformation to the observation must have implicitly resulted from perturbation of some iese-ireleevatt source 8 E E.
Proposition 2: For any given z = fo()))) and any perturbation to true source 8j E E,j E [A + 1, n, resulting in a new
latent = (g)))))) the following must be true for an optimality invariant optimal Q-function:
(z, a) = (z'.za))
(10)
We can rewrite z' as z' fo((()())) i.c. o= g({(s)), and by Proposition 1, g(((s)) h[ v). Essentially, an optimality invariant Q-function immune to variations of task-irrelevant sources from the set E, since they correspond optimality- invariant state transformations.
Theorem 1: For any z = fo()(s)), and for any dimension of z, the following must be true for an optimality invariant Q-function:
cov(s;, 8y12z) = 0 Vs, e D,sy e E,i [1, e [k 1,n...
(11)
To see why this must be the case, suppose that the covariance is nonzero and suppose that we perturb E E to s'j, giving us a new observation o' = g(s'). Since 8, task- inseevveeeeant, z' = =oo(o') = (h((,,, v)) for some v E If z a function of
14

---

Zero shot Generalization of Vision-Based RL
then by equations 8 and 10, (z', - (z, a). However, if co (=)= A 0 for some 26 E z, then 6 E + E Since & E D, (2'2,,,))) a) Q" (z, a), which is contradiction. Therefore, the conditional covariance between any D and any 8j E E for any given z must be zero, which implies that the approximations of task-relevant and task-irrelevant sources in z are disentangled.
A.2. Latent Trajectory Visualizations
3,00
Thjeclory
Thanciery
3.00
Theastary
Tigpctor
Trpactay
0.25
0.15
0.75
040
0.00
0:00
3.0
-0.75
-875
-375
200
300
100
156
200
225
Criomaatio pigooggs po
Oriensation Trajectory
3.00
Tajectory
Orientacion Talectury
1.89
Tajecttry
0.75
0.73
0.35
2.75
3.75
2.50
0.0
0.00
3.0
-0.20
-3.5
m0.m
Figure 7. Visualizations of few latent/state trajectories through time for the Walker agent. Top: Trajectories of several latent variables from the disentangled latent space. Bottom: Trajectories several rigid body orientations.
Given the low dimensionality of the disentangled latent representations and the fact that they are disentangled, we hypothesize that their trajectories through time may correspond trajectories of proprioceptive state variables such as rigid-body positions and orientations. We visualize the trajectories of individual latents through time for a single episode alongside the trajectories of several proprioceptive state variables in Figure 7. Unfortunately, the mappings sources to the disentangled latent space likely do not correspond 1:1 with the proprioceptive state, given that the learned mappings are arbitrary and not unique. However, upon visual inspection, we find that latent trajectories through time exhibit oscillatory behavior patterns similar to that of rigid body orientations from the state representation. Given that the representation is task optimized and that the agent is the only dynamic object in the scene, suspect that the model learns to disentangle and primarily represent the agent's rigid bodies and their orientations. Recovering all or parts of the proprioceptive state representation via unsupervised learning from high- digh-dimensional data is an interesting future research direction.
15

---

Zero shot Generalization of Vision-Based RL
A.3. Additional Latent Traversal Plots
Latent traversals for the other reported DMControl tasks presented here. We also visualize the latent traversals of ALDA trained directly on the color hard environment. We find that the latent traversals for cartpole balance more discontinuous than other tasks. One reason for this might be the lack of balanced data and data diversity of the cartpole replay buffer. The (near) optimal policy is achieved quite early on, after which most images collected are of the cartpole upright and roughly in the same z-position. The lack of data diversity likely makes it difficult to learn a representation in which the latent traversals are more continuous physically plausible. Interestingly, this phenomenon does not seem to affect performance on the "color hard" evaluation environment, although we suspect there are performance gains to be had on DistractingCS the latent interpolations were smoother. We leave an investigation into the effects of data balancing and data diversity on downstream generalization performance future work.
Latent Dimension
Latent Traversal
Figure 8. Latent traversals of the disentangled latent vector when training ALDA directly on the "color hard" environment.
Figure Latent traversals for cartpole balance. Each row corresponds latent dimension that is traversed via lincar interpolation while all other dimensions are held fixed.
16

---

Zero shot Generalization of Vision-Basse Based RL
Figure 10. Latent traversals for the finger spin task. Each row corresponds to latent dimension that is traversed via linear interpolation while all other dimensions are held fixed.
Figure 11. Latent traversals for the ball in cup catch task. Each row corresponds to latent dimension that is traversed via linear interpolation while other dimensions are held fixed.
A.4. Beta Study
We perform an analysis of the effects of different 33 values on ALDA's performance. The memory retrieval dynamics are reintroduced here for the reader's reference:
2d, 1oftmax(33L1(fo000) vy)) @ o.y
Small values of 3 result in a more even distribution of the probability mass between latent values per codebook, which implies that the output will be weighted sum of different latent values (or memories under the Hopfield interpretation). We choose three different values, B = (1, 10, 50), and compare with the main result (B- 100) presented in the paper on the Walker domain. shown in 12. To our surprise, lower 3 values have little to no effect on generalization performance and, in fact, increase training performance. This perhaps challenges the assumptions made in (Hsu et al., 2023) about the requirements disentanglement via latent quantization, but admittedly requires further analysis, which we leave to future work.
17

---

Zero shot Generalization of Vision-Based RL
Walker Walk Train
Walker Walk Color Hard
Walker Walk Distracting
1000
1000
1000
800
800
800
Reward
600
600
600
400
400
400
200
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
environment steps (1e6)
environment steps (1e6)
environment steps (1e6)
ALDA (B = 1)
ALDA (B = 10)
ALDA (3 = 50)
ALDA (B 100)
Figure 12. ALDA is trained the Walker domain with different values of 83 with evaluations periodically performed the "color hard" and DistractingC'S environments. Average of 4 seeds, shaded region represents a 95% bootstrapped confidence interval.
A.5. Framestack Ablation
We provide an ablation comparing against version of ALDA where the encoder receives as input, and the decoder predicts a stack = 3 of frames, i.e., the observation size is (9 x 64 x 64). Since the downstream 1D-CNN layer is no longer necessary, we remove this layer from the variant. We refer to this variant as "ALDA (framestack)" and present results on the Walker domain in Figure 13. We also provide latent traversal visualizations of AL.D (framestack), shown Figure 14.
Walker Walk Train
Walker Walk Color Hard
Walker Walk Distracting
1000
1000
1000
800
800
800
Reward
600
600
600
400
400
400
200
200
200
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
environment steps (1e6)
environment steps (1e6)
environment steps (1e6)
ALDA
ALDA (framestack)
Figure 13. Ablation comparing AL.DA to ALDA (framestack) on the Walker domain. Average of seeds, shaded region represents a 95%
confidence interval.
From the plots, we can see that while the generalization performance ALDA (framestack) does not degrade entirely, does suffer when compared to ALDA. The latent traversal visualizations show that ALDA (framestack) does not disentangle the dynamics of individual bodies of the agent the way ALDA does. Many latent dimensions, for example, latent dim 8 (5th row from the bottom), affect is that ALDA (framestack) does better at capturing and disentangling temporal information, given that it sees stack of consecutive frames, but struggles to disentangle sources of singular images corresponding to non-temporal information, whereas ALDA excels at the latter.
aspects of the agent when interpolating the latent values. One possible explanation
18

---

Zero shot Generalization of Vision-Basce RL
Figure 14. Latent traversal visualizations for ALDA (framestack).
A.6. Implementation Details
Our SAC implementation is based on (Yarats & Kostrikov, 2020).
A.7. DMControl Generalization Benchmark
The DMControl Generalization Benchmark (DMCGB) (Hansen & Wang, 2021) is extension of the DeepMind Control Suite (DMControl) (Tassa et al., 2018). DMControl is a software stack consisting of physics simulation for suite of common reinforcement learning, continuous control environments. DMCGB introduces visual distribution shifts to benchmark visual generalization performance of models trained on these environments from image observations.
A.7.1. ACTOR AND CRITIC NETWORKS
Following (Yarats & Kostrikov, 2020), we use double Q-networks, each of which is 3-layer multi-layer perceptron (MLP) with 1024 hidden units per hidden layer and GeLU activations after all except the final layer. The actor network is similarly 3 layer MLP with 1024 hidden units per layer and GeLU activations on all but the final layer.
A. ENCODER, DECODER, AND LATENT MODEL
We use the same encoder/decoder architectures as (Hsu et al., 2023), with the exception that we replace all leaky ReLU activations with GeLUU. We instantiate the codebooks for the latent model with values evenly spaced between [-1, 1].
A. A.7.3. HYPERPAREMETERS
We list of common hyperparameters that are used in all domains.
19

---

Zero shot Generalization of Vision-Based RL
Parameter
Value
Replay buffer capacity
lc6
Batch size
128
Latent model temperature
100
Number of latents |zu|
Number of values per latent V,
12
Encoder weight decay A
0.1
Decoder weight decay Ao
0.1
Frame stack
3
Action repeat
2 for finger spin otherwise 4
Episode length
100
Observation space
(9 64 x 64)
Optimizer
Adam
Actor/Critic learning rate
le-3
Encoder/Decoder learning rate
lc-3
Latent model learning rate
le-3
Temperature learning rate
le-4
Actor update frequency Critic update frequency
Discount y
0.99
Table Common hyperparameters for SAC and ALDA.
A.8. ALDA Pseudocode
Algorithm 1 ALDA Forward Pass
Input: Observation o, encoder fa, latent model Lg, history encoder hy-
E BBXCKKHH + +0 E BBXXXXXXW // rearrange the framestack dimension
Zcont E ER Bsx, fo(
3d RBIxn= (z association step using the latent model
2d BBfxn, + 2d TR. xxn, // rearrange the framestack dimension
za (2(2)) encode temporal information
return 2d
Algorithm contains pseudocode for ALDA's forward pass. The observation o can be an in-distribution sample during training, or an OOD sample during evaluation. Post-training, when presented with in-distribution samples, the association step unlikely to significantly change Zcont, since Zcont will map very close to the values learned by the latent model Ig. However, when presented with OOD samples, zcont more likely change since certain dimensions of zcone may map far away from the corresponding dimensions of the latent model. Algorithm 2 shows how the latent model Lp maps potentially OOD continuous latent vectors to in-distribution values using modern Hopfield retreival mechanisms.
20

---

Zero shot Generalization of Vision-Based RL
Algorithm 2 Associative Latent Dynamies
Input: Continuous latent vector Zpont BB latent model 4p
for i - to n2 do
W. - (oftmax(BBL1<cnnl) V:l) @ V; // compute weights for how similar Scont, is to each value the 1'th codebook.
Ed. Eiv, W;
end for return Zd
A.9. Latent Study
Walker Walk Train
Walker Walk Color Hard
Walker Walk Distracting
1000
1000
1000
800
800
800
Reward
600
600
600
400
400
400
200
200
200
0.0
0.1
0.2
0.4
0.5
0.0
0.2
0.3
0.5
0
0.0
0.1
0.3
0.4
environment steps (1e6)
environment steps (1e6)
environment steps (1e6)
ALDA (Zg = 12)
ALDA (Zg 32)
ALDA (z, 64)
ALDA (z, = 128)
Figure Training (left) and evaluation (middle and right) results of increasing the dimensionality ALDA's latent space on the Walker Walk task. Results averaged over seeds. Shaded region represents 95% bootstrapped confidence interval.
We examine the effects increasing the dimensionality the latent space on the performance of the Walker "Walk" task and present the results in Figure 15. The baseline model (zæ = 12) performs the best on the training and "Color Hard" evaluation tasks, and that performance drops as the dimensionality of the latent space increases. Since disentangled representation learning methods try approximate the number of ground truth sources of variation of the data distribution, it is possible that setting& 2d to values far away from the number of true sources cause the performance degradation we observe in Figure 15.
21

---

Zero shot Generalization of Vision-Based RL
A.10. Critic Gradient Ablation
Walker Walk Train
Walker Walk Color Hard
Walker Walk Distracting
1000
1000
800
800
Reward
600
600
600
400
400
400
200
200
200
0
0.0
0.1
0.2
0.3
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
environment steps (1e6)
environment steps (1e6)
environment steps (1e6)
ALDA
ALDA (CG)
Figure Ablation backpropogating the critic gradients to the encoder (ALDA CG) compared to the standard model. Results are averaged over 4 seeds. Shaded area represents 95% bootstrapped confidence interval.
In this experiment, we examine the effects of backpropagating the critic's gradients through the latent model and back the encoder. The results presented in Figure 16. We refer to the ALDA variant with critic gradients enabled as "ALDA (CG)" and compare on the Walker "Walk" task. ALDA (CG) performs worse on all training and evaluation tasks. We suspect that backpropagating the critic gradients to the encoder affects the ability for the encoder-decoder pair to disentangle sources of the image distribution, since disentangled representation learning methods typically study disentanglement (Variational) Autoencoders without competing auxiliary objectives.
22

---

Zero shot Generalization of Vision-Based RL
A.11. SVEA with Other Data Augmentations
Walker Walk Train
Walker Walk Color Hard
Walker Walk Distracting
1000
1000
1000
800
800
800
Reward
600
600
600
400
400
400
200
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
0.0
0.1
0.2
0.3
0.4
0.5
environment steps (1e6)
environment steps (1e6)
environment steps (1e6)
ALDA
SVEA (image overlay)
SVEA (color easy)
Figure 17. Comparison of ALDA to SVEA (image overlay) and SVEA (color easy), where SVEA (color easy) is trained directly the color easy environment as data augmentation.
SVEA (image overlay) leverages the Places dataset of 1 million realistic images, which likely puts the DMCGB evaluation tasks in-distribution. However, how well model handles OOD data i.. how well it performs extrapolative generalization is an important consideration for generalist agents. In this study, we train SVEA directly on the color easy evaluation environment in DMCGB as form of weaker data augmentation (DA). Since the color easy environment only slightly randomizes the scene RGB values, it puts the color hard and Distracting CS environments cut-off-dsttrbbus and allows us to tease out how well data augmentation helps with extrapolative generalization. The results are presented in Figure 17.
for SVEA,
We find that SVEA (color easy) performs worse than both SVEA (image overlay) and ALDA the training and Color Hard environments, and equals ALDA in performance Distracting CS. This suggests that not all forms of DA are equally useful, and that DA should be used with caution e.g. when the range of possible variations on the data is bounded and it is computationally tractable to expose the model to all variations. While DA has shown to be useful on toy benchmarks, datasets, or in constrained lab settings, the results presented here suggest that this approach will not scale to, for example, in-the-wild deployments in the real world where the data scale and variance are practically unbounded. The results also highlight the need for models that are robust distribution shifts without knowing or training on the distribution shifts apriori.
23


Reviews:
[{"review_id": "Uv1j4UtXRS", "reviewer": "Reviewer_Tie6", "summary": "The authors present ALDA - an approach for training disentangled representations along with off-policy learning for OOD generalization.\n\nThey build upon existing QLAE-based latent model, which is SOTA disentanglement method, which uses latent space dimensions, each having their own codebook. \n\nThe prove that data augmentation is weak disentanglement and derive a novel loss function for their approach for simultaneous representation and policy training. For temporal sequences, they feed in batches of data to the model, fed to a  1-D convolution for the actor-critic networks, and batch reconstruction is done directly on the latent space. \nThey compare DARLA, SAC+AE, RePo, SVEA with ALDA on DeepMind Control Generalization Benchmark  for \"color hard\" and \"distracting cs\" cases.\n\nALDA beats all baselines except SVEA which uses data augmentation. The authors argue that the test set visuals are also covered with external data augmentation in SVEA, thereby not showing true generalization performance. The authors also present interesting visuals to showcase disentanglement of their model using the produced latents.\n\n### Update After Rebuttal\nThanks to the authors for the rebuttal. I enjoyed reading the rebuttal, the additional experiments and the arguments made in the rebuttal.\nThe paper has its merits. I think the idea that the authors present is interesting.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "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. \n\n**I think it would be nice to discuss this for SVEA vs ALDA since that would be an interesting statistic to see.**\n\nThe 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. \n\nWe 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. \n\n[1] \"Maniskill3: Gpu parallelized robotics simulation and rendering for generalizable embodied ai.\"\n\n**Why did the authors not try SVEA + ALDA?...**\n\nALDA 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.\n\nIn  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.\n\n[1] \"Progress and limitations of deep networks to recognize objects in unusual poses.\" \n\n[2] \"On the Ability of Deep Networks to Learn Symmetries from Data: A Neural Kernel Theory.\"\n\n**The environments the authors use are toy control environments...**\n\nWe 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.", "claims_and_evidence": "- **Data augmentation is weak disentanglement**: I think this is a fair idea and the authors discuss a theoretical proof in detail.\n- **Data augmentation requires larger models, training data, longer training times, and have greater training instability**: While the authors discuss this in the introduction, I think it would be nice to discuss this for SVEA vs ALDA since that would be an interesting statistic to see.\n- **If a data-driven model can generalize better with less data, then it will scale better with more data**: This is a fair point. The authors do provide some reasoning behind this based on their results on ALDA. However, if this is actually true, why did the authors not try SVEA + ALDA (i.e. data augmentation and disentanglement together)? Is it hard to implement such an approach? If not, this would directly prove this proposition. In a world where data is cheap to generate/obtain, data augmentation shouldn't not be used.", "methods_and_evaluation_criteria": "**Methods**\n- The authors train ALDA approach on four tasks for DeepMind Control Suite.  They evaluate on \"color hard\" and \"distracting cs\" environments. This is a fair way to evaluate the method. However, the authors should look into potentially harder tasks (e.g. navigation/rearrangement/manipulation). The environments the authors use are toy control environments, which are far from the real-world environments. There might be other benchmarks that the authors could explore.\n- The baselines are of varied kinds including data augmentation, disentanglement, etc. This provides a nice overview/comparison of their approach with other approaches.\n\n\n**Evaluation Criteria**: They use episode reward for comparison, which is standard in control tasks. The authors mention that it is not possible to directly evaluate disentanglement, and therefore show qualitative examples in Fig 6, which makes sense.", "theoretical_claims": "- The authors prove a theorem showing that data augmentation is weak disentanglement of latent space. The discussion is sound, and shows that in order to achieve a latent representation only relevant to the task, we would have to gather data from all task-irrelevant source, which would be unrealistic in the real world.\n- There is also discussion in Section 4 on how the loss function for ALDA is created based on theory on attention-based Hopfield networks and QLAE dynamics. I have a high level intuition of this idea, but I am not entirely certain on the correctness of the entire discussion.\n- There is a set of proofs in appendix too, which I did not read in detail.", "experimental_designs_or_analyses": "- They show that QLAE achieves better performance over the course of training compared to BioAE (another disentanglement method), showing better OOD generalization.\n- The analyze the results from various training and representation strategies in Figure 5, discussing the performance of ALDA against other approaches.\n- The analysis of different latents in Figure 6 is also interesting and showcases disentanglement.", "relation_to_broader_scientific_literature": "I think the approach is overall interesting. However, the authors could present results on more environments, tasks for a comprehensive study of the approach. Otherwise, if this approach is only applicable/useful in toy control problems, within a limited number of scenarios, then it might not have a lot of impact in the field.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "A well-written paper overall. I like the motivation, the buildup of proofs, and the choice of baselines. The authors could have tried on more benchmarks and set of problems, however.", "other_comments_or_suggestions": "- Line 96 - incorrect quotes around \"weak\"\n- Line 145 - incorrect quotes around \"random convolution\"\n- Line 142, Col2 - incorrect quotes around \"irrelevant\n- Line 157 - don't -> do not\n- Line 295 - we've -> we have\n- Line 290, Col 2 : incorrect quotes around \"color hard\"\n- Inconsistent formatting of color hard and distracting cs throughout the paper.\n- Line 410 - Incorrect quotes\n- Line 418 - Incorrect quotes\n- Line 428 - That's, isn't - informal usage.", "questions_for_authors": "- How easy or hard is it to transfer the approach to on-policy algorithms?", "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": "", "extra_scores": {}}, {"review_id": "5zCijZh4yK", "reviewer": "Reviewer_cwqS", "summary": "The paper introduces Associative Latent Disentanglement (ALDA), an approach to zero-shot generalization in vision-based reinforcement learning (RL) without relying on data augmentation. ALDA leverages disentangled representations and associative memory mechanisms to enable RL agents to generalize to novel environments by factorizing latent spaces into modular components, allowing for independent adaptation of task-relevant and task-irrelevant features.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "Thank you for your feedback and suggestions on how we could improve the paper's clarity. We respond to individual comments and questions below.\n\n**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.**\n\nBecause 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. \n\n**How does ALDA perform if the associative memory component is removed or replaced with a simpler alternative?**\nALDA 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.\n\n**Have you tested ALDA on more complex or real-world RL tasks beyond DeepMind Control Suite? How well does it scale?**\nPlease see our response to Reviewer Tie6 under *\"The environments the authors use are toy control environments...\"*.", "claims_and_evidence": "The claims in the paper are generally supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are well-aligned with the problem of zero-shot generalization in vision-based RL.", "theoretical_claims": "The theoretical claims, particularly the connection between data augmentation and weak disentanglement, appear well-reasoned.", "experimental_designs_or_analyses": "he experimental design is generally sound, with well-chosen benchmarks, multiple baselines, and ablation studies. The inclusion of latent traversals and β-study strengthens the analysis.", "relation_to_broader_scientific_literature": "The paper builds on existing work in vision-based RL, disentangled representation learning, and associative memory. It extends prior research on disentanglement in RL (e.g., DARLA) by integrating modern Hopfield networks for associative memory.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n-- The integration of disentangled representation learning with associative memory is novel in the context of vision-based RL. The theoretical perspective connecting data augmentation to weak disentanglement is insightful.\n-- The work addresses a crucial challenge in RL generalization, providing a potential alternative to data augmentation that could improve scalability and efficiency.\n-- Strong baseline comparisons, ablations, and visualization techniques support the claims.\n\nWeaknesses:\n--  Some theoretical discussions, particularly regarding disentanglement and associative memory mechanisms, could be more clearly explained for a broader audience.\n-- While ablations exist, a more controlled comparison of ALDA with and without associative memory would further clarify its unique benefits.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. How does ALDA perform if the associative memory component is removed or replaced with a simpler alternative?\n2. Have you tested ALDA on more complex or real-world RL tasks beyond DeepMind Control Suite? How well does it scale?", "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": "", "extra_scores": {}}, {"review_id": "QfPp15pIoo", "reviewer": "Reviewer_CC4m", "summary": "This paper proposes Associative Latent DisentAnglement (ALDA) that builds on standard off-policy RL towards zero-shot generalization. It learns a disentangled representation from the training data and then uses an associative memory model to recover data points in the original training distribution given OOD data.The authors also prove that data augmentation methods can be considered as a weak disentanglement. Experiments show that the proposed methods can outperform most baselines.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "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": "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.\n\n**However, I have some concerns in its novelty...**\n\nAs 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. \n\nOur other novel contributions are as follows: \n- 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.\n\n- 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.\n\n**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...**\n\nWe 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.\nThe “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.\nOur 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.\n\n**Besides, it would also be helpful to provide some failure cases to show the limitation of the OOD generalization.**\n\nThe 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.\n\n**Since disentanglement is an emphasis of this paper, it is great to have some analysis...**\n\nThe 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.", "claims_and_evidence": "I think all the claims are well supported.", "methods_and_evaluation_criteria": "1. The proposed algorithm is simple and well-motivated. The Association strategy is also elegant and effective.\n\n2. However, I have some concerns in its novelty. Although disentanglement is highlighted throughout the paper, it looks like a simple application of existing QLAE algorithm. There are no specific adaptions for the concrete Vision-based RL scenario. Especially, the authors do not conduct any disentanglement in the temporal domain. This makes it impossible to analyze the dynamic cues from vision inputs despite its importance in many RL tasks.", "theoretical_claims": "The authors provide a proof that data augmentation is a weak disentanglement of the latent space. This provides a good insight to building the connection between previous literature and this work, and also gives a good motivation for learning a disentangled representation. \n\nI do not see any mistakes in the proof.", "experimental_designs_or_analyses": "1. It is great that the authors conduct extensive experiments to compare with multiple baselines on different tasks. \n\n2. 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. I am not sure whether each experiment is in-distribution or out-of-distribution. If they are OOD, how are they OOD? Do they have OOD visual appearances or dynamics? Is the OOD interpolation or extrapolation?\n\n3. Besides, it would also be helpful to provide some failure cases to show the limitation of the OOD generalization. \n\n4. Since disentanglement is an emphasis of this paper, it is great to have some analysis about how the representation is disentangled and what is the physical meaning of each component if possible.", "relation_to_broader_scientific_literature": "There is a concern about novelty mentioned in \"Methods And Evaluation Criteria\".", "essential_references_not_discussed": "[1] shares a very similar insight with this paper about disentangled representation for generalization. Since it is a key contribution of the proposed algorithm, I believe the authors should have some discussions related to [1].\n\n[1] Wu, Zheng, et al. \"Zero-shot policy transfer with disentangled task representation of meta-reinforcement learning.\" 2023 IEEE International Conference on Robotics and Automation (ICRA). IEEE, 2023.", "other_strengths_and_weaknesses": "Please refer to the previous parts.", "other_comments_or_suggestions": "Please refer to the previous parts.", "questions_for_authors": "The authors can consider replying to my concerns mentioned in previous parts. I am not an expert in the area of RL, so feel free to point them out if I have any misunderstandings.", "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": "", "extra_scores": {}}, {"review_id": "ijS77UGEsv", "reviewer": "Reviewer_N37F", "summary": "The authors propose a method for generalizing vision-based RL policies in the presence of image distribution changes or distractors, such as modifying the color scheme of the image. There method aims to learn a representation that disentangles the different components of an image, specifically into two types: task-relevant and task-irrelevant features. Ideally, if a given feature is either only task relevant info or only task irrelevant info, then generalization is possible because for some downstream image with distractors, those distractors should get filtered into the irrelevant features while the relevant features still capture only the needed info.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "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": "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. \n\n**Bi-simulation may be worth looking into.**\n\nThank 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^*(\\textbf{z}, a) = Q^*(\\textbf{z}', a)$ for two states $\\textbf{z}$ and $\\textbf{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. \n\nSince 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.\n\n**A figure showing the architecture and...**\n\nThank you for your recommendation. We will update Figure 1 to include a diagram showing the inference procedure of the associative latent model. \n\n**A figure illustrating the learned representation...**\n\nHere is the first figure requested by the reviewer: https://drive.google.com/file/d/1P0zkA4s4bRJoNWCNhLWmqKLDOUI-GmbW/view?usp=sharing.\nWe 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.\n\n**Additionally, show if you can use this approach to modify the images...**\n\nRegarding 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. \n\n**What happens if you use to large of a latent space |z_d|? Does this break things?**\n\nThis 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.\n\n**Since each dimension is treated independently...**\n\nThe 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.", "claims_and_evidence": "Their claims, that disentangled visual representations are necessary for generalization, are well-motivated both intuitively and by making some arguments from biology. Their experiments show that their proposed approach outperforms other approaches without using data augmentation, except for a baseline that uses a large pre-training dataset.", "methods_and_evaluation_criteria": "To my knowledge, the benchmark environments they use are standard for this type of RL.", "theoretical_claims": "I checked the proofs in the appendix and it seems correct to me.", "experimental_designs_or_analyses": "The described experiments are reasonable given the problem formulation, though I did not check the code.", "relation_to_broader_scientific_literature": "This work relates to broader literature on image processing, RL, image based RL, and representation learning. This work seems to be mostly connecting some prior work in disentangled representation learning with image-based RL.", "essential_references_not_discussed": "I am not sure if this is strictly necessary, but bi-simulation (e.g. work from Amy Zhang et al.) may be worth looking into.", "other_strengths_and_weaknesses": "Strengths:\n- Paper is well written\n- The introduction and motivation is very clear\n\nWeaknesses:\n- Some figures providing qualitative examples would be helpful. Specifically, the following two:\n    - A figure showing the architecture and inference procedure for the representation, especially relating to the codebook. This part of the paper was harder to parse, and a figure may clarify things.\n    - A figure illustrating the learned representation. For example, show the same trajectory but with two different coloring schemes, and plot the representations as this trajectory progresses. If the algorithm is working correctly, all relevant features should closely match, while irrelevant features may be wildy different. Additionally, show if you can use this approach to modify the images. For example, modify the observation to be a different color, compute its representation, set its irrelevant features to match the training data, then regenerate the image. You should be able to effectively change the image visuals using this approach without losing the relevant information, such as joint position.", "other_comments_or_suggestions": "Minor:\n- Notation in the background section is a bit unclear. For example, you use R(s_t, a_t) for reward, but you also use r_t. r_t is not defined anywhere, and I assume this is reward at time t. Also, the equation for the bellman residual J(Q) has reward listed twice and a missing parenthesis in the subscript notation. Lastly, the policies objective is to maximize value, so this should either be indicated or a negative sign should be included. \n- Quotation marks are all double end quotes.", "questions_for_authors": "1) What happens if you use to large of a latent space |z_d|? Does this break things? \n2) Since each dimension is treated independently, the scaling of the number of the codebook is exponential in |z_d|. Is this correct? Is there a way around this somehow? Please discuss.", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper addresses zero-shot generalization in RL, leveraging disentangled / composable representations in a latent space learned using the QLAE algorithm. The subject of generalization (zero-shot or otherwise) is timely in RL, and it's nice to see this being directly addressed through representation learning. The comparison to and analysis of data augmentation is good, as this is a common technique for improving generalization, and the conclusions are convincing.

There were some concerns about clarity, but these seem to have been sufficiently addressed. Some additional concerns were raised wrt some of the claims wrt disentanglement in the model (and when in the data augmentation experiments), but I'm satisfied with the answers and the reviewers did not raise any additional concerns.

I therefore recommend the paper is accepted to the conference.
Decision: Accept (poster)