Bootstrap Advantage Estimation for Policy Optimization in Reinforcement Learning

TL;DR

This paper introduces Bootstrap Advantage Estimation (BAE), a novel data augmentation technique for policy optimization in reinforcement learning that improves sample efficiency and generalization across diverse environments.

Contribution

The paper presents a new bootstrap advantage estimation method that enhances policy and value function updates using data augmentation, outperforming existing techniques like GAE, RAD, and DRAC.

Findings

BAE reduces policy and value loss more effectively than GAE.

BAE improves cumulative return across multiple benchmarks.

BAE demonstrates superior sample efficiency and generalization in unseen environments.

Abstract

This paper proposes an advantage estimation approach based on data augmentation for policy optimization. Unlike using data augmentation on the input to learn value and policy function as existing methods use, our method uses data augmentation to compute a bootstrap advantage estimation. This Bootstrap Advantage Estimation (BAE) is then used for learning and updating the gradient of policy and value function. To demonstrate the effectiveness of our approach, we conducted experiments on several environments. These environments are from three benchmarks: Procgen, Deepmind Control, and Pybullet, which include both image and vector-based observations; discrete and continuous action spaces. We observe that our method reduces the policy and the value loss better than the Generalized advantage estimation (GAE) method and eventually improves cumulative return. Furthermore, our method performs better than two recently proposed data augmentation techniques (RAD and DRAC). Overall, our method performs better empirically than baselines in sample efficiency and generalization, where the agent is tested in unseen environments.