[Paper Review] Efficient Exploration via State Marginal Matching

Lisa Lee, Benjamin Eysenbach+4 more|arXiv (Cornell University)|Jun 12, 2019

Reinforcement Learning in Robotics66 references95 citations

TL;DR

The paper reframes exploration in reinforcement learning as State Marginal Matching (SMM), a distribution-matching objective where the policy’s state visitation matches a target state distribution. It introduces a two-player/zero-sum game between a density model and a policy, uses fictitious play for optimization, and shows faster, broader exploration and better adaptation to new tasks, including a mixture-of-policies extension (SM4).

ABSTRACT

Exploration is critical to a reinforcement learning agent's performance in its given environment. Prior exploration methods are often based on using heuristic auxiliary predictions to guide policy behavior, lacking a mathematically-grounded objective with clear properties. In contrast, we recast exploration as a problem of State Marginal Matching (SMM), where we aim to learn a policy for which the state marginal distribution matches a given target state distribution. The target distribution is a uniform distribution in most cases, but can incorporate prior knowledge if available. In effect, SMM amortizes the cost of learning to explore in a given environment. The SMM objective can be viewed as a two-player, zero-sum game between a state density model and a parametric policy, an idea that we use to build an algorithm for optimizing the SMM objective. Using this formalism, we further demonstrate that prior work approximately maximizes the SMM objective, offering an explanation for the success of these methods. On both simulated and real-world tasks, we demonstrate that agents that directly optimize the SMM objective explore faster and adapt more quickly to new tasks as compared to prior exploration methods.

Motivation & Objective

• Define State Marginal Matching (SMM) as a principled objective for exploration and show how it yields a task-agnostic exploration policy.
• Propose a practical optimization framework via a two-player, zero-sum game between a state density model and a policy using fictitious play.
• Extend SMM to mixtures of policies to handle multimodal target distributions and accelerate exploration.
• Connect SMM to prior exploration methods and explain their near-MMM behavior and historical averaging importance.

Proposed method

• Define the state marginal distribution rho_pi(s) visited by a policy and the target distribution p*(s).
• Formulate the SMM objective as minimizing KL(rho_pi(s) || p*(s)), equivalently maximizing E[r(s)] with r(s)=log p*(s) - log rho_pi(s) plus a state entropy term.
• Develop a practical algorithm using fictitious play to alternate between fitting a density model q(s) to historical policy states and updating the policy to maximize the pseudo-reward r(s).
• Introduce a historical averaging mechanism over policies and densities to ensure convergence and prevent oscillations.
• Extend to a mixture-of-policies (SM4) with a discriminator over latent components and a mixture-state marginal, enabling multimodal distribution matching.

Experimental results

Research questions

• RQ1 Can exploration be cast as a distribution-matching problem over state marginals?
• RQ2 Does maximizing state entropy via SMM yield a single, robust exploration policy that generalizes across tasks?
• RQ3 Can a mixture of policies improve exploration for multimodal target state distributions?
• RQ4 How does SMM relate to and unify prior prediction-error based exploration methods?
• RQ5 Does the proposed fictitious-play optimization converge and outperform existing exploration strategies on complex tasks?

Key findings

• SMM yields faster exploration and better adaptation than prior exploration methods on simulated and real tasks.
• Prediction-error based methods approximately optimize the SMM objective when averaged over time, but can exhibit oscillatory dynamics without historical averaging.
• A historical averaging (fictitious play) mechanism is crucial for convergence and effective exploration.
• Mixtures of policies (SM4) further accelerate test-time exploration and improve downstream task performance.
• On Fetch and D’Claw experiments, SMM achieves wider state coverage and explores a broader range of object angles and knob rotations than baselines.
• SMM provides a task-agnostic exploration prior that can solve downstream tasks more quickly than baselines.