🤖 machine learning

Goal-Conditioned Agents that Learn Everything All at Once

This paper introduces Learning Everything all at Once (LEO), a method that enables efficient, parallel all-goals off-policy updates in goal-conditioned reinforcement learning by jointly outputting values and actions for every goal, thereby achieving significant performance gains and massive speed-ups over traditional relabelling techniques.

Original authors: Michael Matthews, Matthew Jackson, Michael Beukman, Thomas Foster, Alistair Letcher, Scott Fujimoto, Cédric Colas, Jakob Foerster

Published 2026-05-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Michael Matthews, Matthew Jackson, Michael Beukman, Thomas Foster, Alistair Letcher, Scott Fujimoto, Cédric Colas, Jakob Foerster

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to navigate a giant, ever-changing maze. In the old way of doing this (called Goal-Conditioned Reinforcement Learning), you would give the robot a specific instruction, like "Go to the kitchen." The robot would wander around, bump into walls, and eventually maybe find the kitchen.

When the robot finishes, the teacher looks at the path it took. If the robot was told to go to the kitchen but accidentally walked past the library on the way, the old method would say, "That library part was a mistake; throw it away. We only care about the kitchen."

The Problem: This is a waste of information! The robot learned a lot about the library while trying to find the kitchen. By throwing that data away, the robot has to re-learn about the library later if you ever ask it to go there.

The Old "Fix" (Hindsight Experience Replay):
Previously, researchers tried to fix this by looking at the robot's path and saying, "Oh, it ended up at the library! Let's pretend we wanted it to go to the library all along." This helps a little, but it's like trying to re-label a whole movie after watching it. It's slow, and you can only do it for a few different endings at a time.

The New Solution: "Learning Everything All at Once" (LEO)
The authors of this paper propose a smarter way. Instead of teaching the robot one goal at a time, they build a "super-brain" that learns every possible goal simultaneously.

Think of it like a chef who is cooking a massive banquet.

The Old Way: The chef cooks one dish (the kitchen), tastes it, throws away the notes on how they chopped the onions, and then starts over to cook the next dish (the library).
The LEO Way: The chef cooks the entire banquet in one go. They have a giant board with 512 different recipes (goals) written on it. As they chop an onion, they instantly update their knowledge for all 512 recipes at the same time. They realize, "Hey, chopping onions is useful for the soup, the stew, and the salad!"

How it works technically (but simply):
Usually, a computer network asks, "What is the best move for this specific goal?"
LEO changes the network so it asks, "What is the best move for every single goal right now?" It outputs a giant list of answers for every possible destination in one single pass. This makes learning incredibly fast (over 250 times faster than the old "re-labeling" method) because the robot doesn't have to replay the same journey hundreds of times.

The Catch: The "Late Fusion" Problem
The authors found a small flaw. Because the robot has to think about all goals at once, it sometimes gets confused. It's like a student trying to study for 500 different exams at the exact same moment. They might get a general sense of the material, but they might not be as sharp at answering a specific question as a student who only studied for that one exam.

In the paper, this meant LEO was great at solving hard, complex goals but sometimes worse at simple ones because it was "diluting" its focus.

The Ultimate Fix: "Dual LEO" (The Teacher-Student System)
To solve this, the authors created a team-up:

The Teacher (LEO): This network is the "data sponge." It learns everything about everything, even if it's a bit messy. It provides a rough map of the world.
The Student (Standard Network): This network is the specialist. It focuses only on the specific goal you are currently asking it to do.

The Teacher says to the Student, "Hey, I know the general direction to the library, even if I'm not perfect. Here's a hint." The Student takes that hint and refines it to become an expert at getting to the library.

The Results
The team tested this on a complex video game called Craftax (a bit like Minecraft) and some robot movement tasks.

On the big, hard version of the game with 512 different goals, the new Dual LEO method was the clear winner, beating all other methods.
It learned much faster and could handle a huge variety of tasks that other methods gave up on.
They also showed this works for continuous movements (like a robot arm moving smoothly), not just video game steps.

In Summary
This paper introduces a method where an AI learns to solve every possible problem at the same time, rather than one by one. When it gets too broad and loses focus, they pair it with a specialist "student" network that uses the broad knowledge as a guide. This allows robots and game-playing AIs to learn massive amounts of information much faster and more efficiently than before.

Technical Summary: Learning Everything All at Once (LEO)

Problem Statement

Goal-Conditioned Reinforcement Learning (GCRL) aims to train agents that can execute a diverse range of tasks upon command. A fundamental challenge in GCRL is data efficiency. When an agent explores an environment to achieve a specific commanded goal, the resulting trajectory contains rich information relevant to many other potential goals. Standard on-policy updates or methods that only relabel trajectories with a single alternative goal (e.g., Hindsight Experience Replay, HER) discard the vast majority of this information.

While "all-goals" learning—updating the value function with respect to every possible goal for every transition—offers a theoretical path to maximal information extraction, it has historically been computationally infeasible. Naïve goal relabelling scales linearly with the number of goals, making it impractical for large or continuous goal sets. Furthermore, existing approaches often struggle to balance the computational cost of updating against all goals with the representational capacity required to learn high-fidelity policies for specific commands.

Methodology

The paper proposes Learning Everything All at Once (LEO), a method designed to perform efficient, parallel all-goals updates on arbitrary discrete goal sets.

1. The LEO Architecture

Instead of conditioning the value network on a specific goal (as in Universal Value Function Approximators, UVFA), LEO reparameterizes the network to output values for every goal simultaneously.

Currying: The Q-function is transformed from $Q(s, g) \to \mathbb{R}^A$ to $Q(s) \to \mathbb{R}^{G \times A}$ , where $G$ is the set of goals and $A$ is the set of actions.
Efficient Updates: This allows the agent to perform a single forward pass to generate Q-values for all goals and a single backward pass to update the network using the reward vectors for all goals. This eliminates the linear scaling cost of naïve relabelling.
Inference: During execution, the agent simply indexes the output vector corresponding to the commanded goal.

2. The Late Fusion Limitation and Dual LEO

The authors identify a critical trade-off in the LEO architecture. While LEO excels at ingesting information from all goals, it suffers from late fusion. Because the network must produce outputs for all goals before the specific commanded goal is selected, it creates an information bottleneck. The network is forced to learn a representation suitable for all goals simultaneously, which can degrade performance on specific goals, particularly easier ones where a focused representation (early fusion) is superior.

To address this, the paper introduces Dual LEO, which combines the strengths of LEO and standard UVFA-style goal-conditioned networks:

Teacher-Student Dynamic: A LEO network (the teacher) and a UVFA network (the student) are trained independently on the same data stream.
Information Transfer: The LEO network acts as a "data sponge," internalizing maximal information and providing coarse, directional guidance for goals the UVFA network has not yet seen or solved.
Implementation Variants:
- Dual LEO (PQN): The final Q-value is a linear combination of the LEO and UVFA estimates ( $\alpha Q_{LEO} + (1-\alpha) Q_{UVFA}$ ).
- Dual LEO (PPO): The LEO network provides additional loss terms to guide the PPO actor and critic towards the LEO policy and value estimates, particularly for goals where the UVFA network lacks positive signals.

3. Extension to Continuous Control

The authors extend LEO to continuous control environments (e.g., Ant Maze) by discretizing the continuous goal space into a grid. Continuous goals are snapped to the nearest quantized grid point, allowing the all-goals update mechanism to function. The policy network is similarly adapted to output actions for all discrete goal heads.

Key Contributions

LEO Method: A novel approach enabling efficient all-goals updates on arbitrary discrete goal sets by outputting values for all goals in a single network pass, achieving a >250× speed-up over naïve all-goals relabelling.
CraftaxGC Benchmark: The introduction of a challenging, procedurally generated, partially observable GCRL benchmark based on Craftax, featuring a large, heterogeneous set of semantic goals (e.g., inventory counts, specific block adjacencies, dungeon levels) that cannot be solved by simple target-state mapping.
Dual LEO: A hybrid architecture that mitigates the late fusion bottleneck of pure LEO by using LEO as a teacher to bootstrap a standard goal-conditioned student network, achieving state-of-the-art performance.
Versatility: Demonstration that the LEO concept extends beyond discrete environments to continuous control tasks, maintaining competitiveness with existing baselines.

Experimental Results

The authors evaluate their methods on CraftaxGC (with 136 and 512 goals) and JaxGCRL (Ant Maze tasks).

CraftaxGC (Large Goal Set): On the full Craftax benchmark (512 goals), standard UVFA methods (PPO, PQN) and PQN+HER perform poorly due to sparse rewards and the inability to leverage off-policy data effectively. LEO significantly outperforms these baselines. Dual LEO (PPO) achieves the best final performance, massively outperforming all other methods.
CraftaxGC (Small Goal Set): On the smaller Craftax-Classic (136 goals), simple PPO performs relatively well. LEO outperforms its base PQN but is slightly weaker than PQN+HER. However, Dual LEO variants consistently outperform their constituent parts, demonstrating the robustness of the teacher-student approach.
Goal-Specific Performance: Analysis reveals that while LEO excels at hard goals (where it provides the only positive signal), it can underperform on easy goals compared to UVFA due to the late fusion bottleneck. Dual LEO resolves this, maintaining high performance on hard goals without sacrificing performance on easy ones.
Continuous Control: On Ant Maze tasks, SAC+LEO outperforms baselines on smaller mazes. On larger mazes, performance is comparable to SAC+HER and Contrastive RL (CRL).
Speed: LEO achieves a 34% slowdown compared to single-goal learning but is 264× faster than naïve all-goals relabelling.

Significance and Claims

The paper claims that LEO unlocks the potential of all-goals learning at scale, providing a practical tool for RL practitioners in complex environments. The authors posit that the core trade-off in GCRL lies between:

UVFA-style approaches: Which generalize well via early fusion but may struggle to ingest massive amounts of off-policy data efficiently.
LEO-style approaches: Which efficiently ingest maximal information but struggle with high-fidelity representation due to late fusion.

Dual LEO is presented as a solution that bridges this gap, allowing agents to internalize coarse information from all goals while refining specific policies through a dedicated goal-conditioned network. The authors emphasize that by unlocking efficient all-goals learning, LEO enables agents to reframe prior experience in the context of new goals, effectively multiplying the information extracted from every trajectory. The work is open-sourced to facilitate further research in complex, goal-conditioned domains.