ViSA: Visited-State Augmentation for Generalized Goal-Space Contrastive Reinforcement Learning

2026-03-16

AI Summary

This work addresses the limited generalization of conventional contrastive reinforcement learning in goal-conditioned tasks, particularly to rare or hard-to-reach goal states, which stems from inaccuracies in value function estimation due to insufficient goal coverage. To overcome this, the authors propose ViSA, a method that leverages state data augmentation to generate diverse training samples and incorporates mutual information regularization alongside embedding space consistency constraints. This approach constructs a goal-sensitive yet structurally coherent representation space, substantially enhancing the policy’s ability to estimate values across a broad goal distribution. Empirical results demonstrate that ViSA achieves effective generalization and accurate evaluation for challenging goals in both simulated and real-world robotic tasks.

Technology Category

• Intelligent Robots: State Estimation
• Machine Learning: Imitation Learning & Inverse Reinforcement Learning
• Computer Vision: Visual Reasoning & Symbolic Representations

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Goal-Conditioned Reinforcement Learning (GCRL) is a framework for learning a policy that can reach arbitrarily given goals. In particular, Contrastive Reinforcement Learning (CRL) provides a framework for policy updates using an approximation of the value function estimated via contrastive learning, achieving higher sample efficiency compared to conventional methods. However, since CRL treats the visited state as a pseudo-goal during learning, it can accurately estimate the value function only for limited goals. To address this issue, we propose a novel data augmentation approach for CRL called ViSA (Visited-State Augmentation). ViSA consists of two components: 1) generating augmented state samples, with the aim of augmenting hard-to-visit state samples during on-policy exploration, and 2) learning consistent embedding space, which uses an augmented state as auxiliary information to regularize the embedding space by reformulating the objective function of the embedding space based on mutual information. We evaluate ViSA in simulation and real-world robotic tasks and show improved goal-space generalization, which permits accurate value estimation for hard-to-visit goals. Further details can be found on the project page: \href{https://issa-n.github.io/projectPage_ViSA/}{\texttt{https://issa-n.github.io/projectPage\_ViSA/}}

Problem

• Goal-Conditioned Reinforcement Learning
• Contrastive Reinforcement Learning
• Value Function Estimation
• Goal-Space Generalization
• Visited-State Augmentation

Innovation

• Visited-State Augmentation
• Contrastive Reinforcement Learning
• Goal-Conditioned Reinforcement Learning
• Embedding Regularization
• Mutual Information