Constructing an Optimal Behavior Basis for the Option Keyboard
Authors: Lucas N. Alegre
Ana L. C. Bazzan
Andre Barreto
Bruno Castro da Silva
Venue: NeurIPS 2025 poster ()
Keywords: reinforcement learning
option keyboard
generalized policy improvement
transfer learning
multi-task reinforcement learning
Primary area: reinforcement_learning
TL;DR: A method for constructing an optimal behavior basis for the Option Keyboard, enabling zero-shot identification of optimal solutions for any linear-reward task.

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

Constructing an Optimal Behavior Basis for the Option Keyboard
Lucas N. Alegre Institute of Informatics
Ana L. C. Bazzan
Institute of Informatics Federal University of Rio Grande do Sul Porto Alegre, RS, Brazil bazzan@inf, ufrgs.
Federal University of Rio Grande do Sul Porto Alegre, RS, Brazil Inalegreßinf ufrgs.br
André Barreto
Bruno C. da Silva University of Massachusetts Amherst, MA, USA bsilvaβcs. wmass. edu
Google Didd
London, UK andrebarreto@googlle..oo com
Abstract
Multi-task reinforcement learning aims to quickly identify solutions for new tasks with minimal no additional interaction with the environment. Generalized Policy Improvement (GPI) addresses this by combining set of base policies to produce a new one that at least as wood-thoouh 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 good-and often better. It achieves this through learned meta-policy that dynamically combines base policies. However, its performance critically depends the choice of base policies. This raises key question: is there optimal set of base policies-an optimal behavior basis-thar enables zero-shot identification optimal solutions for any linear tasks? We solve this open problem by introducing a novel method that efficiently constructs such an optimal behavior basis. We show that significantly reduces the number of base policies needed to ensure optimality in new tasks. We also prove that is strictly more expressive than a CCS, enabling particular classes of non-linear tasks be solved optimally. We empirically evaluate our technique in challenging domains and show that it outperforms state- of-the ff-the-art art approaches, increasingly so task complexity increases.
Introduction
Reinforcement learning (RL) methods have been successfully used to solve complex sequential decision-making problems (Silver et al., 2017; Bellemare et al., 2020). However, traditional RL algorithms typically require thousands millions of interactions with the environment learn single policy for a single task. Multi-task RL methods address this limitation by enabling agents to additional interactions. Such methods
quickly identify solutions for new tasks with minimal often achieve this by learning set of specialized policies (or behavior basis) designed for specific tasks and subsequently combining them to more rapidly solve novel tasks.
powerful approach for combining policies to solve new tasks in a zero-shot manner leverages succes- features (SFs) and generalized policy improvement (Barreto et al., 2018, 2020). GPI can combine base policies to solve a new task, producing a policy that least as good as any individual base pol-
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

icy. 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 convex coverage set (CCS) (A (Alegre et al., 2022). A CCS is 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 (Roijers, 2016) 6), making existing algorithms computationally expensive and difficult to scale to complex domains (Yang et al., 2019; Alegre et al., 2022, 2023b).
The option keyboard (OK) method improves upon GPI by producing policies that are at least as good and often better. achieves this through learned meta-policy that dynamically combines base policies by assigning state-dependent linear weights to reward features (Barreto et al., 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 (Barreto et al., 2019; Carvalho et al., 2023b). This raises a key question and open problem: is there an optimal set of base policies for the OK-an an optimal behavior basis-that enables zero-shot identification of optimal solutions for any linear tasks?
We solve this open problem by introducing Option Keyboard Basis (OKB), a novel method with strong formal guarantees that efficiently identifies an optimal behavior basis for the OK. OKB provably identi- fies a set of policies that allows the OK to optimally solve any linear task: i..., it can express all policies in 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 ex- pressive than a CCS, enabling particular classes of non-linear tasks to be solved optimally. We empiri- cally evaluate our method in challenging high-dimensional RL problems and show that it consistently outperforms offteooftthe-ar GPI-based approaches. Importantly, we also observe that the performance gain over competing methods becomes more pronounced as the number of reward features increases.
2 Background
2.1 Reinforcement Learning
An RL problem (Sutton and Barto, 2018) is typically modeled as a Markov decision process (MDP). An MDP is defined
tuple M 2 (S.A. T, 14, 7), where S is a state space, A is an action space, p(-|s, a) describes the distribution over next states given that the agent executed action a in state s, r: S AxSxAx5SS + R is reward function, p is an initial state distribution, and 7 [0, 1) is discounting factor. Let S,, A., and R2 - r(S,, Ar, S2+1) ) be random variables corresponding to the state, action, and reward, respectively, at time step t. The goal of an RI. agent is to learn a policy ι S - A that maximizes the expected discounted sum of rewards (return), G. Ei Reti. The action-
value function of a policy ग is defined as q (((ss)) a) _ EEI[G, = 8, A. = a). where E, [:] denotes the expectation over trajectories induced by Given q", one can define a greedy policy x' <sss) E arg max, (s, a). It is guaranteed that 9d (s,a) 2 q" q" V(s,a) E S x A. The processes of computing qº and nT' are known, respectively, as the policy evaluation and policy improvement steps. Under certain conditions, repeatedly executing the policy evaluating and improvement steps leads to an optimal policy TT(s) (s) E arg max, q" (s, a) (Puterman, 2014). Let (S, Á. p. H4, 7) be a Markov control process (MPC) (Puterman, 2014), i.e., an MDP without reward function. Given an MPC, we define a family of MDPs M & {(S, A. p.r. 14, 7) r S x A S IR) } that only differ in their reward function. We refer to any MDP M E M (and its corresponding reward function r) as task.
2.2 Generalized Policy Evaluation and Improvement
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 set of policies. In particular, GPE and GPI are used, respectively, (i) to evaluate a policy on multiple tasks; and (ii) to con- struct a policy, capable of solving a particular novel task, by improving on an existing set of policies.
Definition 2.1 (Barreto al. (2020)). GPE is the computation of the action-value function of policy π. of. (s, a), for a set of tasks. Moreover, given a set of policies II and reward function of an arbitrary task, T, GPI defines policy, n', such that q.¹ (s, a) max qi (s, a) for all (s.a) e S A.
2

---

Based on the latter definition, for any reward function r, a GPI policy can be constructed based on set of policies II as
GPI (s; LI) E arg max max 9, (s,
aEA wEll
a) be the action- value function of GPI under the reward function r. The GPI theorem (Bar- Let gr (s. (s.
reto et al., 2017 7) ensures that x GGI in Eq. (1) satisfies Def. 2.1; i... that 9f a) 2 ga 9f (s, a) for all (s, a) E S x A. This implies that Eq. (1) can be used to define a policy guaranteed to perform at least well as all other policies ग, IT w.r.t. any given reward function, r. The GPI theorem can also be extended to the case qna is replaced with an approximation, q": (Barreto et al., 2018).
2.3 GPE&GPI via Successor Features
Successor features (SFs) allows us to perform GPE&GPI efficiently (Barreto et al., 2017). Assume that the reward functions of interest are linear w.r.t. reward features, Φ : x A
- Rd. That is, reward function, Tw. can be expressed Iw (w a, s') Φ(s, a, s') w, where E F is weight vector. Then, the SFs of policy for given state-action pair (s, a), 20" (s, a) Rd, are defined as
(s 2 EST
Σ 1==
S = s, A. =
where ob A (S,,, At, S1+1). Notice that the definition of SFs corresponds to form of action-value function, where the features &, play the role of rewards. Thus, SFs can be learned through any temporal-differeee
(TD) learning algorithm (e.g.. Q-learning (Watkins, 1989; Mnih et al., 2015) for learning sb" ). We refer to 2b" A Eso=p [&b" (So, π (So))] as the SF vector associated with x, where the expectation is with respect to the initial state distribution.
Given the SFs up" (s, a) of policy x, it is possible directly compute the action-value func- tion q (s, a) of R,, under any lineably-epreessible reward functions, Tw. as follows: g (s, a) = S: = s, A₁ = a| = *b (s, w. That is, given any set of reward weights, {w,)=_ GPE can be performed efficiently via inner-products between the SFs and the reward weights: Tw, (s, = 2b" (8, a) w,- Note that the value of policy under any task w can be expressed as ww 2b" w. Let II = {Tr } be set of policies with correspond- ing SFs = {1b in 1-=1- Based on the definition of GPI (Def. 2.1) and the reward decomposition Tw (8, a, s') (5s5, a, s'). w, generalized policy, GP SS W + A, can then be defined follows:
E, 1E: p'rw (Se+s: A4tsss+441) S145+1)
GPI (s, w: II) E arg max max 20" (s, w. aEA sell
2.4 GPI via the Option Keyboard
The Option Keyboard (OK) (Barreto et al., 2019, 2020) is a method that generalizes GPI (Eq. (1)) by allowing agents to use a different weight vector at each time step. does so by learning a meta-policy w S + Z that outputs weights, z E Z, used when following the GPI policy π Gi'l z; II) at each state s. Intuitively, given a set of policies II, w modulates the behavior induced by GPI by controlling the preferences (s)) over features for each state s. The OK policy, TOKK at each state s, is given by:
OK Th (s;; 11)
e arg max max 2/"' vb" (s. wss)
(4)
aEA
FEll
Note that, by construction, OK (s; 11) न. GGP (()) I1). Intuitively, in the "keyboard" analogy. w(s) (a "chord") combines base policies in II ("keys") to generate more complex behaviors. A key property of OK is that learning a policy over Z is often easier than over A, as solutions typically involve repeating the same action Z across multiple timesteps, similar to temporally extended options.
3 Optimal Behavior Basis
In this section, we formalize the problem of identifying a behavior basis i.e., a set of base policies-thar enables the OK to optimally solve all tasks within given family of MDPs. We are
Without loss of generality, let 2 be the space of d-dimensional unit vectors. i.e., Z = [z add ||/2|22 1).
3

---

interested in incremental methods for constructing a behavior basis, IIg. that provably enables the OK to solve any MDP in M. Specifically, we consider approaches where an agent first iteratively learns a behavior basis for solving some subset of tasks, M' C M. Then, the method should be capable of leveraging GPI or OK to identify additional specialized policies for optimally solving novel, unseen tasks M & .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.
Let M in S 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
M Ii
- ((S, A, Tw: (. (42) I'w = -....
(5)
Each element M E Mi (or, equivalently, its weight vector w) is called task. In what follows, we consider weight vectors that induce convex combinations of features; that is, W = (w Σ, wy = 1, wy 2 0, Vi}. This is common practice, e.g.. in the multi-objective RL literature (Hayes et al., 2022). Our analyses naturally extend to linear combinations of features in App. B, 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 Min often identify a set of policies, 1Ig { such that their associated SF vectors, V = {zb in J==1" form& convex coverage set (CCS) (Alegre al., 2022). A CCS is set of SF vectors that allows the optimal value function, w = z/"".. w, for any given task w E W, to be directly identified.2
CCS A {25" 3w BWEWSSK s.t. Va 26 w 2 25
(6)
Unfortunately, methods that compute the complete CCS for M do not scale to complex tasks since its size can grow exponentially with co-ounning corner weights is equivalent to Vertex Enumeration (Roijers, 2016). Hence, it becomes critical to develop novel techniques capable of recovering all solutions induced by 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 1I, smaller than the CCS; that is, [11,] I |CCS].
We start by extending the definition of the OK (Eq, (4)) to include the description of the task being solved, w E W, as one of its arguments. That is, we consider meta-policies w S x W + 2, and OK policies defined as
T S. W: II) arg max max 20" (s. a) w(s w).
(7)
aEA
SENl
The meta-policy (s, w) enables the OK decompose the optimal policy for a task w by dynamically assigning state-dependent linear weights to the SFs of its various basc policies. Note that if w(s, w) = w for all s E S. we recover GPI. Given an arbitrary set of base policies 1lx, we define the space of all policies expressible by the OK, IT (ITg), ITOK(II&)) and their associated SF vectors, & OOK((Ta), respectively, as
ITKK (IT,) A (= w; II.)
w E W.W:W.S. SxW- - 2} and OKK (IIx) & (xb" | π E IT ITTKK((..)) (Te)))
We are now ready to mathematically define our goal:
Goal: Learn a set of policies (the behavior basis) 1l, and meta-policy w such that (1) [1x]] S TK JCCS|: and (2) Th (-; IIx) is optimal for any task M e Mi The latter condition implies that
COS C (LL(( i.e. the OK is at least as expressive as CCS.
As observed by Barreto et al. (2019, 2020), learning a meta-policy, w. is often easier than learning base policies over the original action space, A. Consider, e.g., that if task w E W can be solved by switch- ing between two base policies in IIg, then an optimal meta--olicy for w, w(, w), only needs to output two vectors (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.
A CCS may not be unique since tasks can have multiple optimal policies with distinct SF vectors. In what fol- lows, CCS refers specifically to the minimal set satisfying Eq. (6), which uniquely defined (Roijers al., 2013).

---

Constructing an Optimal Behavior Basis
We now introduce Option Keyboard Basis (OKB), novel method to solve the goal introduced in the previous section. The OKB (Alg. 1) learns a set base policies, llg, and a meta- veeaappolic, w, such that the induced OK policy (1; IIx)) is provably optimal w.r.t any given task E W.
The algorithm starts with single base policy optimized (e.g., "/dd in its set of policies II (lines 1-2). OKB's initial partial CCS, yokk and weight support set, Wrup, are initialized as empty sets (line 3). Wdup 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, w, will be trained so that the policies expressible by ToK iteratively approximate (Alg. 2), an algorithm inspired by the SFs Optimistic Linear Support (SFOLS) method (Alegre et al., 2022) but that operates over a meta-policy rather than the space of base policies.
solve an arbitrary initial task
CCS. This process is implemented by OK-LS
Algorithm 1: Option Keyhoard Basis (OKB)
Next, in lines 8-10, OKB identifies a set of can- didate tasks, that the OK cannot optimally solve with its current set of policies, I[;; i.c. set of tasks for which Tw is not included in (TT.(TT.). We discuss how to check this condi- tion in Section 4.1 If C is empty (line 11), the OKB terminates and returns a meta- veta-policy (w) and base policies (Ila) capable ensuring that CCS C (oK This is due to Thm. 4.3, which we discuss later. C is not empty, OKB correspond-
Input: MPC with features (88 a) R Soapooocy/=-Dittaaaaaad
Tw; Th
To < {Tw); uuu Tw]
3 Yok + {}: W/ + {}
Initialize meta- pootic wb
5 for l = 0, 1,
Update meta-policy WK+1, v W
OK-LS(wR,, OOK War lTx)
selects task from it and adds ing optimal base policy to its behavior basis, II; (lines 14-17). In line 18, RemoveDoming nated removes redundant not strictly required to solve at least one task.
7
C C+
for ( U Cornerw (
8
& Check if task w is solvable with OK and TTg
-poiiciss that are
9
if nw is not in (11c) then C C [w}
In Algorithms and the function CornerW( takes set of SF vectors, V, as input and returns set corner weights (see App. F.1). Intu- itively, both Alg. 1 and Alg. 2 rely on the fact that (i) a task w E W that maximizes A A(w) A bw max, ETI, U is guaranteed to be a corner weight, and (ii) if OKB has optimal policy for all corner weights, then it must have identified CCS. These results were shown in prior work on constructing a CCS (Roijers, 2016; Alegre et al., 2022) and can also be proven using the fundamental theorem of linear programming.
C is empty then
& Found optimal basis 1l and meta-policy return wk+1, IT, Vit
12
else
D. Learn new base policy w f select a task from C Tw, T + NeuPolicy(w)
14
1l +- 1l U {Tw}
144 <- the U [bb*w)
T 10 hise +
18
RemoveeDominated((THt) buse
OK-LS (Alg. 2) trains the meta-policy, on selected tasks Vew (line 2) using the base policies 1lg, that is partial CCS, yok iteratively approximates a CCS. It stores the tasks it has already trained on, as well as the corner weights of the current iteration, in the weight support set, Waup (line 5). If Wew 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 Wew W and trains the meta-policy w on the tasks w E Wyuu using the Train0K subroutine (Alg. 3). Finally, in line OK-LS computes the SF vectors of the policies induced by OK for each w Woup, updating its partial CCS, yok In App. D, discuss how train
Algorithm 2: OK Linear Support (OK- -SS)
Input:
Meta-policy w. partial CCS Oo weight support duu base policies lTg-
1. while True do
Wisomer CornerW(bOK) Wru
3 if W empty then return w. OK W
Wru + 1urPP Wisomer
6
w t- TrainOK(w, Wa 11.
(Alg. 3)
7 OK + 2 26 (T))) w w"m
OK < RemovaDominatedd(yW
8
Notice that SFs only need to be computed for vectors newly added to Wru the current iteration (lines 6-7).
5

---

using an actor-critic RL method (Alg. 3). Finally, note that to accelerate the learning of w, the corner weights in Wew can be prioritized, similar Alegre al. (2022).
4.1 Condition for OK Optimality
In this section, we address the following question: Given an arbitrary task with reward function r and a set of base policies Il, is n E II (II+)? In other words, does a meta-policy w exist such that the OK policy, π w(; Ils), can represent the optimal policy x,? Determining whether this holds is essential for implementing the OKB step in line 9.
Proposition 4.1. Let II = (Ti be a set of base policies with corresponding SFs V = ( Given an arbitrary reward function function r, an optimal OK policy OK "x an OK meta veta-ppoiccy w : S + Z. such that for all S.
(:;II) can only exist if there exists
arg max max "b, (8, -0())= arg max 0f (8,
aEA TElly
aEA
This proposition provides a sufficient condition for the existence of an optimal OK meta- meta-policy w for given reward function r. Intuitively, it implies that the OK policy should be able to express all optimal actions through w and V. Next, we show how to verify this condition without requiring access to of-
Let Hgg be a set of base policies, yl [9." (((H1, S, = s, A. = π (S,, z; [[x]] be the meta-policy 's action-value function for task r. Given a state-action pair (s, a), let the advantage function of for executing action a in state s be
be an OK meta-policy, and (s z) r(S,. Ar)
(s, a) A r(s,a) + > /af (Se41, (S.+1)) IS.
xf" ( (s)).
is well-known that an optimal policy's 's advantage function zero when evaluated at an optimal action; i.e., A= (s, a²)= = for all E S. Thm. 4.2 uses this insight to introduce principled way to verify if ,, E 11 (IIs).
Theorem 4.2. Let II, be a set of base policies and w" be a meta-policg trained to convergence to solve a given task r. If A² (s, > 0 for some (s,a), then action a is not expressible by n (:: 11x). Consequently, the OK with base policies II, cannot represent an optimal policy for r.
Note that the procedure for verifying whether the OK is optimal for a given task r (Thm. 4.2) is a sufficient, but not necessary, condition to guarantee optimality. It serves two key purposes: (i) to avoid training on tasks w whose optimal policies have either already been learned or that can be reconstructed from II; and w, and (ii) to prioritize which tasks to train in order to accelerate learning. In other words, the procedure for assessing OK optimality is primarily tool to reduce training effort; method does not rely on exact test to ensure the correctness of the behavior basis it constructs. For a detailed discussion of how Thm. 4.2 can be applied in practical implementations of OKB to verify whether ग. e 11 (IIx), see App. E.
4.2 Theoretical Results
In this section, we present the theoretical guarantees of OKB (Alg. 1). Proofs of the theorems can be found in App. A.
Theorem 4.3. OKB returns a set of base policies, Il, and a meta-policy, w. such that [M&. S |CCS) and - (:; 11) is optimal for any task M E M This implies that CCS C (Io((I i.e. the OK matches or exceeds the expressiveness of CCS and ensures zero-shot optimality for all M E M
This is a key result of this paper: OKB can provably identify a behavior basis, Ig. that enables an option keyboard to optimally solve any task M E MM Furthermore, it is capable of ensuring zero-shot optimality using a set of base policies potentially smaller than CCS. This is the first method to offer such a guarantee-until now, all existing approaches required the number of base policies to be strictly equal the size of the CCS. This limitation posed major scalability challenge: computing 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 Vertex Enumeration (Roijers, 2016). 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 e M Min See below for more details.

---

Proposition 4,4, Let IIg be the set of policies learned by OKB (Alg. 1) when solving tasks that are linearly expressible terms of reward features Φ(s, a) E R Let r be an arbitrary reward function that is non-linear with respect to &. Suppose the optimal policy for can be expressed by alternating, as a function of the state, between policies the CCS induced by Φ. That is, suppose that for all E S, there exists a policy nw (optinal for some w E W) such that x (s) = T (8). Then, the OK can represent the optimal policy for r using the set of base policies 1lg; i.e. ...- E 110x
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 (1) conditions under which the OKB matches the expressiveness of a CCS while relying on strictly fewer policies; and (2) broad class of non-linear tasks that can be optimally solved by an OK using the behavior basis learned by OKB.
First, Prop. 4.4 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. Building on this result and Thm. 4.3, we formally characterize, in App. A.5, the conditions under which the OKB is as expressive as a CCS while relying on strictly fewer policies.
Second, Proposition 4.4 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 M even when the task itself is non-linear. 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 without having learn all policies in the CCS. We further discuss the properties of OK under non-linear reward functions in App. C.
To state our next theoretical result, we first recall Thm. 2 by Barreto et al. (2017):
Theorem 4.5 (Barreto et al. (2017)). Let II= (TT I be a set of optimal policies w.r.f. tasks {w:) (m:
Let (3b Fi be approximations to their SFs and w E W be a task. Let (aw (s.a) Iw (s, a)l S e forr all (s, a,,a) SxA, and F; E Let &max= - *maxs,a (| a)l.. Then, holds that:
-- La (s, a) - Iw (s. a) S I 1-y ( (omax min -aapl//wwl+@@
w.ll+ for all (s, e (10)
Thm 4.5 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: (1) Since the OK generalizes GPI policies, follows that
Ii a) 9 (s, a) S Iw a)
-10))- ti .a (s. a); and (2) when Eq. (8) (Prop. 4.1) is satisfied, ow (s, a)
fw (s,a) O. 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. 4.3) but also naturally lead to transfer learning strategies that are significantly more robust approximation errors than GPI.
Experiments
We now empirically evaluate OKB and investigate the following research questions: Q1: Can OKB approximate a CCS more effectively while requiring fewer base policies than competing methods? Q2: Does OKB's performance advantage grow with problem complexity, measured by the number of reward features d? Q3: Can the base policies learned by OKB be used to solve tasks with non-linear reward functions, under the conditions in Prop. 4.4?
Figure 1: Domains used in the experiments: Minecart, FetchPickAndPlace, Item Collec- tion, and Highway.
Fig. 1 depicts the domains used in our experiments. To handle the high-dimensional state space of these domains, we learn base policies using Universal SF
7

---

Approximator (USFA) (Borsa et al., 2019). Furthermore, rather than learning a separate SF for each base policy π E Hg,, we train a single USFA, (b(s, a, w), conditioned on task vectors w. such that (s) z arg aa A 1b(s, a, w) w. Each call to NewPolicy task w. We employ USFAS in our experiments primarily to avoid the memory and computational costs of training separate neural network for each policy's SFs. Moreover, USFAs typically achieve a level of approximation comparable to-or better than-that that of independently trained networks, since in both cases generalization depends mainly on the architecture, capacity, and distribution of training data. See App. F for more details.
in Alg. trains the USFA to solve
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 (Alegre et al., 2022), method that identifies CCS using GPI policies, and OKB-Uniform tasks uniformly from W in line 14 of Alg. 1, rather than from 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 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. 2). This makes the comparison more conservative for OKB, as it has fewer interactions available to learn base policies.
variant of OKB that selects
Fig. 2 depicts each method's 's mean return over a set of test tasks, M E MM as 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-an domain. Both domains have reward functions defined by d = 3 reward features, which are detailed in App. F22 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. 2a, where OKB reaches near-optimal performance with just 2-3 base policies. While OKB-Unifome
Mean Nomolizel Return
(190
0.70
the Highway
0.0
OKB (oaa) OKS Unioas loma)
0.00-
SPULS (UM5-bacd)
Iteration
(a) Minecart
Mian Normoltzed Reaum
0.70
040
also performs well the Minecart domain, its lower performance in the Highway domain (Fig. 2b) highlights the importance of expanding the set of base policies by carefully desks-defined by corner
GKB trad
GKB Uafar (oea)
(CM kood)
selecting promising weights- for training- Across both domains, OKB consistently outperforms all competing methods.
Boration
(b) Highway
To investigate Q2, we evaluate each method in the FetchPickAndPlace domain with varying numbers of reward features, d E {2, 6, 8). In this domain, an agent controls a robotic arm that must grasp block and move it to a specified location. Each of the reward features represents the block's distance different target location (red spheres in Fig. 1). Fig. 3 shows that as the number of target locations (d) increases, the performance gap between OKB and SFOLS increases significantly positivvyy answering Q2. Intuitively, OKB focuses on tasks where the OK cannot express optimal actions (see Thm. 4.2). By learning the corresponding optimal policies, OKB quickly identifies behavior basis that enables the OK to solve tasks across decreases
Figure 2:
Mean normalized return per itera- tion for each method on a set of test tasks.
the entire space of task vectors W. Conversely, the gap between OKB and since a larger number of reward features enhances the expressivity of the OK (Prop. 4.1). The total computation time required by OKB to train meta-policy and base policies for zero-shot optimality
"To generate test task sets W C W for different values of d, we used the method introduced by Takagi et al. (2020). which produces uniformly spaced weight vectors in W.
$ 'Importantly, to the best of our knowledge, most state-of of-te-of-the-art methods are typically evaluated on problems with dimensionality at most d = 4 (c.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.
8

---

Figure 3: Mean normalized return over test tasks as a function of iteration number (i.e., number of base policies learned per method) PeetchPickAndPlace As d increases, OKB's performance advantage over SFOLS (a state-of-the-ar GPI-based algorithm) grows.
is consistently comparable fewer policies at inference time.
lower than that of competing methods, as it learns and evaluates
Next, we investigate Q3 by conducting experiment similar to the one proposed by Alver and Precup (2022). We compare OKB to relevant competitors in an environment with non-linear reward functions. The Item Collection domain (Fig. I; top right) is 10 x 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 combina- torial number of possible states. After OKB learns behavior basis 1lg, the agent trains a aneta-poliqe w: 8 + Z to solve task with non-linear reward function. Specifically, this is sequential task where the agent must collect all instances of one item type before collecting any items another type. We compare OKB with SIP (Alver and Precup, 2022), " method that learns d base policies each maximizing a specific reward feature-as and a meta-policy w over a discrete set of weights in W. We also compare OKB against 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 as- signing equal importance both. Fig. 4 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 A is significantly more difficult.
6.00
Toua
-.00
ORB Jomo
1. 1.4g
PUN
Time Step
T1
Figure 5: Continuous actions from the meta-policy on a sample task. OKB selects base policies in temporally consistent way.
Figure 4: Mean return in the Item Collec- tion domain under non- con-linear reward function.
Finally, we examine the qualitative behavior of OKB policies by visualizing the continuous actions produced by the meta-policy w(s, w) while solving sample task in the Highway domain, where the agent must prioritize driving fast and staying in the rightmost lane. In Fig. 5, the color of each timestep's 's column represents the selected base policy. The agent initially accelerates as quickly as possible by following base policy тं for the first 10 timesteps while the rightmost lane is occupied. At timestep 10. the agent turns right and switches to base policy T2, which optimized for driving in the rightmost lane. At f. = 20, Ti becomes active again, allowing the agent to accelerate while staying the rightmost lane. Note that while the meta-policy w is relatively smooth, the base policies it selects are complex. In particular, although z, remains approximately constant during the first 10 timesteps, the underlying primitive actions 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 options (Sutton al., 1999) that help solve the task.
A (Fig. 5; black numbers at the top of each
'SIP not part of previous experiments as it assumes independent reward features (Alver and Precup, 2022).

---

Related Work
OK and GPL. Previous works have extended the OK in different ways. Carvalho et al. (2023b) introduced neural network architecture for jointly learning reward features and SFs, while Machado et al. (2023) proposed using the OK with base policies that maximize Laplacian-bai basaaad----ddddd--- reward features. However, unlike our work, these methods do not provide theoretical guarantees on optimality for solv- ing specific families tasks. Other works have introduced GPI variants for risk-aware transfer (Gimel- farb et al., 2021), mitigating function approximation emoos (Kim et al., 2022), planning with approxi- mate models (Alegre al., 2023a), planning to solve tasks defined by finite state automata (Kuric et al., 2024), and combining Mon-Maakov policies (Thakoor et al., 2022). These approaches improve GPI in ways that are orthogonal to our contribution and could potentially be combined with our method.
Learning behavior basis. Previous works have addressed the problem of learning an effective behav- ior basis for transfer learning within the SF framework (Nemecek and Parr, 2021; Zahavy et al., 2021). The method introduced by Alver and Precup (2022) assumes that reward features are independent (i.e., can be controlled independently) and that the MDP's transition function is deterministic. These as- sumptions often do not hold in complex RL domains, including most of those studied in Section 5. Ale- gre et al. (2022) proposed method that learns 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 smaller behavior basis.
Learning features. While we focused on identifying an optimal behavior basis given particular set of reward features, other works have addressed the complementary problem of learning more expres- sive reward features (Touati and Ollivier, 2021; Carvalho et al., 2023a; Chua et al., 2024). A promising future research direction is to combine OKB with methods such as the Forward- tation (Touati et al., 2023) to construct an optimal behavior basis under learned reward features.
Barward-Backward Backward represen-
7 Conclusions
We introduced OKB, 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-arr 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 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 (Alegre et al., 2022) open the door to adapting OKB to integrate with MORI. techniques, paving the way for unified representations that improve both sample efficiency and generalization across objectives.
Acknowledgments
This work is partially supported by CAPES (Coordenação de Aperleiçoamento de Pessoal de Nível Superior Brazil, Finance Code 001). Ana Bazzan is partially supported by CNPq under grant number 304932/2021-3.
10

---

References
Abels, A., Roijers, D. M., Lenaerts, T., Nowé, A., and Steckelmacher, D. (2019). Dynamic weights in multi- objective deep reinforcement learning. In Proceedings of the 36th Interaational Conference on Machine Learning, volume 97, pages 11-20, Long Beach, California, USA. International Machine Learning Society (IMLS). F.2
Alegre, L. N., Bazzan, A. L. C., and da Silva, B. C. (2022). Optimistic linear support and successor features basis for optimal policy transfer. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 394 413, Baltimore, Maryland, USA. PMLR. 1, 4, 8, 5, 6, 7, A.1, A.1.
Alegre, L. N., Bazzan, A. L. C., Nowé, A.. and da Silva, B. C. (2023a). Multi-step generalized policy improvement by leveraging approximate models. In Advances in Neural Information Processing Systems. volume 36, pages 38181-38205, New Orleans, USA. Curran Associates, Inc. 6
Alegre, L. N.. Bazzan, A. L. C., Roijers. D. M.. Nowé, A., and da Silva, B. C. (2023b). Sample-efficient multi-objective learning via generalized policy improvement prioritization. In Proceedings of the 2023 International Conference on Autonomous Agents and Multiagent Systems, AAMAS '23, pages 2003-2000 2012, Richland, SC. International Foundation for Autonomous Agents and Multiagent Systems. 1, A.3
Alver, S. and Precup, D. (2022). Constructing good behavior basis for transfer using generalized policy updates. In Proceedings of the Tenth International Conference on Learning Reprexentations, Virtual. OpenReview.net. 6, 6
Barreto, A., Borsa, D., Hou, S., Comanici, G., Aygün,, E., Hamel, P., Toyama, D., Hunt, J., Mourad, S., Silver, D., and Precup, D. (2019). The Option Keyboard: Combining skills in reinforcement learning. In Proceedings of the 33rd International Conference on Neural Information Processing Systems, pages 13052-13062. Red Hook, NY, USA. Curran Associates Inc. 2,2,,,3, F22
Barreto, A., Borsa, D., Quan, J., Schaul, T., Silver, D., Hessel, M., Mankowitz, D., Zidck, A., and Munos, R. (2018). Transfer in deep reinforcement learning using successor features and generalised policy improvement. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 501-510, Stockholm, Sweden. PMLR. 1.2
Barreto, A., Dabney, W., Munos, R., Hunt, J. J., Schaul, T., van Hasselt, H. P., and Silver, D. (2017). Successor features for transfer in reinforcement learning. In Guyon, L, Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R., editors, Proceedingy of the 31st International Conference on Neurul Information Processing Systems, volume 30, pages 4058-4066 Red Hook, NY, USA. Curran Associates, Inc. 2.2, 2.3, 4.2, 4.5
Barreto, A., Hou, S., Borsa, D., Silver, D., and Precup, D. (2020). Fast reinforcement learning with generalized policy updates. Proceedings of the National Academy of Sciences, 117(48):30079 -0080007777-77777- 2.1, 2.4, 3, C
Bellemare, M. G., Candido. S., Castro, S., Gong. J., Machado, M. C., Moitra, S., Ponda, S. S., and Wang. Z. (2020). Autonomous navigation of stratospheric balloons using reinforcement learning. Nature, 588(7836):77- 82.
Bhatt, A., Palenicek, D., Belousov, B.. Argus, M., Amiranashvili, A., Brox. T., and Peters. J. (2024). Crossq: Batch normalization in deep reinforcement learning for greater sample efficiency and simplicity. In Proceed- ings of The Twelfth International Conference on Learning Representations.
Blank, J. and Deb, K. (2020). Pymoo: Multi-objective optimization in python. IEEE Access, 8:89497 89509, F
Borsa, D., Barreto, A., Quan, J., Mankowitz, D. J., Munos, R., Hasselt, H. V., Silver, D., and Schaul, T. (2019). Universal successor features approximators. In Proceedings the 7th International Conference on Learning Representations (ICLR), New Orleans, USA. OpenReview.net. 5
Carvalho, W., Filos, A., Lewis, R., Lee, H., and Singh, S. (2023a). Composing task knowledge with modular successor feature approximators. In Proceedings of The Eleventh International Conference on Learning Representations, Kigali, Rwanda. OpenReview.net. 6
Carvalho, W., Saraiva, A., Filos, A., Lampinen, A., Matthey, L., Lewis, R., Lee, H., Singh, S., Jimenez Rezende, D., and Zoran, D. (2023b). Combining behaviors with the successor features keyboard. In Oh, A., Neumann, T.. Globerson, A., Saenko, K., Hardt, M.. and Levine, S., editors, Advances in Neural Information Processing Systems, volume 36, pages 9956 9983, Curran Associates, Inc. 1,
11

---

Chen, X., Wang. C., Zhou. Z., and Ross, K. W. (2021). Randomized ensembled double q-learnnng Learning fast without model. In Proceedings of the Ninth International Conference Learning Representations (ICLR). Virtual. OpenReview.net.
Chua, R., Ghosh, A., Kaplanis, C.,, Richards, B. A., and Precup. D. (2024). Learning successor features the simple way. In Proceedings of the Thirty-eighth Annual Conference Neural Information Processing Systems.
de Lazcano, R., Andreas, K., Tai, J., Lee, R., and Terry, J. (2023). Gymnasium robotics. F.2
Felten, F., Alegre, L. N., Nowé, A., Bazzan, A. L. C., Talbi, E.-G., Danoy, G., and da Silva, B. C. (2023). A toolkit for reliable benchmarking and research multi-bbjective reinforcement learning. In Proceedings the 37th International Conference on Neural Information Processing Systems, volume 36, pages 23671 1-23700, Red Hook, NY, USA. Curran Associates, Inc. F.2
Fujimoto, van Hoof, H., and Meger, D. (2018). Addressing function approximation error actor-critic methods. In Dy, J. and Krause, A.. editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1587-1596, Stockholm, Sweden. PMLR. F
Gimelfarb, M., Barreto, A., Sanner, S., and Lee, C.-G. (2021). Risk-aware transfer in reinforcement learning using successor features. In Proceedings of the 35th Annual Conference on Advances in Neural Information Processing Systems, pages 17298-17310, Red Hook, NY, USA. Curran Associates Inc. 6
Hayes, C. F., Rādulescu, R., Bargiacchi, E., Källström, J., Macfarlane, M., Reymond, M., Verstraeten, T., Zintgraf, L. M., Dazeley, R.,, Heintz, F., Howley, E.,, Irissappane, A. A., Mannion, P., Nowé, A., Ramos, G., Restelli, M., Vamplew, P. Roijers, D. M. (2022). practical guide to multi-objection reinforcement learning and planning. Autonomous Agents and Aulti-Agent Systems, 36(1):26.
Kim, Park, S., and Kim, G. (2022). Constrained GPI for zero-shot transfer in reinforcement learning. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS'22 pages 4585-4997, Red Hook, NY, USA. Curran Associates Inc. 6
Kingma, D. P. and Ba, (2015). Adam: A method for stochastic optimization. Bengio, Y. and LeCun, Y., editors, Proceeding of the 3rd International Conference on Learning Representations (ICLR), San Diego, CA. USA. OpenReview.net. F
Kuric, D., Infante, G., Gómez, V., Jonsson, A., and van Hoof, H. (2024). Planning with learned policy basis to optimally solve complex tasks. In Proceedings of the International Conference on Automated Planning and Scheduling, pages 331-341,444 Alberta, Canada. AAAI Press.
Leurent, E. (2018). An Environment for Autonomous Driving Decision-Making Python Package. F.2
Machado, M. C., Barreto, A., and Precup, D. (2023). Temporal ahstraction reinforcement learning with the successor representation. Journal of Machine Learning Research (JMLR), 24(80):1-69 6
Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., et al. (2015). Human-level control through deep reinforcement learning. Nature, 518(7540):529-333 2.3
Motzkin, T. S., Raiffa, H. Thompson, G. L., and Thrall, R. M. (1953). The Double Description Method. Princeton University Press. F.I
Nemecek, M. and Parr, R. (2021). Policy caches with successor features. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8025- 8033, Virtual. PMLR. 6
Plappert, M., Andrychowicz, M., Ray, A., McGrew, B., Baker, B., Powell, G., Schneider, J., Tobin, J., Chociej. M., Welinder, P., Kumar, V., and Zaremba, W. (2018). Multi-goal reinforcement learning: Challenging robotics environments and request for research. arXiv preprint arXiv: 10e::8000 8802099444 F.2
Puterman, M. L. (2014). Markov Decision Processes: Discrete Stochastic Dynamic Programming. John Wiley & Sons, New York, NY, USA. 2.1
Roijers, D. (2016). Multi-Objective Decision-Theoretic Planning. PhD thesis, University of Amsterdam. 1. 8. 4.2, A.3,
Rojjers, D. M.. Vamplew, P., Whiteson, S., and Dazeley, R. (2013). A survey of multi-objectio sequential socision-making J. Artificial Intelligence Research, 48(1):67-113 -13.
12

---

Silver, D., Schrittwieser, J., Simonyan, K., Antonoglou, I., Huang, A., Guez. A.. Hubert. T., Baker. L., Lai, M.. Bolton, A., Chen, Y., Lillicrap. T., Hui, F., Sifre, L., den Driessche, G.. Graepel, T.,, and Hassabis, D. (2017). Mastering the game of go without human knowledge. Nature, 550(7676)335-3399
Sutton, R. S. and Barto, A. G. (2018). Reinforcement learning: An introduction. The MIT Press, Cambridge.
MA. USA, second edition.
Sutton, R. S., Precup, D., and Singh, S. (1999). Between MDPs and semi-MDPs: framework for temporal abstraction in reinforcement learning. Artif. Intell., 112(1 2):181 2211
Takagi. T., Takadama, K., and Sato, H. (2020). Incremental lattice design of weight vector Proceedings 2020 Genetic and Evolutionary Computation Conference Companion, GECCO '20, pages 1486 1494, New York, NY, USA. Association for Computing Machinery. 4, F
Thakoor, Rowland, M., Borsa, D., Dabney, W., Munos. R., Barreto, (2022). Generalised policy improvement with geometric policy composition. In Proceedings of the 39th International Conference on Machine Learning. volume 162 of Proceedings of Machine Learning Research, pages 21272-21307, Baltimore, Maryland, USA. PMIR. 6
Touati, A. and Ollivier, Y. (2021). Learning one representation to optimize all rewards. In Proceedings of the 35th International Conference on Neural Information Processing Systems, pages 13 23, Red Hook, NY, USA.
Curran Associates Inc. 6, C
Touati, A., Rapin, J., and Ollivier, Y. (2023). Does zero-shot reinforeement learning exist? In Proceedings of The Eleventh International Conference on Learning Representations, Kigali, Rwanda. OpenReview.net,
Watkins, C. (1989). Learning from Delayed Rewards. PhD thesis, University of Cambridge. 2.3
Yang. R., Sun, X.. and Narasimhan, K. (2019). A generalized algorithm for multi-objective reinforcement learning and policy adaptation. In Proceedings of the 33vd International Conference Neural Information Processing Systems, pages 14610 14621, Red Hook, NY, USA. Curran Associates Inc.
Zahavy, T., Barreto, A., Mankowitz. D. J., Hou, S., O'Donoghue, B., Kemaev, L,, and Singh, S. (2021). Discovering a set of policies for the worst case reward. In Proceedings of the International Conference on Learning Representations, Vienna, Austria. RevnRevie..net.
13

---

Appendix
A Proofs
A.1 Proof of Proposition 4.1
Proposition (4.1). Let 1 {==Y==1 be set of base policies with corresponding SFs v Re-- { Given arbitrary reward function optimal OK policy T (:; (;;1]) can only exist if there exists an OK meta-policy, w S + Z, such that for all s S,
arg max max (s, w(() REA TETTs
arg max (s. acA
(11)
Proof. The proof follows directly from the definition of a deterministic optimal policy for given reward function Recall that (s) E arg 9 (s. and T (s: arg m max & (s(s, a) w(s). for some state E S, there is no value of z = w(s) such (s; [1) = (s). O
that Eq. (11) holds, then is not possible to define a meta policy w such that TOK/
A.2 Proof of Theorem. 4.2
Theorem (4.2), Let M1g be set of base policies and w be a meta-policy trained to convergence to solve given lask r. If Ay (8, >> 0 for some (8, a). then action a is not expressible by rw" (5 Ilx). Consequently, the OK with base policies Tl cannot represent an optimal policy for
_KK
Proof. First, recall the definitions of g. (s, z) and (s,
(s, rS,, A.) + " w
St s, = GPI (S z: II,]],
(12)
(s, a) r(S,, A,) + yB o (S4+1, (S++1))
S s, A, = a] (s. (s)).
(13)
(* a) measures the relative benefit of executing action in state compared following the current OK policy TTK
If (5 a) > 0. then executing in provides a higher expected return than the action currently selected by OK Hence, if there exists an (s, a) such that A (s, a) > 0, this implics that the OK policy fails to represent at least one optimal action, meaning it is suboptimal.
Recall from Prop. 4.1 that for the OK policy to be optimal, must express the optimal action a" for every
arg max max 3p" (s. a) wss) arg max 9, (s. a rEll,
Since w is trained until convergence to solve r, it has reached a stable policy, However, if there exists an (s, a) such that A (s, a) 0, then
(s. > of (8)""(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 N. is restricted to the actions expressible via its base policies IIg, this implies that
a" &
arg max max (s, w(()
fEll.
Thus, possible weighting of the base policies recover the optimal action in state
Since at least one optimal action a" missing from the action space the OK policy, it follows that n TOK is not an optimal policy for and the task cannot be optimally solved using the given base policies TT. Therefore, guarantee optimality, additional base policies must be introduced into llg.
14

---

A.3 Proof of Theorem. 4.3
Lemma A.1 (Alegre et al. (2022)). Let Vx S CCS be subset of the CCS. If. for all comer weights w E Cornerw (.))) it holds that max 20" w =uw then Vx contains all elements of the CCS, i.e., CCS _ x-
This lemma follows directly from the theoretical guarantees of previous methods that construct a CCS, e.g. OLS and SFOLS (Roijers, 2016; Alegre et al., 2022). Intuitively, this result states that to check a given set of SF vectors, kx, is CCS, it is sufficient to check if it contains an optimal policy for all of its comner weights. If this is the case, then it is not possible to identify some alternative weight vector w E W for which it does not contain an optimal policy.
Theorem (4.3). OKB returns set of base policies, Il,, and a meta meta-policy, w, such that [I&s] S |CCS) and =0K (:; T&3) is optimal for any task M E Mir Min This implies that CCS C Too (TTx; Le., the OK matches or exceeds the expressiveness of a CCS.
Proof. We start the proof by stating two assumptions on which the theorem relies:
Assumption 1. NexPolicy(w) returns an optimal policy, Tw. for the given task w.
Assumption 2. For all E Wau,, if an optimal policy Tे (TTK, then TrainOK(w, Wsop, IT) (Alg. 3) returns a meta-policy, w, s.t. m0 (s, w; II) is optimal w.r.t. task w.
Recall that the OKB (Alg. 1) maintains two partial CCSs, D and & which store, respectively, the SFs of the base policies IIg and the policies generated via the OK using IIg. We first will prove that, if in any given iteration k, the set of candidate corner weights C is empty (see line 16 in Alg. 1), then it holds that CCS S '++1 First, note that C contains corner weights of both Thuse and wo If the OK policy xO (.; [I.))
in Alg. 1), then we have that CCS E OK (TI,) due to Lemma A. Note that assumptions 1 and 2, above, are necessary because Lemma A. I requires the partial CCS V to be subset of the CCS. If contains c-optimal policies, then it is possible to prove convergence to e-CCSs instead (Alegre et al., 2023b).
(;:Te) is able to optimally solve all tasks C (line
Now, conclude the proof, we only need to show that [[g < |CCS). The set of base policies is never larger than the CCS due to line 24 in Alg. 1, in which dominated base policies (i.e., redundant policies that are not exclusively optimal w.r.t any w) are removed from I1, by the procedure RemoveDominated Importantly, the set of base policies returned by OKB, IIg, will only have the
same of the CCS (II& = (CCS)) in domains where no single task in Mg can be optimally solved by combining policies optimal for other tasks in M fin
A.4 Proof of Proposition 4.4
Proposition (4.4). Let Hl be the set of policies learned by OKB (Alg. 1) when solving tasks that are linearly expressible in terms of the reward features (18 a) E Rd. Let r be an arbitrary reward function that is non-linear with respect to O. Suppose the optimal policy for r can be expressed by alternating,
function of the state, between policies in the CCS induced by d. That is, suppose that for all s E S, there exists a policy Th (optimal for some w E W) such that n. (s) = T (s). Then, the OK can represent the optimal policy for r using the set of base policies llxi i.e., n.. E 11 (II2).
Proof. We have that, for all s S. W such that:
(8) = nW (s), where Th E CCS,
(14)
arg max (s, (s, w.
(15)
aeAA
Due to Thm. 4.3, we know that by employing the set of base policies returned by OKB, we have that w E 11 (IIx). Equivalently. π (s) = T OK (s, w; ITg) = arg max max (s, w(s, w).
(16)
REA xell.
7 An example of such a domain is an MDP composed of independent corridors, in which optimal policies for different weights E W must traverse different corridors.
15

---

Combining Eq. (15) and Eq. (16), we have that
7,(s) = < w: ITg), which concludes the proof, demonstrating that ,,, IIOK
(17)
A.5 Conditions for OKB's Behavior Basis to be Strietly Smaller Than the CCS
Transfer learning techniques dating back to the early 1990s often leverage the principles of compo- sitionality and modularity the idea that related tasks can be solved by reusing parts of previously learned solutions. 8 Our formal analysis of when OKB's ' behavior basis is strictly smaller than the CCS draws the principles of compositionality and modularity- 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 specific Set Cover problem.
Intuitively, the OKB only requires a behavior basis as large as the CCS degenerate cases where the solution to all possible tasks involves executing at least one type of behavior that does not occur in the solutions to 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 ss solution can be expressed as (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 Ti might require sub-behaviors 01. O2, Os, and os (opening doors, driving, turning off lights, and finding exits), while task 72 might require different combination of lub-behaviors- -somiiorsssssmeeeeeeee of which it may share with 71, since both come from the same distribution.
Assume distribution of related tasks that induce CCS whose policies use some subset of the available sub-behaviors, O = {or,.... ON J. with n. drawing on O2 C 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 {T,, Tk, ......) such that Oj U O2 UO,, U. ... O. This reduces to the classic Set Cover problem and corresponds to identifying the minimum number of CCS policies containing all uub-behaviorr that might be required to solve any tasks from the distribution of interest.
Let SC C CCS denote the solution to this Set Cover problem. By Proposition 4.4, 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 |SC| |CCS).
In general, a Set Cover solution will have [SC| = 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 never contributes to solving any other possible task. This situation is arguably at odds with the assumptions of transfer learning. Only in such cases-foo example, when every possible task uniquely requires a specialized sub-behavior that no other task ever uses would |SC| = |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., [SC[) will be strictly smaller than the CCS. In such common transfer learning settings, the OKB is as expressive as CCS while requiring strictly fewer policies.
B Transforming General Linear Reward MDPs into Convex Reward MDPs
Let M CM 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
M {(S A,p, rw: 14, (,,7) Fw Φ w]
'These principles underpin a long line of research on hierarchical reinforcement learning. temporal abstraction, macro-acto cnccro-atiioss. skill discovery, and studies in neuroscience and cognitive science exploring how animal behavior is organized hierarchically.
16

---

where w E F is a d-dimensional weight vector that may include negative elements.
Standard methods for computing convex coverage sets (CCS) require reward functions be expressed as convex combinations of features, i.e., using weights in the d-dimensional simplex, Ad:
d W WED Rd
0,
1=1
To use methods that learn a CCS -yypccall designed for convex reward functions- solve linear tasks in Mi we first need to show that any MDP in Min can be transformed into an equivalent one where rewards expressed convex combinations over new set of features. In other words, given family of MDPs M we show how construct corresponding family of MDPs with conver rewards (defined over a transformed feature space 3) that induces the same set of optimal policies. This family of MDPs can be defined as follows:
()s ( a)
M cb cone
- {(S, P, Tw: rw = Φ W.WE Add), where ((s =
ER2d
(18)
We now show that this transformed family of MDPs induces the same set of optimal policies M As result, learning CCS for M allows recover the optimal policy for any task in Imn Let rss, a) (ss, be arbitrary linear reward function with weight vector Rd.. One can decompose w into its non-negative and non-positive components as follows:
w
where W max max(w;; 0), Wi max ws, 0).
We can now define an augmented feature vector क
Φbs,a) a) (0s a)
e(s,a)
E Rd2d
and non-negative weight vector w:
w+ 122d
This allows to rewrite the original reward function as follows:
r(s, = ( w p(s,))
Finally, let us normalize the weights w (e.g., by dividing by their norm) to ensure they lie in the simplex A 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 a transformed feature set-without altering the corresponding optimal policy-it follows that CCS for M contains the optimal solutions for all tasks in the original family, M fin
C Beyond Linear Rewards
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 6(s, a, s') e Rd
We start by arguing that, in many cases, the linear reward assumption does not pose any limitations.
As discussed by Barreto al. (2020), when both S and A are finite, reward function by defining d = | x [A] features Φi. each being an indicator function associated with specific transition (s, s'). This result shows that it is possible to define reward features & E [ such that every task in M is linearly expressible. A challenge, however, is to find alternative features with this property, but with dimension < [|] (Touati and Ollivier, 202
recover any possible
17

---

When examining how the OK is defined (Eq. 4), natural idea is to use it to solve tasks defined by reward functions with state-dependent weights w(s). In particular, let M sd-linear be family of non-linear tasks defined as
M 0 sd-lincar
rw:
( rw(ss,, = Φ(s, a, s w(s)}.
(19)
Note that it is easy to show that M C M Id-iiuuu S.M
Unfortunately, the OK may be unable to solve all tasks in the family M sd-linear even when given access to set of policies II, forming CCS.
Proposition C.1. There exists task M E MM such that no meta-policy, w S + 2, and set of policies, II, results in OK (s: II) being optimal with respect to M.
Proof. Assume a family of MDPs, M E M linear with states, actions, transitions, and features defined in Figure 6. Let task M M
ad-linear be defined via the following reward function:
((;))
(52) [2, 1]
(sa)
(5,)
[0.
[1.]]
[1]]]
S2
S3
Q2
Q3
a1
So
Figure 6: Counterexample MDP.
r(8)== (()))) ((>))
3, r(xg)(()))))))))))
-3 r
((s4) . [[, 2, with optimal policy x² (*a) = 04 and whose SF vector rb" = [1, We now
show that there is no meta-policy w(s) that results in (s) w(s) being equal the optimal policy
T.. This is proved by noticing the impossibility of the following equality:
GPI (so, w(s))
arg max
ab" (su, a) w(su) 04.
(20)
ac (a1.033/3))4)) nell
lineur Notice that the result above implies that there are tasks in M lineur that not be solved by OK even when assuming access to a CCS.
While this result may be negative, we highlight that Ma sd-linear is a family of tasks more general than it may seem at first glance.
Proposition C.2. Let M' C M such that for all M E M', r(s,a,())) r(s). If ((s) + 0 for all E S, then M sd-linear M'.
r(s).
Proof. If ((8) * 0 for all E S, then for any reward function r there exist a function w(s) such that r(s) (ss) w(s).
The proposition above implies that, under mild conditions, M si-linear is able represent all reward functions and tasks in M.
D Training the OK Meta-Policy
Alg. 3 (Train-OK) is the sub-routine for training meta-policy omega under an actor-critic training framework, given the selected tasks by Alg. 2. In Alg. 3, we show how our method optimizes the OKB meta-policy w for given set of tasks woup
18

---

Algorithm 3: Train Option Keyboard (Train0K)
Input: Meta-policy w, weight support WouP, base policies II. Initialize replay buffer B
Let ap" z. w) be the critic of the meta-policy 3 Wo Woup; So ~ μ
for & 0.1,2..... do Se is terminal then
5
Wr Woup
SInp
8
2e f ((,,,w) w)
b Exploration clipped Gaussian noise, in TD3
+ 2, + clip(c. -0.5, 0.5), where ~ WO0, 0.23)
10
2, 4 z//1zz/l2
& Follow OK policy
A, + (S, z;; II)
11
12
Execute A,, observe St+1, and Φ;
13
Add (S,, 21, Φ,, S(41) to B
14
Update 4ith by minimizing
== (s, 2s w) - (Φ + ( w((s,, w), w)))
15
Update via the policy gradient Vz2b" (8, w) (lz=w( Vww w)
return w
E Checking the Condition for OK Optimality
Our practical implementation of the OKB test the condition in Thm. 4.2 by randomly sampling from replay buffer, (= at ((si. az), s/11 ) which contains experiences observed during the training of the base policies in 1lg. Given a candidate task w E C (line 12 of Alg. 1), we compute the mean positive advantage of as:
[8 Σ
B max Am (si; az), 0),
(21)
In line 14 of Alg. we select the task w E with the highest value for Eq. (21). Intuitively, we select the task w for which the OK can improve the most its performance. By adding T to Hg and retraining w, the OK will be more expressive in the subsequent iteration. We highlight that OKB theoretical guarantees (e.g., Thm. 4.3) are independent of the heuristic used select the task C (line 14), and other strategies could be used instead.
where AW (si, a;) (bss, a.) w (s/, (s), (s())
(si,
Experimental Details
The code and scripts necessary reproduce experiments will be made publicly available upon acceptance.
The USFAs 1b(8, w) used for encoding the base policies II were modeled with multi-layer perceptron (MLP) neural networks with layers with 256 neurons. We use an ensemble of 10 neural networks, similar to Chen et al. (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-specins hyperparameters (e.g.. network architectures, OKB settings, and those of competing baselines) via grid search.
The budget environment interactions per iteration (i.e., call to NewPolicy in Alg. 1) used was 25000, 50000, 50000 and 100000 for the Minecart, FetchPickAndPlace Item Collection, and High- way domains, respectively. At each iteration, (b(8, a, w) trained with the current task w, as well as with the tasks from previous iterations in order to avoid catastrophic forgetting.
19

---

The meta meta-ppoicy wsss, w) was modeled with an MLP with 3 layers with 256 neurons. We employed the techniques introduced by Bhatt al. (2024), i.e. batch normalization and removal of target networks, which increased the training efficiency. We used Adam (Kingma and Ba, 2015) as the first-order optimizer used to train all neural networks with mini-batches of size 256.
When training the base policies, we used e-greedy exploration with a linearly decaying schedule. For training the meta-policy, we added clipped Gaussian noise (see line 11 of Alg. 3), as done in other actor-critic algorithms, e.g., TD3 (Fujimoto et al., 2018).
Since running OK-LS (Alg. 2) until large number of iterations, in the experiments we ran OK-LS for 5 iterations, which we found to be enough for learning well- well-performing meta-policy
comnrr weights are identified (see line 4) may require
To generate sets of test tasks W C W given different values of d, we employed the method introduced by Takagi et al. (2020) available on pymoo (Blank and Deb, 2020), which produces uniformly-spacee weight vectors in the simplex, W.
All experiments were performed in a cluster with NVIDIA A A100-PCIE-40GB GPUs with 32GB of RAM. Each individual run of each method took approximately 2.5 hours (Item Collection), hours (Minecart), 10 hours BetchhiikAAdPPlee and 25 hours (Highway).
F.1 Corner Weights.
Below, we define the concept of corner weights used by OKB (Alg. 1) and OK-LS (Alg. 2), as defined in previous works (Roijers, 2016; Alegre et al., 2022).
Definition F1. Let = < be set of SF vectors of n policies. Corner weights are the weights contained in the vertices of polyhedron, P, defined as:
p {x E Ted+1
V S 0, E, u. = w, 20220 0},
(22)
where V is matrix whose rows store the elements of V and is augmented by a column vector of -.ss. Each vector x=.,,...., Wd, Uw) in composed weight vector and its scalarized value.
To compute corner weights, in Def. F.I. we used pycddlib (https://gi thub com/ searofffaes/ycdddibb implementation of the Double Description Method (Motzkin et al., 1953) to efficiently enumerate the vertices of the polyhedron P.
F22 Domains
In this section, we describe in detail the domains used in the experiments, which are shown in Fig. 1.
Minecart domain. The Minecart domain is widely-used benchmark in the multi-objective reinforcement learning literature (Abels et al., 2019). We used the implementation available on MO- 1O--Gymasiumm (Felten et al., 2023). This domain consists of cart that must collect two different ores and return them to the base while minimizing fuel consumption. The agent' observation S C IR7 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: S = [-1, 1]5 x [0, 1 The agent has the choice between actions: A {MINE, LEFT, RIGHT, ACCELERATE, BRAKE, DO NOTHING). Each mine has a different distribution over two types of ores. Fuel is consumed at every time step, and extra fuel consumed when the agent accelerates or selects the mine action. The reward features of this domain, Φ(s, a, s') E R is defined as:
Φ (8, a, 8') = quantity of ore 1 collected if s' is inside the base, else 0, 02(5, a, s = quantity of ore 2 collected if s' is inside the base, else O,
Φ3(8, a, = -0.005 0. 0..221 {a ACCELERATE] - 0.051 (a = MINE].
We used y = 0.98 in this domain.
FetchPickAndPher domain. We extended the PetchPickAndPPa domain (Plappert et al., 201 8), which consists of 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. 1). Our implementation
20

---

of this domain is an adaptation of the one available in Gymnasium Robooiss (de Lazcano et al., 2023).
We note that the state space of this domain, S R is high-dimensional. The action space consists of discretized Manhattan-styy [-1.0, 0.0, 0.0].3 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 actions. The reward features ((s, a, s')
movements for the three movement axes, i.e., (1 [1.0, 0.0. 0.0), 2
Rd correspond to the negative Euclidean distances between the square block and the d target locations (shown in red in Fig. 1). We used y 0.95 in this domain.
Item Collection domain. This domain consists of a 10 x 10 grid world, in which an agent (depicted as a yellow circle in the rightmost panel of Fig. moves along the directions, A {UP, DOWN, LEFT, RIGHT), and must collect items of two different types (denoted by red triangles and green diamonds, respectively). The reward features Φ(s, a, ) functions indicating whether the agent collected one of the items in state In each episode, 5 items of each type are randomly placed in the grid. The agent perceives its observation as 10 x 10 image with 2 channels (one per item), which is flattened into 200-dimensional vector. We make the observations and dynamics toroidal-thhe is, the grid "wraps around" connecting cells on opposite similarly done by Barreto et al. (2019). We used 0.95 in this domain.
R² are indicator
Highway domain. This domain is based on the autonomous driving environment introduced by Leurent (2018). The agent controls vehicle 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 array of size V F, where V represents the number of vehicles in the environment and are the features describing each vehicle, c.g., velocity. position, angle. The agent's 's actions consist in changing lanes, accelerating or decelcrating, doing nothing, ie.., A = (TURN_ domain, Φ(s,a, a, 8') E R is defined
IDLE, TURN_ JURN_RIGHT, FASTER, SL.OEER) The reward features of this
Φ (s. s') normalized forward speed of the vehicle, 02 a,)) = 0.5 if driving in the rightmost lane, else O, 03(8,a, 1 if a LEFT or a TURN_ TURN_RIGHT else 0.
All three reward features above are zeroed if the agent is off the road and penalized by -10 the agent crashes into another vehicle. We used 0.99 in this domain.
21

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We support the theoretical claims in Section 4.2 and the empirical claims in Section 5.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Yes, we discuss all theoretical assumptions in Section 3 and Section 4.2, as well as future research directions for improving the method Section 7.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
22

---

Answer: [Yes]
Justification: The proof of all theoretical results can be found in Appendix A. All assumptions made stated in the introduced theorems in Section 4.2.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide all the details on how the algorithms were implemented and the hyperparameters selected in Appendices D, E, and F.
Guidelines:
The answer NA means that the paper does not include experiments.
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.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution 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 accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which 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 have some path reproducing or verifying the results.
5. Open access to data and code
23

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: All the code required to reproduce our experiments is available in the Supple- mental Material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might 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 new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ //nips. wwtpccccppccc.eec.cc...ccc.cctlt g/guidgusss./ Cooeeuumiis wwtammssssooni ssliciolicyy for more details.
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.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide details about the experimental setup. hyperparameters and algo- rithm implementations in Section 5, Appendix F.2, and Appendix F.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We state Section 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.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
24

---

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).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
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 96% CI, if the hypothesis of Normality of errors is not verified.
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 range (e.g. negative error rates).
If error bars 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.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We specify all the compute resources in Appendix F.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: / /ttp//nuuripppeee cctppbicccc//cpcccccccccccccccct cew.iiieeeleeeeeeeeeee
Answer: [Yes]
Justification: No human subjects were involved in this research, and all analyses were based publicly available simulators. All required details were documented, and algorithms and code are provided.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper is foundational research and is not tied to particular applications. We do not anticipate any negative social impact.
Guidelines:
25

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious 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.
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.
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 intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not release any dataset or model with risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have 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 usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
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 best faith effort.
12. Licenses for existing assets
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?
Answer: [Yes]
Justification: We cite and credit the papers in which the original domains used in Section 5 were introduced. All licenses are respected.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
26

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
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.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Our extensions the original domains used in the experiments are detailed in Appendix F.2, and their implementations available in the Supplementary Material.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/modee part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
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 details about compensation (if any)?
Answer: [NA]
Justification:
This paper does involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS 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.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
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?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
27

---

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.
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.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if 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, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
This work does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: httppp://neuripete http:pppc.cccoooeeeeeeeee/e//e// for what should or should not be described.
28


Reviews:
[{"review_id": "dTisrw60kA", "reviewer": "Reviewer_rDD6", "summary": "This paper formalizes the idea of an optimal behavior basis for use with the Option Keyboard (OK), i.e., a set of base policies which are sufficient for OK to solve a family of MDPs with optimality. Further, the authors propose Option Keyboard Basis (OKB), an incremental method for contructing a behavior basis and meta-policy which combined are provably optimal with respect to any of the tasks in the family. The paper includes the results of several experiments intended to answer specific research questions and demonstrate the effectiveness of OKB compared to baselines in different domains.", "questions": "* Can you discuss how accurate the OK optimality test is in practice? Including results from domains where this can be checked exactly but are still large would be helpful.\n\n* What other approaches did you consider for the OK optimality test and why did you choose the one that you did?\n\n* Does using the OK framework specifically more constraining than a CCS? In other words, does this limit the problems OKB can be used to solve compared to previous methods?", "limitations": "* The authors admit that the OK optimality test must be approximated, but do not expand much on how their approximation method was chosen or how much suboptimality it might introduce.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "# Strengths\n* The issue of discovery of a set of policies to use with Successor Feature-based methods which are sufficient to behave optimally in all tasks in a family is interesting and an important one to solve.\n* The formalization of the optimal behavior basis is simple but still meaningful.\n* OKB appears to be a strong approach to producing a set of base policies and meta-policy which are optimal with respect to the family of tasks while requiring fewer base policies than existing methods.\n* The experiments are nicely targeted to answer the specific research questions that are proposed in Section 5.\n\n# Weaknesses\n* Only a single baseline that is not an ablation of OKB is used in each experiment. While this is sufficient to answer the specific research questions that the authors presented, it would be nice to see a broader comparison to compare this method to others that don't necessarily guarantee optimality across the family.\n* For guaranteed optimality, the check for OK optimality requires knowledge of (or solving for) the optimal policy for the task in question, which is obviously not practical. While the sample-based method used in the experiments appears to be sufficient to allow OKB to be effective, there is little discussion of the best way to do this and no way to know how much suboptimality this approximation can introduce.", "was_revised": true, "final_justification": "The authors addressed the questions that I had regarding OKB. These were all things that I would like to see discussed somewhere in the paper, so I am happy that they intend to adjust their text to include them.\n\nMy rating remains Accept.", "rebuttal": "We thank the reviewer for their positive assessment of our paper and for their encouraging comments on the relevance of the problem, the strength of our method, and the quality of our experiments. Below, we address the reviewer’s questions:\n\n---\n\n### Comparison with methods that do not guarantee optimality\n\nThe reviewer highlighted that our comparisons were sufficient to positively answer all specific research questions and wondered how our method would perform when compared to weaker methods that do not necessarily guarantee optimality.\n\nWe appreciate this question and agree that broader comparisons are valuable. For a more detailed discussion—including additional baselines from the unsupervised skill discovery literature—please see our response to Reviewer 1Yjh. Please do not hesitate to let us know if you have any thoughts or further questions based on those results: we would be happy to discuss them further if helpful.\n\n---\n\n### Do approximations to the OK optimality test affect correctness or optimality guarantees?\n\nThank you for raising this important point. Your question helped us recognize the need to further clarify that our method does *not* rely on an exact test to guarantee the correctness of the behavior basis it constructs. As briefly noted in line 226 of Section 4.1, the test introduced in Eq. (9) for assessing OK optimality is a **sufficient (but *not* necessary) condition** for guaranteeing optimality. This test is designed to serve two key purposes: **(1)** to avoid training on tasks $w$ whose optimal policies have either already been learned or that can be reconstructed from $\\Pi_k$ and $\\omega$, and **(2)** to prioritize which tasks to train on in order to accelerate learning (see Appendix E for further discussion). In other words, the test is primarily a tool to reduce training effort, not a prerequisite for correctness.\n\nMore concretely, even if the test is imperfect, as long as the algorithm avoids re-solving previously trained tasks, the algorithm is still guaranteed to identify a behavior basis that supports zero-shot optimal synthesis for all tasks in the family. In cases where the test is imperfect, the algorithm may take longer to converge, potentially solving more tasks than strictly necessary, but all correctness guarantees remain intact. These guarantees follow from our argument in lines 189–200 and from established results in the literature (e.g., Theorem 3 in Roijers, 2016). In short, although approximate versions of Eq. (9) may slightly increase the number of training iterations required, they do not affect the method’s correctness or result in suboptimal solutions.\n\nWe appreciate your suggestion to make this point more explicit. In the final version, we will revise Section 4.1 to more clearly state that this test is a sufficient, but not necessary, condition for optimality.  We will also move portions of the discussion from Appendix E into the main text to better emphasize that its primary role is to accelerate the identification of a behavior basis.\n\n---\n\n### Is the OK framework more constraining than a CCS or prior methods?\n\nThis is a great question. The set of problems solvable by the OKB is no more limited than that of previous approaches; in fact, it is *broader*. As shown, for example, by Alegre et al. (2022), a Coverage Convex Set (CCS) is guaranteed to contain the optimal solution to every and all tasks in the family defined in Eq. (5). Since our method can directly reconstruct any element of the CCS (and, importantly, without having to train an individual policy for each), it follows that it can optimally solve all linear problems in that family.\nThis result is established formally in Theorem 4.3, where we prove that the OKB is at least as expressive as a CCS; that is, it can reconstruct all policies in the CCS, and potentially more. Proposition 4.4 goes further, showing that the OKB is in fact *strictly more expressive* than a CCS. In particular, it supports the direct reconstruction of optimal solutions not only for all linear tasks, but also for a relevant class of non-linear problems. This means that while existing methods that compute a CCS are limited to linearly expressible tasks, our method is not only guaranteed to return zero-shot optimal solutions to all such tasks, but it can also—importantly—solve a broader class of non-linear problems that previous approaches could not handle.\n\n---\n\nWe hope the comments above address your questions. Please feel free to reach out with any further thoughts—we would be glad to continue the discussion. Thank you once again for your thoughtful feedback and for taking the time to carefully read our work and share your insights."}, {"review_id": "pHPU5J4puk", "reviewer": "Reviewer_1Yjh", "summary": "The paper proposes a novel multi-task reinforcement learning method that uses a bi-level policy structure. In the linear-reward setting, prior work has shown that it is possible to compute a Convex Coverage Set (CCS), a set of base policies where the optimal policy for any reward function (by the linear reward definition) can be expressed as a convex combination of policies in the CCS. However, constructing a CCS can be expensive and it is hard to scale it to more complex domains. The proposed method additionally uses a meta policy that is allowed to dynamically choose the convex combination weights of the base policies depending on the state (unlike in prior work the weights are constant across states). This allows them to express policies that are optimal for more complex reward functions beyond the linear assumption. The proposed method is evaluated on four simulated environments. Compared to prior methods, the proposed method needs fewer number of base policies to achieve the same performance of the baselines with more base policies.", "questions": "N/A", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 2, "originality": 3, "strengths_and_weaknesses": "*Strengths*\n\n- The paper is well-written and easy to read.\n- The proposed method is novel and theoretically sound. \n\n*Weaknesses*\n\n- The authors consider only a limited number of baselines (SIP and DQN on one environment and SFOLS on others). The paper can benefit from including more baselines to demonstrate the effectiveness of the proposed method much better. For example, how does the proposed method compare to using existing unsupervised skill discovery methods that learn a latent-conditioned policies? (e.g., [1]). Also, comparing methods in options framework would also be relevant.\n- One of the main motivations of the paper is that CCS \"are computationally expensive and do not scale to complex domains\". However, as far as I could tell. All the tasks considered are relatively simple where their CCS can be computed exactly. It would greatly strengthen the paper if the authors could showcase the propose method on some of the environments where computing a CCS is completely intractable. \n\n\n[1] Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. \"Diversity is all you need: Learning skills without a reward function.\" arXiv preprint arXiv:1802.06070 (2018).", "was_revised": true, "final_justification": "The authors have addressed all of my concerns. The new experiments provide evidence that the proposed method OKB can scale to more complex domains. I am raising my score to 5.", "rebuttal": "We thank the reviewer for the positive comments regarding our method’s novelty and theoretical soundness, as well as the clarity and quality of the writing. The reviewer’s questions focused on two key aspects: **(1)** the set of baseline methods; and **(2)** the complexity of the domains used in the experiments. Below, we address both points.\n\n---\n\n### Comparison with Unsupervised Skill Discovery and Options Baselines\n\nWe thank the reviewer for bringing up this important point. As the reviewer correctly noted, other methods from the unsupervised skill discovery and options literature are indeed relevant to the broader goal of more rapidly solving novel tasks. Examples include *Diversity is All You Need* (DIAYN), which learns a diverse set of skills without access to task rewards; *Set-Max Policy* (SMP), which constructs a set of policies that maximize worst-case performance over a task set; and *DSP*, which encourages diversity in the space of successor features while maintaining near-optimality (Zahavy et al., 2021).\n\nIn Section 5, we chose SFOLS and SIP as our primary baselines, as both have been repeatedly shown in the literature to outperform more task-agnostic methods across a variety of settings. Moreover, they are representative of approaches that aim to construct policy sets with theoretical guarantees—a central objective of our work.\n\nBelow, we present results comparing SIP to the three methods discussed above (DIAYN, SMP, and DSP), which belong to the family of approaches highlighted by the reviewer. These results are adapted from Alver and Precup (2022) and show that SIP consistently achieves better downstream task coverage (measured by mean normalized return over seventeen tasks uniformly covering the task space) in an experimental setting similar to ours. Since our results demonstrate that OKB systematically outperforms SIP—and SIP is well known to consistently outperform the other baselines in this class—these findings further support our choice of SIP as the most competitive and informative point of comparison.\n\n|  w  |  SIP  |  DIAYN  |  SMP  |  DSP  |\n|:---:|:-----:|:-------:|:-----:|:-----:|\n| 1   | 0.81  | 0.19    | 0.47  | 0.37  |\n| 2   | 0.86  | 0.23    | 0.50  | 0.39  |\n| 3   | 0.87  | 0.31    | 0.51  | 0.43  |\n| 4   | 0.87  | 0.44    | 0.54  | 0.50  |\n| 5   | 0.90  | 0.67    | 0.60  | 0.56  |\n| 6   | 0.90  | 0.84    | 0.62  | 0.67  |\n| 7   | 0.89  | 0.89    | 0.67  | 0.68  |\n| 8   | 0.91  | 0.91    | 0.68  | 0.68  |\n| 9   | 0.93  | 0.93    | 0.70  | 0.69  |\n| 10  | 0.94  | 0.94    | 0.73  | 0.71  |\n| 11  | 0.93  | 0.93    | 0.70  | 0.69  |\n| 12  | 0.95  | 0.90    | 0.67  | 0.66  |\n| 13  | 0.92  | 0.77    | 0.60  | 0.69  |\n| 14  | 0.94  | 0.65    | 0.58  | 0.51  |\n| 15  | 0.94  | 0.53    | 0.56  | 0.47  |\n| 16  | 0.93  | 0.38    | 0.54  | 0.44  |\n| 17  | 0.86  | 0.22    | 0.49  | 0.37  |\n| **Mean norm. return →**   | **0.90** | **0.63** | **0.60** | **0.56** |\n| **StdDev norm. return →** | **0.0388** | **0.2857** | **0.0832** | **0.1304** |\n\n*Table 1: Normalized returns (mean and standard deviation shown at the bottom) for various baselines, including SIP, DIAYN, SMP, and DSP across different weight vectors (tasks).*\n\nFor completeness, we will include these comparisons in the paper. The findings above—namely, that SIP consistently outperforms more task-agnostic baselines such as DIAYN, DSP, and SMP in settings like the one we study—have been repeatedly observed in other prior work (e.g., Alegre et al., 2022). Based on the broader literature comparing these techniques, we thus fully expect to observe the same trend here, further reinforcing that SIP is the strongest and most informative point of comparison. All supporting tables and graphs will be added to Section 5 and the appendix.\n\n---\n\n### Complexity of the Tasks\n\nThe second point raised by the reviewer concerns whether the domains used in our experiments are too simple and whether their CCS could be computed exactly.\nWe have empirically evaluated OKB in challenging high-dimensional RL problems, including Minecart, Fetch-PickAndPlace (a robotic arm task), Item Collection, and Highway (an autonomous driving environment). These domains are well-established benchmarks in multi-task and multi-objective RL. Our results show that OKB consistently outperforms state-of-the-art GPI-based approaches, with its performance advantage becoming increasingly pronounced as task complexity (measured by the number of reward features $d$) increases. These results demonstrate that OKB is indeed applicable and highly effective in complex and challenging scenarios relevant to real-world applications, such as robotics and autonomous driving, going beyond simple simulated environments.\n\nRegarding tractability, prior to our paper, to the best of our knowledge, existing methods had been evaluated in problems with dimensionality at most $d=4$ (e.g., SFOLS and SIP). In our paper, we double the number of reward features and show that our method remains effective—outperforming all baselines even on problems an order of magnitude more challenging.\n\nFinally, we highlight that the CCS of some domains (e.g., FetchPickAndPlace and Highway) *cannot be computed exactly*. In all domains, we resorted to function approximation via neural networks to handle continuous state spaces of up to 25 dimensions. We agree with the reviewer on the importance of further discussing these points, and we will update the text to more clearly emphasize that our experiments were conducted on domains whose size and complexity had not been previously addressed in the literature.\n\n---\n\nWe hope our responses address your main questions. We will revise the text to incorporate a discussion on all the points brought up by the reviewer. If you believe the clarifications above address the key points you raised, we would be grateful if you would consider revisiting your score. If any points remain unclear, we welcome further feedback and would be happy to continue the discussion. Thank you again for your thoughtful comments!"}, {"review_id": "g3J3NvSegD", "reviewer": "Reviewer_QsaG", "summary": "The paper proposes a novel approach to multitask reinforcement learning with successor features. Concretely, the contribution consists in computing a policy basis that enables zero-shot identification of an optimal policy for any linear task. The policy basis can be significantly smaller than a convex coverage set, and generalization occurs by considering linear combinations of the policies in the basis. The contribution is empirically tested in a set of benchmark domains.", "questions": "Proposition 4.4 states that the OK *can* represent an optimal policy for a non-linear reward function. However, how would you actually compute the function \\omega(s,w) in this setting, if there is no linear task w?\n\nIn the experiments, the authors use a USFA parameterized on the task w to learn the successor features. This seems like a peculiar choice given the problem at hand. On one hand, this will likely introduce approximation errors. On the other hand, this allows representing the successor features of *any* policy, not just those in a small set.\n\nThere seems to be a typo in Theorem 4.2: in A^{w^*}r(s,a), r should supposedly be a quantifier of A.", "limitations": "Notably, a discussion around the number of successor feature vectors computed by the algorithm is lacking.", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "The paper is clearly written, and achieving efficient zero-shot identification of an optimal policy in a family of MDPs is an important research question. It is also true that the policy basis is indeed smaller than that in previous work that also guarantee optimality.\n\nOn the negative side, I am not convinced that the proposed algorithm leads to more efficient learning than previous work. The reason is that the computational cost is dominated by the number of successor feature vectors, not the number of base policies. Specifically, learning the successor feature vector of a single policy is as hard as solving a single MDP in the given family. From the description of the algorithm (on page 4 and in Algorithm 1) it appears that the number of successor feature vectors in the set \\Psi^{OK} is *larger* than that of a convex coverage set, which would imply that the overall computational complexity is greater. The authors do not discuss why this is the case, but my intuition is that we cannot linearly compose successor feature vectors of different policies, making it necessary to explicitly compute a successor feature vector for each policy considered (including many linear combinations of the base policies).", "was_revised": true, "final_justification": "After discussion with the authors I have raised my score, since some of my concerns were clarified (specifically, the number of successor feature vectors computed by the algorithm). I believe that the final version of the paper needs to be more honest regarding the contribution, and include the theoretical justification provided by the authors in the discussion.", "rebuttal": "We thank the reviewer for the insightful review. We appreciate their positive comments acknowledging that our method is capable of producing a basis significantly smaller than a CCS while ensuring zero-shot optimality—this was precisely the primary goal of our paper. \n\nBelow, we address the reviewer's main concerns and questions:\n\n---\n\n### On the dominant computational cost in our approach\n\nThank you for your questions. Your comments helped us realize that we should more clearly emphasize the distinct roles played by the two sets central to our method—namely, $\\Pi_k$ and $\\Psi^{OK}$. \n\nThe reviewer makes three claims: *(a)* learning an SF vector is as hard as solving an MDP; *(b)* the computational cost of our method is dominated by the number of successor feature (SF) vectors rather than the number of base policies; and *\\(c\\)* the overall cost of learning $\\Psi^{OK}$ exceeds that of learning a CCS.\n\nTo carefully address these, we begin with a broader overview of our method and its key properties, followed by a more technical and formal discussion.\n\n### Overview and definitions\n\n**(1)** Two sets of policies are central to our method—both for formally characterizing its properties and for designing the algorithm itself: $\\Pi_k$ and $\\Psi^{OK}$. The set $\\Pi_k$ is a *behavior basis*—a small, compact set of policies (much smaller than the CCS) that our method combines to solve new tasks (line 153). In contrast, $\\Psi^{OK}$ denotes the set of all optimal policies (more precisely, their SF vectors) that can be directly reconstructed by our method by combining the policies in $\\Pi_k$ (line 151). These are the strategies our method can deploy to solve novel tasks zero-shot, without the need to train a specialized policy for each task—a process that is often costly, time-consuming, or infeasible.\n\n**(2)** A key property of our method—central to its effectiveness and flexibility—is its provable ability to directly reconstruct all policies in a CCS, as well as policies for a challenging family of non-linear tasks, without ever explicitly having to compute or optimize them. Instead, it is capable of reconstructing them on demand using a meta-policy $\\omega$ that combines policies from the behavior basis $\\Pi_k$. These formal guarantees are established in Theorem 4.3 and extended by Proposition 4.4.\n\n**(3)** In our paper, we formally express this property by stating that the CCS is a subset of $\\Psi^{OK}$; that is, CCS $\\subseteq \\Psi^{OK}$. Intuitively, this means that our method is not only capable of reconstructing all optimal solutions for linear tasks, but it generalizes beyond the CCS and directly synthesizes optimal policies for particular non-linear tasks.\n\n**(4)** Importantly, note that the formal statement CCS $\\subseteq \\Psi^{OK}$ is a claim about the ***expressiveness*** of our method—specifically, which optimal policies it is capable of directly reconstructing, zero-shot, without requiring any training. It is ***not*** a statement about which (or how many) policies must be learned to construct the behavior basis that enables such zero-shot capability.\n\n### Addressing the three primary points raised by the reviewer\n\nGiven the overview above, we now respond to the three key claims raised by the reviewer.\n\n---\n**(a)** ***”Learning an SF vector is as hard as solving an MDP”****.\n\nUnlike what the reviewer claims, learning SF vectors is *not* as hard as solving an MDP. \n\nEstimating the successor features associated with a policy is a *policy evaluation* problem. Solving an MDP, by contrast, is a *policy optimization* problem. The sample complexity of solving an MDP and identifying an $\\varepsilon$-optimal policy is approximately $\\tilde{O}\\Big( \\frac{|S| \\, |A|}{(1-\\gamma)^3 \\varepsilon^2} \\Big)$. The sample complexity of evaluating a policy (specifically, computing its successor features with respect to the MDP’s initial state, as done in our paper; see line 97) is $\\tilde{O}\\Big( \\frac{1}{(1-\\gamma)^2 \\varepsilon^2}\\Big)$. \n\nThese well-known results imply that solving an MDP requires $\\tilde{O}\\Big( \\frac{|S| \\, |A|}{(1-\\gamma)} \\Big)$ more samples than evaluating a fixed policy (e.g., computing its SFs). As $\\gamma$ approaches 1 (as is common in the literature), solving an MDP becomes orders of magnitude harder than computing its SFs. As a concrete example, consider that in deterministic environments, evaluating a policy—computing its SFs—can be done by simply executing the policy *once*, over a *single episode*, and recording the resulting rewards. By contrast, learning an optimal policy for standard RL benchmarks often requires thousands or millions of environment interactions. \n\nHence, we respectfully point out that the reviewer’s claim—that learning SFs is as hard as solving an MDP—is incorrect. This misunderstanding also appears to underlie the two subsequent points they raised; please see below.\n\n---\n**(b)**  ***“The computational cost of [our method] is dominated by the number of SF vectors, not the number of base policies”***.\n\nWe believe this claim stems from the reviewer’s observation that $\\Psi^{OK}$ stores SF vectors, along with our statement that CCS $\\subseteq \\Psi^{OK}$. However, as clarified in point **(4)** above, the cardinality of $\\Psi^{OK}$ is not related to the number of base policies (or SFs) the agent must learn, and the notation CCS $\\subseteq \\Psi^{OK}$ does not imply that our method needs to compute more policies than those in a CCS. Rather, these statements merely indicate that our method is capable of *reconstructing* a larger set of optimal policies than those in a CCS—specifically, by combining policies from the behavior basis $\\Pi_k$, where $|\\Pi_k| \\leq$ $|$CCS$|$ (line 153).\n\nThe observation that the size of $\\Psi^{OK}$ is unrelated to the number of base policies (or SFs) the agent must learn—combined with the result from point **(a)** above, which states that learning base policies requires orders of magnitude more samples than learning SFs—implies that the computational cost of our method is *not* dominated by the number of SF vectors. Instead, the primary computational bottleneck lies in training the base policies that make up the behavior basis $\\Pi_k$.\n\n\n---\n**(c)** ***The overall cost of learning $\\Psi^{OK}$ exceeds that of learning a CCS***.\n\nTo address this claim, we first note the set $\\Psi^{OK}(\\Pi_k)$, as defined in Section 3 line 151, is *never explicitly computed* by Algorithm 1. Instead, as discussed earlier, this set *characterizes* the optimal policies (more precisely, their SF vectors) that can be directly reconstructed by our method by combining the policies in $\\Pi_k$.\nThe set $\\Psi^{OK}$ that appears in Algorithm 1, on the other hand, is updated iteratively and stores the SF vectors of the CCS policies our method becomes capable of reconstructing. Hence, it is *never larger than the CCS* itself. Our algorithm terminates when this set converges to the CCS (line 12), at which point $\\Psi^{OK} = CCS$ and $|\\Pi_k| \\leq |CCS|$. \nThus, contrary to the reviewer's claim, our method does not have to learn all policies in a CCS, and certainly not more. This is formally established in Theorem 4.3, which is central to our theoretical contributions.\n\n---\n\n### Using a USFA to learn SFs\n\nWe note that the use of USFAs is common practice in the SF literature—see, e.g., Borsa et al. (2019), Kim et al. (2022), and Carvalho et al. (2023a,b). Our motivation for using a USFA is to avoid the memory and computational costs associated with training a separate neural network for each policy’s SFs. Importantly, a USFA does not necessarily introduce more approximation error than training separate networks—one per policy—as generalization in both cases depends primarily on the network architecture, capacity, and the distribution of training data.\n\nFinally, although a USFA is theoretically capable of approximating the SFs of any policies for linear tasks, in practice we train it (and rely on its predictions) only for the tasks $w$ on which the OKB is trained on. By limiting its use in this way, we reduce potential generalization errors. We agree this is an important point and will clarify it further in the updated version of the paper.\n\n\n### Proposition 4.4\n\nThe reviewer asked how to learn a meta-policy $\\omega$ in settings where the OKB is used to reconstruct the optimal policy for non-linear tasks. The reviewer is correct that when generalizing over linear tasks $w$, the meta-policy takes the form $\\omega(s, w)$. In contrast, when solving potentially non-linear tasks (as in Proposition 4.4), the meta-policy is conditioned only on the state; that is, it takes the form $\\omega(s)$. In this case, it reduces to the formulation used in the original Option Keyboard paper; see Section 2.4 and Eq. (4) for details.\n\nTo train $\\omega(s)$, we use an actor-critic policy gradient approach, following the methodology of Barreto et al. (2019). Specifically, we follow Algorithm 3 (Appendix D), replacing $\\phi$ with the scalar non-linear reward $r$, and $\\psi^{\\omega}(s, a)$ with a standard scalar Q-function $q^{\\omega}(s, a)$. We will clarify these points in the updated version of the paper.\n\n---\n\n### Minor Typo\n\nThank you for pointing out this minor typo. The correct notation is indeed $A^{w^*}_{r}(s,a)$ (i.e., $r$ subscripted), as we defined in Eq. (9). This will be fixed in the updated version of the paper.\n\n---\n\nWe hope our responses address the main point you raised regarding the computational cost of the method, as well as the other minor clarification questions. We will revise the text to incorporate a discussion about all these points. If you believe the clarifications above address the key points you raised, we'd be grateful if you would consider revisiting your score. If any points remain unclear, we welcome further feedback and would be happy to continue the discussion. Thank you again for your thoughtful comments!"}, {"review_id": "W4D5OteHUC", "reviewer": "Reviewer_g75S", "summary": "This paper is about a method for efficiently constructing an optimal behavior basis (a set of so-called base policies) for solving multi-task reinforcement learning. Based this optimal behavior basis a dynamic combination of these policies is used to solve new RL tasks. The underlying concept that make this work are successor features. The paper provides strong formal guarantees that this optimal basis allows the options keyboard (a generalization of GPI with different weights in each state) to optimally solve any new linear task. One of the theoretical contributions is a proof that this can express a strictly larger set of policies than the convex combination baseline. The approach is evaluated on high-dimensional RL problems and shows favorable performance relative to baselines.", "questions": "- See weaknesses. \n\n- How do you designed suitable successor features?", "limitations": "yes", "rating": 5, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n- The paper is well-written and addresses an important and interesting problem. \n- From what I understand, the paper is correct and provides a novel and interesting solution. However, I am not an expert in this area and cannot claim to understand the theoretical results of the paper completely. \n- The research questions are clearly phrased in the experiments section which helps understanding the evaluation. \n- Given the results, the method seems to extend to non-linear reward functions to a certain degree.\n- Evaluation is done with 15 seeds.\n- The method is shown to significantly reduce the number of needed base policies which is very relevant. \n- I am sure the paper is of high quality and I would like to see an expert reviewer make a statement about the theoretical contributions which are extensive (considering the appendix).\n\nWeaknesses\n- I a bit puzzled by the question of when such an approach is actually practical. The assumption that this approach makes seem challenging for a real-world RL tasks / settings. Perhaps, the authors can argue more for practical suitability beyond the simulated robot example.", "was_revised": true, "final_justification": "I thank the authors for their responses (also regarding the comments of the other reviews). At this point, I do not have any further questions and I remain positive about this paper and maintain my ratings.", "rebuttal": "We thank the reviewer for the positive and encouraging feedback. We are glad you found the paper to be clear, novel, and impactful. We also appreciate your comments on the correctness of our approach, the strength of our empirical evaluation, and the significance of our method’s ability to produce an optimal behavior basis significantly smaller than a CCS: this was precisely the central goal of the work.\n\nBelow, we address the reviewer’s questions and comments:\n\n---\n\n### Questions regarding the practical applicability of our method\n\nThe reviewer raised an excellent question about the practical suitability of our approach beyond robotics examples. This is indeed a crucial point, since Successor Features (SFs)—the core mathematical object underlying our method—form the foundation of many recent advances in reinforcement learning.\n\nA key reason SFs scale well to large problems is that computing them reduces to estimating a vectorized value function, which allows direct application of state-of-the-art TD learning algorithms. These techniques have been shown to scale to domains with over a million state features, including computer Go. Accordingly, SF-based methods that build upon such techniques have successfully extended to high-dimensional settings such as 3D image-based navigation and robotic locomotion (Carvalho et al., 2023a; Chua et al., 2024). Our method can directly leverage all of these advances.\n\nMoreover, please note that our approach is the first to ensure zero-shot optimality without requiring explicit construction of exponentially large Convex Coverage Sets (CCS), a key limitation in existing methods. This results in a substantial improvement in scalability over prior approaches.\n\nFinally, we would like to highlight that:\n\n**(1)** Our results go beyond demonstrating that OKB consistently outperforms state-of-the-art competitors in the most relevant and challenging benchmarks in the field (e.g., the autonomous driving environment in Figure 5). They also show that our method’s performance advantage over existing techniques becomes *increasingly pronounced* as task complexity increases.\n\n**(2)** Prior methods have only been evaluated on problems with dimensionality $d \\leq 4$ (e.g., in the original SFOLS and SIP papers). Our paper is the first to demonstrate effective, scalable, and provably optimal zero-shot transfer in problems with twice as many dimensions. This shows not only that our technique outperforms all relevant baselines, but also that it does so in problems an order of magnitude more complex than those addressed in prior work.\n\nThis is a great question and helped us realize the need to better emphasize that the core components of our approach (SFs) have already been systematically shown to scale to problems with millions of features, and that our own results provide the first empirical evidence of consistent outperformance over state-of-the-art competitors on domains with *twice* the dimensionality of those studied in prior work.\n\n---\n\n### Designing Successor Features\n\nThe paper details how the reward features $\\phi$ are designed for each domain in Appendix F.2. For example, in the Minecart domain, $\\phi$ includes quantities of collected ores and fuel consumption. In Fetch-PickAndPlace, $\\phi$ represents negative Euclidean distances to target locations. For Item Collection, $\\phi$ is composed of indicator functions for collected items, and in Highway, $\\phi$ includes normalized forward speed, lane position, and turning actions. These are hand-crafted features based on domain knowledge, which is a common practice in the SF literature.\nDetails about the architecture and training of the neural networks used to model the SFs follow standard practices in deep RL and are described in Appendix F.\n\n---\n\nWe thank the reviewer again for the positive and supportive review. Given your primary questions regarding scalability of the approach, we hope to have clarified the practical applicability of our method and SFs in general. \nWe note that the reviewer said in their review that they are “sure the paper is of high quality” and they would like to see an expert reviewer make a statement about the theoretical contributions. \nOther reviewers also pointed out that: “OKB produces a set of base policies and meta-policy which are optimal with respect to the family of tasks while requiring fewer base policies than existing methods” (Reviewer rDD6); “the proposed method is novel and theoretically sound” (Reviewer 1Yjh); “the policy basis is indeed smaller than that in previous work that also guarantee optimality” (Reviewer QsaG). \n\nIf you believe the clarifications above address the questions raised, and in light of the positive comments of the other reviewers regarding our theoretical contributions, we would be grateful if you would consider revisiting your score and confidence level. If any points remain unclear, we welcome further feedback and would be happy to continue the discussion. Thank you again for your thoughtful comments!"}]

Author rebuttal:

Meta-review: The paper proposes a method to construction an option-keyboard basis, a basis of behaviors that can be used to express the solution for all MDPS within a family, within the options keyboard framework. The argument is centered around the idea of a CCS (convex coverage set), which is a set of policies that include an optimal policy for any linear task. The challenge with a CCS is that it often grows exponentially with the number of reward features. The proposed methodology incrementally constructs a basis of behavior that can express all policies in a CCS, but significantly more compactly. The authors provide empirical and theoretical results to this effect and validate on a number of simulated problems. 

Reviewers brought up several key concerns during discussion - 1) limited baselines, 2) limited task complexity, 3) theoretical justification of complexity. Through the rebuttal period and discussion the authors were able to include more thorough comparisons to unsupervised RL methods, increase task complexity to other tasks like MO-supermario and provide an argument on the theoretical merits of their approach. The method does still need further validation and comparisons to make it really stand out in the literature, but it is a good addition to the community.
Decision: Accept (poster)