PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient Exploration via State Marginal Matching

Lisa Lee, Benjamin Eysenbach, Emilio Parisotto, Eric Xing, Sergey Levine, Ruslan Salakhutdinov

arXiv:1906.05274v3cs.LGcs.AIcs.ROstat.ML

TL;DR

Reinforcement learning requires effective exploration for long-horizon tasks with sparse or delayed rewards, while prior methods lack a clear objective and often target single tasks. The paper formulates exploration as State Marginal Matching, optimizes it through a policy–density-model game, and analyzes prior methods within this framework. Across simulated and real-world tasks, SMM and historical averaging support efficient exploration, faster adaptation, and convergence to policies matching target state distributions.

  • Problem

    RL exploration lacks a mathematically grounded objective with clear properties, while many existing methods are difficult to compare and repurpose beyond single tasks.

  • Method

    State Marginal Matching learns a policy whose state marginal matches a target distribution and optimizes this objective as a two-player, zero-sum game between a policy and density model.

  • Results

    SMM learns exploration policies that explore faster and adapt more quickly to new tasks, while historical averaging helps prior approximate distribution-matching methods converge.

  • Takeaways & Limitations

    A target-distribution policy can amortize exploration across downstream tasks and provide a reusable prior for tasks with potentially sparse rewards.

  • Takeaways & Limitations

    The approach assumes a target state distribution is given and that every policy's state distribution can be represented exactly by a density model in the chosen class.

Abstract

from arXiv · show

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.

1. Introduction

The paper recasts exploration as State Marginal Matching, providing an explicit objective for learning policies whose state distributions match desired targets. It develops an optimization method, unifies prior approaches as approximate distribution matching, and reports faster exploration and adaptation than prior methods.

  • RL needs exploration mechanisms for long-horizon tasks with limited or delayed rewards, especially where human supervision is expensive.
  • Prior exploration methods lack an explicit objective for defining or quantifying good exploration and often target only a single task.These limitations make it difficult to understand why methods work, compare them, and repurpose them across tasks.
  • State Marginal Matching learns a policy whose state marginal matches a target distribution, which can encode safety constraints, preferences, reward shaping, or state-dimension importance.Without prior information, the objective reduces to maximizing marginal state entropy and encourages visiting all states.
  • SMM produces a converged stochastic policy that visits states according to a target distribution and can serve as a reusable exploration policy for downstream tasks.The paper reports faster exploration and adaptation than prior state-of-the-art methods on complex RL tasks.
  • The SMM objective can be optimized as a two-player, zero-sum game between a state density model and a parametric policy using fictitious play.
  • Prior methods approximately optimize the SMM objective, but predictive-error approaches omit historical averaging, potentially preventing convergence to an exploratory policy.

2. State Marginal Matching

State Marginal Matching (SMM) defines exploration as matching a policy’s finite-episode state distribution to a target distribution. The framework explains why prediction-error methods fail to retain a single exploratory policy and extends matching to mixtures for complex targets.

  • Why Prediction Error is Not Enough: Prediction-error exploration methods can lose their exploratory effect at convergence, because intrinsic bonuses vanish or the converged RL solution becomes deterministic.In deterministic MDPs, ICM prediction error and Pseudocount bonuses approach zero; more broadly, each converged intrinsic-reward RL step admits a deterministic policy.
  • The State Marginal Matching Objective: SMM learns a policy whose state marginal distribution matches a specified target distribution, providing an explicit objective for task-agnostic exploration.The state marginal is the distribution of states visited during a finite-length episode, rather than a trajectory distribution or necessarily the infinite-horizon stationary distribution.
  • The State Marginal Matching Objective: The target distribution can encode task knowledge, including safety constraints, state preferences, reward shaping, or relative state importance.For goal-reaching tasks, the target distribution can also be derived analytically.
  • The State Marginal Matching Objective: SMM uses the pseudo-reward log p∗(s) − log ρπ(s), increasing utility for under-visited states and decreasing it for over-visited states.Optimizing only log p∗(s) would converge to the target mode rather than match the full distribution; the state-entropy term is therefore essential.
  • Better SMM with Mixtures of Policies: For multimodal targets, SMM decomposes the target into easier-to-learn components and learns a corresponding mixture of policies.The mixture state marginal averages conditional policy state distributions under a prior over latent components, while preserving a separate historical-averaging mechanism.
  • Better SMM with Mixtures of Policies: The mixture objective rewards target-density states, novel states for each component, distinguishable component assignments, and exploration across components.This decomposition resembles mutual-information objectives, and goal-conditioned RL is described as a special case when goal sampling is learned jointly with the policy.

3. A Practical Algorithm

The practical SMM algorithm replaces the policy-dependent density with a learned state-density model and optimizes the resulting max-min game. Fictitious play and historical averaging provide the algorithmic mechanism, while mixtures add component-specific density models, discriminators, and policies.

  • Optimizing the State Marginal Matching Objective: A parametric density model qψ(s) approximates the policy’s state marginal, and sufficient expressiveness makes optimization with q equivalent to the original SMM objective.The equivalence assumes that every policy’s state marginal can be represented exactly by some density model in the chosen class.
  • Optimizing the State Marginal Matching Objective: The reformulated objective is a two-player, zero-sum game in which the policy chooses π and the density model chooses q.The desired solution corresponds to a Nash equilibrium of the policy-density game.
  • Optimizing the State Marginal Matching Objective: Fictitious play is used instead of greedy alternating updates because it converges to a Nash equilibrium in finite games, whereas naive alternation is not generally guaranteed to converge.Each player responds to the historical behavior of the other player rather than only its latest strategy.
  • Optimizing the State Marginal Matching Objective: The exploration policy is the historical average over policy iterates, not the final policy produced by the last update.The historical average samples one policy at the start of each episode and follows it throughout that episode.
  • Optimizing the State Marginal Matching Objective: Historical averaging implicitly supplies the prior term omitted by prior predictive-error objectives, which the authors suggest may explain continued exploration at convergence.The comparison is summarized in Table 1 and linked to the SMM objective’s prior term.
  • Extension to Mixtures of Policies: The mixture extension learns component-specific density models, a discriminator d(z | s), and policies updated with component-specific rewards.Updates for different components can be conducted in parallel, and the algorithm returns the historical average of mixtures containing n · m policies.

4. Prediction-Error Exploration is Approximate State Marginal Matching

The paper frames prediction-error exploration as approximate state marginal matching, while showing that greedy optimization can oscillate and fail to produce uniform exploration without historical averaging.

  • Objective comparison: SMM and prediction-error methods can be expressed as related min-max objectives involving auxiliary density or predictive models.The SMM formulation includes a regularizer omitted by the compared prediction-error objectives.
  • Optimization: Historical averaging modifies alternating optimization by evaluating the policy against the historical average of auxiliary models rather than only the current model.This is implemented through fictitious play to address convergence problems in greedy optimization.
  • Interpretation of prior methods: Prior methods can still succeed because their policy iterates collectively visit diverse states, producing a diverse replay buffer despite lacking a single exploratory policy.The historical average over policies, rather than the converged policy, accounts for this behavior.
  • Didactic experiments: SMM converges to a uniform state distribution, whereas inverse, forward, count-based, and action-entropy methods are biased toward the stochastic state in the didactic gridworld.The experiment computes marginal state entropy exactly while varying environment stochasticity.
  • Didactic experiments: SMM maintains high state entropy as stochasticity increases, while the inverse model’s entropy decreases and the other methods fail to remain highly entropic across environments.The comparison uses the same didactic gridworld trajectory data across stochasticity settings.
  • Historical averaging: Without historical averaging, the policy and density players alternate probability mass between environment halves, producing oscillatory learning dynamics.Historical averaging is therefore identified as important for preventing oscillation.

5. Experimental Evaluation

Experiments in simulated and real-world manipulation tasks show that SMM achieves broad state coverage and enables faster downstream-task exploration than the evaluated baselines.

  • Experimental setup: SMM is evaluated in simulated Fetch and real-world D’Claw manipulation environments using target distributions designed for broad state or object-angle coverage.Fetch additionally assigns low probability to fallen-block states and undesirable arm configurations.
  • State Coverage at Convergence: SMM maximizes Fetch marginal state entropy at least as effectively as, and potentially better than, the compared exploration methods after 1e6 environment steps.Exploration is evaluated by discretized marginal state entropy, which is the objective SMM maximizes.
  • State Coverage at Convergence: SMM turns the D’Claw knob through a wider range of angles than baselines, with p = 0.046 for number of rotations, p = 0.019 for maximum clockwise angle, and p = 0.001 for maximum counter-clockwise angle.The evaluation averages metrics across 100 episodes and repeats training for 5 independent runs.
  • State Coverage at Convergence: On the real robot, SMM explores a wider angle range than SAC during 1e5 environment steps, while its exploration increases throughout training.The result suggests that the SMM objective is correlated with real-world exploration metrics.
  • Test-time Exploration: SMM and SM4 adapt faster than other methods, achieving a success rate 20% higher than the next best method and matching ICM’s performance in 4x fewer episodes.The comparison evaluates how quickly the acquired exploration policies find downstream test-time goals.
  • Test-time Exploration: Increasing mixture components improves test-time exploration, while historical averaging reduces the additional benefit of mixture modelling.The authors conclude that efficient exploration may require either historical averaging or mixture modelling, but not necessarily both.

6. Related Work

The paper distinguishes SMM from action-, parameter-, goal-, and state-space exploration methods, while positioning it as a principled state-matching policy prior related to intrinsic motivation and meta-RL.

  • Exploration spaces: Exploration methods operate in action, policy-parameter, goal, or state spaces, with MaxEnt RL and noise-based methods primarily exploring through actions or parameters.The paper places SMM among methods that directly target state or goal exploration.
  • State and goal exploration: SMM extends prior state marginal matching by analyzing distribution matching as a two-player game and applying historical averaging from fictitious play.The analysis also unifies several intrinsic-motivation methods as approximate state marginal matching.
  • Intrinsic motivation: Prediction-error bonuses reward states where an auxiliary predictive model performs poorly, while count-based methods directly encourage visits to novel states.These approaches explore during single-task learning but differ in how novelty is represented.
  • Meta-RL: SMM resembles meta-learning by learning a policy prior for downstream tasks, but the paper focuses specifically on learning to explore rather than the broader meta-RL problem.The connection is framed through the shared goal of supporting downstream-task solution.
  • Maximum entropy RL: MaxEnt RL maximizes action entropy rather than state entropy, and the exact relationship between trajectory and state distributions is generally intractable in MDPs.This distinguishes action-space entropy objectives from SMM’s direct state-marginal objective.
  • Inverse RL: SMM alternates policy and density-model learning like some inverse RL methods, but obtains its reward function from a density model rather than a discriminator.The paper groups adversarial games and density estimation among shared methodological ingredients.

7. Conclusion

The paper formalizes exploration through State Marginal Matching, whose converged policy visits states according to a target distribution. This policy amortizes exploration across tasks, while historical averaging supports proper distribution matching and prior-method analysis.

  • At convergence, SMM yields a policy that visits states in proportion to their density under a target distribution.
  • The resulting policy can serve as a prior for multi-task settings, amortizing exploration and enabling faster adaptation to new, potentially sparse reward functions.
  • Historical averaging explains how to perform distribution matching properly and guarantees convergence when added to prior exploration methods.
  • SMM unifies prior exploration methods as approximate distribution matching and offers a clearer perspective for improving exploration with adversarial games and density estimation.

A. Proofs

The proof rewrites the SMM objective as a policy-dependent entropy-like term plus a KL divergence to the density model. Under the stated assumption, the divergence can be made zero, yielding the desired result.

  • The objective decomposes into E_ρπ(s)[log p∗(s) − log ρπ(s)] plus D_KL(ρπ(s) ∥ q(s)).
  • With the divergence term eliminated, the proof obtains the desired result for the remaining policy objective.
  • Assumption 1 ensures that some q in Q makes D_KL(ρπ(s) ∥ q(s)) equal to zero.

B. Choosing p∗(s) for Goal-Reaching Tasks

For goal-reaching tasks, the target state distribution is chosen from the test-time goal distribution and an ε-ball reaching criterion. The resulting optimal policy matches a transformed goal distribution, while goal-conditioned implementations depend on sampling and supervision.

  • Choosing the target distribution: The target distribution p∗(s) depends on the distribution of test-time tasks, with goals sampled from a known distribution p_g(g).
  • Choosing the target distribution: Goal reaching is defined as visiting a state within an ε ball of the goal, using known ε > 0 and a known distance metric.
  • Deriving the objective: The analysis bounds expected episodes to reach a random goal by considering whether any state in a trajectory reaches that goal.
  • Optimal target distribution: For nonzero ε, the optimal target distribution is a box-filter-smoothed version of the goal distribution before taking its square root.
  • Optimal target distribution: The optimal policy therefore performs distribution matching to a function of the goal distribution, which need not itself be a normalized probability distribution.
  • Connections to Goal-Conditioned RL: Goal-conditioned RL is a special case of SMM with a mixture policy whose component maps bijectively to goal states.
  • Connections to Goal-Conditioned RL: Goal-conditioned distribution matching is favored when goals are easy to sample or test-time supervision provides a goal, whereas latent-conditioned policies may explore better without test-time goals.

C.1. Navigation experiments

Navigation and Fetch experiments compare state-space exploration, historical averaging, non-uniform targets, and computational cost. SMM explores more broadly, matches non-uniform targets well, and has comparable wall-clock cost to prior methods.

  • Navigation experiments: In the three-hallway navigation environment, SAC explores only one hallway, whereas SMM explores all three during training.
  • Historical averaging: Historical averaging improves SMM and prior methods, supporting the view that prior exploration methods approximately optimize the SMM objective.
  • Navigation experiments: The navigation task varies hallway length and hallway count to control difficulty, with the target distribution placing uniform mass at hallway ends.
  • Historical averaging: SMM achieves higher state entropy than other methods in Fetch, while historical averaging increases state coverage for previous exploration methods.
  • Non-uniform exploration: With a non-uniform target distribution, SMM matches the target better than SAC and performs on par with Count across four random seeds.
  • Computational cost: SMM's computational cost is comparable with prior work in average wall-clock time per epoch on Fetch.

C.5. SMM Ablation Study

The ablation study examines how SM4’s entropy components contribute to exploration and visualizes its latent mixture-state marginals across benchmark environments. Experiments use target distributions that are uniform or encode task-specific preferences.

  • Ablation results: SM4 relies heavily on both conditional state entropy and latent conditional action entropy in the Navigation task.The comparison removes either entropy term or both from the SM4 objective, and training performance is evaluated on Navigation with 3 halls of length 50.
  • Entropy diagnostics: The study tracks latent action entropy H[z | s] and latent state entropy H[s | z] per epoch while varying the number of mixture components.The former is associated with the discriminator and the latter with the density model.
  • Target distributions: The target distributions include uniform coverage of task-relevant states and shaped preferences for selected regions or low-cost configurations.Navigation uses uniform mass at hall ends, Fetch-Half favors block states on the table’s left side, and D’Claw uses uniform object angles with shaping for velocity and joint-position deviations.
  • Mixture components: For Fetch, SM4 visualizes log state marginals log ρπz(s) over block XY-coordinates separately for eight latent mixture components.Results are averaged over 1000 epochs, and the policy is trained with a uniform target distribution.
  • Experimental settings: The experiments cover Navigation, simulated Fetch, and simulated or hardware D’Claw environments with task-specific state, action, and episode parameters.The reported setup includes 3 halls in Navigation, a 28-dimensional Fetch state, and a 12-dimensional D’Claw observation space with 9 actions.

D.2. GAIL ablation

The GAIL ablation studies how discriminator inputs and synthetic target-state sampling affect comparison with exploration methods. The implementation also details density estimation, training schedules, hyperparameters, and evaluation metrics.

  • GAIL setup: GAIL is trained on synthetic states sampled from the target distribution because the exploration methods do not require expert demonstrations.The discriminator receives states only, excluding actions, to make the comparison more comparable.
  • State sampling: Fetch and D’Claw GAIL training uses 1e4 synthetic states, with task-specific sampling for object positions, gripper positions, and valve angles.Other state dimensions are sampled uniformly in [−10, 10], while Fetch compares alternative gripper-position distributions.
  • Discriminator ablation: Restricting the GAIL discriminator to reachable object, gripper, or angle dimensions addresses states sampled from the target distribution that may be unreachable from the initial state.Fetch uses object or gripper positions, while D’Claw uses object angle; the best ablation model is used for baseline comparisons.
  • Density estimation: The SMM density model estimates p(x) through VAE reconstruction likelihood under a unit-variance Gaussian.The resulting log-likelihood is represented by reconstruction mean-squared error plus a constant independent of x.
  • Training and evaluation: GAIL training is more unstable than the exploration baselines, so evaluation uses training and hyperparameter procedures specified separately from the final iterate.The broader experiments train algorithms for task-dependent step budgets and tune weights for entropy, count-based, pseudocount, and SMM loss components.
Loading 1906.05274v3…