[Paper Review] Reinforcement Learning with Automated Auxiliary Loss Search

Tairan He, Yuge Zhang+6 more|arXiv (Cornell University)|Oct 12, 2022

Reinforcement Learning in Robotics4 citations

TL;DR

This paper proposes Automated Auxiliary Loss Search (A2LS), a method that uses evolutionary search over a vast space of auxiliary loss functions to automatically discover optimal representations for reinforcement learning. A2LS achieves state-of-the-art performance across diverse tasks—image-based, vector-based, partially observable, and even Atari games—demonstrating strong generalization and sample efficiency beyond handcrafted baselines.

ABSTRACT

A good state representation is crucial to solving complicated reinforcement learning (RL) challenges. Many recent works focus on designing auxiliary losses for learning informative representations. Unfortunately, these handcrafted objectives rely heavily on expert knowledge and may be sub-optimal. In this paper, we propose a principled and universal method for learning better representations with auxiliary loss functions, named Automated Auxiliary Loss Search (A2LS), which automatically searches for top-performing auxiliary loss functions for RL. Specifically, based on the collected trajectory data, we define a general auxiliary loss space of size $7.5 imes 10^{20}$ and explore the space with an efficient evolutionary search strategy. Empirical results show that the discovered auxiliary loss (namely, A2-winner) significantly improves the performance on both high-dimensional (image) and low-dimensional (vector) unseen tasks with much higher efficiency, showing promising generalization ability to different settings and even different benchmark domains. We conduct a statistical analysis to reveal the relations between patterns of auxiliary losses and RL performance.

Motivation & Objective

• To address the sub-optimality of handcrafted auxiliary losses in reinforcement learning that rely on expert intuition and may not align with final RL performance.
• To develop a universal, automated method for discovering high-performing auxiliary loss functions that improve state representation learning.
• To evaluate the generalization capability of discovered losses across diverse RL settings, including different observation modalities, network architectures, and benchmark domains.
• To uncover statistical patterns linking auxiliary loss structures to RL performance through analysis of the full search process.

Proposed method

• A2LS formulates the auxiliary loss search as a bi-level optimization: the outer loop searches for optimal auxiliary losses using an evolutionary algorithm, while the inner loop trains an RL agent with the candidate loss.
• The method defines a comprehensive search space of size $7.5 \times 10^{20}$, covering diverse auxiliary loss configurations derived from trajectory data.
• Evolutionary search is applied on a small set of simulated environments (DeepMind Control Suite) to discover a top-performing loss, named A2-winner.
• The discovered A2-winner loss is evaluated on unseen tasks, including high-dimensional (image) and low-dimensional (vector) observations, as well as partially observable and discrete control settings.
• Statistical analysis is performed on the entire evolutionary search trajectory to correlate loss patterns with RL performance.

Experimental results

Research questions

• RQ1 Can an automated method discover auxiliary loss functions that significantly outperform handcrafted designs in reinforcement learning?
• RQ2 Does the discovered auxiliary loss generalize across different observation modalities (e.g., images vs. vectors) and RL environments?
• RQ3 What structural patterns in auxiliary losses correlate with improved RL performance, and can these be systematically identified?
• RQ4 How does the performance of the A2-winner loss compare to strong baselines across diverse RL benchmarks, including Atari?

Key findings

• The A2-winner loss, discovered via automated search, significantly improves sample efficiency and final performance on both high-dimensional (image) and low-dimensional (vector) RL tasks.
• The method generalizes effectively to unseen environments, including different robots in continuous control, partially observable settings, and even the discrete control domain of Atari 2600 games.
• The A2-winner loss outperforms strong baselines such as intrinsic curiosity, prediction-based losses, and standard auxiliary objectives across all evaluated benchmarks.
• Statistical analysis reveals that loss functions emphasizing temporal consistency and multi-step prediction patterns correlate strongly with higher RL performance.