P3: Probabilistic Policy Propagation for Stable VAE-Based Robot Learning

2026-07-28

AI Summary

This work addresses a critical mismatch in the VAE-PPO framework, where the stochasticity of the encoder’s latent space is incompatible with policy optimization, leading to high-variance and biased policy gradient estimates due to single-sample approximation—thereby undermining learning stability and data efficiency. The authors propose the P³ framework, which systematically uncovers the theoretical shortcomings of this approximation mechanism and introduces a distribution-aware optimization strategy. By integrating moment propagation with sampling calibration, P³ enables robust modeling of latent uncertainty. Empirical results demonstrate that this approach substantially improves data efficiency—increasing it from 64.6% to over 96%—and accelerates convergence by reducing training steps by more than 20%, with strong validation on challenging humanoid parkour tasks.

Abstract

Variational Autoencoders are widely used to encode high-dimensional and noisy observations in robotics. However, their stochastic latent creates a mismatch with Proximal Policy Optimization (PPO): an effective policy marginalizes over the latent distribution, whereas former implementations estimate its probability ratio and KL divergence using only one latent sample. We identify a fundamental but overlooked theoretical cause: naive single-sample approximations in stochastic latent space induce significant variance and bias in the surrogate loss. To address this, we introduce P^3 (Probabilistic Policy Propagation), a distribution-aware optimization framework for VAE-based policies. $P^3$ couples moment-based probabilistic method for stable and efficient learning with sampling-based calibration for robust policy behavior under latent uncertainty. In our experiments, P^3 boosts data efficiency from 64.6% to >96%, reduces convergence steps by >20%. Furthermore, P^3 is evaluated on challenging humanoid parkour tasks and shows an effective foundation for VAE-based PPO. Code is available at https://github.com/ylyem9x/P3_Open.

Problem

• Variational Autoencoders
• Proximal Policy Optimization
• stochastic latent space
• policy learning
• surrogate loss

Innovation

• Probabilistic Policy Propagation
• VAE-based reinforcement learning
• distribution-aware optimization
• stochastic latent space
• policy gradient variance reduction