Source-linked AI summary
Reinforcement Learning with Unsupervised Auxiliary Tasks
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, Koray Kavukcuoglu
TL;DR
Deep reinforcement learning typically relies on extrinsic reward, although environments provide many other training signals and rewards may be sparse. The paper introduces UNREAL, which jointly learns auxiliary pseudo-reward control and reward-prediction tasks through a shared representation. UNREAL more than doubled prior state-of-the-art results on Labyrinth, reaching over 87% of human scores, while also improving learning speed and robustness across Atari.
Problem
Deep reinforcement learning usually maximizes extrinsic reward despite sparse rewards and the availability of many other training signals.
Method
UNREAL augments A3C with auxiliary pixel-control, network-feature, and reward-prediction tasks that share the agent’s representation and use replayed experience.
Results
UNREAL more than doubled prior state-of-the-art results on Labyrinth, averaging over 87% of human scores, and significantly improved A3C learning speed and robustness across 57 Atari games.
Takeaways & Limitations
Auxiliary control and reward-prediction tasks can substantially improve reinforcement-learning data efficiency, robustness, and final performance across challenging 3D and Atari environments.
Takeaways & Limitations
Reward prediction uses a rebalanced training distribution to oversample rare rewarding events, while its predictor is restricted to shaping features rather than the policy or value function.
Abstract
from arXiv · showhide
Deep reinforcement learning agents have achieved state-of-the-art results by directly maximising cumulative reward. However, environments contain a much wider variety of possible training signals. In this paper, we introduce an agent that also maximises many other pseudo-reward functions simultaneously by reinforcement learning. All of these tasks share a common representation that, like unsupervised learning, continues to develop in the absence of extrinsic rewards. We also introduce a novel mechanism for focusing this representation upon extrinsic rewards, so that learning can rapidly adapt to the most relevant aspects of the actual task. Our agent significantly outperforms the previous state-of-the-art on Atari, averaging 880\% expert human performance, and a challenging suite of first-person, three-dimensional \emph{Labyrinth} tasks leading to a mean speedup in learning of 10$\times$ and averaging 87\% expert human performance on Labyrinth.
1 RELATED WORK
Prior work used pseudo-rewards, factored value representations, environment models, and auxiliary predictions for reinforcement learning. UNREAL instead uses pseudo-reward policies as auxiliary objectives for representation learning rather than as direct temporal abstractions.
- Pseudo-rewards: UNREAL’s auxiliary pseudo-reward policies develop a shared representation instead of directly serving as temporal abstractions.The learned policies are used solely as auxiliary objectives, not directly for planning or acting.
- Value functions: Horde learned many pseudo-reward value functions, but trained each separately with distinct weights rather than using them for representation learning.
- Factored representations: UVFA combines state features with pseudo-reward embeddings to represent a continuous set of optimal value functions.A pre-trained UVFA representation was transferred to novel pseudo-rewards in a simple task.
- Factored representations: Successor representations factor expected value functions for a fixed policy and have supported transfer across pseudo-rewards or reward scales.
- Auxiliary prediction: Environment models and auxiliary predictions had been explored, but model-based auxiliary tasks had not yet been demonstrated in rich visual environments.Related 3D work predicted emulator features or depth as auxiliary targets.
2 BACKGROUND
The paper builds on standard reinforcement learning, where agents receive observations and rewards, choose actions, and estimate returns, values, and policies. A3C uses asynchronous agents with an LSTM to jointly approximate policy and value from experience history.
- Reinforcement learning setting: At each discrete time step, an agent receives an observation and reward, produces an action, and forms a state from its experience history.
- Returns and values: The n-step return is a discounted reward sum, while Vπ and Qπ represent expected returns from a state and state-action pair, respectively.
- Value-based learning: Value-based methods approximate Q(s, a; θ) and update parameters by minimizing mean-squared error on n-step lookahead targets.
- Policy gradients: Policy-gradient methods maximize expected reward by adjusting the policy, using approximations of the true value functions in practice.
- A3C: A3C jointly approximates policy and value with an LSTM over the full experience history, while parallel environment instances accelerate and stabilize learning.
3 AUXILIARY TASKS FOR REINFORCEMENT LEARNING
The auxiliary-task framework augments A3C with pseudo-reward control, reward-focused prediction, and replay-based updates. These tasks share representations with the base agent so learning can continue from experience beyond extrinsic rewards and emphasize sparse rewarding events.
- Overview: The complete UNREAL agent combines auxiliary control tasks, reward-focused auxiliary tasks, and replay mechanisms with A3C.The framework targets faster training, robust learning, and higher final performance.
- 3.1 AUXILIARY CONTROL TASKS: Auxiliary control tasks define additional pseudo-reward functions over states and actions, with shared parameters balancing base-task and auxiliary performance.The underlying state includes observation and reward history plus network hidden-unit activations.
- 3.1 AUXILIARY CONTROL TASKS: UNREAL uses off-policy value-based reinforcement learning to maximize many pseudo-rewards simultaneously from a single experience stream.
- 3.1 AUXILIARY CONTROL TASKS: Pixel-control tasks maximize changes in cells of an image grid, while network-feature tasks maximize activations of hidden units.
- 3.1 AUXILIARY CONTROL TASKS: The pixel-control head predicts action-conditioned discounted changes for each image cell using a spatial deconvolutional network.
- 3.2 AUXILIARY REWARD TASKS: Reward prediction addresses sparse rewards by predicting immediate reward from historical observation sequences, while its predictor shapes features without changing the policy or value function.Sampling rebalances zero and non-zero rewards, with non-zero rewards represented at probability 0.5.
- Replay: Experience replay supplies extra value and auxiliary-control updates and replays rewarding and nonrewarding samples equally to oversample rare rewarding states.
4 EXPERIMENTS
Experiments evaluate UNREAL against A3C and ablated auxiliary-task variants on Labyrinth and Atari. Across Labyrinth, auxiliary tasks improve final performance, learning speed, data efficiency, and hyperparameter robustness; control-based self-supervision outperforms reconstruction and prediction baselines.
- Experimental setup: UNREAL combines auxiliary control and reward tasks with an A3C CNN-LSTM baseline, using the same training setup and replay mechanism.The experiments compare the full agent with ablated variants and A3C.
- Labyrinth evaluation: Labyrinth evaluates agents on 13 first-person 3D levels spanning fruit gathering, navigation, and laser-tag abilities.The platform has rich dynamic visuals, realistic physics, fine-grained 3D pointing, and continuous motion.
- Labyrinth results: 87% human-normalised mean performance more than doubles A3C’s 54% on Labyrinth, while every auxiliary component improves upon the baseline.UNREAL reaches 116% on lt hallway slope and 100% on nav maze random goal 02.
- Labyrinth results: 10× mean speedup across Labyrinth levels reduces the data required to reach A3C’s final performance to less than 10%.The median speedup is 11× and the maximum reported speedup is 18× on nav maze random goal 02.
- Auxiliary-task comparison: Control-based self-supervision performs best: predicting input changes beats reconstruction, while learning to control pixel changes beats prediction.Reconstruction improves initial learning but worsens final scores in the reported comparison.
- Atari results: UNREAL attains 880% mean and 250% median human-normalised performance across 57 Atari games and is more robust to hyperparameter settings than A3C.The evaluation varies learning rate and entropy cost across 50 settings per game.
5 CONCLUSION
The conclusion reports that auxiliary control and reward-prediction tasks substantially improve deep reinforcement learning. UNREAL improves data efficiency and hyperparameter robustness, with strong results on Labyrinth and Atari.
- Auxiliary control and reward-prediction tasks drastically improve data efficiency and robustness to hyperparameter settings.
- UNREAL more than doubles prior state-of-the-art results on 3D Labyrinth, averaging over 87% of human scores.
- The same architecture significantly improves A3C’s learning speed and hyperparameter robustness across 57 Atari games.
A ATARI GAMES
Figure 6 presents learning curves for three Atari games while showing variation across random seeds and hyperparameter settings. Population means and the best final-performing agents are distinguished from individual runs.
- Learning curves cover three example Atari games and distinguish individual agents, population means, and best final-performing agents.Semi-transparent lines represent different seeds and hyperparameters; bold lines are population means and dotted lines are best agents.
B IMPLEMENTATION DETAILS
The implementation combines CNN/LSTM-based agents with auxiliary control and reward-prediction architectures, replay, and asynchronous optimization. These components specify how auxiliary tasks are trained and integrated with the base agent.
- Base architecture: The agent processes 84 × 84 RGB images through two convolutional layers, a 256-unit fully connected layer, and an LSTM.The convolutional layers use 16 8 × 8 filters with stride 4 and 32 4 × 4 filters with stride 2.
- Auxiliary control: Pixel control divides the central 80 × 80 input crop into a 20 × 20 grid of 4 × 4 cells.Each cell’s instantaneous reward is its average absolute frame-to-frame pixel-and-channel difference.
- Auxiliary control: Feature control learns to control a 32 × 9 × 9 spatial feature map using a deconvolutional network to produce auxiliary values.Its architecture is analogous to pixel control and exploits spatial structure in the learned representation.
- Reward prediction: Reward prediction encodes three observations with separate CNN instances, concatenates them, and classifies the next reward as positive, negative, or zero.The classifier uses a 128-unit ReLU layer followed by a three-class linear classifier and softmax.
- Replay: Auxiliary tasks run every 20 environment steps after the replay buffer fills, using its most recent 2k observations, actions, and rewards.This schedule corresponds to each update of the base A3C agent.
- Optimization: Training uses 32 asynchronous threads with shared RMSprop, while learning rates, entropy costs, and task weights are sampled from log-uniform ranges.The pixel-control task-weight range is 0.01–0.1 for Labyrinth and 0.0001–0.01 for Atari.
C LABYRINTH LEVELS
The Labyrinth-level figures depict the environments from complementary viewpoints: top-down layouts and egocentric observations. The layout figure distinguishes representative fixed mazes from episodes with newly randomized mazes.
- Level layouts: Top-down renderings show each Labyrinth level’s layout, with navigation mazes represented by one example maze and all-random episodes by newly generated mazes.The figure provides structural context for comparing the level layouts.
- Visual observations: Egocentric example images show the visual input presented to the agent for each Labyrinth level.Compare the first-person appearance of the different levels alongside their top-down layouts.