[Paper Review] Uncertainty Weighted Actor-Critic for Offline Reinforcement Learning

Yue Wu, Shuangfei Zhai+5 more|arXiv (Cornell University)|May 17, 2021

Reinforcement Learning in Robotics43 references20 citations

TL;DR

This paper proposes Uncertainty Weighted Actor-Critic (UWAC), a method that improves offline reinforcement learning by detecting out-of-distribution (OOD) state-action pairs using dropout-based uncertainty estimation and down-weighting their contribution in the Bellman update. UWAC achieves state-of-the-art performance on benchmark offline RL environments, especially on sparse human demonstration datasets, with improved training stability and reduced error propagation from OOD backups.

ABSTRACT

Offline Reinforcement Learning promises to learn effective policies from previously-collected, static datasets without the need for exploration. However, existing Q-learning and actor-critic based off-policy RL algorithms fail when bootstrapping from out-of-distribution (OOD) actions or states. We hypothesize that a key missing ingredient from the existing methods is a proper treatment of uncertainty in the offline setting. We propose Uncertainty Weighted Actor-Critic (UWAC), an algorithm that detects OOD state-action pairs and down-weights their contribution in the training objectives accordingly. Implementation-wise, we adopt a practical and effective dropout-based uncertainty estimation method that introduces very little overhead over existing RL algorithms. Empirically, we observe that UWAC substantially improves model stability during training. In addition, UWAC out-performs existing offline RL methods on a variety of competitive tasks, and achieves significant performance gains over the state-of-the-art baseline on datasets with sparse demonstrations collected from human experts.

Motivation & Objective

• To address the instability and performance degradation in offline reinforcement learning caused by out-of-distribution (OOD) state-action pairs during bootstrapping.
• To reduce destructive error propagation from inaccurate Q-value estimates on OOD samples, which commonly destabilizes actor-critic training.
• To develop a lightweight, effective method that improves generalization without adding extra models or loss terms.
• To demonstrate that uncertainty estimation can stabilize offline RL and enhance performance on sparse expert demonstrations.

Proposed method

• Uses dropout during inference to estimate model uncertainty, treating the variance across multiple forward passes as an uncertainty score for each state-action pair.
• Applies a weighted Bellman update where targets with high uncertainty are down-weighted using a learned scaling factor β.
• Integrates uncertainty weighting directly into the standard actor-critic loss, modifying the Bellman loss to reduce influence of high-uncertainty targets.
• Employs spectral normalization on the Q-network for improved stability in high-dimensional control tasks like Adroit.
• Tunes the uncertainty weighting hyperparameter β to balance uncertainty sensitivity and training convergence.
• Maintains the same network architecture and hyperparameters as baseline methods (e.g., BEAR), ensuring fair comparison.

Experimental results

Research questions

• RQ1 Can dropout-based uncertainty estimation effectively detect out-of-distribution state-action pairs in offline reinforcement learning?
• RQ2 Does down-weighting high-uncertainty targets in the Bellman update improve training stability and final performance in offline RL?
• RQ3 Can UWAC outperform existing offline RL methods, especially on datasets with sparse human demonstrations?
• RQ4 Is the uncertainty weighting mechanism robust across different uncertainty estimation techniques, such as ensembles?

Key findings

• UWAC significantly improves training stability, as shown by smoother return curves and more controlled Q-value targets compared to baseline methods.
• On the narrow human demonstration dataset (relocate-expert), UWAC achieves a 20% higher average return than the baseline BEAR, which fails to outperform random behavior.
• UWAC outperforms state-of-the-art methods like BEAR and REM on standard offline RL benchmarks, including AntMaze and Adroit environments.
• The method remains robust when replacing dropout with ensemble-based uncertainty estimation, indicating generalization across uncertainty estimation techniques.
• Down-weighting by standard deviation or variance yields no significant performance difference, validating the robustness of the uncertainty weighting mechanism.
• Spectral normalization improves stability but does not resolve OOD issues alone; UWAC’s uncertainty weighting is essential for performance on narrow datasets.