Reproduce “ViSA: Visited-State Augmentation for Generalized Goal-Space Contrastive Reinforcement Learning” on OpenResearch

About this paper

Issa Nakamura, Tomoya Yamanokuchi, Yuki Kadokawa, Jia Qu, Shun Otsub, Ken Miyamoto, Shotaro Miwa, Takamitsu Matsubara

ViSA (Visited-State Augmentation) provides a framework for Goal-Conditioned Reinforcement Learning (GCRL) that mitigates sample bias in Contrastive Reinforcement Learning (CRL) through structured data augmentation and mutual information factorization. This approach enables more accurate value function estimation across diverse goals by regularizing the embedding space. Experiments showed ViSA improved learning performance and goal-space generalization in various robotic tasks, including real-world applications, achieving up to 1.5 times higher success rates in some real-robot tasks.

Key insights

• Artificially augmenting samples for inherently reachable but rarely visited states is an effective strategy to address sample bias in GCRL, provided the augmentation is structured and relevant.
• Simply augmenting data is insufficient; a principled regularization method, such as the proposed mutual information factorization, is necessary to ensure the learned embedding space remains consistent and accurately reflects goal reachability.
• The specific design of the augmentation distribution is critical; it must maintain relevance to the original visited states while intensively targeting hard-to-visit states to enhance both task performance and sample coverage.

Summary by alphaXiv.