Constructing an Optimal Behavior Basis for the Option Keyboard
Extracted source text · arXiv 2505.00787
\documentclass{article}
% if you need to pass options to natbib, use, e.g.:
% \PassOptionsToPackage{numbers, compress}{natbib}
% before loading neurips_2025
% ready for submission
% \usepackage{neurips_2025}
% to compile a preprint version, e.g., for submission to arXiv, add add the
% [preprint] option:
%\usepackage[preprint]{neurips_2025}
% to compile a camera-ready version, add the [final] option, e.g.:
\usepackage[final]{neurips_2025}
% to avoid loading the natbib package, add option nonatbib:
% \usepackage[nonatbib]{neurips_2025}
\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc} % use 8-bit T1 fonts
%\usepackage{hyperref} % hyperlinks
\usepackage{url} % simple URL typesetting
\usepackage{booktabs} % professional-quality tables
\usepackage{amsfonts} % blackboard math symbols
\usepackage{nicefrac} % compact symbols for 1/2, etc.
\usepackage{microtype} % microtypography
%\usepackage{xcolor} % colors
\usepackage[backref]{hyperref} % hyperlinks
\usepackage[dvipsnames]{xcolor}
\definecolor{mydarkblue}{rgb}{0,0.08,0.45}
\definecolor{myblue}{RGB}{86,120,255}
\definecolor{citecolor}{HTML}{0b64c5}
%\hypersetup{
% colorlinks=true,
% allcolors=mydarkblue
%}
\hypersetup{
colorlinks=true,
linkcolor=Mulberry,
citecolor=NavyBlue,
}
\usepackage{amsmath,bm}
\usepackage{amssymb}
\usepackage{mathtools}
\usepackage{amsthm}
\usepackage{bbm}
\usepackage[linesnumbered,ruled,vlined,noend]{algorithm2e}
\definecolor{mydarkblue}{rgb}{0,0.08,0.45}
\definecolor{myblue}{RGB}{86,120,255}
\definecolor{citecolor}{HTML}{0b64c5}
%\newcommand{\mycomment}[1]{\textcolor{citecolor}{\textit{\(\triangleright\) #1}}}
\SetCommentSty{mycommentstyle}
\newcommand{\mycommentstyle}{\footnotesize\color{NavyBlue}\ttfamily}
\SetAlCapFnt{\footnotesize} % for the main caption line
\SetAlCapNameFnt{\footnotesize} % for "Algorithm" + number
\usepackage{soul}
\usepackage{cancel}
%\usepackage{ulem}
\usepackage{graphicx}
\usepackage{wrapfig}
%\usepackage{subfigure}
\usepackage{subcaption}
\usepackage[export]{adjustbox} % for valign options with includegraphics
\usepackage{etoolbox}
\makeatletter
% Remove right hand margin in algorithm
\patchcmd{\@algocf@start}% <cmd>
{-1.5em}% <search>
{0pt}% <replace>
{}{}% <success><failure>
\makeatother
\theoremstyle{plain}
\newtheorem{theorem}{Theorem}[section]
\newtheorem{proposition}[theorem]{Proposition}
\newtheorem{lemma}[theorem]{Lemma}
\newtheorem{corollary}[theorem]{Corollary}
\theoremstyle{definition}
\newtheorem{definition}[theorem]{Definition}
\newtheorem{assumption}[theorem]{Assumption}
\theoremstyle{remark}
\newtheorem{remark}[theorem]{Remark}
\theoremstyle{plain}
\newtheorem*{theorem*}{Theorem}
\newtheorem*{proposition*}{Proposition}
\newtheorem*{lemma*}{Lemma}
\DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\argmin}{arg\,min}
\newcommand{\vect}[1]{\boldsymbol{\mathbf{#1}}}
\newcommand{\Ex}{\mathbb{E}}
\newcommand{\ours}{NAME\_ALGO}
\newcommand{\pigpi}{\pi^{\text{GPI}}}
\newcommand{\piok}{\pi^{\text{OK}}}
\newcommand{\pismp}{\pi^{\text{SMP}}}
\newcommand{\vpsi}{\vect{\psi}}
\newcommand{\vphi}{\vect{\phi}}
\newcommand{\ccs}{\mathrm{CCS}}
\newcommand{\setcover}{\mathrm{SC}}
\newcommand{\Aspace}{\mathcal{A}}
\newcommand{\Sspace}{\mathcal{S}}
\newcommand{\Wspace}{\mathcal{W}}
\newcommand{\Zspace}{\mathcal{Z}}
\newcommand{\Mspace}{\mathcal{M}}
\newcommand{\Mlinear}{\mathcal{M}_{\text{lin}}^{\vect{\phi}}}
\newcommand{\Wsupport}{\mathcal{W}^{\text{sup}}}
\newcommand{\Wcorner}{\mathcal{W}^{\text{cw}}}
\newcommand{\PsiOK}{\Psi^{\text{OK}}}
\newcommand{\PsiBase}{\Psi^{\text{base}}}
\newcommand{\w}{\vect{w}}
\newcommand{\z}{\vect{z}}
\newcommand{\bellmanop}{\mathcal{T}}
\newcommand{\defas}{\triangleq}
\title{Constructing an Optimal Behavior Basis for the Option Keyboard}
% The \author macro works with any number of authors. There are two commands
% used to separate the names and addresses of multiple authors: \And and \AND.
%
% Using \And between authors leaves it to LaTeX to determine where to break the
% lines. Using \AND forces a line break at that point. So, if LaTeX puts 3 of 4
% authors names on the first line, and the last on the second line, try using
% \AND instead of \And before the third author name.
\author{%
Lucas N.~Alegre\\
%\thanks{Use footnote for providing further information about author (webpage, alternative address)---\emph{not} for acknowledging funding agencies.} \\
Institute of Informatics\\
Federal University of Rio Grande do Sul\\
Porto Alegre, RS, Brazil \\
\texttt{lnalegre@inf.ufrgs.br} \\
\And
Ana L.~C.~Bazzan\\
Institute of Informatics\\
Federal University of Rio Grande do Sul\\
Porto Alegre, RS, Brazil \\
\texttt{bazzan@inf.ufrgs.br} \\
\And
André Barreto\\
Google DeepMind\\
London, UK\\
\texttt{andrebarreto@google.com}
\And
Bruno C.~da Silva\\
University of Massachusetts\\
Amherst, MA, USA\\
\texttt{bsilva@cs.umass.edu}
}
\begin{document}
\maketitle
\begin{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 \textit{behavior basis}---that enables zero-shot identification of optimal solutions for \textit{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 \textit{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.
\end{abstract}
\section{Introduction}
Reinforcement learning (RL) methods have been successfully used to solve complex sequential decision-making problems~\citep{Silver+2017, Bellemare+2020}.
However, traditional RL algorithms typically require thousands or millions of interactions with the environment to learn a \textit{single} policy for a \textit{single} task.
Multi-task RL methods address this limitation by enabling agents to quickly identify solutions for new tasks with minimal or no additional interactions.
Such methods often achieve this by learning a \textit{set} of specialized policies (or \textit{behavior basis}) designed for specific tasks and subsequently combining them to more rapidly solve novel tasks.
\looseness-1
A powerful approach for combining policies to solve new tasks in a zero-shot manner leverages \textit{successor features} (SFs) and \textit{generalized policy improvement} (GPI)~\citep{Barreto+2018,Barreto+2020}.
%
GPI can combine base policies to solve a new task, producing a policy that is at least as good as any individual base policy. Importantly, however, GPI policies are not guaranteed to be optimal.
%
Optimality can be ensured if tasks can be expressed as a linear combination of reward features using techniques that compute a \textit{convex coverage set} (CCS)~\citep{Alegre+2022}. A CCS is a set of policies that includes an optimal policy for any linear task. Unfortunately, the number of policies in a CCS often grows exponentially with the number of reward features~\citep{Roijers2016}, making existing algorithms computationally expensive and difficult to scale to complex domains~\citep{Yang+2019,Alegre+2022,Alegre+2023}.
The \textit{option keyboard} (OK) method 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 by assigning state-dependent linear weights to reward features~\citep{Barreto+2019}.
%
This allows OK to express a larger spectrum of behaviors than GPI and potentially solve tasks beyond linear rewards, enabling it to produce policies that GPI cannot represent.
%
However, its performance critically depends on the choice of base policies available.
%
Existing OK-based methods assume a predefined set of base policies or rely on expert knowledge to construct one, without addressing how to identify a good behavior basis~\citep{Barreto+2019,Carvalho+2023neurips}.
%
This raises a key question and open problem: is there an optimal set of base policies for the OK---an optimal \textit{behavior basis}---that enables zero-shot identification of optimal solutions for \textit{any} linear tasks?
%
\looseness-1
We solve this open problem by introducing \textbf{O}ption \textbf{K}eyboard \textbf{B}asis (OKB), a novel method with strong formal guarantees that efficiently identifies an optimal behavior basis for the OK.
%
OKB provably identifies a set of policies that allows the OK to optimally solve any linear task; i.e., it can express all policies in a CCS.
%
We show that it significantly reduces the number of base policies required to ensure zero-shot optimality in new tasks.
%
Furthermore, we prove that the set of policies it can express is strictly more expressive than a CCS, enabling particular classes of non-linear tasks to be solved optimally.
%
We empirically evaluate our method in challenging high-dimensional RL problems and show that it consistently outperforms state-of-the-art GPI-based approaches. Importantly, we also observe that the performance gain over competing methods becomes more pronounced as the number of reward features increases.
\section{Background}
\subsection{Reinforcement Learning}
\label{sec:rl}
An RL problem \citep{Sutton&Barto2018} is typically modeled as a \textit{Markov decision process} (MDP). An MDP is defined as a tuple $M \defas (\mathcal{S},\mathcal{A},p,r,\mu,\gamma)$, where $\mathcal{S}$ is a state space, $\mathcal{A}$ is an action space, $p(\cdot|s,a)$ describes the distribution over next states given that the agent executed action $a$ in state $s$, $r : \mathcal{S} \times \mathcal{A} \times \mathcal{S} \to \mathbb{R}$ is a reward function, $\mu$ is an initial state distribution, and $\gamma \in [0,1)$ is a discounting factor.
Let $S_t$, $A_t$, and $R_t \defas r(S_t,A_t,S_{t+1})$ be random variables corresponding to the state, action, and reward, respectively, at time step $t$.
The goal of an RL agent is to learn a policy $\pi:\mathcal{S} \to \mathcal{A}$ that maximizes the expected discounted sum of rewards (\textit{return}), $G_t = \sum_{i=0}^{\infty} \gamma^{i} R_{t+i}$.
The action-value function of a policy $\pi$ is defined as $q^\pi(s,a) \defas \Ex_\pi [G_t | S_t = s, A_t = a]$, where $\Ex_\pi [\cdot]$ denotes the expectation over trajectories induced by $\pi$.
%
Given $q^\pi$, one can define a \textit{greedy} policy $\pi'(s) \in \argmax_a q^\pi(s,a)$.
It is guaranteed that $q^{\pi'}(s,a) \geq q^\pi(s,a), \forall (s,a) \in \mathcal{S} \times \mathcal{A}$.
The processes of computing $q^\pi$ and $\pi'$ are known, respectively, as the \textit{policy evaluation} and \textit{policy improvement} steps. Under certain conditions, repeatedly executing the policy evaluating and improvement steps leads to an optimal policy $\pi^*(s) \in \argmax\nolimits_a q^*(s,a)$ \citep{Puterman2014}.
%
Let $(\mathcal{S},\mathcal{A},p,\mu,\gamma)$ be a \textit{Markov control process} (MPC)~\citep{Puterman2014}, i.e., an MDP without a reward function. Given an MPC, we define a family of MDPs $\mathcal{M} \defas \{(\mathcal{S},\mathcal{A},p,r,\mu,\gamma) \mid r : \Sspace\times\Aspace\times\Sspace \to\mathbb{R}\}$ that only differ in their reward function.
We refer to any MDP $M\in\mathcal{M}$ (and its corresponding reward function $r$) as a \emph{task}.
\subsection{Generalized Policy Evaluation and Improvement}
\looseness-1
Generalized Policy Evaluation (GPE) and Generalized Policy Improvement (GPI) generalize the policy evaluation and improvement steps to the case where an agent has access to a \emph{set} of policies.
In particular, GPE and GPI are used, respectively, \textit{(i)} to evaluate a policy on multiple tasks; and \textit{(ii)} to construct a policy, capable of solving a particular novel task, by improving on an existing set of policies.
%
\begin{definition}[\citet{Barreto+2020}] \label{def:gpi}
\looseness-1
\textit{GPE} is the computation of the action-value function of a policy $\pi$, $q^\pi_{r}(s,a)$, for a set of tasks.
Moreover, given a set of policies $\Pi$ and a reward function of an arbitrary task, $r$, \textit{GPI} defines a policy, $\pi'$, such that
%\begin{equation}
$
q^{\pi'}_{r}(s,a) \geq \max_{\pi\in\Pi} q^{\pi}_{r}(s,a) \ \text{ for all } (s,a) \in \Sspace \times \Aspace.
$
%\end{equation}
\end{definition}
Based on the latter definition, for any reward function $r$, a \textit{GPI policy} can be constructed based on a set of policies $\Pi$ as
\begin{equation}
\label{eq:sf-gpi}
\pigpi(s;\Pi) \in \underset{a \in \mathcal{A}}{\arg\max} \,\, \underset{\pi \in \Pi}{\max} \,\, q^{\pi}_{r}(s,a) .
\end{equation}
Let $q^{\text{GPI}}_{r}(s,a)$ be the action-value function of $\pigpi$ under the reward function $r$. %for any state-action pair $(s,a)$.
The GPI theorem \citep{Barreto+2017} ensures that $\pigpi$ in Eq.~\eqref{eq:sf-gpi} satisfies Def.~\ref{def:gpi}; i.e., that
$q^{\text{GPI}}_{r}(s,a) \geq \max_{\pi \in \Pi} q_{r}^{\pi}(s,a)$ for all $(s,a) \in \mathcal{S} \times \mathcal{A}$.
This implies that Eq.~\eqref{eq:sf-gpi} can be used to define a policy guaranteed to perform at least as well as all other policies $\pi_i \in \Pi$ w.r.t. any given reward function, $r$.
%
The GPI theorem can also be extended to the case $q^{\pi_i}$ is replaced with an approximation, $\tilde{q}^{\pi_i}$ \citep{Barreto+2018}.
\subsection{GPE\&GPI via Successor Features}
\textit{Successor features} (SFs) allows us to perform GPE\&GPI efficiently~\citep{Barreto+2017}.
Assume that the reward functions of interest are linear w.r.t. \textit{reward features}, $\vphi: \Sspace\times\Aspace\times\Sspace\to\mathbb{R}^d$. That is, a reward function, $r_{\w}$, can be expressed as $r_{\w}(s,a,s') = \vphi(s,a,s')\cdot\w$, where $\w\in\mathbb{R}^d$ is a weight vector.
Then, the SFs of a policy $\pi$ for a given state-action pair $(s,a)$, $\vpsi^\pi(s,a) \in \mathbb{R}^d$, are defined as
\begin{equation}
\label{eq:psi}
\vpsi^\pi(s,a) \defas \Ex_\pi \left[ \sum_{i=0}^{\infty} \gamma^i \vphi_{t+i} \mid S_t=s, A_t=a \right],
\end{equation}
where $\vphi_t \defas \vphi(S_t,A_t,S_{t+1})$.
Notice that the definition of SFs corresponds to a form of action-value function, where the features $\vphi_t$ play the role of rewards. Thus, SFs can be learned through any temporal-difference (TD) learning algorithm (e.g., Q-learning~\citep{Watkins1989,Mnih+2015} for learning $\vpsi^{\pi^*}$).
We refer to
%\begin{equation}
$
\vpsi^\pi \defas \Ex_{S_0 \sim \mu} \left[ \vpsi^\pi(S_0, \pi(S_0)) \right]
$
as the \textit{SF vector} associated with $\pi$, where the expectation is with respect to the initial state distribution.
%\end{equation}
\looseness-1
Given the SFs $\vpsi^{\pi}(s,a)$ of a policy $\pi$, it is possible to \textit{directly} compute the action-value function $q^\pi_{\vect{w}}(s,a)$ of $\pi$, under \textit{any} linearly-expressible reward functions, $r_{\vect{w}}$, as follows:
%\begin{align}
$
q_{\vect{w}}^\pi(s,a)
= \Ex_\pi \left[ \sum_{i=0}^{\infty} \gamma^i r_{\w}(S_{t+i},A_{t+i},S_{t+i+1}) \mid S_t=s, A_t=a \right]
%
= \vpsi^\pi(s,a)\cdot \vect{w} .
\label{eq:SF_dot_product}
$
%\end{align}
That is, given any set of reward weights, $\{\w_i\}_{i=1}^n$, GPE can be performed efficiently via inner-products between the SFs and the reward weights: $q_{\w_i}^\pi(s,a) = \vpsi^\pi(s,a) \cdot \w_i$.
Note that the value of a policy $\pi$ under any task $\w$ can be expressed as $v^{\pi}_{\w} = \vpsi^\pi \cdot \w$.
%
Let $\Pi = \{\pi_i\}_{i=1}^{n}$ be a set of policies with corresponding SFs $\Psi = \{\vpsi^{\pi_i}\}_{i=1}^{n}$.
Based on the definition of GPI (Def.~\ref{def:gpi}) and the reward decomposition $r_{\w}(s,a,s') = \vphi(s,a,s')\cdot \w$, a \textit{generalized policy}, $\pigpi: \mathcal{S}\times\mathcal{W} \to \mathcal{A}$, can then be defined as follows:
\begin{equation}
\label{eq:gpi-sf}
\pigpi(s, \w; \Pi) \in \underset{a \in \mathcal{A}}{\arg\max} \,\, \underset{\pi \in \Pi}{\max} \,\, \vpsi^\pi(s,a) \cdot \w.
\end{equation}
\subsection{GPI via the Option Keyboard}\label{sec:ok}
\looseness-1
The \textit{Option Keyboard} (OK)~\citep{Barreto+2019,Barreto+2020} is a method that generalizes GPI (Eq.~\eqref{eq:sf-gpi}) by allowing agents to use a different weight vector at each time step. It does so by learning a meta-policy $\omega : \Sspace \to \Zspace$ that outputs weights, $\z\in\Zspace$, \footnote{Without loss of generality, let $\Zspace$ be the space of $d$-dimensional unit vectors, i.e., $\Zspace=\{\z \in \mathbb{R}^d \mid ||\z||_2 = 1\}$.} used when following the GPI policy $\pigpi(s, \z; \Pi)$ at each state $s$.
%
Intuitively, given a set of policies $\Pi$, $\omega$ modulates the behavior induced by $\pigpi$ by controlling the preferences $\omega(s)$ over features for each state $s$.
The OK policy, $\piok_{\omega}$, at each state $s$, is given by:
\begin{equation}\label{eq:ok-policy}
\pi_{\omega}^{\text{OK}}(s;\Pi) \defas \underset{a \in \mathcal{A}}{\arg\max} \,\, \underset{\pi \in \Pi}{\max} \,\, \vpsi^\pi(s,a) \cdot \omega(s).
\end{equation}
Note that, by construction, $\pi_{\omega}^{\text{OK}}(s;\Pi) = \pigpi(s,\omega(s);\Pi)$.
Intuitively, in the ``keyboard'' analogy, $\omega(s)$ (a ``chord'') combines base policies in $\Pi$ (``keys'') to generate more complex behaviors.
%
A key property of OK is that learning a policy over $\Zspace$ is often easier than over $\Aspace$, as solutions typically involve repeating the same action $\z$ across multiple timesteps, similar to \textit{temporally extended options}.
\section{Optimal Behavior Basis}\label{sec:problem-formulation}
\looseness-1
In this section, we formalize the problem of identifying a behavior basis---i.e., a set of base policies---that enables the OK to optimally solve all tasks within a given family of MDPs.
%
We are interested in incremental methods for constructing a behavior basis, $\Pi_k$, that provably enables the OK to solve any MDP in $\Mspace$. Specifically, we consider approaches where an agent first iteratively learns a behavior basis for solving some subset of tasks, $\mathcal{M}' \subset \mathcal{M}$.
%
Then, the method should be capable of leveraging GPI or OK to identify additional specialized policies for optimally solving novel, unseen tasks $M \not\in \mathcal{M}'$. This process should progressively expand the behavior basis until it converges to a small but sufficient set that guarantees zero-shot optimality.
%
The core problem investigated in this paper is, then, how to identify an optimal set of base policies that an agent should learn to facilitate transfer learning within specific classes of MDPs.
\looseness-1
Let $\Mlinear \subseteq \mathcal{M}$ be the---possibly infinite---set of MDPs associated with all linearly expressible reward functions. This set, typically studied in the SFs literature, can be defined as
%
\begin{equation}
\label{eq:linear-reward-mdps}
\Mlinear \defas \{ (\mathcal{S}, \mathcal{A}, p, r_{\vect{w}}, \mu, \gamma) \mid r_{\vect{w}} = \vphi \cdot \vect{w} \}.
\end{equation}
Each element $M\in\Mlinear$ (or, equivalently, its weight vector $\w$) is called a \textit{task}.
%
In what follows, we consider weight vectors that induce convex combinations of features; that is, $\mathcal{W} = \{\w \mid \sum_{i} w_i = 1$, $w_i \geq 0, \forall i\}$.
%
This is common practice, e.g., in the multi-objective RL literature~\citep{Hayes+2022}.
Our analyses naturally extend to \textit{linear} combinations of features---in App.~\ref{ap:linear-to-convex}, we show how linear-reward MDPs can be transformed into equivalent convex-reward MDPs (w.l.o.g.) so that our results and algorithms apply directly.
%
Existing algorithms capable of optimally solving all tasks in $\Mlinear$ often identify a set of policies, $\Pi_k=\{\pi_i\}_{i=1}^{n}$, such that their associated SF vectors, $\Psi=\{\vpsi^{\pi_i}\}_{i=1}^{n}$, form a \textit{convex coverage set} (CCS)~\citep{Alegre+2022}. A CCS is a set of SF vectors that allows the optimal value function, $v^{*}_{\w} = \vpsi^{\pi^{*}_{\w}} \cdot \w$, for any given task $\w\in\Wspace$, to be directly identified:\footnote{\looseness-1
% BCS ADD
A CCS may not be unique since tasks can have multiple optimal policies with distinct SF vectors. In what follows, CCS refers specifically to the minimal set satisfying Eq.~\eqref{eq:ccs}, which is uniquely defined~\citep{Roijers+2013}.
}
%
\begin{align}
\ccs & \defas
\{\vpsi^\pi \ | \ \exists \vect{w}\in\Wspace \ \text{s.t.}\ \forall \pi', \vpsi^\pi \cdot \vect{w} \geq \vpsi^{\pi'} \cdot \vect{w} \}
.
\label{eq:ccs}
\end{align}
Unfortunately, methods that compute the complete CCS for $\Mlinear$ do not scale to complex tasks since its size can grow exponentially with $d$---counting corner weights is equivalent to Vertex Enumeration~\citep{Roijers2016}.
%
Hence, it becomes critical to develop novel techniques capable of recovering all solutions induced by a CCS without incurring the cost of learning the corresponding full set of policies. This could be achieved, for instance, by methods capable of expressing all solutions in a CCS by combining policies in a behavior basis $\Pi_k$ smaller than the CCS; that is, $|\Pi_k| \leq |\ccs|$.
We start by extending the definition of the OK (Eq.~\eqref{eq:ok-policy}) to include the description of the task being solved, $\w\in\Wspace$, as one of its arguments.
%
That is, we consider meta-policies $\omega: \Sspace\times\Wspace\to\Zspace$, and OK policies defined as
\begin{equation}
\piok_{\omega}(s,\w; \Pi) \triangleq \argmax_{a\in\Aspace} \max_{\pi\in\Pi} \vpsi^\pi(s,a) \cdot \omega(s,\w).
\end{equation}
\looseness-1
The meta-policy $\omega(s, \w)$ enables the OK to decompose the optimal policy for a task $\w$ by dynamically assigning state-dependent linear weights to the SFs of its various base policies. Note that if $\omega(s,\w) = \w$ for all $s\in\Sspace$, we recover GPI.
%
Given an arbitrary set of base policies $\Pi_k$, we define the space of all policies expressible by the OK, $\Pi^{\text{OK}}(\Pi_k)$, and their associated SF vectors, $\Psi^{\text{OK}}(\Pi_k)$, respectively, as
\begin{align}
\nonumber
\Pi^{\text{OK}}(\Pi_k) \defas \{ \piok_{\omega}(\cdot,\w;\Pi_k) \mid \w \in \Wspace, \omega : \Sspace \times \Wspace \to \Zspace \} \textrm{ and} \
\Psi^{\text{OK}}(\Pi_k) \defas \{ \vpsi^\pi \mid \pi \in \Pi^{\text{OK}}(\Pi_k)\}.
\end{align}
We are now ready to mathematically define our goal:
\textbf{Goal}: Learn a set of policies (the \textit{behavior basis}) $\Pi_k$ and a meta-policy $\omega$ such that \textbf{(1)} $|\Pi_k| \leq |\ccs|$; and \textbf{(2)} $\piok_{\omega}(\cdot;\Pi_k)$ is optimal for any task $M \in \Mlinear$. The latter condition implies that $\ccs \subseteq \Psi^{\text{OK}}(\Pi_k)$; i.e., the OK is at least as expressive as a CCS.
\looseness-1
As observed by \citet{Barreto+2019,Barreto+2020}, learning a meta-policy, $\omega$, is often easier than learning base policies over the original action space, $\Aspace$. Consider, e.g., that if a task $\w\in\Wspace$ can be solved by switching between two base policies in $\Pi_k$, then an optimal meta-policy for $\w$, $\omega(\cdot,\w)$, only needs to output two vectors $\z$ (one for each base policy) for any given state. Thus, solving the goal above is bound to be more efficient than constructing a complete CCS since it requires learning fewer optimal base policies.
\section{Constructing an Optimal Behavior Basis}
\vspace{-0.13cm}
\looseness-1
We now introduce \textbf{O}ption \textbf{K}eyboard \textbf{B}asis (\textbf{OKB}), a novel method to solve the goal introduced in the previous section. The OKB (Alg.~\ref{algo:okb}) learns a set of base policies, $\Pi_k$, and a meta-policy, $\omega$, such that the induced OK policy $\piok_{\omega}(\cdot;\Pi_k)$ is provably optimal w.r.t. any given task $\w\in\Wspace$.
\looseness-1
The algorithm starts with a single base policy optimized to solve an arbitrary initial task (e.g., $\w{=}[1/d,...,1/d]^\top$) in its set of policies $\Pi$ (lines 1--2).
%
OKB's initial partial CCS, $\PsiOK$, and weight support set, $\Wsupport$, are initialized as empty sets (line 3). $\Wsupport$ will store the weights of tasks the meta-policy has been trained on.
%
At each iteration $k$, OKB carefully selects the weight vectors on which its meta-policy, $\omega$, will be trained so that the policies expressible by $\PsiOK$ iteratively approximate a CCS.
%
This process is implemented by OK-LS (Alg.~\ref{algo:okls}), an algorithm inspired by the SFs Optimistic Linear Support (SFOLS) method~\citep{Alegre+2022} but that operates over a meta-policy $\omega$ rather than the space of base policies.
\IncMargin{0.76em}
\begin{wrapfigure}{l}{0.5\textwidth}
\begin{minipage}{0.50\textwidth}
\begin{algorithm}[H]
\footnotesize{
\SetKwInput{KwInput}{Input}
\SetKwComment{Comment}{$\triangleright$\ }{}
\DontPrintSemicolon
\caption{\small{\textbf{O}ption \textbf{K}eyboard \textbf{B}asis (OKB)}}
\label{algo:okb}
\KwInput{MPC with features $\vphi(s,a) \in\mathbb{R}^d$}
$\pi_{\w}, \vpsi^{\pi_{\w}} \leftarrow \texttt{NewPolicy}( \w{=}\texttt{InitialTask()})$\;
$\Pi_0 \leftarrow \{\pi_{\w}\}; \PsiBase_0 \leftarrow \{\vpsi^{\pi_{\w}}\}$ \;
$\PsiOK_0 \leftarrow \{\}; \Wsupport_0 \leftarrow \{\}$ \;
Initialize meta-policy $\omega_0$ \;
\For{$k = 0,1,2,\dots$}{
\Comment*[l]{Update meta-policy $\omega$}
$\omega_{k+1}, \PsiOK_{k+1}, \Wsupport_{k+1} \leftarrow \texttt{OK-LS}(\omega_k,\PsiOK_k,\Wsupport_k,\Pi_k)$ \;
$\mathcal{C} \leftarrow \{\}$ \;
\For{$\w \in \textnormal{\texttt{CornerW}}(\PsiBase_k) \cup \textnormal{\texttt{CornerW}}(\PsiOK_{k+1})$}{
\Comment*[l]{Check if task $\w$ is solvable with OK and $\Pi_k$}
\If{$\pi^*_{\w}$ is not in $\Pi^{\textup{OK}}(\Pi_k)$}{
$\mathcal{C} \leftarrow \mathcal{C} \cup \{\w\}$ \;
}
}
\eIf{$\mathcal{C}$ is empty}{
\Comment*[l]{Found optimal basis $\Pi$ and meta-policy $\omega$}
\Return{$\omega_{k+1}, \Pi_k, \PsiOK_{k+1}$}
}{
\Comment*[l]{Learn a new base policy}
$\w \leftarrow$ select a task from $\mathcal{C}$ \;
$\pi_{\w}, \vpsi^{\pi_{\w}} \leftarrow \texttt{NewPolicy}(\w)$ \;
$\Pi_{k+1} \leftarrow \Pi_k \cup \{\pi_{\w}\}$ \;
$\PsiBase_{k+1} \leftarrow \PsiBase_k \cup \{\vpsi^{\pi_{\w}}\}$ \;
$\Pi_{k+1}, \PsiBase_{k+1} \leftarrow \texttt{RemoveDominated}(\Pi_{k+1}, \PsiBase_{k+1})$ \;
}
}
}
\end{algorithm}
\end{minipage}
\begin{minipage}{0.50\textwidth}
\begin{algorithm}[H]
\footnotesize{
\SetKwInput{KwInput}{Input}
\SetKwComment{Comment}{$\triangleright$\ }{}
\DontPrintSemicolon
\caption{\small{OK - Linear Support ($\texttt{OK-LS}$)}}
\label{algo:okls}
\KwInput{Meta-policy $\omega$, partial CCS $\PsiOK$, weight support $\Wsupport$, base policies $\Pi_k$.}
\While{$\textnormal{\texttt{True}}$}{
$\Wspace^{\text{corner}} \leftarrow \texttt{CornerW}(\PsiOK) \setminus \Wsupport$ \;
\If{$\Wcorner$ is empty}{
\Return{$\omega, \PsiOK, \Wsupport$}
}
$\Wsupport \leftarrow \Wsupport \cup \Wspace^{\text{corner}}$ \;
$\omega \leftarrow \texttt{TrainOK}(\omega, \Wsupport, \Pi_k)$ \hfill\Comment{(Alg.~\ref{algo:train-ok})}
$\PsiOK \leftarrow \{ \vpsi^{\piok_{\omega}(\cdot,\w;\Pi_k)} \mid \w \in \Wsupport\}$ \;
$\PsiOK \leftarrow \texttt{RemoveDominated}(\PsiOK)$ \;
}
}
\end{algorithm}
\end{minipage}
\vspace{-20pt} % BCS ADDED
\end{wrapfigure}
\DecMargin{0.76em}
Next, in lines 8--10, OKB identifies a set of candidate tasks, $\mathcal{C}$, that the OK cannot optimally solve with its current set of policies, $\Pi_k$; i.e., a set of tasks for which $\pi_{\w}^{*}$ is not included in $\Pi^{\text{OK}}(\Pi_k)$. We discuss how to check this condition in Section~\ref{sec:check-ok-optimality}
%
If $\mathcal{C}$ is empty (line 11), the OKB terminates and returns a meta-policy ($\omega$) and base policies ($\Pi_k$) capable of ensuring that $\ccs \subseteq \PsiOK(\Pi_k)$. This is due to Thm.~\ref{th:ok-ccs}, which we discuss later.
If $\mathcal{C}$ is not empty, OKB selects a task from it and adds a new corresponding optimal base policy to its behavior basis, $\Pi_k$ (lines 14--17).
%
In line 18, \texttt{RemoveDominated} removes redundant policies---policies that are not strictly required to solve at least one task.
\looseness-1
In Algorithms~\ref{algo:okb} and \ref{algo:okls}, the function $\texttt{CornerW}(\Psi)$ takes a set of SF vectors, $\Psi$, as input and returns a set of \textit{corner weights} (see App.~\ref{sec:corner-weights}).
%
Intuitively, both Alg.~\ref{algo:okb} and Alg.~\ref{algo:okls} rely on the fact that \textit{(i)} a task $\w\in\Wspace$ that maximizes $\Delta(\w) \defas v^{*}_{\w} - \max_{\pi\in\Pi_k} v^{\pi}_{\w}$ is guaranteed to be a corner weight, and \textit{(ii)} if OKB has an optimal policy for all corner weights, then it must have identified a CCS. These results were shown in prior work on constructing a CCS~\citep{Roijers2016,Alegre+2022} and can also be proven using the fundamental theorem of linear programming.
\looseness-1
OK-LS (Alg.~\ref{algo:okls}) trains the meta-policy, $\omega$, on selected tasks $\Wcorner$ (line 2) using the base policies $\Pi_k$, so that is partial CCS, $\PsiOK$, iteratively approximates a CCS.
It stores the tasks $\w$ it has already trained on, as well as the corner weights of the current iteration, in the weight support set, $\Wsupport$ (line 5).
%
If $\Wcorner$ is empty in a given iteration, no tasks remain to be solved, and the algorithm returns the updated meta-policy.
%
Otherwise, OK-LS adds the corner weights $\Wcorner$ to $\Wsupport$ and trains the meta-policy $\omega$ on the tasks $\w \in \Wsupport$ using the \texttt{TrainOK} subroutine (Alg.~\ref{algo:train-ok}).
%
Finally, in line 7, OK-LS computes the SF vectors of the policies induced by OK for each $\w \in \Wsupport$, updating its partial CCS, $\PsiOK$.\footnote{\looseness-1
%
Notice that SFs only need to be computed for vectors newly added to $\Wsupport$ in the current iteration (lines 6-7).
}
%
In App.~\ref{sec:train-ok}, we discuss how to train $\omega$ using an actor-critic RL method (Alg.~\ref{algo:train-ok}).
%
Finally, note that to accelerate the learning of $\omega$, the corner weights in $\Wcorner$ can be prioritized, similar to \citet{Alegre+2022}.
\subsection{Condition for OK Optimality}\label{sec:check-ok-optimality}
\looseness-1
In this section, we address the following question: \textit{Given an arbitrary task with reward function $r$ and a set of base policies $\Pi_k$, is $\pi^{*}_{r} \in \Pi^{\textup{OK}}(\Pi_k)$?} In other words, does a meta-policy $\omega$ exist such that the OK policy, $\piok{\omega}(\cdot;\Pi_k)$, can represent the optimal policy $\pi_{r}^{*}$?
Determining whether this holds is essential for implementing the OKB step in line 9.
\begin{proposition}\label{th:existence-omega}
Let $\Pi_k=\{\pi_i\}_{i=1}^{n}$ be a set of base policies with corresponding SFs $\Psi=\{\vpsi^{\pi_i}\}_{i=1}^{n}$. Given an arbitrary reward function $r$, an optimal OK policy $\piok_{\omega}(\cdot;\Pi)$ can only exist if there exists an OK meta-policy, $\omega : \Sspace \to \Zspace$, such that for all $s\in\Sspace$,
\begin{equation}\label{eq:existence-omega}
\argmax_{a\in\Aspace} \max_{\pi\in\Pi_k} \vpsi^\pi(s,a) \cdot \omega(s) = \argmax_{a\in\Aspace} q^*_{r}(s,a).
\end{equation}
\end{proposition}
%
This proposition provides a sufficient condition for the existence of an optimal OK meta-policy $\omega$ for a given reward function $r$. Intuitively, it implies that the OK policy should be able to express all optimal actions through $\omega$ and $\Psi$. Next, we show how to verify this condition without requiring access to $q^{*}_{r}$.
Let $\Pi_k$ be a set of base policies, $\omega$ be an OK meta-policy, and $q^{\omega}_{r}(s,\z) = r(S_t, A_t) + \gamma \Ex_{\omega}[q^{\omega}_{r}(S_{t+1},\omega(S_{t+1})) \mid S_t=s, A_t=\pigpi(S_t,\z;\Pi_k)]$ be the meta-policy's action-value function for task $r$.
Given a state-action pair $(s,a)$, let the \textit{advantage function} of $\omega$ for executing action $a$ in state $s$ be
\begin{align}\label{eq:advantage-optimality}
A^{\omega}_{r}(s,a) \defas r(s,a) + \gamma\Ex_{\omega}[q^{\omega}_{r}(S_{t+1},\omega(S_{t+1})) \mid S_t{=}s, A_t{=}a ] - q^{\omega}_{r}(s,\omega(s)).
\end{align}
It is well-known that an optimal policy's advantage function is zero when evaluated at an optimal action; i.e., $A^{\pi^*}(s,a^*) = 0$ for all $s\in\Sspace$.
Thm.~\ref{th:advantage} uses this insight to introduce a principled way to verify if $\pi^{*}_{r} \in \Pi^{\textup{OK}}(\Pi_k)$.
\begin{theorem}\label{th:advantage}
Let $\Pi_k$ be a set of base policies and $\omega^*$ be a meta-policy trained to convergence to solve a given task $r$. If $A^{\omega^*}_{r}(s,a) > 0$ for some $(s,a)$, then action $a$ is not expressible by $\piok_{\omega^*}(\cdot;\Pi_k)$. Consequently, the OK with base policies $\Pi_k$ cannot represent an optimal policy for $r$.
\end{theorem}
%\textcolor{blue}{
Note that the procedure for verifying whether the OK is optimal for a given task $r$ (Thm.~\ref{th:advantage}) is a sufficient, but not necessary, condition to guarantee optimality. It serves two key purposes: \textit{(i)} to avoid training on tasks $\omega$ whose optimal policies have either already been learned or that can be reconstructed from $\Pi_k$ and $\omega$, and \textit{(ii)} to prioritize which tasks to train on in order to accelerate learning. In other words, the procedure for assessing OK optimality is primarily a tool to reduce training effort; our method does not rely on an exact test to ensure the correctness of the behavior basis it constructs.
For a detailed discussion of how Thm.~\ref{th:advantage} can be applied in practical implementations of OKB to verify whether $\pi^{*}_{r} \in \Pi^{\textup{OK}}(\Pi_k)$, see App.~\ref{sec:checking-ok-condition}.
%}
\subsection{Theoretical Results}\label{sec:theoretical-results}
In this section, we present the theoretical guarantees of OKB (Alg.~\ref{algo:okb}). Proofs of the theorems can be found in App.~\ref{ap:proofs}.
\begin{theorem}\label{th:ok-ccs}
OKB returns a set of base policies, $\Pi_k$, and a meta-policy, $\omega$, such that $|\Pi_k| \leq |\ccs|$ and $\piok_{\omega}(\cdot;\Pi_k)$ is optimal for any task $M \in \Mlinear$. This implies that $\ccs \subseteq \Psi^{\text{OK}}(\Pi_k)$; i.e., the OK matches or exceeds the expressiveness of a CCS and ensures zero-shot optimality for all $M \in \Mlinear$.
\end{theorem}
%
This is a key result of this paper: OKB can provably identify a behavior basis, $\Pi_k$, that enables an option keyboard to \textit{optimally solve} any task $M \in \Mlinear$.
%
Furthermore, it is capable of ensuring zero-shot optimality using a set of base policies potentially smaller than a CCS. This is the first method to offer such a guarantee---until now, all existing approaches required the number of base policies to be \textit{strictly equal} to the size of the CCS. This limitation posed a major scalability challenge: computing a full CCS becomes intractable as the number of reward features $d$ grows, since its size can grow exponentially with $d$---counting corner weights is equivalent to Vertex Enumeration~\citep{Roijers2016}. OKB overcomes this bottleneck by avoiding the need to construct a full CCS, thereby reducing the total cost of computing optimal policies for all $M \in \Mlinear$. See below for more details.
\begin{proposition}\label{th:non-linear-tasks}
Let $\Pi_k$ be the set of policies learned by OKB (Alg.~\ref{algo:okb}) when solving tasks that are linearly expressible in terms of reward features $\vphi(s,a) \in \mathbb{R}^d$. Let $r$ be an arbitrary reward function that is non-linear with respect to $\vphi$. Suppose the optimal policy for $r$ can be expressed by alternating, as a function of the state, between policies in the CCS induced by $\vphi$. That is, suppose that for all $s \in \Sspace$, there exists a policy $\pi^*_{\w}$ (optimal for some $\w \in \Wspace$) such that $\pi^*_r(s) = \pi^*_{\w}(s)$. Then, the OK can represent the optimal policy for $r$ using the set of base policies $\Pi_k$; i.e., $\pi^*_r \in \Pi^{\textup{OK}}(\Pi_k)$.
\end{proposition}
This deepens the theoretical understanding of the OKB by showing it is strictly more expressive than existing methods. In particular, it allows us to characterize \textbf{(1)} conditions under which the OKB matches the expressiveness of a CCS while relying on strictly \emph{fewer} policies; and \textbf{(2)} a broad class of \textit{non-linear} tasks that can be optimally solved by an OK using the behavior basis learned by OKB.
First, Prop.~\ref{th:non-linear-tasks} establishes the optimality of the OKB in transfer learning scenarios where related tasks can be solved by reusing parts of existing solutions---that is, the widely studied setting in which transfer is possible due to shared structure across tasks.
%
Concretely, the OKB is guaranteed to optimally solve, in a zero-shot manner, tasks whose solutions decompose into sequences of sub-policies. \emph{Building on this result and Thm.~\ref{th:ok-ccs}, we formally characterize, in App.~\ref{ap:strictly-smaller}, the conditions under which the OKB is as expressive as a CCS while relying on strictly \textbf{fewer} policies}.
\looseness-1
Second, Proposition~\ref{th:non-linear-tasks} precisely characterizes a class of non-linear tasks that can be optimally solved by an OK using the behavior basis learned by OKB.
%
Intuitively, OKB can solve any task whose optimal policy can be constructed by alternating between the optimal policies for tasks in $\Mlinear$, even when the task itself is non-linear.
%
\textit{This implies, importantly, that while existing methods that compute a CCS are limited to linearly expressible tasks, OKB can handle a broader class of non-linear tasks.}
%
Furthermore, as discussed above, the OKB achieves this \textit{without} having to learn all policies in the CCS. We further discuss the properties of OK under non-linear reward functions in App.~\ref{sec:beyond-linear}.
To state our next theoretical result, we first recall Thm. 2 by \citet{Barreto+2017}:
\begin{theorem}[\citet{Barreto+2017}]\label{th:ok-gpi-theorem}
Let $\Pi{=}\{\pi_i\}_{i=1}^{n}$ be a set of optimal policies w.r.t. tasks $\{\w_i\}_{i=1}^{n}$. Let $\{\hat{\vpsi}^{\pi_i}\}_{i=1}^{n}$ be approximations to their SFs and $\w\in\Wspace$ be a task. Let $|q^{\pi_i}_{\w}(s,a) - \hat{q}^{\pi_{i}}_{\w}(s,a)| \leq \epsilon \text{ for all } (s,a) \in\Sspace{\times}\Aspace, \text{ and }\, \pi_i\in\Pi$.
Let $\vphi_{\mathrm{max}} \defas \max_{s,a}||\vphi(s,a)||$.
%
Then, it holds that:
\begin{align}
q^{*}_{\w}(s,a) - q^{\pigpi}_{\w}(s,a)
\leq \frac{2}{1-\gamma}\left( \vphi_{\mathrm{max}} \min_{i}||\w - \w_i|| + \epsilon \right) \quad \textrm{for all } (s,a)\in\Sspace{\times}\Aspace.
\end{align}
\end{theorem}
\looseness-1
Thm~\ref{th:ok-gpi-theorem} describes the optimality gap of GPI policies, how it depends on the available base policies, and how it is affected by function approximation errors.
%
It does not, however, characterize the optimality gap of option keyboards, which generalize GPI policies.
%
We introduce and highlight two generalizations of this theorem: \textbf{(1)} Since the OK generalizes GPI policies, it follows that $q^{*}_{\w}(s,a) - q^{\piok_{\omega}}_{\w}(s,a) \leq q^{*}_{\w}(s,a) - q^{\pigpi}_{\w}(s,a)$; and \textbf{(2)} when Eq.~\eqref{eq:existence-omega} (Prop.~\ref{th:existence-omega}) is satisfied, $q^{*}_{\w}(s,a) - q^{\piok_{\omega}}_{\w}(s,a)=0$.
%
That is, OK can entirely avoid optimality gaps from approximation errors in the base policies' SFs. Thus, OK and OKB are not only more expressive than GPI (Thm.~\ref{th:ok-ccs}) but also naturally lead to transfer learning strategies that are significantly more robust to approximation errors than GPI.
\section{Experiments}\label{sec:experiments}
\begin{wrapfigure}{r}{0.45\textwidth}
\centering
\vspace{-0.5cm}
\begin{minipage}{0.45\textwidth}
\includegraphics[width=0.3\textwidth,valign=t]{figs/minecart-domain.png}\hfill
\includegraphics[width=0.3\textwidth,valign=t]{figs/fetchpickandplace.png}\hfill
\includegraphics[width=0.3\textwidth,valign=t]{figs/item-collection2.png}\\[0.5ex]
\includegraphics[width=1.0\textwidth]{figs/sf-highway.png}
\caption{Domains used in the experiments: Minecart, FetchPickAndPlace, Item Collection, and Highway.}
\label{fig:domains}
\end{minipage}
\vspace{-0.55cm}
\end{wrapfigure}
\looseness-1
We now empirically evaluate OKB and investigate the following research questions:
\textbf{Q1}: Can OKB approximate a CCS more effectively while requiring fewer base policies than competing methods?
\textbf{Q2}:
%
Does OKB’s performance advantage grow with problem complexity, as measured by the number of reward features $d$?
%
\textbf{Q3}:
%
Can the base policies learned by OKB be used to solve tasks with non-linear reward functions, under the conditions in Prop.~\ref{th:non-linear-tasks}?
\looseness-1
Fig.~\ref{fig:domains} depicts the domains used in our experiments. To handle the high-dimensional state space of these domains, we learn base policies using a Universal SF Approximator (USFA)~\citep{Borsa+2019}.
%
Furthermore, rather than learning a separate SF for each base policy $\pi \in \Pi_k$, we train a single USFA, $\vpsi(s,a,\w)$, conditioned on task vectors $\w$, such that $\pi_{\w}(s) \approx \argmax_{a\in\Aspace}\vpsi(s,a,\w)\cdot\w$.
%
%
Each call to \texttt{NewPolicy}($\w$) in Alg.\ref{algo:okb} trains the USFA to solve task $\w$.
%
We employ USFAs in our experiments primarily to avoid the memory and computational costs of training a separate neural network for each policy’s SFs. Moreover, USFAs typically achieve a level of approximation comparable to---or better than---that of independently trained networks, since in both cases generalization depends mainly on the architecture, capacity, and distribution of training data. %}
%
See App.~\ref{ap:experiments-details} for more details.
In all experiments, we report the mean normalized return of each method (normalized with respect to the minimum and maximum returns observed for a given task) along with the 95\% bootstrapped confidence interval over 15 random seeds.
%
We compare OKB to SFOLS~\citep{Alegre+2022}, a method that identifies a CCS using GPI policies, and to OKB-Uniform, a variant of OKB that selects tasks uniformly from $\Wspace$ in line 14 of Alg.~\ref{algo:okb}, rather than from $\mathcal{C}$.
%
An iteration corresponds to each method training a base policy within a fixed budget of environment interactions.
%
To ensure fair comparisons, since OKB must also train a meta-policy while SFOLS does not, we restrict OKB to using only half of its budget for learning base policies, allocating the other half to training the meta-policy (Alg.~\ref{algo:okls}). This makes the comparison more conservative for OKB, as it has fewer interactions available to learn base policies.
\begin{wrapfigure}{l}{0.45\textwidth}
\centering
\vspace{-0.1cm}
\begin{minipage}{0.45\textwidth}
\centering
\begin{subfigure}[t]{0.9\linewidth}
\includegraphics[width=\linewidth]{figs/minecart.pdf}
\subcaption{Minecart}
\label{fig:minecart-results}
\end{subfigure}
%\vspace{0.5em}
\begin{subfigure}[t]{0.9\linewidth}
\includegraphics[width=\linewidth]{figs/highway.pdf}
\subcaption{Highway}
\label{fig:highway-results}
\end{subfigure}
\caption{\looseness-1
Mean normalized return per iteration for each method on a set of test tasks.}% in (a) Minecart and (b) Highway.}
\label{fig:mean-performance}
\end{minipage}
\vspace{-0.5cm}
\end{wrapfigure}
\looseness-1
Fig.~\ref{fig:mean-performance} depicts each method's mean return over a set of test tasks, $M \in \Mlinear$,\footnote{To generate test task sets $\Wspace' \subset \Wspace$ for different values of $d$, we used the method introduced by \citet{Takagi+2020}, which produces uniformly spaced weight vectors in $\Wspace$.} as a function of the iteration number (i.e., the number of base policies learned).
%
We report results for the Minecart domain---a classic multi-objective RL problem---and the Highway domain. Both domains have reward functions defined by $d = 3$ reward features, which are detailed in App.~\ref{ap:domains}.
%
These results demonstrate that OKB achieves strong performance across test tasks with a small number of base policies, positively answering research question Q1: OKB can approximate a CCS more effectively using fewer base policies. This is particularly evident in Fig.~\ref{fig:minecart-results}, where OKB reaches near-optimal performance with just 2--3 base policies. While OKB-Uniform also performs well in the Minecart domain, its lower performance in the Highway domain (Fig.~\ref{fig:highway-results}) highlights the importance of expanding the set of base policies by carefully selecting promising tasks—defined by corner weights—for training. Across both domains, OKB consistently outperforms all competing methods.
\looseness-1
To investigate \textbf{Q2}, we evaluate each method in the FetchPickAndPlace domain with varying numbers of reward features, $d\in\{2,4,6,8\}$. In this domain, an agent controls a robotic arm that must grasp a block and move it to a specified location. Each of the $d$ reward features represents the block's distance to a different target location (red spheres in Fig.~\ref{fig:domains}).
%
Fig.~\ref{fig:pickandplace-results} shows that as the number of target locations ($d$) increases, the performance gap between OKB and SFOLS increases significantly---positively answering \textbf{Q2}.
%
\footnote{
%
Importantly, to the best of our knowledge, most state-of-the-art methods are typically evaluated on problems with dimensionality at most $d=4$ (e.g., SFOLS and SIP). In our experiments, by contrast, we double the number of reward features, thus showing that our method remains effective and continues to outperform all baselines even on problems an order of magnitude more challenging.%}
}
%
Intuitively, OKB focuses on tasks where the OK cannot express optimal actions (see Thm.~\ref{th:advantage}). By learning the corresponding optimal policies, OKB quickly identifies a behavior basis that enables the OK to solve tasks across the entire space of task vectors $\Wspace$.
%
Conversely, the gap between OKB and OKB-Uniform decreases since a larger number of reward features enhances the expressivity of the OK (Prop.~\ref{th:existence-omega}).
%
The total computation time required by OKB to train a meta-policy and base policies for zero-shot optimality is consistently comparable to or lower than that of competing methods, as it learns and evaluates fewer policies at inference time.
\looseness-1
Next, we investigate \textbf{Q3} by conducting an experiment similar to the one proposed by~\citet{Alver&Precup2021}. We compare OKB to relevant competitors in an environment with non-linear reward functions.
%
The Item Collection domain (Fig.~\ref{fig:domains}; top right) is a $10 \times 10$ grid world where agents must collect two types of items. Reward features are indicator functions that signal whether the agent has collected a particular item in the current state. This domain requires function approximation due to the combinatorial number of possible states.
%
After OKB learns a behavior basis $\Pi_k$, the agent trains a meta-policy $\omega: \Sspace \to \Zspace$ to solve a task with a \textit{non-linear reward function}.
%
Specifically, this is a sequential task where the agent must collect all instances of one item type before collecting any items of another type.
%
We compare OKB with SIP~\citep{Alver&Precup2021},\footnote{\looseness-1
%
SIP was not part of previous experiments as it assumes independent reward features~\citep{Alver&Precup2021}.}
%
a method that learns $d$ base policies---each maximizing a specific reward feature---and a meta-policy $\omega$ over a discrete set of weights in $\Wspace$.
%
We also compare OKB against a baseline DQN agent that learns tasks from scratch and against GPI policies (horizontal blue lines), which are optimal for either exclusively prioritizing one item type or assigning equal importance to both.
%
Fig.~\ref{fig:sequential-task} shows the mean total reward achieved by each method in this non-linear task as a function of the total number of environment interactions required to train the meta-policies. Both OKB and SIP successfully solved the task, with OKB doing so more quickly.
%
The DQN baseline failed to solve the task as directly learning a policy over $\Aspace$ is significantly more difficult.
\begin{figure*}[!t]
\centering
\includegraphics[width=0.24\linewidth]{figs/pickandplaced2.pdf}
\includegraphics[width=0.24\linewidth]{figs/pickandplaced4.pdf}
\includegraphics[width=0.24\linewidth]{figs/pickandplaced6.pdf}
\includegraphics[width=0.24\linewidth]{figs/pickandplaced8.pdf}
\caption{\looseness-1
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.}
\vspace{-0.5cm}
\label{fig:pickandplace-results}
\end{figure*}
\begin{figure}[ht]
\centering
\begin{minipage}{0.4\textwidth}
\centering
\includegraphics[width=1.0\textwidth]{figs/sequential-item-collection.pdf}
\caption{Mean return in the Item Collection domain under a \textit{non-linear} reward function.}
\label{fig:sequential-task}
\end{minipage}
\hfill
\centering
\begin{minipage}{0.585\textwidth}
\centering
\hspace{0.55cm}
\includegraphics[width=0.28\textwidth]{figs/frame_0004.png}
\includegraphics[width=0.28\textwidth]{figs/frame_0008.png}
\includegraphics[width=0.28\textwidth]{figs/frame_0023.png}
\includegraphics[width=0.95\textwidth]{figs/meta-policy-ok.pdf}
\caption{
\looseness-1
Continuous actions from the meta-policy on a sample task. OKB selects base policies in a temporally consistent way.}%, suggesting they encode useful temporally extended behaviors.}
\label{fig:meta-policy}
\end{minipage}
%\vspace{-0.5cm}
\end{figure}
Finally, we examine the qualitative behavior of OKB policies by visualizing the continuous actions produced by the meta-policy $\omega(s, \w)$ while solving a sample task in the Highway domain, where the agent must prioritize driving fast and staying in the rightmost lane.
%
In Fig.~\ref{fig:meta-policy}, the color of each timestep’s column represents the selected base policy.
%
The agent initially accelerates as quickly as possible by following base policy $\pi_1$ for the first 10 timesteps while the rightmost lane is occupied.
%
At timestep $t = 10$, the agent turns right
%($a{=}2$)
and switches to base policy $\pi_2$, which is optimized for driving in the rightmost lane. At $t = 20$, $\pi_1$ becomes active again, allowing the agent to accelerate while staying in the rightmost lane.
%
Note that while the meta-policy $\omega$ is relatively smooth, the base policies it selects are complex. In particular, although $\z_t$ remains approximately constant during the first 10 timesteps, the underlying primitive actions $a \in \Aspace$ (Fig.~\ref{fig:meta-policy}; black numbers at the top of each column) continuously alternate between turning left or right, accelerating, and idling.
%
Finally, notice that OKB selects base policies in a temporally consistent manner, suggesting that it learns to identify temporally extended behaviors---akin to \textit{options}~\citep{Sutton+1999}---that help solve the task.
\section{Related Work}
\looseness-1
\textbf{OK and GPI.} Previous works have extended the OK in different ways. \citet{Carvalho+2023neurips} introduced a neural network architecture for jointly learning reward features and SFs, while \citet{Machado+2023} proposed using the OK with base policies that maximize Laplacian-based reward features.
%
However, unlike our work, these methods do not provide theoretical guarantees on optimality for solving specific families of tasks.
%
Other works have introduced GPI variants for risk-aware transfer~\citep{Gimelfarb+2021}, mitigating function approximation errors~\citep{Kim+2022neurips}, planning with approximate models~\citep{Alegre+2023neurips}, planning to solve tasks defined by finite state automata~\citep{Kuric+2024}, and combining non-Markovian policies~\citep{Thakoor+2022}.
%
These approaches improve GPI in ways that are orthogonal to our contribution and could potentially be combined with our method.
\looseness-1
\textbf{Learning behavior basis.} Previous works have addressed the problem of learning an effective behavior basis for transfer learning within the SF framework~\citep{Nemecek&Parr2021,Zahavy+2021}.
%
The method introduced by \citet{Alver&Precup2021} assumes that reward features are independent (i.e., can be controlled independently) and that the MDP's transition function is deterministic. These assumptions often do not hold in complex RL domains, including most of those studied in Section~\ref{sec:experiments}.
%
\citet{Alegre+2022} proposed a method that learns a set of policies corresponding to a CCS and combines them with GPI. In contrast, we use option keyboards, which enable a broader range of policies to be expressed, allowing our method to approximate a CCS more effectively with a smaller behavior basis.
\looseness-1
\textbf{Learning features.} While we focused on identifying an optimal behavior basis given a particular set of reward features, other works have addressed the complementary problem of learning more expressive reward features~\citep{Touati&Ollivier2021,Carvalho+2023,Chua+2024}.
%
A promising future research direction is to combine OKB with methods such as the Forward-Backward representation~\citep{Touati+2023} to construct an optimal behavior basis under learned reward features.
\section{Conclusions}\label{sec:conclusions}
\looseness-1
We introduced OKB, a principled method with strong theoretical guarantees for identifying the optimal behavior basis for the Option Keyboard (OK).
%
Our theoretical results, supported by thorough empirical analysis, show that OKB significantly reduces the number of base policies required to achieve zero-shot optimality in new tasks compared to state-of-the-art methods. In particular, OKB constructs a behavior basis efficiently by carefully selecting base policies that iteratively improve its meta-policy’s approximation of the CCS.
%
We empirically evaluate OKB in challenging high-dimensional RL problems, demonstrating that it consistently outperforms GPI-based approaches. Notably, its performance advantage becomes increasingly pronounced as the number of reward features grows.
%
Finally, we prove that OKB's expressivity surpasses that of a CCS, enabling it to optimally solve specific classes of non-linear tasks.
%
Among several possible research directions, one promising direction is to extend OKB with temporally extended meta-policies that incorporate learned termination conditions. The resulting expressiveness could further expand OKB’s flexibility and efficiency in reusing behaviors across tasks.
%
Furthermore, recently identified mathematical connections between multi-objective RL (MORL) and SFs~\citep{Alegre+2022} open the door to adapting OKB to integrate with MORL techniques, paving the way for unified representations that improve both sample efficiency and generalization across objectives.
%
\section*{Acknowledgments}
This work is partially supported by CAPES (Coordena\c c\~ao de Aperfei\c coamento de Pessoal de N\'ivel Superior - Brazil, Finance Code 001).
Ana Bazzan is partially supported by CNPq under grant number 304932/2021-3.
\newpage
{\small
\bibliography{bibs/bib}
%\bibliographystyle{plain}
\bibliographystyle{apalike}
}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\newpage
\appendix
\section*{\LARGE{{Appendix}}}
\vspace{0.25cm}
\section{Proofs}\label{ap:proofs}
\subsection{Proof of Proposition~\ref{th:existence-omega}}
\begin{proposition*}[\ref{th:existence-omega}]
Let $\Pi_k=\{\pi_i\}_{i=1}^{n}$ be a set of base policies with corresponding SFs $\Psi=\{\vpsi^{\pi_i}\}_{i=1}^{n}$. Given an arbitrary reward function $r$, an optimal OK policy $\piok_{\omega}(\cdot;\Pi)$ can only exist if there exists an OK meta-policy, $\omega : \Sspace \to \Zspace$, such that for all $s\in\Sspace$,
\begin{equation}\label{eq:existence-omega2}
\argmax_{a\in\Aspace} \max_{\pi\in\Pi_k} \vpsi^\pi(s,a) \cdot \omega(s) = \argmax_{a\in\Aspace} q^*_{r}(s,a).
\end{equation}
\begin{proof}
The proof follows directly from the definition of a deterministic optimal policy for a given reward function $r$.
Recall that $\pi^*_{r}(s) \in \argmax_{a\in\Aspace} q^*_{r}(s,a)$ and $\piok_{\omega}(s;\Pi) \in \argmax_{a\in\Aspace} \max_{\pi\in\Pi} \vpsi^\pi(s,a) \cdot \omega(s)$.
If for some state $s\in\Sspace$, there is no value of $\z = \omega(s)$ such that Eq.~\eqref{eq:existence-omega2} holds, then it is not possible to define a meta policy $\omega$ such that $\piok_{\omega}(s;\Pi) = \pi^*_{r}(s)$.
\end{proof}
\end{proposition*}
\subsection{Proof of Theorem.~\ref{th:advantage}}
\begin{theorem*}[\ref{th:advantage}]
Let $\Pi_k$ be a set of base policies and $\omega^*$ be a meta-policy trained to convergence to solve a given task $r$. If $A^{\omega^*}_{r}(s,a) > 0$ for some $(s,a)$, then action $a$ is not expressible by $\piok_{\omega^*}(\cdot;\Pi_k)$. Consequently, the OK with base policies $\Pi_k$ cannot represent an optimal policy for $r$.
\begin{proof}
First, recall the definitions of $q^{\omega}_{r}(s,\z)$ and $A^{\omega}_{r}(s,a)$:
\begin{align}
q^{\omega}_{r}(s,\z) &\defas r(S_t, A_t) + \gamma \Ex_{\omega}[q^{\omega}_{r}(S_{t+1},\omega(S_{t+1})) \mid S_t=s, A_t=\pigpi(S_t,\z;\Pi_k)],
\\
A^{\omega}_{r}(s,a) &\defas r(S_t,A_t) + \gamma\Ex_{\omega}[q^{\omega}_{r}(S_{t+1},\omega(S_{t+1})) \mid S_t=s, A_t=a ] - q^{\omega}_{r}(s,\omega(s)).
\end{align}
$A^{\omega}_{r}(s,a)$ measures the relative benefit of executing action $a$ in state $s$ compared to following the current OK policy $\piok_{\omega}$.
If $A^{\omega}_r(s, a) > 0$, then executing $a$ in $s$ provides a \textit{higher expected return} than the action currently selected by $\piok_{\omega}$.
Hence, if there exists an $(s, a)$ such that $A^{\omega}_r(s, a) > 0$, this implies that the OK policy fails to represent at least one optimal action, meaning it is suboptimal.
Recall from Prop.~\ref{th:existence-omega} that for the OK policy to be optimal, it must express the optimal action $a^*$ for every $s$:
\[
\arg\max_{a \in A} \max_{\pi \in \Pi_k} \vpsi^{\pi}(s, a) \cdot \omega(s) = \arg\max_{a \in A} q^*_r(s, a).
\]
Since $\omega^*$ is trained until convergence to solve $r$, it has reached a stable policy. However, if there exists an $(s, a)$ such that $A^{\omega^*}_r(s, a) > 0$, then
\[
q^*_r(s, a) > q^{\omega^*}_{r}(s, \omega^*(s)).
\]
This means that in at least one state-action pair, the OK policy does not select an action that maximizes the expected return.
%
Since $\piok_{\omega^*}$ is restricted to the actions expressible via its base policies $\Pi_k$, this implies that
\[
a^* \not\in \left\{\arg\max_{a} \max_{\pi \in \Pi_k} \vpsi^{\pi}(s, a) \cdot \omega^*(s) \right\}.
\]
Thus, \textbf{no possible weighting of the base policies can recover the optimal action} in state $s$.
Since at least one optimal action $a^*$ is missing from the action space of the OK policy, it follows that $\pi^{OK}_{\omega^*}$ is not an optimal policy for $r$, and the task $r$ cannot be optimally solved using the given base policies $\Pi_k$. Therefore, to guarantee optimality, additional base policies must be introduced into $\Pi_k$.
\end{proof}
\end{theorem*}
\subsection{Proof of Theorem.~\ref{th:ok-ccs}}
\begin{lemma}[\citet{Alegre+2022}]\label{lemma:corner-weights}
Let $\Psi_k \subseteq \ccs$ be a subset of the CCS.
If, for all corner weights $\w \in \texttt{\textup{CornerW}}(\Psi_k)$, it holds that
$\max_{\vpsi^\pi \in \Psi_k} \vpsi^\pi \cdot \w = v^{*}_{\w}$,
then $\Psi_k$ contains all elements of the CCS, i.e.,
$\ccs \subseteq \Psi_k$.
\end{lemma}
This lemma follows directly from the theoretical guarantees of previous methods that construct a CCS, e.g., OLS and SFOLS~\citep{Roijers2016,Alegre+2022}.
Intuitively, this result states that to check if a given set of SF vectors, $\Psi_k$, is a CCS, it is sufficient to check if it contains an optimal policy for all of its corner weights. If this is the case, then it is not possible to identify some alternative weight vector $\w\in\Wspace$ for which it does not contain an optimal policy.
\begin{theorem*}[\ref{th:ok-ccs}]
OKB returns a set of base policies, $\Pi_k$, and a meta-policy, $\omega$, such that $|\Pi_k| \leq |\ccs|$ and $\piok_{\omega}(\cdot;\Pi_k)$ is optimal for any task $M \in \Mlinear$. This implies that $\ccs \subseteq \Psi^{\text{OK}}(\Pi_k)$; i.e., the OK matches or exceeds the expressiveness of a CCS.
\begin{proof} We start the proof by stating two assumptions on which the theorem relies: \\ \\
\textbf{Assumption 1.} \texttt{NewPolicy}($\w$) returns an optimal policy, $\pi_{\w}^{*}$, for the given task $\w$.\\ \\
\textbf{Assumption 2.} For all $\w\in\Wsupport$, if an optimal policy $\pi_{\w}^* \in \Pi^{\text{OK}}(\Pi)$, then \texttt{TrainOK}($\omega,\Wsupport,\Pi$) (Alg.~\ref{algo:train-ok}) returns a meta-policy, $\omega$, s.t. $\piok_{\omega}(\cdot,\w;\Pi)$ is optimal w.r.t. task $\w$.\\
Recall that the OKB (Alg.~\ref{algo:okb}) maintains two partial CCSs, $\PsiBase_k$ and $\PsiOK_k$, which store, respectively, the SFs of the base policies $\Pi_k$ and the policies generated via the OK using $\Pi_k$.
We first will prove that, if in any given iteration $k$, the set of candidate corner weights $\mathcal{C}$ is empty (see line 16 in Alg.~\ref{algo:okb}), then it holds that $\ccs \subseteq \PsiOK_{k+1}$.
First, note that $\mathcal{C}$ contains corner weights of both $\PsiBase_k$ and $\PsiOK_k$.
If the OK policy $\piok_\omega(\cdot;\Pi_k)$ is able to optimally solve all tasks in $\mathcal{C}$ (line 12 in Alg.~\ref{algo:okb}), then we have that $\ccs \in \PsiOK(\Pi_k)$ due to Lemma~\ref{lemma:corner-weights}.
%
Note that assumptions 1 and 2, above, are necessary because Lemma~\ref{lemma:corner-weights} requires the partial CCS $\Psi_k$ to be a subset of the CCS. If it contains $\epsilon$-optimal policies, then it is possible to prove convergence to $\epsilon$-CCSs instead~\citep{Alegre+2023}.
Now, to conclude the proof, we only need to show that $|\Pi_k| \leq |\ccs|$.
The set of base policies is never larger than the CCS due to line 24 in Alg.~\ref{algo:okb}, in which dominated base policies (i.e., redundant policies that are not exclusively optimal w.r.t. any $\w$) are removed from $\Pi_k$ by the procedure $\texttt{RemoveDominated}$.
Importantly, the set of base policies returned by OKB, $\Pi_k$, will only have the same of the CCS ($|\Pi_k| = |\ccs|$) in domains where no single task in $\Mlinear$ can be optimally solved by combining policies optimal for other tasks in $\Mlinear$.\footnote{An example of such a domain is an MDP composed of independent corridors, in which optimal policies for different weights $\w\in\Wspace$ must traverse different corridors.}
\end{proof}
\end{theorem*}
\subsection{Proof of Proposition~\ref{th:non-linear-tasks}}
\begin{proposition*}[\ref{th:non-linear-tasks}]
Let $\Pi_k$ be the set of policies learned by OKB (Alg.~\ref{algo:okb}) when solving tasks that are linearly expressible in terms of the reward features $\vphi(s,a) \in \mathbb{R}^d$. Let $r$ be an arbitrary reward function that is non-linear with respect to $\vphi$. Suppose the optimal policy for $r$ can be expressed by alternating, as a function of the state, between policies in the CCS induced by $\vphi$. That is, suppose that for all $s \in \Sspace$, there exists a policy $\pi^*_{\w}$ (optimal for some $\w \in \Wspace$) such that $\pi^*_r(s) = \pi^*_{\w}(s)$. Then, the OK can represent the optimal policy for $r$ using the set of base policies $\Pi_k$; i.e., $\pi^*_r \in \Pi^{\textup{OK}}(\Pi_k)$.
\begin{proof}
We have that, for all $s\in\Sspace$, $\exists \w\in\Wspace$ such that:
\begin{align}
\pi^*_{r}(s) &= \pi^*_{\w}(s), \quad \text{where } \vpsi^{\pi^*_{\w}} \in \ccs,\\
&= \argmax_{a\in\Aspace} \vpsi^{\pi^{*}_{\w}}(s,a) \cdot \w. \label{eq-th-non-linear-step1}
\end{align}
Due to Thm.~\ref{th:ok-ccs}, we know that by employing the set of base policies returned by OKB, we have that $\pi^{*}_{\w} \in \Pi^{\textup{OK}}(\Pi_k)$. Equivalently,
\begin{equation}\label{eq-th-non-linear-step2}
\pi^{*}_{\w}(s) = \piok_{\omega}(s,\w;\Pi_k) = \argmax_{a\in\Aspace} \max_{\pi\in\Pi_k} \vpsi^\pi(s,a) \cdot \omega(s,\w).
\end{equation}
Combining Eq.~\eqref{eq-th-non-linear-step1} and Eq.~\eqref{eq-th-non-linear-step2}, we have that
\begin{equation}
\pi^*_{r}(s) = \piok_{\omega}(s,\w;\Pi_k),
\end{equation}
which concludes the proof, demonstrating that $\pi^{*}_{r} \in \Pi^{\textup{OK}}(\Pi_k)$.
\end{proof}
\end{proposition*}
\subsection{Conditions for OKB’s Behavior Basis to be Strictly Smaller Than the CCS}\label{ap:strictly-smaller}
\looseness-1
Transfer learning techniques dating back to the early 1990s often leverage the principles of compositionality and modularity—the idea that related tasks can be solved by reusing parts of previously learned solutions.\footnote{These principles underpin a long line of research on hierarchical reinforcement learning, temporal abstraction, macro-actions, skill discovery, and studies in neuroscience and cognitive science exploring how animal behavior is organized hierarchically.}
%
Our formal analysis of when OKB’s behavior basis is strictly smaller than the CCS draws on the principles of compositionality and modularity---the idea that related tasks can be solved by combining reusable components (``sub-behaviors'') from previously acquired policies. We show that identifying these conditions reduces to analyzing the cardinality of a specific Set Cover problem.
Intuitively, the OKB only requires a behavior basis as large as the CCS in degenerate cases where the solution to \emph{all} possible tasks involves executing at least one type of behavior that does not occur in the solutions to \emph{any} other tasks. Such cases are fundamentally at odds with the assumptions that make transfer learning viable. In contrast, as long as policies in the CCS (which solve related tasks drawn from a shared distribution) produce similar behaviors in shared regions of the state space, the OKB matches the expressiveness of the CCS while requiring strictly fewer policies.
Consider a transfer learning scenario where agents must solve tasks drawn from some distribution, and each task’s solution can be expressed as a (possibly stochastic, closed-loop) sequence of sub-behaviors---often referred to as options, skills, policy modules, controllers, motor primitives, or sub-policies. For example, task $\tau_1$ might require sub-behaviors $o_1$, $o_2$, $o_5$, and $o_6$ (opening doors, driving, turning off lights, and finding exits), while task $\tau_2$ might require a different combination of sub-behaviors---some of which it may share with $\tau_1$, since both come from the same distribution.
Assume a distribution of related tasks that induce a CCS whose policies use some subset of the available sub-behaviors, $O = \{o_1, \dots, o_N\}$, with $\pi_i$ drawing on $O_i \subseteq O$. The question is: What is the smallest subset of CCS policies that together cover all sub-behaviors in $O$? Formally, we seek the smallest collection $\{\pi_j, \pi_k, \pi_r, \dots\}$ such that $O_j \cup O_k \cup O_r \cup \dots = O$. This reduces to the classic Set Cover problem and corresponds to identifying the minimum number of CCS policies containing all sub-behaviors that might be required to solve any tasks from the distribution of interest.
Let $\setcover \subseteq \ccs$ denote the solution to this Set Cover problem. By Proposition~\ref{th:non-linear-tasks}, if the solution to a novel task consists of switching between policies in a behavior basis, then a corresponding optimal meta-policy exists that sequences such policies. To formally characterize when the OKB is as expressive as a CCS while relying on a strictly smaller behavior basis, we need to understand when $|\setcover| < |\ccs|$.
In general, a Set Cover solution will have $|\setcover| = |\ccs|$ only in degenerate cases where no CCS policy can be discarded without leaving some sub-behavior uncovered. One such case arises when every policy in the CCS contains at least one unique sub-behavior---a behavior that \emph{never} contributes to solving any other possible task. This situation is arguably at odds with the assumptions of transfer learning. Only in such cases---for example, when every possible task uniquely requires a specialized sub-behavior that no other task ever uses---would $|\setcover| = |\ccs|$.
As long as some sub-behaviors are shared between elements of the CCS, the set of policies needed to cover all sub-behaviors (i.e., $|\setcover|$) will be strictly smaller than the CCS. In such common transfer learning settings, the OKB is as expressive as a CCS while requiring strictly fewer policies.
\section{Transforming General Linear Reward MDPs into Convex Reward MDPs}\label{ap:linear-to-convex}
Let $\Mlinear \subseteq \mathcal{M}$ be the (possibly infinite) set of MDPs associated with all linearly expressible reward functions. This set, commonly studied in the SFs literature, can be defined as
\[
\Mlinear \defas \{ (\mathcal{S}, \mathcal{A}, p, r_{\vect{w}}, \mu, \gamma) \mid r_{\vect{w}} = \vphi \cdot \vect{w} \},
\]
where $\w \in \mathbb{R}^d$ is a $d$-dimensional weight vector that may include negative elements.
Standard methods for computing convex coverage sets (CCS) require reward functions to be expressed as \textit{convex} combinations of features, i.e., using weights in the $d$-dimensional simplex, $\Delta^d$:
\[
\Delta^d = \left\{ \w \in \mathbb{R}^d \mid w_i \geq 0, \sum_{i=1}^d w_i = 1 \right\}.
\]
% BCS ADD
To use methods that learn a CCS---typically designed for convex reward functions---to solve all \textit{linear} tasks in $\Mlinear$, we first need to show that any MDP in $\Mlinear$ can be transformed into an equivalent one where rewards are expressed as convex combinations over a new set of features. In other words, given a family of MDPs $\Mlinear$, we show how to construct a corresponding family of MDPs with \textit{convex} rewards (defined over a transformed feature space $\tilde{\vphi}$) that induces the same set of optimal policies.
%
This family of MDPs can be defined as follows:
\begin{equation}
\label{eq:convex-reward-mdps}
\mathcal{M}_{conv}^{\tilde{\vphi}} \defas \{ (\mathcal{S}, \mathcal{A}, p, r_{\vect{w}}, \mu, \gamma) \mid r_{\vect{w}} = \tilde{\vphi} \cdot \vect{w}, \w \in \Delta^{2d} \}, \ \text{where} \ \tilde{\vphi}(s,a) = \begin{bmatrix} \vphi(s,a) \\ -\vphi(s,a) \end{bmatrix} \in \mathbb{R}^{2d}.
\end{equation}
% BCS ADD
We now show that this transformed family of MDPs induces the same set of optimal policies as $\Mlinear$. As a result, learning a CCS for $\mathcal{M}_{\text{conv}}^{\tilde{\vphi}}$ allows us to recover the optimal policy for any task in $\Mlinear$.
%
Let $r(s,a) = \vphi(s,a) \cdot \w$ be an arbitrary linear reward function with a weight vector $\w \in \mathbb{R}^d$.
One can decompose $\w$ into its non-negative and non-positive components as follows:
\[
\w = \w^+ - \w^-, \quad \text{where} \quad w^+_i = \max(w_i, 0), \quad w^-_i = \max(-w_i, 0).
\]
We can now define an augmented feature vector $\tilde{\vphi}$:
\[
\tilde{\vphi}(s,a) = \begin{bmatrix} \vphi(s,a) \\ -\vphi(s,a) \end{bmatrix} \in \mathbb{R}^{2d},
\]
and a non-negative weight vector $\tilde{\w}$:
\[
\tilde{\w} = \begin{bmatrix} \w^+ \\ \w^- \end{bmatrix} \in \mathbb{R}_{\geq 0}^{2d}.
\]
This allows us to rewrite the original reward function as follows:
\[
r(s,a) = \w \cdot \phi(s,a) = \tilde{\w} \cdot \tilde{\phi}(s,a).
\]
% BCS ADD
Finally, let us normalize the weights $\tilde{\w}$ (e.g., by dividing by their norm) to ensure they lie in the simplex $\Delta^{2d}$. Recall that scaling a reward function by a positive constant does not affect its optimal policy, so this transformation is valid.
%
Since we showed that any linear reward function can be expressed as a convex combination of a transformed feature set—without altering the corresponding optimal policy—it follows that a CCS for $\mathcal{M}_{\text{conv}}^{\tilde{\vphi}}$ contains the optimal solutions for all tasks in the original family, $\Mlinear$.
%
\section{Beyond Linear Rewards}\label{sec:beyond-linear}
In this section, we further discuss the use of the OK to solve tasks defined by non-linear reward functions; that is, reward functions that are not linear-expressible under reward features $\vphi(s,a,s') \in \mathbb{R}^d$.
We start by arguing that, in many cases, the linear reward assumption does not pose any limitations.
As discussed by \citet{Barreto+2020}, when both $\Sspace$ and $\Aspace$ are finite, we can recover any possible reward function by defining $d=|\Sspace|^2\times|\Aspace|$ features $\phi_i$, each being an indicator function associated with a specific transition $(s,a,s')$.
%
This result shows that it is possible to define reward features $\vphi \in \mathbb{R}^{|\Sspace|^2\times|\Aspace|}$, such that every task in $\mathcal{M}$ is linearly expressible.
A challenge, however, is to find alternative features with this property, but with dimension $d \ll |\Sspace|$~\citep{Touati&Ollivier2021}.
When examining how the OK is defined (Eq.~\ref{eq:ok-policy}), a natural idea is to use it to solve tasks defined by reward functions with state-dependent weights $\w(s)$.
In particular, let $\mathcal{M}_{\text{sd-linear}}^{\vphi}$ be a family of non-linear tasks defined as
\begin{equation} \label{eq:state-dependent-reward-mdps}
\mathcal{M}_{\text{sd-linear}}^{\vphi} \defas \{ (\mathcal{S}, \mathcal{A}, p, r_{\vect{w}}, \mu, \gamma) \mid r_{\vect{w}}(s,a,s') = \vphi(s,a,s') \cdot \vect{w}(s) \}.
\end{equation}
Note that it is easy to show that $\Mlinear \subseteq \mathcal{M}_{\text{sd-linear}}^{\vphi} \subseteq \mathcal{M}$.
Unfortunately, the OK may be unable to solve all tasks in the family $\mathcal{M}_{\text{sd-linear}}^{\vphi}$, even when given access to a set of policies $\Pi_k$ forming a CCS.
\begin{proposition}
There exists a task $M\in\mathcal{M}_{\text{sd-linear}}^{\vphi}$ such that no meta-policy, $\omega : \Sspace \to \Zspace$, and set of policies, $\Pi$, results in $\piok_{\omega}(s;\Pi)$ being optimal with respect to $M$.
\begin{proof}
Assume a family of MDPs, $M\in\mathcal{M}_{\text{sd-linear}}^{\vphi}$, with states, actions, transitions, and features defined as in Figure~\ref{fig:counterexample1}.
\begin{figure}[ht!]
\centering
\includegraphics[width=0.4\linewidth]{figs/counterexample1.pdf}
\caption{Counterexample MDP.}
\label{fig:counterexample1}
\end{figure}
Let a task $M\in\mathcal{M}_{\text{sd-linear}}^{\vphi}$ be defined via the following reward function: $r(s_1) = \vphi(s_1) \cdot [1,0] = 0, r(s_2) = \vphi(s_2) \cdot [-1,-1]=-3, r(s_3) = \vphi(s_3) \cdot [-1,-1]=-3, r(s_4) = \vphi(s_4) \cdot [1,1]=2$, with optimal policy $\pi^*(s_0)=a_4$ and whose SF vector is $\vpsi^{\pi^*}=[1,1]$.
We now show that there is no meta-policy $\omega(s)$ that results in $\pigpi(s;\omega(s))$ being equal to the optimal policy $\pi^*$.
This is proved by noticing the impossibility of the following equality:
\begin{align}
\pigpi(s_0,\omega(s_0)) = \underset{a \in \{a_1,a_2,a_3,a_4\} }{\arg\max} \,\, \underset{\pi \in \Pi}{\max} \,\, \vpsi^\pi(s_0,a) \cdot \omega(s_0) = a_4.
\end{align}
\end{proof}
\end{proposition}
Notice that the result above implies that there are tasks in $\mathcal{M}_{\text{sd-linear}}^{\vphi}$ that can not be solved by $\piok$ even when assuming access to a CCS.
While this result may be negative, we highlight that $\mathcal{M}_{\text{sd-linear}}^{\vphi}$ is a family of tasks more general than it may seem at first glance.
\begin{proposition}
Let $\mathcal{M}' \subset \mathcal{M}$ such that for all $M \in \mathcal{M}'$, $r(s,a,s') = r(s)$.
%
If $\vphi(s) \neq \vect{0}$ for all $s\in\Sspace$, then $\mathcal{M}_{\text{sd-linear}}^{\vphi} = \mathcal{M}'$.
\begin{proof}
If $\vphi(s) \neq \vect{0}$ for all $s\in\Sspace$, then for any reward function $r$ there exist a function $\vect{w}(s)$ such that $r(s) = \vphi(s) \cdot \vect{w}(s)$.
\end{proof}
\end{proposition}
The proposition above implies that, under mild conditions, $\mathcal{M}_{\text{sd-linear}}^{\vphi}$ is able to represent \textit{all} reward functions and tasks in $\mathcal{M}$.
\section{Training the OK Meta-Policy}\label{sec:train-ok}
Alg.~\ref{algo:train-ok} (Train-OK) is the sub-routine for training a meta-policy omega under an actor-critic training framework, given the selected tasks by Alg.~\ref{algo:okls}.
In Alg.~\ref{algo:train-ok}, we show how our method optimizes the OKB meta-policy $\omega$ for a given set of tasks $\Wsupport$.
\IncMargin{0.9em}
\begin{algorithm}[!ht]
\SetKwInput{KwInput}{Input}
\SetKwComment{Comment}{$\triangleright$\ }{}
\DontPrintSemicolon
\caption{Train Option Keyboard (\texttt{TrainOK})}
\label{algo:train-ok}
\KwInput{Meta-policy $\omega$, weight support $\Wsupport$, base policies $\Pi$.}
Initialize replay buffer $\mathcal{B}$ \;
Let $\vpsi^{\omega}(s,\z,\w)$ be the critic of the meta-policy $\omega$ \;
$\w_0 \sim \Wsupport$; $S_0 \sim \mu$ \;
\For{$t = 0, 1, 2, \dots, T$}{
\If{$S_t$ is terminal}{
$\w_t \sim \Wsupport$ \;
$S_t \sim \mu$ \;
}
$\z_t \leftarrow \omega(S_t,\w)$ \;
\Comment{Exploration clipped Gaussian noise, as in TD3}
$\z_t \leftarrow \z_t + \mathrm{clip}(\epsilon, -0.5, 0.5)$, where $\epsilon \sim \mathcal{N}(0, 0.2^2)$ \;
$\z_t \leftarrow \z_t / ||\z_t||_2$ \;
\Comment{Follow OK policy}
$A_t \leftarrow \pigpi(S_t,\z_t;\Pi)$ \;
Execute $A_t$, observe $S_{t+1}$, and $\vphi_t$ \;
Add $(S_t, \z_t, \vphi_t, S_{t+1})$ to $\mathcal{B}$ \;
Update $\vpsi^{\omega}$ by minimizing $\Ex_{(s,\z,\vphi,s')\sim\mathcal{B},\w\sim\Wsupport} \left[ \left( \vpsi^{\omega}(s,\z,\w) - \left( \vphi + \gamma \vpsi^{\omega}(s', \omega(s', \w), \w) \right) \right)^2 \right]$ \;
Update $\omega$ via the policy gradient $\nabla_{\z} \vpsi^{\omega}(s,\z,\w) \cdot \w |_{\z = \omega(s,\w)} \nabla_{\omega}\omega(s,\w)$ \;
}
\Return{$\omega$} \;
\end{algorithm}
\DecMargin{0.9em}
\section{Checking the Condition for OK Optimality}\label{sec:checking-ok-condition}
Our practical implementation of the OKB test the condition in Thm.~\ref{th:advantage} by randomly sampling from a replay buffer, $\mathcal{B}{=}\{(s_i,a_i,\vphi(s_i,a_i),s'_i,)\}_{i=1}^{n}$, which contains experiences observed during the training of the base policies in $\Pi_k$.
%
Given a candidate task $\w\in\mathcal{C}$ (line 12 of Alg.~\ref{algo:okb}), we compute the \textit{mean positive advantage} of $\omega$ as:
\begin{equation}\label{eq:sum-advantage}
\left(\sum_{i=1}^{|\mathcal{B}|} \mathbbm{1}_{\{\hat{A}_{\w}^{\omega}(s_i,a_i) > 0\}}\right)^{-1}
\sum_{i=1}^{|\mathcal{B}|}
\max(\hat{A}_{\w}^{\omega}(s_i,a_i), 0),
\end{equation}
where $\hat{A}_{\w}^{\omega}(s_i,a_i)=\vphi(s_i,a_i)\cdot\w + \gamma \vpsi^{\omega}(s'_i,\omega(s'_i))\cdot \w - \vect{\psi}^{\omega}(s_i,\omega(s_i))\cdot\w$.
In line 14 of Alg.~\ref{algo:okb}, we select the task $\w\in\mathcal{C}$ with the highest value for Eq.~\eqref{eq:sum-advantage}.
Intuitively, we select the task $\w$ for which the OK can improve the most its performance. By adding $\pi_{\w}$ to $\Pi_k$ and retraining $\omega$, the OK will be more expressive in the subsequent iteration.
We highlight that OKB theoretical guarantees (e.g., Thm.~\ref{th:ok-ccs}) are independent of the heuristic used to select the task $\w\in\mathcal{C}$ (line 14), and other strategies could be used instead.
\section{Experimental Details}\label{ap:experiments-details}
The code and scripts necessary to reproduce our experiments will be made publicly available upon acceptance.
The USFAs $\vpsi(s,a,\w)$ used for encoding the base policies $\Pi_k$ were modeled with multi-layer perceptron (MLP) neural networks with 4 layers with 256 neurons. We use an ensemble of $10$ neural networks, similar to \citet{Chen+2021}, and compute the minimum value over two randomly sampled elements when computing the Bellman update targets.
We used Leaky ReLU activation functions and layer normalization for improved training stability.
For a fair comparison, we independently optimized all method-specific hyperparameters (e.g., network architectures, OKB settings, and those of competing baselines) via grid search.
%
The budget of environment interactions per iteration (i.e., call to \texttt{NewPolicy} in Alg.~\ref{algo:okb}) used was $25000$, $50000$, $50000$ and $100000$ for the Minecart, FetchPickAndPlace, Item Collection, and Highway domains, respectively.
At each iteration, $\vpsi(s,a,\w)$ is trained with the current task $\w$, as well as with the tasks from previous iterations in order to avoid catastrophic forgetting.
The meta-policy $\omega(s,\w)$ was modeled with an MLP with 3 layers with 256 neurons.
We employed the techniques introduced by \citet{Bhatt+2024}, i.e., batch normalization and removal of target networks, which increased the training efficiency.
We used Adam~\citep{Kingma&Ba2015} as the first-order optimizer used to train all neural networks with mini-batches of size $256$.
When training the base policies, we used $\varepsilon$-greedy exploration with a linearly decaying schedule.
For training the meta-policy, we added a clipped Gaussian noise (see line 11 of Alg.~\ref{algo:train-ok}), as done in other actor-critic algorithms, e.g., TD3~\citep{Fujimoto+2018}.
Since running OK-LS (Alg.~\ref{algo:okls}) until no more corner weights are identified (see line 4) may require a large number of iterations, in the experiments we ran OK-LS for $5$ iterations, which we found to be enough for learning a well-performing meta-policy $\omega$.
To generate sets of test tasks $\Wspace' \subset \Wspace$ given different values of $d$, we employed the method introduced by \citet{Takagi+2020} available on pymoo~\citep{blank&deb2020}, which produces uniformly-spaced weight vectors in the simplex, $\Wspace$.
All experiments were performed in a cluster with NVIDIA A100-PCIE-40GB GPUs with 32GB of RAM. Each individual run of each method took approximately 2.5 hours (Item Collection), 6 hours (Minecart), 10 hours (FetchPickAndPlace), and 25 hours (Highway).
\subsection{Corner Weights.}\label{sec:corner-weights}
Below, we define the concept of corner weights used by OKB (Alg.~\ref{algo:okb}) and OK-LS (Alg.~\ref{algo:okls}), as defined in previous works~\citep{Roijers2016,Alegre+2022}.
\begin{definition}\label{def:corner-weights}
Let $\Psi = \{\vpsi^{\pi_i}\}_{i=1}^{n}$ be a set of SF vectors of $n$ policies. \textit{Corner weights} are the weights contained in the vertices of a polyhedron, $P$, defined as:
\begin{equation}
\label{eq:polyhedron-corner-weights}
P = \{ \vect{x} \in \mathbb{R}^{d+1} \ | \ \vect{V}^{+}\vect{x} \leq \vect{0}, \textstyle\sum_i w_i = 1, w_i \geq 0\},
\end{equation}
where $\vect{V}^{+}$ is a matrix whose rows store the elements of $\Psi$ and is augmented by a column vector of $-1$'s.
Each vector $\vect{x} {=} (w_1, ..., w_d, v_{\w})$ in $P$ is composed of a weight vector and its scalarized value.
\end{definition}
To compute corner weights, as in Def.~\ref{def:corner-weights}, we used pycddlib (\url{https://github.com/mcmtroffaes/pycddlib}) implementation of the Double Description Method~\citep{Motzkin+1953} to efficiently enumerate the vertices of the polyhedron $P$.
\subsection{Domains}\label{ap:domains}
In this section, we describe in detail the domains used in the experiments, which are shown in Fig.~\ref{fig:domains}.
\paragraph{Minecart domain.}
The Minecart domain is a widely-used benchmark in the multi-objective reinforcement learning literature~\citep{Abels+2019}. We used the implementation available on MO-Gymnasium~\citep{Felten+2023}.
This domain consists of a cart that must collect two different ores and return them to the base while minimizing fuel consumption. The agent' observation $\Sspace\subset\mathbb{R}^7$ contains the agent $x, y$ position, the current speed of the cart, its orientation ($\sin$ and $\cos$), and the percentage of occupied capacity in the cart by each ore: $\mathcal{S} = [-1, 1]^5 \times [0, 1]^2$.
The agent has the choice between 6 actions:
$\mathcal{A} = \{\text{MINE}, \text{LEFT}, \text{RIGHT}, \text{ACCELERATE}, \text{BRAKE}, \text{DO NOTHING}\}$.
Each mine has a different distribution over two types of ores. Fuel is consumed at every time step, and extra fuel is consumed when the agent accelerates or selects the mine action.
The reward features of this domain, $\vect{\phi}(s,a,s')\in\mathbb{R}^3$, is defined as:
\begin{align*}
\phi_1(s,a,s') &= \text{quantity of ore 1 collected if } s' \text{ is inside the base, else }0, \\
\phi_2(s,a,s') &= \text{quantity of ore 2 collected if } s' \text{ is inside the base, else }0, \\
\phi_3(s,a,s') &= -0.005 -0.025 \mathbbm{1}\{a=\mathrm{ACCELERATE}\} -0.05\mathbbm{1}\{a=\mathrm{MiNE}\}.
\end{align*}
We used $\gamma=0.98$ in this domain.
\paragraph{FetchPickAndPlace domain.}
We extended the FetchPickAndPlace domain~\citep{Plappert+2018}, which consists of a fetch robotic arm that must grab a block on the table with its gripper and move the block to a given target position on top of the table (shown in the middle of Fig.~\ref{fig:domains}).
Our implementation of this domain is an adaptation of the one available in Gymnasium-Robotics~\citep{Lazcano+2023}.
%
We note that the state space of this domain, $\Sspace \subset \mathbb{R}^{25}$, is high-dimensional.
The action space consists of discretized Manhattan-style movements for the three movement axes, i.e., $\{1: [1.0, 0.0, 0.0], 2: [-1.0, 0.0, 0.0], 3: [0.0, 1.0, 0.0], 4: [0.0, -1.0, 0.0], 5: [0.0, 0.0, 1.0], 6: [0.0, 0.0, -1.0]\}$, and two actions for opening and closing the gripper, totaling $8$ actions.
The reward features $\vect{\phi}(s,a,s') \in \mathbb{R}^d$ correspond to the negative Euclidean distances between the square block and the $d$ target locations (shown in red in Fig.~\ref{fig:domains}).
We used $\gamma=0.95$ in this domain.
\paragraph{Item Collection domain.}
This domain consists of a $10\times10$ grid world, in which an agent (depicted as a yellow circle in the rightmost panel of Fig.~\ref{fig:domains}) moves along the 4 directions, $\Aspace = \{\mathrm{UP},\mathrm{DOWN},\mathrm{LEFT},\mathrm{RIGHT}\}$, and must collect items of two different types (denoted by red triangles and green diamonds, respectively).
The reward features $\vphi(s,a,s')\in\mathbb{R}^2$ are indicator functions indicating whether the agent collected one of the items in state $s'$.
In each episode, $5$ items of each type are randomly placed in the grid.
The agent perceives its observation as a $10\times10$ image with $2$ channels (one per item), which is flattened into a $200$-dimensional vector.
We make the observations and dynamics toroidal---that is, the grid
``wraps around'' connecting cells on opposite edges---similarly as done by \citet{Barreto+2019}.
We used $\gamma=0.95$ in this domain.
\paragraph{Highway domain.}
This domain is based on the autonomous driving environment introduced by \citet{Leurent2018}. The agent controls a vehicle on a multilane highway populated with other vehicles. Its observations includes its coordinates as well as coordinates from other vehicles.
Formally, the state space is given by an array of size $V \times F$, where $V$ represents the number of vehicles in the environment and $F$ are the features describing each vehicle, e.g., velocity, position, angle. The agent's actions consist in changing lanes, accelerating or decelerating, or doing nothing, i.e., $\mathcal{A} = \{\text{TURN\_LEFT}, \text{IDLE}, \text{TURN\_RIGHT}, \text{FASTER}, \text{SLOWER}\}$.
The reward features of this domain, $\vect{\phi}(s,a,s')\in\mathbb{R}^3$, is defined as:
\begin{align*}
\phi_1(s,a,s') &= \text{normalized forward speed of the vehicle}, \\
\phi_2(s,a,s') &= 0.5 \text{ if driving in the rightmost lane, else } 0, \\
\phi_3(s,a,s') &= -1 \text{ if } a=\mathrm{TURN\_LEFT} \text{ or } a=\mathrm{TURN\_RIGHT} \text{ else } 0.
\end{align*}
All three reward features above are zeroed if the agent is off the road and penalized by $-10$ if the agent crashes into another vehicle.
We used $\gamma=0.99$ in this domain.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\newpage
\section*{NeurIPS Paper Checklist}
%%% BEGIN INSTRUCTIONS %%%
% The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: {\bf The papers not including the checklist will be desk rejected.} The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page
% limit.
% Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
% \begin{itemize}
% \item You should answer \answerYes{}, \answerNo{}, or \answerNA{}.
% \item \answerNA{} means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
% \item Please provide a short (1–2 sentence) justification right after your answer (even for NA).
% % \item {\bf The papers not including the checklist will be desk rejected.}
% \end{itemize}
% {\bf The checklist answers are an integral part of your paper submission.} They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
% The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "\answerYes{}" is generally preferable to "\answerNo{}", it is perfectly acceptable to answer "\answerNo{}" provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "\answerNo{}" or "\answerNA{}" is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer \answerYes{} to a question, in the justification please point to the section(s) where related material for the question can be found.
% IMPORTANT, please:
% \begin{itemize}
% \item {\bf Delete this instruction block, but keep the section heading ``NeurIPS Paper Checklist"},
% \item {\bf Keep the checklist subsection headings, questions/answers and guidelines below.}
% \item {\bf Do not modify the questions and only use the provided macros for your answers}.
% \end{itemize}
%%% END INSTRUCTIONS %%%
\begin{enumerate}
\item {\bf Claims}
\item[] Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: We support the theoretical claims in Section~\ref{sec:theoretical-results} and the empirical claims in Section~\ref{sec:experiments}.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the abstract and introduction do not include the claims made in the paper.
\item The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
\item The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
\item It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
\end{itemize}
\item {\bf Limitations}
\item[] Question: Does the paper discuss the limitations of the work performed by the authors?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: Yes, we discuss all theoretical assumptions in Section~\ref{sec:problem-formulation} and Section~\ref{sec:theoretical-results}, as well as future research directions for improving the method in Section~\ref{sec:conclusions}.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
\item The authors are encouraged to create a separate "Limitations" section in their paper.
\item The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
\item The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
\item The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
\item The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
\item If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
\item While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
\end{itemize}
\item {\bf Theory assumptions and proofs}
\item[] Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: The proof of all theoretical results can be found in Appendix~\ref{ap:proofs}. All assumptions made are stated in the introduced theorems or in Section~\ref{sec:theoretical-results}.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not include theoretical results.
\item All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
\item All assumptions should be clearly stated or referenced in the statement of any theorems.
\item The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
\item Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
\item Theorems and Lemmas that the proof relies upon should be properly referenced.
\end{itemize}
\item {\bf Experimental result reproducibility}
\item[] Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: We provide all the details on how the algorithms were implemented and the hyperparameters selected in Appendices~\ref{sec:train-ok}, \ref{sec:checking-ok-condition}, and \ref{ap:experiments-details}.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not include experiments.
\item If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
\item If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
\item Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
\item While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
\begin{enumerate}
\item If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
\item If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
\item If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
\item We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
\end{enumerate}
\end{itemize}
\item {\bf Open access to data and code}
\item[] Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: All the code required to reproduce our experiments is available in the Supplemental Material.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that paper does not include experiments requiring code.
\item Please see the NeurIPS code and data submission guidelines (\url{https://nips.cc/public/guides/CodeSubmissionPolicy}) for more details.
\item While we encourage the release of code and data, we understand that this might not be possible, so “No” is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
\item The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (\url{https://nips.cc/public/guides/CodeSubmissionPolicy}) for more details.
\item The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
\item The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
\item At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
\item Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
\end{itemize}
\item {\bf Experimental setting/details}
\item[] Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: We provide details about the experimental setup, hyperparameters and algorithm implementations in Section~\ref{sec:experiments}, Appendix~\ref{ap:domains}, and Appendix~\ref{ap:experiments-details}.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not include experiments.
\item The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
\item The full details can be provided either with the code, in appendix, or as supplemental material.
\end{itemize}
\item {\bf Experiment statistical significance}
\item[] Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: We state in Section~\ref{sec:experiments} that, in all experiments, we report the mean normalized return of each method (normalized with respect to the minimum and maximum returns observed for a given task) along with the 95\% bootstrapped confidence interval over 15 random seeds.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not include experiments.
\item The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
\item The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
\item The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
\item The assumptions made should be given (e.g., Normally distributed errors).
\item It should be clear whether the error bar is the standard deviation or the standard error of the mean.
\item It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96\% CI, if the hypothesis of Normality of errors is not verified.
\item For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
\item If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
\end{itemize}
\item {\bf Experiments compute resources}
\item[] Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: We specify all the compute resources in Appendix~\ref{ap:experiments-details}.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not include experiments.
\item The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
\item The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
\item The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
\end{itemize}
\item {\bf Code of ethics}
\item[] Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics \url{https://neurips.cc/public/EthicsGuidelines}?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: No human subjects were involved in this research, and all analyses were based on publicly available simulators. All required details were documented, and algorithms and code are provided.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
\item If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
\item The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).
\end{itemize}
\item {\bf Broader impacts}
\item[] Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
\item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: This paper is foundational research and is not tied to particular applications. We do not anticipate any negative social impact.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that there is no societal impact of the work performed.
\item If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
\item Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
\item The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
\item The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
\item If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
\end{itemize}
\item {\bf Safeguards}
\item[] Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
\item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: We do not release any dataset or model with risk for misuse.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper poses no such risks.
\item Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
\item Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
\item We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
\end{itemize}
\item {\bf Licenses for existing assets}
\item[] Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: We cite and credit the papers in which the original domains used in Section~\ref{sec:experiments} were introduced. All licenses are respected.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not use existing assets.
\item The authors should cite the original paper that produced the code package or dataset.
\item The authors should state which version of the asset is used and, if possible, include a URL.
\item The name of the license (e.g., CC-BY 4.0) should be included for each asset.
\item For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
\item If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, \url{paperswithcode.com/datasets} has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
\item For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
\item If this information is not available online, the authors are encouraged to reach out to the asset's creators.
\end{itemize}
\item {\bf New assets}
\item[] Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
\item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: Our extensions of the original domains used in the experiments are detailed in Appendix~\ref{ap:domains}, and their implementations are available in the Supplementary Material.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not release new assets.
\item Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
\item The paper should discuss whether and how consent was obtained from people whose asset is used.
\item At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
\end{itemize}
\item {\bf Crowdsourcing and research with human subjects}
\item[] Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
\item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: This paper does not involve crowdsourcing nor research with human subjects.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
\item Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
\item According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
\end{itemize}
\item {\bf Institutional review board (IRB) approvals or equivalent for research with human subjects}
\item[] Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
\item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: This paper does not involve crowdsourcing nor research with human subjects.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
\item Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
\item We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
\item For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
\end{itemize}
\item {\bf Declaration of LLM usage}
\item[] Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
%this research?
\item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
\item[] Justification: This work does not involve LLMs.
\item[] Guidelines:
\begin{itemize}
\item The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
\item Please refer to our LLM policy (\url{https://neurips.cc/Conferences/2025/LLM}) for what should or should not be described.
\end{itemize}
\end{enumerate}
\end{document}