ViSA: Visited-State Augmentation for Generalized Goal-Space Contrastive Reinforcement Learning
Authors: Issa Nakamura, Tomoya Yamanokuchi, Yuki Kadokawa, Jia Qu, Shun Otsub, Ken Miyamoto, Shotaro Miwa, Takamitsu Matsubara
Organizations: I. Nakamura, T. Yamanokuchi, Y. Kadokawa, and T. Matsubara are with the Graduate School of Information Science, Nara Institute of Science and Technology (NAIST), Nara, Japan. · J. Qu, S. Otsubo, Y. Susumu, S. Miwa are with the Advanced Technology R&D Center, Mitsubishi Electric Corporation, Hyogo, Japan.
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: https://issa-n.github.io/projectPage_ViSA/
Figures & tables
Fig. 1: Overview of ViSA. This approach has two components: 1) generating augmented state samples to augment 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. The embedding space is trained to pull a visited state closer to a state-action pair (anchor) and push a random state farther away, encoding goal reachability from the anchor. ViSA suppresses estimation bias by considering relative distances of augmented states and augmented random states with respect to the anchor, enabling accurate value estimation and action selection for diverse goals in CRL.
Fig. 2: Framework of ViSA: (a) Generating augmented state samples. Using augmentation distribution p(sa∣sv) , hard-to-visit states from on-policy rollouts are artificially augmented. Crosses indicate samples. Dotted region is inherently reachable state space, blue region contains samples from visited state distribution p(sv∣s,a) , and red region includes additional samples from p(sa∣sv) . (b) Learning consistent embedding space. To regularize embedding space and prevent overfitting to visited states sv , we use not only anchors (s,a) and visited states sv+,sv− from the previous CRL method but also augmented states sa+,sa− as auxiliary information. Objective function of embedding space is reformulated based on mutual information factorization, and encoders ψ,ϕ,ϕ^ are trained to maximize mutual information ISaFE((s,a);sv) estimated using augmented states sa . Mutual information terms inside dashed boxes are used to compute ISaFE((s,a);sv) , where red and green text indicate upper and lower bounds, respectively.
Fig. 3: Experimental environments
Fig. 4: Learning results from robot tasks. Plots show success rate during learning for each task. Solid lines indicate the mean over three trials, and shaded areas represent variance.
Fig. 5: Visited state distributions and learning results for a naive sample diversification CRL method. (a) Visualization of visited state distributions. CRL is modified by adjusting discount factor γ to collect visited state samples sv more broadly from on-policy rollouts. For visualization, sampling probabilities are normalized so that the maximum value is 1. (b) Learning curves of success rates. Solid lines show mean success rates over three trials for proposed method and baselines, and shaded regions designate variance.
Fig. 6: Augmentation distributions and learning results for ablations with different generation ranges of augmented states sa : (a) Visualization of augmentation distributions. Each distribution is heuristically defined by adjusting time steps from the anchor that are likely to be selected as augmented states. For visualization, sampling probabilities are normalized so that the maximum value is 1. (b) Learning curves of success rates. Dashed red line denotes maximum success rate achieved by previous CRL method, while solid lines show mean success rates over three runs for each method. Shaded regions indicate variance. (c) Visualization of sample coverage with respect to reachability for visited and augmented states. Black lines show range from minimum to maximum values, and shaded regions represent sample distribution. Values closer to 1 correspond to hard-to-visit states requiring more steps, whereas values closer to 0 indicate easy-to-visit states.
Fig. 7: Real-world learning results: (a) Learning curves of success rates. Dashed lines denote maximum success rates achieved in simulation for each method, and solid lines show success rates during real-robot learning. (b) Control trajectories for Robel Turn using the policy trained by proposed method. Cases 1–3 correspond to trajectories toward different goals. Black, green, and purple lines indicate true valve angle, measured angle, and goal angle, respectively. (c) t-SNE visualization of embedding space. Each circle represents a state-action pair, dotted circles denote representative points corresponding to the snapshot, stars denote goals, and colors indicate time steps. Distances between circles and stars reflect value, with shorter distances indicating easy-to-reach states (higher value).
Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.