Low-Variance Policy Gradient Estimation with World Models
This work addresses sample efficiency in reinforcement learning for researchers and practitioners, though it appears incremental as it builds on existing methods like AC and MAC.
The paper tackles the problem of high variance in policy gradient estimates by introducing World Model Policy Gradient (WMPG), which uses learned world models to imagine trajectories for gradient estimation and baseline calculation, resulting in improved sample efficiency demonstrated on environments like CartPole, LunarLander, and Pong.
In this paper, we propose World Model Policy Gradient (WMPG), an approach to reduce the variance of policy gradient estimates using learned world models (WM's). In WMPG, a WM is trained online and used to imagine trajectories. The imagined trajectories are used in two ways. Firstly, to calculate a without-replacement estimator of the policy gradient. Secondly, the return of the imagined trajectories is used as an informed baseline. We compare the proposed approach with AC and MAC on a set of environments of increasing complexity (CartPole, LunarLander and Pong) and find that WMPG has better sample efficiency. Based on these results, we conclude that WMPG can yield increased sample efficiency in cases where a robust latent representation of the environment can be learned.