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.
Standard approaches to goal-conditioned reinforcement learning (GCRL) that rely on temporal-difference learning can be unstable and sample-inefficient due to bootstrapping. While recent work has explored contrastive and supervised formulations to improve stability, we present a probabilistic alternative, called survival value learning (SVL), that reframes GCRL as a survival learning problem by modeling the time-to-goal from each state as a probability distribution. This structured distributional Monte Carlo perspective yields a closed-form identity that expresses the goal-conditioned value function as a discounted sum of survival probabilities, enabling value estimation via a hazard model trained via maximum likelihood on both event and right-censored trajectories. We introduce three practical value estimators, including finite-horizon truncation and two binned infinite-horizon approximations to capture long-horizon objectives. Experiments on offline GCRL benchmarks show that SVL combined with hierarchical actors matches or surpasses strong hierarchical TD and Monte Carlo baselines, excelling on complex, long-horizon tasks. Webpage and Code: https://simple-robotics.github.io/publications/survival-value-learning/
Franki Nguimatsia Tiofack, Fabian Schramm, Théotime Le Hellard +1
Inria and École Normale Supérieure, PSL Research University, Paris, France.
Offline goal-conditioned reinforcement learning (GCRL) provides a practical framework for obtaining goal-reaching policies from fixed datasets. However, learning a reliable goal-conditioned value function in long-horizon tasks remains challenging. In this paper, we identify erroneous generalization in goal-conditioned value functions as a fundamental bottleneck, and demonstrate that appropriate inductive bias in the value function is crucial for addressing the bottleneck. Building on these findings, we propose Latent-Aligned Value Learning (LAVL), an offline GCRL algorithm that integrates latent-representation-based value generalization with hierarchical planning in a unified framework. Extensive experiments on OGBench demonstrate that LAVL consistently outperforms existing offline GCRL methods, achieving the highest performance on 20 out of 22 datasets. Notably, LAVL exhibits strong performance in long-horizon tasks and trajectory stitching datasets, where prior methods suffer significant performance degradation. Our code is available at https://github.com/oh-lab/LAVL.git.
Existing approaches to offline goal-conditioned reinforcement learning (GCRL) struggle with long-horizon tasks. Discounting shrinks value differences between distant states until they fall below the function approximation error, leaving the agent with no signal for ranking states. Temporal abstraction, which treats k environment steps as a single transition, restores this signal at long range, but no single fixed k suits all state-goal distances: large k preserves value differences across long temporal distances while collapsing distinctions between nearby states, and small k does the reverse. We make this trade-off explicit and introduce Generalized Implicit Temporal Abstraction (GITA), which conditions a single value function on k. GITA trains one policy by aggregating advantage-weighted supervision across multiple k values, so scales assigning larger positive advantages to a state-goal pair contribute more strongly to its update. GITA does not need to choose between local resolution and long-range signal; it retains both without committing to a single k. On OGBench, GITA outperforms a broad range of offline GCRL baselines, raising average success rate across all tasks by 25 percentage points (73% relative improvement) over HIQL. It also improves over the strongest fixed-k method, OTA, by 7 percentage points (14% relative).
Pedro Robles Dutenhefner, Dikshant Shehmar, Wagner Meira +1
Universidade Federal de Minas Gerais · University of Alberta · Alberta Machine Intelligence Institute (Amii) +1