Constructing an Optimal Behavior Basis for the Option Keyboard

TL;DR

The paper tackles zero-shot transfer in multitask reinforcement learning by addressing the exponential cost of convex coverage sets (CCS) and the limitations of Generalized Policy Improvement (GPI). It introduces Option Keyboard Basis (OKB), a method that incrementally builds an optimal behavior basis $\Pi_k$ and a task-conditioned meta-policy $\omega$ so that the Option Keyboard (OK) is optimal for any task in the linear-reward family $\mathcal{M}_{\text{lin}}^{\boldsymbol{\phi}}$, with $|\Pi_k| \le |\mathrm{CCS}|$ and $\mathrm{CCS} \subseteq \Psi^{\text{OK}}(\Pi_k)$. The authors prove that OKB is strictly more expressive than CCS in transfer settings and can solve certain non-linear tasks efficiently, while offering strong guarantees that zero-shot optimality is achieved without constructing the full CCS. Empirically, OKB outperforms state-of-the-art GPI-based methods across challenging high-dimensional domains, with larger advantages as the reward-feature dimensionality $d$ grows, and demonstrates temporally extended base-policy selection in the learned meta-policy. These results suggest that OKB enables scalable, robust reuse of sub-policies for a broad class of linear and some non-linear tasks, with practical implications for efficient multitask RL in complex environments.

Abstract

Multi-task reinforcement learning aims to quickly identify solutions for new tasks with minimal or no additional interaction with the environment. Generalized Policy Improvement (GPI) addresses this by combining a set of base policies to produce a new one that is at least as good -- though not necessarily optimal -- as any individual base policy. Optimality can be ensured, particularly in the linear-reward case, via techniques that compute a Convex Coverage Set (CCS). However, these are computationally expensive and do not scale to complex domains. The Option Keyboard (OK) improves upon GPI by producing policies that are at least as good -- and often better. It achieves this through a learned meta-policy that dynamically combines base policies. However, its performance critically depends on the choice of base policies. This raises a key question: is there an optimal set of base policies -- an optimal behavior basis -- that enables zero-shot identification of optimal solutions for any linear tasks? We solve this open problem by introducing a novel method that efficiently constructs such an optimal behavior basis. We show that it significantly reduces the number of base policies needed to ensure optimality in new tasks. We also prove that it is strictly more expressive than a CCS, enabling particular classes of non-linear tasks to be solved optimally. We empirically evaluate our technique in challenging domains and show that it outperforms state-of-the-art approaches, increasingly so as task complexity increases.

Figures (7)

• Figure 1: Domains used in the experiments: Minecart, FetchPickAndPlace, Item Collection, and Highway.
• Figure 2: Mean normalized return per iteration for each method on a set of test tasks.
• Figure 3: Mean normalized return over test tasks as a function of iteration number (i.e., number of base policies learned per method) [FetchPickAndPlace]. As $d$ increases, OKB's performance advantage over SFOLS (a state-of-the-art GPI-based algorithm) grows.
• Figure 4: Mean return in the Item Collection domain under a non-linear reward function.
• Figure 5: Continuous actions from the meta-policy on a sample task. OKB selects base policies in a temporally consistent way.

Theorems & Definitions (20)

• Definition 2.1: Barreto+2020
• Proposition 4.1
• Theorem 4.2
• Theorem 4.3
• Proposition 4.4
• Theorem 4.5: Barreto+2017
• Proposition : \ref{['th:existence-omega']}
• proof
• Theorem : \ref{['th:advantage']}
• proof