Source-linked AI summary

Reinforcement Learning with Deep Energy-Based Policies

Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, Sergey Levine

arXiv:1702.08165v2cs.LGcs.AI

TL;DR

The paper addresses the difficulty of learning expressive stochastic policies in continuous domains, where general energy-based policy learning had previously been feasible mainly in tabular settings. It introduces soft Q-learning with an amortized Stein variational sampler for Boltzmann policies, and reports improved multimodal exploration and compositional skill initialization in simulated robots.

  • Problem

    Maximum-entropy policy learning needs expressive stochastic distributions for continuous domains, but existing approaches generally operate on tabular representations or face intractable sampling.

  • Method

    Soft Q-learning represents the policy as an energy-based model linked to a soft Q-function and trains a separate sampling network with Stein variational gradient descent.

  • Results

    The method captures complex multimodal behavior and improves exploration and initialization for new skills across simulated swimming and walking robot tasks.

  • Takeaways & Limitations

    Energy-based policies provide a tractable route to multimodal stochastic control in continuous domains and connect soft Q-learning with actor-critic approximate inference.

  • Takeaways & Limitations

    The method’s broader capability for complex behavioral repertoires and composability remains an avenue for future study.

Abstract

from arXiv · show

We propose a method for learning expressive energy-based policies for continuous states and actions, which has been feasible only in tabular domains before. We apply our method to learning maximum entropy policies, resulting into a new algorithm, called soft Q-learning, that expresses the optimal policy via a Boltzmann distribution. We use the recently proposed amortized Stein variational gradient descent to learn a stochastic sampling network that approximates samples from this distribution. The benefits of the proposed algorithm include improved exploration and compositionality that allows transferring skills between tasks, which we confirm in simulated experiments with swimming and walking robots. We also draw a connection to actor-critic methods, which can be viewed performing approximate inference on the corresponding energy-based model.

1. Introduction

The paper develops tractable energy-based stochastic policies for continuous domains, addressing the difficulty of maximum-entropy policy learning and sampling. Soft Q-learning supports multimodal exploration and policy compositionality in simulated control tasks.

  • Motivation: Maximum-entropy control learns the full range of low-cost behaviors rather than only one deterministic solution.This can support exploration, robustness, and transfer through multiple possible ways of performing a task.
  • Problem: Existing maximum-entropy methods generally rely on tabular representations that are difficult to apply to continuous or high-dimensional domains.
  • Approach: The paper formulates stochastic policies as energy-based models whose energy corresponds to a soft Q-function, then uses an approximate sampling network for action selection.This connects soft Q-learning with actor-critic methods and probabilistic inference.
  • Contribution: The proposed algorithm is a tractable, efficient method for optimizing arbitrary multimodal stochastic policies represented by energy-based models.
  • Evaluation: Experiments show improved exploration on multimodal reward landscapes and better initialization for learning new skills than random or conventionally pretrained policies.

2. Preliminaries

The paper defines maximum-entropy reinforcement learning for continuous state and action spaces and relates its optimal policy to soft value functions and energy-based models. This formulation motivates soft Bellman backups while exposing the need for approximation in continuous domains.

  • Problem setup: The reinforcement-learning problem is policy search in an infinite-horizon MDP with continuous states, continuous actions, stochastic transitions, and bounded rewards.
  • Maximum-entropy objective: Maximum-entropy reinforcement learning augments reward with an entropy term weighted by α, encouraging high-entropy policies at visited states.The discount factor γ is introduced for infinite-horizon objectives so expected rewards and entropies remain finite.
  • Maximum-entropy objective: Unlike greedy Boltzmann exploration, the maximum-entropy objective optimizes the entropy of the entire trajectory distribution, including future states.
  • Energy-based policies: Energy-based policies can represent complex multimodal behaviors, with a universal energy function capable of representing any action distribution.The energy is linked to the soft Q-function through E(s_t, a_t) = -1/α Q^soft(s_t, a_t).
  • Soft value functions: The optimal maximum-entropy policy is connected to the soft Q-function and soft value function through a Boltzmann form, with the value providing the log-partition function.As α approaches zero, the soft Bellman equation approaches the conventional hard maximum.

3. Training Expressive Energy-Based Models via Soft Q-Learning

Soft Q-learning turns maximum-entropy policy optimization into a practical continuous-domain algorithm by combining soft Q-function updates with an amortized stochastic sampler. Stein variational gradient descent trains the sampler to approximate the energy-based policy.

  • Soft Q-iteration: Soft Bellman backups provide the fixed-point update for learning the soft Q-function, but exact backups and policy sampling are generally intractable in continuous or large spaces.
  • Soft Q-learning: Importance sampling and stochastic optimization convert soft Q-iteration into an empirical procedure using sampled states and actions.The current policy can provide an unbiased estimate of the soft value, while uniform sampling may scale poorly in high dimensions.
  • Approximate sampling: A state-conditioned network maps Gaussian noise to action samples and is trained so its induced distribution approximates the target energy-based policy.The induced distribution is optimized using KL divergence and Stein variational gradient directions.
  • Actor-critic connection: The sampling network functions as the actor in an actor-critic interpretation of the method.
  • Algorithm: The algorithm alternates environment interaction with soft Q-function and sampling-network updates using replay-memory minibatches and delayed target values.Optimization uses empirical gradients and ADAM.

4. Related Work

Related work connects maximum-entropy policies to probabilistic inference, Boltzmann exploration, energy-based policy representations, and actor-critic algorithms. The paper’s implementation uses a soft Q-learning procedure with replay memory and stochastic optimization.

  • Maximum-entropy methods: Maximum-entropy policies arise from casting optimal control as probabilistic inference and have been studied in linear-quadratic, discrete, and entropy-regularized settings.
  • Algorithmic procedure: The algorithm samples actions through a stochastic network, interacts with the environment, stores transitions in replay memory, and updates parameters from minibatches.
  • Energy-based policies: Prior energy-based policy methods used models such as restricted Boltzmann machines, whereas this work targets general energy functions with approximate sampling.
  • Actor-critic connections: The soft Q-learning structure resembles actor-critic methods, with the actor approximating the maximizer or sampler associated with the critic’s Q-values.

5. Experiments

The experiments test whether soft Q-learning captures multimodal policies, improves exploration in tasks requiring multiple modes, and provides useful initialization for downstream tasks. Across point-mass, swimming, maze, and quadrupedal-robot settings, the method represents diverse behaviors, explores more reliably, and accelerates fine-tuning.

  • 5.1. Didactic Example: Multi-Goal Environment: The multi-goal point-mass policy follows complex Q-value landscapes and samples diverse trajectories reaching all four goals, whereas DDPG commits to one goal.At states with different Q-value geometries, the learned policy produces unimodal, convex, or bimodal action distributions aligned with the energy landscape.
  • 5.2. Learning Multi-Modal Policies for Exploration: Soft Q-learning explores both modes in the swimmer task before committing forward, while DDPG prematurely selects a direction.Only 80% of DDPG policies converge to forward motion; 20% choose the suboptimal backward mode.
  • 5.2. Learning Multi-Modal Policies for Exploration: In the quadrupedal maze, all soft Q-learning runs reach the successful threshold, compared with 60% of DDPG runs.The threshold corresponds to the minimum possible distance if the robot chooses the blocked upper passage, so success requires discovering the lower passage.
  • 5.3. Accelerating Training on Complex Tasks with Pretrained Maximum Entropy Policies: Maximum-entropy pretraining teaches a quadrupedal robot to locomote in arbitrary directions before fine-tuning it on hallway and maze tasks.The downstream environments include wide, narrow, and U-shaped hallways requiring increasingly task-specific directional or curved trajectories.
  • 5.3. Accelerating Training on Complex Tasks with Pretrained Maximum Entropy Policies: The broadly exploratory pretrained policy enables faster downstream learning than DDPG from random initialization, while deterministic pretraining provides a poor initialization.Deterministic pretraining selects an arbitrary but consistent direction, limiting its usefulness for task-specific fine-tuning.

6. Discussion and Future Work

The paper presents stochastic energy-based policies learned with approximate inference and demonstrates their use for multimodal exploration and compositional skill learning. It identifies broader composability and behavioral-repertoire representation as important directions for future work.

  • 6. Discussion and Future Work: The method learns stochastic energy-based policies using approximate inference via Stein variational gradient descent.The approach is described as a form of soft Q-learning that obtains complex multimodal policies.
  • 6. Discussion and Future Work: Experiments show that the method captures complex multimodal behavior from toy point-mass tasks to torque-controlled simulated walking and swimming robots.
  • 6. Discussion and Future Work: Stochastic policies support improved exploration for multimodal objectives and compositionality through pretrained general-purpose policies that can be efficiently fine-tuned.
  • 6. Discussion and Future Work: Future work should further study energy-based policies' ability to represent complex behavioral repertoires and their potential for composability.The paper suggests extending policy composition from tractable representations to complex, highly multimodal deep neural network models.

A. Policy Improvement Proofs

The appendix defines a discounted maximum-entropy objective and establishes policy-improvement and convergence properties for soft Q-learning. Under boundedness and regularity assumptions, repeated soft policy improvement converges to a unique energy-based optimal policy.

  • A. Policy Improvement Proofs: The discounted objective weights future rewards and policy entropy by the current policy's state-action visitation probability.The formulation retains entropy at future states rather than only maximizing current-step entropy.
  • A. Policy Improvement Proofs: The soft Q-value is the expected discounted sum of rewards and entropy under a policy.The trajectory begins from a specified state-action tuple, with entropy parameter α set to 1 for convenience.
  • A. Policy Improvement Proofs: Soft Q-learning defines policy improvement through a new policy that greedily maximizes one-step entropy plus value.
  • A. Policy Improvement Proofs: Soft policy iteration improves the soft Q-function monotonically, and under regularity conditions the policy converges to a Boltzmann energy-based form.The limiting policy satisfies π∞(a|s) proportional to exp(Qπ∞(s,a)).
  • A. Policy Improvement Proofs: The soft value-iteration operator is a contraction, implying a unique soft Q-value and unique optimal policy.The contraction bound is ||TQ1 − TQ2|| ≤ γ||Q1 − Q2||.

B. Connection between Policy Gradient and Q-Learning

The appendix connects entropy-regularized policy gradients with soft Q-learning by parameterizing policies as energy-based models. With an appropriate baseline and soft-Q estimate, the policy-gradient update recovers a soft Q-learning update, while additive energy shifts leave the Boltzmann policy unchanged.

  • B. Connection between Policy Gradient and Q-Learning: Entropy-regularized policy gradient can be viewed as performing soft Q-learning on the maximum-entropy objective.
  • B. Connection between Policy Gradient and Q-Learning: A stochastic policy is parameterized as an energy-based model whose energy function can represent any action distribution with a universal approximator.The corresponding partition function normalizes the exponential-form policy.
  • B. Connection between Policy Gradient and Q-Learning: Choosing the state-dependent baseline as the log partition function plus one yields a simplified policy-gradient form.
  • B. Connection between Policy Gradient and Q-Learning: The resulting gradient recovers the policy gradient from a Bellman-error formulation using an empirical soft Q-function estimate.The derivation uses an empirical soft advantage so the target is independent of the soft value.
  • B. Connection between Policy Gradient and Q-Learning: Approximating the soft Q-value up to an additive constant still produces the correct policy because the Boltzmann distribution is invariant to constant energy shifts.

C. Implementation

The implementation computes amortized SVGD updates by replacing expectations with empirical averages over sampled actions and states. The final policy-update direction averages per-state gradients over a minibatch.

  • C. Implementation: Amortized SVGD approximates its optimal descent direction with an empirical average over sampled actions.
  • C. Implementation: The samples used to evaluate the update may be the same as or different from the sampled actions generated by the network.
  • C. Implementation: Substituting the SVGD direction into the sampling-network update and differentiating produces an empirical gradient estimate.
  • C. Implementation: The overall policy-update direction is the average of per-state update directions over states drawn from a minibatch.

C.2. Computing the Density of Sampled Actions

The method computes soft values and policy updates from sampled actions, using a learned sampling network after an initial uniform-sampling phase. Experiments use fixed sampling and optimization settings, with entropy annealing during fine-tuning.

  • Soft values can be computed by sampling actions from a distribution proportional to the relevant quantity.
  • The sampling network directly produces actions from noise and the current state.
  • Uniform action sampling is used initially because the learned sampler's Jacobian may be singular before training.The method switches to the learned sampler later in training.
  • The experiments use a replay pool of one million samples, minibatches of 64, and 10,000 time steps per training iteration.Training begins after the replay pool contains at least 10,000 samples.
  • Soft Q-learning uses 32 action samples for policy updates and 50 additional samples for computing the soft value.The multi-goal experiment uses 100 policy-update samples instead.
  • Fine-tuning anneals the entropy coefficient log-linearly to 0.001 within 20 epochs and reduces the noise-sample count linearly to one.

D.3. Additional Results

The additional figures visualize exploration across swimming and quadrupedal locomotion settings. They show directional exploration and how entropy controls the breadth and diversity of pretrained trajectories.

  • Swimming policies explore positive and negative directions before committing to the better direction.The plot uses training iteration on x, distance on y, and marks the finish line in red.
  • Quadrupedal trajectories become more concentrated as the entropy coefficient decreases.Columns correspond to α = 10, 1, 0.1, and 0.01; rows represent different random seeds.
Loading 1702.08165v2…