Imagine Beyond ! Distributionally Robust Autoencoding for State Space Coverage in Online Reinforcement Learning
Authors: Nicolas Castanet
Olivier Sigaud
Sylvain Lamprier
Venue: NeurIPS 2025 poster ()
Keywords: Goal-Conditionned Reinforcement Learning
States representation learning
Intrinsic motivation
Exploration
Primary area: reinforcement_learning

Abstract: Goal-Conditioned Reinforcement Learning (GCRL) enables agents to autonomously acquire diverse behaviors, but faces major challenges in visual environments due to high-dimensional, semantically sparse observations. In the online setting, where agents learn representations while exploring, the latent space evolves with the agent's policy, to capture newly discovered areas of the environment. However, without incentivization to maximize state coverage in the representation, classical approaches based on auto-encoders may converge to latent spaces that over-represent a restricted set of states frequently visited by the agent. This is exacerbated in an intrinsic motivation setting, where the agent uses the distribution encoded in the latent space to sample the goals it learns to master.
To address this issue, we propose to progressively enforce distributional shifts towards a uniform distribution over the full state space, to ensure a full coverage of skills that can be learned in the environment.
We introduce DRAG (Distributionally Robust Auto-Encoding for GCRL), a method that combines the $\beta$-VAE framework with Distributionally Robust Optimization (DRO).  
DRAG leverage an adversarial neural weighter of training states of the VAE, to account for the mismatch between the current data distribution and unseen parts of the environment. This allows the agent to construct semantically meaningful latent spaces beyond its immediate experience. Our approach improves state space coverage and downstream control performance on hard exploration environments such as mazes and robotic control involving walls to bypass, without relying on pre-training nor prior environment knowledge.

Imagine Beyond! Distributionally Robust Auto-Encoding for State Space Coverage in Online Reinforcement Learning
Nicolas Castanet Sorbonne Université, CNRS, ISIR, F-75005 Paris, France nicolas. castanet@isir. upmc fr
Olivier Sigaud
Sorbonne Université, CNRS, ISIR, F-75005 Paris, France alivier.sigauddistt wpw.siir...p.m....mmmmmmmmmm upmc. fr
Sylvain Lamprier
Univ Angers, LERIA, Angers, France sylvain 1amprier fr
Abstract
Reinforcement Learning (GCRL) enables agents to au- tonomously acquire diverse behaviors, but faces major challenges in visual envi- ronments due high-dimensional, semantically sparse observations. In the online setting, where agents learn representations while exploring, the latent space evolves with the agent's policy, to capture newly discovered areas of the environment. However, without incentivization to maximize state coverage in the representation, classical approaches based auto-encoders may converge to latent spaces that over-represent restricted set of states frequently visited by the agent. This is exacerbated in intrinsic motivation setting, where the agent uses the distribution encoded in the latent space to sample the goals it learns to master. To address this issue, propose progressively enforce distributional shifts towards uniform distribution over the full state space, to ensure full coverage of skills that can be learned in the environment. We introduce DRAG (Distributionally Robust Auto-Encoding for GCRL), a method that combines the 3-VAE framework with Distributionally Robust Optimization. DRAG leverages an adversarial neural weighter of training states of the VAE, to account for the mismatch between the current data distribution and unseen parts of the environment. This allows the agent to construct semantically meaningful latent spaces beyond its immediate experience. Our approach improves state space coverage and downstream control performance on hard exploration environments such as mazes and robotic control involving walls bypass, without pre-training nor prior environment knowledge.
Introduction
Coal-Conditioned Reinforcement Learning (GCRL) enables agents master diverse behaviors in complex environments without requiring predefined reward functions. This capability is particularly valuable for building autonomous systems that can adapt to various tasks, especially navigation and robotics manipulation environments (Plappert et al. 2018 Rajeswaran et al., 2018 Tassa et al.) 2018 Yu et al. 2021]. However, when working with visual inputs, agents face significant challenges:
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

observations are high- fieh--dmensional and lack explicit semantic information, making intrinsic goal generation for exploration, reward calculation, and policy learning substantially more difficult.
A common approach to address these challenges involves learning a compact latent representation of the observation space, that captures semantic information while reducing dimensionality [Nair iet al., 2018 Colas et al. 20118] Pong et al. 2019 Hafner et al. 22119a Laskin et al., 2020 Gallouèdec and Dellandréa 2023]. J. Assuming a compact
representation that encodes the main variation factors from the whole state space, agents can leverage latent codes as lower- dimensional inputs. In the GCRL setting, agents are conditioned with goals encoded as latent codes, usually referred to as skills [Campos et al., 2020), which reduces control noise and enables efficient training. Thus, many works build agents on such pre-trained representations of the environment Mendoncca et al. 2 training data, by assuming the availability of a state distribution from which sampling efficient. Without such knowledge, some methods use auxiliary exploration policies for data collection [Campos) Yarats et al., 202T Mendonca et al. 202T], such as maximum entropy strategies Hazan 2019 or curiosity- driven exploration Pathak et al.) [2017]. but these often struggle in high-
pnfformattion-preervinn
Zhou et al. [025], I, but usually leave aside the question of the collection of
et al. et al. dimensional or stochastic environments due to density estimation and dynamics learning difficulties.
2020
An alternative, which we focus on in this work, is the online setting; the representation is learned jointly with the agent's policy, using rollouts to train an -ecoeerd This allows the representa- tion to evolve with the agent's 's progress and potentially cover the full state space. Unlike auxiliary exploration, GCRL-driven representation learning aligns training with meaningful behavioral distribution, which naturally acts as a curriculum. A representative approach is RIG [Nair et al., 2018], where a VAE encodes visited states, and latent samples from the prior are used as "imagined" goals-creatin feedback loop between representation and policy learning.
However, this process suffers from key limitations. common critique is that continual encoder training leads to distributional shift-a well-known issue in machine learning-which destabilizes policy learning and reduces exploration diversity. In this collaborative setting, we identify two distinct sources of shift: one from the agent's 's perspective, where the meaning of the latent codes it receives as inputs continuously evolves; and one from the encoder's 's perspective, in the distribution of visited states to be encoded during rollouts. While policy instability caused by distributional shift from the agent's 's perspective can be mitigated using a delayed encoder, we argue that distributional shift in the encoder's input data-i.e.. the states reached during rollouts- is not only desirable, but essential for exploring the environment and expanding the representation. Rather than limiting such shift, we propose to anticipate and deliberately steer it using a principled method, ensuring that it benefits exploration and learning rather than undermining them.
Our main contribution is to leverage Distributionally Robust Optimization (DRO) Delagee and Ye 2010 to guide the evolution of the representation. By integrating DRO with VAE [Higgins et al. [207aa]. we introduce DRAG (Distributionally Robust Auto-Encoding for GCRL), which uses an ad udversarial weighter to emphasize underrepresented states. This allows the agent to build latent spaces that generalize beyond its current experience, progressively covering the state space.
Our contributions are:
We introduce DRO-based VAE framework tailored to GCRL.
We reinterpret SKEW-FIT (Pong et al. 2019 as a non-parametric instance of DRO-VAE.
We propose DRAG, a more stable- parametric DRO-VAE approach to encourage state coverage through adversarial neural weighting.
We show that when encoder learning anticipates distributional shift, explicit exploration strategies become unnecessary in RIG-like methods; the ingful goals. This enables focusing on selecting goals of intermediate difficulty (GOIDs Florensa et al. (018)) to improve sample efficiency.
latent prior alone generates mean-
Our approach improves state space coverage and downstream control performance on hard exploration environments such as mazes and robotic control involving walls to bypass, without pre-training nor prior environment knowledge.

---

Representation
VAE training
Goal sampling
Learning
DRO sampling
Latent priorp(:)
Reinforcement Learning
goal
VAE Encoder
Environment
Pelicy
reward lunction
12
Latent
Replay Buffer
r Fat Rg]
Policy
Action
To(- [Ets #g))
General framework of online VAE representation learning in RI. Green: RL loop using the VAE encoder convert high-dimensional states 8, to latent states žt- Blue: latent goal 2g sampling (from prior distribution or replay buffer) and selection. Red: Representation Learning with VAE, using data from the replay buffer combined with Distributionally Robust Optimisation (DRO).
Figure 1:
2 Background & Related Work
2.1 Problem Statement: Unsupervised Goal-conditioned Reinforcement Learning
In this work, we consider the multi-goal reinforcement learning (RL) setting, defined as an extended Markov Decision Process (MDP) M =< S, T, A, G, Rg So X >, where S is set of continuous states, 7's is the transition function, A the set of actions, So the distribution of starting states, X the observation function and R, the reward function parametrized by goal g E G. In our unsupervised setting, we are interested in finding control policies that are able to reach any state in S from the distribution of starting states Sg. Thus, we consider that G S. being continuous, we set the reward function Rg as depending on a threshold distance from the goal g: for any state S. we consider the sparse reward function Rs(s) [[]|s gll2 < 8). Rs) >> only possible once per trajectory. For simplicity, we also consider that any pixel observation æ E X corresponds single state s E S. Thus, given an horizon T, the optimal policy is defined as n = max, Eg Et=o (s)]], where T = (sg, 01, $1, ..., ST) is a trajectory, and π (Tl9) the distribution of trajectories given g in the MDP when following policy
2.2 Multi-task Intrinsically Motivated Agents
Multitask Intrinsically Motivated Agents provide powerful framework in GCRL to tackle supervised settings by enabling agents to self-generate and pursue diverse tasks without external prior knowledge of the environment via intrinsic goal distribution. This approach has proven effective for complex problems such as robotic control and navigation, and has also shown benefits in accelerating learning in supervised tasks where goals known in advance [Colas al. 2018
Ren et al.] 2019] [Hartikainen et al.) 22000] Gallouédec and Dellandréa (023). Various criteria have been investigated for the formulation of the intrinsic goal distribution. Many of them focus on exploration, to encourage novelty or diversity in the agent's 's behavior [Warde-Farles Farley et al. 2018
Pong al. 2019 Pitis al. 2020 Gallouédec and Dellandréa, 2023 Kim et al.) 2023 Among them, MEGA [Pitis et al. 2020 defines a density estimator p from the buffer (e.g.. via a KDE) and samples goals at the tail of the estimated distribution to foster exploration. SKEW-FIT Pong al. 2019 maximizes the entropy of the behavior distribution. It performs goal sampling from skewed distribution Pskewed, designed as an importance resampling of samples from the buffer with a rate 1/r to simulate sampling from the uniform distribution. Other approaches are focused on control success and agent progress, by looking at goals that mostly benefit improvements the trained policy. This includes learning progress criteria [Colas et al. 2011]] or the selection of goals of intermediate difficulty (GOIDs) (Sukhbaatar al. [2017] Florensa et al. 2018 Zhang et al.) 2020 Castanet et al.. 2022], not too easy or too hard to master for the agent, depending on its current level. Approaches
from that family, such as GOAL.GAN (Florensa et al., 2018 or SVGG [Castanet et al. 2022 usually
rely auxiliary network that produces GOIDs distribution based success predictor.

---

2.3 Online Representation Learning with GCRL
Variational Auto Encoders (VAEs) present appealing properties when it comes to learning latent state representations in RL. With their probabilistic formulation, the observation space can be represented by the latent prior distribution, which enables several operations to take place, such goal sampling in GCRL [Nair et al., 2018)] Pong et al. 2019 Gallouédec and Dellandréa 2023] and having access to the likeliiol of trajectories for model-based RL and planning [Higgins et al. 2017b Hafner ict al.] 2 Lee et al., 2020b]. The seminal work RIG [Nair et al. 2018). which is the foundation of our paper, is an online GCRL method that jointly trains latent representation and policy T (a 2z, 2g), where 2g is a goal sampled in the latent space, and 2n = (xx) is a VAE encoding of observation x: of the current state. During training, the agent samples goal 2y ~ p(2), with p(z) the prior (typically (0, f)), and performs a policy rollout during I'' steps or until the latent goal and the encoded current state are close enough. The policy is then optimized via policy gradient, using e.g. a sparse reward in the latent space, and the visited states are inserted in a training buffer for the VAE. This framework enables the agent to autonomously acquire diverse behaviors without extrinsic rewards, by aligning representation learning and control. The policy collects new examples for the VAE training, which in turn produces new goals to guide the policy. implicitly defining an automatic exploration curriculum. To avoid exploration bottlenecks, which is the main drawback of RIG, the SKEW-FIT principle introduced in the previous section for the sampling of uniform training goals was also applied in the context of GCRL representation learning, on top of RIG. SKEW-FIT for visual inputs [Pong et al. 2019 is, to our knowledge, the most related approach to our work, which can be seen as an instance of our framework, as we show below.
Beyond generative models based on VAE, other types of encoder-decoder approaches have been introduced in the context of unsupervised RL, including normalizing flows [Lee et al., [020a] and diffusion models [Emami et al., 2023), each offering different trade-offs in terms of expressivity, stability, and sample quality. In addition, contrastive learning methods [Oord ct al. 20018] Srinivas cet al., 2020, Stooke et al., 202T Lu al., 2019 Li et al., [02T]] Aubret et al. [023] have been employed to learn compact and dynamic- aware representations, without relying on reconstruction- based objectives. Some methods rely on pre-trained generalistic models such as Dino to compute
semantically meaningful features from visual observations [Zhou et al., 2025 although usually inducing additional computational cost.
In this work, we use the 3- VAB framework (Higgins et al. 2017a for simplicity and to follow the main trend initiated by RIG Nair et al. [2018. However, the principle introduced in Section could easily be applied to many other representation learning frameworks. The general framework of online VAE representation learning in GCRL is depicted in Figure Compared to RIG, it includes a DRO resampling component, which we discuss in the following.
2.4 Distributionally Robust Optimization
This section introduces the general principles of Distributionally Robust Optimization (DRO) [Delage
and Ye 2010 Ben-Tal et al. 2013 Duchi et al. 2021], developed in the context of supervised machine learning to address the problem of distributional shift, which happens when a model is deployed on a data distribution different from the one used for its training. DRO proposes to anticipate possible shifts by optimizing model performance against the worst-case distribution within a specified set around the training distribution. Formally, given a family of possible data distributions Q. DRO considers the following adversarial risk minimization problem: mingge max,ee B(zw)>~ [e( (4), v)], with é a specified loss function which compares the prediction fo(r) with a given ground truth y.
In the absence of a predefined uncertainty set Q, DRO methods strive to define such an uncertainty set relying on heuristics. This has been the subject of many research papers, see [Rahimian and Mehrotra 2019 for a broad and comprehensive review of these approaches. In the following, we build on the formulation proposed in [Michel et al., 2022], which considers Q as the set of distributions whose KL-divergence w.r.t. the training data distribution p is upper-bounded by given threshold 8.

---

Likelihood Ratios Reformulation Assuming Q as a set of distributions that are absolutely continu- ous with respect to A the inner maximization problem of DRO can be reformulated using importance weights r v) such that q = rp |Michel et al. 2022]. In that case, we have:
B( [((f(())))) v)] ((, |r v) (1)), v)], which is convenient training data is as- sumed to follow p.
Given training dataset = {(z; ))) sampled from p. the optimization problem considered in [Michel et al. 2022] is then defined as:
N
N ( vi =
min max 0
Σ r(,,, () - logr(22 S.t. N
(I)
where the constraint ensures that the q function keeps valid integration property for a distribution (i.e.. Jx.y ((x, y)drdy = 1). The term À logr is a relaxation of a KI. constraint, which ensures that q does not diverge too far from x is an hyper-paramen trade-off between generalization to shifts (low X) and accuracy on training distribution (high A).
that acts
regularizer ensuring
From this formulation, can see that the risk associated to shift of test distribution can be mitigated simply by associating adversarial weights r; := r(r,, yt) to every sample (Ii, yt) from the training dataset, respecting F : x Σ Ti = 1. That said, r can be viewed as proportional to a categorical distribution defined on the components of the training set.
Analytical solution:
Given any function h: X' -> TR, the distribution q that maximizes E. (h(c)) +
XHg. with Hg the Shannon entropy of q, is the maximum entropy distribution g(z) xx caa// Thus, can easily deduce that the inner maximization problem of [ has analytical solution
in r N
(proof in Appendix C.I.. The spread of Q is controlled with 1= 11
temperature weight X, which can be seen as the weight of a Shannon entropy regularizer defined on discrepancies of g regarding p.
Solution based on likelihood ratios: While appealing, it is well-known that the use of this analytical solution for may induce an unstable optimization process in DRO, weights may vary abruptly for even very slight variations of the classifier outputs. Moreover, it implies individual weights, only interlinked via the outputs from the classifier, while one could prefer smoother weight allocation regarding inputs. This is particularly true for online processes like our RL setting, with new training samples periodically introduced in the learning buffer.
Following Michel et al. [202 2021]. likelihood ratios defined as functions (o, parameterized by neural network fe, where set:
we rather focus in our contribution in the next section on
(1 y:) = 71 jj=1 expFo(35.3))
V mini-batch (1,,,;)))7
where fø is periodically trained on mini-batches of n samples from the training set, using fixed current 0 parameters, according to the unconstrained inner maximization problem of (T) for a given number of gradient steps. This parameterization enforces the validity constraint at the batch-level, through batch normalization hard-coded in the formulation of rg.. Though it does not truly respect the full validity constraint from (I) in the case of small batches, this performs well for commonly used batch sizes in many classification benchmarks (Michel al. 2022]. Classifiers obtained through the alternated min-max optimization of (1) are more robust to distribution shifts than their classical counterparts. Using shallow or regularized networks fu advised, as strong Lipschitz-ness of r(x, y) allows to treat similar samples similarly in the input space, which guarantees better generalization and stability of the learning process. These generalization and stability properties lack to non- non- aon-aarmetric versions of DRO, such as version using the analytical solution for maximizaainiii presented above, which could be viewed as the optimal rw based on an capaciyy neural network fa- In the next section, we build on this framework to set representation learning process for RL, that encourages the agent to explore.
'In the situation where all distributions subset C X x Y and all E Q (A) > 0 only p( A) > 0.
are absolutely continuous with respect to p. for all measurable
"This can be seen easily, observing that: K KL (p) q())) ((g(()))))()))) log)
5

---

3 Distributionally Robust Auto-Encoding for GCRL
To anticipate distributional shift naturally arising in GCRL with online representation learning, we first propose the design of a DRO- VAE approach, which was never considered in the literature to the best of our knowledge Then, we include it in our GCRL framework, named DRAG, see Figure |
3.1 DRO-VAE
Classic VAE learning aims at minimizing the negative log-likelihood: 4. = -Bgep((a) log (2)) with po.e(z) the predictive posterior, which can be written as: Pa.d (z( (1))p)(r)2)dd where p=z) is a prior over latent encoding of the data I, commonly taken as N 7). and po [x] is the likelihood of x knowing z and the parameters of the decoding model 0. Given that this marginalization can be subject to very high variance, the idea is to use an encoding distribution go (z|xx) to estimate this generation probability [Kingma and Welling 2013 For any distribution qo such that (=(2)z)
for any z with p(z) > O, we have: (0))(z)= = = (==_) (/ol/////////////// (z(rr.
In our instance of the DRO framework, we thus consider the following optimization problem:
min 1maxx-Br-fa -Bx-((s) s
(3)
0,0 EE=
where E is the uncertainty set of distributions of our DRO- VAE approach. As standard DRO, we introduce weighting function r X' - IR+ which aims at modeling P and respects both validity (i.e., Exr(z) = 1) and shape constraints (i.e., K ((lll)) S for given pre-defined e 0). Relaxing the KL constraint by introducing À hyper-parametent we can get similar optimization problem as in classical DRO. However, as log po.o (x) is intractable directly, we consider a slightly different objective:
& for any distribution & E E,
ming 6 - Exup r (z) log (z))(z),
(4)
with arg max EEep r(x) ((z) AB r(x) logr(x), r:Ep==1
where the only difference is that the inner maximization considers an approximation (o so log (1r). (x))(a)) is estimated via Monte-Carlo importance sampling, as (1)( log Σe= exp(log po (x)2²) + lo") log (2 [z)) log(M) given M samples 2d from qs (2 (r) for any which can be computed accurately (without loss of low log values) using the LogSumExp trick.
This formulation suggests a learning algorithm which alternates between updating the weighting function and optimizing the VAE. At each VAE step, the considering a weighted version of the classical ELBO. Denoting as r the weighting function adapted for current VAE parameters via (4), we have:
networks optimized
E (x) log p(=)po(z(2) 4o (2(1)))
E (xxp(e)((z)) (o
r(x;) 1
log po (2/12) K (L(o (5)
n
m
where this approximated lower-bound CDRO-VA Carlo based mini-batches of data points (x:) from the training buffer and latent codes (zf) ) for each data point 24. Optimization performed using the reparameterization trick, where each latent code zi is obtained from deterministic transformation of white noise ef ~N(00 I).
{2; H can be estimated each step via Monte-
3.2 DRAG
Plugging our DRO-VAF in our GCRL framework as depicted in Figure I thus simply comes down to weight (of resample) each sample v; taken from the replay buffer with weight ri-
This is not surprising, in classical VAE settings, the aim is to model with the highest fidelity.

---

As shown in Section 2.4 classical DRO maximization in Equation 44 has a closed-fore solution: show that in our GCRL setting, this reduces the SKEW-FIT method, where VAE training samples are resampled based on their Pskewed distribution.
(a.))a(a.)A In Appendix C.2,
We claim that the instability of non-parametric DRO, well-known in the context of supervised ML, amplified in our online RL setting, where the sampling distribution p depends on the behavior of constantly evolving RL agent. Our DRAG method thus considers the parametric version of the weighting function, implying a neural weighter fø: X - R as defined Equation (2) trained periodically for given number of gradient steps the inner maximization problem of
In our experiments, we use for our weighter fa a similar CNN architecture as the encoder of the VAE, but with a greatly smaller learning rate for stability (as it induces a regularizing lag behind the encoder, and hence enforces a desirable smooth weighting w.r.t. the input space). We also use a delayed copy of the VAE to avoid instabilities of encoding from the agent's perspective. The full pseudo-code of our approach is given Algorithm| Appendix B
Experiments
Our experiments seek to highlight the impact of DRAG on the efficiency of GCRL from pixel inpu we structure this section around two experimental steps that seek to answer the two following research questions in isolation:
As depicted in Figure 2
Representation Learning strategy: Does DRAG helps overcoming exploration bottlenecks of RIG-like approaches? (Figure 2a
Latent goal sampling strategy: What is the impact of additional intrinsic motivation when using the representation trained with DRAG? (Figure 2b
Representation Learning
Random
Goal sampling Latent prior p(a)
Representation Learning
Goal sampling
Sampling
VAE traisino.
training
Sieew
Sampling
Sampling
DRO
sheu sntestonn Porier
Leart pea
Sampling
Pige
Reinforcement Leaminng
poliy
Proeps
Papkay Butter
Learning
pollcy
Replay
Paill
Buller
Pyni
(a) Representation Learning strategy experiments.
(b) Latent Goal selection strategy experiments
Figure 2: Our two questions: (a) how does DRAG perform representation learning approach? (b) how does DRAG impact goal sampling approaches from the literature?
In all experiments, the policy we (-12, 2g) is trained using the TQC off off-policy RL algorithm (Kzznetsvy et al., 2200)), conditioned on the latent state and goal. Learning is guided by a sparse reward in the latent space, defined as Rg (s) r(2z, 2g) 2| zyla < 8). where æ stands the pixel mapping of s and Zμ its encoding. Experimental details are given in Appendix
Evaluation Our main evaluation metric is the success coverage, which measures the control of any policy π on the entire space of states, defined as:
((T) Σ [[ E E 2lla < 8|T T[-1=0 = (x(())))); (X)))) Sol], 9EG
where g is test set of goals evenly spread on S, and X() stands for the projection of the input state its pixel representation. Note that goal achievement is measured in the true state space. The knowledge of S is only used for evaluation metrics, remaining hidden to the agent.
'The code is available at http://giitabccmmmiooolaassttttt

---

Environments We consider two kinds of environments, with observations as images of size 82 x 82. Additional results on image size 128 x 128 are also provided in appendix [.5] In Pixel continuous PoinMazes, we evaluate the different algorithms over 4 hard-to-explore continuous point mazes. The action space a continuous vector (óz, 9y) [0. 1]2. Episodes start at the bottom left corner of the maze. Reaching the farthest area requires at least 40 steps in any maze. States and goals are pixel top-down view of the maze with a red dot highlighting the corresponding xy position. Pixel Reach-Harrd-Wal is adapted from the Reach-v2 Meta World benchmark (Yu et al. [021]. We add 4 brick walls that limit the robotic arm's s ability to move freely. At the start of every episode, the robotic arm is stuck between the walls.
4.1 Representation Learning strategy
In this initial stage of our experiments, we set aside the intrinsic motivation component GCRL and adopt the standard practice of sampling goals from the learned prior of the VAE, i.c. 2g ~N(00 I). Our objective is to compare DRAG, which trains the VAE on data sampled from a distribution proportional to r (2)) Pro (x), with the classical RIG approach, which samples uniformly from the
replay buffer, i.e. Prig (x) x Pao (z)) We also include a variant taken from SKEW-FIT, where the VAE trained samples drawn from skewed distribution defined Pskewed (2) Pmg (ar) with
an hyper-parameter that acts analogously to À from DRAG (with & = - 1/X, see Appendix C.2
DRAG (ours)
Skew-fit
RIG
Maze
Maze
Ma.20
Maze
Reach Walls
02
Env steps Ie61
Median
IQM
Mean
Optimality Gap
RIG
Skew-Fit
DRAG (ours)
0.45
0.60
0.75
0.90
0.45
0.60
0.75
0.4
0.5
0.6
0.7
0.8
0.2
0.3
0.4
0.5
0.6
Success Coverage
Figure 3: Evolution of the success coverage over PointMazes and Reach-Hard-Wai Hall- environments (6 seeds each) for 4M steps (shaded Mean, Mean and Optimality Gap of success coverage across the all runs after 4M steps. We plot these metrics and confidence intervals using the Rliable library [Agarwal al. 2021
standard deviation). Bottom: Median, Interquartile
Success Coverage evaluation Results in Figure show the evolution of the success coverage 4M steps. We that DRAG significantly outperforms RIG and SKEW-FIT. These results corroborate that online representation learning with RIG is unable to overcome an exploration bottleneck. Therefore, RIG agent can only explore and control very small part the environment. Furthermore, the success coverage of RIG is systematically capped to certain value. SKEW-FIT is often able overcome the exploration bottleneck but suffers from high instability, which indeed corresponds to the main drawback of -orammemettiic DRO, highlighted in Michel et al. [2021] Therefore, SKEW-FIT is unable reliably maximize the success coverage. DRAG is more stable due to the use of parametric likelihood ratios and is able to maximize the success coverage. Additional results and visualizations on these experiments are presented in Appendix D.I. In particular, they show greatly better organized latent space with DRAG than with other approaches. We also show in appendix [D.2 that DRAG obtains better latent representations in terms of the trustworthiness score |Venna and Kaski 2001), which measures the preservation of
the other hand,

---

the local neighborhood structure in the input space. Metrics regarding computational runtime are reported in appendix D.6
4.2 Latent Goal selection strategy
Our second question is on the impact of goal selection on the maximization of success coverage. As depicted in Figure 2b] we compare several goal selection criteria from the literature, on top of DRAG, as follows. The selection of each training latent goal is performed as follows. First, we sample set of candidate goals Cg { ji= from the latent prior (0, I). Then, the selected goal resampled among such pre-sampled candidates Cg. using one of the following strategies. Among them, MEGA and LGE only focus on exploration, GoalGan and SVGG look at the success of control:
MEGA Minimum Density selection, from PPiis et al. 2020]: Pmega (zn) x 8e(2g), where 8e is Dirac distribution centered on c, which corresponds to the code from C with minimal density (according KDE estimator trained on latent codes from the buffer);
Minimum density geometric sampling, from [Gallouédec and Dellandréa 2023 :: Pige (Fg) =
(1 (1-p) p) (<(,,)<<pppp>>) p, where R(2g) stands for the density rank of Eg (according trained KDE) among candidates Cg. and is the parameter of a geometric distribution.
GoalGan - Goals of Intermediate Difficulty selection, from [Florensa et al. 2018 Pgoid (z₂) = u(GOIDs), GOIDs = {gg E Cgl. LplPmin < (==g) < P Pmea], J, where )(ag) is success prediction
model, and thresholds are arbitrarily set as min = 0.3 and Pmaz 0.7, following recommended values in Florensa et al. 2018);
(sxl)((2g) x exp (fa. (D(g)), where a success prediction model trained simultaneously from rollouts, fa.8 is beta distribution controlling the target difficulty, = B = 2, smoothly emphasize goals such that D(zg) R 1/2.
SVGG Control of goal difficulty, from [Castanet et al.
2022
DRAG
GRAG + LGG
DRAG MEGA
DRAG SWGG
DRAG GOAL GAN
Beach-MarcAkell
Sucoess Coverage
Eru staps leb)
Median
K0M
Mcan
Optimality Gap
DRAG CON--GAN
DRAS Sugg
DRAG Mega
DRAG
0.45
0.50
0.75
0.30
0.80
0.75
0.50
0.6
0.7
0.8
0.5
0.15
0.30
0.45
SucEMIN Coverage
Figure 4: Impact of goal resampling on DRAG. Evolution of the success coverage for different goal sampling methods (6 seeds per run). DRAG directly uses goals sampled from the prior (i.e., same results as figure (3). DRAG + X includes an additional goal resampling method X, taken among the four strategies: LGE, MEGA, GOALGAN or SVGG.
Success Coverage evaluation The success coverage results in Figure reveal an interesting pattern: methods that incorporate diversity-based goal selection on top of DRAG, such as MEGA and LGE, do not lead to any improvement over the original DRAG approach using goals sampled from the prior distribution p(z). This suggests a redundancy between these diversity-basee criteria and our core DRO-based representation learning mechanism, which already inherently fosters exploration. GOID selection criterion degrades DRAG 's performance, likely due to its overly restrictive goal selection strategy, which hinders exploratory behaviors hard goals sampled from the prior must be given a chance to be selected for rollouts in order to support exploration.
Integrating a GOAL GOALGAN-IK
In contrast, the SVGG resampling distribution which leverages the control success predictor in smoother and more adaptive manner significantly outperforms direct sampling from the trained prior. In general, control-based goal selection is ineffective when using classical VAE training in GCRL (as

---

exemplified by RIG), since goals that are not well mastered tend to be poorly represented in the latent space. However, the representation learned with DRAG enables goal selection to focus entirely on control improvement, as it ensures more structured and meaningful latent space. Additional results on these experiments are presented Appendix D33
4.3 Conclusion
In this work, we introduced DRAG, an algorithm leveraging Distributional Robust Optimization, to learn representation from pixel observations in the context of intrinsically motivated Foal-Conditionee agents, in online RL, without requiring any prior knowledge. We showed that by taking advantage of the DRO principle, we are able overcome exploration bottlenecks in environments with discontinu- ous goal spaces, setting us apart from previous methods like RIG and SKEW-FIT.
might As future work, DRAG is agnostic to the choice of representation learning algorithm, consider alternatives such as other reconstruction-bal
techniques [Van Den Oord et al. 2017 7
Razavi 2019 Gregor al. 2019) contrastive learning objectives [Oord al. 2018 Henafi 2020 He et al. 20200] Zbontar et al. 2021]. Besides, DRAG does not leverage pre-trained visual representations. though they could greatly improve performance
complex visual observations [Zhou et al., 2025). In particular, we may incorporate pre-trained representations from models specific RL tasks as VIP [Ma al. [022] and R3M (Nair et al., 2022] well general- purpose visual encoders such as CLIP Radford et al. 202T] or DINO models [Caron et al. [202]] Oquab et al.) [2024).
DRAG also opens promising avenues for discovering more principled and effective goal resampling strategies, made possible by better anticipation of distributional shifts that previously constrained
the potential of the behavioral policy.
Acknowledgements
This work was granted access the HPC resources of IDRIS under the allocation AD01061 and ADO1 101432R22 made by GENCI. We acknowledge funding from the European Commission's Horizon Europe Framework Programme under grant agreement No 101070381 (PILLAR-robots project).
10

---

References
R. Agarwal, M. Schwarzer, Castro, A. Courville, and M. G. Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in Neural Information Processing Systems, 2021.
A. Aubret, Matignon, and S. Hassas. DisTop: Discovering a topological representation learn diverse and rewarding skills. IEEE Transactions on Cognitive and Developmental Systems, 15(4): 1905-1915, 2023.
A. Ben-Tal, D. Den Hertog, A. De Waegenaere, B. Melenberg, and G. Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Science, 59(2)3341-357, 2013.
V. Campos, A. Trott, C. Xiong, R. Socher, X. G. i Nieto, and Torres. Explore, discover and learn: Unsupervised discovery state-covering skills, 2020. URL httpsttss:///xxi org /tt////xiiv..rr////////////0/ abs/2002 03647
M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and A. Joulin. Emerging properties in self-supervised vision transformers, 2021. URL htthh:ttt/ttaiiiii //rpiiv www.:/or///bb//.///// 14294]
N. Castanet, S. Lamprier, and O. Sigaud. Stein variational goal generation for adaptive exploration in multi-goal reinforcement learning. arXiv preprint arXiv:2206.06711 2022.
C. Colas, Fournier, Sigaud, and P.-Y. Y. Oudeyer. CURIOUS: intrinsically motivated multi-task
multi-goal reinforcement learning. In ICML, 2018
C. Colas, P. Fournier, M. Chetouani, O. Sigaud, and P.-Y. Oudeyer. CURIOUS: intrinsically motivated modular multi-goal reinforcement learning. International conference on machine learning. pages 1331-1340. PMLR, 2019.
E. Delage and Y. Ye. Distributionally robust optimization under moment uncertainty with application to data-driven problems. Operations research, 58(3):595-66 612, 2010.
J. C. Duchi, P. W. Glynn, and H. Namkoong. Statistics of robust optimization: A generalized empirical likelihood approach. Mathematics of Operations Research, 46(3):946-969, 2021.
B. Emami, D. Ghosh, Zeng, and S. Levine. Goal-conditioned diffusion policies for robotic manipulation. In Conference on Robot Learning (CoRI.), 2023.
C. Florensa, D. Held, X. Geng, and P. Abbeel. Automatic goal generation for reinforcement learning agents. In International conference on machine learning, pages 1515-1528. PMLR, 2018.
Q Gallouédec and E. Dellandréa. Cell-free latent go-explore. In International Conference Machine Learning, pages 10571-10586. PMLR, 2023.
K. Gregor, G. Papamakarios, Besse, L. Buesing, and Weber. Temporal difference variational auto-encoder, 2019. URL https: :/tpss://///xxi org/ abss/1006 03107
D. Hafner, T. Lillicrap, J. Ba, and M. Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 1xXiv:9922.01033 2019a.
D. Hafner, T. Lillicrap. I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pages 2555-2565, PMLR, 2019b.
K. Hartikainen, X. Geng, T. Haarnoja, and S. Levine. Dynamical distance distance learning for semi-supervised and unsupervised skill discovery, 2020. URL /ttp::ppp//.aaatoooo///////////0 08225
E. Hazan, S. Kakade, K. Singh, and A. Van Soest. Provably efficient maximum entropy exploration.
In International Conference Machine Learning, pages 2681-2691. PMLR, 2019.
K. He, H. Fan, Y. Wu, S. Xie, and R. Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEFF/CVF conference on computer vision and pattern recognition, pages 9729-9738, 2020.

---

O. Henaff. Data-efficient image recognition with contrastive predictive coding. In International conference on machine learning. pages 4182-4192 PMLR, 2020.
I. Higgins, Matthey, A. Pal, C. P. Burgess, X. Glorot, M. M. Botvinick, Mohamed, and A. Lerchner. bet:-vae: Learning basic visual concepts with a constrained variational framework. ICLR (Poster), 3, 2017a.
I. Higgins, A. Pal, A. Rusu, L. Matthey, C. Burgess, A. Pritzel, M. Botvinick, C. Blundell, and A. Lerchner. Darla: Improving zero-shot transfer in reinforcement learning. In International Conference on Machine Learning, pages 1480-1490. PMLR, 2017b.
S. Kim, K. Lee, and J. Choi. Variational curriculum reinforcement learning for unsupervised discovery of skills. arXiv preprint MXNI::2310.1942: 2023.
D. Kingma and M. Welling. Auto-encoding variational bayes. arXiv preprint urXiv:1312.611 2013.
A. Kuznetsov, P. Shvechikov, A. Grishin, and D. Vetrov. Controlling overestimation bias with truncated mixture of continuous distributional quantile critics, 2020. https:///axx abs/2005 04269.
M. Laskin, A. Srinivas, and P. Abbeel. CURL: Contrastive unsupervised representations for rein- forcement learning. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research. pages 5639-5650. PMLR, 13-18 Jul 2020. URL https://pproceedings mlr pat..p/eess///////// 1askin20a.html
A. X. Lee, B. Eysenbach, Parisotto, and Levine. Stochastic latent actor-critic: Deep reinforce- ment learning with a latent variable model. In Advances in Neural Information Processing Systems (NeurIPS), 2020a.
A. X. Lee, A. Nagabandi, P. Abbeel, and S. Levine. Stochastic latent actor-critic: Deep reinforcement learning with latent variable model, 2020b. BolLatttss///rrri.orrr///////
00953
S. Li, L. Zheng, J. Wang, and C. Zhang. Learning subgoal representations with slow dynamics. In International Conference on Learning Representations, 2021.
X. Lu, S. Tiomkin, and P. Abbeel. Predictive coding for boosting deep reinforcement learning with sparse rewards. arXiv preprint /91ivv19/2.1344 2019.
Y. J. Ma, S. Sodhani, D. Jayaraman, O. Bastani, V. Kumar, and A. Zhang. VIP: Towards universal visual reward and representation via value-implicit pre-training. arXiv preprint arXiv:2210.00030, 2022.
R. Mendonca, O. Rybkin, K. Daniilidis, D. Hafner, and D. Pathak. Discovering and achieving goals via world models. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 24379-24391 Cur- ran Associates, Inc., 2021. URL http:///rrre/ www.gis........... cwcppper_files//ppeeeeepeeee 2DL//5111cc445555522263355566666
pdf
R. Mendonca, S. Bahl, and D. Pathak. Structured world models from human videos, 2023. URL https://arxiv.. http:/..gga/bbbb//b/.a 10900
P. Michel, T. Hashimoto, and G. Neubig. Modeling the second player in distributionally robust optimization, 2021.
P. Michel, T. Hashimoto, and G. Neubig. Distributionally robust models with parametric likelihood ratios, 2022.
A. Nair, V. Pong, M. Dalal, S. Bahl, S. Lin, and S. Levine. Visual reinforcement learning with imagined goals. arXiv preprint 1007.v/:000.00 04742, 2018.
Nair, A. Rajeswaran, V. Kumar, C. Finn, and Gupta. R3M: universal visual representation for robot manipulation, 2022. URL //arxiv org/ /2t.. 12601
12

---

A. v. d. Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. arXiv preprint Xrvi 18071007733344. 2018.
M. Oquab, T. Darcet, T. Moutakanni, H. Vo, M. Szafraniec, V. Khalidov, P. Fernandez, D. Haziza, F. Massa, A. El-Nouby, M. Assran, N. Ballas, W. Galuba, R. Howes, P.-Y. Huang, S.-W. Li, I. Misra, M. Rabbat, V. Sharma, G. Synnaeve, H. Xu, H. Jegou, J. Mairal, P. Labatut, A. Joulin, and Bojanowski. Dinov2: Learning robust visual features without supervision, 2024. URL
ftttp:://arxv.orgg abs/ 2304 07193
D. Pathak, P. Agrawal, A. A. Efros, and T. Darrell. Curiosity-driven exploration by self-supervised prediction. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 2778-2787 PMLR, 06-11 Aug 2017. URL
mlr. axress/v70/pathak117 html
N. Perrin-Gillbert. xpag: modular reinforcement learning library with jax agents, 2022. URL htctps///ggthau..cco/prrrri----p
Pitis, H. Chan, Zhao, B. Stadie, and Ba. Maximum entropy gain exploration for long horizon multi-goal reinforcement learning. In H. D. III and A. Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 7750-7761, PMLR, 13-18 Jul 2020. URL v119/pitis2aa. html
mwr.
M. Plappert, M. Andrychowicz, A. Ray, B. McGrew, B. Baker, G. Powell, J. Schneider, J. To- bin, M. Chociej, P. Welinder, et al. Multi-goal reinforcement learning: Challenging robotics environments and request for research. arXiv preprint arXiv: 10X.iv:8002099444 2018.
V. H. Pong, M. Dalal, S. Lin, A. Nair, S. Bahl, and S. Levine. Skew-fit: State-covering self-supervised reinforcement learning. arXiv preprint arXiv:1903.000 19030036988 2019.
A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, Mishkin, J. Clark, G. Krueger, and I. Sutskever. Learning transferable visual models from natural language supervision, 2021. Utttttps::/arrr httppp////aaooooo/aaaa////////// 220./21110000000
H. Rahimian and S. Mehrotra. Distributionally robust optimization: A review. arXiv preprint 2019.
A. Rajeswaran, V. Kumar, A. Gupta, G. Vezzani, J. Schulman, E. Todorov, and S. Levine. Learning complex dexterous manipulation with deep reinforcement learning and demonstrations, 2018. URL ftaap://arrxi.orrggg. 1008
A. Razavi, A. Van den Oord, and O. Vinyals. Generating diverse high-fidelity images with va-vae-2. Advances in neural information processing systems, 32, 2019.
Z. Ren, K. Dong, Y. Zhou, Q. Liu, and Peng. Exploration via hindsight goal generation. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bue, E. Fox, and R. Garnett, ed- itors, Advances in Neurul Information Processing Systems, volume 32. Curran Associates, httppsscccppppppp Inc., 2019. URL https: https://proceddings..oem .pdf 17d17688ss35522255554400ddddddd3
A. Srinivas, M. Laskin, and P. Abbeel. Curl: Contrastive unsupervised representations for reinforce- ment learning. In International Conference on Machine Learning (ICML), 2020.
A. Stooke, K. Lee, P. Abbeel, and M. Laskin. Decoupling representation learning from reinforcement learning. In International conference on machine learning, pages 9870-9879, PMLR, 2021.
S. Sukhbaatar, Z. Lin, I. Kostrikov, G. Synnaeve, A. Szlam, and R. Fergus. Intrinsic motivation and automatic curricula via asymmetric self-play. arXiv preprint arXiv: I 700.05407,7 2017.
Y. Tassa, Y. Doron, A. Muldal, T. Erez, Y. Li, D. de Las Casas, D. Budden, A. Abdolmaleki, Merel, A. Lefrancq, T. Lillicrap, and M. Riedmiller. Deepmind control suite, 2018. URL https://arxil wrw./oo/////////// 00690
13

---

A. Touati and Y. Ollivier. Learning one representation to optimize all rewards. Advances in Neural Information Processing Systems, 34:13-23 2021.
A. Trott, S. Zheng, C. Xiong, and R. Socher. Keeping your distance: Solving sparse reward tasks using self-balancing shaped rewards, 2019.
A. Van Den Oord, O. Vinyals, al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Venna and S. Kaski. Neighborhood preservation nonlinear projection methods: An experimental study. In International conference on artificial neural networks, pages 485-491. Springer, 2001.
D. Warde-Farley, T. Van de Wiele, T. Kulkarni, C. Ionescu, S. Hansen, and V. Mnih. Unsupervised control through non-parametric discriminative rewards. arXiv preprint arXiv:1811.11339. 2018.
D. Yarats, R. Fergus, A. Lazaric, and L. Pinto. Reinforcement learning with prototypical representa- tions. M. Meila and T. Zhang, editors, Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pages 1920-11931. PMLR, 18-24 Jul 2021. URL httph: htttps://proceedingg.o mlr (eress/v139/yarattee
html
T. Yu, D. Quillen, Z. He, R. Julian, A. Narayan, H. Shively, A. Bellathur, K. Hausman, C. Finn, and S. Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning, 2021. /ttt::::: //apiv..r//o/o/////////// abs/ 1910. 10897
J. Zbontar, L. Jing, I. Misra, Y. LeCun, and S. Deny. Barlow twins: Self-supervised learning via redundancy reduction. In International conference on machine learning, pages 12310-1233 12320.. PMLR, 2021.
Y. Zhang, P. Abbeel, and L. Pinto. Automatic curriculum learning through value disagreement. arXiv preprint arXiv:2006. 09641, 2020. doi: 10.8550/ARXX1220000666// URL https://arxiv
org/ org/aba//e abs/2006.660000 2006.09641.
G. Zhou, H. Pan, Y. LeCun, and L. Pinto. Dino-wm: World models pre-trained visual features enable zero-shot planning, 2025. https:/ www.rii.wg//.....eeeeeeee.oooeee org/abg. 04983
14

---

A Experimental details
A.1 Environments
In both environments below, transform 2D states and goals into 82 82 pixel observations.
Pixel continuous point maze: This continuous 2D maze environment is taken from [Trott et al. 2019]. The action space is a continuous vector (őx, σy) = [0, 1Ja. Original states and goals are 2D (x. y) positions in the maze and success is achieved if the L.2 distance between states and goal coordinates is below = 0.15 (only used during the evaluation of success coverage), while the overall size of the mazes is 6 x 6. The episode rollout horizon is T 50 steps. Examples of pixel
observation goals are shown in Figure 5
Figure 5: Example of decoded pixel goals in environment: sample latent goals from the latent prior z ~ p(2) = (0. /) and plot the corresponding decoded pixel goals (o(2))). Images obtained using the decoder trained with DRAG.
Pixel Reach hard walls: This environment is adapted from the Reach-v2 Meta World benchmark
Yu et al., 202T] where the gripper initially stuck between four walls and has navigate carefully between them to reach the goals. The original observations are 49-dimensional vectors containing the gripper position as well as other environment variables, the actions and the goals are 3-dimensional corresponding (z, y, coordinates. Success is achieved if the L2 distance between states and goal
15

---

coordinates is less than 8 = 0.1 (only used during the evaluation of success coverage). The episode rollout horizon T 300 steps.
We transform states and goals into pixels observation using the Mujoco rendering function with the following camera configuration:
REFAULT_CAMERA_CC CONFIG {
"distance": 2-,
"azimuth": 270,
"elevation": -30.0,
"lookat"; ap.array(IO) 0.5, 0)),
}
The walls configuration is obtained with the addition of the following bodies into the "worldbody" of the xml file of the original environment:
<body name="sall_1" pos="O. 15 0.55 .2">
<geon material="Mall_brice" brick" type="box" size=" 000 .24 .2" rgba="O 1 0 1"/>
<geom class="vall_con col" type="box" size=" .055 24 .2" rgba="0 1"/>
</ /odyy
<body name="Mall_2" pos-" 0.15 0.55 .2">
<geom material-"vall brick" tbpe="box" size=". .24 .2" rgba-"0 1 0 1"/>
<geom class="vall_col type="box" size=" 005 24 .2" rgba="O 1"/>
</ /bodyy>
<body name="Mall_3" 0o 0.65 .2>>
<geon material="vall brick" type="box" size=" .005 .2" rgba="O 0 1"/>
<geom class="rall_com type="box" size=" .00 .2" rgba="0 1 0 1"/>
</ <boyy
<body name="wall," pos-"0.0 0.35 .2>
<geom material="vall brick" type=""ox" size=" .005 .2" rgba="O 0 1"/>
<geom class="vall_co type="box" size=". .005 .2" rgba="0 1 0 1"/> </body>
Examples of pixel observation goals are shown in Figure
A.2 3-VAE
A.2.1
Training schedule
During the first 300k steps of the agent, we train the VAE every 5k agent steps for 50 epochs of 10 optimization steps (on dataset of 1000 inputs uniformly sampled from the buffer, divided in 10 minibatches of 100 examples). Afterward, we train it every 10k agent steps.
The following other schedules have been experimented, each getting worse average results for any algorithm:
1. During the first 300k steps of the agent, train the VAE every 10k agent steps. Afterward, train it every 20k steps.
2. During the first 100k steps of the agent, train the VAE every 5k agent steps. Afterward, train it every 10k steps.
3. During the first 100k steps of the agent, train the VAE eveevery 2k agent steps. Afterward, train it every 5k steps.
A.2.2 Encoder smooth update
To enhance the stability the agent's input representations, actions are selected based on smoothly updated version of the VAE encoder, denoted by parameters G. This encoder is refreshed after each VAE training phase and used to produce latent states:
a; Te(. St, Zg); 2 ==(1)))
16

---

Figure 6: Example of decoded pixel goals in fetch environment: we sample latent goals from the latent prior ~ 2) (0, Z) and plot the corresponding decoded pixel goals (z(2))). Images were obtained using the decoder trained with DRAG.
Analogous to the use of a target network for Q-function updates in RL, the delayed encoder 43 is updated using an exponential moving average (EMA) the primary encoder's weights qa:
T + (1 T)
(6)
Figure illustrates how the smoothing coefficient т influences success coverage with DRAG. We observe that 0.05 provides good balance, yielding stable performance. In contrast, setting == (no smoothing) leads to less stability, while 7 0.01 results updates that are too slow. This value was observed to provide the best average results for other approaches (i.e., RIG and SKEW-FIT). We use it in any experiment reported in other sections of this paper.
A.3 Methods papper-prrameerss
The hyper-paranee paper-parameters of our DRAG algorithm are given in Table 2 Notations refer those used in the main paper or the pseudo-code given in Algorithm [ Hyper-parameters that are common to any approach were set to provide best average results for RIG. RIG, SKEW-FIT and DRAG share the same values for these hyper-parameter reported in the tables below, is set to a = -1 This value was tuned following grid search for
The skewing temperature for SKEW-FIT, which is not
[-100,- -50, -10,- -5, -1, -05,, -0.1].
17

---

Maze
Mazo
Maze
Mazo 3
Reach Harc Walls
Coverage
Eny stens leb1
Median
Mean
Optimality Cep
T- 0.01
7- 0.05
-0.1
m=1
0.72
0.80
0.88
0.72
0.78
0.84
0.04
0.72
0.80
0.16
0.24
0.32
0.40
Coverage
Figure 7: Impact of the exponential moving average smoothing coefficient (T in Equation experiment success coverage (6 seeds per run).
Table 1: Hyper-pane parameters used for the VAE used in the experiments (same for every approach, top) and values used specifically in DRAG, for the specification our DRO weighter (bottom).
VAE Hyper- Parameters
Symbol
Value
3-VAEE
Latent dim
[aze env :2, Fetch: 3]
Prior distribution
p(z)
\O Ia) 2
Regularization factor
3
Learning rate
le-3
CNN channels
(3. 32, 64, 128, 256) (512,128)
Dense layers
Activation Function
ReLu
CNN Kernel size
Training batch size
100
Optimization interval (in agent steps) Nb of training steps
freqOpt
10e3
nbEpochs
50
Size of training buffer
R
le6
Number of samples per epoch
N
le3
DRO Weighter Tip
Learning rate
2e-6
Convolutional layers channels
(3, 32, 64, 128, 256) (512,128)
Dense layers
Activation Function Temperature
ReLu
À
0,01
18

---

Table 2: Off-policy RL algorithm TQC parameters
TQC
Value
Batch size for replay buffer
2000
Discount factor 7
0.98
Action L2 regularization
0.1
(Gaussian) Action noise max std
0.1
Warm up steps before training
2500
Actor learning rate
1e-3
Critic learning rate
1e-3
Target network soft update rate
0.05
Actor & critic networks activation
ReLu
Actor & critic hidden layers sizes
5123
Replay buffer size (l5)
le6
Table 3: goal criterion hyper-parameters
Goal criterion Hyper-parametery
Symbol
Value
Kernel density Estimation for MEGA and LGE
RBF kernel bandwidth
0.1
KDE optimization interval (in agent steps)
Nb of state samples for KDE optim.
10.000
Nb of sampled candidate goals from p(z)
100
Agent's skill model Da for SVGG and GOALGAN Hidden layers sizes Gradient steps per optimization Learning rate Training batch size Training history length (episodes) Optimization interval (in agent steps) Nb training steps
(64, 64)
100
le-3
100
500
5000
100
Activations
Relu
A.4 Compute ressources & code assets
This work was performed with 35,000 GPU hours on NVIDIA V100 GPUs (including main experi- ments and ablations).
Algorithms were implemented using the GCRL library XPAG [Perrin-Gilber 2022 designed for intrinsically motivated RL agents.
B DRAG algorithm
Algorithm reports the full pseudo-code of our DRAG approach. RIG and SKEW-FIT follow the same procedure, without the DRO weighter update loop (line to 12), and replacing L0,0r ([x:)] in
DR
line 15 by:
(classical ELBO):
CYAE ((
71
log (gpol(212)) (2(])()(
N
m
19

---

SKEW-FIT:
12
m og pa (x: K -
CYAE-S 8.0
((1))=))
Pakemed(2i)
N
where Pakewed (2:) ) is the skewed distribution of SKEW-FIT, that uses an estimate Le (2x²) of the generative posterior of x; from the current VAE, obtained from M codes sampled from go (2(2))) More details about Pskewed are given in section C.2
For comparison, recall, for DRAG we take:
DRAG:
((ri)
1 N
(m;)
Σ
m
log po (x))= -------K KT(Ga((laa)lll(z))
0,000
explole))
where the weighting function Ti is defined, following equation 2 (w(x,) =
2 exp fo(z77 for
any minibatch {xi)
Algorithm Distributionally Robust Exploration
1: Input: GCP wa, VAE: encoder (0(2)x)) and smoothly updated version (2(rr)) decoder Pe (ar(),, latent prior p(z), DRO Neural Weighter re, buffers of transitions B, reached states R. train size N, batch-size number m of sampled noises for each VAE training input, number M of Monte Carlo samples used to estimate Lo.s for each input image, temperature X, number of optimization epochs n Eppchss frequence of VAE and policy optimization freqOpt.
2: while not stop do
3: D Data Collection (during freqOptim steps): Perform rollouts of To 2g) in the latent space, conditioned on goals sampled from prior F9 p(z) the buffer (with possible resampling depending on the goal selection strategy), and latent state 2; = 93 (zz), with x pixel observation;
4:
Store transitions in B, visited states in R;
5:
6: > Learning Representations with VAE
7:
for n Epochs epochs do
8:
Sample train of N states I from R for every mini-batch {o:}=_ from I do
9:
DRO Weighter Update
11:
Update weighter by one step of Adam optimizer, for the maximization problem from with temperature X, using T.o.s estimated from samples from go for each Ty.
12:
end for
for every mini-batch {= from I' do
14:
D Weighted VAE Update
Update encoder go and decoder ps by one step of Adam optimizer on 20,0.r ((x:)
as defined in (5). with m sampling noises (c/) for each z;.
16:
Perform smooth update of & as a function of & according to equation (6)
17:
end for end for
18:
19:
20:
D Agent Improvement
21:
Improve agent with any Off-Policy RL algorithm (e.g.. TOC, DDPG, SAC...) using transitions from B;
end while
C Skew-Fit is a non-parametric DRO
In this section show that SKEW-FIT is special of the non-parametric version of DRO.
20

---

C.1 Von-parametric solution of DRO
We start from the inner maximization problem stated in ( for a given fixed 8:
N
max N
Σ ((o.))))((
r Uo) (
N N
st
r(z, 1.
From this formulation, we can see that the risk associated to shift of test distribution can be mitigated by simply associating adversarial weights r; r(z,, y4) to every sample (æ y:) from the training dataset, respecting F := N Ex=rr = 1. This can be viewed as an infinite capacity function r, able to over-specify every training data point. Equivalently to (7). thus have:
N
N
max N
r44
r; log Ti
(8)
N
1=1
N
88
Σ
1,
N
where 4: :== ((f((1;))))))))))))))))()))))) v:)). The Lagrangian corresponding to this constrained maximization given by:
1 N
r4
Σ
ra log ri 7( r 1--1)
Σ
1=1
where y is an unconstrained I Lagrangian coefficient.
Following the Marush-Kuhn Kuhn-Tucker conditions applied to the derivative of the Lagrangian function L of this problem in r; for any given i [[I, VV]. we obtain:
ac = 0 + 4 - X(log r. - 1) - - - 2e
(10)
Or,
with == -1
The KKT condition on the derivative in gives: results, we thus obtain:
ac = 0 +
1 N
= 1. Combining these two
8y
N
N
r
=
N Σ
ex
Which again gives, reinjecting this result in Equation
10
T: N
j= e
This leads to the form of Boltzmann distribution, which proves the result.
C.2 Application GCRL with VAE and Relation Skew-Fit
SKEW-FIT resamples training data points from a batch (w,),_ using a skewed distribution defined, for any sample in that batch, as:
Pekewed (x)
Wt.ar (x),
(11)
Za
Zo
Σ
Wt.o (xi)
i=1
21

---

where Wto is an importance sampling coefficient given as: Wt, (x) = (x. a < 0,
(12)
with (2) the generative distribution of samples given current parameters (0, Φ).
Applied to generative model defined as VAE, we have:
p(z)po(z(z) dz,
po.o(1) (=lg
2 (z(x)))
where p(z) is the prior over latent encodings of the data 2, po(x|2) is the likelihood x know- ing and (²(2)z) the encoding distribution of data points. As stated in Section 3.1 this can be estimated on set of m samples for each data point using the log-approximator: Co
log Σ exp(log pa ) + log pa (z²) - log ( [z)) log(M).
Thus, this is equivalent as associating any from the data batch with a weight r; defined as:
calo,o(a.)
Pskewed (xi)
(13)
Z
Za
Σ cL)))(z))
for any a O. Setting a we get Pskewed (x1) x (=))) for any temperature > Reusing the result from Section C. this is fully equivalent to the analytical closed-form solution of DRO when applied to - -Le (1r)) for weighting training points of a VAE thus exactly corresponds to the non coramaraettiic version our DRAG algorithm.
use in our DRO-VAF approach. Using Pekewed with =-
D Aditional results
D.1 Visualization of Learned Latent Representations
Latent Representations (Learned with (-VAE)
Pixel Observations
Decoded prior M/0.1 sampling
Figure 8: Learned Representation DRAG after million training steps in Maze Left: every colored dot corresponds to the pixel observation of its specific coordinates. Middle: Every pixel observation æ on the left is processed by the VAE encoder to get the learned latent posterior distribution (2(2)) (2))))(())
Colored ellipsoids correspond to these 2-dimensional Gaussian distributions. Right: we sample latent goals from the latent prior ~ p(2) A A I) and decode the corresponding pixel observations Po (2) (2) (red dots correspond to the xy coordinates of the pixel observations).
Figure X presents our methodology to study latent representation learning. We uniformly sample data points in the maze and process them iteratively from 2D points to pixels, then from pixels to the latent code of the VAE, Using the same color for the source data points and the latent code, this process allows us visualize the 2D latent representation of the VAE in the environment (Figure g left and
22

---

middle). In addition, to get a sense of what part of the environment is encoded in the latent prior, we sample latent codes from p(2) and plot the 2D coordinates of the decoded observations using po(2)) which corresponds to the red dots. The blue distribution corresponds to KDE estimate fitted to the red dots.
Step: 10K
50K
100K
500K
1M
4M
RIG
Success: 0.0 %
22.4 %
30.5 %
43.0 %
47.4 %
46.5 %
222
09
61
Skew-Fi Fit
80
DRAG (ours)
Success: 0.0 %
36.5 %
44.9 %
54.0 %
70.9
86.0 %
Figure 9: First row of each method: evolution of learned representations. Second row of each method: evolution of the intrinsic goal distribution when sampling from the latent prior p(2) NO, 1). Third row of each method: evolution of the success coverage. (See Figure 8 for details on how we obtain these plots).
23

---

Step
Step
Step
0.5
0.5
0.5
0.4
0.4
RIG
0.3
0.3
0.5
0.8
0.9
0.7
0.0
0.6
0.0
0.6
0.6
x
Skew-Fit
0.2
a1
21
0.1
05
-0A -0
0.7
0.7
0.
02
0.5
02
0.4
0.4
0.4
0.
0.4
0.
0.5
0.5
0.5
DRAG (ours)
0.4
0.4
0.3
0.3
0.3
0.2
0.2
011
A1
1.0
1.0
1.0
0.8
0.8
-04
-2
0.0
0.5
-04
-02
0.5
0.6
-04 .02
0.5
0.4
0.4
0.4
0.0
Figure 10: Evolution of the prior distribution in the Fetch environment for the DRAG, SKEW-FIT, and RIG methods: We sample latent goals from the latent prior 2 ~ p(2) = NO, 1) and we decode the corresponding pixel observations po (z]2z) (black dots correspond the 3D xyz coordinates of the pixel observations).
In order to gain deeper insight into the performance of RIG, SKEW-FIT, and DRAG, we show in Figure the parallel evolution of the prior sampling NO, I), and the corresponding learned 2D representations for the maze environment.
One clearly that RIG is stuck in exploration bottleneck (which in this case corresponds to the first U-turn of the maze): the VAE cannot learn meaningful representations poorly explored areas (red part of the maze in Figure 9. As consequence, the prior distribution p(z) only encodes small subspace of the environment. On the other hand, SKEW-FIT and DRAG manage to escape these bottlenecks and incorporate organized representation of nearly every area of the environment, with the difference that DRAG is more stable and therefore reliably learns well organized representations.
In order to quantify the evolution of latent representations to highlight the differences in terms of latent distribution dynamics between RIG, SKEW-FIT, and DRAG, we introduce the following measurement:
- 1...T d((x)
Hor (x:) 1 ll.
(14)
where x = [ is batch of pixel observations uniformly sampled from the environment state space using prior knowledge (only for evaluation purposes). With this metric, we measure the evolution of the embedding of every point z,, using the movement of the expectation pz(x;) from the
24

---

latent posterior distribution (z(2)) N (2))) throughout updates of VAE parameters
Ø.
DRAG (ours)
Skew-fit
RIG
Embedding evalutton
Maze
Maze
Maze
Maze
Figure 11: Evolution of the embedding over 4 different PointMazes (6 seeds each) for 4M steps (shaded areas correspond to standard deviation). Every point corresponds to the shift of representation between step and step + of VAE training: 5 Ma (2:) (xz) - Po (Tx) ||. For every pixel observation 24 and timestep representation shifts between and t + every 40, 000 training steps.
have qst xz) N (zlm (zi), (2:)). We compute
Figure 11 shows that the embedding movement d of DRAG is higher and less variable across seeds, indicating that the learned representations evolve more consistently. Meanwhile, the VAE training process of SKEW-FIT is prone variability, and the evolution of the embedding in RIG is close null after a certain number of training steps.
D.2 Representation quality
We assess the quality of the obtained representations in term of the trustworthiness score [Venna and Kaski [200]], which measures to what extent the local neighborhood structure in the input space preserved in the latent space. Higher trustworthiness indicates better preservation of tale-relevant information across the encoding. We computed this trustworthiness metric (with nearest neighbors) using a batch of 1K uniformly sampled observations from the valid state space, identical to the batch used for success coverage. Our experiments show that DRAG consistently improves this score over baseline encoders across environments.
Table 4: Mean trustworthiness of obtained representations across environments after 4M training steps.
Methods
Trustworthiness [Venna Kaski, 2001]
DRAG
98.7%
SKEWFIT
93.3%
RIG
88.5%
D.3 Ablations
We study the impact of the main hyper-parameters, namely À and M.
D.3.1 Impact of X
illustrates the effect varying the regularization parameter À on the performance DRAG. As discussed in the main paper, lower values of À bias the training distribution to emphasize samples from less covered regions of the state space. Conversely, higher values of X lead to flatter weighting distributions across batches, eventually resembling the behavior of standard VAE (as used in RIG) when À becomes very large. In fact, setting À = oo makes DRAG behave identically to RIG.
Figure 12
The reported results show that DRAG achieves the highest success coverage for À values between 10 and 100, with a slight edge at À 100. This range represents a good trade-off: too small À can lead to unstable training, where the model places excessive weight on underrepresented samples; too large leads to overly strong regularization toward the marginal distribution p(x), limiting generalization, and hence exploration.
25

---

le
A- 16-
1e
1e-
A-1
(RIG)
Success Coveroge
Ery steas les)
Median
1QM
Mean
Optimaliyy Gap
(MGG)
tend
len1
0.0
0.75
0.45
Coverage
Figure 12: Impact of the regularization parameter À on the performances of DRAG (6 seeds per run). Results obtained with goals directly selected from the prior (as in Section 4.1 Note that À 0o comes down to the RIG approach, as weights converge to a constant (over-regulariiatioo))
For comparison, SKEW-FIT performs best at a -1, which corresponds to À = 1 in the DRO muon-parametrio formulation (see Section C for theoretical equivalences). This much lower value reflects a key difference: SKEW-FIT relies on pointwise estimations of the generative posterior, while DRAG uses smoothed estimates provided by a neural weighting function. As result, SKEW-FIT requires less aggressive skewing to avoid instability.
D.3.2 Impact of M (number of samples for the estimation of Le.s)
Drag
Drag
Mare
Maze
Mazz
Reach Hand Walls
Coverace
Ern stack (x leb)
Modian
Mean
Optimality Gap
Skow Rit RMM 10
Drag 10
Skow FiL =1
Drag =1
Success
Figure 13: Impact of the number of samples M, used for the estimation of Le.s in DRAG and SKEW-FIT (6 seeds per run). Results obtained with goals directly selected from the prior (as in section 4.1
Both SKEW-FIT and DRAG need to estimate the generative posterior of inputs in order to build their VAE weighting schemes. This estimator, denoted in the paper as Lo (x) for any input æ, is obtained via Monte Carlo samples of codes from ( (r). The number M of samples used impacts the variance of this estimator. The higher M is, the more accurate the estimator is, at the cost of an increase of computational resources (M samples means M likelihood computations through the decoder). This section inspects the impact of M on the overall performance. Figure 13 presents the results for SKEW-FIT and DRAG using M (as in the rest of the paper) and M 10. While one might expect more accurate estimates of Lg, (Ir) with M = 10, this improvement does not translate into better success coverage for the agent. According to the results, the value of M does not appear to significantly impact the agent's performance for either algorithm. In fact, on average, increasing M even slightly decreases success coverage.
26

---

This is noteworthy finding, as it suggests that the improved stability of DRAG compared to SKEW- FIT is not due more accurate pointwise likelihood estimation (which could benefit DRAG through the inertia introduced by using parametric predictor), but rather due to greater spatial smoothness. This smoothness arises from the T.-I.ischhtzz continuity of the neural network: inputs located in the same region of the visual space are assigned similar weights by DRAG's neural weighting function. In contrast, SKEW-FIT may overemphasize specific inputs, with abrupt weighting shifts, particularly when those inputs are poorly represented in the latent space, despite being located in familiar visual regions.
D.4 RIG+Goal selection criterion
RG
RIG 1 LGE
RIG 1 MEGA
NIG SVGG
RG GOAL GAN
Success Coverage
Matelippy Ahw
Enn steps (x leb)
Median
10M
Mean
Optimality Gap
Goal-GAN
RIG
RIG Mega
RIG LGE
INC
Success Coverage
Figure 14: Impact of goal resampling with classical (unbiased) VAE training, as in RIG. Evolution of the success coverage for different goal sampling methods (6 seeds per run). RIG directly uses goals sampled from the prior (i.e., same results as RIG in figure RIG resampling method X, taken among the four strategies: LGE, MEGA, GOALGAN or SVGG. This figure presents the same experiment as in Figure 4 but using a standard VAE instead of our proposed DRO-VAE.
includes an additional goal
Figure 14 presents the results of combining the RIG representation learning strategy (i.e. without biasing VAE training) with goal selection criterion, following the same experimental setup as in Figure 2b These results highlight two key insights.
First, the results show that the exploration limitations inherent to the RIG strategy cannot be effectively addressed by improved goal selection alone. Even the best-performing combination (RIG LGE) achieves less than 60% success coverage on average across environments, while DRAG alone reaches 80%. This highlights the critical role of DRAG's 's representation learning capability in overcoming exploration bottlenecks in complex environments. When relying solely on the latent space of a standard VAE, sampling even when guided by intrinsic motivation- remains limited to regions already well-represented in the training data. The model cannot generate goals in poorly explored areas, as no latent codes exist that decode to such states.
Second, we observe a reversal in the relative effectiveness of goal selection criteria compared to the DR experiments in Figure 2b In the case of RIG, both LGE and MEGA outperform SVGG, which contrasts with the pattern observed with AG. This can be explained by the fact that RIG inherently limited in its ability explore, and thus benefits more from goal selection strategies that explicitly promote exploration. In contrast, strategies based on intermediate difficulty, such as SVGG, are less effective when the agent is confined codes associated with intermediate difficulty typically decode to well-known states, while those corresponding to poorly explored areas often lead to posterior distributions with higher variance. As result, the latter are more likely to be classified as too difficult and filtered out. Therefore, these strategies tend to foster learning around familiar areas, without actively pushing the agent toward under-explored or novel regions that are critical for improving coverage. This type of goal selection can therefore only be effective when built on top of representations such as the one learned by DRAG- that explicitly encouraged to include marginal or rarely visited states.
limited region of the environment. Latent
27

---

D.5 Image size study
To complement our experiments, we ran the three methods (i.e. RIG, SKEW-FIr FIT and DRAG) on all environments with higher image resolution (128x128 instead of 82x82 in the main experiments). We observe that all methods degrade slightly with higher resolution, but DRAG retains a clear advantage.
Table 5: Mean success coverage across maze and robotic environment for 128x 128 pixels observations (success coverage for 82x82 given for comparison).
Methods
1M
2M
3M
4M steps
DRAG
42%
58%
71 %
79%
82x82 comparison
46%
66%
74%
81%
SKEWFIT
28%
44%
57%
63%
82x82 comparison
33%
51%
59%
66%
RIG
22%
29%
35%
38%
82x82 comparison
28%
34%
39%
42%
D.6 Runtime
To assess the impact of the overhead induced by the additional likelihood estimation (for SKEWFIT and DRAG compared to RIG) and the use neural weighter (for DRAG), (bble/) mean execution runtime for Maze and Metaworld environments averaged over 6 seeds each environment. We observe that DRAG and SKEWFIT require on average about 200 seconds for 20k steps, against 176 seconds for RIG. This slight overhead is negligible compared the time spent in environment interactions. This is because we only use one sample (i.e. for the likelihood estimation in DRAG and SKEWFIT (see appendix D.3.2 for comparative study with more samples, which does not impact the results of both methods). Also, no notable difference in runtime is observed when comparing DRAG SKEWFIT.
Table 6: Mean execution runtime on V100 GPU (32GB) for MAZE and METAWORLD environments.
Method
Time per 20K steps (s)
Time for 4M steps (s)
DRAG
211
42,200 (11.7 h)
SKEWFIT
208
41,600 (11.55 h)
RIG
176
35,200 (9.77 h)
To better link performance to runtime, Table 7 also shows success coverage over wall-clock time (82x82, same runs Fig. 3):
Table 7: Success coverage metric hour.
Hour
DRAG (%)
SKEWFIT (%)
RIG (%)
34
31
26
45
42
28
5
60
50
32
70
58
36
9
77
63
40
81
66
42
E Limitations
The main limitations of our work are the following.
28

---

Latent space reward definition While our study makes progress on learning representations online and generating intrinsic goals from high-dimensional observations, it does not address how to measure when a goal has truly been achieved. Throughout our experiments, we relied on simple sparse reward rt [||ar; Zg lg < 5) which, although common goal-conditioned RI,, sidesteps the challenges of defining dense feedback signal. In particular, the Euclidean metric used in dense rewards often fails to reflect the true topology of the environment and can mislead the agent.
Representation Learning algorithm Our DRO-based approach is agnostic to the choice of repre- sentation learning algorithm, suggesting future work should benchmark alternatives such as other
reconstruction-basee techniques [Van Den Oord al. 2017 Razavi al. 2019 Gregor et al. 2019]]
or contrastive learning objectives [Oord et al., 2018
Henafi 2020 He et al. 2020 Zbontar et al.
2021
Notably, contrastive methods (Stooke al. 2021 and Borward-Backwerr approaches Touati and
Ollivier, 2021] aim to incorporate dynamics by bringing temporally adjacent states closer or by modeling universal rewards. However, these methods generally assume access transitions from representative part of the environment. To our knowledge, they do not include any explicit mechanism to avoid collapse onto narrow parts of the state space- a critical issue in hard exploration settings without expert priors.
Addressing this limitation is precisely the aim of our DRO-based reweighting, which promotes state space coverage even in sparse reward regimes. While our implementation focuses on 3-VAEs for interpretability and disentanglement, our DRO framework general and not tied to specific representation architecture.
In fact, any representation learning method trained from replay buffer samples using loss ((fo(x)) can be reweighted using the DRO objective:
N
1 N
min max N
((ss(((fe(x;))) - (r(2;)))g((z.))
s.t.
r(æ))=
=1
1
1=1
This includes diffusion-based decoders, contrastive losses (e.g.. InfoNCE), and temporal-difffrence- based objectives (e.g.. in Forward-Backward RL). For generative losses, in VAEs diffusion models, the optimization must rely on likelihood estimates (i.e., log log().)(()) ), similarly to our alternate optimization strategy described in (4)
Leveraging Pre-trained Representations Our study did not leverage pre-trained visual representa- tions, that could greatly improve performance on complex visual observations as demonstrated in [Zhou et al. 2025). In particular, future work should explore incorporating into our setting pre-trained
representations from models specific to RL tasks VIP [Ma et al. 2022 and R3M [Nair al. 2022
as well as paneral-purpose visual encoders such as CLIP [Rafford et al. 2021 or DINO models
[Caron et al. 202 Oquab et al. 2024
29

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our paper's contribution clearly reflects the claim made in the abstract: we design new GCRL method to learn representations online foster exploration and agent performance.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We briefly address the key limitation of our work the conclusion and develop them in more details in the appendix.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
30

---

Answer: [Yes]
Justification: All theoretical results are either justified in the main paper or a complete proof is provided in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: All experimental details needed to reproduce the results are provided in the appendix. The codebase will be made publicly available upon acceptance of the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
31

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The github repository of the code for experiments reproducibility is provided in the main paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr cc/
Sublic/guidess Sosgiicgguides/Coeuuutiii Cooeeouaiissioooooiiiiiiiioeeiie
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to reproduce the results. See the NeurlPS code and data submission guidelines (https
cttpss.ccppubblci gwwdeg/ideeeee www.oo/wwwwswooooooooooooooo
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All experimental details needed to reproduce the results are either included in the main paper provided in the appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
The number of runs with different seeds for each run are provided in the main paper. Furthermore, to assess the statistical significance of results, we compute the median, interquartile mean, mean and optimality gap metrics and confidence intervals using the Rliable library [Agarwal al. 2021), specialized statistical analysis of Deep RL methods.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
32

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτor rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute ressources used to conduct our experiments are listed in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics httph: atttss//neuurps
Answer: [Yes]
Justification:
All authors are familiar and have respected the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Due the theoretical nature of this work, we believe, the best of our knowledge, that no societal impact is at stake. Furthermore, RI. does not work with data that could include bias with societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
33

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work does not make use of pretrained language models, image generators, nor scraped datasets.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging. and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All existing assets used this work are listed in the appendix.
Guidelines:
The answer NA means that the paper does
existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
34

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: the paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
35

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM for what should or should not be described.
36


Reviews:
[{"review_id": "mQ3drA8eNx", "reviewer": "Reviewer_13ri", "summary": "The paper introduces DRAG, an online goal-conditioned RL method that couples a beta-VAE with distributionally robust optimisation (DRO). A small neural weighter assigns adversarial importance weights to replay-buffer images so that VAE training emphasizes under-represented states, effectively steering the latent space toward a uniform coverage of the environment. The authors derive a DRO-VAE objective and show that SKEW-FIT is its non-parametric special case, propose the parametric DRAG variant for greater stability and integrate DRAG into a RIG-like GCRL loop with latent-goal sampling. They empirically demonstrate markedly higher success-coverage and improved downstream control on four pixel PointMazes and a Reach-Hard-Walls robot task, compared with RIG and SKEW-FIT.", "questions": "- How sensitive is DRAG to the temperature of the KL relaxation and to the beta of the VAE? A small study varying these would clarify robustness.\n- What is the wall-clock overhead of training the neural weighter versus plain RIG? Please report GPU hours or frames per second.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Quality:\n- Strengths: sound formulation with DRO-VAE derivation following standard DRO theory, experiments run six seeds and report confidence intervals with RLiable statistics.\n- Weaknesses: limited scope with small-scale visual tasks (82 × 82 images) and a single robotic manipulation variant, no runtime or sample-efficiency analysis is provided, computational overhead of the weighter network is not quantified.\n\nClarity:\n- Strengths: well organised paper, the relation between SKEW-FIT and DRAG is explicitly proved, figures 3-4 summarize results effectively.\n- Weakness: a clear presentation of the exact ELBO and hyperparamters used would help reproducibility\n\nSignificance:\n- Strengths: tackling exploration bottlenecks in online visual GCRL without separate explorers is an important practical problem, the idea of casting state-coverage as DRO over VAE training data is novel in this context and could inspire future work in representation learning for RL.\n- Weakness: impact is moderated by the modest experimental breadth and by missing comparisons with other contemporary representation learners\n\nOriginality:\n- Strengths: first application of parametric DRO to auto-encoder training in RL, the analytical link to SKEW-FIT clarifies prior work.\n- Weakness: the individual ingredients (beta-VAE, DRO and RIG) are established techniques and novelty lies mainly in their combination.", "was_revised": true, "final_justification": "After rebuttal, my main concerns are addressed. The methods remains technically sound but the evaluation is still limited in scope. I keep 4.", "rebuttal": "Many thanks for your thorough and constructive feedback. We address the main concerns below.\n\n**Weakness: Quality (and Question about runtime overhead)**\n\n- ***Limited scope (image resolution, task diversity):*** \n\nThanks for the suggestion for considering higher resolution images to assess the ability of the method to deal with larger inputs. To complement our experiments, we ran the three methods (i.e. RIG, SKEWFIT and DRAG) on all environments with a higher image resolution (128x128 instead of 82x82 in the paper). We obtained the following results, averaged on every environment from the paper: \n\n| Methods | 1M  | 2M | 3M |4M steps\n|----------|----------|----------|-------|----------|\n| **DRAG**    | **42%**   | **58%** | **71%**  | **79%**\n| **↳ 82x82 comparison**    | **46%**   | **66%** | **74%**  | **81%**\n| SKEWFIT    | 28%   | 44%  |57% | 65%\n| ↳ 82x82 comparison    | 33%   | 51%  |59% | 66%\n| RIG    | 22%   | 29%   |35%| 38%\n|  ↳ 82x82 comparison    | 28%   | 34%   |39%| 42%\n\n**Table 1: Mean success coverage across maze and robotic environment for 128x128 pixels observations (success coverage for 82x82 is given for comparison)**\n\n\nAll methods degrade slightly with higher resolution, but DRAG retains a clear advantage.\n\nWe agree that broader task diversity would further validate generality. While our study focuses on goal-conditioned robotic tasks, the DRO-based representation principle is generic and could be extended to navigation or multi-agent setups, which we leave for future work.\n\n\n- ***Sample-efficiency and runtime analysis***\n\nFirst, we would like to respectfully highlight that Figures 3 and 4 already inform about sample efficiency as the x-axis shows the number of steps (or environment interactions). Our approach appears significantly more sample efficient than baselines. \n\nHowever, we agree with the reviewer that we should also include runtime measurement, to assess the impact of the overhead induced by the additional likelihood estimation (for SKEWFIT and DRAG compared to RIG) and the use of a neural weighter (for DRAG). With additional experiments reported in Table 2 below (values are averaged over 6 seeds on each environment), we observe  that DRAG and SKEWFIT require on average about 200 seconds for 20k steps, against 176 seconds for RIG. This slight overhead is negligible compared to the time spent in environment interactions. This is because we only use one sample (i.e. $M=1$) for the likelihood estimation in DRAG and SKEWFIT (please refer to the experiments with $M=10$ in Appendix D.2.2, which shows that using more samples does not impact the results of both methods). Also, no notable difference in runtime is observed when comparing DRAG to SKEWFIT. \n\n**Table 2 : Mean execution runtime on V100 GPU (32GB) for MAZE and METAWORLD environments**\n| Method          | Time per 20K steps (s)  | Time for 4M steps (s) |\n|-----------------|------------------------|------------------------|\n| DRAG  | 211                    | 42,200  (11.72 h)                   |\n| SKEWFIT  | 208                    | 41,600  (11.55 h)                   |\n| RIG             | 176                    | 35,200  (9.77 h)                   |\n\n\n\nTo better link performance to runtime, Table 3 shows success coverage over wall-clock time (82×82, same runs as Fig. 3):\n\n**Table 3 : Success coverage / runtime for MAZE and METAWORLD environments**\n| Hour | **DRAG** (%) | SKEWFIT (%) | RIG (%) |\n|------|----------|--------------|---------|\n| 1    | **34**       | 31           | 26      |\n| 3    | **45**       | 42           | 28      |\n| 5    | **60**       | 50           | 32      |\n| 7    | **70**       | 58           | 36      |\n| 9    | **77**       | 63           | 40      |\n| 11   | **81**       | 66           | 42      |\n\n\nWe will add this additional analysis in Appendix D.\n\n\n**Weakness: Clarity on ELBO and hyperparameters:** \"a clear presentation of the exact ELBO and hyperparamters used would help reproducibility\"\n\nWe appreciate the request for more clarity, but we are not sure to fully understand what the reviewer means when asking for the exact ELBO. The ELBO used for VAE training is given in Eq. (5), which is the standard form for a VAE, but weighted by the learned reweighting function $r(x)$. \n\n- ***Method Derivation***\n\nIf the reviewer wants more explanation about the derivation of the method, we can complement the paper starting from Eq. (3), that we relax with a KL term as in (1), to get:  \n$\\min_{\\theta, \\phi} \\max_{r:E_{p}r = 1} - E_{x \\sim p} \\ r(x) \\log{p_{\\theta,\\phi}(x)} - \\lambda E_{x \\sim p} \\ r(x) \\log r(x)$\n\n$=\\max_{\\theta, \\phi} \\min_{r:E_{p}r = 1}  E_{x \\sim p} \\ r(x) \\log{p_{\\theta,\\phi}(x)} + \\lambda E_{x \\sim p} \\ r(x) \\log r(x)$\n\n$=\\max_{\\theta, \\phi} \\min_{r:E_{p}r = 1}  L_{\\theta,\\phi,r}$\n\nwhere $L_{\\theta,\\phi,r}=E_{x \\sim p} \\ r(x) \\log{p_{\\theta,\\phi}(x)} + \\lambda E_{x \\sim p} \\ r(x) \\log r(x)$ \n\nThen, we could write a lower bound as: \n$L_{\\theta,\\phi,r} \\geq E_{x \\sim p(x)} r(x) E_{z \\sim q_\\phi(z|x)} \\log \\frac{p(z) p_\\theta(x|z)}{q_\\phi(z|x)} +  \\lambda E_{x \\sim p} \\ r(x) \\log r(x)$\n\nHowever, using this quantity to optimize the weighter $r$ would not be valid, as minimizing a lower bound makes no sense. This led us to an alternate learning scheme that optimizes $r$ on an approximation of the true likelihood of any sample $x$ (which we construct with M samples of code $z$ for each $x$), as defined in Eq. (4), while the weighted VAE is optimized on the weighted ELBO given in Eq. (5) for a fixed weighter. \n\n- ***Reproducibility (Algorithm and Hyper-parameters)***\n\nMaybe the reviewer is more interested by the full algorithm that we employ in our experiments (as the global weakness refers to reproducibility). Please refer to **Appendix B for the full pseudo-code of our method**, and to **Appendix A.3 for all hyper-parameters** used in our method and the baselines considered in our experiments. \n\n\n**Weakness: Originality**\n\nWhile we build upon existing components (β-VAE, DRO, RIG), our key contribution is the novel application of DRO to VAE training in online RL, which—unlike classification—presents unique challenges due to dynamic data distributions and non-stationarity (induced by the progress of the agent). This connection is non-trivial, and to the best of our knowledge, has not been explored before.\n\nOur theoretical analysis also highlights how SKEWFIT can be seen as a particular case of DRO under specific choices of $r$, allowing us to generalize and improve it using the DRO framework. The use of DRO in the setting of a VAE is also a contribution, for which we proposed a specific learning scheme, alternating the optimization of the weighter on an estimate of the true likelihood a depicted in Eq. (4) (since minimizing a lower bound for $r$ would make no sense) and optimizing the encoder-decoder on a weighted ELBO (Eq. (5)). \n\n**Weakness Significance: Comparison with other representation learning paradigms**\n\nWe agree that comparing against other contemporary representation learning approaches could enrich the analysis. Notably, contrastive methods [Stooke et al., 2021] and Forward-Backward approaches [Touati and Ollivier, 2021] aim to incorporate dynamics by bringing temporally adjacent states closer or by modeling universal rewards. However, these methods generally assume access to transitions from a representative part of the environment. To our knowledge, they do not include any explicit mechanism to avoid collapse onto narrow parts of the state space—a critical issue in hard exploration settings without expert priors.\n\nAddressing this limitation is precisely the aim of our DRO-based reweighting, which promotes state space coverage even in sparse reward regimes. While our implementation focuses on β-VAEs for interpretability and disentanglement, our DRO framework is general and not tied to a specific representation architecture.\n\nIn fact, any representation learning method trained from replay buffer samples using a loss $\\ell(f_\\theta(x))$ can be reweighted using the DRO objective:\n\n$\\min_{\\theta} \\max_{r} \\frac{1}{N} \\sum_{i=1}^N r(x_i) (\\ell(f_\\theta(x_i))) - \\lambda  r(x_i) \\log r(x_i) \\qquad \\text{s.t.   } \n\\frac{1}{N} \\sum_{i=1}^N r(x_i) = 1,$ \n\nThis includes diffusion-based decoders, contrastive losses (e.g., InfoNCE), and temporal-difference-based objectives (e.g., in Forward-Backward RL). For generative losses, as in VAEs or diffusion models, the optimization must rely on likelihood estimates (i.e., $\\log p_{\\theta,\\phi}(x)$), similarly to our alternate optimization strategy described in Eq. (4).\n\nExtending our DRO approach to contrastive or dynamic-aware representation learning is a promising direction we aim to pursue in future work.\n\n[Touati and Ollivier, 2021] Touati, Ahmed, and Ollivier, Yann. \"Learning one representation to optimize all rewards.\" Advances in Neural Information Processing Systems 34 (2021): 13-23.\n\n[Stooke et al., 2021] Stooke, Adam, et al. \"Decoupling representation learning from reinforcement learning.\" International conference on machine learning. PMLR, 2021.\n\n**Question – Sensitivity to hyperparameters (KL temperature, β)**\n\nWe report a detailed sensitivity analysis over the DRO parameter λ in Appendix D.2.1. We show that values between 10 and 100 yield the most consistent results across environments. We will highlight this better in the main paper.\n\nFor β (VAE KL weight), we performed ablations over the range β ∈ [0.5, 3]. While β = 2 yields the best performance on average across DRAG, SKEWFIT and RIG, we observed that the relative performance ranking remains unchanged across this range. We will clarify this in Appendix D. \n\nAdditional ablations in the supplementary material also explore sensitivity to $M$ (the number of latent samples per input) and $\\tau$ (the EMA coefficient for the delayed encoder)."}, {"review_id": "Tcn8eEkAEX", "reviewer": "Reviewer_xUWH", "summary": "This work contributes DRAG, a novel method for online, intrinsically-motivated, goal-conditioned reinforcement learning from high-dimensional inputs (such as images). In particular, the method leverages the framework of distributional robust optimization (DRO), which allows supervised learning models to deal with distributional shift by optimizing it against worst-case data distributions. The authors extend this framework for online reinforcement learning tasks by integrating a VAE model, thus allowing the sampling of novel goals in the latent space during training of the RL agent. The authors evaluate DRAG on 2D mazes and a robotic tasks across two different dimensions: (i) whether DRAG outperforms other methods is reaching a wide range of goals at execution time; (ii) how do different selection strategies of latent goals improve (or not) the performance of DRAG. The results show (i) how DRAG is able to achieve a higher coverage of different goals, while being more stable than the baselines; (ii) the performance of DRAG can be improved with particular resampling strategies.", "questions": "I am willing to further revise my score if the authors positively address my comments in \"Strengths And Weaknesses\", especially focused on:\n1) Can the authors clarify the scope of their contribution and its novelty?\n2) How does DRAG compare against the other baseline methods in sample and time complexity?", "limitations": "The authors don't provide an explicit limitations section or subsection in the main document, only in appendix, which is unfortunate. As previously mentioned, I would consider discussing the computational complexity of the method against other standard methods in online goal-conditioned RL, as it provide a clearer view of pros and cons of the method for a practitioner or researcher in the field.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**:\n- Well written paper, with a (welcomed) background section on DRO;\n- Interesting results, demonstrating a clear improvement over the selected baselines;\n- Extensive evaluation (even if most in Appendix).\n\n**Weaknesses**:\n- Limited novelty (see comments).\n\n**Comments**\n-  Overall I thoroughly enjoyed this paper, which I found well written and without any major typos. However, I would like to point out some minor details that are currently missing (see \"Typos and Minor Comments\" below).\n- The authors do a great job at describing all necessary background information to understand the paper, in particular the discussion on DRO (Section 2.4) which was most welcomed. However, this also leads to my biggest concern with the paper, which resides in its apparent lack of novelty: the current version of the work seems to pass the idea that the core contribution of the paper is the integration of a VAE (to process high-dimensional image observations and sample goals from the prior) in the (already existing) DRO framework. The related works of Section 2.2 and Section 2.3 also do not highlight clearly the differences between previous works and the proposed method. This is exacerbated by a confusing \"contributions\" summary part of the introduction, where point number 1 seems exactly the same as point number 3. My recommendation would be to have a dedicated related work section, where the authors clarify the differences between the existing methods and their proposed approach, instead of mixing background and related work. Moreover, I would streamline the \"contributions\" summary of the introduction.\n- Despite not being evaluated on a wide range of goal-conditioned tasks, the results on the available tasks show a significant improvement over other, DRO-like, approaches. I would recommend the authors to highlight some of the main results of the experiments available in the appendix in the main text, in particular of the hyperparameter ablation studies and the latent space visualizations, which I found interesting. However, there is one evaluation that I believe is missing from the paper: what is the time complexity of the method against the baselines? Since DRAG requires Monte-Carlo sampling of latent codes for each data-point in the mini-batch, I would assume that it's a more expensive approach than the baselines. There are some results regarding the performance of DRAG with different values of $m$ (Appendix D2) but not regarding computational time. \n\n**Typos and Minor Comments**\n- What do the authors mean with \"collaborative setting\" (line 53)? Collaborative settings in regards to RL has a specific meaning which I don't believe is what the authors what to convey.\n- (Line 234) $\\mathcal{L}_{\\theta, \\phi, r}^{\\text{DRO-VAE}}$ is not defined in any equation.\n- The use of $S$ for the state space in the extended MDP formulation and as the symbol for the success coverage metric (line 267) is unfortunate.\n- A lot of in-line equations are broken (line 178, 225, 226, 308, 311).\n- The caption of Figure 1 is quite generic... Does it represent your DRAG framework? If so, it should be mentioned in the caption.\n- In Appendix, the title of Table 3 is not in Camel-Case. Also, the structure is a bit confusing: there is an example of hyperparameter ablation (Figure 7) in Appendix A.2.2. and then a whole section dedicated to hyperparameter ablations in Section D.2. Perhaps, it can be optimized.", "was_revised": true, "final_justification": "The authors introduce a novel method for intrinsically-motivated goal-conditioned reinforcement learning from high-dimensional inputs. My main concerns with the paper (inconclusive novelty) were addressed by the authors. As such, I improved by score to 5.", "rebuttal": "We thank the reviewer for the thorough and constructive feedback. We address the main concerns below.\n\n**1. On Novelty and Related Work Presentation**\n\nWe thank the reviewer for raising concerns about the novelty and the structure of our related work discussion. While it is true that our method builds upon established components (β-VAEs, DRO principles, and goal-conditioned RL agents), we believe our contribution is novel in both formulation and application.\n\nOur key innovation lies in the integration of DRO into the online training of VAEs within reinforcement learning, which—to the best of our knowledge—has not been explored before, even outside of RL. Applying DRO in this context is non-trivial, as it requires careful adaptation of the min-max objective to variational training with approximate likelihoods, as well as efficient sampling and optimization strategies in the presence of non-stationary replay buffers. This results in a principled alternative to prior  approaches, such as SKEWFIT, which we show to be a particular instance of our more general framework.\n\nWe also agree that the distinctions between our approach and related methods could be highlighted more clearly in Sections 2.2 and 2.3. We will revise these sections as follows:\n\n- Section 2.2, which focuses on intrinsically motivated agents and goal resampling strategies, will explicitly point out that these methods are difficult to scale to high-dimensional observations such as images, due to the curse of dimensionality. A natural idea is to perform such resampling in a compressed latent space. However, we show that this is far from trivial when the representation is learned online, as it can collapse or become biased, especially when exploration induces significant distribution shifts. This motivates the need for principled control over representation learning during training.\n- Section 2.3, on online representation learning with VAEs, positions our method in the landscape of prior VAE-based methods like SKEWFIT. We will emphasize that naively training the encoder via the ELBO objective fails to account for the non-stationarity of the data, and leads to representations that underfit the full diversity of visited states. SKEWFIT proposes to mitigate this via a handcrafted reweighting scheme designed to approximate a uniform goal distribution. However, as we show, this corresponds to a particular instantiation of a more general DRO objective, which can lead to unstable training. One of our contributions is to formalize this connection and leverage the theoretical tools of DRO to derive a more stable and general reweighting strategy, grounded in optimization theory.\n\nRegarding the clarity of the stated contributions, we acknowledge that Contribution 1 (DRO formulation for VAEs) and Contribution 3 (practical algorithm) may appear overlapping. We will revise the introduction to clarify the distinction:\n\n- Contribution 1 provides the theoretical formulation that adapts DRO principles to VAE-based representation learning, establishing a new optimization objective that subsumes prior methods.\n- Contribution 3 describes the algorithmic realization of this framework in the DRAG agent, which combines VAE-based goal encoding, neural importance weighting, and online training from replay buffers.\n\nFinally, we take note of the suggestion to better separate background and related work. While our current structure aimed to streamline exposition, we agree that a clearer separation (e.g., explicitly distinguishing background from prior art, or deferring part of it to an appendix) may improve readability and will consider this in the camera-ready version.\n\n\n**2. Experimental Scope and Ablations in Main Text**\n\nWe agree that highlighting key results from the appendix in the main paper would improve clarity. We will use the extra page to incorporate in the main text:\n- A short discussion of latent space visualizations.\n- A summary of the ablation study findings (e.g., optimal values for $\\lambda$, stability across $\\beta$, the number of samples $M$ per estimation, etc.).\n\n**3. Runtime and Computational Cost**\n\nWe agree with the reviewer that we should also include runtime measurement, to assess the impact of the overhead induced by the additional likelihood estimation (for SKEWFIT and DRAG compared to RIG) and the use of a neural weighter (for DRAG). With additional experiments reported in table 1 below (values are averaged over 6 seeds on each environment), we observe  that DRAG and SKEWFIT require on average about 200 seconds for 20k steps, against 176 seconds with RIG. This slight overhead is negligible compared to the time spent in environment interactions. This is because we  only use one sample (i.e. $M=1$) for the likelihood estimation in DRAG and SKEWFIT (please refer to the experiments with $M=10$ in Appendix D.2.2, which shows that using more samples does not impact the results of both methods). Also, no notable difference in runtime is observed when comparing DRAG to SKEWFIT. \n\n**Table 1 : Mean execution runtime on V100 GPU (32GB) for MAZE and METAWORLD environments**\n| Method          | Time per 20K steps (s)  | Time for 4M steps (s) |\n|-----------------|------------------------|------------------------|\n| DRAG  | 211                    | 42,200  (11.72 h)                   |\n| SKEWFIT  | 208                    | 41,600  (11.55 h)                   |\n| RIG             | 176                    | 35,200  (9.77 h)                   |\n\n\n\nTo better link performance to runtime, Table 2 shows success coverage over wall-clock time (same runs as Fig. 3):\n\n**Table 2 : Success coverage / runtime for MAZE and METAWORLD environments**\n| Hour | **DRAG** (%) | Skew-FIT (%) | RIG (%) |\n|------|----------|--------------|---------|\n| 1    | **34**       | 31           | 26      |\n| 3    | **45**       | 42           | 28      |\n| 5    | **60**       | 50           | 32      |\n| 7    | **70**       | 58           | 36      |\n| 9    | **77**       | 63           | 40      |\n| 11   | **81**       | 66           | 42      |\n\n\nWe will add this additional analysis in Appendix D.\n\n**4. Clarification on “Collaborative Setting”**\n\nWe thank the reviewer for pointing this out. Our intention was to emphasize the mutual interaction between the representation learner and the RL agent. However, we acknowledge that “collaborative” may be misleading in the RL context, where it often implies multi-agent cooperation. We will rephrase the sentence to avoid confusion.\n\n**5. Definition of $L_{\\theta, \\phi, r}^{\\text{DRO-VAE}}$**\n\n$L_{\\theta, \\phi, r}^{\\text{DRO-VAE}}$ refers to the quantity introduced in Equation (5). We agree that the phrasing at line 234 could be clearer and will revise it to:\n\n“…where this approximated lower bound, which we call hereafter $L_{\\theta, \\phi, r}^{\\text{DRO-VAE}}({x_i}_{i=1}^n)$, can be estimated at each step via Monte Carlo…”\n\n**6. Notation and Formatting Issues**\n\n- We agree that using $S$ both for state space and success coverage metric is unfortunate. We will change the success metric notation (e.g., to $C$) for clarity.\n- Broken in-line equations (lines 178, 225, 226, 308, 311): thank you, we will fix these in the final version.\n- Figure 1 caption: we confirm that it illustrates the DRAG framework, as it includes \"DRO sampling\" in the red part. We will clarify this. \n- Appendix Table 3 title and section organization: thanks for pointing this out. We will improve the formatting and merge/restructure ablation-related sections for clarity."}, {"review_id": "XJQL9dJb2F", "reviewer": "Reviewer_u4AE", "summary": "This paper introduces DRAG, a method for goal-conditioned reinforcement learning that addresses exploration challenges in high-dimensional state spaces by combining a $\\beta$-VAE with DRO. Unlike traditional approaches where auto-encoders overfit to frequently visited states, DRAG uses an adversarial neural weighter to prioritize underrepresented states during training, ensuring broader state-space coverage and more effective goal sampling. The method eliminates the need for explicit exploration strategies by leveraging DRO to guide latent-space learning, enabling the policy to generalize to unseen goals. Experiments in maze navigation and robotic control tasks demonstrate that DRAG outperforms baselines like RIG and Skew-Fit in both state coverage and downstream task performance, even without pre-training or environment priors.", "questions": "1. The paper claims DRAG improves \"state space coverage and downstream control performance\"—could you clarify what \"downstream control\" specifically refers to? Is it the policy’s ability to reach arbitrary goals sampled from the latent space? Or does it include task-specific performance (e.g., maze exit) when the policy is fine-tuned with sparse rewards?\n\n2. The background/related work section is somewhat extensive. Have you considered moving parts of it to the appendices to improve readability? Meanwhile, Figures 8 and 9 seem critical for understanding DRAG’s contributions. Why were they placed in the appendix instead of the main text? Their visual impact could help readers grasp the method’s advantages immediately.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\nThis work advances RL by addressing the valuable yet understudied challenge of improving state coverage for VAEs in RL, which has broad implications for both model-free and model-based RL. The paper provides both theoretical derivations and experiments for adversarial reweighting, demonstrating its effectiveness in hard-exploration tasks. By ensuring the VAE’s latent space comprehensively represents the environment, DRAG unlocks more future study on reliable goal-conditioned policies.\n\n\n**Weaknesses:**\n\n1. Unclear contribution statement: The introduction’s contribution list is somewhat disorganized, making it hard to discern the core novelty (DRAG via DRO-VAE) from secondary insights (e.g., reinterpreting Skew-Fit). A clearer structure would first present DRAG as the main contribution, then contextualize it against baselines (Skew-Fit as a non-parametric DRO-VAE instance), and finally summarize empirical results.\n\n2. The experimental validation is conducted exclusively in controlled, simplistic environments (2D mazes and single robotic arm tasks with static backgrounds), which raises significant questions about the method's generalizability to more realistic scenarios. Some limitations emerge from this narrow scope:\n    - The environments use clean, synthetic visual inputs without real-world challenges like camera viewpoint changes, dynamic backgrounds, visual distractors (e.g., moving objects), etc. This leaves untested whether the adversarial weighter can maintain robust performance when the state representation must handle noisy, high-variance visual inputs common in real applications.\n    - The static environments don't evaluate DRAG's ability to handle time-varying visual features, moving targets or obstacles, and non-stationary dynamics. Such conditions could particularly challenge the DRO component's assumption of a coherent uncertainty set.", "was_revised": true, "final_justification": "I'll keep my original score", "rebuttal": "We thank the reviewer for their insightful comments and suggestions. We address their points below:\n\n**Regarding experimental validation and generalization:**\n\nWe thank the reviewer for these insightful suggestions. To partially address them within the rebuttal timeframe, we performed additional experiments where input images were perturbed by noise, turning each pixel black with probability $p=0.2$ according to a Bernoulli distribution. The results show that while all methods (DRAG, SKEWFIT, DRO) experience some slight performance degradation under noise, the relative comparative trends remain unchanged, with DRAG maintaining its advantage.\n\n**Table 4: Mean Success Coverage across Noisy Maze and Robotic Environments after 4M training steps**\n| Methods | Success coverage (with noise) | Success coverage (without noise)\n| --------- | ---------------- | ---------------- |\n| **DRAG**  | **72%**          | **81%**          |\n| Skew-FIT  | 59%              | 66%              |\n| RIG       | 36%              | 42%              |\n\nRegarding distractors, moving targets, and dynamic obstacles, these are indeed challenging aspects. However, there is no clear reason why DRAG would lose its edge in such scenarios. Our method operates in a fully unsupervised setting where any observed state is a potential goal to reach, enabling preparation for a broad range of future tasks. The aim of this paper is not to propose methods for state selection or skill decomposition, but rather to maximize coverage given a set of states. Introducing uncontrollable distractors naturally expands the state and goal space, increasing the problem complexity. While interesting, this lies somewhat beyond the scope of the current study.\n\n\nRegarding the handling of time-varying visual features and non-stationary dynamics, we agree that this is an interesting direction of analysis, which can pose challenges to many RL approaches based on learned representations. Testing DRAG’s ability to recover after abrupt changes in dynamics—such as adding a wall in a maze, as done in SVGG [Castanet et al., 2023]—is a promising direction for future work to further validate robustness. However, we emphasize that the vanilla version of DRAG is not specifically designed to handle non-stationary dynamics. That said, when combined with a dedicated goal selection strategy like the SVGG criterion used in Section 4.2 of our paper (which leverages the entropy of a success predictor), there is reason to believe it could recover effectively.\n\n**On the meaning of “downstream control”:**\n\nBy “downstream control,” we specifically refer to the policy’s ability to reach arbitrary goals sampled from the learned latent goal space. This captures goal-conditioned control performance (accuracy and efficiency in reaching goals). We do not include task-specific fine-tuning performance with sparse rewards in this definition. We will clarify this point in the revised manuscript (we agree that “downstream control” may be misleading).\n\n**On related work placement and appendix figures:**\n\nFigures 8 and 9 were placed in the appendix primarily due to their size (occupying a full page each). We agree these figures are important to visually convey DRAG’s advantages. As the final version allows one extra page, we will prioritize moving these figures to the main text to improve accessibility and impact.\n\nThank you again for your constructive feedback."}, {"review_id": "UzLBxDzlWm", "reviewer": "Reviewer_pMkP", "summary": "This paper proposes a method to improve representation learning in online goal-conditioned reinforcement learning, particularly in complex visual environments. The key idea is to guide the learning process toward better state space coverage by integrating principles from distributionally robust optimization into a VAE framework. By adaptively reweighting training samples, the method encourages learning from underrepresented states, leading to more effective exploration and goal-reaching capabilities. The approach is compatible with existing GCRL frameworks and demonstrates improved performance on tasks requiring broad exploration, without relying on prior knowledge or pretraining.", "questions": "1. Can the authors clarify how reweighting VAE training samples directly addresses the distributional shift problem? While the method uses DRO-inspired sample weighting, the theoretical or empirical justification for why this resolves or leverages distributional shift is not sufficiently discussed.\n2. Extending experiments and discussions about the approach beyond the β-VAE setup could help demonstrate its general applicability and broaden its impact.\n3. Could the authors provide additional evaluation metrics to assess the quality of the learned representations? For example, visualizations of the latent space or metrics such as clustering quality, mutual information could help support the claim of improved representation learning.\n4. It would be helpful to include ablation hyper-parameters analysis to better understand the contribution of each component in DRAG.\n5. The paper would benefit from a clearer explanation of how the proposed method directly addresses the challenges outlined in the introduction. Currently, the transition from the problem motivation to the technical formulation of DRAG lacks sufficient intermediate discussion. Rather than immediately diving into optimization objectives and formulas, it would be helpful to first provide illustrative examples that connect the theoretical design to the practical issues raised earlier. This could significantly improve readability and help the reader understand why the proposed weighting mechanism is a suitable and effective solution to the stated problem.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strength:\n1. The paper identifies a critical limitation in prior GCRL methods, that VAE-based representations tend to overfit to frequently visited states. By addressing this issue, the proposed approach enables broader and more effective exploration of the environment.\n2. The authors introduce a principled method that integrates distributionally robust optimization into VAE training. This design effectively mitigates distributional bias and encourages learning from underexplored states, leading to improved coverage of the state space.\n3. The method is evaluated across two environments, demonstrating its generality and robustness across different task types.\n\nWeakness:\n1. The paper emphasizes the importance of addressing distributional shift in the introduction. However, the proposed method focuses on reweighting training samples in the VAE using DRO. The connection between this sample reweighting and the mitigation or utilization of distributional shift remains underexplained. It is not fully clear why optimizing VAE weights in this manner would directly resolve the shift-related issues discussed earlier.\n2. While the method is evaluated on two distinct environments, the paper does not sufficiently explore the generalizability of the approach to broader representation learning paradigms. In particular, it remains unclear whether the proposed distributionally robust framework could be extended to alternative representation learning methods such as contrastive learning or diffusion-based models, which are increasingly used in modern reinforcement learning.\n3. The evaluation relies solely on success coverage, without additional analysis of representation quality (e.g., latent space visualization, clustering metrics). Moreover, the experimental design lacks ablation studies or hyper-parameters analyses on key components. This limits the understanding of the contribution and robustness of individual components in the proposed framework.", "was_revised": true, "final_justification": "The rating has been updated based on the authors' rebuttal.", "rebuttal": "Many thanks for your thorough and constructive feedback. We appreciate the opportunity to clarify the key motivations, theoretical underpinnings, and empirical design of our work. Below, we address each of your concerns in detail.\n\n***Weakness 1, Question 1, Question 5: Clarification of the connection between distributional shift and DRO-based sample reweighting***\n\nThank you for the insightful comments. While Section 2.4 provides an  explanation, we agree that the connection between DRO and distributional shift could be more clearly emphasized and that we should do so earlier. We will add a clarifying sentence in the introduction.\n\nIn short, Distributionally Robust Optimization (DRO) is a family of approaches specifically designed to deal with distributional shifts that can arise at test time. This was originally applied to supervised ML tasks (e.g., classification), by optimizing performance under worst-case distributions within a certain neighborhood of the empirical data distribution, by reweighting training samples. This naturally lends to anticipate test-time distributions that differ from those encountered during training, by giving more emphasis on training data points located on less populated areas of the feature space (we refer the reviewer for instance to [Rahimian and Mehrotra, 2019] or [Michel et al., 2022] for further justification on this connection). The strength of this robustness is controlled by a KL-regularized constraint (via λ), which interpolates between focusing on worst-case regions (low λ) and recovering the standard objective (high λ).\n\nIn the VAE setting, applying DRO leads to reweighting samples with high reconstruction loss—typically those in sparse or high-variance latent regions—thus preventing the model from collapsing to dominant modes. While DRO has never been used in generative modeling (since VAEs are usually employed to reproduce a given data distribution, not a larger distribution around it), we argue it is especially valuable in online representation learning for GCRL, where the data distribution shifts continuously and lacks clear support boundaries.\n\nThis shift creates a “vicious circle”: poor representation for underrepresented states induce narrow policies, which in turn induce narrow data. Our method breaks this loop by reweighting toward underrepresented states using the VAE’s own generative distribution—effectively steering the encoder toward better support for future exploration.\n\nWe agree that a visual illustration (that we already have, but  unfortunately we cannot include it in this rebuttal due to new NeurIPS policies) would clarify this intuition, and will include it in the final version.\n\n***Weakness 2 & Question 2: On generalization to other representation learning paradigms***\n\nWe agree this is an important direction. While the paper focuses on β-VAEs due to their interpretability and controlled disentanglement properties, the DRO-based reweighting principle is not tied to any specific representation learning architecture.\n\nIn fact, following the DRO principle, any method that learns representations from replay buffer samples using a loss function can be recast as:\n\n$\\min_{\\theta} \\max_{r} \\frac{1}{N} \\sum_{i=1}^N r(x_i) (\\ell(f_\\theta(x_i))) - \\lambda  r(x_i) \\log r(x_i) \\qquad \\text{s.t.   } \n\\frac{1}{N} \\sum_{i=1}^N r(x_i) = 1,$ \n\nwhere the sum is computed on $N$ samples $x_i$ from the buffer, and $\\ell(f_\\theta(x_i))$ corresponds to a quantity to be minimized for outputing $f_\\theta(x_i)$ for sample $x_i$. For Forward-Backward approaches [Touati and Ollivier, 2021], this loss stands as a Bellman temporal difference on the successor measure. For state diffusion methods, this corresponds to a negative log-likelihood of the diffusion decoder output. Note that, similarly to our DRO-VAE, as diffusion models imply a lower bound, $r$ must also be optimized from an estimation of the true likelihood of $\\tilde{ L_{\\theta,\\phi}}(x)\\approx \\log{p_{\\theta,\\phi}(x)}$ rather than the ELBO, in an alternate optimization scheme as we propose in Eq. (4). Other extensions, using contrastive learning to include dynamics in the representation, could incorporate an InfoNCE term alongside reconstruction in a VAE-like setting, to bring temporally adjacent states closer in latent space (as in [Stooke et al., 2021]) while encouraging the representation to cover the full state distribution from the buffer. We leave this for future work. \n\n\n[Touati and Ollivier, 2021] Touati, Ahmed, and Ollivier, Yann. \"Learning one representation to optimize all rewards.\" Advances in Neural Information Processing Systems 34 (2021): 13-23.\n\n[Stooke et al., 2021] Stooke, Adam, et al. \"Decoupling representation learning from reinforcement learning.\" International conference on machine learning. PMLR, 2021.\n\n***Weakness 3, Question 3: On representation quality metrics and additional evaluation***\n\nWe appreciate the reviewer’s suggestion. While we did not report clustering metrics or mutual information, we provide several visualizations and analyses of the latent space evolution over training in the appendix (please refer to the submitted supplementary material). Specifically:\n\n- **Figure 9** shows the evolution of the learned representation, alongside the intrinsic goal distribution and the success coverage, for DRAG, SKEWFIT, and RIG in 2D Maze tasks;\n- **Figure 10** depicts the evolution of the learned prior distribution in the Fetch environment for the same three methods;\n- **Figure 11** presents the magnitude of representation change over training, measured through latent distance evolution.\n\nFrom these figures, we observe that DRAG consistently yields more structured and well-distributed representations, leading to better state space coverage. Additionally, DRAG's representation evolves more steadily and with less variability across seeds. In contrast, SKEWFIT displays unstable training dynamics, while RIG exhibits almost no change in the latent space beyond a certain training point.\n\nTo complement these results from the paper, we followed the reviewer's suggestion by assessing the obtained representations in term of the trustworthiness score [Venna & Kaski, 2001], which measures to what extent the local neighborhood structure in the input space is preserved in the latent space. Higher trustworthiness indicates better preservation of task-relevant information across the encoding. We computed this trustworthiness metric (with 5 nearest neighbors) using a batch of 1K uniformly sampled observations from the valid state space, identical to the batch used for success coverage. Our experiments show that DRAG consistently improves this score over baseline encoders across environments. These results will be added in the appendix. \n\n\n\n**Table 1: Mean trustworthiness of obtained representations across environments after 4M training steps**\n| Methods | trustworthiness [Venna & Kaski. 2001]\n|----------|----------|\n| **DRAG**    | **98.7%**   | \n| SKEWFIT    | 93.3%   | \n| RIG    | 88.5%   |\n\n\n\n\n***Question 4: On ablations and hyperparameter sensitivity***\n\nWe thank the reviewer for highlighting this point. Ablation studies and sensitivity analyses are indeed important to assess the contribution of individual components in our framework.\n\nWe provide these analyses in the appendix:\n\n- **Figure 12** presents a sensitivity analysis over the main hyperparameter λ, which controls the strength of distributional robustness. DRAG performs best for λ values between 10 and 100, with optimal coverage at λ=100. Smaller values can lead to instability due to excessive focus on rare samples, while larger values overly regularize toward the marginal distribution p(x), limiting exploration. For comparison, SKEWFIT performs best at α=−1, corresponding to λ=1 in the DRO formulation (see Appendix C). This difference stems from SKEWFIT’s pointwise posterior estimation, which requires less aggressive reweighting than the smoother, neural-weighted distribution used by DRAG\n- **Figure 13** compares the performance of DRAG and SKEWFIT when varying the number of Monte Carlo samples M used to estimate the generative posterior. While increasing M (from 1 to 10) improves the accuracy of the likelihood estimate, it does not lead to better agent performance—in fact, it slightly reduces success coverage on average. This suggests that the robustness of DRAG is not due to more accurate pointwise estimations, but rather to the smoother weighting function induced by its neural predictor, which better generalizes across similar inputs. In contrast, SKEWFIT’s pointwise estimates can lead to abrupt weight changes, especially for poorly represented inputs.\n- **Figure 7** studies the influence of smoothing the VAE encoder used during action selection. We employ a delayed encoder updated via EMA to stabilize latent representations, using hyper-parameter τ. The best performance is obtained with τ=0.05; larger values (e.g., τ=1) make the representation unstable, while smaller values (e.g, τ=0.01) slow down adaptation excessively. \n\nWe will make sure to better highlight these results in the main paper in the final version."}]

Author rebuttal:

Meta-review: This paper integrates distributionally robust optimization into VAE-based representation learning to encourage more uniform visitation of states and improve exploration in online RL settings with complex "semantically sparse" visual input. The proposed method is motivated by the inherent nature of distribution shifts within online settings. This method highlights and addresses an important problem of representation collapse in online RL, and the baselines are appropriate and the results convincing. There were concerns about novelty -- being a combination of well-known methods -- but the authors successfully convinced the reviewers otherwise. Additional experiments wrt wall-clock time were provided and also seemed to satisfy reviewers.

I therefore recommend the paper is accepted as a poster.
Decision: Accept (poster)