Generalized Advantage Estimation for Distributional Policy Gradients

TL;DR

This paper introduces Distributional GAE (DGAE), a new advantage estimation method for distributional RL that uses optimal transport theory to better handle stochasticity and improve policy gradient estimates.

Contribution

It proposes DGAE, integrating Wasserstein-like metrics into advantage estimation for distributional RL, enhancing robustness and low variance in policy updates.

Findings

DGAE outperforms traditional GAE in various environments.

DGAE provides low-variance advantage estimates with controlled bias.

The method improves policy gradient stability in stochastic systems.

Abstract

Generalized Advantage Estimation (GAE) has been used to mitigate the computational complexity of reinforcement learning (RL) by employing an exponentially weighted estimation of the advantage function to reduce the variance in policy gradient estimates. Despite its effectiveness, GAE is not designed to handle value distributions integral to distributional RL, which can capture the inherent stochasticity in systems and is hence more robust to system noises. To address this gap, we propose a novel approach that utilizes the optimal transport theory to introduce a Wasserstein-like directional metric, which measures both the distance and the directional discrepancies between probability distributions. Using the exponentially weighted estimation, we leverage this Wasserstein-like directional metric to derive distributional GAE (DGAE). Similar to traditional GAE, our proposed DGAE provides a low-variance advantage estimate with controlled bias, making it well-suited for policy gradient algorithms that rely on advantage estimation for policy updates. We integrated DGAE into three different policy gradient methods. Algorithms were evaluated across various OpenAI Gym environments and compared with the baselines with traditional GAE to assess the performance.