PaperlayerView exact revision on arXiv

Source-linked AI summary

Uncertainty Weighted Actor-Critic for Offline Reinforcement Learning

Yue Wu, Shuangfei Zhai, Nitish Srivastava, Joshua Susskind, Jian Zhang, Ruslan Salakhutdinov, Hanlin Goh

arXiv:2105.08140v1cs.LG

TL;DR

Offline RL must learn from static datasets when online data collection is costly, but robustly handling uncertainty around out-of-distribution state-action pairs remains important. UWAC estimates uncertainty with dropout and down-weights uncertain Bellman targets, improving training stability and benchmark performance.

  • Problem

    Offline RL seeks to use previously collected datasets when large-scale online data collection is infeasible, while uncertainty handling remains central to robust learning.

  • Method

    UWAC is a model-free actor-critic method that uses dropout-based uncertainty estimates to down-weight high-uncertainty target state-action pairs in Bellman updates.

  • Results

    UWAC achieves superior performance and improved training stability on standard offline RL benchmarks, including significant gains on narrow human demonstrations.

  • Takeaways & Limitations

    Uncertainty estimation can detect OOD samples and stabilize offline RL without adding another model or loss term.

  • Takeaways & Limitations

    Applying UWAC to stabilize other actor-critic methods and combining model-based with model-free offline RL remain future work.

Abstract

from arXiv · show

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.

1. Introduction

Offline RL aims to learn from static datasets, but OOD states and actions can destabilize bootstrapped value learning. UWAC estimates uncertainty and down-weights high-uncertainty targets, improving training stability and benchmark performance.

  • Offline RL is motivated by the high cost of active data collection and the need to reuse previously collected datasets.
  • OOD states and actions can produce arbitrary value estimates whose errors propagate through Bellman losses and destabilize training.
  • UWAC uses dropout-based uncertainty estimation to identify high-uncertainty state-action pairs and down-weight their contribution to Bellman updates.
  • Uncertainty-based down-weighting substantially improves training stability over BEAR and achieves state-of-the-art performance across standard offline RL benchmarks.
  • UWAC is presented as a simple, efficient technique requiring no additional loss terms or models, with gains on narrow human demonstrations.

2. Related Work

Related offline RL methods address distribution shift through models, pessimism, or policy constraints, while UWAC uses model-free uncertainty weighting. Its MC-dropout estimator offers a simpler alternative to explicit ensembles.

  • Offline RL is vulnerable to inadequate dataset coverage, actor-critic distribution shifts, and neural-network variance that accumulate as Bellman errors.
  • Model-based methods estimate uncertainty through environment models, whereas model-free methods use pessimism or constrain the actor toward the training distribution.
  • UWAC down-weights the Bellman loss by the inverse uncertainty of the critic target using Monte Carlo dropout.
  • MC-dropout behaves similarly to explicit ensembles while being substantially simpler.
  • BEAR and CQL avoid the model-performance limitation affecting model-based methods, while MOReL and MOPO can be limited by data-dependent model quality.

3. Preliminaries

The preliminaries formulate offline RL as an MDP and describe Q-learning, actor optimization, and BEAR’s support constraint. They identify OOD maximization as a source of bootstrapping instability and constrained generalization.

  • 3.1. Notations: The environment is modeled as an MDP with state and action spaces, transitions, rewards, an initial-state distribution, and discount factor γ.
  • 3.1. Notations: Q-learning estimates expected cumulative discounted reward through the Bellman optimality operator and Bellman squared error.
  • 3.1. Notations: In continuous action spaces, an actor is learned to select actions that maximize the critic function.
  • 3.1. Notations: Naively maximizing Q over next actions can select unseen OOD actions, causing large errors to propagate through Bellman bootstrapping.
  • 3.2. Baseline Algorithm: BEAR restricts policy actions toward the training distribution using an approximate support constraint based on sampled MMD.
  • 3.2. Baseline Algorithm: The BEAR constraint can limit Q-function generalization beyond the dataset and is associated with numeric instability on narrower offline datasets.

4. Uncertainty weighted offline RL

UWAC estimates uncertainty with MC dropout and incorporates it into both critic and actor training to reduce the influence of OOD backups and actions. The resulting objectives are designed to stabilize learning while preserving policy generalization beyond the dataset.

  • Uncertainty estimation: UWAC uses dropout-based Bayesian inference to estimate model uncertainty for detecting OOD state-action pairs.The Q function is trained with dropout before each weight layer, and the model-uncertainty component is used for OOD detection.
  • Uncertainty estimation: The uncertainty-weighted policy distribution π′ modifies the original policy using Q-function uncertainty and a state-dependent normalization factor Z(s).The paper states that π′ is defined with respect to π, the Q function, and Z(s).
  • Training procedure: UWAC’s theoretical analysis reports better convergence properties against OOD training samples, while its algorithm initializes Q networks with MC dropout and estimates target variance from stochastic forward passes.The pseudocode uses T = 100 stochastic forward passes and does not propagate gradients through the uncertainty estimate.
  • Critic objective: The critic is trained on π′, with Bellman loss down-weighted by the inverse uncertainty of the target Q estimate.This reduces the effect of high-uncertainty OOD backups on the overall training process.
  • Actor objective: The actor is optimized using samples from π′, reducing the probability of maximizing Q on OOD samples without constraining the actor distribution.The resulting actor loss is intended to discourage Q-function explosion and stabilize Q estimates.

5.1. Dropout Uncertainty Estimation for Reinforcement Learning

UWAC uses dropout-based uncertainty to identify sparse or OOD state-action pairs, then down-weights uncertain backups during offline actor-critic training. Experiments show improved OOD detection, training stability, and performance across MuJoCo and Adroit datasets, with benefits extending beyond dropout uncertainty estimation.

  • Dropout Uncertainty Estimation: MC dropout assigns higher Q-function uncertainty to sparse or removed LunarLander observations, especially OOD states.The uncertainty estimates are higher where observations were removed and lower where training coverage is denser.
  • Dropout Uncertainty Estimation: MC dropout distinguishes training state-action pairs from random-action pairs with a ROC/AUC score of 0.845 for identifying OOD actions.The uncertainty regions overlap because many random actions remain in-distribution in the small, noisy robotic-control action space.
  • Standard Benchmarking Datasets: UWAC achieves strong performance across 15 MuJoCo benchmark settings, particularly on medium-expert datasets, while performing less well than MOPO on medium-replay datasets.The authors attribute MOPO’s advantage on medium-replay to model-based methods performing well when state/action diversity is high.
  • Adroit Hand Dataset: UWAC significantly improves over BEAR on all Adroit human-demonstration datasets and achieves state-of-the-art performance on the other Adroit datasets.The human datasets contain narrow demonstrations and therefore expose the agent to many OOD backups during training.
  • Training Dynamics and Ablations: Down-weighting OOD backups breaks the Q-value-explosion cycle that destabilizes BEAR, allowing UWAC to continue training and maintain positive performance.Spectral normalization alone does not handle OOD backups on narrow Adroit datasets, while UWAC also works with ensemble-based uncertainty estimates and shows no significant variance-versus-standard-deviation difference.

6. Conclusion and Future Work

UWAC uses uncertainty estimation to detect and down-weight OOD backups in offline RL, improving performance and training stability without additional models or losses. The work also identifies applications for uncertainty estimation and future extensions to other actor-critic methods.

  • 6. Conclusion and Future Work: UWAC detects and down-weights OOD backups in the Bellman squared loss for offline RL.The method improves training stability without introducing an additional model or loss.
  • 6. Conclusion and Future Work: UWAC achieves superior performance and improved training stability without introducing any additional model or losses.
  • 6. Conclusion and Future Work: Uncertainty estimation may guide model-based versus model-free training decisions and curiosity-based learning in future work.These directions are presented as possible extensions rather than evaluated contributions of this work.

A.1. Analysis for Convergence Properties

The convergence analysis bounds offline RL Bellman approximation error by using uncertainty-weighted policy constraints under boundedness and normalization assumptions. This addresses the otherwise unbounded OOD error term and yields convergence with respect to OOD samples.

  • A.1. Analysis for Convergence Properties: Theorem A.1 alone does not ensure convergence because OOD estimates can make the Bellman approximation error arbitrarily large.
  • A.1. Analysis for Convergence Properties: Theorem A.1 bounds approximation error using Bellman approximation error, a suboptimality constant, and a concentrability coefficient.The proof modifies a contraction argument from prior convergence analyses.
  • A.1. Analysis for Convergence Properties: The suboptimality constant measures how far the optimal policy is from the policy set, while concentrability measures its distance from the training distribution.These quantities capture actor suboptimality and exposure to OOD actions and states, respectively.
  • A.1. Analysis for Convergence Properties: Under mild assumptions, Theorem A.2 bounds the Bellman error term by any constant with arbitrarily high probability through uncertainty-weighted policy optimization.The analysis assumes Z(s) ≥ 1 and bounded Q values; boundedness can be imposed with spectral normalization.
  • A.1. Analysis for Convergence Properties: The theorem analyzes inverse standard-deviation weighting, whereas the practical algorithm uses inverse-variance weighting with clipping for numerical stability.Variance weighting is chosen to simplify tracing through repeated max, min, and summation operations.
  • A.1. Analysis for Convergence Properties: The resulting constant Bellman-error bound allows convergence of the proposed framework with respect to OOD samples.The proof uses triangle inequality, Chebyshev’s inequality, and bounds derived from the Q-function assumptions.

A.2. Training Time of MC Dropout

UWAC with dropout requires less than 1.5 times BEAR’s training time while estimating uncertainty using batches 100 times larger than the original batch size on a single GPU.

  • A.2. Training Time of MC Dropout: Less than 1.5 times BEAR’s training time is required by UWAC with dropout using 100 times the original batch size for uncertainty calculation.The enlarged uncertainty batch is processed in parallel on a single GPU.

A.3. Observations on the Q Value Divergence of BEAR

UWAC complements BEAR’s action-support constraint by further penalizing OOD actions that remain after MMD-based support enforcement.

  • A.3. Observations on the Q Value Divergence of BEAR: UWAC further penalizes OOD actions that survive BEAR’s imperfect MMD-based action-support constraint.The authors report that changing BEAR’s learning rate among 10^-3, 10^-4, and 10^-5 did not change its divergence behavior.

B. Figures

The figures show that UWAC improves training stability and performance on D4RL Adroit data, while ablations indicate that uncertainty reweighting is central to the gains. Additional plots and demonstrations illustrate uncertainty detection, comparisons with alternative methods, and learned policies across tasks.

  • UWAC assigns higher uncertainty to states made out of distribution by removing observations with vertical displacements above 0.8.The figure visualizes higher uncertainty with brighter colors and lower uncertainty with colder colors.
  • UWAC achieves significantly better overall performance and training stability than BEAR on D4RL Adroit data.BEAR’s performance degrades over time while its target Q value explodes; UWAC avoids this pattern.
  • Spectral normalization stabilizes Q estimates but often produces significantly worse training performance in average return.The comparison uses BEAR and BEAR with spectral-normalized Q functions on D4RL Adroit data.
  • Dropout alone does not improve performance without the UWAC reweighting loss.This comparison is made between BEAR with spectral normalization and BEAR with dropout plus spectral normalization.
  • Removing uncertainty reweighting causes performance to drop drastically, despite retaining dropout in the Q function.The ablation compares UWAC with a version using dropout but no uncertainty weighting.
  • The figures also compare UWAC with Averaged-DQN ensembles and REM, and examine variance versus standard-deviation down-weighting.Separate figures depict learned policies from expert and human demonstrations across MuJoCo and Adroit tasks.
Loading 2105.08140v1…