Value Function Spaces: Skill-Centric State Abstractions for Long-Horizon Reasoning

1 INTRODUCTION . For an agent to perform complex tasks in realistic environments , it must be able to effectively reason over long horizons , and to parse high-dimensional observations to infer the contents of a scene and its affordances . It can do so by constructing a compact representation that is robust to distractors and suitable for planning and control . Consider , for instance , a robot rearranging objects on a desk . To succcessfully solve the task , the robot must learn to sequence a series of simple skills , such as picking and placing objects and opening drawers , and interpret its observations to determine which skills are most appropriate . This requires the ability to understand the capabilities of these simpler skills , as well as the ability to plan to execute them in the correct order . Hierarchical reinforcement learning ( HRL ) aims to enable this by leveraging abstraction , which simplifies the higher-level control or planning problem . Typically , this is taken to mean abstraction of actions in the form of primitive skills ( e.g. , options ( Sutton et al. , 1999 ) ) . However , significantly simplifying the problem for the higher level requires abstraction of both states and actions . This is particularly important with rich sensory observations , where standard options frameworks provide a greatly abstracted state space , but do not simplify the perception or estimation problem . The nature of the ideal state abstraction in HRL is closely tied to the action abstraction , as the most suitable abstraction of state should depend on the kinds of decisions that the higher-level policy needs to make , which in turn depends on the actions ( skills ) available to it . This presents a challenge in designing HRL methods , because it is difficult to devise a state abstraction that is both highly abstracted ( and therefore removes many distractors ) and still sufficient to make decisions for long-horizon tasks . This challenge differs markedly from representation learning problems in other domains , like computer vision and unsupervised learning ( Schroff et al. , 2015 ; van den Oord et al. , 2018 ; Chen et al. , 2020 ) , since it is intimately tied to the capabilities exposed to the agent via its skills . We therefore posit that a suitable representation for a higher-level policy in HRL should depend on the capabilities of the skills available to it . If this representation is sufficient to determine the abilities and outcomes of these skills , then the high-level policy can select the skill most likely to perform the desired task . This concept is closely tied to the notion of affordances , which has long been studied in cognitive science and psychology as an action-centric representation of state ( Gibson , 1977 ) , and has inspired techniques in robotics and RL ( Zech et al. , 2017 ; Xu et al. , 2021 ; Mandikal & Grauman , Open Drawer Pick Blue Place Blue in Drawer Close Drawer = = = = = t0 t3t2t1 t4 ( a ) A trajectory through the skill value function space for the task “ Place blue block in drawer ” . VFS , visualized on top of corresponding scene , captures positional information about the contents of the scene , preconditions for interactions , and the effects of executing a feasible skill , making it suitable for high-level planning . Open Drawer Pick Blue Place Blue in Drawer Close Drawer = = = = = ( b ) VFS learns a skill-centric representation of the scene and ignores factors like arm pose , task-centric object positions , color of the desk , or additional objects , which do not affect the values of the low-level skills . All configurations shown above are functionally equivalent and hence , map to the same VFS representation . Figure 1 : Visualizing VFS embeddings in an example desk rearrangement task . VFS can capture the affordances of the low-level skills while ignoring exogenous distractors . 2021 ) . Given a set of skills that span the possible interactions in an environment , we propose that the value functions corresponding to these policies can provide a representation that suitably captures the capabilities of the skills in the current state and thus can be used to form a compact embedding space for high-level planning . We call this embedding space Value Function Spaces ( VFS ) . Note that we intend to use these skills as high-level actions without modifications – the important problems of skill discovery and test-time adaptation are beyond the scope of our work . Figure 1a illustrates the state abstraction constructed by VFS for the desk rearrangement example discussed above : VFS captures the affordances of the skills and represents the state of the environment , along with preconditions for the low-level skills , forming a functional representation to plan over . Since VFS constructs a skill-centric representation of states using the value functions of the low-level skills , it captures functional equivalence of states in terms of their affordances : in Figure 1b , states with varying object or arm positions ( i-iii ) , different background textures ( iv ) , and distractor objects ( v ) are functionally equivalent for planning , and map to the same embedding in our representation . This significantly simplifies the high-level planning problem , making it easier for the higher level policy to generalize to novel environments to the limit of the skills themselves . Statement of Contributions . The primary contribution of this work is VFS , a novel state representation derived from the value functions of low-level skills available to the agent . We show that VFS leverages the properties of value functions , notably their ability to model possible skills and completed skills , to form an effective representation for high-level planning , and is compatible with both model-free and model-based high-level policies . Empirical evaluations in maze-solving and robotic manipulation demonstrate that the skill-centric representation constructed by VFS outperforms representations learned using contrastive and information-theoretic objectives in long-horizon tasks . We also show that VFS can generalize to novel environments in a zero-shot manner . 2 RELATED WORK . Hierarchical RL has been studied extensively in the literature , commonly interpreted as a temporal abstraction of the original MDP . Early works have interpreted the hierarchy introduced in this setting as an abstraction of state and action spaces ( Sutton et al. , 1999 ; Dietterich , 2000 ; Thomas & Barto , 2012 ) . The popular options framework Sutton et al . ( 1999 ) ; Precup ( 2000 ) provides a natural way of incorporating temporally extended actions into RL systems . An agent that possesses the transition model and reward model for such a SMDP ( known as an option model ) is capable of sample-based planning in discrete ( Kocsis & Szepesvári , 2006 ) and continuous state spaces ( Konidaris et al. , 2014 ; Gopalan et al. , 2017 ) . However , doing so in environments with high-dimensional observations ( such as images ) is challenging . In this work , we explore the efficacy of learned skills operating on highdimensional observations for long-horizon control in realistic environments . To improve the quality of lower-level policies , recent work in HRL has studied various facets of the problem , including discovery of skills ( Konidaris & Barto , 2009 ; Zhang et al. , 2021b ; Florensa et al. , 2017 ; Warde-Farley et al. , 2018 ) , end-to-end training of both levels ( Kulkarni et al. , 2016 ; Tessler et al. , 2017 ) and integrating goal-conditioned behaviors ( Ghosh et al. , 2019 ; Nachum et al. , 2019 ) . In this work , we assume that the low-level skills are given , and do not focus on discovering them . Instead , we focus on how skills can simplify the higher-level control problem by providing a representation that is suitable for either model-based or model-free control . The high-level policy does not rely on any hand-crafted rewards and only receives a sparse outcome reward . The problem of learning meaningful long-horizon behavior by combining a set of lower-level skills , or options , has been studied extensively in prior work . Popular approaches to composing skills involve combining them in the space of value functions ( Todorov , 2009 ; Hunt et al. , 2019 ; da Silva et al. , 2009 ; Ziebart , 2010 ; Haarnoja et al. , 2018 ) , and more recently , using pseudo-rewards ( Barreto et al. , 2019 ) . While these methods have been very effective at composing low-level skills , they do not address the question of learning a meaningful state abstraction using these skills . We take inspiration from these works and use the value function space to construct a skill-centric state abstraction , along with a high-level policy that composes these skills for temporally extended tasks . Our method can hence be interpreted as a representation learning approach . Representation learning techniques have been employed extensively in model-free RL by augmenting auxiliary tasks based on reconstruction losses ( Lange et al. , 2012 ; Higgins et al. , 2017 ; Yarats et al. , 2019 ) and predicting the future conditioned on past observations ( Schmidhuber , 1990 ; Jaderberg et al. , 2017 ; van den Oord et al. , 2018 ; Shelhamer et al. , 2016 ) . Contrastive learning has also been used in recent works to discover a meaningful latent space and extract reward signals for RL ( Sermanet et al. , 2017 ; Warde-Farley et al. , 2018 ; Dwibedi et al. , 2018 ; Laskin et al. , 2020 ) . Another popular approach to learning useful representation is to optimize for effective transfer of the lower-level skills and the learned state abstraction ( Gupta et al. , 2017 ; Konidaris & Barto ; Goyal et al. , 2019 ) . Unlike these prior works , our aim is specifically to learn a representation that is grounded in the capabilities of the low-level skills , which gives us a skill-centric abstraction of high-dimensional observations . While we do not explicitly optimize for transfer like Gupta et al . ( 2017 ) , our experiments show that VFS can compactly represent affordances of the scene and generalize to novel environments while being robust to functionally-irrelevant distractors . Alongside developments in model-free RL , prior work has also sought to learn predictive models of the environment for sampling and planning , in a model-based RL framework . This has been demonstrated by learning dynamics using future predictions ( Watter et al. , 2015 ; Oh et al. , 2017 ; Ebert et al. , 2017 ; Banijamali et al. , 2018 ; Ha & Schmidhuber ; Hafner et al. , 2018 ; Ichter & Pavone , 2019 ; Hafner et al. , 2020 ; Zhang et al. , 2019 ) , learning belief representations ( Gregor et al. , 2019 ; Lee et al. , 2019 ) and representing state similarity using the bisimulation metric ( Castro , 2020 ; Zhang et al. , 2021a ; Agarwal et al. , 2021 ) . The combination of learned model-free policies with structures like graphs ( Savinov et al. , 2018 ; Eysenbach et al. , 2019 ) and trees ( Ichter et al. , 2021 ) to plan over extended horizons has also been demonstrated to improve generalization and exploration . Although our method can utilize a model-based high-level controller to plan over temporally extended skills , our objective is not to develop better model-based RL algorithms . Rather , we focus on developing suitable state abstractions in a hierarchical RL framework , and then evaluate these representations by employing them in both model-based and model-free higher-level controllers .

Summary: This paper presents Value Function Spaces (VFS), an abstract "skill-centric" representation for reinforcement learning and long-horizon planning. The core idea of VFSs is to leverage the learned value functions that are often trained alongside skills during skill learning to abstract knowledge about the world. Each element of this low-dimensional representation corresponds to the value estimated from one of the learned value function, so that the representation (by construction) tends to ignore distractor information unnecessary for accomplishing the skills. Using this representation on a set of pre-trained skills, the authors demonstrate improved performance on long-horizon planning tasks in grid-based maze and "locked door" environments, showing improved success probability on a number of such tasks compared to competitive baselines. They also show the ability of their abstraction to facilitate model-based goal-directed planning by learning a state-transition model (in VFS-space) that allows one to predict the outcome of executing a skill.