Task diversity produces systematic transfer but inhibits continual reinforcement learning

Technical Summary: Banyan – A Procedurally-Generated Environment for Studying the Effect of Task Diversity on Transfer

1. Problem Statement

Continual reinforcement learning (RL) aims to develop agents that not only improve on current tasks but also adapt as task distributions shift. While previous research has established that training on diverse tasks can induce zero-shot generalization (evaluated with frozen weights), it remains unclear whether task diversity improves an agent's ability to continue learning through weight updates as distributions change. Existing benchmarks for continual RL typically involve sequential learning over at most hundreds of tasks, while large-scale procedurally generated environments (e.g., XLand, AdA) contain millions of tasks but lack the control necessary to isolate how the amount and type of diversity affect learning across controlled distribution shifts.

2. Methodology: The Banyan Environment

The authors introduce Banyan, a GPU-accelerated, procedurally generated domain designed to study continual RL over compositional task spaces. Banyan factorizes a task $\tau = \langle L, T, O \rangle$ into three independently controllable axes:

• Layout ($L$): The environment map, including wall configurations and initial object placements, which defines the navigation problem.
• Task-Tree Topology ($T$): The abstract structure of the goal, defining the depth, branching, and dependency order of sub-goals (transformations vs. combinations).
• Object Assignment ($O$): The concrete identities of objects assigned to the nodes of the task tree.

Task Mechanics:

• Goal: The agent must navigate the grid, pick up, drop, and combine objects (or transform single items) to produce a specific root object defined by the task tree.
• Observations: The agent sees the full grid and the goal object description but must infer the task tree structure and solution sequence.
• Dynamics: Object combination and transformation rules are fixed; only the spatial layout and object identities vary.
• Curriculum: Tasks span depths 1 through 6. Deeper tasks act as a curriculum for shallower ones, and distributions are nested (depth-6 tasks contain all shallower sub-tasks).

Experimental Protocol:
The authors evaluate two algorithms: Proximal Policy Optimization (PPO) and Parallelised Q-Network (PQN). They measure:

• Forward Transfer: The "transfer gap" ($\Delta_i$) between the terminal success rate on distribution $d_{i-1}$ and the initial success rate on $d_i$.
• Backward Transfer: The change in performance on an earlier distribution $d_i$ after training on subsequent distributions.

3. Key Contributions

1. Banyan Benchmark: A controlled, many-task domain allowing independent variation of layouts, task-tree instances, and topologies, enabling the study of continual RL over billions of potential tasks.
2. Evidence of Systematic Transfer: Experimental proof that increasing diversity along any axis (layouts, object assignments, or topologies) induces systematic transfer, allowing agents to start new distributions near their previous performance levels.
3. Evidence of Diversity-Induced Inhibition: Experimental proof that while moderate diversity aids transfer, excessive diversity inhibits long-run continual learning, causing agents to plateau on new distributions while continuing to refine performance on earlier ones.

4. Key Results

A. Diversity Produces Systematic Transfer (Single Shift)
In experiments involving a single shift from distribution $d_1$ to $d_2$:

• Forward Transfer: Increasing the number of layouts ($|L|$), object assignments ($|O|$), or topologies ($|T|$) in $d_1$ reduces the transfer gap ($\Delta_2$) toward zero. Agents trained on highly diverse $d_1$ begin $d_2$ with performance levels comparable to their final performance on $d_1$.
• Backward Transfer: Increasing $|L|$ or $|O|$ had little effect on backward transfer. However, increasing the number of topologies ($|T|$) significantly improved backward transfer, suggesting that varying the structural dependencies of tasks is crucial for retaining knowledge of prior distributions.
• Mechanism: Analysis of PQN's TD-error at the $d_1 \to d_2$ boundary revealed that higher diversity yields lower initial TD-error and better initial policies. This suggests diversity shapes value features that generalize well off-distribution.
• Generalizability: These effects were replicated in a continuous-control "Point Mass" domain, indicating the phenomenon is not an artifact of discrete grid-world control.

B. Diversity Inhibits Long-Run Continual Learning (Sequence of Shifts)
In experiments involving a sequence of ten distributions ($d_1 \dots d_{10}$):

• The Paradox: While high diversity ($n=256$ unique layouts/topologies) successfully closed the transfer gap at every single boundary (agents started each new distribution at the previous peak), it caused the agent to plateau in its ability to improve within the sequence.
• Performance Trajectory:
  • At intermediate diversity ($n=16$), agents continued to learn, improving success rates from $\approx 0.74$ on $d_1$ to $\approx 0.95$ on $d_{10}$.
  • At high diversity ($n=256$), the success rate on new distributions remained flat near $0.80$ across all ten distributions. The endpoint performance on $d_{10}$ actually declined monotonically as diversity increased beyond $n=16$.
• Interference: High diversity led to strong backward transfer (performance on $d_1$ improved from $\approx 0.80$ to $\approx 0.95$ after 10 rounds) but prevented specialization on new distributions. The authors attribute this to interference: the network learns the general task structure shared across distributions but is inhibited from learning the specific structures required by new distributions.

5. Significance and Claims

The paper claims a significant, somewhat paradoxical finding: Task diversity is a double-edged sword for continual RL.

• Motivation: Diversity is essential for inducing systematic transfer, allowing agents to adapt immediately to new distribution shifts without catastrophic forgetting of the general task structure.
• Limitation: However, the same diversity that smooths individual transitions inhibits the compounding of learning over long sequences. The authors argue that "learning and scale, on their own, may not be enough" to enable continual RL; the trade-off between generalization (transfer) and specialization (learning new specifics) appears intrinsic to weight-sharing function approximation in the many-task regime.

The work highlights a critical gap in current RL research: while diversity is celebrated for zero-shot generalization, it may actively hinder the ability of agents to optimize performance on new distributions when learning continuously. The authors release Banyan as a tool for the community to further investigate these trade-offs and develop algorithms that can mitigate the interference caused by high task diversity.