Learning Actionable Representations with Goal Conditioned Policies

1 INTRODUCTION . Representation learning refers to a transformation of an observation , such as a camera image or state observation , into a form that is easier to manipulate to deduce a desired output or perform a downstream task , such as prediction or control . In reinforcement learning ( RL ) in particular , effective representations are ones that enable generalizable controllers to be learned quickly for challenging and temporally extended tasks . While end-to-end representation learning with full supervision has proven effective in many scenarios , from supervised image recognition ( Krizhevsky et al. , 2012 ) to vision-based robotic control ( Levine et al. , 2015 ) , devising representation learning methods that can use unlabeled data or experience effectively remains an open problem . Much of the prior work on representation learning in RL has focused on generative approaches . Learning these models is often challenging because of the need to model the interactions of all elements of the state . We instead aim to learn functionally salient representations : representations that are not necessarily complete in capturing all factors of variation in the observation space , but rather aim to capture factors of variation that are relevant for decision making – that are actionable . How can we learn a representation that is aware of the dynamical structure of the environment ? We propose that a basic understanding of the world can be obtained from a goal-conditioned policy , a policy that can knows how to reach arbitrary goal states from a given state . Learning how to execute shortest paths between all pairs of states suggests a deep understanding of the environment dynamics , and we hypothesize that a representation incorporating the knowledge of a goal-conditioned policy can be readily used to accomplish more complex tasks . However , such a policy does not provide a readily usable state representation , and it remains to choose how an effective state representation ∗dibya.ghosh @ berkeley.edu should be extracted . We want to extract those factors of the state observation that are critical for deciding which action to take . We can do this by comparing which actions a goal-conditioned policy takes for two different goal states . Intuitively , if two goal states require different actions , then they are functionally different and vice-versa . This principle is illustrated in the diagram in Figure 1 . Based on this principle , we propose actionable representations for control ( ARC ) , representations in which Euclidean distances between states correspond to expected differences between actions taken to reach them . Such representations emphasize factors in the state that induce significant differences in the corresponding actions , and de-emphasize those features that are irrelevant for control . While learning a goal-conditioned policy to extract such a representation might itself represent a daunting task , it is worth noting that such a policy can be learned without any knowledge of downstream tasks , simply through unsupervised exploration of the environment . It is reasonable to postulate that , without active exploration , no representation learning method can possibly acquire a dynamics-aware representation , since understanding the dynamics requires experiencing transitions and interactions , rather than just observations of valid states . As we demonstrate in our experiments , representations extracted from goal-conditioned policies can be used to better learn more challenging tasks than simple goal reaching , which can not be easily contextualized by goal states . The process of learning goal-conditioned policies can also be made recursive , so that the actionable representations learned from one goalconditioned policy can be used to quickly learn a better one . Actionable representations for control are useful for a number of downstream tasks : as representations for task-specific policies , as representations for hierarchical RL , and to construct well-shaped reward functions . We show that ARCs enable these applications better than representations that are learned using unsupervised generative models , predictive models , and other prior representation learning methods . We analyze structure of the learned representation , and compare the performance of ARC with a number of prior methods on downstream tasks in simulated robotic domains such as wheeled locomotion , legged locomotion , and robotic manipulation . 2 PRELIMINARIES . Goal-conditioned reinforcement learning . In RL , the goal is to learn a policy πθ ( at|st ) that maximizes the expected return Rt = Eπθ [ ∑ t rt ] . Typically , RL learns a single task that optimizes for a particular reward function . If we instead would like to train a policy that can accomplish a variety of tasks , we might instead train a policy that is conditioned on another input – a goal . When the different tasks directly correspond to different states , this amounts to conditioning the policy π on both the current and goal state . The policy πθ ( at|st , g ) is trained to reach goals from the state space g ∼ S , by optimizing Eg∼S [ Eπθ ( a|s , g ) ( Rg ) ) ] , where Rg is a reward for reaching the goal g. Maximum entropy RL . Maximum entropy RL algorithms modify the RL objective , and instead learns a policy to maximize the reward as well as the entropy of the policy ( Haarnoja et al. , 2017 ; Todorov , 2006 ) , according to π ? = arg maxπ Eπ [ r ( s , a ) ] +H ( π ) . In contrast to standard RL , where optimal policies in fully observed environments are deterministic , the solution in maximum entropy RL is a stochastic policy , where the entropy reflects the sensitivity of the rewards to the action : when the choice of action has minimal effect on future rewards , actions are more random , and when the choice of action is critical , the actions are more deterministic . In this way , the action distributions for a maximum entropy policy carry more information about the dynamics of the task . 3 LEARNING ACTIONABLE REPRESENTATIONS . In this work , we extract a representation that can distinguish states based on actions required to reach them , which we term an actionable representation for control ( ARC ) . In order to learn state representations φ that can capture the elements of the state which are important for decision making , we first consider defining actionable distances DAct ( s1 , s2 ) between states . Actionable distances are distances between states that capture the differences between the actions required to reach the different states , thereby implicitly capturing dynamics . If actions required for reaching state s1 are very different from the actions needed for reaching state s2 , then these states are functionally different , and should have large actionable distances . This subsequently allows us to extract a feature representation ( φ ( s ) ) of state , which captures elements that are important for decision making . To formally define actionable distances , we build on the framework of goal-conditioned RL . We assume that we have already trained a maximum entropy goal-conditioned policy πθ ( a|s , g ) that can start at an arbitrary state s0 ∈ S in the environment , and reach a goal state sg ∈ S. Although this is a significant assumption , we will discuss later how this is in fact reasonable in many settings . We can extract actionable distances by examining how varying the goal state affects action distributions for goal-conditioned policies . Formally , consider two different goal states s1 and s2 . At an intermediate state s , the goal-conditioned policy induces different action distributions πθ ( a|s , s1 ) and πθ ( a|s , s2 ) to reach s1 and s2 respectively . If these distributions are similar over many intermediate states s , this suggests that these states are functionally similar , while if these distributions are different , then the states must be functionally different . This motivates a definition for actionable distances DAct as DAct ( s1 , s2 ) = Es [ DKL ( π ( a|s , s1 ) ||π ( a|s , s2 ) ) +DKL ( π ( a|s , s2 ) ||π ( a|s , s1 ) ) ] . ( 1 ) The distance consists of the expected divergence over all initial states s ( refer to Section B for how we do this practically ) . If we focus on a subset of states , the distance may not capture action differences induced elsewhere , and can miss functional differences between states . Since maximum entropy policies learn unique optimal stochastic policies , the actionable distance is well-defined and unambiguous . Furthermore , because max-ent policies capture sensitivity of the value function , goals are similar under ARC if they require the same action and they are equally “ easy ” to reach . We can use DAct to extract an actionable representation of state . To learn this representation φ ( s ) , we optimize φ such that Euclidean distance between states in representation space corresponds to actionable distances DAct between them . This optimization yields good representations of state because it emphasizes the functionally relevant elements of state , which significantly affect the actionable distance , while suppressing less functionally relevant elements of state . The problem is : min φ Es1 , s2 [ ‖φ ( s1 ) − φ ( s2 ) ‖2 −DAct ( s1 , s2 ) ] 2 ( 2 ) This objective yields representations where Euclidean distances are meaningful . This is not necessarily true in the state space or in generative representations ( Section 6.4 ) . These representations are meaningful for several reasons . First , since we are leveraging a goal-conditioned policy , they are aware of dynamics and are able to capture local connectivity of the environment . Secondly , the representation is optimized so that it captures only the functionally relevant elements of state . Requirement for Goal Conditioned Policy : A natural question to ask is whether needing a goalconditioned policy is too strong of a prerequisite . However , it is worth noting that the GCP can be trained with existing RL methods ( TRPO ) using a sparse task-agnostic reward ( Section 6.2 , Appendix A.1 ) – obtaining such a policy is not especially difficult , and existing methods are quite capable of doing so ( Nair et al . ( 2018 ) ) . Furthermore , it is likely not possible to acquire a functionalityaware state representation without some sort of active environment interaction , since dynamics can only be understood by observing outcomes of actions , rather than individual states . Importantly , we discuss in the following section how ARCs help us solve tasks beyond what a simple goalconditioned policy can achieve . 4 USING ACTIONABLE REPRESENTATIONS FOR DOWNSTREAM TASKS . A natural question that emerges when learning representations from a goal-conditioned policy pertains to what such a representation enables over the goal-conditioned policy itself . Although goalconditioned policies enable reaching between arbitrary states , they suffer from fundamental limitations : they do not generalize very well to new states , and they are limited to solving only goalreaching tasks . We show in our empirical evaluation that the ARC representation expands meaningfully over these limitations of a goal-conditioned policy - to new tasks and to new regions of the environment . In this section , we detail how ARCs can be used to generalize beyond a goalconditioned policy to help solve tasks that can not be expressed as goal reaching ( Section 4.1 ) , tasks involving larger regions of state space ( Section 4.2 ) , and temporally extended tasks which involve sequences of goals ( Section 4.3 ) .

Summary: The paper presents a method to learn representations where proximity in euclidean distance represents states that are achieved by similar policies. The idea is novel (to the best of my knowledge), interesting and the experiments seem promising. The two main flaws in the paper are the lack of details and missing important experimental comparisons.