GrASP: Gradient-Based Affordance Selection for Planning

The ability to plan using a model learned from experience is arguably a key component of intelligence . Repeated planning can help both in the better selection of the action to execute in the current state as well as in providing better targets for updating policies and value functions . A landmark in the success of the use of planning in model-based Deep RL ( RL that uses neural-networks as function approximators ) was AlphaGo ( Silver et al. , 2017 ) in which a simulator of the game of Go was used by the planner as a model . More recently MuZero ( Schrittwieser et al. , 2020 ) has shown the capability of learning models from experience simultaneously with planning in solving an impressive number of large-scale RL problems ( albeit on problems where a simulator is available and thus sample complexity is not a big concern ) . Nevertheless , there remain interesting challenges in making planning with learned models become a routine part of Deep RL agents across the spectrum of RL domains . One such challenge , the subject of interest in this paper , is dealing effectively with continuous actions spaces where using tree-search based planning is not straightforward ( e.g. , just considering every action at the root node of the tree becomes infeasible ) . Accordingly , we consider RL domains with continuous action or option spaces ( where options are temporally extended actions ) . Our approach is to select a small number of affordances—to discover what actions/options are useful to consider in tree-based planning . The notion of affordances is a rich one in psychology that refers to many kinds of organism-environment relations ( Gibson , 1977 ) . We take inspiration in our work from the idea of directly perceived high-level actions that environment states afford , such as object interactions ( e.g. , chairs afford sitting , cups afford grasping ) or navigation actions ( doorways afford passage ) , all restricted subsets of the vast space of possible actions the organism may perform . We implement affordances as parametric mappings from states ( and goals in multi-task settings ) to actions/options . So selecting K affordances means learning K such mappings , and then there are at most K choices at each node in the planning tree . What action/option the ith affordance corresponds to depends on the state corresponding to a node . We also consider simply selecting K actions/options that are available in all states to evaluate how much conditioning on state matters . Our main contributions are twofold : ( 1 ) the idea of selecting affordances to deal with continuous actions/options in planning , and ( 2 ) GrASP an algorithm for Gradient-based Affordance Selection for Planning . GrASP computes gradients through the planning procedure to update the parameters of the functions that represent affordances . We do not claim or show that GrASP is state of the art for model-based RL in continuous domains . Instead , we show empirically that despite the nonstationarity of learning to select affordances whilst simultaneously learning a value-equivalent model that predicts their value-consequences , GrASP is fast enough at these two tasks for planning with the discovered affordances and their learned model to often be more sample efficient than model-free approaches in benchmark RL control problems . A comparison with state of the art awaits further development of this new idea of using gradients to discover affordances for planning . 1 RELATED WORK . Other approaches to affordances in RL . Abel et al . ( 2014 ) defines affordances as a subset of actions available to an agent in a state if the state meets certain preconditions defined through propositional functions . They introduce an approach for learning affordances in a scaffolded multitask learning environment , where affordances learned from simple tasks with smaller state spaces are transferred onto difficult tasks with larger state spaces . With this approach , they demonstrate the utility of affordances to reduce the planning time in a learning agent . Abel et al . ( 2015 ) extends this to the setting of goal-based priors : probability distributions over the optimality of each action for a given state and goal pair . Both works are limited in their applicability as they rely on discovering affordances in Object Oriented-MDPs , which assumes objects , object classes , and their attributes as part of the state space ( Diuk et al. , 2008 ) . It is also limited in the sense that their affordance discovery approach heavily relies on a carefully designed multi-task environment that provides a curriculum to the learning agent . Khetarpal et al . ( 2020 ) defined affordances for MDPs and in simple discrete domains showed that hand-specified affordances allow for faster planning in an agent as they reduce the number of actions available at a state and also allow for efficient learning of partial transition models of the environment . Sherstov & Stone ( 2005 ) improved the time required to learn an optimal policy by transferring a reduced action set obtained from the policy of a source task to a new but similar target task . Rosman & Ramamoorthy ( 2012 ) provide a method for learning Dirichlet action priors from a set of related tasks that were then shown to improve the planning speed . Other work ( Even-Dar et al. , 2006 ; Cruz et al. , 2014 ; 2016 ; 2018 ; Zahavy et al. , 2018 ) have also demonstrated the utility of action-pruning to allow an agent to scale up to environments with large action spaces . In summary , none of the above approaches developed an approach to the discovery of affordances based on gradients through the planning procedure . Gradient through a trajectory in a model to update policies . Another body of work , summarized below , can be seen as effectively using gradients through a trajectory generated from a learned model to update a policy . Our approach to learning affordances using gradients through a planning tree is different in at least two ways . First , we learn affordances that provide a set of actions as a function of state that an agent should consider during planning . Second , we compute gradients through a planning tree rather than a single trajectory . In our empirical work , we consider an ablation of GrASP in which we learn a single affordance mapping , in effect learning a policy via gradients through a trajectory as a means of comparing against an instance of this approach . PILCO ( Deisenroth & Rasmussen , 2011 ) is an approach that learned state transition models parametrized as Gaussian processes . The policy was derived from gradients obtained by differentiating through a trajectory from the transition model and was shown to learn optimal policies in simple control tasks such as cartpole and mountaincar . Many approaches extended the idea from PILCO ( Gal et al. , 2016 ; Higuera et al. , 2018 ; Chua et al. , 2018 ; Amos et al. , 2018 ) to other kinds of models and domains . Dreamer , introduced by Hafner et al . ( 2020 ) , used a recurrent architecture for learning latent transition models and directly learned the policy parameters from the world model by differentiating the unrolled model with the policy parameters . Byravan et al . ( 2020 ) introduced an approach that learns latent deterministic transition models and learned a policy by backpropagating through transition model expansions , similar to that of Dreamer . Continuous actions in RL . Stochastic Value Gradients ( SVG ) ( Heess et al. , 2015 ) allow learning of continuous control policies by obtaining value gradients of one-step model predictions and these value gradients have been shown to have reduced variance , thus producing successful learning algorithms such as Deterministic Policy Gradients ( Silver et al. , 2014 ) , Deep Deterministic Policy Gradients ( Casas , 2017 ) and Soft Actor-Critic ( Haarnoja et al. , 2018 ) . These methods use gradients of one-step trajectories to update policies and thus are quite different from GrASP . Planning with continuous action spaces . Progressive widening ( Couëtoux et al. , 2011 ) is an early approach for planning in continuous-action domains when a simulator is available . It was recently extended to planning with value-equivalent models by Moerland et al . ( 2018 ) and Yang et al . ( 2020 ) , but both of these have only been shown to work on domains with 1-D actions . MuZero ( Schrittwieser et al. , 2020 ) is considered to be a state-of-the-art algorithm for planning with value-equivalent models in RL but was limited to discrete-action domains . More recently an extension of MuZero , called Sampled MuZero ( Hubert et al. , 2021 ) , was introduced for dealing with continuous actions . Sampled MuZero is different from GrASP in two important respects . Sampled MuZero uses the planning process to update a policy prior distribution over continuous action space and uses samples from this distribution to construct the planning tree , while GrASP uses gradients through the planning process to discover a discrete set of affordance mappings which are used for building the tree . Like in MuZero , our GrASP agents learn value-equivalent models rather than observation prediction models . But unlike MuZero , our GrASP agents also learn option-models in addition to primitive-action models . ( We could not directly compare our method against MuZero because no code was released with the paper , and our attempts at a local implementation efficient enough to be usable did not succeed ) . 2 GRASP DETAILS . Key idea . Our RL agent is a planning agent that uses an affordance module that maps state representations to a small set of K actions or options from a continuous space for use in expanding a look-ahead search tree . The affordance module is represented by a network with parameters θafford and so the planner is , in effect , parameterized by θafford . The key idea underlying our method of discovering useful affordances for planning is that we can compute the gradient of performance loss with respect to θafford through the computations of the planner . The planners we explore here use the affordance module and a learned value-equivalent model to expand a lookahead tree . In principle any tree-expansion procedure could be used including various MCTS algorithms . The key requirement is that performance-loss gradients w.r.to θafford can be computed through the planner ’ s computations , then GrASP can discover affordances online . In our empirical work we will use two tree-expansion procedures : ( 1 ) shallow-depth complete trees with depth as a parameter . We use learned value functions at the leaf nodes to bootstrap value estimates at the root node and so shallow depth can be sufficient , and the ability of GrASP to discover useful affordances means that we can have very few action choices at each state and thus afford to build complete lookahead trees . ( 2 ) UCT-based MCTS with number and depth of trajectories as parameters ; here we choose affordances from the learned affordances based on an upper-confidence bonus . High level overview . Figure 1 provides an overview of how the planning agent selects actions or options to execute in the environment using the planning tree . The current observation from the environment x is encoded into an abstract state s which becomes the root node in the planning tree expanded using fafford to generate actions/options at each abstract state node that is expanded , and using the value equivalent model to generate next abstract states and predictions of rewards and values . The reward and value predictions are backed-up to produce Q-values for each action or option , which are used to construct a policy over the actions/options at the root node . An action/option is sampled from this policy and executed to produce a transition in the environment . The next observation is encoded to a new abstract-state , from which the agent repeats its planning procedure . There are thus five learned components : the novel affordance module fafford , and the four components of the value equivalent model , an observation encoder f encode , a state-next state dynamics function fdynamic , a reward predictor f rew , and a value function approximator fV . In what follows we specify the key aspects of the algorithm and how the module parameters are learned ; full details are in the Appendix . Agent-environment interaction . We assume a discrete time setting in which an agent interacts with a continuous-observation environment using actions or options selected from a continuous space . For the most part we will treat options and actions as if they are equivalent by overloading some terms such as reward and next observation . In the case of selecting actions , at each time step t the agent receives an observation xt ∈ Ro , selects an action at ∈ A = Ra and then receives reward rt+1 ∈ R and new observation xt+1 . In the case of selecting options , if option at ∈ A ( noteA is used to denote both the space of actions and options ) is selected at time t , its policy is followed until termination and the agent receives reward ( defined as the discounted sum of rewards from the start of the option until termination ) , and new observation xt+n ( if the option lasts for n time steps ) . The next option is then chosen at time t+ n. The agent ’ s goal is the usual RL goal of maximizing discounted sum of rewards over an episode ( all our tasks use a discount factor of 0.99 ) . Affordance module . The agent encodes each observation xt into an abstract state st using an encoding module describe below . The affordance module is a network with parameters θafford that maps abstract-states s to K actions or options a1 · · ·aK selected from A ; i.e. , fafford : s → { a1 , a2 · · ·aK } . In our experiments , we use a simple feedforward neural net architecture for the affordance module with K separate output heads , each head outputting a different action/option . Value-equivalent model representation . The agent plans with a value-equivalent model represented by four neural networks parameterized by θM = { θencode , θdynamic , θrew , θV } . These four modules are in common with the modules in MuZero ’ s value equivalent model ( Schrittwieser et al. , 2020 ) ; the difference is that our agent does not have a policy-prediction module and that we allow for option-transitions that are temporally extended while MuZero was restricted to action-transitions . The functions computed by each module are defined as follows : The representation or encoding module f encode : x → s , parameterized by θencode , maps an observation x to an abstract state s ∈ S = Rm . The abstract state representations are learned ; they are not latent environment states or observation predictions . The value module fV : s → R parameterized by θV estimates the value of an abstract state . The reward module f rew : s , a → r̂ ∈ R , parameterized by θrew , predicts the discounted sum of rewards until termination after executing an action or option a given state s. For options , the agent learns to predict the expected duration of an option given that it is started in encoded state s. The dynamics module fdynamic : s , a → s′ , parameterized by θdynamic maps an abstract state s and action/option a to a predicted next abstract-state s′ ∈ S. Value-equivalent model training . The parameters of the representation , dynamics and value networks are trained so that the resulting model estimates the value of the environment observation as accurately as possible ( hence the term value-equivalent ) . The reward network is trained to predict the rewards that the agent experiences from a given observation-action/option pair . Training operates by sampling a sequence of transitions from the replay buffer , as described in pseudocode in Appendix , then using these transitions to update the neural networks of the learning agent . We briefly describe here the learning update of θM ; the update of θafford is described below . Given a sequence of state-action/option-reward transitions { x1 , a1 , r2 , x2 , a2 , r3 , x3 · · ·xn } , the agent first constructs abstract-state s1 = f encode ( x1 ) , then predicts subsequent abstract states si = f dynamic ( si−1 , ai−1 ) . The model parameters θM are then updated to minimize the loss : Lmodel = n∑ i=1 ( ri − f rew ( si , ai ) ) 2 + ( v̂i − fV ( si ) ) 2 where , v̂i is the discounted sum of observed rewards bootstrapped by the predicted value function of the last abstract state in the sequence from a target value function network ( the value function network is copied into the target value function network periodically ) . Note that the model loss function is not used to adapt the parameters of the affordance mappings . Model-learning and affordance-discovery are thus independent , though they do constrain each other indirectly because model-learning leads to learning of the abstract state representations which are input to the affordance mappings and similarly the affordance mapping ’ s choice of actions/options produce the trajectories for the replay buffer that are used as data for model-learning . Lookahead tree expansion using affordances and value-equivalent model . In the case of complete trees , the planning procedure expands a K-ary search tree to depth D , expanding each abstract state node with all K actions/options computed by the affordance module for each abstract state . In the case of UCT , the planner rolls out H trajectories to depth D ; the trajectories then implicitly define a partial tree . Each trajectory starts at the root note . If there are one or more affordances that have not yet been selected at the root node , we choose one of those at random . Else , an expansion procedure selects greedily based on the sum of the current Q-values at the root node ( see below for backup procedure ) and a UCT bonus that boosts the chances of selecting less frequently chosen affordances . At a non-root node whose abstract state has previously been encountered we follow the same expansion procedure ( except there is no requirement to select all the affordances at least once ) . At a non-root node whose abstract state has not been encountered we pick an affordance at random . Action selection via planning tree . Regardless of whether we have a complete tree or a partial tree produced by UCT the following recursive procedure is applied . The value of a leaf-node is the value of the corresponding abstract state from the value network . The value of a non-leaf node is a weighted sum of the Q-values for all the actions/options taken in the partial tree from that node . Suppose the abstract state of the non-leaf node under consideration is s , the actions/options in the tree are a1 , a2 , · · · , af , the Q-value of ( s , ai ) is the sum of the reward from the reward network for ( s , ai ) summed with the appropriately-discounted value of the successor node in the tree to ( s , ai ) . The value of node s then is ∑ i π ( ai|s ) Q ( s , ai ) , where π is a probability distribution over actions selected in the tree . For complete-tree based planning π ( ai|s ) = exp ( Q ( s , ai ) /τ ) ∑ j exp ( Q ( s , aj ) /τ ) , where τ is the temperature parameter , while for the UCT-based tree the weight π ( ai|s ) = n ( s , ai ) ∑ j n ( s , aj ) , where n ( s , ai ) is the count of the number of times action ai is chosen in state s in the partial tree . Finally , the action executed at the current observation as a result of planning is selected by sampling from the distribution defined by π at the root node . Gradient update for the affordance module . The agent updates the affordance module parameters by learning to maximize the value-estimate at the root node V ( s0 ) . As described above , the value-estimate is obtained via a value backup procedure through the tree , which yields the following expression : V ( s0 ) = ∑ b∈ { a1 , a2 , ··· , aK } π ( b|s0 ) Q ( s0 , b ) . where the summation is over the action/option-selections produced by the affordance module . Note that for both the complete tree and the partial tree produced by UCT , each of the affordances is selected at least once at the root . The recursive form of the value function leads to the following recursion to compute the needed gradient : ∂V ( s0 ) ∂θafford = ∂ ∂θafford [ ∑ b∈ { a1 , a2 , ··· , aK } π ( b|s0 ) Q ( s0 , b ) ] = ∑ b∈ { a1 , a2 , ··· , aK } [ ∂π ( b|s0 ) ∂θafford Q ( s0 , b ) + π ( b|s0 ) ∂Q ( s0 , b ) ∂θafford ] Recall that π ( ai|s0 ) is a function of θafford through the action/option-selections ai that are outputs of the affordance module . Finally , ∂Q ( s0 , ai ) ∂θafford can be computed via a recursive computation very similar to the value backup procedure : ∂Q ( s0 , ai ) ∂θafford = ∂f rew ( s0 , ai ) ∂θafford + γn ∂V ( s1 ) ∂θafford , where n is the expected duration of option ai . This recursive gradient term can be efficiently computed with existing auto-differentiation packages ( Griewank & Walther , 2008 ) with minor additional computational complexity .

Summary: The paper proposes a gradient-based affordance selection, i.e., action/option selection, in addition to the value equivalent modules for tree-expansion procedure(s) in planning. The claim is that GrASP can learn primitive-action and option selection and plan in a continuous state and action space to outperform model-free RL. The affordance module maps the state to K actions or options from continuous space. These affordances are later used for expanding the look-ahead search tree. The key idea is to compute the gradient of performance loss with respect to the affordance model's parameters through the planner. The method uses two expansion procedures: (1) shallow-depth complete trees, (2) UCT-based MCTS. The performance of TD3 is compared to different variants of the proposed algorithm on some tasks from the DM Control Suite.