Self-Composing Policies for Scalable Continual Reinforcement Learning

TL;DR

CompoNet is a modular neural network architecture for Continual Reinforcement Learning (CRL) that avoids catastrophic forgetting by adding a new "frozen" module for every new task. Unlike previous "growing" architectures that scale quadratically, CompoNet uses attention mechanisms to selectively compose the outputs of previous policies, achieving linear parameter growth and superior knowledge transfer across diverse robotic and visual tasks.

What Problem Does This Paper Address?

Reinforcement Learning (RL) agents typically struggle with Continual Learning: the ability to learn a sequence of tasks without forgetting old ones (catastrophic forgetting) or having new tasks hindered by old knowledge (interference).

While "growing" neural networks—which add new neurons or modules for each task—naturally prevent forgetting by freezing old parameters, they usually suffer from a scalability crisis. For example, the well-known Progressive Neural Networks architecture scales quadratically because every new layer must connect to every previous layer's hidden states. Other "compositional" approaches often require a separate, complex "manager" network to decide how to combine old skills. The authors sought a method that scales linearly, reuses knowledge autonomously, and remains robust even when previous tasks are completely unrelated to the current one.

Methods

The authors introduce CompoNet, where each task is assigned a specific Self-Composing Policy Module. When a new task arrives, the previous module is frozen, and a new one is added. The "magic" happens inside the module's three-part architecture:

1. Output Attention Head: This block looks at the current state and the outputs of all previous modules. It uses an attention mechanism to produce a "tentative" action vector—essentially a weighted combination of what previous policies would do in the current situation.
2. Input Attention Head: This block gathers relevant contextual information from previous modules to guide the internal decision-making process.
3. Internal Policy: A small feed-forward network that takes the current state and the attention-derived context to either refine, overwrite, or simply pass through the tentative action.

Because these modules attend to the outputs (actions) of previous modules rather than their internal hidden representations, the number of parameters added per task remains constant, leading to linear scalability.

Datasets

The researchers evaluated CompoNet on three distinct benchmarks:

• Meta-World: A sequence of 20 robotic manipulation tasks (10 unique tasks like "faucet-close" or "peg-unplug," repeated twice) involving continuous control.
• SpaceInvaders (ALE): 10 different playing modes of the classic Atari game, requiring visual processing of RGB frames.
• Freeway (ALE): 7 different playing modes of the Atari game "Freeway," characterized by extremely sparse rewards.

Key Insights and Ideas

• Autonomous Composition: Unlike prior work, CompoNet doesn't need a "master" network to orchestrate modules. Each module learns how to use its predecessors as part of its own training process.
• Linear vs. Quadratic Growth: By focusing on policy composition (outputs) rather than representation sharing (hidden layers), the authors break the quadratic scaling law that plagues many growing networks.
• The "Residual" Nature: The architecture acts similarly to a residual connection. If a previous policy already solves the task, the Output Attention Head can simply "pick" it, and the Internal Policy can output zeros to preserve that result.
• Needle in a Haystack: In "extreme" scalability tests, the authors showed that CompoNet could successfully identify and use one helpful policy even when it was buried among 511 "junk" modules (policies that output random noise).

Key Results

• Superior Performance: CompoNet outperformed or matched all baselines (including ProgressiveNet and PackNet) across all benchmarks. In the difficult Freeway sequence, it achieved a success score of 0.94, significantly higher than the next best method (FT-N at 0.81).
• Positive Forward Transfer: In the Meta-World benchmark, CompoNet was the only method to achieve positive forward transfer (0.01), whereas all other methods showed negative transfer (interference).
• Efficiency: Empirical tests showed that CompoNet's inference time scales much better than ProgressiveNet. While theoretically quadratic in time (due to sequential execution), in practice, it scaled effectively up to 300 tasks on a single GPU.
• Memory Footprint: The authors demonstrated that 10,000 CompoNet modules could fit into the VRAM of a single high-end consumer GPU (NVIDIA A5000).

Limitations and Future Work

• Task Boundaries: Like most CRL research, this paper assumes the agent is told exactly when a task changes and is given a task ID. Real-world applications would require the agent to detect these boundaries and IDs automatically.
• Sequential Execution: Because each module's input depends on the previous ones, they must be executed sequentially. While the authors showed this is manageable for hundreds of tasks, it could become a bottleneck for "never-ending" learning scenarios involving millions of tasks.
• Future Directions: The authors suggest exploring quantization and policy distillation to further compress the network and improve computational efficiency for even longer task sequences.