Contrastive Initial State Buffer for Reinforcement Learning

TL;DR

This work tackles the challenge of sample-efficient reinforcement learning by reusing past experiences to steer data collection through an Initial State Buffer (ISB). It introduces a Contrastive Learning Buffer (CL-Buffer) that learns an embedding space where states with similar learning experiences are grouped together, enabling adaptive, diverse state sampling via K-Means clustering. Across quadruped locomotion and drone racing tasks, the CL-Buffer accelerates convergence and boosts final performance (e.g., up to 18.3% improvement on the quadruped task and a 0.9 vs 0.2 success-rate advantage in drone racing) without altering the underlying RL algorithm. The approach offers a general, prior-free mechanism to improve data efficiency in robotics and can be extended with priors or prioritized sampling for further gains.

Abstract

In Reinforcement Learning, the trade-off between exploration and exploitation poses a complex challenge for achieving efficient learning from limited samples. While recent works have been effective in leveraging past experiences for policy updates, they often overlook the potential of reusing past experiences for data collection. Independent of the underlying RL algorithm, we introduce the concept of a Contrastive Initial State Buffer, which strategically selects states from past experiences and uses them to initialize the agent in the environment in order to guide it toward more informative states. We validate our approach on two complex robotic tasks without relying on any prior information about the environment: (i) locomotion of a quadruped robot traversing challenging terrains and (ii) a quadcopter drone racing through a track. The experimental results show that our initial state buffer achieves higher task performance than the nominal baseline while also speeding up training convergence.

Figures (6)

• Figure 1: Initial State Buffer for Reinforcement Learning. Our method uses a network to project observations $o_{t_k}$ to an embedding space, in which we apply K-means clustering. In a next step, we add states $s_{t_i}$ close to the cluster center to an Initial State Buffer, which sets the initial states of the robot in the environment $s^i_0$ during the roll-out.
• Figure 2: sub-MDP. A standard MDP problem can be divided into multiple sub-MDPs.
• Figure 3: Quadrupedal Locomotion. The mean validation performance at different iteration steps obtained with five different training seeds for all of the tested methods.
• Figure 4: State Distribution. (a) The quadruped states during a roll-out phase in the middle of the training in a top-down view. (b) For flying through a race track, if the agent can not fly yet through the complete racetrack, (c) the initial state clusters from the CL-Buffer are located around the struggling gate. (d) Once the agent can finish the racetrack, the CL-Buffer clusters are more spread out while still focusing on the difficult parts around the gates and less on the straight lines.
• Figure 5: Cluster State Visualization. Our proposed projection network can be used to cluster the embeddings of different states according to the corresponding experience. It can be observed that each cluster represents a specific skill, i.e., Walking Forward, Walking Sideways Up, and Failure State.