Active Reinforcement Learning Strategies for Offline Policy Improvement

Background

Reinforcement Learning (RL) is a paradigm for sequential decision-making where agents learn optimal behaviors by interacting with an environment to maximize cumulative rewards. Traditional RL relies on online exploration, where agents iteratively collect data through trial-and-error. While effective in simulated settings, this approach becomes impractical in real-world applications—such as clinical trials, autonomous navigation, or robotics—where interactions are costly, risky, or constrained by safety protocols.

To address these limitations, offline RL has emerged as a promising alternative. Here, agents learn from a fixed dataset of pre-collected interactions, eliminating the need for risky online exploration. However, offline RL faces a critical challenge: the dataset may lack coverage of critical states or actions, especially if the environment changes between data collection and deployment. For instance, an autonomous vehicle trained on urban driving data might struggle in rural settings absent from the dataset. This limitation necessitates strategic online exploration to complement offline data, but under strict budget constraints.

This problem mirrors Active Learning in supervised learning, where agents select the most informative unlabeled data points to annotate, maximizing model performance with minimal labeling effort. However, in RL, the challenge is more complex. Instead of selecting static data points, agents must decide:

• Where to explore (e.g., under-sampled regions).
• How to explore (e.g., trajectory length).
• When to stop (e.g., avoiding redundant states).

Prior work in RL exploration has focused on uncertainty-driven methods, such as:

• Count-based exploration: Encouraging visits to rarely seen states.
• Intrinsic rewards: Using curiosity or prediction errors to guide exploration.
• Ensemble disagreement: Leveraging model variance to identify uncertain regions.

While these methods excel in pure online RL, they do not account for pre-existing offline datasets. Conversely, offline RL algorithms like CQL (Conservative Q-Learning) and IQL (Implicit Q-Learning) regularize policies to avoid overestimation errors but lack mechanisms for targeted online data collection. Bridging this gap requires a framework that actively guides exploration using uncertainty estimates derived from offline data.

Our work builds on these foundations to address a key question: How can agents efficiently collect new trajectories—within a limited interaction budget—to augment an offline dataset and learn robust policies? We propose that representation-aware uncertainty estimation is critical to identifying under-explored regions while reusing knowledge from offline data. By combining insights from active learning and offline RL, we design a method that strategically allocates exploration resources, ensuring every interaction maximizes learning progress.

Proposed Method:

To address the challenge of enhancing an agent’s performance using limited online interactions alongside offline data, we propose a representation-aware uncertainty-based active exploration strategy. Our method efficiently guides the agent to collect high-value trajectories in under-explored regions of the environment, minimizing redundant data collection.

1. Representation-Aware Uncertainty Estimation

The core idea is to quantify epistemic uncertainty (model uncertainty due to lack of data) in the state-action space using an ensemble of representation models.

Model Architecture

We design two encoders:

• State encoder \( E_s \): Maps states to latent embeddings.
• Action encoder \( E_a \): Maps actions to latent embeddings.

State-action pairs are encoded as:

\[ E(s,a) = E_s(s) + E_a(a) \]

Training Objectives

• Clustering: We ensure that similar states are clustered in the latent space.
• Transition Dynamics Modeling: We enforce that the embedding of a state-action pair aligns with the next state’s embedding.

Uncertainty Calculation

For a given state or state-action pair, we compute the disagreement among the ensemble’s predictions.

We use the maximum squared difference between embeddings across ensemble members as our uncertainty metric.

2. Active Exploration Policy

Using the uncertainty estimates, we strategically determine where to start exploring and how to collect trajectories.

Initial State Selection

We sample candidate initial states:

\[ C \sim {\rho} \]

We prioritize states with the highest uncertainty, focusing on regions that are poorly covered by the offline dataset.

Trajectory Collection

• At each step, we sample multiple actions from the current policy (e.g., TD3+BC or IQL) and select the action leading to the most uncertain state-action pair.
• We employ an \( \epsilon \)-greedy strategy:
• With probability \( \epsilon \), we explore using uncertainty-driven actions.
• Otherwise, we follow the offline policy.

Trajectory Termination

We stop collecting data in regions where uncertainty falls below a threshold to avoid redundancy.

3. Handling Restricted Initial States

In environments where initial states cannot be freely chosen , we employ a two-stage policy:

• Reach Phase: We use a goal-conditioned policy (trained offline) to navigate from the default initial state to a high-uncertainty region in the current offline dataset.
• Explore Phase: Once the target region is reached, we switch to our uncertainty-driven exploration policy.