Source-linked AI summary
Multi-task Deep Reinforcement Learning with PopArt
Matteo Hessel, Hubert Soyer, Lasse Espeholt, Wojciech Czarnecki, Simon Schmitt, Hado van Hasselt
TL;DR
The paper addresses how to train one reinforcement-learning agent across multiple tasks despite task-dependent reward scales, sparsity, and learning competence. It uses PopArt to make actor-critic updates invariant to these factors, achieving strong multi-task results on Atari-57 and DmLab-30.
Problem
Deep RL agents have typically been trained one task at a time, while multi-task updates can be affected by reward scale, reward sparsity, and agent competence.
Method
The paper applies PopArt normalisation to derive a scale-invariant actor-critic update for parallel multi-task learning.
Results
A single agent achieved a median normalised score of 110% on Atari-57 and a mean score of 72.8% on DmLab-30.
Takeaways & Limitations
PopArt-IMPALA surpassed median human performance across all 57 Atari games using a single trained agent.
Takeaways & Limitations
Parallel learning can be complicated to implement outside of simulation.
Abstract
from arXiv · showhide
The reinforcement learning community has made great strides in designing algorithms capable of exceeding human performance on specific tasks. These algorithms are mostly trained one task at the time, each new task requiring to train a brand new agent instance. This means the learning algorithm is general, but each solution is not; each agent can only solve the one task it was trained on. In this work, we study the problem of learning to master not one but multiple sequential-decision tasks at once. A general issue in multi-task learning is that a balance must be found between the needs of multiple tasks competing for the limited resources of a single learning system. Many learning algorithms can get distracted by certain tasks in the set of tasks to solve. Such tasks appear more salient to the learning process, for instance because of the density or magnitude of the in-task rewards. This causes the algorithm to focus on those salient tasks at the expense of generality. We propose to automatically adapt the contribution of each task to the agent's updates, so that all tasks have a similar impact on the learning dynamics. This resulted in state of the art performance on learning to play all games in a set of 57 diverse Atari games. Excitingly, our method learned a single trained policy - with a single set of weights - that exceeds median human performance. To our knowledge, this was the first time a single agent surpassed human-level performance on this multi-task domain. The same approach also demonstrated state of the art performance on a set of 30 tasks in the 3D reinforcement learning platform DeepMind Lab.
Introduction
Deep reinforcement learning has achieved strong results on individual tasks, but multi-task learning remains difficult because tasks compete for shared resources and produce differently scaled updates. The paper uses PopArt to adapt task contributions and reports strong performance across Atari and DeepMind Lab.
- Motivation: Deep RL successes have generally required training a new agent instance for each task.Prior results span Go, chess, continuous control, Atari, and 3D environments.
- Problem: Multi-task reinforcement learning includes several settings, and this work focuses on learning multiple tasks in parallel.The other settings include multiple predictions, continual learning, and distillation of task-specific experts.
- Prior work: 59.7% median human-normalised Atari score and 49.4% mean human-normalised DeepMind Lab score made IMPALA state of the art but remained below individually trained agents.These results illustrate the gap between parallel multi-task learning and single-task performance.
- Challenge: Shared representation capacity forces tasks to compete, while reward scale and sparsity can make some tasks more salient and dominate updates.Reward clipping can alter the objective and does not ensure equal update magnitudes when reward sparsity differs.
- Challenge: Reward sparsity and magnitude are non-stationary because the agent actively learns to maximize collected rewards, preventing a priori update normalization.The changing learning dynamics make fixed normalization difficult even with substantial domain knowledge.
- Contribution: PopArt normalisation produces an actor-critic update invariant to reward scale, sparsity, and agent competence.The approach achieved a 110% median normalised score on Atari-57 and a 72.8% mean score on DmLab-30.
Background
The paper formulates reinforcement learning with policies, value estimates, and discounted returns, then extends distributed actor-critic learning to multiple environments. IMPALA uses asynchronous actors and a learner, with importance sampling and clipped v-trace corrections for slightly off-policy data.
- Background: Reinforcement learning models an agent interacting with an environment to maximize discounted future rewards under uncertainty.The formalization uses an MDP with states, actions, transition dynamics, and a discount factor.
- Multi-task setting: Multi-task reinforcement learning requires one agent to master multiple environments with distinct dynamics, while tasks may share action spaces or physical rules.The environments can be represented as a larger MDP containing task-indexed states.
- Multi-task setting: The policy is not conditioned on task identity at test time, although the task index is used during training for value estimates.This yields a single general policy whose action selection depends only on the environment state.
- Actor-critic method: Actor-critic learning updates the policy with a REINFORCE-style advantage term and updates the value function toward truncated, bootstrapped returns.The value baseline reduces policy-gradient variance, while multi-step returns permit earlier updates.
- Actor-critic method: Both policy and value updates depend linearly on return scale, which varies with reward scale, sparsity, and agent competence.This dependence motivates scale-invariant normalization in the multi-task setting.
- IMPALA architecture: IMPALA distributes rollout generation across actors and batches their experience on a shared learner using a convolutional ResNet and LSTM network.Actors generate fixed-length rollouts, enqueue them, and receive updated parameters before subsequent rollouts.
- Off-policy corrections: Importance sampling corrects policy mismatch between acting and learning, while clipped ratios reduce the resulting variance through the v-trace return.The clipped correction preserves the validity of the off-policy return in the described setting.
Adaptive normalisation
PopArt normalises value estimates and updates using task-specific statistics, then extends this design to multi-task actor-critic learning. The resulting updates preserve unnormalised values while allowing a shared policy to learn across environments.
- Single-task normalisation: PopArt parameterises each value estimate as a scaled and shifted normalised prediction, using statistics μ and σ to rescale values and updates.The value estimate is represented as v_μ,σ,θ(s) = σ · n_θ(s) + μ, with n_θ produced by a final linear layer.
- Stable normalisation: Non-stationary normalised targets are handled by adjusting the final value-network layer so unnormalised value estimates remain unchanged when μ or σ changes.This preserves v_μ,σ,θ under updates from μ,σ to μ′,σ′.
- Multi-task extension: In the multi-task extension, a task-agnostic policy shares one action mapping while the value function uses N outputs and vector-valued normalisation statistics.Each environment D_i has a corresponding value head, while the policy operates over the shared action space.
- Multi-task extension: Each rollout updates only its task-specific value head, while the shared policy receives appropriately rescaled updates and final-layer parameters preserve unnormalised outputs.The value statistics and final-layer parameters are updated consistently for the active task.
- Task-agnostic execution: Task conditioning is restricted to value estimates during training, so the learned policy remains fully task agnostic during action selection.Values reduce policy-update variance but are not needed to choose actions at test time.
Experiments
The experiments evaluate PopArt-IMPALA across Atari-57 and DmLab-30, including clipped and unclipped reward settings. PopArt-IMPALA achieves strong multi-task performance, with results attributed to adaptive rescaling and improved data efficiency.
- Atari-57: 110% median human normalised score with reward clipping and 101% without clipping on Atari-57, achieved by one trained PopArt-IMPALA agent.Both results surpass median human performance across all 57 games.
- Atari-57: IMPALA reaches barely 60% with clipping and close to 0% without it, whereas PopArt-IMPALA performs almost equally in both regimes.This contrast indicates greater robustness to reward-scale and sparsity differences across games.
- Atari-57: MultiHead-IMPALA performs slightly worse than IMPALA with and without clipping, supporting adaptive rescaling rather than extra value-prediction capacity as the source of PopArt-IMPALA’s boost.MultiHead-IMPALA uses task-specific value predictions but no PopArt normalisation.
- Normalisation statistics: PopArt-IMPALA’s normalisation statistics vary by orders of magnitude across games and adapt during training alongside undiscounted episode returns.In crazy climber the shift exceeds 2500, while in bowling it remains below 15; qbert statistics also track performance instabilities.
- Extensions: Pixel control lets PopArt-IMPALA match vanilla PopArt-IMPALA at approximately 2B frames and use less than one-tenth the data needed to match the authors’ IMPALA baseline.The added auxiliary task has a small computational cost and therefore reduces training cost in this setting.
Discussion
The paper presents scale-invariant actor-critic learning as a way to improve parallel multi-task reinforcement learning, where sharing representations can support learning across sparse-reward tasks. It also identifies implementation challenges and possible combinations with other multi-task methods.
- The paper proposes a scale-invariant actor-critic algorithm for improved multi-task reinforcement learning performance.
- Sharing neural-network representations across tasks can help bootstrap learning when individual-task rewards are sparse.
- Parallel execution of task policies may support deep exploration and policy composition.
- The approach can be complicated to implement outside simulation because it requires parallel learning resources.
- PopArt’s adaptive normalisation could be combined with policy distillation and active task-distribution sampling.
Appendix
The appendix provides supplementary results, per-level score breakdowns, and training hyperparameters for the baseline and PopArt-IMPALA agents.
- The appendix reports additional DmLab-30 experiments and details supporting the main-text results.
- It includes score breakdowns by level for IMPALA and PopArt-IMPALA on Atari-57 and DmLab-30.
- It documents the hyperparameters used to train baseline agents and PopArt-IMPALA.
Hyper-parameter tuning
The appendix compares hyperparameter tuning with PBT populations of 8 and 24 instances on DmLab-30. Final performance is generally similar, although discrepancies appear for a less-performing IMPALA variant.
- PBT tuning in the DmLab-30 experiments used smaller populations than the original IMPALA study.
- The larger-population comparison reports one tuning experiment per agent because larger populations increase cost.
- Using 24 rather than 8 PBT instances produced very similar final performance for IMPALA and PopArt-IMPALA.
- Larger discrepancies between population sizes were observed for the less-performing IMPALA agent with a more limited action set.
Pixel Control
Adding pixel control improves data efficiency for both IMPALA and PopArt-IMPALA on DmLab-30, while PopArt retains a performance advantage. The auxiliary task is learned off-policy from actor-generated experience.
- Pixel control increases performance for both PopArt-IMPALA and the IMPALA baseline.
- Approximately 10 points separate Pixel-PopArt-IMPALA from Pixel-IMPALA in median human normalised score.
- Pixel control defines cell-wise pseudo-rewards from absolute consecutive-frame pixel-intensity changes and trains multi-step Q-values for each cell.
- The pixel-control head is evaluated only by the learner, while actors continue executing the main policy.
- ∼2B frames were sufficient for pixel-control PopArt-IMPALA to match vanilla PopArt-IMPALA’s final performance.
Atari-57 Score breakdown
The Atari-57 breakdown evaluates one frozen policy across all 57 games and compares PopArt-IMPALA with IMPALA, including PopArt-IMPALA conditions with and without reward clipping.
- Evaluation protocol: The final trained agent is evaluated with a frozen policy on each Atari level for 200 episodes.The same trained policy is used across all levels without being told which task is being evaluated.
- Evaluation protocol: The evaluation uses a single policy across all levels rather than separate task-specific policies.The policy is not provided information about the task during evaluation.
- Reported comparisons: Figure 8 reports human normalised scores for IMPALA and PopArt-IMPALA across the 57 Atari games.PopArt-IMPALA scores are shown both with and without reward clipping.
DeepMind Lab action discretisation
DeepMind Lab uses a continuous seven-dimensional native action space, but the experiments follow the platform’s typical coarse discretisation approach. The authors introduce a modified discretisation to unlock levels that the original IMPALA discretisation cannot solve.
- Native action space: DeepMind Lab’s native action space has seven continuous dimensions covering movement, rotation, tagging, crouching, and jumping.The dimensions include horizontal and vertical rotation, strafing, forward and backward movement, tagging, crouching, and jumping.
- Discretisation approach: The experiments use a coarse discretisation of the continuous action space, following previous work on DeepMind Lab.This matches the approach typically used on the platform.
- Discretisation variants: The authors introduce a new action discretisation alongside the original IMPALA discretisation.The new version is intended to unlock DmLab-30 levels unsolvable under the original discretisation.
- Discretisation variants: Table 1 documents the action discretisation used by the original IMPALA agent.The reference is the discretisation from Espeholt et al.
- Discretisation variants: Table 2 documents the action discretisation used by the improved IMPALA baseline and PopArt-IMPALA.These agents share the listed action-space discretisation.
Fixed Hyperparameters
The experiments hold PopArt-specific hyperparameters fixed and use shared baseline settings to support fair comparisons. Preprocessing, network, optimization, and population-based training configurations are specified in the accompanying tables.
- PopArt settings: PopArt-specific hyperparameters are held fixed during training and only lightly tuned.A lower bound is used to avoid numerical issues when rewards are extremely sparse.
- Preprocessing: DeepMind Lab preprocessing uses 72 × 96 observations, four action repeats, and optimistic asymmetric reward rescaling.The reward rescaling method was introduced in Espeholt et al. for exploration.
- Preprocessing: Atari experiments use standard Atari preprocessing and report conditions both with and without reward clipping.The preprocessing follows settings established as standard practice since DQN.
- Shared training configuration: Population-based training tunes shared agent hyperparameters, while network and other settings are specified to match prior IMPALA work.Table 9 states that tuned hyperparameters are the same for all compared baseline agents to ensure fair comparisons.
- Hyperparameter search: The optimization search includes entropy cost and learning rate ranges, RMSProp epsilon choices, and a maximum gradient norm distribution.The listed distributions include log-uniform, categorical, and uniform searches.