Compositional Automata Embeddings for Goal-Conditioned Reinforcement Learning

Authors: beyazit  yalcinkaya
niklas  lauffer
marcell  vazquez-chanlatte
sanjit a seshia
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 01:51:44

Abstract: Goal-conditioned reinforcement learning is a powerful way to control an AI agent's behavior at runtime. That said, popular goal representations, e.g., target states or natural language, are either limited to Markovian tasks or rely on ambiguous task semantics. We propose representing temporal goals using compositions of deterministic finite automata (cDFAs) and use cDFAs to guide RL agents. cDFAs balance the need for formal temporal semantics with ease of interpretation: if one can understand a flow chart, one can understand a cDFA. On the other hand, cDFAs form a countably infinite concept class with Boolean semantics, and subtle changes to the automaton can result in very different tasks, making them difficult to condition agent behavior on. To address this, we observe that all paths through a DFA correspond to a series of reach-avoid tasks and propose pre-training graph neural network embeddings on "reach-avoid derived" DFAs. Through empirical evaluation, we demonstrate that the proposed pre-training method enables zero-shot generalization to various cDFA task classes and accelerated policy specialization without the myopic suboptimality of hierarchical methods.

Main paper: 
 introduction: 
 Goal-conditioned reinforcement learning (RL) [33] has proven to be a powerful way to create AI agents whose task (or goal) can be specified (conditioned on) at runtime. In practice, this is done by learning a goal encoder, i.e., a mapping to dense vectors, and passing the encoded goals as inputs to a policy, e.g., a feedforward neural network. This expressive framework enables the development of flexible agents that can be deployed in a priori unknown ways, e.g., visiting states never targeted during training. The rise of large language models has popularized leveraging natural language as an ergonomic means to specify a task, e.g., "pick up the onions, chop them, and take them to stove." While incredibly powerful, tasks specified by target states or natural language have a number of shortcomings. First and foremost, target states are necessarily limited to non-temporal tasks. On the other hand, natural language is, by definition, ambiguous, providing little in the way of formal guarantees or analysis of what task is being asked of the AI agent.To this end, we consider conditioning on tasks in the form of Boolean combinations of deterministic finite automata (DFAs). We refer to this concept class as compositional DFAs (cDFAs). The choice of cDFAs as the concept class is motivated by three observations. First and foremost, DFAs offer simple and intuitive semantics that require only a cursory familiarity with formal languages-if one can understand a flow chart, one can understand a DFA. Moreover, recent works have demonstrated that DFA and cDFA can be learned in a few shot manner from expert demonstrations and natural language descriptions [41] . As such, DFAs offer a balance between the accessibility of natural language and rigid formal semantics. The addition of Boolean combinations to cDFA, e.g., perform task 1 AND task 2 AND task 3, provides a simple mechanism to build complicated tasks from smaller ones.Second, DFAs represent temporal tasks that become Markovian by augmenting the state-space with finite memory. Further, they are the "simplest" family to do so since their finite states are equivalent to having a finite number of sub-tasks, formally Nerode congruences [19] . This is particularly important for goal-conditioned RL which necessarily treats temporal tasks differently than traditional RL. For traditional RL, because the task is fixed, one can simply augment the state space with the corresponding memory to make the task Markovian. In goal-conditioned RL, this is not, in general, possible as it is unclear what history will be important until the task is provided. Instead, the encoded task must relay to the policy this temporal information. Third, existing formulations like temporal logics over finite traces and series of reach-avoid tasks are regular languages and thus are expressible as DFAs [11] . This makes DFAs a natural target for conditioning an RL policy on temporal tasks.The expressivity of DFAs introduces a number of challenges for goal-conditioned RL. First, DFAs form a countably infinite and exponentially growing concept class where subtle changes in the DFA can result in large changes in an agent's behavior. Notably, this means that any distribution over DFAs is necessarily biased with many "similar" DFAs having drastically different probability. Thus, to generally work with DFAs, one cannot simply match finite patterns, but need to learn to encode details necessary for planning. Second, as with traditional goal-based objectives, DFAs provide a very sparse binary reward signal-did you reach the accepting state or not? Together with non-trivial dynamics, naïve applications of RL become infeasible due to the lack of dense reward signal. Finally, the exponentially-expanding concept class presents computational limitations for encoding. For example, many interesting DFAs may be too large to be feasibly processed by a graph neural network.To address these issues of reward sparsity and the need to encode planning, we introduce a distribution of DFAs, called reach-avoid derived (RAD). This concept class is inspired by the observation that all paths through a DFA correspond to a series of (local) reach-avoid problems. We argue in Section 4 that RAD encourages learning to navigate a DFA's structure. Our first key result is that pre-training DFA encoders on RAD DFAs enables zero-shot generalization to other DFAs. Next, we treat the problem of DFA size. Many problems are naturally expressed compositionally, e.g., a sequence of rules that must all hold or a set of tasks of which at least one must be accomplished. Due to their Boolean semantics, i.e., did you reach the accepting state or not, DFAs offer well-defined semantics under Boolean compositions. Our second key insight is to encode conjunctions 2 of DFAs (called cDFAs). This is done by using a graph attention network [9] (GATv2) to encode the individual DFA graph structure as well as the conjunction (AND) relationships between a collection of tasks. Recalling that the conjunction of any two DFAs grows (worst-case) quadratically, we observe that cDFAs offer an exponential reduction in the size of temporal tasks passed to GATv2.1.1 Contributions. Our main contributions are: (1) we propose compositional DFAs (cDFAs), which balance formal semantics, accessibility, and expressivity, as a goal representation for temporal tasks in goal-conditioned RL; (2) we propose encoding cDFAs using graph attention networks (GATv2); (3) we propose pre-training the cDFA encoder on reach-avoid derived (RAD) DFAs, and (4) we perform experiments demonstrating strong zero-shot generalization of pre-training on RAD DFAs.1.2 Related Work. Our work explores a temporal variant of goal-conditioned RL [33] where goals are represented as automata. While traditionally focused on goals as future target states, this has since been extended to tasks such as natural language [7, 21, 26, 37] and temporal logic [38] . While not as expressive as natural language, we believe cDFAs offer a balance between the ergonomics of language while maintaining unambiguous semantics-something of increasing importance due to the seemingly inevitable proliferation of AI agents to safety-critical systems. Moreover, DFA inference has a rich literature [6, 15, 17, 24, 27, 28, 45] with recent works have even shown the ability to learn DFA and cDFA from natural language and expert demonstrations-bridging the gap even more between natural language and automata specified goals [41, 42, 43] .Operating in a similar space, previous work, LTL2Action [38] , has shown success with using (finite) linear temporal logic (LTL), a modal extension of propositional logic, to condition RL policies on. In fact, the pre-training and test domains used in this paper are directly derived from that work.Our choice to focus on DFA rather than LTL is two-fold. First and foremost, over finite traces, LTL is strictly less expressive than DFA. For example, LTL cannot express tasks such as "the number of times the light switch is toggled should be even." Second, like DFA, LTL tasks constitute a countably infinite set of tasks. This again means that any distribution over LTL is necessarily biased to certain subclasses. On the one hand, the syntactic structure makes separation of subclasses very easy. On the other, it remains unclear how to generalize to "common" LTL formula. By contrast, we argue that the local reach-avoid structure of DFAs offers a direct mechanism for generalization. Finally, we note that while LTL is exponentially more succinct than DFA, this is largely mitigated by supporting boolean combinations of DFAs (cDFAs).Moving away from goal-conditioned RL, recent works have proposed performing symbolic planning on the DFA and cleverly stringing together policies to realize proposed paths [16, 22, 23, 30] . However, their method can still suffer from sub-optimality due to their goal-conditioned value functions being myopic. Specifically, if there are multiple ways to reach the next sub-goal of a temporal task, the optimality of the next action depends on the entire plan, not just the next sub-goaldestroying the compositional structure used for planning. For example, Figure 4 shows a simple example in which hierarchical approaches will find a suboptimal solution due to their myopia. In contrast, conditioning on cDFA embeddings allows our policies to account for the entire task.On the single-task LTL-constrained policy optimization side, several works have tackled the problem in varying settings [10, 13, 36, 44] . Adjacent to these efforts, various approaches have explored LTL specifications and automaton-like models as objectives in RL [2, 3, 8, 12, 14, 20, 29, 32, 46, 48] . A different line of work considers leveraging quantitative semantics of specifications to shape the reward [1, 22, 25, 36] . However, all of these lines of work are limited to learning a single, fixed goal.Finally, in our previous work [47] , we used hindsight experience replay [5] to solve the reward sparsity problem for DFA-conditioned off-policy learning of DFA task classes. We observe that our RAD pre-training pipeline has similar sample efficiency while generalizing to larger classes of DFAs.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: In the introduction, we clearly list our contributions, and in the experiments section, we list the research questions we explore. We then relate these research questions and their answers to the listed contributions while discussing the experiment results. 
preliminaries: 
 In the next sections, we will develop a formalism for conditioning agent behavior on temporal tasks represented as DFAs. To facilitate this exposition, we quickly review goal-conditioned RL and DFA.2.1 Goal-Conditioned RL. We start with technical machinery for goal-conditioned RL. Definition 2.1. A Markov Decision Process 3 (MDP) is a tuple M = (S, A, T ), where S is a set of states, A is a set of actions, and T is a map determining the conditional transition probabilities, i.e. Pr(s ′ | s, a) def = T (s ′ , a, s) for s, s ′ ∈ S and a ∈ A. We assume that M contains a sink state, $, denoting the end of an episode. Any MDP, M, can be given γ-temporal discounting semantics by asserting that each transition can end the episode with probability 1 -γ. Formally, T ′ (s ′ , a, s) = (1 -γ) • T (s ′ , a, s), for s, s ′ ̸ = $. Definition 2.2. Let G be a set of states called goals. A goal-conditioned policy4 , π : S ×G → ∆A, is a map from state-goal pairs to distributions over actions. Together, an MDP M and a goal-conditioned policy, π, define a distribution of sequences of states, called paths.A goal-conditioned reinforcement learning problem is a tuple consisting of a (possibly unknown) MDP and a distribution over goals, Pr(g) for g ∈ G. The objective of goal-conditioned RL is to maximize the probability of generating a path containing g.2 Deterministic Finite Automata. Next, we provide a brief refresher on automata. Definition 2.3. A Deterministic Finite Automaton (DFA) is a tuple A = (Q, Σ, δ, q 0 , F ), where Q is a finite set of states, Σ is a finite alphabet, δ : Q × Σ → Q is the transition function, q 0 ∈ Q is the initial state, and F ⊆ Q are the accepting states. The transition function can be lifted to strings via repeated application, δ * : Q × Σ * → Q. One says A accepts x if δ * (q 0 , x) ∈ F . Finally for simplicity, we abuse notation and denote by DFA the set of all DFA.Visually, we represent DFAs as multi-graphs whose nodes represent states and whose edges are annotated with a symbol that triggers the corresponding transition. Omitted transitions are assumed to stutter, i.e., transition back to the same state. DFAs have a number of interesting properties.Up to a state isomorphism, all DFAs can be minimized to a canonical form, e.g., using Hopcroft's algorithm [18] , denoted by minimize(A). Further, DFAs are closed under Boolean combinations. For example, the conjunction of two DFAs, A 1 and A 2 , denoted A 1 ∧ A 2 , is constructed by taking the product of the states, Q = Q 1 × Q 2 , and evaluating the transition functions element-wise, δ((q 1 , q 2 ), σ) = (δ 1 (q 1 , σ), δ 2 (q 2 , σ)). Also, (q 1 , q 2 ) ∈ F if q 1 ∈ F 1 and q 2 ∈ F 2 . While DFAs are closed under conjunction, each pair-wise operation results in a O(n 2 ) increase in the number of states. Definition 2.4. A compositional DFA (cDFA) is a finite set of DFA, C = {A 1 , . . . , A n }. The monolithic DFA associated to cDFA, C, is given by monolithic(C) def = A∈C A. The semantics of a cDFA are inherited from its monolithic DFA, i.e., a string x is accepted by cDFA C if and only if it is accepted by monolithic(C). Graphically, cDFA are represented as trees of depth 1 where the root corresponds to conjunction and each leaf is itself a DFA.Observe that DFAs can be viewed as deterministic MDPs leading to a natural cascading composition we refer to as DFA-augmention of a MDP. Informally, one imagines that the MDP and DFA transitions are interleaved. First, the MDP transitions given an action. The new state is then mapped to the DFA's alphabet. This symbol then transitions the DFA. The goal is to reach an accepting state in the DFA before the MDP reaches its end-of-episode state.Formally, let A and M be a DFA and MDP, respectively and let L : S → Σ denote a labelling function mapping states to symbols in the DFAs alphabet. The new MDP is constructed as a quotient of the cascading composition of A and M. First, take the product of the state spaces, S × Q. The action space is M's action set. The transition of the MDPs state is as before and the DFA transitions are given by: δ(q, L(s)). In order for there to be a unique goal and end-of-episode state, we quotient S × Q as follows: (i) ($, q) are treated as a single end of episode state and (ii) for all accepting states, q ∈ F , the product states (s, q) are treated as a single accepting (goal) state.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: Throughout the paper and explicitly in the conclusion, we clearly state the limitations of the work.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • 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 (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • 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 fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
compositional automata-conditioned reinforcement learning: 
 We are now ready to formally define (compositional) automata-conditioned RL as a variant of goal-conditioned RL where goals are defined in the augmented MDPs. Definition 3.1. A DFA-conditioned policy, π : S × DFA → ∆A is a mapping from state-DFA pairs to action distributions. A DFA-conditioned RL problem is a tuple (M, P ), where M is an MDP and P is a distribution over DFAs. The objective is to maximize the probability of generating a path containing the accepting state. cDFA-conditioned RL is defined via the underlying monolithic DFA. q 0 q 2 q 1 q 0 q 2 q 1 AND g t Figure 2 : Message passing illustration on the featurization of the cDFA given in Figure 1 .The proposed architecture for conditioning RL policies on temporal tasks encoded as cDFA is given in Figure 1 . At step t, given an observation o t and a cDFA C t , we compute an embedding s t of the observation o t using a convolutional neural network and embedding g t of C t (constituting a goal) using a message-passing neural network (MPNN). We then concatenate s t and g t and feed it to a policy network to get an action a t . After taking a t , we use the next observation o t+1 to update the initial state of each A t ∈ C t based on the transition taken. We then minimize each updated DFA to form the next cDFA C t+1 , which is used to get the next action a t+1 .3.1 cDFA Featurization. In order to pass a cDFA C t to an MPNN, we construct a featurization ĜCt = ( VCt , ÊCt , h Ct ). Details of this process are given in Appendix A. Essentially, we apply four sequential modifications to the graph structure trivially induced by C t : (i) add new nodes for every transition in each DFA in the composition, remove these transitions to connect the new nodes to the source and target nodes of the removed transitions, and in the feature vectors of the new nodes, encode all symbols triggering that transition using positional one-hot encodings, (ii) reverse all the edges, (iii) add self-loops to all nodes, and (iv) add edges from the nodes corresponding to the initial states of DFAs to the "AND" node. Figure 2 shows the featurization of the cDFA given in Figure 1 .3.2 cDFA Embedding. Given a cDFA featurization ĜCt = ( VCt , ÊCt , h Ct ), we construct an embedding of the cDFA using a graph attention network (GATv2) [9] performing a sequence of message passing steps to map each node to a vector. For ease of notation, we refer to the features of a node v ∈ VCt as h v def = h Ct (v). At each message passing step, node features are updated as:h ′ v = u∈N (v) α vu W tgt h u ,where N (v) denotes neighbors of v, W tgt is a linear map, and α vu is the attention score between v and u computed as:α vu = softmax i (e vu ) e vu = a • LeakyReLU (W src h v + W tgt h u ) ,where a is a vector and W src is a linear map. After message passing steps5 , the feature vector of the cDFA's "AND" node (which has received messages from all A t ∈ C t ) represents the latent space encoding g t of the temporal task given by C t .4 Pre-training on Reach-Avoid Derived (RAD) Compositional AutomataIn this section, we introduce pre-training on reach-avoid derived (RAD) cDFAs. This pre-training is designed to solve three important problems. First, we wish to avoid simultaneously learning a control policy and embeddings for the cDFAs. Second, we wish to learn a domain-independent cDFA encoder that performs well across cDFA distributions. Third, we wish to have either a domain-specific or a general cDFA-conditioned policy that performs well across different cDFA distributions.To alleviate the first problem, i.e., decoupling of learning of control and cDFA embeddings, we follow [38] and pre-train our cDFA encoding GNN on a "dummy" MDP containing a single nonend-of-episode state. To solve this dummy MDP, we connect a single linear layer to the output of the GNN which maps cDFA embeddings to the symbols in its alphabet. We then take these symbols and take a transition (either a stuttering one or one that changes the state of the DFA) in each DFA of the composition. The episode ends either when all DFAs end up in their corresponding accepting states (reward 1) or when one of the DFAs in the composition ends up in a non-accepting sink state (reward -1). The cDFA-augmentation then only models the dynamics of the cDFA being conditioned on.q 0 q 1 q 2 q 3 q 0 q 3 q 1 q 2 q 4 SRA DFA RAD DFAThe resulting cDFA-conditioned RL problem can thus focus on only learning to encode the cDFAs.4.1 Sequential reach-avoid as the basis for planning. What remains then is to support generalization to other cDFA distributions. In particular, recall that cDFAs (and DFAs) lack a canonical "unbiased" distribution making it a priori difficult to pre-select a cDFA distribution. To begin, observe that realizing any individual state transition q → q ′ through a DFA corresponds to a reach-avoid task: (i) eventually transition to state q ′ and (ii) avoid symbols that transition to a state different than q and q ′ . Further, a path through a DFA, q 1 , . . . , q n , corresponds to a sequence of reach-avoid (SRA) tasks. Notably, as illustrated in Figure 3 , SRA tasks can be represented as DFA. Thus, in order to generally reason about satisfying a DFA, our neural encoding must be able to solve SRA tasks.4.2 Reach-avoid derived DFAs. For general DFAs, there can be arbitrarily many (often infinitely many) paths leading to an accepting state, resulting in multiple SRA tasks interacting and interleaving. Thus, the overall structure of a DFA might be much richer than an SRA task even though it consists of local reach-avoid problems. Building on the observation that paths of a DFA correspond to SRA tasks, we define reach-avoid derived DFAs as a generalization of SRA. Definition 4.1. We define a mutation of a DFA A to be the process of randomly changing a transition of A, removing the outgoing edges of the accepting state, and then minimizing. To change a transition, we uniformly sample a state-symbol-state triplet (q, σ, q ′ ) and define δ(q, σ) def = q ′ . A DFA A is said to be m-reach-avoid derived (m-RAD) if it can be constructed by mutating an SRA DFA m-times.We define the (m, p)-RAD distribution over DFAs as follows: (i) sample a SRA DFA, A, with k +2 states where k ∼ Geometric(p); each transition stutters with probability 0.9 and is otherwise made a reach or avoid transition with equal probability; (ii) sequentially apply m mutations where mutations leading to trivial one-state DFAs are rejected. To avoid notational clutter, we shall often suppress m, p and say RAD DFAs. Finally, we define the RAD cDFA distribution by sampling n RAD DFAs, where n is also drawn from a geometric distribution. See Appendix B for the pseudocode. Figure 3 demonstrates how mutations and minimization can turn an SRA DFA into a RAD DFA. To get the RAD DFA, we apply two mutations to the SRA DFA, adding two new transitions: (i) from q 0 to q 3 on green and (ii) from q 4 to q 1 on blue, and then the minimization collapses q 0 and q 4 to a single state.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA] Justification: The paper does not include any theorems or proofs. But, we still give a rigorous formalism mapping cDFA-conditioned RL to the well-known goal-conditioned RL problem.Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • 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. 
experimental result reproducibility: 
 Question: Does the paper fully disclose all the information needed to reproduce the experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: We provide all the information needed to reproduce the main experimental results. In the corresponding section, we define all the RL environments, all the task classes, the pretraining procedure, the training procedure, and the algorithms used in the experiments. We also provide all the hyperparameter values in the appendix. 
experiments: 
 We explore the following 6 questions on several task classes in discrete and continuous environments: RQ1 Do our policies overcome the limitations of hierarchical approaches? RQ2 Does pre-training of the cDFA encoder on RAD cDFAs accelerate cDFA-conditioned RL? RQ3 Does freezing the cDFA encoder negatively impact cDFA-conditioned RL performance? RQ4 Do the cDFA-embeddings pre-trained on RAD generalize to other task classes? RQ5 How does the RAD pre-trained cDFA encoder represent the corresponding cDFA? RQ6 Do RAD pre-trained cDFA-conditioned policies generalize across task classes? Figure 4 : An example in which the myopia of hierarchical approaches causes them to find a suboptimal solution. If the task is to first go to orange and then green, the hierarchical approaches will choose the closest orange which takes them further from green whereas our approach finds the optimal solution.Letterworld Environment. Introduced in [4, 38] , Letterworld is a 7x7 grid where the agent occupies one square at a time. Some squares have randomly assigned letters which cDFA tasks are defined over. The agent can move in the four cardinal directions to adjacent squares. Moving into a square with a letter satisfies the corresponding symbol. The agent observes its position and the letters in each grid. Each layout has 12 unique letters (each appearing twice) and a horizon of 75.Zones environment. Introduced in [38] (building on Safety Gym [31] ), Zones is as an extension to Letterworld with continuous state and action spaces. There are 8 colored (2 of each color) circles randomly positioned in the environment. cDFA tasks are defined over colors. The agent's action space allows it to steer left and right and accelerate forward and backward while observing lidar, accelerometer, and velocimeter data. The environment has a horizon of 1000.Task classes. In addition to RAD, we consider the following task classes in our experiments.• Reach (R): Defines a total order over k ∼ Uniform(1, K) symbols, e.g., q 0 and q 1 in Figure 1 .• Reach-Avoid (RA): Defines an SRA task of length k ∼ Uniform(1, K), e.g., first DFA in Figure 1 .• Reach-Avoid with Redemption (RAR): Similar to RA, but they have transitions from their avoid state to the previous state, e.g., the sub-DFA with states q 0 , q 1 , q 3 in Figure 3 . • Parity: A sub-class of RAR where the symbol that leads to the avoid state also triggers a transition back to the previous state, e.g., the DFA on the right in Figure 1 . These tasks are not expressible in LTL due to their languages not being star-free regular [40] .Given a bound N , using these task classes, we generate their compositional counterparts by sampling n ∼ Uniform(1, N ) many DFAs to form a cDFA. As a short-hand notation, for example, we refer to RAR cDFAs with a maximum number of conjunctions 3 where each DFA has a task length of also 2 as cRAR.1.2.1.3 and refer to its associated monolithic (see Definition 2.4) class as RAR.1.2.1.3, i.e., without the "c" prefix. We use the same abbreviation convention for other task classes as well. We use cRAD for RAD cDFAs with truncated distributions and NT-cRAD for the not-truncated version.Finally, for comparisons with [38] , we include the task classes defined in that work, i.e., partially ordered (PO) and avoidance (AV) finite LTL tasks, represented as cDFAs.Pre-training procedure. We pre-train GATv2 on RAD cDFAs in the dummy MDP as described in Section 4. We set both geometric distribution parameters of the RAD cDFA sampler to 0.5 and truncate these distributions so that n ≤ 5 and 2 + k ≤ 10, i.e., we make sure that the maximum number of DFAs in a composition is 5, and the maximum number of states of each DFA is 10. As a baseline for GATv2, we also pre-train a relational graph convolutional network (RGCN) [34] .Training procedure. We experiment with training policies three ways: Training without pretraining (no pretraining), with frozen MPNNs pre-trained on RAD as described in Section 4 (pretraining (frozen)), and allowing the pre-trained MPNN to continue receiving gradient updates (pretraining).We use proximal policy optimization (PPO) [35] for all reinforcement learning setups, giving a reward of +1 when a rollout reaches the accepting state of the cDFA task, -1 when a rollout reaches RQ1 cDFA-conditioned policies do not suffer from myopia. First, we present the experiment results highlighting the myopia of hierarchical approaches. To this end, we construct a simplified variation of the Letterworld environment, given in Figure 4 , where the task is to visit the orange square before the green one. Observe that any hierarchical approach, like [23, 30] , based on high-level planning over some automaton-like structure suffers from myopia. Specifically, such approaches do not account for the overall task while planning for the intermediate goals. In this example, hierarchical approaches fail to differentiate between the two orange squares in the context of the given task. On the other hand, our approach quickly converges to the optimal policy as given in Figure 4 .RQ2,RQ3: Pretraining helps policies perform better faster. Figure 5 shows training curves for learning on RAD tasks using RGCN (dashed curves) and GATv2 (solid curves) as the MPNN module in Letterworld. Without pre-training in the dummy environment (blue curves), the policies barely learn how to accomplish the tasks, likely because of the difficulty in trying to simultaneously learn an embedding of the cDFA task and learn the dynamics of the environment required to reach certain symbols. Frozen pre-training (green curve) helps tremendously in learning the RAD task class in the Letterworld regardless of the MPNN module architecture. The policies learned are able to satisfy sampled RAD tasks on average over 95% of the time and within 5 timesteps. Non-frozen pre-training (red curve) also helps the policies learn more quickly, but with a significant difference between GATv2 and RGCN. Notably, non-frozen pre-training performs significantly worse for RGCN than its frozen counterpart, showing that freezing the MPNN module after pre-training actually helps performance. We further note that the wall clock training time was significantly faster due to not propagating gradients through the MPNN (about 20% faster).RQ4: Pretrained cDFA embeddings generalize well. We tested pre-trained GATv2 and RGCN models on various new cDFAs and their associated monolithic task distributions in the dummy DFA environment. Figure 10 in Appendix C.2 presents experiment results averaged over 10 seeds and 50 episodes per seed. The left y-axis illustrates the satisfaction likelihood while the right y-axis presents the number of steps needed to complete the task. The figure shows that GATv2 achieves approximately 90% accuracy across all tasks, with nearly 100% accuracy in most cases. Although RGCN generalizes comparably to GATv2 in most instances, it achieves around 80% accuracy for both cRA.1.5.1.5 and cRA.1.1.1.10 whereas GATv2 performs at nearly 90% and 100%, respectively.Considering that the maximum number of DFAs in a composition seen during training is 5 (sampled with low probability) and the DFA with the maximum task length is 9 (with a maximum of 10 states, sampled with even lower probability), these generalization results are quite remarkable. GATv2 can generalize to 10 DFAs in composition as well as DFAs with a task length of 10 with nearly 100% accuracy. Furthermore, Figure 10 shows that the embeddings produced by GATv2 lead to shorter episodes compared to RGCN, indicating that GATv2 learns better representations for cDFAs. Overall, regardless of the MPNN model used, pre-training on the RAD cDFAs provides generalization. The results are given in Figure 6 . We first observe that modifications like 2-conjunction collapse and 1step advance result in similar embeddings, showcasing the robustness provided by pre-training GATv2 on RAD. Moreover, due to the 1-step advance operation, some cDFAs end up in their accepting states, and we see that accepting cDFAs are well-isolated inside the green region in Figure 6a . Similarly, we observe that parity samples are mapped very close to RAR samples (its parent class).The same figure also demonstrates that GATv2 successfully separates different cDFA classes into distinct clusters. This separation is further confirmed by the cosine similarity (c.f. Figure 6b ) and Euclidean distance heat maps (c.f. Figure 6c ). On the other hand, both heat maps also show that the RAD cDFA class is very rich in terms of various DFAs it contains since it has both similar and different samples compared to other task classes. Notably, we see that the vast majority of RAD cDFA we sample are out-of-distribution for the other classes implying robust generalization.RQ6: Policies Trained on RAD generalize well. Figure 7 shows that RAD pre-training allows policies trained in Letterworld to zero-shot generalize far out of distribution. If you take a RAD pre-trained cDFA encoder and train a policy on some small subclass (cRA.1.1.1.5 and cRAR.1.2.1.2), the policies can still handle much longer tasks within their class and even tasks outside of their class. In the case of cRA.1.1.1.5, the training of the policy module never gets to see a multistage task, only multiple conjunctions, yet it's able to satisfy tasks from the full RAD distribution 90% of the time.We also compare against policies trained using LTL syntax embeddings [38] . [38] ). Figure 8 shows how well these policies and our RAD pre-trained policy generalize in the Letterworld to harder version of PO and AV. The results show that our method generalizes well across all tasks. Furthermore, in some cases, it outperforms the LTL policy on the task it is specifically trained on, e.g., even though the LTL policy shown by the green bar is trained on cAV.1.3.1.2, our method outperforms it. Similarly, observe that although our method was not specifically trained on cPO.1.5.1.4 tasks, it achieves 100% accuracy just like the LTL policy that is specifically trained on this task class. In addition to all of the earlier generalization results, this figure shows that RAD pre-training generalizes to these out of distribution tasks just as well as the policies from [38] even though it is not specifically trained on these tasks. 
conclusion: 
 We introduced compositional deterministic finite automata (cDFA) as an expressive and robust goal representation for temporal tasks in goal-conditioned reinforcement learning (RL). Addressing the limitations of hierarchical approaches and the challenge of extending goal-conditioned RL to non-Markovian tasks, we proposed using cDFAs to leverage formal semantics, accessibility, and expressivity. By employing graph attention networks (GATv2) to construct embeddings for cDFA tasks and conditioning the RL policy on these representations, our method encourages the policy to take the entire temporal task into account during decision-making. Additionally, we introduced reach-avoid derived (RAD) cDFAs to enable the learning of rich task embeddings and policies that generalize to unseen task classes. Our experiments show that pre-training on RAD cDFAs provides strong zero-shot generalization and robustness. Using cDFAs to encode temporal tasks enhances explainability and enables automated reasoning for interpretable and safe AI systems.6.1 Limitations. A limitation of this work is the assumption of a labeling function mapping MDP states to the task alphabet. Also, further research is needed to understand how cDFA-encoders learn to represent tasks and how much further their generalization capabilities can be pushed. 2 shows the featurization of the cDFA in Figure 1 . 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: We discuss the societal impacts of our work in the context of explainable and interpretable AI systems.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • 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 as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • 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 a 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 clearly credit and cite all the original papers that produced the assets we used as baseline, benchmark, etc.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 a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 are encouraged to reach out to the asset's creators. • 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. 
: 
 Figure 1 : Given a (conjunctive) composition of deterministic finite automata (shown on the left), we construct its embedding using a graph neural network (GATv2) and use this embedding as a goal to condition the reinforcement learning policy. 
 
 Appendix: Algorithm 1 shows the pseudocode for RAD cDFA sampling.Require: Geometric distribution parameters p n and p k and maximum number of mutations m1: n ∼ Geometric(p n ) 2: C = {} 3: for i ← 1 to n do 4: k ∼ Geometric(p k ) 5:A ← Sample a reach-avoid DFA with k + 2 states by sampling a reach and an avoid symbol uniformly from the alphabet for each state and adding these random transitions; for other symbols, add a stuttering transition with probability 0.9, or, make it reach or avoid with equal probability 6:for i ← 1 to m do 7:A ′ ← A 8:Mutate A ′ by sampling a state pair (s, s ′ ) and a symbol a and adding s a -→ s ′ 9:Make A ′ 's accepting state a sink by removing its out transitions and minimize A ′ C.3 Embedding space analysis: GATv2 vs RGCN. Figure 11 presents the embedding space analysis results for both GATv2 and RGCN. Results show that RGCN demonstrates competence in effectively separating different classes, albeit with certain limitations. While it succeeds in distinguishing reach from other classes, it struggles to differentiate RA tasks from parity and RAR. Even though parity and RAR cDFAs share structural similarities with RA cDFAs, this lack of separation is concerning as there is a semantic difference between these task classes-RA cDFAs have rejecting states but parity and RAR cDFAs do not. However, it still successfully clusters cDFAs that are 2-conjunction collapsed and 1-step advanced while also isolating accepting ones in the green region shown in Figure 11d , illustrating the robustness provided by pre-training on RAD cDFAs.Another point to note is the scales of the Euclidean distance heat maps given in Figure 11c and Figure 11f . The distance scale for RGCN is from 0 to 1.5 whereas the scale for GATv2 is from 0 to 3. Moreover, we see darker colors in Figure 11c , suggesting that GATv2 captures semantic nuances by effectively leveraging the embedding space for robust representations of cDFAs. The analysis results suggest that GATv2 learns a richer representation of the cDFA embedding space, likely due to its attention mechanism, which enables it to attend to relevant features and relationships within the data. C.4 Generalization of policies: pre-trained vs not pre-trained. Figure 12 compares the generalization strengths of policies with not-pre-trained, pre-trained-and-frozen, and pre-trained GATv2 cDFA encoders. It shows a similar pattern to the learning curves given in Figure 5 . Specifically, we see that the frozen one outperforms others while the pre-trained one performs comparably. However, we see that not-pre-trained one does not come close to the others, indicating that learning embedding spaces through pre-training is essential for good performance and generalization to different task classes. Figure 13 presents the number of steps taken in Letterworld to solve the task for the same experiment. We see a similar trend where without pre-training model needs to take a lot of steps to end the episode (either successfully or unsuccessfully). Both pre-trained models take significantly fewer steps while the frozen one solves the task fastest. Combining both figures, we see that the frozen version generalizes well while also learning to solve tasks faster. We first present the training curves for RA.1.1.1.5 and RAR.1.2.1.2 using RAD pre-trained cDFA encoders. Figure 14 shows that policies trained on RAD and RAR.1.2.1.2 converge relatively faster than the one trained on RA.1.1.1.5. Notice that RA.1.1.1.5 samples cDFAs with more than 3 DFAs with high probability (since the number of conjunctions sampled uniformly at random from the given interval) whereas the probability of sampling such big compositions is lower in RAD (as it sampled from a geometric distribution). Thus, the figure indicates that it is relatively harder to learn big conjunctions.In Figure 7 , we shared the satisfaction generalization results for these policies, testing how well these policies trained on specialized task distributions generalize to other task distributions. Here, we provide the number of steps taken to finish the episode during the generalization experiments in Figure 15 . Observe that training the policy on RAD cDFAs results in shorter episodes. However, it is still worth noting that other policies perform comparably and provide good generalization on unseen tasks. C.6 Generalization of pre-trained policies: GATv2 vs RGCN. Figure 16 present the results. It compares policies with frozen and pre-trained GNNs (both GATv2 and RGCN). In line with the training curves, we can see that GATv2 outperforms RGCN across the board, with the most significant differences on harder task classes with many symbols or many conjunctions. Figure 17 shows the number of steps taken in the same experiments, indicating that GATv2 also learns to complete the tasks with fewer steps compared to RGCN. The results highlight the superior performance of GATv2. C.7 Generalization of pre-trained policies in continuous domains. Figure 18 shows the generalization results for a policy (with a frozen and pre-trained GATv2) trained on RAD DFAs in Zones. The results show that our method can generalize to cDFAs with up to 10 DFAs in the composition as long as each DFA's task length is relatively short. We see that it does not perform as well in tasks with length 5 and more. This is probably due to the fact that control and navigation are much harder in continuous domains compared to discrete ones. Therefore, the low accuracies we are seeing for tasks with more length might be due to timeouts. It is worth noting that the policies can generalize up to 10 conjunctions in a cDFA, which is completely out of distribution for the policy. Overall, the results show that training on RAD DFAs provide zero-shot generalization in continuous domains as well. Figure 19 shows the number of steps taken to complete the episodes in the same experiment. We see that when the policy satisfies a task class with a high probability, it can also complete it with fewer steps. On the other hand, in almost all cases where it could not satisfy the task with high probability, it also took more steps in those episodes, suggesting that the policy might have timed out. C.8 Generalization of pre-trained policies: RAD cDFAs vs LTL2Action. To compare our approach to the closest work in the literature, we design an experiment with LTL2Action [38] . Recall that one of our main contributions is our RAD pre-training procedure that provides generalization across different tasks. Defining such a general LTL task class is not a trivial problem and is a contribution on this own. Moreover, the intuitions used to develop the RAD cDFA class are not applicable to LTL as it is not as easy to generate interesting and satisfiable LTL formulas. Therefore, for LTL2Action, we define a joint task distribution consisting of partially ordered and avoidance tasks, which are task classes defined in [38] . To sample LTL tasks, we flip a coin and pick a samplers. We then take the pre-trained RGCN and use it to train a policy that is exclusively trained on one of these task classes, i.e., partially ordered or avoidance LTL tasks. Figure 21 presents the learning curves for this training. The figure shows that if we let the RGCN get gradient updates during training, then pre-training does not help as much as freezing the network. We saw a similar behavior for RGCN in Figure 5 , which further justifies our choice of using a network with an attention mechanism. Furthermore, in [38] , authors show that pre-training on specific task classes helps training even when we let the RGCN continue getting gradient updates. Combining these results from [38] with what we are seeing on Figure 21 suggests that when pre-trained on rich classes with multiple distinct task patterns, RGCN need to forget about the other task classes it learned during pre-training and has to focus its capacity on the specific task class it is being trained on. We do not see such a gap in the learning curves when we let GATv2 get gradient updates, see Figure 5 , suggesting that GATv2 does not suffer from the same issue as RGCN, and it learns more robust cDFA embeddings. Both for the RGCN and GATv2 cDFA encoders, we use a hidden dimension size of 32 and perform 8 message passing steps. Also, for GATv2, we use 4 heads for the multi-head attention.For both Letterworld and Zones, we implemented the same actor and critic architectures. The actor network comprises three fully connected layers, each with 64 units and ReLU activations, while the critic network has three fully connected layers with 64, 64, and 1 units, using a Tanh activation. For pretraining, we simplified the actor and critic to single-layer models without hidden units, allowing the cDFA encoder to learn a rich embedding space. For environments with discrete actions, the actor's output was passed through a log softmax and a logit layer. In continuous action settings, we used a Gaussian distribution, parameterizing the mean and standard deviation with two separate linear layers from the actor's output. To construct observation embeddings, for Letterworld, we employed a 3-layer convolutional network with 16, 32, and 64 channels, 2×2 kernels, and a stride of 1. In Zones, a 2-layer fully connected network with 128 units per layer and ReLU activations was used.Table 1 shows the hyperparameters used for every training run in the experiments section. Each seed in the experiments section was run as an individual Slurm job with access to 4 cores of an AMD EPYC 7763 running at 2.45GHz and access to at most 20gb of memory • 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 in 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 as long as it is clear that these goals are not attained by the paper.• 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 a 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • 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 or the standard error of the mean. Justification: All the details regarding computer resources are given in the appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or 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 as 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 or failed experiments that didn't make it into the paper).Justification: We confirm that the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).With negations and disjunctions omitted as straightforward extensions.We omit a reward function since we are in a goal-conditioned setting with varying objectives/rewards.Note that in practice, RL (and by extension goal-conditioned RL) policies often consume a state observation rather than the state itself. This is a simple extension to the above which is left out for simplicity of exposition.Note that the same weights, Wsrc and Wtgt, are used in each message-passing step. 


Decision: The paper proposes a method for goal-conditioned reinforcement learning using a collection of deterministic finite automata (cDFAs) to define temporally-extended tasks. It introduces a graph neural network to embed cDFAs into vectors and a pretraining procedure using a Reach Avoid Derived (RAD) task distribution. This pretraining improves the generalization and performance of policies by leveraging these embeddings. Experiments demonstrate that freezing the pretrained embeddings results in better performance and showcases generalisation to several task classes.

There was much discussion on the paper, in particular with the positioning of the proposed approach in relation to existing work (most notably LTL2Action). This was addressed to the satisfaction of the reviewers during the rebuttal phase. There are still some outstanding issues with the paper: more empirical evaluations comparing against GCRL-LTL and understanding the tradeoffs between this style of approach and one involving UVFAs only, would strengthen the paper, and there are still concerns about how the RAD pretrained embeddings degrade depending on the task distribution.

Despite these shortcomings, the reviewers generally agree on acceptance, and authors are encourage to incorporate additional experiments and suggestions as provided by the reviewers.

Reviews: review: paper_summary: paper_summary: This paper considers the multi-task setting where each task is a temporal logic task specified by a conjunction of deterministic finite automata (cDFA). To address the sample efficiency and generalisation problems present in this setting, they propose a method for generating good cDFA embeddings which can then be used by a universal value function approximator to learn general policies. Precisely, they propose converting a given cDFA into a graph which can then be passed to a graphical neural network (GNN) like GATv2 to create a vector embedding of the task. Similarly to prior work that demonstrated the benefit of pretraining GNN embeddings on LTL tasks (LTL2Action), this paper also proposes the *reach avoid derived* DFA (RAD-DFA) task distribution to pretrain the GNN embeddings. They then demonstrate that using frozen pretrained embeddings leads to better performance than finetuning, and no pretraining leads to abysmal performance. Finally, they also demonstrate that their pretraining procedure leads to reasonable cDFA embeddings and help with generalisation.
strength_weakness: strength_weakness: # Major

## Originality
- Overall, this work is very similar to LTL2Action (Vaezipoor et al. 2021) and does not deviate particularly far from the work presented there beyond considering DFA specified tasks instead of LTL ones, which is just a choice of task specification language (regular/DFA vs LTL/BuchiAutomata). Also the paper claims that LTL2Action is limited to "finite" LTL (line 82). It is not clear what is meant by "finite" here. If it means finite trace, then that is incorrect. LTL2Action is applicable to any LTL specification.
- In addition, relying only on a UVFA for sample efficiency and generalisation is a widely known idea that has also been explored in prior works. While the exact implementation here may be new (such as the specific way the goal/task embedding is obtained), the ideas are familiar from a function approximation front. 
- The paper introduces the cDFAs, but this does not seem novel since they don't seem meaningfully different from DFAs and the distinction does not seem relevant for the proposed approach.  Hence this just adds more notations and terminologies. Also, the footnote on page 2 says "With negations and disjunctions omitted as straightforward extensions". It is not clear how they are straightforward, unless cDFAs are indeed just DFAs as I mentioned previously.

## Quality and clarity
### Mathematical formulations
- I found it very hard to fully understand the proposed framework and judge the soundness because there are a number of incorrect statements and missing information.
- The MDP definition (Def 2.1) is wrong. An MDP is a five tuple M = (S,A,T,R,$\gamma$) where $R$ is the reward function and  $\gamma$ is the discount factor (optional when $\gamma=1$). It is not a triple M = (S,A,T). The objective of goal-conditioned RL is not to maximize the probability of generating a path containing g (lines 127-128), but instead to maximise the value function (the expected cumulative rewards). 
- It is fine if the authors want to focus on the maximisation of the success probability, but then the paper needs to describe how that will be achieved in the RL setting defined.
- The reward function and discount factor for the problem formulation are never defined. The only reward function mentioned in the mathematical formulations is the one used for pretraining the GNN. Since the experiments uses the same reward scheme for learning policies, I suggest the authors formally state that in their problem formulation. Additionally, Table 1 shows that the experiments use a discounting, meaning that the learned policies *do not* correspond to the maximisation of success probability.
- The paper says that goals are defined in the augmented MDPs (line 164), but the DFA-conditioned policy uses the environment states and not the augmented ones (Def 3.1).   
 - The paper uses message-passing neural network (MPNN) (line 173) as a core factor in their approach, but never define nor describe what these are.
- The paper proposes cDFAs as a novel goal-representation for temporal tasks in goal-conditioned RL. Beyond the fact that I am not convinced that they are meaningfully different from DFAs, the paper gives no theory nor conducts any experiment to demonstrate that they "balance formal semantics, accessibility, and expressivity" as claimed in their first contribution. In general, since the authors claim it as a main contribution of this work, they should justify why it is better than other representations like RMs, LTL, GTL, TLTL, SPECTRL, etc. I think such rigorous justifications are not needed if they did not claim it as a contribution of this work.   

### Experiments
- It is not clear why the main paper contains no experiment comparing with prior works. The only such experiment is left to the appendix, and only uses LTL2Action as baseline which is not state-of-the-art (there is also no sample-efficiency comparison). Without state-of-the-art baselines like [1] or [2] in the main paper, it is hard to tell how good the proposed approach is. Given that [2] drastically outperforms LTL2Action in their experiments in the same domains, this work does not look promising as it only shows similar performance to LTL2Action. 
- The paper claims that the proposed approach is applicable to any DFA, but the DFAs used for pretraining the goal encoder (Algorithm 1) and the ones used for the experiments (page 7) only have one atomic proposition per edge (e.g. $red$, $square$, etc). These are the simplest DFA transitions possible. In general the edge of a DFA is a Boolean expression over the atomic propositions (e.g. $red \wedge \neg square$). It is unclear if the approach generalises to more complex DFAs (where the edges/transitions can be arbitrary Boolean expressions).  
- All the example DFAs used are generally not satisfiable in the environments considered since their states lack self-transitions (e.g. in Fig 1). Since all the task distributions used for the experiments are defined similarly, it is unclear why the results show high success rates on them (they should be generally unsolvable). This suggests that important details of the implementation of the approach for the experiments may be different from what has been described in the paper.
- It is also unclear what the tasks sampled from the defined distributions look like. It would have helped if the authors included examples of the simplest and most complex DFAs sampled during training and during testing. 
- It is unclear why freezing the pretrained weights performs better than finetuning. This is counter-intuitive but is never explained in the paper.
- Given how the performance of the proposed approach is heavily reliant on the goal embeddings obtained from pretraining (the performance is abysmal without it), it is concerning that the paper did not stress test it to understand when it fails and why. For example, what if the number of conjunctions goes up to 100 where each DFA has a task length that also goes up to 100 (relates to the temporal curse of dimensionality)? What if the DFA transitions are more complex, that is when they are arbitrary Boolean expressions over the atomic propositions (relates to the spatial curse of dimensionality)?

## Significance
- This work does not make a very large step forward from prior work (the addition of state-aware planning), but it is one. So I do not want to be overly pessimistic on the front of significance, but it is necessary to note that the contributions of this work are relatively incremental. If I understood the paper correctly, the only main contribution is the approach for pretraining a *good* goal (DFA) embedding. Additionally Figures 20-21 also only shows similar performance to non-state-of-the-art prior work (LTL2Action). 
- The absence of theory does not help. Theory guaranteeing generalisation and convergence to optimal policies (as claimed) under some reasonable assumptions would have helped. The lack of experiments comparing with state-of-the-art baselines like [1,2] also does not help.
- Relying only on a UVFA for generalisation has clear limitations (e.g. in terms of the spatial and temporal curses of dimensionality present in temporal logic tasks), since no optimality and generalisation guarantees can be given even in finite MDPs (like grid-worlds). This is why the field has largely moved away from such approaches. More recent works such as [1,2] (and the large majority of prior works like [3,4]) break down the problem into a set of reach-avoid sub-tasks (goals), then use planning (or HRL) to determine how to sequence those goals (addressing the temporal curse). These works use UVFA for the sub-goal policies only when the goal space is sufficiently large. When the goal space isn't too large, such as in the letter grid-world, these methods can still learn and generalise without relying on a UVFA (by learning the policies/options for each goal independently).
- This reliance on end-to-end function approximation for generalisation hurts the significance of this work.


[1] Tasse, Geraud Nangue, et al. "Skill Machines: Temporal Logic Skill Composition in Reinforcement Learning." The Twelfth International Conference on Learning Representations. (2024)

[2] Qiu, Wenjie, Wensen Mao, and He Zhu. "Instructing goal-conditioned reinforcement learning agents with temporal logic objectives." Advances in Neural Information Processing Systems 36 (2024).

[3] Araki, Brandon, et al. "The logical options framework." International Conference on Machine Learning. PMLR, 2021.

[4] Jason Xinyu Liu, Ankit Shah, Eric Rosen, George Konidaris, and Stefanie Tellex. Skill transfer for temporally-extended task specifications. arXiv preprint arXiv:2206.05096, 2022.

# Minor

- It is not clear what "concept class" means or how it is different from a "task class". Are they the same?
- It is not clear what GATv2 in the caption of Figure 1 means. Please add in the caption that it is a GNN.
- It is seems like "task" and "goal" are mean the same thing in this work (line 16), i.e. they mean a DFA. Is that correct? That wasn't clear. 
- Adding the cDFA in Fig 1 to Fig 2 will
- There are a couple minor typos or writing mistakes. For example:
 - There are numerous places where shorthands like "i.e." and "e.g." are used (e.g. first paragraph of the introduction). It is not recommended in a formal writing. Please spell them out.
 - The sentence on Line 184 doesn't end.
questions: questions: Please refer to the weaknesses above.
limitations: limitations: The authors have discussed some of the limitations. However, there seems to be several other ones. For example,

- Just like *some* prior works, the paper assumes that all tasks are defined over a fixed set of given atomic propositions. If the set increases, everything needs to be retrained.
- The paper demonstrates their approach only on DFAs with one atomic proposition per edge. It is unclear how it will behave when the DFA transitions are more complex (i.e. arbitrary Boolean expressions over the atomic propositions)?
- The paper relies on UVFAs to learn general policies. Hence, it is likely that it can only generalise to tasks *similar* to those it is trained on, and not to tasks that are *significantly* out of distribution. The paper does not investigate when the model fails to generalise and why. That would have helped the reader have some justified intuition of what types of tasks are *significantly* out of distribution for function approximation to fail for proposed approach.


I recommend the authors carefully think of the various significant limitations of their approach, such as the one I mentioned above, and explicitly state and discuss them in the paper.
paper_summary: paper_summary: This work focuses on goal-conditioning policies using DFAs. This leverages the ability of DFAs to be composed - in this case the work focuses on conjunction. Two main pieces are introduced in leverging DFAs: 1) a GATv2 model which provides task embeddings from the DFA, 2) pre-training of the GATv2 on Reach Avoid Derived compositional automata which trains the embeddings with a basis representation in terms of the RAD embeddings which can then be used to represent other DFA embeddings. Once the  DFA embeddings is obtained it is appended to the environment state space and used to train a policy. Experiments demonstrate the benefit of GATv2 over an alternative RGCN (when pre-training and not freezing the embeddings) and also demonstrate the benefit of pre-training the task embeddings, particularly when freezing the embeddings after.
strength_weakness: strength_weakness: # Clarity
Figure captions could elaborate more. As it stands they are vague and do not add any explanation or support understanding of the figure. Secondly, Section 3.1 is unclear due to the omission of an explanation for the DFA modifications. As a result it is unclear why any of these steps are needed, but this seems to be an important and early part of the pipeline. Thus, it leaves the entire pipeline feeling unclear or at least on an uneasy grounding.

# Significance
Unfortunately, I think the significance of this work may be its biggest weakness. Firstly, it is not clear to me that it is trivial to extend the findings on conjunctions will generalize to general boolean combinations as is mentioned in lines 144 and 145. Secondly, the curse of dimensionality which results from conjunctions of DFAs is never addressed or touched on. Perhaps the vector representation of task embedding does support generalization - and the visualizations would even support this - but it is not discussed. This is even more worrying when considering that a fixed dimension for the task embeddings is used. Thirdly, the experiments do support the answering of the research questions posed. However, the experiments do not compare to any other baselines from the literature, which makes it difficult to contextualize and build from this work. Finally, and I acknowledge that the authors point this out in the limitations section and so I will take this point lightly, but there is no insight into how the task embeddings support generalization. This does unfortunately limit the significance of the work.
questions: questions: 1. How would the results of the generalized agents with task embeddings compare to specialized agents?
2. In Figure 4, why do the Zones results only have two models?
3. At a high level, how would you extend the approach to incorporate general boolean combinations

If the authors could address my concerns with regard to the curse of dimensionality, and answer these questions I would likely advocate for acceptance.
limitations: limitations: Limitations are mentioned and I do not see any omissions.
paper_summary: paper_summary: This paper extends the framework of goal-conditioned reinforcement learning to support temporally-extended goals specified as (compositions of) deterministic finite automata (cDFA), avoiding the limitations of state-based goals and the ambiguities of natural language specifications. The authors introduce a graph neural network architecture for embedding cDFAs to vectors via message passing. They then develop a pretraining distribution for learning cDFA embeddings, which is constructed by leveraging the insight that paths in a DFA correspond to a series of reach-avoid tasks, motivating a distribution of cDFAs that are derived from sequential reach-avoid DFAs. In experiments, the authors show that pretraining on this distribution produces generalizable embeddings that capture the structure of DFA space, while also enabling the training of generalizable cDFA-conditioned policies. The authors run these experiments on two labeled MDPs, Letterworld and Zones, showing that embeddings or policies trained on only RAD DFAs or other DFA subsets often achieve high goal satisfaction rates even when presented with cDFAs with longer task lengths and higher composition numbers. cDFA conditioned policies using a frozen cDFA embedder trained on the RAD distribution also lead to high satisfaction rates across the board in Letterworld, with reasonable strong performance in Zones (a harder continuous MDP).
strength_weakness: strength_weakness: As the authors themselves note, probably the main weakness / limitation of work that relies on symbolic specifications like cDFAs or LTL formulae is that it requires MDPs with a labeling function. This limitation isn't insurmountable --- for e.g., it should be possible to use computer vision algorithms to segment and label states of the world from pixels -- but in the final version of the paper, I think it'd be good to discuss this a bit more, since it might help readers appreciate how work like this could be applicable without a oracle labeling function.

Another slight weakness of the paper is that most of the experiments focus on Letterworld, with weaker average performance in the more realistic environment of  Zones. This raises questions about how applicable the method will be to more realistic tasks. With the extra page in the final version, I think it'd be good to show images of both environments, and examples of harder cDFAs in those environments, just to illustrate how difficult it might be to solve those tasks. Otherwise it might be easy to write-off e.g. Letterworld as a toy domain even though with cDFAs the tasks can be quite long-horizon. It might also be good to say a bit bout how performance could be improved for harder tasks like Zones.

Finally, I think one thing that bears discussing in the final version is how this framework might enable safety/correctness of the policy, not just interpretability of the specification. While it's no doubt a safety benefit to have more interpretable specifications that say exactly what correct behavior is, this on its own doesn't guarantee the resulting learned policies are safe. I wonder if the authors could say a bit more on this point, perhaps by connecting to work on formal verification of NNs, or by combining cDFA-conditioned RL with constrained planning methods, using the learned policy to accelerate provably safe planning.
questions: questions: There were a couple of minor things I didn't fully grasp, which it'd be great to clarify:

- When constructing/sampling a SRA DFA (e.g. RHS of Figure 3) how are reach or avoid transitions created? Is the reach transition always a transition to the accepting state? Or is it to a a state that eventually leads to an accepting state? Is it the case that there is a single long path to the accepting state? And is the avoid transition always to a single sink state? I couldn't fully understand this from Appendix B.

- In Figure 4, in the panel for Zones, where is the (non-frozen) GATv2 pretraining line baseline? And am I right to say that the RGCN baseline was not trained for Zones?

- For each MDP with its own labeling function and alphabet, do you basically to do the RAD pretraining phase for specifically that alphabet of symbols? Or is there hope of training a cDFA embedder that generalizes across different MDPs with different labeling functions and alphabets?
limitations: limitations: The authors have adequately discussed a number of limitations of their work. It would be good to also discuss the fact that interpretable specifications (in the form of cDFAs) do not immediately lead to verifiably safe behavior.
paper_summary: paper_summary: This paper tackles Goal-conditioned Reinforcement Learning (RL) when used with Temporal Logic Objectives. The benefits of directly considering Deterministic Finite Automatons (DFAs) as the task definition leads to a new class of Compositional DFAs introduced that cover a conjunction of several tasks. These cDFAs can be encoded using a Graphical Neural Network that is pretrained on a novel task distribution (RAD). Experimental results show the approach generalizes to several task classes.
strength_weakness: strength_weakness: - The approach follows the similar structure as LTL2Action [41] but replaces the LTL Module with a different GNN (GATv2) and considers cDFAs or compositions/conjunctions of DFAs. The modifications w.r.t. LTL2Action apart from this and the pre-training step over a new randomly sampled (RAD) task distribution are novel yet arguably incremental.
- Baselines appear lacking which bring into question the merits of the approach. A direct comparison to LTL2Action [41] and GCRL-LTL [32] using the monolithic DFA of the cDFA in the experiment could be possible and if not, the reason should be justified.
- The paper is at times hard to follow and overloaded with notations and conventions. Some improvements could be made in the description of RAD pre-training (see questions) and the task nomenclature 1.X.1.X in the experiments.
questions: questions: 1. Should the statement in L35 be better quantified or validated with references? "DFAs represent all temporal tasks that become Markovian by augmenting the state-space
with finite memory”. Does this hold when Signal Temporal Logic (STL) is considered a temporal task language as well? 
2. Is Reference [7] a concurrent submission or a typo?
3. In Figure 5, what are the black elements in the Legend? Why are they not present in the figures?
4. What do the envelope colors in Figure 2 indicate? Does the encoding generalize to conjunctions of the predicates on edges (e.g., Red & Green)? Are these considered during pretraining?
5. Defn. 4.1. could use some clarity. How is a random transition defined in a mutation? Is this simply changing the target node of a given transition to a random state in the DFA? What is being “minimized" here?
6. Minor grammatical error - Should DFA and DFA(s) be used interchangeably? I would think the following references to DFA should be changed to DFAs (L31, L33, L80, L85-87, and more)
limitations: limitations: The limits of cDFA encoding are not fully explored (mentioned).

Comments: title: Comment by Reviewer A6oy
Thank you for the response.
Comment by Reviewer Ww4e
comment: Thank you to the authors for their detailed response and clarifications. I particularly appreciate the comments regarding the originality and significance of the proposed approach, and the clarifications regarding DFAs vs cDFAs. The intuition behind why frozen embeddings are giving the best results is also interesting. I think a detailed discussion of this will improve the quality of the paper, given that it is the second main contribution of the paper.

I believe the RAD pretraining approach is interesting and as I mentioned in my original review, it could be a useful contribution to the field. But I am not fully convinced by the provided emperical evidence. I have the following main outstanding concerns regarding this paper: 

- GCRL+LTL/DFA and end-to-end function approximation (UVFA-only) methods do have their tradeoffs. As the authors have highlighted, GCRL based methods are often very sample-efficient but often suffer from suboptimality which requires planning while taking into account the environment states (e.g. using Dijkstra). Similarly, UVFA-only methods have the potential to learn optimal policies, but lack any optimality and generalisation guarantees due to the reliance on function approximation. Hence, it is important to understand emperically or theoretically when a proposed UVFA-only approach is likely to suffer from sub-optimality. My main concern here is that there's not enough investigation done on the tradeoffs induced by the RAD pretraining approach, and the reliance on UVFAs for generalisation over cDFAs. Such investigations are particularly important in RL research where the statistical significance of results is often very low (for example this paper shows results over only 10 seeds, which is common, but likely statistical insignificant).

- Regarding DFAs with Boolean expressions on edges, I meant Boolean expressions over atomic propositions (AP). I.e. when your set of symbols is $2^{AP}$, as is commonly done. The experiments only have edges over AP, which is a severely restricted class of DFAs.

- I see what you mean for the conjunction stress test. But you can still define such tasks. An example in the letter world is the conjunction of all 12*11 unique SRA tasks: Reach Letter$\_i$ while avoiding Letters$\_{j \neq i}$ . You can also have a sequencial stress test. E.g. SRA task of length k ∼ Uniform(1,100). 

-  As I mentioned previously, it will help to include examples of the simplest and most complex DFAs from the defined distributions.

I have increased my score to reflect the clarifications provided and my outstanding concerns.
_We thank the reviewer for their response and for arguing that this is a strong paper._

We will certainly move the comparison to LTL2Action into the main text in the final version of the paper to make sure that it is clear to readers.
_We thank the reviewer for their response and for raising their score._

We will make sure to include all of the suggested presentation refinements into the final version and believe it will significantly help the clarity of the paper. Space and time permitting, we plan on running empirical comparisons against the related approaches (GCRL-LTL) on the types of problems described in the meta-rebuttal to backup the qualitative argument.
I thank the authors for the clarifications. The main contributions are clear (RAD pretraining/embeddings) and I understand how DFAs can encode many tasks captured by different temporal logics. I will raise my score accordingly since I beilieve the method presents a useful technique to generalize over logical tasks.

Rather than a qualitative argument as to why other related approaches (GCRL-LTL) are different, I would have been interested to see simulations on a sample task/specification. Further, I believe the paper could use some presentation refinements (as suggested) and look forward to seeing the final version.
Thank you for the answers to my questions. I continue to think this is a strong paper, and will be maintaining my score. Having read the other reviews, I also think it would be valuable to make the comparison with LTL2Action more explicit and move it into the main text.
_We thank the reviewer for their response and for raising their score._

We are glad to hear that you agree that the comparison to LTL2Action is the most appropriate baseline and that the additional clarity in Section 3.1 addresses your concerns. We agree that further investigation into the task embedding space to understand its scalability (in more complex environments such as Zones and with larger tasks) is an exciting direction of future research. We will make sure to address these limitations and future directions clearly in the final version of the paper.

It is true that the fixed representation size does provide a theoretical limitation on its representation capacity as the size of the task space increases. We have two reasons to be hopeful that this will not pose a major limitation in practice:
1. This exponential blowup is a theoretical limitation for any sequence modeling task and yet, fixed size embedding have proven to be successful in computer vision (e.g., CLIP) and natural language (e.g., word2vec) domains.
2. DFAs (compared to unstructured representations such as natural language) enable a natural finite approximation of the embedding space: _planning over the DFA up to a limited depth_. A fixed size embedding space might naturally converge on this approximation.
I thank the authors for their thorough response.

* For the additions to improve clarity of Section 3.1. These additions are what I had in mind and will address my concern.
* For the concern on the curse of dimensionality: I note that the cDFA means the exponential blowup in states does not explicitly occur. However, the vector representation of the task space must still be able to disambiguate the composed tasks from all other tasks. This is my main concern, that the fixed vector size will either need to be excessively large to begin with or will rapidly become exhausted. This is an empirical question of how quickly this vector space representation will become difficult to use due to the exponential increase in task space from composition. The cDFA point above does not seem to address this, unless I am missing something.
* A comparison to LTL2Action seems useful and to my knowledge is the most appropriate baseline. This would correct the original omission in my opinion.
* I agree that investigating the task embedding space would support future work. However, it is unfortunately related to my above point that it is unclear how quickly the empirical results of this work would decay without such an understanding. To me this isn't a difference between acceptance or rejection but does limit the significance of the work at the upper end. Particularly, because many benchmarks which do not use task embeddings are guaranteed to generalise arbitrarily. The task embedding having a theoretical limit in its representation is a weakness of the proposed method which is very difficult to assess - I admit this - but is still a limitation none the less. However, I acknowledge that this is pointed out by the authors and so I emphasise that I am not using this to decide on acceptance.
* I think this concession on the zones experiment needs to be stated explicitly. This could also be of importance in assessing the scaling capacity of the proposed method. However, once again I can accept these concerns being left to future work as it is stated in the limitations section.
* I acknowledge that I have read the reviewers comments in the general rebuttals.

Overall I think that if the changes to clarity are made as proposed then I would be comfortable seeing the work accepted. Unless I am missing something, the score is limited due to the difficulty of assessing the scalability of the model which is a weakness unique to using task embeddings compared to approaches which do not such as those cited by Reviewer A6oy. This needs to be stated too. However, in acknowledgement of the proposed changes and some helpful clarification I will raise my score to a 6. However, I think the discussion with the other reviewers appear useful and given their in-depth responses I would defer the outcome to be weighted towards their reviews. Thus, I will be lowering my confidence.