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 deals with the problem of how to enable the generalization of discrete action policies to solve the task using unseen sets of actions. The authors develop a general understanding of unseen actions from their characteristic information and train a policy to solve the tasks using the general understanding. The challenge is to extract the action's characteristics from a dataset. This paper presents the HVAE to extract these characteristics and formulates the generalization for policy as the risk minimization.