Research map / Reinforcement learning / Visual and representation RL
Imagine Beyond ! Distributionally Robust Autoencoding for State Space Coverage in Online Reinforcement Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DRAG, a distributionally robust beta-VAE with an adversarial neural weighter that reweights replay-buffer states so online GCRL latent spaces progressively cover the full state space from pixels.
Problem
In online goal-conditioned RL from pixels, VAE representations learned from the agent's own rollouts over-represent frequently visited states and fail to cover the full state space, creating exploration bottlenecks, especially when latent prior goals drive intrinsic motivation.
Approach
Cast online VAE training as a DRO problem with a KL-constrained uncertainty set and learn a parametric neural weighter (adversarial, batch-normalized softmax weights) that upweights underrepresented states; shows Skew-Fit is a non-parametric instance of this DRO-VAE and that a parametric weighter is more stable.
Key results
- DRAG significantly outperforms RIG and Skew-Fit in success coverage over 4M steps on pixel mazes, with RIG capped by an exploration bottleneck and Skew-Fit showing high instability.
- Trustworthiness of learned latent spaces after 4M steps: DRAG 98.7%, Skew-Fit 93.3%, RIG 88.5%.
- Combining DRAG with SVGG goal selection significantly outperforms sampling goals from the prior, while MEGA and LGE give no improvement and GoalGAN-style GOID selection degrades performance.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use DRAG instead of RIG for online representation learning | GCRL with pixel observations · Environments after 4M training steps | trustworthiness: 88.5% -> 98.7% | Helped |
| Use DRAG instead of Skew-Fit for online representation learning | GCRL with pixel observations · Environments after 4M training steps | trustworthiness: 93.3% -> 98.7% | Helped |
| Add MEGA goal selection to DRAG (ablation) | DRAG · Success coverage | success coverage | No clear effect |
| Add LGE goal selection to DRAG (ablation) | DRAG · Success coverage | success coverage | No clear effect |
| Add GoalGan-like GOID goal selection to DRAG (ablation) | DRAG · Success coverage | success coverage | Hurt |
| Add SVGG goal selection to DRAG (ablation) | DRAG · Success coverage | success coverage | Helped |
| Set the delayed-encoder smoothing coefficient to 1 (no smoothing) (ablation) | DRAG · Success coverage | success coverage | Hurt |
| Set the delayed-encoder smoothing coefficient to 0.01 (ablation) | DRAG · Success coverage | success coverage | Hurt |
| Use a VAE training schedule of every 10k agent steps initially, then every 20k (ablation) | DRAG, RIG, and Skew-Fit · All environments | average results | Hurt |
| Use a VAE training schedule of every 5k agent steps for the first 100k, then every 10k (ablation) | DRAG, RIG, and Skew-Fit · All environments | average results | Hurt |
| Use a VAE training schedule of every 2k agent steps for the first 100k, then every 5k (ablation) | DRAG, RIG, and Skew-Fit · All environments | average results | Hurt |
| Increase the Monte Carlo estimator sample count from M=1 to M=10 (ablation) | DRAG and Skew-Fit · Success coverage | success coverage | No clear effect |
Surprising findings
- Adding diversity-based goal selection (MEGA, LGE) on top of DRAG yields no improvement, suggesting redundancy with the DRO representation learning.
- Using M=10 Monte Carlo samples for the likelihood estimate does not translate into better success coverage than M=1.
Limitations
- Evaluated only on two simple environments (pixel point mazes at 82x82 and a Reach variant with walls); no pre-trained visual encoders used.
- Parametric weighter requires a very small learning rate and careful temperature tuning (lambda), with instability at low lambda.
- Monte Carlo likelihood estimation adds decoder computation per sample.
Experimental setup
Datasets: online rollouts from pixel PointMaze (82x82) online rollouts from pixel Reach-Hard-Walls (MetaWorld-based)
Benchmarks: Pixel continuous PointMazes Pixel Reach-Hard-Walls
Metrics: success coverage trustworthiness
Baselines: RIG Skew-Fit MEGA LGE GoalGAN SVGG
Builds on: RIG Skew-Fit beta-VAE TQC Distributionally Robust Optimization
Closest papers by meaning
- Reasoning with Latent Diffusion in Offline Reinforcement Learning (ICLR 2024 · similarity 0.88)
- Leveraging Skills from Unlabeled Prior Data for Efficient Online Exploration (ICML 2025 · similarity 0.87)
- Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models (NeurIPS 2025 · similarity 0.87)
- VIGOR: Zero-Shot Visual Generalization via Latent-Space Consistency in Model-Based Reinforcement Learning (NeurIPS 2026 · similarity 0.87)
- Towards Principled Representation Learning from Videos for Reinforcement Learning (ICLR 2024 · similarity 0.87)
- Flattening Hierarchies with Policy Bootstrapping (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.