Source-linked AI summary
Scalable trust-region method for deep reinforcement learning using Kronecker-factored approximation
Yuhuai Wu, Elman Mansimov, Shun Liao, Roger Grosse, Jimmy Ba
TL;DR
Deep RL needs more sample-efficient optimization because interaction and simulation can be costly, while existing natural-gradient approaches are difficult to scale. The paper introduces ACKTR, which applies K-FAC-based trust-region natural-gradient updates to both actor and critic. Across Atari and MuJoCo, ACKTR improves sample efficiency and final performance relative to A2C and TRPO, with average sample-efficiency gains of 2- to 3-fold.
Problem
Deep RL requires more sample-efficient optimization because real-world interaction and simulated-environment computation are costly, while scalable natural policy-gradient methods remain unavailable.
Method
ACKTR uses a Kronecker-factored approximation to compute trust-region natural-gradient updates for both the actor and critic.
Results
ACKTR substantially improves sample efficiency and final performance on Atari and MuJoCo compared with A2C and TRPO.
Takeaways & Limitations
ACKTR provides a scalable trust-region natural-gradient algorithm for actor-critic methods, with average sample-efficiency improvements of 2- to 3-fold.
Takeaways & Limitations
The K-FAC approximation assumes that activation and backpropagated-derivative second-order statistics are uncorrelated.
Abstract
from arXiv · showhide
In this work, we propose to apply trust region optimization to deep reinforcement learning using a recently proposed Kronecker-factored approximation to the curvature. We extend the framework of natural policy gradient and propose to optimize both the actor and the critic using Kronecker-factored approximate curvature (K-FAC) with trust region; hence we call our method Actor Critic using Kronecker-Factored Trust Region (ACKTR). To the best of our knowledge, this is the first scalable trust region natural gradient method for actor-critic methods. It is also a method that learns non-trivial tasks in continuous control as well as discrete control policies directly from raw pixel inputs. We tested our approach across discrete domains in Atari games as well as continuous domains in the MuJoCo environment. With the proposed methods, we are able to achieve higher rewards and a 2- to 3-fold improvement in sample efficiency on average, compared to previous state-of-the-art on-policy actor-critic methods. Code is available at https://github.com/openai/baselines
1 Introduction
Deep RL can solve challenging control tasks but relies on inefficient first-order optimization, making sample efficiency a central concern. The paper motivates K-FAC-based trust-region optimization as a scalable alternative for actor-critic learning.
- Motivation: Deep RL uses neural networks for control policies, yet commonly trains them with SGD and related first-order methods.These methods explore weight space inefficiently and can take days to master continuous and discrete control tasks.
- Motivation: Sample efficiency matters because real-world robotic interaction and simulated-environment computation are typically costly.Advanced optimization techniques for gradient updates are presented as one way to reduce the required sample size.
- Limitations of existing methods: TRPO avoids explicit Fisher-matrix inversion but requires many conjugate-gradient steps and large batches for curvature estimation.These requirements make TRPO impractical for large models and contribute to sample inefficiency.
- Proposed direction: K-FAC approximates natural-gradient curvature, makes updates comparable in cost to SGD, and uses running curvature averages that support small batches.These properties suggest that applying K-FAC to policy optimization could improve deep-RL sample efficiency.
- Proposed direction: ACKTR applies Kronecker-factored trust-region optimization to actor-critic methods and efficiently inverts the gradient covariance approximation.The reported per-update computation cost is 10% to 25% higher than SGD-based methods.
2 Background
This section introduces policy-gradient reinforcement learning, actor-critic value estimation, and natural-gradient updates. K-FAC approximates curvature to make these updates tractable for large neural networks, with substantial sample-efficiency gains illustrated in Atlantis.
- Reinforcement learning and actor-critic methods: Policy-gradient methods parameterize a policy and update its parameters to maximize expected discounted return.The agent selects actions from πθ(a|st), receives rewards, and transitions according to the environment dynamics.
- Reinforcement learning and actor-critic methods: The advantage function measures an action’s relative value at a state and is used to weight policy-gradient updates.The paper follows A3C by defining the advantage with k-step returns and function approximation.
- Reinforcement learning and actor-critic methods: The critic estimates expected cumulative rewards from a state and is trained with temporal-difference updates minimizing squared prediction error.The target uses bootstrapped k-step returns.
- Natural gradient using Kronecker-factored approximation: Natural gradient uses the Fisher information matrix as a parameterization-independent metric, but exact Fisher inversion is impractical for million-parameter networks.The natural-gradient step corresponds to a constrained steepest-descent update under a Fisher-based norm.
- Natural gradient using Kronecker-factored approximation: K-FAC approximates each Fisher block with Kronecker factors, reducing required computations to matrices comparable in size to the layer weights.The approximation assumes activation and backpropagated-derivative second-order statistics are uncorrelated.
3 Methods
ACKTR applies K-FAC natural-gradient updates to both actor and critic networks within a trust-region framework. The method defines reinforcement-learning curvature from policy trajectories and uses Gauss-Newton curvature for the critic.
- Natural gradient in actor-critic: ACKTR applies Kronecker-factored natural-gradient updates to both the actor and the critic.The paper presents this as a scalable and sample-efficient natural-gradient algorithm for actor-critic methods.
- Natural gradient in actor-critic: The actor’s Fisher metric is defined from the policy’s action distribution and averaged over trajectory distributions collected during training.The trajectory expectation is approximated using sampled experience.
- Natural gradient in actor-critic: The critic uses a Gauss-Newton curvature matrix, which is equivalent to the Fisher matrix under a Gaussian observation model.This equivalence permits K-FAC to be applied to value-function optimization.
- Natural gradient in actor-critic: Shared actor-critic representations use a joint policy-and-value distribution whose output distributions are assumed independent.Disjoint networks can instead receive separate K-FAC updates.
- Natural gradient in actor-critic: ACKTR uses factorized Tikhonov damping and asynchronous computation of curvature statistics and matrix inverses to reduce computation time.These implementation choices follow earlier K-FAC approaches.
- Natural gradient in actor-critic: Trust-region scaling limits policy-distribution changes in KL divergence because unconstrained natural-gradient steps can cause premature convergence to near-deterministic policies.The method tunes trust-region radius and learning-rate parameters for the actor and critic settings.
4 Related work
Prior work established natural-gradient policy optimization and explored computationally cheaper trust-region variants. ACKTR addresses scalability and complements approaches based on replay or auxiliary objectives.
- Related work: Earlier actor-critic natural-gradient work used natural policy gradients for the actor and LSTD for the critic, but computational costs limited scalability.The main challenges were storing and inverting the Fisher matrix, leading prior methods toward compatible linear function approximators.
- Related work: TRPO reduced explicit Fisher inversion through conjugate-gradient Fisher-vector products, but repeated products hinder scaling to larger architectures.Other variants used linearized KL constraints or soft KL penalties instead of the hard trust-region constraint.
- Related work: Experience replay and auxiliary-objective actor-critic methods improve sample efficiency through mechanisms orthogonal to ACKTR.The paper notes that these approaches could potentially be combined with ACKTR.
5 Experiments
Experiments compare ACKTR with A2C and TRPO across Atari and MuJoCo, including state and pixel observations, sample efficiency, critic norms, computational cost, and batch-size scaling.
- 5.1 Discrete control: ACKTR significantly outperformed A2C in sample efficiency across all six Atari games, while TRPO learned only Seaquest and Pong within 10 million timesteps.TRPO also performed worse than A2C in sample efficiency on those games.
- 5.1 Discrete control: 2.7, 3.5, 5.3, and 3.0 times more episodes were required by A2C than ACKTR to reach human performance on Beamrider, Breakout, Pong, and Q-bert, respectively.ACKTR also achieved 26%, 35%, and 67% larger episode rewards than A2C on Breakout, Q-bert, and Beamrider.
- 5.2 Continuous control: 16.4 times more sample efficiency was achieved by ACKTR than TRPO on Ant, while TRPO was 4.1 times more sample efficient than ACKTR on Swimmer.ACKTR reached the specified threshold faster on all other listed MuJoCo tasks, and mean reward scores were generally comparable.
- 5.2 Continuous control: 1.6, 2.8, and 1.7 times greater final rewards were achieved by ACKTR than A2C on pixel-based Reacher, HalfCheetah, and Walker2d, respectively, after 40 million timesteps.The experiments used image observations without low-dimensional state input.
- 5.3 A better norm for critic optimization?: A Gauss-Newton norm produced more substantial sample-efficiency and final-reward improvements than a Euclidean norm for critic optimization.The Gauss-Newton norm also stabilized training by reducing variance across random seeds.
- 5.4 How does ACKTR compare with A2C in wall-clock time?: 25% was the maximum increase in computing time per timestep for ACKTR over A2C across six Atari games and eight MuJoCo environments.With larger batch sizes, ACKTR performed as well as with smaller batches, whereas A2C experienced significant sample-efficiency degradation.
6 Conclusion
The paper proposes ACKTR, a scalable trust-region method using K-FAC natural-gradient approximations for both actor and critic, and reports improved sample efficiency across Atari and MuJoCo.
- ACKTR applies K-FAC to approximate natural-gradient updates for actor-critic methods within trust-region optimization.The method optimizes both the actor and critic using natural-gradient updates.
- 2- to 3-fold improvements in sample efficiency were observed on average against A2C and TRPO.
- ACKTR trained several non-trivial continuous-control tasks directly from raw pixel observations.
- The results suggest extending Kronecker-factored natural-gradient approximations to other reinforcement-learning algorithms.
A Experimental details
The experiments use distinct network architectures and training protocols for Atari, low-dimensional state inputs, and pixel-based inputs, with Table 4 reporting raw game scores.
- Atari experiments: Atari experiments use a shared policy-value network with three convolutional layers followed by a fully connected component.The passage specifies filter sizes and strides for the convolutional layers.
- Atari experiments: Atari agents are trained for 50 million time steps, or 200 million frames, using a linear learning-rate schedule and entropy regularization.The maximum learning rate is tuned on Breakout with δ set to 0.001, then reused across Atari experiments.
- State-space experiments: Low-dimensional state experiments use separate two-layer networks with 64 hidden units per layer for policy and value functions.Tanh is used in the policy network and ELU in the value network, except at the output layer.
- Pixel-based experiments: Pixel-based experiments use 42 × 42 RGB images with the previous frame and separate policy and value networks.The networks contain two convolutional layers followed by a 256-unit fully connected layer.
- Reported metrics: Table 4 reports raw scores across all games, starting each episode with 30 no-op actions.
B Results on the remaining Atari games
On the remaining Atari games, ACKTR is evaluated with a common score protocol and a single random seed, achieving results comparable to reported Q-learning methods while training faster.
- ACKTR scores are the mean of the last 100 episode rewards after 50 million time steps, with each episode beginning with 30 no-op actions.
- ACKTR is compared with Q-learning results because matching A3C and A2C scores were unavailable.
- Despite evaluation on only a subset of games with one random seed and untuned hyperparameters, ACKTR results are reported as on par with Q-learning methods.
- ACKTR training takes 16 hours on a modern GPU, compared with Q-learning methods that usually take days.
C MuJoCo results with comparisons to OpenAI baselines
MuJoCo comparisons use OpenAI-provided A2C and TRPO baselines under closely matched protocols; performance varies by task, while ACKTR’s results are shown across seven environments.
- ACKTR is compared with OpenAI implementations of A2C and TRPO using closely matched experimental protocols.The baselines use two-layer architectures with 64 hidden units per layer and batch sizes of 2500 for A2C and 25000 for TRPO.
- OpenAI A2C performs better on Hopper, InvertedPendulum, Swimmer, and Walker2d, but worse on Reacher and HalfCheetah, relative to the authors’ A2C implementation.
- OpenAI TRPO performs worse on Hopper and achieves the same performance on the remaining tasks compared with the authors’ TRPO training.
- Figure 6 compares performance on seven MuJoCo environments after 1 million timesteps, with shaded regions representing standard deviation over three random seeds.One timestep equals four frames.
D Adaptive Gauss-Newton?
The section tests whether adaptive Gauss-Newton improves critic training over vanilla Gauss-Newton across Atari and MuJoCo tasks. Adaptive Gauss-Newton estimates the Bellman-error standard deviation through the critic output distribution.
- Adaptive Gauss-Newton keeps an estimate of the Bellman-error standard deviation as the standard deviation of the critic output distribution.
- The comparison uses adaptive and vanilla Gauss-Newton, both defined in Section 3.1.
- The experiments cover all six standard Atari games and eight MuJoCo tasks, with results reported in Figures 7 and 8.
E How well does the Kronecker-factored quadratic approximation match the exact KL?
The experiments assess whether ACKTR’s Kronecker-factored quadratic model tracks exact KL changes during trust-region optimization. Across the tested settings, exact KL stays close to the trust-region radius.
- The study measures exact KL changes during ACKTR training to indirectly assess the accuracy of its Kronecker-factored curvature approximation.
- The comparison uses HalfCheetah and Reacher, with approximated and exact KL values shown in Figure 9.
- Figure 7 compares adaptive and vanilla Gauss-Newton critics on six Atari environments over 10 million timesteps.
- Figure 8 compares adaptive and vanilla Gauss-Newton critics on eight MuJoCo environments over 1 million timesteps.
- The exact KL is close to the trust-region radius, indicating effective trust-region optimization through the Kronecker-factored approximation.
- Figure 9 plots exact KL changes during trust-region optimization using ACKTR.