Generalizing Reinforcement Learning to Unseen Actions

1 INTRODUCTION . Imagine visiting your friend for the first time , and you decide to cook your favorite dish there . But since you have never been in their kitchen before , there could be certain tools you have never seen , like an odd-shaped sponge . However , by looking at its porous texture or observing its interaction with water , you can understand that this object can absorb liquid . Later during cooking when you want to clean the table , you can select that sponge since you can relate its absorbing characteristics with another tool you have used for cleaning . Just like in this scenario , our tasks often involve making selections from novel or unseen entities . When we encounter such choices , we examine them to first understand their functionality which informs our selection process while solving a task . Can machines also understand previously unseen choices and subsequently use them for solving tasks ? From a reinforcement learning perspective , this brings an interesting question of how to enable generalization of discrete action policies to solve tasks using unseen sets of actions . Prior work in deep reinforcement learning has explored generalization over environments ( Cobbe et al. , 2018 ; Nichol et al. , 2018 ) , and tasks ( Finn et al. , 2017 ; Parisi et al. , 2018 ) . However , action space generalization is relatively unexplored and is crucial for agents to be flexible in the face of novel circumstances , like selecting an unseen sponge for a known task of cleaning in above example . In this work , our goal is to develop a framework that reflects the two phases of solving action generalization : ( 1 ) general understanding of unseen discrete actions from their characteristic information ( like appearance or behaviors ) , and ( 2 ) training a policy to solve tasks by utilizing this general understanding . However , an action can have diverse behaviors and hence requires a collection of data ( e.g . different viewpoints , videos or state trajectories of how it effects on environment ) to sufficiently express this diversity . Hence , the primary challenge is to develop a generalizable unsupervised learning method which can extract an action ’ s characteristics from a dataset constituting its diverse effects . To this end , we propose to embed actions ’ datasets by extending the work on hierarchical variational autoencoders ( Edwards & Storkey , 2017 ) . The obtained embeddings reflect an action ’ s general utility , and can be used as action representations in the downstream task of reinforcement learning . However , conventional reinforcement learning algorithms utilize the available actions in a way that best optimizes a reward . This directly incentivizes a policy to overfit to the actions seen during training , just like the problem of overfitting to training data in supervised learning . To address this challenge , we formulate this problem as risk minimization ( Vapnik , 1992 ) for reinforcement learning , and propose regularization objectives to enforce generalization of policy to unseen actions . The main contributions of this paper are : ( 1 ) introducing the problem and a proposed solution to enable action space generalization in reinforcement learning , ( 2 ) representing an action with a dataset reflecting its diverse characteristics , and employing a generalizable unsupervised learning approach to embed these datasets . ( 3 ) a method to use learned action representations in reinforcement learning , and regularization methods to enable learning of generalizable policies . 2 RELATED WORK . Generalization in reinforcement learning In typical deep reinforcement learning ( RL ) settings ( Mnih et al. , 2015 ; 2016 ; Lillicrap et al. , 2015 ; Schulman et al. , 2017 ) , a policy or value network learns to act over an action space of fixed dimensionality . By taking states or observations as input to neural networks , these methods are able to generalize to unseen environment states drawn from a similar distribution as training ( Cobbe et al. , 2018 ; Nichol et al. , 2018 ) . Similarly , prior works have explored generalization in RL for unseen instructions ( Oh et al. , 2017 ) , new sequences of subtasks ( Andreas et al. , 2017 ) , manipulation of unseen tools ( Fang et al. , 2018 ; Xie et al. , 2019 ) , task demonstrations ( Xu et al. , 2017 ) , and agent morphologies ( Wang et al. , 2018 ; Sanchez-Gonzalez et al. , 2018 ; Pathak et al. , 2019 ) . In contrast , our framework enables zero-shot generalization of RL policies when the agent gets a previously unseen action set . Unsupervised representation learning for downstream tasks Bengio et al . ( 2013 ) state representation learning of data makes it easier to extract useful information when building predictors . Prior works show that such representations have been useful for a variety of downstream tasks , like classification and video prediction ( Denton et al. , 2017 ) , visually representing objects for relational reasoning tasks ( Steenbrugge et al. , 2018 ) , representing image-states for domain adaptation in RL ( Higgins et al. , 2017 ) , and , representing goals for better exploration ( Laversanne-Finot et al. , 2018 ) and sample efficiency ( Nair et al. , 2018 ) in RL . In this paper , we show how unsupervised representation learning over datasets ( Edwards & Storkey , 2017 ) can be used for embedding discrete actions , and enable generalization in the downstream task of reinforcement learning . Action Representations Using continuous representations of discrete actions , a policy can be trained through a combined Q-function over state and action representations ( He et al. , 2015 ) , or in an actor-critic architecture by selecting the nearest neighbor action vector to the policy ’ s continuous output ( Van Hasselt & Wiering , 2009 ; Dulac-Arnold et al. , 2015 ) . Unlike our work , these prior works assume access to ground truth action representations , which are usually not readily available . In other related work , action representations are learned implicitly through inverse model on a fixed action space to ease learning in large discrete action spaces ( Chandak et al. , 2019 ) or for intrinsic reward ( Kim et al. , 2019 ) . In contrast , we do not have the assumption of fixed action space as we learn action representations separately , and hence are able to incorporate new actions for the same policy . While Tennenholtz & Mannor ( 2019 ) pre-learn action representations explicitly using co-occurrence of actions in task-specific demonstrations , our generic embedding method applies to various modalities of datasets to represent actions , which are task-independent and hence suited for generalization to unseen actions . Skill and Trajectory Embeddings In reinforcement learning , variational autoencoders ( VAE ) ( Kingma & Welling , 2014 ) are often used for learning an abstraction for continuous entities like skills and state-action trajectories . Specifically , Co-Reyes et al . ( 2018 ) utilize a trajectory autoencoder for hierarchical RL , and Lynch et al . ( 2019 ) learn a latent space of trajectories and employ a goal-conditioned planner over it . Hausman et al . ( 2018 ) learn an embedding space of skills through a shared policy for different tasks , and utilize this space for solving other related tasks . In this paper , we extend the framework of hierarchical VAE ( Edwards & Storkey , 2017 ; Achille et al. , 2019 ) to trajectories , so as to embed even sequential datasets which are better indicative of action behavior . In general , an action can be a discrete skill choice , and an action ’ s behavior can be represented as the trajectory of effects it causes on the environment . Since individual trajectories are incapable of capturing the diverse effects of actions , we propose to use datasets for representing actions . 3 GENERALIZATION TO UNSEEN ACTIONS . Our approach is based on the intuition that when humans encounter previously unseen discrete entities , we examine them to understand their functionality through visual inspection or physical interaction , before deciding what to select for a task . Once the general functionality is inferred , these discrete objects can be used as actions in decision-making tasks , like selecting a tool for cooking or furniture assembly . In this paper , we incorporate these two phases ( Figure 2 ) to enable agents to utilize previously unseen actions : ( 1 ) extracting representations of actions from datasets of unstructured information ( e.g . image , videos ) , and ( 2 ) training a reinforcement learning policy to utilize these action representations with the joint objective of generalization and reward maximization . In order to represent actions , we note that an action can have diverse behaviors like how it interacts with its environment . Further , there can be various ways an agent observes this dataset . In the sponge example , the action exhibits diverse properties like absorption or compression , and the agent can observe this through porous texture ( image ) or through interacting with it ( states trajectory ) . Therefore , in its most general form , information about an action can be expressed in the form of a diverse collection of unstructured data like images , videos or trajectories . To learn action representations in an unsupervised and generalizable , we use a hierarchical VAE and extend it to sequence data like videos ( Section 3.2 ) . Next , we show how a policy is trained to use these action representations as input , and propose training objectives for enabling generalization ( Section 3.3 ) . 3.1 PRELIMINARIES . For a learning agent , we denote the entire set of possible discrete actions as A . For evaluation , we assume an episodic setting , where the agent only has a subset A ⊂ A of actions available to it . Each action a ∈ A has an associated dataset D = { x1 , . . . , xL } of observable samples xn ∼ P ( x|a ) which are characteristic of the action a . During training , the agent only has access to a subset AK ⊂ A of known actions . During evaluation , action set A can even be totally unseen for the agent . The action set A constitutes the discrete action space of an episodic Markov decision process ( MDP ) . Formally , { S , A , T , R , γ } defines the set of states , actions , transition probability , reward function , and discount factor of an MDP . Given a set of available actions A ⊂ A at any time step t , the core problem is to learn parameters θ of a policy πθ ( at|st ) , which defines a probability distribution over actions at ∈ A for a state st . Since the available action sets A are stochastically sampled and the environments are in general stochastic , we primarily consider stochastic policies in this paper . The performance of πθ is evaluated based on a discounted return R = ∑T−1 t=0 γ tr ( st , at ) where r is the reward function and T is the episode horizon . The aim is to train a policy which only has access to the known actions AK and its datasets , but generalizes to maximize reward on unseen actions . 3.2 UNSUPERVISED LEARNING OF ACTION REPRESENTATIONS . We represent the diverse characteristics of an action with a dataset of observed information . To extract usable information from these action datasets , we propose an unsupervised representation learning method to learn action embeddings . Our key insight is that the common information underlying different samples of an action ’ s dataset best represents the general properties of that action . Therefore , we aim to learn an action encoder to map each discrete action ’ s entire dataset to a continuous representation . For unsupervised learning of this encoder , we can use a variational autoencoder ( VAE ) with reconstruction objective ( Kingma & Welling , 2014 ) . However , since the input to VAE is in the form of a dataset , it should capture the information shared across multiple data samples . Therefore we encode both , the action datasets and the sample within each action ’ s dataset into a hierarchy of connected latent spaces . Such a hierarchical VAE ( HVAE ) architecture has been explored by Edwards & Storkey ( 2017 ) for few-shot classification and clustering of datasets . We use it for the purpose of encoding action datasets and using them for generalization ( Figure 2 ) . HVAE is composed of an action VAE over datasets and an instance VAE over samples . The encoders and decoders of the instance VAE are conditioned on its parent action latent vector . For each action a and its associated dataset D = { x1 , . . . , xL } , the action encoder qa ( c|D ) is used to sample an action latent c , while regularized by an action prior pa ( c ) . For each action sample x ∈ D , the instance encoder qs ( z|x , c ) is used to sample a latent z encoding the sample instance x , while conditioned on c. The prior distribution ps ( z|c ) as well as the decoder p ( x|z , c ) are also conditioned on the action latent . For each action dataset , ELBO comprises of reconstruction over data samples and the two KL divergence terms ( Edwards & Storkey , 2017 ) : LD = Ec∼qa ( .|D ) [ ∑ x∈D Ez∼qs ( .|x , c ) [ log p ( x|z , c ) ] −D [ qs ( z|x , c ) ||ps ( z|c ) ] ] −D [ qa ( c|D ) ||pa ( c ) ] ( 1 ) We further extend this framework to incorporate sequential data like state trajectories and videos , as that is more suitable to express behaviors of actions . For a dataset of trajectories τ , we use a Bi-LSTM encoder for qa ( c|D ) , and LSTM decoder p ( τ |z , c , s ) which also takes the initial state s of τ and reconstructs the rest of it ( Schuster & Paliwal , 1997 ; Wang et al. , 2017 ; Co-Reyes et al. , 2018 ) . For the case of video datasets , we also incorporated temporal skip connections ( Ebert et al. , 2017 ) from s by predicting an extra mask channel , to weigh contributions from the predicted frame and the first frame s. For getting representations of any action a ∈ A ( seen or unseen ) through a trained HVAE , we use the action dataset encoder qa ( c|Da ) output ’ s mean as the representation ca ( Figure 2 ) . This choice of using mean as representation follows prior work like Higgins et al . ; Steenbrugge et al . ( 2018 ) , but one could also use sampling from the output distribution as representation , as done in Locatello et al . ( 2019 ) . The generalizability of these representations to unseen actions depends on whether the action ’ s behaviors lie in the distribution of behaviors of known actions . Hence , the hierarchy in HVAE makes it an expressive encoder for actions , since even seemingly new discrete actions can have characteristics which belong to the distribution of previously seen effects .

Summary: This paper addresses the very interesting problem of generalising to new actions after only training on a subset of all possible actions. Here, the task falls into different contexts, which are inferred from an associated dataset (like pixel data). Having identified the context, it is used in the policy which therefore has knowledge of which actions are available to it.