Planning with a Learned Policy Basis to Optimally Solve Complex Tasks

Venue: icaps-conference_org_ICAPS_2024_Conference
Published: 2023-12-07 11:46:09

Abstract: Conventional reinforcement learning (RL) methods can successfully solve a wide range of sequential decision problems. However, learning policies that can generalize predictably across multiple tasks in a setting with non-Markovian reward specifications is a challenging problem. We propose to use successor features to learn a set of local policies that each solves a well-defined subproblem. In a task described by a finite state automaton (FSA) that involves the same set of subproblems, the combination of these local policies can then be used to generate an optimal solution without additional learning. In contrast to other methods that combine local policies via planning, our method asymptotically attains global optimality, even in stochastic environments.

Main paper: 
 introduction: 
 Autonomous agents that interact with an environment usually face tasks that comprise complex, entangled behaviors over long horizons. Conventional RL methods have successfully addressed this. However, in cases when the agent is meant to perform several tasks across similar environments, training a policy for every task separately can be time-consuming 20 and requires a lot of data. In such cases, the agent can utilize a method that has built-in generalization capabilities. One such method relies on the assumption that reward functions of these tasks can be decomposed into a linear combination of successor features (Barreto et al. 2017) . When a new task is presented, it is possible to combine previously learned policies and their successor features to solve a new task. While the combination of such policies is guaranteed to be better than any previously learned policy, it need not be optimal. However, as was recently shown by Alegre, Bazzan, and Da Silva (2022) , one can utilize recent advancement in multiobjective RL to learn a set of policies that contains an optimal policy for any linear combination of successor features.While these and other conventional RL methods use Markovian reward functions, expressing a task with such reward function can be difficult and may not even be possible in some cases (Whitehead and Lin 1995) . In settings where the reward function cannot be expressed in Markovian terms, task specification has raised especial interest in the last few years (Toro Icarte et al. 2018a; Camacho et al. 2019) . In our work we focus on developing a method that utilizes generalization capabilities of successor features in a settings with non-Markovian reward functions.Prior techniques for such settings have been proposed in contexts where the existence of a set of propositional symbols allowed for the definition of high-level tasks by the means of logics (Vaezipoor et al. 2021; Toro Icarte et al. 2019) or finite state automata (Toro Icarte et al. 2018a ). Similar to several hierarchical reinforcement learning methods, they are often based on decomposing tasks into sub-tasks and solving each sub-task independently (Dietterich 2000; Sutton, Precup, and Singh 1999) . While this has an advantage that solving parts of a task in isolation is often simpler, the disadvantage of such approach is that the combination of optimal solutions to sub-tasks can lead to a sub-optimal overall policy. This is commonly referred to as recursive optimality (Dietterich 2000) or myopic policy (Vaezipoor et al. 2021) .To alleviate this issue, one can consider methods that condition the policy or the value function on the specification of the whole task (Schaul et al. 2015) and such approaches were recently also proposed for tasks with non-Markovian reward functions (Vaezipoor et al. 2021) . However, the methods that specify the whole task usually rely on a blackbox neural network (NN) for planning when determining which sub-goal to reach next. This makes it hard to interpret the plan to solve the task and although they achieve promising empirical results it is unclear whether and when these approaches will generalize to a new task.Instead, our work aims to use task decomposition without sacrificing global optimality to achieve predictable generalization. The method we propose learns a set of local policies in sub-tasks such that their combination forms a globally optimal policy for a large collection of problems described with FSAs. A new policy that solves any new task can then be created without additional learning by planning on a given FSA task description. Our contributions are:• We propose to use successor features to learn a policy basis that is suitable for planning in stochastic domains.• We develop a planning framework that uses such policy bases for zero-shot generalization to complex temporal tasks described by an arbitrary FSA.• We prove that if the policies in this basis are optimal, our framework produces a globally optimal solution even in stochastic domains.Given a finite set X , let ∆(X ) = {p ∈ R X : x p(x) = 1, p(x) ≥ 0 (∀x)} denote the probability simplex on X . Given a probability distribution q ∈ ∆(X ), let supp(q) = {x ∈ X : q(x) > 0} ⊆ X denote the support of q.Reinforcement learning problems commonly assume an underlying Markov Decision Process (MDP). We define an MDP as the tuple M = ⟨S, E, A, R, P 0 , P, γ⟩ where S is the set of states, E is the set of exit states, A is the action space, R : S ×A×S → R is the reward function, P 0 ∈ ∆(S) is the probability distribution of initial states, P : S × A → ∆(S) is the transition probability function and 0 ≤ γ < 1 is the discount factor. The goal of the agent is to find an optimal policy π * that maximizes the expected discounted return, for any stateaction pair (s, a) ∈ S × A,EQUATIONwhereR i = R(S i , A i , S i+1 ). Hence, an optimal policy is 110 π * ∈ arg max π Q π (s, a) ∀(s, a) ∈ S × Awith ties broken arbitrarily. The action value function defined in Equation (1) satisfies the recursive Bellman equationQ π (s, a) = E s ′ ∼P(•|s,a) R(s, a, s ′ ) + V π (s ′ ) (2)for any (s, a) ∈ S × A. The state value function is obtained by maximizing the action value function over the actions, V π (s) = max a Q π (s, a) ∀s ∈ S. Throughout the paper we use Q * and V * to refer to the optimal, respectively, action and state value functions.Successor features (SFs) (Dayan 1993; Barreto et al. 2017 ) is a widely used RL representation framework that assumes the reward function is linearly expressible with respect to a feature vector,EQUATIONHere, ϕ : S ×A×S → R d maps transitions to feature vectors and w ∈ R d is a weight vector. Every weight vector w induces a different reward function and, thus, a task. The SF vector of a state-action pair (s, a) ∈ S × A under a policy π is the expected discounted sum of feature vectors following ,EQUATIONwhere ϕ i = ϕ(S i , A i , S i+1 ). The action value function for a state-action pair (s, a) under policy π can be efficiently represented using the SF vector. Due to the linearity of the reward function, the weight vector can be decoupled from the Bellman recursion. Following the definition of Equations ( 1) and(3), the action value function in the SF framework can be rewritten asQ π w (s, a) = E π ∞ i=t γ i-t w ⊺ ϕ i S t = s, A t = a = w ⊺ E π ∞ i=t γ i-t ϕ i S t = s, A t = a . (5)Once the SF representation is obtained, this allows generalized policy evaluation (GPE) as a consequence of Equation (5) (Barreto et al. 2020) . Simultaneously, generalized policy improvement (GPI) can be used to obtain new better policies (Barreto et al. 2017) .A family of MDPs is defined as the set of MDPs that share all the components, except the reward function. This set can be properly defined asM ϕ ≡ {⟨S, E, A, R w , P 0 , P, γ⟩|R w = w ⊺ ϕ, ∀w ∈ R d }.Transfer learning on families of MDPs is possible thanks to GPI. Given a set of policies Π, learned on the same family M ϕ , for which their respective SF representations have been computed, and a new task by w ′ ∈ R d , a GPI policy π GPI is derived asEQUATIONThe new policy π GPI is known to perform no worse than any policy π ∈ Π in the new task w ′ (Barreto et al. 2017, cf. Theorem 1). However, there is no guarantee with respect to the optimal behavior for such task. A fundamental question to solve the so-called optimal transfer learning problem is which policies should be included in the set of policies Π so the optimal policy for any weight vector w ∈ R d can be derived.The recent work of Alegre, Bazzan, and Da Silva (2022) solves the optimal transfer learning problem. They draw the connection between the SF transfer learning problem and multi-objective RL (MORL). The pivotal fact is that the SF representation (4) can be interpreted as a multidimensional value function and then solved using multiobjective optimization techniques. They use the optimistic linear support (OLS) algorithm to learn a set of policies that constitutes a convex coverage set (CCS) (Roijers, Whiteson, and Oliehoek 2015) . In terms of its multi-objective value,EQUATIONIntuitively, all policies in Π CCS are optimal in at least one task w ∈ R d . This also means that for every task w d there exists π ∈ Π CCS such that π is optimal for w.We assume that environments are endowed with set of highlevel, boolean-valued propositional symbols P and that they are associated with the set of exit states E of a low-level MDP M. Every transition ∈ S × A × S induces some propositional valuation (assignment of truth values) 2 P . Such a valuation depends on the new state and occurs under a mapping O : S → 2 P that is known to the agent. Nonetheless, only exit states ε ∈ E make propositions true under O. We assume that that propositional symbols are mutually exclusive, and the agent cannot observe two symbols in the same transition. We say that a valuation Γ satisfies a propositional symbol p, formally Γ ⊨ p, if p is true in Γ. E.g., in the office domain depicted in Figure 1a , the propositional symbols are P = { , , o} while the exit statesE = { 1 , 2 , 1 , 2 , o 1 , o 2 }. Consequently, O( 1 ) ⊨ and O( 2 ) ⊨ .Task instructions can be specified via a finite state automaton.These are tuples F = ⟨U, u 0 , T , L, δ⟩ where U is the finite set of states, u 0 ∈ U is the initial state, T is the set of terminal states with U ∩ T = ∅, L : U × (U ∪ T ) → 2 P is a labelling function that maps FSA states transitions to truth values for the propositions and δ : U → {0, 1} is a highlevel reward function. Each transition among FSA states (u, u ′ ) defines a subgoal. The agent has to observe some propositional valuation L(u, u ′ ) in order to achieve it and FSA states can only be connected by a subgoal. Non-existing transitions (u, u ′ ) get mapped to L(u, u ′ ) = ⊥. In Figure 4c , the FSA state u 0 has two outgoing subgoals: getting mail (labeled as ) and getting coffee (labeled as ). The reward function δ gives a reward larger than 0 only to terminal states.In other words, such a reward function is δ(u) = 0 ∀u ∈ U and δ(t) = 1 ∀t ∈ T .We focus on the setting in which a low-level MDP is equipped with a reward structure like in Equation (3). We let the lowlevel be represented by a family of MDPs M ϕ , where each weight vector w ∈ R d specifies a low-level task. The agent receives high-level task specifications in the more flexible form of an FSA which permits the specification of non-Markovian reward structures. The combination of a low-level family of MDPs and a high-level FSA gives rise to a product MDP M ′ = F × M ϕ , that satisfies the Markov property and where the state space is augmented to be U × S.A product MDP M ′ is a well-defined MDP and can be solved with conventional RL methods such as Qlearning (Watkins and Dayan 1992) . This is, however, impractical since policies should be retrained every time a new high-level task is specified. Exploiting the problem structure is essential for tractable learning, where components can be reused for new task specifications. The special reward structure of the low-level MDPs and our particular choice of feature vectors, later introduced, allow us to define an algorithm able to achieve a solution by simply planning in the space of augmented exit states U × E. This inherently makes obtaining the optimal policy more efficient than solving the whole product MDP, as we reduce the number of states on 230 which it is necessary to compute the value function.When presented with different task specifications (e.g. Figure 4 ), the agent may have to perform the same subtask at different moments of the plan or in different FSAs. We aim to provide agents with a collection of base behaviors that can 235 be combined to retrieve the optimal behavior.In line with the previous reasoning, we introduce a twostep algorithm in which the agent first learns a Π CCS (a set of policies that constitute a CCS) on a well-specified representation of the environment. Then these subpolicies are 240 used to solve efficiently any FSA task specification on the propositional symbols of the environment. In what follows, we motivate the design of the feature vectors, explain our high-level dynamic programming algorithm and prove that it achieves the optimal solution.Feature vectors: For a family of MDPs M ϕ , feature vectors ϕ(s, a, s ′ ) are |E|-dimensional. Each feature component ϕ j is associated with an exit state ε j ∈ E = {ε 1 , . . . , ε |E| }. Such vectors are built as follows. At terminal transitions (s, a, ε i ) ∈ T , ϕ j = 1 when j = i and ϕ j = 0 when 250 j ̸ = i. For non-terminal transitions, we just require that w ⊺ ϕ(s, a, s ′ ) < 1. In the case that ϕ(s, a, s ′ ) = 0 ∈ R |E| for non-terminal transitions, the SF representation in Equation (4) of each policy consists of a discounted distribution over the exit states. This indicates how likely it is to reach 255 each exit state following such a policy. Furthermore, we require that E ⊂ supp(P 0 ) so the value functions at exit states are well-defined.Algorithm 1: SF-FSA-VIInput: Low-level MDP M ϕ , task specification F 1: Obtain Π CCS on M ϕ . 2: Initially w 0 (u) = 0 ∈ R |E| ∀u ∈ U. 3: while not done do 4:for u ∈ U do 5:Update each w k+1 j (u) with Equation (10).6: return {w * (u) ∀u ∈ U} Algorithm The solution to an FSA task specification implies solving a product MDP M ′ = F × M ϕ . Since we 260 have the CCS, the optimal Q-function can be represented by a weight vector w * :EQUATIONfor all (u, s, a) ∈ U × S × A. Here, w * j (u) indicates the optimal value from exit state ε j ∈ E for FSA state u. Thus, we observe that finding the optimal weight vectors w * (u), 265 ∀u ∈ U, is enough for retrieving the optimal action value function of the product MDP M ′ . We can obtain this vector using a dynamic-programming approach similar to value iteration:EQUATION= max a,πEQUATIONwhere τ (u, O(ε)) ∈ U is the FSA state that results from achieving the valuation O(ε) in u. We know that, w k j (u) = 1 if τ (u, O(j)) = t, per definition, since the high-level reward function δ(t) = 1. As k → ∞, our algorithm converges to the optimal set of weight vectors {w * (u)} u∈U and, as a result, to the optimal value function in Equation ( 8).We first restate the following theorem from Alegre, Bazzan, and Da Silva (2022) .Theorem 1 (Alegre, Bazzan, and Silva, 2022) . Let Π be a set of policies such that the set of their expected SFs, Ψ = {ψ π } π∈Π , constitutes a CCS. Then, given 280 any weight vector w ∈ R d , the GPI policy π GP I w (s) ∈ arg max a∈A max π∈Π Q π w (s, a) is optimal with respect to w :V GP I w = V *w . Applied to our setting, once the set of policies Π CCS and associated SFs have been computed, we can define an arbitrary vector w of rewards on the exit states, and use the CCS to obtain an optimal policy π * w and an optimal value function V * w without learning. We can then use composition by setting the reward of the exit states equal to the optimal value.We aim to show that for each augmented state (u, s) ∈ U × S, the value function output by our algorithm equals the optimal value of (u, s) in the product MDPM ′ = F × M ϕ , i.e. that V w(u) (s) = V * M ′ (u, s).To do so, it is sufficient to show that the weight vectors {w(u)} u∈U are optimal.Each element of w(u) is recursively defined as w j (u) = V w(τ (u,O(εj ))) (ε j ). If all weight vectors are optimal, it holds that V w(τ (u,O(εj ))) (ε j ) = V * M ′ (w(τ (u, O(ε j ))), ε j ) for each such exit state. Due to the above theorem, the value function V w(u) is optimal for w(u). Due to composition that follows GPE and GPI, this means that the value of each internal state s is optimal, i.e. that V w(u) (s) = V * M ′ (u, s).It remains to show that the weight vectors {w(u)} u∈U returned by the algorithm are indeed optimal. To do so it is sufficient to focus on the set of augmented exit states U × E. We can state a set of optimality equations on the weight vectors as follows:w * j (u) = V w * (τ (u,O(ε))) (ε j ) = max a Q * (τ (u, O(ε)), ε j , a) = max a max π ψ π (ε j , a) ⊺ w * (τ (u, O(ε))),where ψ π (ε j , a) = s ′ P(s ′ |ε j , a)ψ π (ε j , a, s ′ ). Due to the discount factor γ, we have ∥ψ(ε j , a)∥ 1 < 1. Hence the update rule in ( 10) is a contraction and converges to the set of optimal weight vectors.We test our method in two complex discrete environments. At test time, we change the reward to -1 for every timestep and use the cumulative reward as performance metric. We report two types of results. First, we are interested in how fast our method can build a global solution for the high-level task specification as it learns the base subpolicies. For this, we let the algorithm learn the base sub-subpolicies and after some interactions fully retrain the high-level policy and test its performance. Second, once the base behaviors are learned, we measure how many of iterations Algorithm 1 needs to converge to an optimal solution for different task specifications.In any case, we compare against existing baselines.We use the Delivery domain (Araki et al. 2021 ) and a modified version of the Office domain (Toro Icarte et al. 2018a) as test beds for our algorithm. Both environments are depicted in Figure 1 and present a propositional vocabulary that is rich enough to build complex tasks. In the Delivery domain there is a single low-level state associated with each of the propositional symbols, implying that E = P = {A, B, C, H}. The feature vectors are consistent with our design choice. For terminal transitions, they correspond to their one-hot encodings of the terminal states. There exist obstacle states (in black) for which, upon entering, the feature vector is ϕ(s, a, s ′ ) = -1000 ∈ R 4 . This transforms in a large negative reward when multiplied with a corresponding weight vector w ∈ R 4 . For regular grid cells (in white) ϕ(s, a, s ′ ) = 0 ∈ R 4 . The Office domain is more complex since there are three propositional symbols 340 P = { , , o} which can be satisfied at different locations, namelyE = { 1 , 2 , 1 , 1 , o 1 , o 2 }.Here, there are no obstacle states and ϕ(s, a, s ′ ) = 0 ∈ R 6 for non-terminal transitions.For each of the environments we define three different tasks: sequential, disjunction and composite (all described in Figure 4 ). The sequential task is meant to show how our algorithm can indeed be effectively used to plan over long horizons, when the other two tasks show the ability of our method to optimally compose the base subpolicies in complex settings. In natural language, the tasks in the Delivery domain correspond to: "go to A, then B, then C and finally H" (sequential), "go to A or B, then C and finally H" (disjunction) and "go to A and B in any order, then B, then C and finally H" (composite). The agent has to complete the tasks by avoiding obstacles. The counterpart of these tasks in the Office environment are: "get a coffee, then pick up mail and then go to an office" (sequential), "get a coffee or mail, and then go to an office" (disjunction) and "get a coffee and mail in any order, and then go to an office" (composite).Our agent never learns how to solve these tasks, but rather learns the set of subpolicies that constitutes the CCS. At test time, we provide the agent with the FSA task specification, extract the high-level value function and test its performance on solving the task.In the literature, we find the most similar approach to ours in the Logical Options Framework (LOF, (Araki et al. 2021) ). We thus use LOF and flat Qlearning on the product MDP as baselines. LOF trains one option per exit state, which are trained simultaneously using intra-option learning, and then uses a high-level value iteration algorithm to train a meta-policy that decides which option to execute in each of the MDP states. On the other hand, the latter learns the action value function in the flat product MDP, from which it extracts the policy. Under certain conditions, flat Q-learning converges to the optimal value function but, especially for longer tasks, it may take a large number of samples. Additionally, it is trained for a specific task, so it is not able to generalize to other task specifications. For LOF, we followed the implementation details prescribed by the authors.Empirical results are shown in Figure 2 . The plots on the left column reflect how LOF, flat Q-learning and our method perform at solving an FSA task specification during the learning phase. Results are averaged over the three tasks (sequential, disjunction and composite) previously described for each environment. In both environments, our method is the first to reach optimal performance. There exist, however, some differences between LOF and our method. LOF trains all options simultaneously with intra-option learning. This means that, every transition (s, a, s) is used to update all options' value functions and policies. In our case, the learning of a Π CCS is done sequentially, with a fixed sample budget per subpolicy, we use a total of 8 • 10 3 samples per subpolicy in both environments. Both options and the SF representation of subpolicies are learned using Q-learning. Though it is true that our method uses experience replay to speed up the sub-policies' learning, the sequential learning may be seen as a disadvantage with respect to LOF since it needs to have fully learned some base behaviors to be able to build a feasible solution. This is clearly observed in the Delivery domain (Figure 2 , top left), where at the early stages of the interaction, our method achieves very low cumulative reward. This is because it takes some time for the agent to learn enough subpolicies and can offers a plausible solution. Nonetheless, it eventually converges to the optimal policy. Similarly, LOF converges to the optimal policy albeit it takes slightly longer to learn. In the more complex Office environment, results follow the same pattern. However, this environment breaks one of the of LOF requirements for optimality: to have a single exit state associated with each predicate. In this problem, for each predicate there exist two exit states that can satisfy them. This makes LOF prone to converge to suboptimal solutions when our algorithm attains optimality. This is the case for the composite task, where LOF is short-sighted and returns a longer path (in red, Figure 1a ) in contrast to ours that retrieves the optimal solution (in green, Figure 1a ). This means that our method is more flexible in the task specification. In this environment, our algorithm also converge faster with a more obvious gap with LOF. In any case, learning sub-policies or options is faster than learning on the flat product MDP, as flat Q-learning takes the longest to converge. The right column in Figure 2 shows that our algorithm needs less iterations to learn an optimal representation of the global optimal value function via planning when compared to LOF. In LOF, the cost of each iteration of value iteration is |U| × |S| × |K|, where K is the set of options, while for the Algorithm 1 we propose it is |U| × |E| × |Π CCS |. By definition, the number of options is equivalent to the number of exit states |K| = |E|, so a single iteration of our method is more efficient than LOF whenever |Π CCS | ≪ |S|. In our experiments, the sizes of the CCS are 15 and 12 for the Delivery and Office domains, respectively, while the sizes of the state spaces are of 225 and 121. Therefore, since our algorithm needs less shorter iterations during planning, it outperforms LOF in terms of planning speed in both domains when composing the global solution.In deterministic environments, it is sufficient to learn the subpolicies associated with the extrema weights ( i.e. those subpolicies that reach each of the exit states individually) to find a globally optimal policy via planning. In such cases, it may not be necessary to learn a full CCS. That is why, approaches that use the options framework such as LOF traditionally define one option per subgoal. However, there are scenarios, in which these approaches will not find optimal policy. This is the case for most stochastic environments. For example, consider the very simple domain of Double Slit in and the FSA task specification in Figure 3 . In this environment, there are two exit states E = {blue, red}. The agent starts in the leftmost column and middle row. At every timestep, the agent chooses an action amongst {UP, RIGHT, DOWN} and is pushed one column to the right in addition to moving in the chosen direction, except in the last column. If the agent chooses RIGHT, he moves an extra column to the right.At every timestep there is a random wind that can blow the agent away up to three positions in the vertical direction. The FSA task specification represents a task in which the agent is indifferent between achieving either of the goal states. Since the RIGHT action brings the agent closer to both goals, the 460 optimal behavior in this case is to commit to either goal as late as possible. In this setting, methods that use one policy per sub-goal, such as LOF, train two policies to reach both goals. This means that the agent has to commit to one of the goals from the very beginning, which hurts the performance 465 as it has to make up for the consequences of the random noise. On the other hand, the CCS used by our method will contain an additional policy that is indifferent between two goals. This leads to a performance gap as our approach achieves an average accumulated reward of -19.7 ± 3.65 and LOF 470 -22.70 ± 5.72. One of the key distinctions in our research compared to prior studies is the optimality of the final solution. As noted by Dietterich (2000) , hierarchical methods usually have the capa-475 bility to achieve hierarchical, recursive, or global optimality.The challenge that often arises when sub-task policies are trained in isolation is that the combination of these locally optimal policies does not lead to a globally optimal policy but a recursively (Dayan and Hinton 1992) or hierarchically opti-480 mal policy (Sutton, Precup, and Singh 1999; Mann, Mannor, and Precup 2015; Araki et al. 2021) . To tackle this challenge, our approach relies on acquiring a set of low-level policies for each sub-task and employing planning to identify the optimal combination of low-level policies when solving a particular 485 task. By learning the CCS with OLS (Roijers, Whiteson, and Oliehoek 2014) in combination with high-level planning our approach ensures that globally optimal policy is found. In this regard, the work of Alegre, Bazzan, and Da Silva ( 2022) is of particular interest as it was the first work that used OLS and 490 successor features (Barreto et al. 2017) for optimal transfer learning. However, this method has only applied in a setting with Markovian reward function and has not been used with non-Markovian task specifications or high-level planning.On the other hand, many recent approaches proposed to use 495 high-level task specifications in the form of LTL (Toro Icarte et al. 2018b; Kuo, Katz, and Barbu 2020; Vaezipoor et al. 2021; Jothimurugan et al. 2021) , or similar formal language specifications (Toro Icarte et al. 2019; Camacho et al. 2019; Araki et al. 2021; Toro Icarte et al. 2022 ) to learn policies. 500 From these works, our approach is the most similar to the Logical Options Framework (Araki et al. 2021) . The main difference is that LOF trains a single policy for each sub-goal, resulting in a set of learned policies that is either smaller than or equal to the set acquired through our method. While employing one policy per sub-goal proves sufficient for obtaining a globally optimal policy through planning in deterministic environments (Wen et al. 2020) , this may not hold true in stochastic environments, as our experiments 515 demonstrate. In such instances, the policies generated by LOF are hierarchically optimal but fall short of global optimality.u 0 start u 1 u 2 u T o (a) u 0 start u 1 u 2 u T o o (b) u 0 start u 1 u 2 u 3 u T o (c) u 0 start u 1 u 2 u 3 u T A B C H (d) u 0 start u 1 u 2 u 3 u T A B C C H (e) u 0 start u 1 u 2 u 3 u 4 u T A B B A C H (f)Two notable examples from aforementioned works on multi-task learning with formal language specifications are the works of Toro Icarte et al. (2018b) and Vaezipoor et al. (2021) . The former struggles with generalizing to unseen tasks, because it uses LTL progression to determine which sub-tasks need to be learned to solve given tasks. The Qfunctions that are subsequently learned for each LTL sub-task will therefore not be useful for a new task if its sub-tasks were 525 not part of the training set. Such limitation does not apply to the latter as it instead encodes the remaining LTL task specification using a neural network and conditions the policy on this LTL embedding. While this approach may be more adaptable to tasks with numerous propositions or sub-goals, it risks generating sub-optimal policies as it relies solely on the neural network to select the next proposition to achieve, without incorporating planning. Additionally, since the planning is implicitly done by the neural network, the policy is less interpretable than when explicit planning is used.The method we propose can be viewed as a method for composing value functions through successor features, akin to previously proposed approaches for composition of value functions and policies (van Niekerk et al. 2019; Barreto et al. 2019; Nangue Tasse, James, and Rosman 2020; Infante, Jonsson, and Gómez 2022) . In the work of Infante, Jonsson, and Gómez (2022), which is the closest to our work, the authors propose to learn a basis of value functions that can be combined to form an optimal policy. However, unlike our method, their approach only works in a restricted class of linearly-solvable MDPs. Lastly, since our approach uses the values of exit states for planning it is also related to planning with exit profiles (Wen et al. 2020) . The CCS that we propose to use as a policy basis in our work can be seen as a collection of policies that are optimal for all possible exit profiles.In this work, we address the problem of finding optimal behavior for new non-Markovian goal specifications in known environments. To do so, we introduce a novel approach that uses successor features to learn a policy basis, that can subsequently be used to solve any unseen task specified by an FSA with the set of given predicates P by planning. Our method is the first that can provably generalize to such new task specification without sacrificing optimality in both deterministic and stochastic environments.The experiments show that our method offers several advantages over previous methods. First, due to the use of SF, it allows for faster composition of the high-level value function since it drastically reduces the number of states to plan on. Secondly, thanks to using a CCS over a set of options our method achieves optimality even in stochastic environments (as shown in the Double Slit example). Lastly, we do not require that there exists a single exit state per predicate which permits more flexible task specification while at the same time allowing deployment in more complex environments.A limitation of our approach could be the need to construct a full CCS if one wants to attain global optimality. While the construction of CCS is not timecomsuming for environments with several exit states presented in our work, the computation cost of finding the full CCS could become too large for environments with many exit states. In such case one could instead learn a partial CCS at the cost of a bounded decrease in performance (Alegre, Bazzan, and Da Silva 2022) or consider splitting the environment into smaller parts with less exit states. While our experiments only considered discrete environments, our method should also be applicable in continuous environments with minor adjustments. These include: using an contiguous set of states instead of a single exit state and using reward shaping to facilitate learning in sparse reward setting.Whitehead, S. D.; and Lin, L.-J. 1995. Reinforcement learning of non-Markov decision processes. Artificial intelligence, 73(1-2): 271-306. 
 
 Appendix:  


Decision: The basic idea of the paper is a technique for learning local policies for sub-tasks that can be composed to obtain optimal global policies for non-markovian composite tasks where the reward is a linear combination of the sub-tasks.

The reviewers agree the method is interesting and there is novelty in the idea.

Two major points were discussed by the reviewers.

1. Clarity: The paper is quite dense in the technical explanations and the CCS should be better reported in the main text to make the paper self-contained. In general, a high-level description of the approach at the beginning would help the understanding of the whole article.
2. The experiments are not very comprehensive though, but they do show a reasonable amount of potential.

All in all, the general agreement was that the paper has no major flaws and introduces interesting ideas, even if there is room for improvement especially in experiments and presentation. We recommend the authors to report the rebuttal answers in the main text and make any effort to improve the clarity of the presentation.

Reviews: review: strength_weakness: strength_weakness: I found the abstract/introduction initially hard to reconcile with the rest of the paper. The abstract and introduction talk about local policies, and I had difficulty seeing the connection with the rest of the paper. I eventually came to understand these as the policies in the CCS which one would expect to attempt to reach the terminal states. But this connection was left largely implicit for the rest of the paper, and I think this should be better connected.

I also am unclear about some of the limitations of this method. Constructing the CCS is listed as one, but it also seems to rely on the fact that the set of tasks that are possible can only be captured by an (arbitrarily complex) non-Markovian task that involves visiting a small number of exit states. In particular, since a successor feature is created per exit state (instead of per proposition), it seems like that must be kept low to make this feasible. However, I would be interested in hearing more comment on this topic.

While I suppose this is left to future work, there is also the big question of how this would extend to a function approximation setting.

I had some issues understanding the experimental design, partially because I am not very familiar with either CCS or LOF. In both cases, it feels like a bit more background is needed to understand the described method and how they were evaluated. In the case of CCS, I think further information is needed on the how the policies are generated. The results section suggests these are generated sequentially, but are elements of that set added to or discarded? Or is it some sort of process where one is added at a time until some sort of "complete set" is found. Or are the policies generated first, then evaluated? The new method also seems like it would involve the initial learning phase, following by the planning phase after which learning is no longer needed. Its unclear to me in the charts where this crossover happens.

Because of the distinct learning vs planning phase, I am unsure how to think about runtime here, especially in comparison to LOF (which I assume is a standard non-planning method?). How long does the planning take, for example?

In general, the experiments do not feel especially comprehensive, as they focus on two relatively small problems, and a small amount of tasks between them. They do show what you would hope for, but they do not go in especially great detail.

Minor Points
- Related to this point about learning vs planning time, it might be helpful to see how long it takes to converge to the optimal policy on a single task after the CCS has been computed. How this compares to LOF for the different tasks could be interesting
- Is the CCS really explainable? It finds 12 policies for Office World and 15 for Delivery. Those are non-intuitive numbers, and I'm curious how one could extract meaning from the found policies.
questions: questions: 1) Please clarify how the charts correspond to a planning phase versus a learning phase. Perhaps a summary of how the CCS and the successor feature vectors for these is computed would be helpful here. 
2) Please clarify how the method would scale with a larger set of exit states vs the number of propositions.
strength_weakness: strength_weakness: This paper presents a theoretical description of the problem that is not completely clear on some points. For instance:
1) Authors mentioned that In line with the previous reasoning, we introduce a two-step algorithm in which the agent first learns a Π CCS (a set of policies that constitute a CCS) on a well-specified representation of the environment.” → However, it is not clear the difference between the proposed approach with respect to the state-of-the-art methods that pre-train the policies that are tuned for specific tasks.
(2) Please see the other aspects mentioned in the next section.

Another weakness of the paper is the lack of a link/detailed description of the planning framework that uses the proposed policy (such a framework is mentioned among the contributions of the paper). It would be helpful to add the link to the code.
questions: questions: The questions are reported below:

(1) Please explain the domain in E.g., in the office domain depicted in Figure 1a with a sentence under the “Propositional Logic” section. This helps to clarify the logical consequence.

(2) The sentence “This inherently makes obtaining the optimal policy more efficient than solving the
whole product MDP, as we reduce the number of states on which it is necessary to compute the value function” is very strong. It has to be demonstrated. 

(3) The sentences “we can define an arbitrary vector w of rewards on the exit states, and use the CCS to obtain an optimal policy π  and an optimal value function V w without learning. We can then use composition by setting the reward of the exit states equal to the optimal value” needs to be clarified. How the optimal value is computed?

(4) The computation of Equation 11 and the following steps for achieving that it converges to optimal weights should be further explained.

(5) In the experiment section, the ways to train the sub-subpolicies should be mentioned (which algorithms, which rewards, number of epochs, number of steps per epoch, hardware for training, etc.)

(6) Similarly to the previous point, please specify the parameters related to the next re-training. 

(7) “We use a total of 8 · 10 3 samples per subpolicy in both environments” → Please check the number of samples with respect to the results in Figure 2.

(8) Additionally, since the planning is implicitly done by the neural network, the policy is less interpretable than when explicit planning is used. → Explain and motivate how the proposed approach is more explainable.

(9) Given the discussion reported at the end of the paper, it would be interesting to compare the proposed approach with respect to the baseline in terms of computation cost.


Minor:

(10) Please specify the meaning of FSA (finite state automaton (FSA)) also in the introduction not only in the abstract

(11) The figure should be cited in ascendant order in the text. Fig. 4 appears sequentially after Figure 1 and before Figure 1.

(12) Typo: provably
strength_weakness: strength_weakness: I just suggest the authors to include a leading example in the methodological section, e.g., from the office domain, for easier readability.
Also, please clarify in the caption the meaning of lines in figure 1a (office domain sketch).
questions: questions: I just require to fix minor presentation weaknesses as above.
Also, it would be useful for the community to have access to an open source implementation of the methodology. I encourage the authors to share this, upon acceptance.