[Paper Review] Learning to Score Behaviors for Guided Policy Optimization
This paper proposes a novel framework for guided policy optimization in reinforcement learning by using Wasserstein distances in a learned behavioral embedding space to score and steer policies toward or away from desired behaviors. By leveraging the dual formulation of entropy-regularized Wasserstein distances, the method enables efficient stochastic gradient updates through behavioral regularizers, leading to two new on-policy algorithms—Behavior-Guided Policy Gradient and Behavior-Guided Evolution Strategies—that outperform existing methods in challenging continuous control and deceptive reward environments.
We introduce a new approach for comparing reinforcement learning policies, using Wasserstein distances (WDs) in a newly defined latent behavioral space. We show that by utilizing the dual formulation of the WD, we can learn score functions over policy behaviors that can in turn be used to lead policy optimization towards (or away from) (un)desired behaviors. Combined with smoothed WDs, the dual formulation allows us to devise efficient algorithms that take stochastic gradient descent steps through WD regularizers. We incorporate these regularizers into two novel on-policy algorithms, Behavior-Guided Policy Gradient and Behavior-Guided Evolution Strategies, which we demonstrate can outperform existing methods in a variety of challenging environments. We also provide an open source demo.
Motivation & Objective
- To address the challenge of effectively measuring and leveraging policy similarity beyond local action similarity in reinforcement learning.
- To develop a flexible, differentiable method for comparing policies based on their global behavioral trajectories rather than parameters or state visitation frequencies.
- To enable efficient policy optimization by using Wasserstein distances in a latent behavioral space with stochastic gradient updates.
- To improve exploration and generalization in RL by incorporating behavioral repulsion and attraction signals via learned test functions.
- To demonstrate the effectiveness of the approach in challenging environments, including those with deceptive rewards and sparse rewards.
Proposed method
- The paper introduces Behavioral Policy Embeddings (BPEs) that map trajectories into a compact latent behavioral space using a Behavioral Embedding Map (BEM), enabling policy comparison via probability distributions over embeddings.
- It defines policy dissimilarity using Wasserstein distances (WDs) between these embedded distributions, which allows for non-injective BEMs and avoids issues with likelihood-based divergences like KL divergence.
- The dual formulation of the WD is used to derive behavioral test functions that score trajectories or state-policy pairs, guiding optimization toward or away from specific behaviors.
- Entropy regularization of the WD enables efficient, stochastic gradient descent updates via Reproducing Kernel Hilbert Space (RKHS) representations using random feature maps.
- Two new on-policy algorithms are proposed: Behavior-Guided Policy Gradient (BGPG), replacing KL trust regions with WD-based regularizers, and Behavior-Guided Evolution Strategies (BGES), which jointly optimize for reward and behavioral novelty.
- The method supports both on-policy and off-policy embeddings, with off-policy variants using probing state distributions to evaluate policies independently.
Experimental results
Research questions
- RQ1Can a Wasserstein-based metric in a learned behavioral embedding space provide a more effective and interpretable measure of policy similarity than traditional methods like KL divergence?
- RQ2How can the dual formulation of the Wasserstein distance be leveraged to generate meaningful behavioral test functions that guide policy optimization?
- RQ3Can WD-based regularizers enable faster and more sample-efficient policy optimization compared to KL-based trust region methods?
- RQ4Can the proposed framework improve exploration in environments with deceptive rewards by using behavioral repulsion to avoid suboptimal policies?
- RQ5To what extent can the same framework be applied to imitation learning and safety-constrained RL through behavioral scoring?
Key findings
- BGPG outperforms standard TRPO with KL divergence across six continuous control tasks from the DeepMind Control Suite and Roboschool, achieving higher final performance with consistent gains.
- BGPG achieves a normalized reward of 90% of the best result up to 3.5x faster than the baseline method from Zhang et al., with clock time reductions from ~26k to ~4k seconds on Hopper and Walker tasks.
- In deceptive reward environments, BGES successfully learns to navigate past barriers to reach goals, while NSR-ES and NoisyNet-TRPO fail to learn forward locomotion due to suboptimal local rewards.
- For the quadruped environment, BGES achieved rewards exceeding -5000, indicating successful navigation around the wall, whereas baselines remained trapped in local optima.
- For the point-mass environment, BGES achieved rewards above -800, demonstrating effective exploration and goal-reaching despite deceptive rewards.
- The method generalizes to imitation and repulsion learning, showing that behavioral test functions can be used to guide policies toward or away from specific behaviors.
Better researchstarts right now
From reading papers to final review, dramatically reduce your research time.
No credit card · Free plan available
This review was created by AI and reviewed by human editors.