Imagine Beyond! Distributionally Robust Auto-Encoding for State Space Coverage in Online Reinforcement Learning

TL;DR

DRAG introduces Distributionally Robust Auto-Encoding for GCRL by marrying a β-VAE with Distributionally Robust Optimization through a neural weighter to aggressively emphasize underrepresented states. By online updating and a delayed encoder, DRAG builds latent spaces that extend beyond current experience, enabling broader state-space coverage and improved control in hard exploration tasks. The approach reframes Skew-Fit as a DRO-VAE instance and demonstrates superior stability and coverage over baselines such as RIG and Skew-Fit in pixel-based mazes and wall-bounded robotic tasks. Practically, this yields more robust intrinsic goal generation and reduces the need for auxiliary exploration strategies, with broad potential for applying DRO to other online representation-learning RL setups.

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 $β$-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 to bypass, without pre-training nor prior environment knowledge.

Figures (14)

• Figure 1: General framework of online VAE representation learning in RL. Green: RL loop using the VAE encoder to convert high-dimensional states $s_t$ to latent states $z_t$. Blue: latent goal $z_g$ 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 2: Our two questions: (a) how does drag perform as a representation learning approach? (b) how does drag impact goal sampling approaches from the literature?
• Figure 3: Evolution of the success coverage over PointMazes and Reach-Hard-Walls environments (6 seeds each) for 4M steps (shaded areas as standard deviation). Bottom: Median, Interquartile 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 agarwal2021deep.
• 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 in figure \ref{['fig:dro_res']}), drag+ X includes an additional goal resampling method X, taken among the four strategies: LGE, MEGA, GOALGAN or SVGG.
• Figure 5: Example of decoded pixel goals in maze environment: we sample latent goals from the latent prior $z\sim p(z) = \mathcal{N}(0,I)$ and plot the corresponding decoded pixel goals $p_{\theta}(x|z)$. Images were obtained using the decoder trained with drag.