[Paper Review] MERL: Multi-Head Reinforcement Learning

Yannis Flet-Berliac, Pierre‐Marie Preux|arXiv (Cornell University)|Sep 26, 2019

Reinforcement Learning in Robotics4 citations

TL;DR

This paper proposes MERL, a general framework for multi-head reinforcement learning that injects problem-specific knowledge—specifically the fraction of variance explained ($\mathcal{V}^{ex}$) and future state prediction—into policy gradient updates. By optimizing for both reward and these auxiliary performance signals, MERL improves sample efficiency, reduces reward sparsity issues, and enhances transfer learning across continuous control and pixel-based environments, outperforming standard PPO baselines.

ABSTRACT

A common challenge in reinforcement learning is how to convert the agent's interactions with an environment into fast and robust learning. For instance, earlier work makes use of domain knowledge to improve existing reinforcement learning algorithms in complex tasks. While promising, previously acquired knowledge is often costly and challenging to scale up. Instead, we decide to consider problem knowledge with signals from quantities relevant to solve any task, e.g., self-performance assessment and accurate expectations. $\\mathcal{V}^{ex}$ is such a quantity. It is the fraction of variance explained by the value function $V$ and measures the discrepancy between $V$ and the returns. Taking advantage of $\\mathcal{V}^{ex}$, we propose MERL, a general framework for structuring reinforcement learning by injecting problem knowledge into policy gradient updates. As a result, the agent is not only optimized for a reward but learns using problem-focused quantities provided by MERL, applicable out-of-the-box to any task. In this paper: (a) We introduce and define MERL, the multi-head reinforcement learning framework we use throughout this work. (b) We conduct experiments across a variety of standard benchmark environments, including 9 continuous control tasks, where results show improved performance. (c) We demonstrate that MERL also improves transfer learning on a set of challenging pixel-based tasks. (d) We ponder how MERL tackles the problem of reward sparsity and better conditions the feature space of reinforcement learning agents.

Motivation & Objective

• To address the challenge of sample inefficiency and reward sparsity in deep reinforcement learning by incorporating problem-focused, task-agnostic signals.
• To develop a general-purpose framework applicable to any policy gradient method and environment without requiring task-specific redesign.
• To improve policy generalization and transfer learning by learning representations aligned with self-performance assessment and accurate expectations.
• To demonstrate that auxiliary signals like $\mathcal{V}^{ex}$ and future state prediction can regularize learning and enhance feature space conditioning.

Proposed method

• Introduces $\mathcal{V}^{ex}$, the fraction of variance explained by the value function, as a measure of discrepancy between estimated returns and actual returns.
• Designs MERL as a multi-head actor-critic framework where the policy network is jointly optimized using the main reward objective and auxiliary heads for $\mathcal{V}^{ex}$ and future state prediction.
• Modifies policy gradient objectives to include a regularization term based on $\mathcal{V}^{ex}$, improving the alignment of the agent’s internal representations with task-relevant performance metrics.
• Employs a lightweight, task-agnostic auxiliary head for predicting future states, commonly used in auxiliary task learning, to improve feature learning.
• Applies the framework to standard continuous control benchmarks (MuJoCo) and pixel-based Atari environments, using PPO as the base algorithm.
• Uses a multi-head architecture where each head provides a distinct, complementary signal to guide policy optimization, enhancing robustness and generalization.

Experimental results

Research questions

• RQ1 Can injecting a measure of self-performance assessment ($\mathcal{V}^{ex}$) into policy gradient updates improve sample efficiency and learning stability?
• RQ2 Does combining $\mathcal{V}^{ex}$ with a future state prediction head lead to better generalization and transfer learning in complex control tasks?
• RQ3 How does MERL compare to standard PPO in terms of performance and sample efficiency across diverse environments?
• RQ4 To what extent does MERL mitigate the problem of sparse rewards by providing dense, informative supervisory signals?
• RQ5 Can MERL’s auxiliary heads learn representations that are transferable across different tasks, such as from Ms. Pac-Man to other Atari games?

Key findings

• MERL outperforms standard PPO on all 9 continuous control tasks in the MuJoCo benchmark, with consistent improvements in sample efficiency and final performance.
• On Atari 2600 games, MERL demonstrates superior transfer learning, achieving higher average performance when transferring from Ms. Pac-Man to other games compared to standard PPO.
• The ablation study shows that combining $\mathcal{V}^{ex}$ and future state prediction heads yields better performance than using either head alone, with the $\mathcal{V}^{ex}$ head alone degrading performance in HalfCheetah but improving when combined.
• The computational overhead of MERL is minimal—only 5% on MuJoCo and 7% on Atari tasks—making it practical for real-world deployment.
• MERL heads learn task-agnostic, general-purpose representations that improve feature space conditioning and support better generalization across environments.
• The framework successfully regularizes the policy network by aligning optimization with higher-level insights into task performance, reducing reliance on sparse reward signals.