Imagine Beyond ! Distributionally Robust Autoencoding for State Space Coverage in Online Reinforcement Learning

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.

Weakness 1, Question 1, Question 5: Clarification of the connection between distributional shift and DRO-based sample reweighting

Thank 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.

In 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 λ).

In 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.

This 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.

We 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.

Weakness 2 & Question 2: On generalization to other representation learning paradigms

We 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.

In fact, following the DRO principle, any method that learns representations from replay buffer samples using a loss function can be recast as:

$\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. } \frac{1}{N} \sum_{i=1}^N r(x_i) = 1,$

where 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.

[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.

[Stooke et al., 2021] Stooke, Adam, et al. "Decoupling representation learning from reinforcement learning." International conference on machine learning. PMLR, 2021.

Weakness 3, Question 3: On representation quality metrics and additional evaluation

We 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:

• 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;
• Figure 10 depicts the evolution of the learned prior distribution in the Fetch environment for the same three methods;
• Figure 11 presents the magnitude of representation change over training, measured through latent distance evolution.

From 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.

To 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.

Table 1: Mean trustworthiness of obtained representations across environments after 4M training steps

Question 4: On ablations and hyperparameter sensitivity

We 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.

We provide these analyses in the appendix:

• 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
• 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.
• 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.

We will make sure to better highlight these results in the main paper in the final version.

We thank the reviewer for their insightful comments and suggestions. We address their points below:

Regarding experimental validation and generalization:

We 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.

Table 4: Mean Success Coverage across Noisy Maze and Robotic Environments after 4M training steps

Regarding 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.

Regarding 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.

On the meaning of “downstream control”:

By “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).

On related work placement and appendix figures:

Figures 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.

Thank you again for your constructive feedback.

We thank the reviewer for the thorough and constructive feedback. We address the main concerns below.

1. On Novelty and Related Work Presentation

We 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.

Our 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.

We 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:

• 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.
• 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.

Regarding 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:

• Contribution 1 provides the theoretical formulation that adapts DRO principles to VAE-based representation learning, establishing a new optimization objective that subsumes prior methods.
• 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.

Finally, 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.

2. Experimental Scope and Ablations in Main Text

We 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:

• A short discussion of latent space visualizations.
• A summary of the ablation study findings (e.g., optimal values for $\lambda$, stability across $\beta$, the number of samples $M$ per estimation, etc.).

3. Runtime and Computational Cost

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 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.

Table 1 : Mean execution runtime on V100 GPU (32GB) for MAZE and METAWORLD environments

To better link performance to runtime, Table 2 shows success coverage over wall-clock time (same runs as Fig. 3):

Table 2 : Success coverage / runtime for MAZE and METAWORLD environments

We will add this additional analysis in Appendix D.

4. Clarification on “Collaborative Setting”

We 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.

5. Definition of $L_{\theta, \phi, r}^{\text{DRO-VAE}}$

$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:

“…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…”

6. Notation and Formatting Issues

• 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.
• Broken in-line equations (lines 178, 225, 226, 308, 311): thank you, we will fix these in the final version.
• Figure 1 caption: we confirm that it illustrates the DRAG framework, as it includes "DRO sampling" in the red part. We will clarify this.
• 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.

Thank you again for your kind and encouraging comments. However, as far as we can see from the author interface, the score has not been updated as you mentioned in your previous comment. If you still intend to raise it, we’d be very grateful if you could kindly submit the change when you get a chance. Many thanks in advance.

Many thanks for your thorough and constructive feedback. We address the main concerns below.

Weakness: Quality (and Question about runtime overhead)

• Limited scope (image resolution, task diversity):

Thanks 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:

Table 1: Mean success coverage across maze and robotic environment for 128x128 pixels observations (success coverage for 82x82 is given for comparison)

All methods degrade slightly with higher resolution, but DRAG retains a clear advantage.

We 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.

• Sample-efficiency and runtime analysis

First, 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.

However, 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.

Table 2 : Mean execution runtime on V100 GPU (32GB) for MAZE and METAWORLD environments

To better link performance to runtime, Table 3 shows success coverage over wall-clock time (82×82, same runs as Fig. 3):

Table 3 : Success coverage / runtime for MAZE and METAWORLD environments

We will add this additional analysis in Appendix D.

Weakness: Clarity on ELBO and hyperparameters: "a clear presentation of the exact ELBO and hyperparamters used would help reproducibility"

We 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)$.

• Method Derivation

If 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:
$\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)$

$=\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)$

$=\max_{\theta, \phi} \min_{r:E_{p}r = 1} L_{\theta,\phi,r}$

where $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)$

Then, we could write a lower bound as: $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)$

However, 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.

• Reproducibility (Algorithm and Hyper-parameters)

Maybe 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.

Weakness: Originality

While 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.

Our 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)).

Weakness Significance: Comparison with other representation learning paradigms

We 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.

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 β-VAEs for interpretability and disentanglement, our DRO framework is general and not tied to a specific representation architecture.

In fact, any representation learning method trained from replay buffer samples using a loss $\ell(f_\theta(x))$ can be reweighted using the DRO objective:

$\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. } \frac{1}{N} \sum_{i=1}^N r(x_i) = 1,$

This 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).

Extending our DRO approach to contrastive or dynamic-aware representation learning is a promising direction we aim to pursue in future work.

[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.

[Stooke et al., 2021] Stooke, Adam, et al. "Decoupling representation learning from reinforcement learning." International conference on machine learning. PMLR, 2021.

Question – Sensitivity to hyperparameters (KL temperature, β)

We 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.

For β (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.

Additional 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).