Compositional Reinforcement Learning from Logical Specifications

We study the problem of learning control policies for complex tasks given by logical specifications . Recent approaches automatically generate a reward function from a given specification and use a suitable reinforcement learning algorithm to learn a policy that maximizes the expected reward . These approaches , however , scale poorly to complex tasks that require high-level planning . In this work , we develop a compositional learning approach , called DIRL , that interleaves highlevel planning and reinforcement learning . First , DIRL encodes the specification as an abstract graph ; intuitively , vertices and edges of the graph correspond to regions of the state space and simpler sub-tasks , respectively . Our approach then incorporates reinforcement learning to learn neural network policies for each edge ( sub-task ) within a Dijkstra-style planning algorithm to compute a high-level plan in the graph . An evaluation of the proposed approach on a set of challenging control benchmarks with continuous state and action spaces demonstrates that it outperforms state-of-the-art baselines . 1 Introduction . Reinforcement learning ( RL ) is a promising approach to automatically learning control policies for continuous control tasks—e.g. , for challenging tasks such as walking [ 11 ] and grasping [ 6 ] , control of multi-agent systems [ 31 , 22 ] , and control from visual inputs [ 28 ] . A key challenge facing RL is the difficulty in specifying the goal . Typically , RL algorithms require the user to provide a reward function that encodes the desired task . However , for complex , long-horizon tasks , providing a suitable reward function can be a daunting task , requiring the user to manually compose rewards for individual subtasks . Poor reward functions can make it hard for the RL algorithm to achieve the goal ; e.g. , it can result in reward hacking [ 3 ] , where the agent learns to optimize rewards without achieving the goal . Recent work has proposed a number of high-level languages for specifying RL tasks [ 5 , 29 , 24 , 34 , 19 ] . A key feature of these approaches is that they enable the user to specify tasks compositionally—i.e. , the user can independently specify a set of short-term subgoals , and then ask the robot to perform a complex task that involves achieving some of these subgoals . Existing approaches for learning from high-level specifications typically generate a reward function , which is then used by an off-the-shelf RL algorithm to learn a policy . Recent works based on Reward Machines [ 19 , 35 ] have proposed RL algorithms that exploit the structure of the specification to improve learning . However , these algorithms are based on model-free RL at both the high- and low-levels instead of model-based RL . Model-free RL has been shown to outperform model-based approaches on low-level control tasks [ 10 ] ; however , at the high-level , it is unable to exploit the large amount of available structure . Thus , these approaches scale poorly to long horizon tasks involving complex decision making . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . We propose DIRL , a novel compositional RL algorithm that leverages the structure in the specification to decompose the policy synthesis problem into a high-level planning problem and a set of low-level control problems . Then , it interleaves model-based high-level planning with model-free RL to compute a policy that tries to maximize the probability of satisfying the specification . In more detail , our algorithm begins by converting the user-provided specification into an abstract graph whose edges encode the subtasks , and whose vertices encode regions of the state space where each subtask is considered achieved . Then , it uses a Djikstra-style forward graph search algorithm to compute a sequence of subtasks for achieving the specification , aiming to maximize the success probability . Rather than compute a policy to achieve each subtask beforehand , it constructs them on-the-fly for a subtask as soon as Djikstra ’ s algorithm requires the cost of that subtask . We empirically evaluate1 our approach on a “ rooms environment ” ( with continuous state and action spaces ) , where a 2D agent must navigate a set of rooms to achieve the specification , as well as a challenging “ fetch environment ” where the goal is to use a robot arm to manipulate a block to achieve the specification . We demonstrate that DIRL significantly outperforms state-of-the-art deep RL algorithms for learning policies from specifications , such as SPECTRL , TLTL , QRM and HRM , as well as a state-of-the-art hierarchical RL algorithm , R-AVI , that uses state abstractions , as the complexity of the specification increases . In particular , by exploiting the structure of the specification to decouple high-level planning and low-level control , the sample complexity of DIRL scales roughly linearly in the size of the specification , whereas the baselines quickly degrade in performance . Our results demonstrate that DIRL is capable of learning to perform complex tasks in challenging continuous control environments . In summary , our contributions are as follows : • A novel compositional algorithm to learn policies for continuous domains from complex high-level specifications that interleaves high-level model-based planning with low-level RL . • A theoretical analysis of our algorithm showing that it aims to maximize a lower bound on the satisfaction probability of the specification . • An empirical evaluation demonstrating that our algorithm outperforms several state-of-theart algorithms for learning from high-level specifications . Motivating example . Consider an RL-agent in the environment of interconnected rooms in Figure 1 . The agent is initially in the blue box , and their goal is to navigate to either the top-left room S1 or the bottom-right room S2 , followed by the top-right room S3 , all the while avoiding the red block O . This goal is formally captured by the SPECTRL specification φex ( middle top ) . This specification is comprised of four simpler RL subtasks—namely , navigating between the corner rooms while avoiding the obstacle . Our approach , DIRL , leverages this structure to improve learning . First , based on the specification alone , it constructs the abstract graph Gex ( see middle bottom ) whose vertices represent the initial region and the three subgoal regions , and the edges correspond to subtasks ( labeled with a safety constraint that must be satisfied ) . However , Gex by itself is insufficient to determine the optimal path—e.g. , it does not know that there is no path leading directly from S2 to S3 , which is a property of the environment . These differences 1Our implementation is available at https : //github.com/keyshor/dirl . can be represented as ( a priori unknown ) edge costs in Gex . At a high level , DIRL trains a policy πe for each edge e in Gex , and sets the cost of e to be c ( e ; πe ) = − logP ( e ; πe ) , where P ( e ; πe ) is the probability that πe succeeds in achieving e. For instance , for the edge S0 → S1 , πe is trained to reach S1 from a random state in S0 while avoiding O . Then , a naı̈ve strategy for identifying the optimal path is to ( i ) train a policy πe for each edge e , ( ii ) use it to estimate the edge cost c ( e ; πe ) , and ( iii ) run Djikstra ’ s algorithm with these costs . One challenge is that πe depends on the initial states used in its training—e.g. , training πe for e = S1 → S3 requires a distribution over S1 . Using the wrong distribution can lead to poor performance due to distribution shift ; furthermore , training a policy for all edges may unnecessarily waste effort training policies for unimportant edges . To address these challenges , DIRL interweaves training policies with the execution of Djikstra ’ s algorithm , only training πe once Djikstra ’ s algorithm requires the cost of edge e. This strategy enables DIRL to scale to complex tasks ; in our example , it quickly learns a policy that satisfies the specification with high probability . These design choices are validated empirically—as shown in Figure 1 , DIRL quickly learns to achieve the specification , whereas it is beyond the reach of existing approaches . Related Work . There has been recent work on using specifications based on temporal logic for specifying RL tasks [ 2 , 7 , 12 , 18 , 30 , 17 , 40 , 15 , 39 , 23 ] . These approaches typically generate a ( usually sparse ) reward function from a given specification which is then used by an off-the-shelf RL algorithm to learn a policy . In particular , Li et al . [ 29 ] propose a variant of Linear Temporal Logic ( LTL ) called TLTL to specify tasks for robots , and then derive shaped ( continuous ) rewards from specifications in this language . Jothimurugan et al . [ 24 ] propose a specification language called SPECTRL that allows users to encode complex tasks involving sequences , disjunctions , and conjunctions of subtasks , as well as specify safety properties ; then , given a specification , they construct a finite state machine called a task monitor that is used to obtain shaped ( continuous ) rewards . Icarte et al . [ 19 ] propose an automaton based model called reward machines ( RM ) for high-level task specification and decomposition as well as an RL algorithm ( QRM ) that exploits this structure . In a later paper [ 35 ] , they propose variants of QRM including an hierarchical RL algorithm ( HRM ) to learn policies for tasks specified using RMs . Camacho et al . [ 9 ] show that one can generate RMs from temporal specifications but RMs generated this way lead to sparse rewards . Kuo et al . [ 27 ] and Vaezipoor et al . [ 36 ] propose frameworks for multitask learning using LTL specifications but such approaches require a lot of samples even for relatively simpler environments and tasks . There has also been recent work on using temporal logic specifications for multi-agent RL [ 16 , 33 ] . More broadly , there has been work on using policy sketches [ 5 ] , which are sequences of subtasks designed to achieve the goal . They show that such approaches can speed up learning for long-horizon tasks . Sun et al . [ 34 ] show that providing semantics to the subtasks ( e.g. , encode rewards that describe when the subtask has been achieved ) can further speed up learning . There has also been recent interest in combining high-level planning with reinforcement learning [ 1 , 25 , 13 ] . These approaches all target MDPs with reward functions , whereas we target MDPs with logical task specifications . Furthermore , in our setting , the high-level structure is derived from the given specification , whereas in existing approaches it is manually provided . Illanes et al . [ 20 ] propose an RL algorithm for reachability tasks that uses high-level planning to guide low-level RL ; however , unlike our approach , they assume that a high-level model is given and high-level planning is not guided by the learned low-level policies . Finally , there has been recent work on applying formal reasoning for extracting interpretable policies [ 37 , 38 , 21 ] as well as for safe reinforcement learning [ 4 , 26 ] . 2 Problem Formulation . MDP . We consider a Markov decision process ( MDP ) M = ( S , A , P , η ) with continuous states S ⊆ Rn , continuous actions A ⊆ Rm , transitions P ( s , a , s′ ) = p ( s′ | s , a ) ∈ R≥0 ( i.e. , the probability density of transitioning from state s to state s′ upon taking action a ) , and initial states η : S → R≥0 ( i.e. , η ( s ) is the probability density of the initial state being s ) . A trajectory ζ ∈ Z is either an infinite sequence ζ = s0 a0−→ s1 a1−→ · · · or a finite sequence ζ = s0 a0−→ · · · at−1−−−→ st where si ∈ S and ai ∈ A . A subtrajectory of ζ is a subsequence ζ ` : k = s ` a ` −→ · · · ak−1−−−→ sk . We let Zf denote the set of finite trajectories . A ( deterministic ) policy π : Zf → A maps a finite trajectory to a fixed action . Given π , we can sample a trajectory by sampling an initial state s0 ∼ η ( · ) , and then iteratively taking the action ai = π ( ζ0 : i ) and sampling a next state si+1 ∼ p ( · | si , ai ) . Specification language . We consider the specification language SPECTRL for specifying reinforcement learning tasks [ 24 ] . A specification φ in this language is a logical formula over trajectories that indicates whether a given trajectory ζ successfully accomplishes the desired task . As described below , it can be interpreted as a function φ : Z → B , where B = { true , false } . Formally , a specification is defined over a set of atomic predicates P0 , where every p ∈ P0 is associated with a function JpK : S → B ; we say a state s satisfies p ( denoted s |= p ) if and only if JpK ( s ) = true . For example , given a state s ∈ S , the atomic predicate Jreach sK ( s′ ) = ( ‖s′−s‖ < 1 ) indicates whether the system is in a state close to s with respect to the norm ‖ · ‖ . The set of predicates P consists of conjunctions and disjunctions of atomic predicates . The syntax of a predicate b ∈ P is given by the grammar b : := p | ( b1 ∧ b2 ) | ( b1 ∨ b2 ) , where p ∈ P0 . Similar to atomic predicates , each predicate b ∈ P corresponds to a function JbK : S → B defined naturally over Boolean logic . Finally , the syntax of SPECTRL specifications is given by 2 φ : := achieve b | φ1 ensuring b | φ1 ; φ2 | φ1 or φ2 , where b ∈ P . In this case , each specification φ corresponds to a function JφK : Z → B , and we say ζ ∈ Z satisfies φ ( denoted ζ |= φ ) if and only if JφK ( ζ ) = true . Letting ζ be a finite trajectory of length t , this function is defined by ζ |= achieve b if ∃ i ≤ t , si |= b ζ |= φ ensuring b if ζ |= φ and ∀ i ≤ t , si |= b ζ |= φ1 ; φ2 if ∃ i < t , ζ0 : i |= φ1 and ζi+1 : t |= φ2 ζ |= φ1 or φ2 if ζ |= φ1 or ζ |= φ2 . Intuitively , the first clause means that the trajectory should eventually reach a state that satisfies the predicate b . The second clause says that the trajectory should satisfy specification φ while always staying in states that satisfy b . The third clause says that the trajectory should sequentially satisfy φ1 followed by φ2 . The fourth clause means that the trajectory should satisfy either φ1 or φ2 . An infinite trajectory ζ satisfies φ if there is a t such that the prefix ζ0 : t satisfies φ . We assume that we are able to evaluate JpK ( s ) for any atomic predicate p and any state s. This is a common assumption in the literature on learning from specifications , and is necessary to interpret a given specification φ . Learning from Specifications . Given an MDPM with unknown transitions and a specification φ , our goal is to compute a policy π∗ : Zf → A such that π∗ ∈ arg maxπ Prζ∼Dπ [ ζ |= φ ] , where Dπ is the distribution over infinite trajectories generated by π . In other words , we want to learn a policy π∗ that maximizes the probability that a generated trajectory ζ satisfies the specification φ . We consider the reinforcement learning setting in which we do not know the probabilities P but instead only have access to a simulator of M. Typically , we can only sample trajectories of M starting at an initial state s0 ∼ η . Some parts of our algorithm are based on an assumption that we can sample trajectories starting at any state that has been observed before . For example , if taking action a0 in s0 leads to a state s1 , we can store s1 and obtain future samples starting at s1 . Assumption 2.1 . We can sample p ( · | s , a ) for any previously observed state s and any action a .

Summary: This paper studies the problem of using logical specifications of a task to derive an abstract graph, come up with a plan, and train low-level policies more effectively. Given a logical specification written in SPECTRL syntax, it is first converted into an abstract MDP. Dijkstra's algorithm is executed on this abstract MDP, assigning probabilities to edges when they are encountered for the first time. A policy is trained for each edge (S_0, S_1) in the abstract MDP where a random state in S_0 is the initial state and a random state in S_1 is the goal state. The success probability of the policy is used to assign the probability of the corresponding edge. The model is evaluated on two rooms environments and a fetch environment from OpenAI gym. The model is shown to outperform previous models.