Distral: Robust Multitask Reinforcement Learning

Yee Whye Teh; Victor Bapst; Wojciech Marian Czarnecki; John Quan; James Kirkpatrick; Raia Hadsell; Nicolas Heess; Razvan Pascanu

Source-linked AI summary

Distral: Robust Multitask Reinforcement Learning

Yee Whye Teh, Victor Bapst, Wojciech Marian Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, Razvan Pascanu

arXiv:1707.04175v1Jul 13, 2017cs.LGstat.ML

TL;DR

Deep reinforcement learning is data-inefficient and unstable in complex environments, while multitask learning can suffer from gradient interference and task dominance. Distral shares a distilled policy rather than parameters, constrains task policies toward it, and learns the shared policy by distillation; the resulting algorithms transfer efficiently and outperform related methods while being more robust and stable.

Problem
Deep reinforcement learning is data-inefficient and unstable in complex environments, and multitask training can worsen individual-task learning through gradient interference or task dominance.
Method
Distral shares a distilled policy capturing common behavior, constrains each task policy toward it, and trains the shared policy as the centroid of task policies.
Results
Distral learns quicker, achieves better final performance, and is more stable and robust to hyperparameter settings than standard shared-parameter multitask reinforcement learning.
Takeaways & Limitations
Distral supports effective sharing of behavioral structure across tasks and is identified as a promising solution for multitask reinforcement learning in complex 3D environments.
Takeaways & Limitations
Future work is needed to extend Distral to multiple distilled policies or latent variables, continual learning, auxiliary losses, and adaptive KL and entropy costs.

Abstract

from arXiv · show

Most deep reinforcement learning algorithms are data inefficient in complex and rich environments, limiting their applicability to many scenarios. One direction for improving data efficiency is multitask learning with shared neural network parameters, where efficiency may be improved through transfer across related tasks. In practice, however, this is not usually observed, because gradients from different tasks can interfere negatively, making learning unstable and sometimes even less data efficient. Another issue is the different reward schemes between tasks, which can easily lead to one task dominating the learning of a shared model. We propose a new approach for joint training of multiple tasks, which we refer to as Distral (Distill & transfer learning). Instead of sharing parameters between the different workers, we propose to share a "distilled" policy that captures common behaviour across tasks. Each worker is trained to solve its own task while constrained to stay close to the shared policy, while the shared policy is trained by distillation to be the centroid of all task policies. Both aspects of the learning process are derived by optimizing a joint objective function. We show that our approach supports efficient transfer on complex 3D environments, outperforming several related methods. Moreover, the proposed learning process is more robust and more stable---attributes that are critical in deep reinforcement learning.

1 Introduction

Deep reinforcement learning scales to complex environments but remains data-hungry and unstable. Multitask transfer is expected to improve efficiency, yet task interference and dominance often make individual-task learning harder; Distral is presented as a more effective alternative.

Motivation: Deep reinforcement learning uses deep neural networks to scale reinforcement learning to complex and rich environments, but effective policies require substantial time and data.Learning can also remain unstable despite changes to algorithms and optimizers.
Motivation: Multitask learning targets related problems because shared structure could reduce data requirements or improve asymptotic performance.This direction is motivated by the computational cost of increasingly complex tasks and the need for robust algorithms without extensive task-specific tuning.
Problem: In practice, training multiple tasks can hurt individual-task performance because gradients may interfere as noise or one task may dominate the others.These effects create additional challenges for existing multitask and transfer methods.
Contribution: Distral learns faster, reaches better asymptotic performance, and is more robust and stable than multitask A3C baselines across visually rich 3D tasks.The paper evaluates the approach on grid-world illustrations and DeepMind Lab comparisons against A3C.

2 Distral: Distill and Transfer Learning

Distral links task-specific policies through a learned distilled policy, optimizing task returns with KL and entropy regularization. The framework supports task specialization, transfer, and controlled exploration while making the distilled policy a centroid of task policies.

Mathematical framework: Distral optimizes task-specific policies and a shared distilled policy using expected returns, KL regularization, and entropy regularization.The KL term keeps each task policy close to the distilled policy, while entropy encourages exploration.
Soft Q-Learning and Distillation: With the distilled policy fixed, optimization decomposes into separate entropy-regularized expected-return problems with redefined rewards.This alternating maximization procedure optimizes task policies given the distilled policy and then updates the distilled policy given the task policies.
Soft Q-Learning and Distillation: The distilled policy is learned rather than hand-picked and becomes the centroid of the task policies’ state-action distributions.Its update corresponds to fitting the policy to a mixture of discounted state-action distributions from the tasks.
Soft Q-Learning and Distillation: Entropy regularization prevents task policies from becoming greedy and controls exploration through the entropy coefficient.Without sufficient exploration, an easy task can dominate the distilled and task policies before harder tasks produce rewards.
Policy Gradient and a Better Parameterization: A two-column parameterization combines the distilled policy with task-specific soft advantages, making common behavior immediately available for task specialization.The architecture uses one column for the distilled policy and another for the adjustment required by task i.
Policy Gradient and a Better Parameterization: KL regularization makes the distilled policy a policy-space centroid and helps transfer information quickly across tasks and to new tasks.The authors distinguish this policy-space centroid from parameter-space averaging methods and associate the KL term with increased stability.

3 Algorithms

Distral spans algorithmic choices in regularization and network architecture, balancing transfer against exploration and evaluating four variants alongside A3C baselines.

KL divergence vs entropy regularization: α = 0 removes task coupling, α = 1 enables transfer but may prematurely stop exploration, and 0 < α < 1 balances both effects.The purely entropy-regularized objective does not transfer across tasks, whereas mixed regularization retains coupling while encouraging exploration.
Algorithmic instantiations: Table 1 organizes algorithms by architecture columns and relative KL-versus-entropy regularization rows, with the first row containing A3C baselines without KL loss.The table contains seven evaluated algorithms.
Separate vs two-column parameterization: Using the distilled policy in a two-column parameterization can make transferred behavior immediately available, but overly rapid transfer may hinder task-specific exploration.Separate parameterization avoids immediate architectural access to the distilled policy, while two-column parameterization can accelerate transfer.

4 Experiments

Distral is evaluated on grid-world and challenging partially observed 3D tasks, where it generally learns faster, achieves stronger final performance, and is more stable than multitask A3C baselines.

Two-room grid world: Distral learns significantly faster than separate single-task agents in the two-room grid world while learning a corridor policy that preserves travel direction.The distilled policy guides movement between rooms and can help exploration on new tasks.
Mazes: Distral algorithms learn faster and achieve better final performance than all three A3C baselines on the maze tasks.Two-column variants learn faster than corresponding single-column variants; entropy regularization improves final scores despite slower learning.
Navigation: Distral algorithms have greater stability and better final results on navigation and memory tasks, with KL_2col and KL+ent_2col performing best.These tasks use procedurally generated maps and require varied navigation and memory behaviors.
Evaluation: Figure 4 reports task-averaged learning curves, distilled-policy scores, and final performances across 36 runs per algorithm.Curves use the best hyperparameters by area under the curve, with averages over four runs and variability across tasks.
Laser-tag: On laser-tag, Distral clearly outperforms multitask A3C but is not significantly better than independently trained A3C.Single-column Distral variants perform better initially, while algorithms without entropy show early plateauing.

5 Discussion

The discussion presents Distral as a policy-level framework for distilling and transferring common behavior across multitask reinforcement-learning problems.

Discussion: Distral distills common behavior while regularizing each task policy toward the distilled policy.Distillation and task-policy regularization arise as complementary halves of a joint KL-regularized optimization procedure.
Discussion: The paper argues that policies are a more semantically meaningful regularization target than deep-network parameters.The authors note that network parameters do not typically have semantic meaning by themselves.
Future work: Future work includes auxiliary losses, multiple distilled policies or latent variables, continual learning, and adaptive KL and entropy costs.These directions aim to support behavioral diversity, sequential task settings, and more flexible transfer and exploration control.

A Algorithms

The evaluated algorithms vary in whether policies are separate or shared and whether Distral uses one or two network columns with KL and entropy regularization.

A3C baselines: A3C trains a separate policy for each task, while A3C_multitask trains one policy simultaneously across all tasks.These provide independent-task and shared-policy baselines.
A3C baselines: A3C_2col uses task-specific policies with one network column shared across tasks.This baseline shares part of the architecture rather than using a distilled-policy objective.
One-column Distral: KL_1col and KL+ent_1col use one network per policy, with KL-only or combined KL and entropy regularization, respectively.KL_1col sets α = 1; KL+ent_1col sets α = 0.5 without tuning it.
Two-column Distral: KL_2col and KL+ent_2col add a shared network column that also produces the distilled policy.The two variants differ in whether entropy regularization accompanies the KL regularization.

B.1 Two room grid world

The two-room grid-world experiment uses movement costs, wall penalties, and a terminal goal reward to test Distral’s learned behavior.

Task and rewards: Each time step incurs −0.1, wall collisions incur −0.5, and reaching the goal terminates the episode with reward +1.The agent can stay put or move in the four cardinal directions.
Training settings: The experiment uses learning rate 0.1, discount 0.95, β = 5, and distilled-policy pseudocounts of 1 for every action in each state.The reported results are not sensitive to these settings.

B.2 Complex 3D tasks

The experiments use distributed A3C-style agents with shared task-specific infrastructure, extensive hyperparameter settings, and normalized scores to compare tasks with different reward scales.

Implementation: Each task is trained with a distributed agent using 32 workers coordinated through parameter servers.Agents receive 3 × 84 × 84 RGB observations and use convolutional layers, a fully connected layer, and an LSTM.
Hyperparameters: The experiments search nine combinations of entropy cost and initial learning rate while fixing α according to the algorithm variant.The learning rate is linearly annealed from its initial value.
Training setup: Training uses action repetition of 4, with 10^8 and 1.65 · 10^8 environment steps across the respective settings.The implementation also adds L2 regularization to task value estimates with coefficient 0.005, without tuning this parameter.
Evaluation: Results for navigation and laser-tag are normalized task by task using the best performance of a standard A3C agent to account for different reward scales.

C Detailed learning curves

The learning-curve figures compare task-specific and distilled policies across navigation and laser-tag tasks, showing run- or seed-level variability alongside averaged performance over training steps.

Navigation suite: Figure 5 compares task-specific policies with distilled policies across eight navigation tasks.The top two rows show task-specific policies, while the bottom two rows show distilled policies.
Navigation suite: Figure 6 compares task-specific and distilled policies across four navigation tasks using four seeds per algorithm.The average over seeds is shown in bold, and the x-axis reports training environment steps for each task.
Laser-tag suite: Figure 7 reports results for eight laser-tag tasks using task policies only.The figure reports best-hyperparameter results selected by averaged area under the learning curves, with four seed curves and their bold average.