Title: GENERALIZING REINFORCEMENT LEARNING TO UN-SEEN ACTIONS

Abstract: A fundamental trait of intelligence is the ability to achieve goals in the face of novel circumstances. In this work, we address one such setting which requires solving a task with a novel set of actions. Empowering machines with this ability requires generalization in the way an agent perceives its available actions along with the way it uses these actions to solve tasks. Hence, we propose a framework to enable generalization over both these aspects: understanding an action's functionality, and using actions to solve tasks through reinforcement learning. Specifically, an agent interprets an action's behavior using unsupervised representation learning over a collection of data samples reflecting the diverse properties of that action. We employ a reinforcement learning architecture which works over these action representations, and propose regularization metrics essential for enabling generalization in a policy. We illustrate the generalizability of the representation learning method and policy, to enable zero-shot generalization to previously unseen actions on challenging sequential decision-making environments. More training and testing videos can be found at sites.google.com/view/action-generalization/

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. 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.

Skill and Trajectory Embeddings:



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 = {x 1 , . . . , x L } of observable samples x n ∼ P (x|a) 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 π θ (a t |s t ), which defines a probability distribution over actions a t ∈ A for a state s t . 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 γ t r(s t , a t ) 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 A K and its datasets, but generalizes to maximize reward on unseen actions.

l 0 H F 0 M 3 B d Y 3 2 5 h D B j c 1 X L t h C 4 C k = " > A A A B / 3 i c b V A 7 T 8 M w G H R 4 l v I K I L G w W L R I T F X S B c Z K M D A W R B 9 S E 1 W O 4 7 R W H T u y H a Q q Z O C v s D C A E C t / g 4 1 / g 9 N m g J a T L J / u v k 8 + X 5 A w q r T j f F s r q 2 v r G 5 u V r e r 2 z u 7 e v n 1 w 2 F U i l Z h 0 s G B C 9 g O k C K O c d D T V j P Q T S V A c M N I L J l e F 3 3 s g U l H B 7 / U 0 I X 6 M R p x G F C N t p K F 9 X P c C w U I 1 j c 2 V e W O k s + s 8 r w / t m t N w Z o D L x C 1 J D Z R o D + 0 v L x Q 4 j Q n X m C G l B q 6 T a D 9 D U l P M S F 7 1 U k U S h C d o R A a G c h Q T 5 W e z / D k 8 M 0 o I I y H N 4 R r O 1 N 8 b G Y p V k d B M x k i P 1 a J X i P 9 5 g 1 R H l 3 5 G e Z J q w v H 8 o S h l U A t Y l A F D K g n W b G o I w p K a r B C P k U R Y m 8 q q p g R 3 8 c v L p N t s u E 7 D v W 3 W W n d l H R V w A k 7 B O X D B B W i B G 9 A G H Y D B I 3 g G r + D N e r J e r H f r Y z 6 6 Y p U 7 R + A P r M 8 f L u + W P Q = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " A O U V X H D v X I I S L p Z 9 / 3 i A q h 8 J A N I = " > A A A B / 3 i c b V A 7 T 8 M w G H R 4 l v I K I L G w W L R I T F X S B c Z K L I w F 0 Y f U R J X j O K 1 V x 4 5 s B 1 G F D P w V F g Y Q Y u V v s P F v c N o M 0 H K S 5 d P d 9 8 n n C x J G l X a c b 2 t l d W 1 9 Y 7 O y V d 3 e 2 d 3 b t w 8 O u 0 q k E p M O F k z I f o A U Y Z S T j q a a k X 4 i C Y o D R n r B 5 K r w e / d E K i r 4 n Z 4 m x I / R i N O I Y q S N N L S P 6 1 4 g W K i m s b k y b 4 x 0 9 p D n 9 a F d c x r O D H C Z u C W p g R L t o f 3 l h Q K n M e E a M 6 T U w H U S 7 W d I a o o Z y a t e q k i C 8 A S N y M B Q j m K i / G y W P 4 d n R g l h J K Q 5 X M O Z + n s j Q 7 E q E p r J G O m x W v Q K 8 T 9 v k O r o 0 s 8 o T 1 J N O J 4 / F K U M a g G L M m B I J c G a T Q 1 B W F K T F e I x k g h r U 1 n V l O A u f n m Z d J s N 1 2 m 4 N 8 1 a 6 7 a s o w J O w C k 4 B y 6 4 A C 1 w D d q g A z B 4 B M / g F b x Z T 9 a L 9 W 5 9 z E d X r H L n C P y B 9 f k D f l u W c Q = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " T v h j 4 V 3 I L p H x K 4 j M z j X b K d y W S k g = " > A A A B + X i c b V D N T g I x G P w W / x D / V j 1 6 a Q Q T T 2 S X i x 5 J v H h E I 2 A C G 9 L t F m j o t p u 2 S 4 I b 3 s S L B 4 3 x 6 p t 4 8 2 3 s w h 4 U n K T p Z O b 7 0 u m E C W f a e N 6 3 U 9 r Y 3 N r e K e 9 W 9 v Y P D o / c 4 5 O O l q k i t E 0 k l + o x x J p y J m j b M M P p Y 6 I o j k N O u + H k J v e 7 U 6 o 0 k + L B z B I a x H g k 2 J A R b K w 0 c N 1 a P 5 Q 8 0 r P Y X t n T v D Z w q 1 7 d W w C t E 7 8 g V S j Q G r h f / U i S N K b C E I 6 1 7 v l e Y o I M K 8 M I p / N K P 9 U 0 w W S C R 7 R n q c A x 1 U G 2 S D 5 H F 1 a J 0 F A q e 4 R B C / X 3 R o Z j n W e z k z E 2 Y 7 3 q 5 e J / X i 8 1 w + s g Y y J J D R V k + d A w 5 c h I l N e A I q Y o M X x m C S a K 2 a y I j L H C x N i y K r Y E f / X L 6 6 T T q P t e 3 b 9 r V J v 3 R R 1 l O I N z u A Q f r q A J t 9 C C N h C Y w j O 8 w p u T O S / O u / O x H C 0 5 x c 4 p / I H z + Q O R q J O m < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " W c P p u J H h D D Z w a r u 7 V P J 0 k I q q v 5 E = " > A A A B + X i c b V D N T g I x G P w W / x D / V j 1 6 a Q Q T T 2 S X i x 5 J v H h E I 2 A C G 9 L t F m j o t p u 2 S y Q b 3 s S L B 4 3 x 6 p t 4 8 2 3 s w h 4 U n K T p Z O b 7 0 u m E C W f a e N 6 3 U 9 r Y 3 N r e K e 9 W 9 v Y P D o / c 4 5 O O l q k i t E 0 k l + o x x J p y J m j b M M P p Y 6 I o j k N O u + H k J v e 7 U 6 o 0 k + L B z B I a x H g k 2 J A R b K w 0 c N 1 a P 5 Q 8 0 r P Y X t n T v D Z w q 1 7 d W w C t E 7 8 g V S j Q G r h f / U i S N K b C E I 6 1 7 v l e Y o I M K 8 M I p / N K P 9 U 0 w W S C R 7 R n q c A x 1 U G 2 S D 5 H F 1 a J 0 F A q e 4 R B C / X 3 R o Z j n W e z k z E 2 Y 7 3 q 5 e J / X i 8 1 w + s g Y y J J D R V k + d A w 5 c h I l N e A I q Y o M X x m C S a K 2 a y I j L H C x N i y K r Y E f / X L 6 6 T T q P t e 3 b 9 r V J v 3 R R 1 l O I N z u A Q f r q A J t 9 C C N h C Y w j O 8 w p u T O S / O u / O x H C 0 5 x c 4 p / I H z + Q O O n J O k < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " 7 u J l k P y W C R e P + C S 5 b z A 9 3 R h H 8 Q 4 = " > A A A B 9 H i c b V B N T 8 J A E J 3 i F + I X 6 t H L R j D B x J C W i x 5 J v H h E I x 8 J N G S 7 b G H D d l t 2 t 0 S s / A 4 v H j T G q z / G m / / G B X p Q 8 C W T v L w 3 k 5 l 5 X s S Z 0 r b 9 b W X W 1 j c 2 t 7 L b u Z 3 d v f 2 D / O F R Q 4 W x J L R O Q h 7 K l o c V 5 U z Q u m a a 0 1 Y k K Q 4 8 T p v e 8 H r m N 8 d U K h a K e z 2 J q B v g v m A + I 1 g b y S 1 G p Q f 0 h B 4 v E D k v d v M F u 2 z P g V a J k 5 I C p K h 1 8 1 + d X k j i g A p N O F a q 7 d i R d h M s N S O c T n O d W N E I k y H u 0 7 a h A g d U u c n 8 6 C k 6 M 0 o P + a E 0 J T S a q 7 8 n E h w o N Q k 8 0 x l g P V D L 3 k z 8 z 2 v H 2 r 9 y E y a i W F N B F o v 8 m C M d o l k C q M c k J Z p P D M F E M n M r I g M s M d E m p 5 w J w V l + e Z U 0 K m X H L j u 3 l U L 1 L o 0 j C y d w C i V w 4 B K q c A M 1 q A O B E T z D K 7 x Z Y + v F e r c + F q 0 Z K 5 0 5 h j + w P n 8 A L I u Q c A = = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " Z q F 1 D L Q L X P 6 x I 8 K S 6 S O W D H m A c j o = " > A A A B + X i c b V C 7 T s M w F L 0 p r 1 J e A U Y W i x a J q U q 6 w F i J h b E g + p D a q H I c p 7 X q 2 J H t V K q i / g k L A w i x 8 i d s / A 1 O 2 w F a j m T 5 6 J x 7 5 e M T p p x p 4 3 n f T m l r e 2 d 3 r 7 x f O T g 8 O j 5 x T 8 8 6 W m a K 0 D a R X K p e i D X l T N C 2 Y Y b T X q o o T k J O u + H k r v C 7 U 6 o 0 k + L J z F I a J H g k W M w I N l Y a u m 5 t E E o e 6 V l i r 5 z M a 0 O 3 6 t W 9 B d A m 8 V e k C i u 0 h u 7 X I J I k S 6 g w h G O t + 7 6 X m i D H y j D C 6 b w y y D R N M Z n g E e 1 b K n B C d Z A v k s / R l V U i F E t l j z B o o f 7 e y H G i i 2 x 2 M s F m r N e 9 Q v z P 6 2 c m v g 1 y J t L M U E G W D 8 U Z R 0 a i o g Y U M U W J 4 T N L M F H M Z k V k j B U m x p Z V s S X 4 6 1 / e J J 1 G 3 f f q / k O j 2 n x c 1 V G G C 7 i E a / D h B p p w D y 1 o A 4 E p P M M r v D m 5 8 + K 8 O x / L 0 Z K z 2 j m H P 3 A + f w B u n p O P < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " A O N o E 7 W g 1 2 J A H k O N I C J + h 6 5 6 Y N 8 = " > A A A B + 3 i c b V C 9 T s M w G P x S / k r 5 C 2 V k i W i R m K q k C 4 y V W B g L o i 1 S G 0 W O 4 7 R W H T u y H U Q V 9 V V Y G E C I l R d h 4 2 1 w 2 g z Q c p L l 0 9 3 3 y e c L U 0 a V d t 1 v q 7 K x u b W 9 U 9 2 t 7 e 0 f H B 7 Z x / W + E p n E p I c F E / I h R I o w y k l P U 8 3 I Q y o J S k J G B u H 0 u v A H j 0 Q q K v i 9 n q X E T 9 C Y 0 5 h i p I 0 U 2 P X m K B Q s U r P E X D k O 6 L w Z 2 A 2 3 5 S 7 g r B O v J A 0 o 0 Q 3 s r 1 E k c J Y Q r j F D S g 0 9 N 9 V + j q S m m J F 5 b Z Q p k i I 8 R W M y N J S j h C g / X 2 S f O + d G i Z x Y S H O 4 d h b q 7 4 0 c J a p I Z y Y T p C d q 1 S v E / 7 x h p u M r P 6 c 8 z T T h e P l Q n D F H C 6 c o w o m o J F i z m S E I S 2 q y O n i C J M L a 1 F U z J X i r X 1 4 n / X b L c 1 v e b b v R u S v r q M I p n M E F e H A J H b i B L v Q A w x M 8 w y u 8 W X P r x X q 3 P p a j F a v c O Y E / s D 5 / A P U q l G s = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " 1 P g O 7 Y a i c z 5 O P 1 c D F C N / u 8 C a h f Y = " > A A A B + X i c b V D N T g I x G P w W / x D / V j 1 6 a Q Q T T 2 S X i x 5 J 9 O A R j S A J b E i 3 W 6 C h 2 2 7 a L g n Z 8 C Z e P G i M V 9 / E m 2 9 j F / a g 4 C R N J z P f l 0 4 n T D j T x v O + n d L G 5 t b 2 T n m 3 s r d / c H j k H p 9 0 t E w V o W 0 i u V T d E G v K m a B t w w y n 3 U R R H I e c P o W T m 9 x / m l K l m R S P Z p b Q I M Y j w Y a M Y G O l g e v W + q H k k Z 7 F 9 s p u 5 7 W B W / X q 3 g J o n f g F q U K B 1 s D 9 6 k e S p D E V h n C s d c / 3 E h N k W B l G O J 1 X + q m m C S Y T P K I 9 S w W O q Q 6 y R f I 5 u r B K h I Z S 2 S M M W q i / N z I c 6 z y b n Y y x G e t V L x f / 8 3 q p G V 4 H G R N J a q g g y 4 e G K U d G o r w G F D F F i e E z S z B R z G Z F Z I w V J s a W V b E l + K t f X i e d R t 3 3 6 v 5 9 o 9 p 8 K O o o w x m c w y X 4 c A V N u I M W t I H A F J 7 h F d 6 c z H l x 3 p 2 P 5 W j J K X Z O 4 Q + c z x 8 / Z J N w < / l a t e x i t > < l a t e x i t s h a _ b a s e = " w B H e T n Y c T t l s p K y G F n Q T = " > A A A B + i c b V A T M w G H T K q R X K C O L R Y v E V C V d Y K w E A N B C G U e Q T m v V s S P b Q V R R / g o L A w i x k f Y + D c b Q Z o O c n y e P M F C a N K O V d n Y N r e q e W v Y P D o / s p f i V R i s O C C T k M k C K M c t L T V D M y T C R B c c D I I J h d F / g k U h F B X / Q R M Z p w G l G M t J F u c B K F a h b K v x a d Y b T c h a A Q t S Q O U P r z g U O I J g h p U a u k g v Q J T z E h e G e K J A j P I S M D O U o J s r L F t l z e G U E E Z C m s M X K i / N z I U q y K d m Y y R n q p V r x D / a p j q j P I k Y T j U N R y q A W s C g C h l Q S r N n c E I Q l N V k h n i K J s D Z U w J u q X m / X K d l n v X b n T u y z q q B S c g Q v g g k v Q A b e g C o A g y f w D F B m V b L a b E c r V j l z g n A + v z B W y l E w = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " A O N o E 7 W g 1 2 J A H k O N I C J + h 6 5 6 Y N 8 = " > A A A B + 3 i c b V C 9 T s M w G P x S / k r 5 C 2 V k i W i R m K q k C 4 y V W B g L o i 1 S G 0 W O 4 7 R W H T u y H U Q V 9 V V Y G E C I l R d h 4 2 1 w 2 g z Q c p L l 0 9 3 3 y e c L U 0 a V d t 1 v q 7 K x u b W 9 U 9 2 t 7 e 0 f H B 7 Z x / W + E p n E p I c F E / I h R I o w y k l P U 8 3 I Q y o J S k J G B u H 0 u v A H j 0 Q q K v i 9 n q X E T 9 C Y 0 5 h i p I 0 U 2 P X m K B Q s U r P E X D k O 6 L w Z 2 A 2 3 5 S 7 g r B O v J A 0 o 0 Q 3 s r 1 E k c J Y Q r j F D S g 0 9 N 9 V + j q S m m J F 5 b Z Q p k i I 8 R W M y N J S j h C g / X 2 S f O + d G i Z x Y S H O 4 d h b q 7 4 0 c J a p I Z y Y T p C d q 1 S v E / 7 x h p u M r P 6 c 8 z T T h e P l Q n D F H C 6 c o w o m o J F i z m S E I S 2 q y O n i C J M L a 1 F U z J X i r X 1 4 n / X b L c 1 v e b b v R u S v r q M I p n M E F e H A J H b i B L v Q A w x M 8 w y u 8 W X P r x X q 3 P p a j F a v c O Y E / s D 5 / A P U q l G s = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " b z h x O w C b r b 5 V I d j L W 1 U h T Y i P Q Z A = " > A A A B + 3 i c b V C 9 T s M w G P x S / k r 5 C 2 V k i W i R m K q k C 4 y V W B g L o i 1 S G 0 W O 4 7 R W H T u y H U Q V 9 V V Y G E C I l R d h 4 2 1 w 2 g z Q c p L l 0 9 3 3 y e c L U 0 a V d t 1 v q 7 K x u b W 9 U 9 2 t 7 e 0 f H B 7 Z x / W + E p n E p I c F E / I h R I o w y k l P U 8 3 I Q y o J S k J G B u H 0 u v A H j 0 Q q K v i 9 n q X E T 9 C Y 0 5 h i p I 0 U 2 P X m K B Q s U r P E X H k a 0 H k z s B t u y 1 3 A W S d e S R p Q o h v Y X 6 N I 4 C w h X G O G l B p 6 b q r 9 H E l N M S P z 2 i h T J E V 4 i s Z k a C h H C V F + v s g + d 8 6 N E j m x k O Z w 7 S z U 3 x s 5 S l S R z k w m S E / U q l e I / 3 n D T M d X f k 5 5 m m n C 8 f K h O G O O F k 5 R h B N R S b B m M 0 M Q l t R k d f A E S Y S 1 q a t m S v B W v 7 x O + u 2 W 5 7 a 8 2 3 a j c 1 f W U Y V T O I M L 8 O A S O n A D X e g B h i d 4 h l d 4 s + b W i / V u f S x H K 1 a 5 c w J / Y H 3 + A A k h l H g = < / l a t e x i t > < l a t e x i t s h a 1 _ b a s e 6 4 = " D q T / 1 m B l O f X e D W M T s b V l o k G Q L x g = " > A A A B + 3 i c b V C 9 T s M w G P x S / k r 5 C 2 V k i W i R m K q k C 4 y V W B g L o i 1 S E 1 W O 4 7 R W H T u y H U Q V 9 V V Y G E C I l R d h 4 2 1 w 2 g z Q c p L l 0 9 3 3 y e c L U 0 a V d t 1 v q 7 K x u b W 9 U 9 2 t 7 e 0 f H B 7 Z x / W + E p n E p I c F E / I h R I o w y k l P U 8 3 I Q y o J S k J G B u H 0 u v A H j 0 Q q K v i 9 n q U k S N C Y 0 5 h i p I 0 0 s u t N P x Q s U r P E X L m f 0 n l z Z D f c l r u A s 0 6 8 k j S g R H d k f / m R w F l C u M Y M K T X 0 3 F Q H O Z K a Y k b m N T 9 T J E V 4 i s Z k a C h H C V F B v s g + d 8 6 N E j m x k O Z w 7 S z U 3 x s 5 S l S R z k w m S E / U q l e I / 3 n D T M d X Q U 5 5 m m n C 8 f K h O G O O F k 5 R h B N R S b B m M 0 M Q l t R k d f A E S Y S 1 q a t m S v B W v 7 x O + u 2 W 5 7 a 8 2 3 a j c 1 f W U Y V T O I M L 8 O A S O n A D X e g B h i d 4 h l d 4 s + b W i / V u f S x H K 1 a 5 c w J / Y H 3 + A A R 4 l H U = < / l a t e x i t >

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 = {x 1 , . . . , x L }, the action encoder q a (c|D) is used to sample an action latent c, while regularized by an action prior p a (c). For each action sample x ∈ D, the instance encoder q s (z|x, c) is used to sample a latent z encoding the sample instance x, while conditioned on c. The prior distribution p s (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) :

L D = E c∼qa(.|D) x∈D E z∼qs(.|x,c) [log p(x|z, c)] -D[q s (z|x, c)||p s (z|c)] -D[q a (c|D)||p a (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 q a (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 q a (c|D a ) output's mean as the representation c a (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.

3.3. LEARNING POLICIES OVER ACTION REPRESENTATIONS:

While solving tasks with new actions, humans first form a general interpretation of the behaviors of actions, and then utilize it to take appropriate actions. Similarly, once our agent learns actions representations based on observed datasets (section 3.2), it should learn to utilize them for solving tasks. This involves not only extracting the task-specific information from the representations, but also doing so in a generalizable manner so that it can utilize previously unseen action representations.

Here we assume access to an embedder φ, and hence the associated action representations c a = φ(D a ) for each a ∈ A. Our aim is to learn a policy π θ (a|s, A, φ) which maximizes the expected reward under any set of available actions A ⊂ A. We propose to utilize the action representations c a as inputs to the policy, which acts as a function approximator over action representations and states. Specifically, our policy consists of a utility function f θ : S × R d → R, which maps a d-dimensional action embedding and a state to its utility. The probability distribution over actions is simply defined as the Softmax over the utilities of each available action a ∈ A.

π θ (a|s, A) = e f θ (s,ca)

a ∈A e f θ (s,c a )

(2)

We can train the parameters θ using policy gradient methods on π.

3.4. ENABLING GENERALIZATION TO UNSEEN ACTIONS:

The primary objective is to find parameters θ of a policy which maximizes rewards on unseen action sets A ⊂ A. We formulate this generalization problem with statistical learning theory (Vapnik, 1998; 2013) , and propose regularization objectives which aim to satisfy its assumptions. The theory mainly deals with generalization in supervised learning problems with an assumption on training examples to be independent and identically distributed (i.i.d. sampled). In a reinforcement learning setup with action representations c a , the objective becomes minimizing the theoretical risk of the policy:

EQUATION

Here L is a real-valued loss function which measures the optimality of policy hypothesis π θ (equivalently the utility function f θ ) with respect to the output y * of an optimal stochastic policy π * at state s. While L, π * or y * cannot be defined in closed form, the definition of optimal policy (Sutton & Barto, 2018; Sutton et al., 2000) makes this objective equivalent to maximizing the cumulative reward R, given an unseen action space A and their action representations c a . Note that the expectation in Eq. 3 is also over states s drawn from environment, but dropped for readability.

During training, the agent only has access to a limited set of known actions A K ⊂ A. The standard reward maximization objective in RL with training set of actions, A K is equivalent to Empirical Risk Minimization (ERM) of the hypothesis π θ (Vapnik, 1992) . Hence, the ERM training objective is:

EQUATION

However, a policy trained with ERM is prone to overfitting to data seen during training, just like in supervised learning. This problem becomes more severe for on-policy RL because the distribution of input data, (s, c a ) used for training π θ is governed by the actions taken by π θ itself. This means that the policy can bias its own training data distribution towards a small subset of actions, while ignoring other actions, which could actually be more informative about the actions available at test time. Since there is no prior information on the distribution over action space A ⊂ A at test time, it is assumed to be uniform. Therefore, this discrepancy between training and evaluation due to the non-stationarity of RL training, breaks the identical distribution (in i.i.d.) assumption in statistical learning theory (Bousquet et al., 2003) . To address this non-uniformity in training data, the following regularizing techniques are proposed to augment the ERM objective in Eq. 4:

(1) Maximum entropy regularization: Maximum entropy objective (Ziebart et al., 2008) augments Eq. 4 with the stochastic policy's entropy H[π θ (a|s)] with weight β, as in Eq. 5. This makes the policy maximize environment reward, under the constraint of taking diverse actions. This helps generalization in two ways: (a) the input data distribution used for training the policy becomes more uniform over action representations, and (b) the policy outputs maximum entropy distributions which make the least assumptions about the possibly unseen set of available actions A ⊂ A, and hence by the principle of maximum entropy (Jaynes, 1957; Guiasu & Shenitzer, 1985) overfits the least.

EQUATION

(2) Changing action spaces: The training data distribution can be made more uniform by sampling a set of available actions A ⊂ A K , uniformly in every episode. This blocks certain actions, making the policy select appropriate actions only from the available set A. Hence, the experience collected by the policy is uniformly spread over the known actions A K , making the training data distribution more identical to the assumed uniform distribution at test time. Eq. 6 shows this training objective:

EQUATION

(3) Clustering similar actions: The known action space A K can contain several groups of similar actions (e.g. various knives for cutting), and a randomly sampled action space A may contain actions from each group. This can be exploited by the reward-maximizing policy during training to overfit to actions from particular groups, but it will fail to generalize if similar actions are unavailable while testing. To avoid this, we propose to utilize the pre-learned action representations (section 3.2) to partition A K into a set of k groups

G K = {g 1 . . . g k },

where k is a hyperparameter. For every episode during training, an action set A G is built by sampling a subset of groups, G ⊂ G K and then sampling actions from G only. Two-step sampling ensures that certain groups of actions are blocked every episode, encouraging the policy to utilize underused action groups as well, making training data more uniform over the action representation space. We use equal-sized variant of k-means for clustering. The overall objective is formalized in Eq. 7 below:

max θ E a∼π θ (.|s,A G ) [R(s, c a ) + βH[π θ (a|s)]], where A G ⊂ {a|a ∈ g, g ∈ G} and G ⊂ G K (7)

In experiments (Section 5), we perform model selection based on a validation set of actions. We further ablate each regularization techniques and analyze their contribution in different environments.

4. ENVIRONMENTS:



4.1. GRID WORLD:

In GRID WORLD environment (Chevalier-Boisvert et al., 2018) , an agent navigates a 2D 9x9 maze to reach a goal cell for a sparse reward. A column of lava is randomly placed in every episode, touching which ends the episode. The discrete action space consists of all 5-step macro actions, where each macro-action is defined by a 5-length sequence of left, right, up or down movement. The entire action space of size 4 5 = 1024 actions is randomly split into a train and test set of 512 actions. The action datasets are collected on an empty grid where the agent is initialized at random locations. Two kinds of data types are used to represent the state sequence of agent -one-hot vectors and continuous (x,y) grid coordinates.

4.2. RECOMMENDER SYSTEM:

The RECOMMENDER SYSTEM environment (Rohde et al., 2018) simulates how users may respond to product recommendations. Every episode, the agent must recommend items to a new user with the objective of maximizing the click through rate (CTR) for the recommendations. This simulated environment uses randomly initialized embeddings for recommendations (actions), and we use the same to demonstrate policy generalization to new actions. Action space of size 10,000 is randomly split equally into train and test actions.

4.3. CHAIN REACTION TOOL ENVIRONMENT (CREATE):

CREATE is a physics-based puzzle where the goal is to make a specified ball reach a goal position (blue), inspired by the popular video game The Incredible Machine. The agent must place tools in real time to manipulate the path of the ball to reach the goal position. The environment presents a challenging multi-step task, requiring the agent to select the tool to place as well as its position (x, y) on the screen. The agent has access to a subset of diverse tools such as trampolines, see-saws, cannons, funnels, and conveyor belts (Appendix C.2). The position aspect makes this a parameterized action space Hausknecht & Stone (2015) with both discrete and continuous components. Our policy architecture consists of another head to output this continuous vector and it is trained jointly with the discrete action. We solve 3 different CREATE tasks: Push, Navigate and Obstacle. The tools evaluated at test time are completely unseen tool types from those seen during training.

4.4. SHAPE STACKING:

In SHAPE STACKING the agent must drop blocks on a table to build the highest standing tower. Our objective is different from prior works (Groth et al., 2018; Lerer et al., 2016) in that we maximize the tower height in an RL setting, whereas the prior work predicts the stability of the tower. Similarly to CREATE, the action space in Object Stacking, consists of (x, y) coordinates of where the object should be dropped above the table. This environment is shows our ability to generalize problem solving ability to a new action space in a complex 3D task. The action dataset here are images of the objects from various angles (or viewpoints). In this case the visual appearance of the object is sufficient to infer its functionality.

5. EXPERIMENTS:



5.1. BASELINES & ABLATIONS:

Baselines: We compare against two policy architectures which can utilize action representations for generalization to unseen action sets. We also compare against a VAE-based non-hierarchical embedding learning method, to learn action representations from unstructured action data (see Fig. 3 ).

• Nearest Neighbor: During training, a policy is learned over all known actions. Given unseen actions, the policy's output is used to select the nearest available action in embedding space. • Distance Based: Based on Dulac-Arnold et al. ( 2015), a continuous action-space policy outputs in the action embedding space and the closest available action to this output is selected. • Non hierarchical VAE: A shared VAE is trained over the samples across all action datasets. An action's embedding is then computed as the mean over the embeddings of samples in its dataset.

Ablations: We individually ablate each of the three proposed regularization metrics in our method.

• Ours: no entropy: Trained without entropy regularization by setting entropy coefficient to zero.

• Ours: no changing: Trained over the entire set of known actions without any action space sampling.

• Ours: no clustering: Training action-space is uniformly sampled (Eq. 6), no k-means clustering.

Alternate embeddings::

We compare how the embedding learning method (HVAE) applies to various forms of unstructured data (Fig. 3 ). In CREATE, action datasets comprise of state trajectories for tool behavior, except Ours (video) where video datasets are used instead. In Grid World, action datasets contain trajectories of states in one-hot representation, except in Ours (state) where states are real-valued 2D coordinates. Ours (Ground Truth) representations are not learned, but instead uses manually engineered representations for the actions. Detailed descriptions are present in Appendix C.

5.2. QUANTITATIVE RESULTS:

The generalization performance of the policy to unseen actions across all environments and method variations is shown in 3. As seen from the results our method or ablations all of our methods have the strongest ability to generalize to unseen actions across a variety of environments. The difference among our ablations is smaller in simpler environments like Grid World, Recommender systems and Shape Stacking, where the unseen action spaces are very similar to training actions. The effect of clustering-based sampling and entropy regularization can be seen for Obstacle and Navigate environments, which require solving the task with quite different tools at testing. CREATE Push is solvable with a wide variety of tools, and hence the no-entropy policy trains to a higher reward, and is able to generalize as well as many unseen tools can solve the task easily. The performance of our method against its variant with non-hierarchical VAE embeddings shows the importance of hierarchy in latent space to represent actions.

We test the generalizability of our embedder and policy for the task of zero-shot generalization to unseen actions. Specifically, our primary experiments across all four environments, discussed in section 4, train a policy on a fixed set of actions, tune hyperparameters on a separate evaluation set, and then test the ability to generalize to a new set of actions. We further provide qualitative analysis on cases where this generalization succeeds and fails. Finally, we evaluate how our method's generalizability varies with the degree of difference between seen and unseen.

5.3. FURTHER ANALYSIS:

Qualitative results of the policy test performance are shown in 5. The left and middle column contain success cases. In the left column for CREATE we seen the policy, despite never having used on of the tools before, still be able to solve the task. Likewise, for shape stacking we see the policy able to use novel shapes to build a tall and stable pile to maximize the height. We also show cases of failure We also analyze the conditions needed for generalization to unseen actions. We perform all analyses on CREATE Push task because of the large diversity of tool functionalities. We show generalization across changing physical tool parameters with angle and the embeddings the policy is trained and tested on. Finally, we show the effect of unseen versus seen actions on performance.

Figure 5 : Qualitative analysis: shown are two success cases and one failure case for CREATE and Object Stacking. In CREATE the trace of the ball trajectory is outlined. All of the tools or objects in these results the policy is generalizing to select and was not trained over these actions.

6. CONCLUSION:

Generalization to novel circumstances is an important ability to have, for robust and widely applicable artificial agents. In this paper we propose the problem of generalization of reinforcement learning policies to unseen spaces of actions, with the use of action representations learned in an unsupervised manner. Our two-phase framework demonstrates how representation learning can be combined with the downstream task of reinforcement learning, specifically to represent actions. We demonstrate the efficacy of our methods on four challenging environments, and discuss which variants work when.

The key takeaway is that when unseen actions are quite different from known actions, then more regularization helps to train generalizable policies.

A ALGORITHM (Edwards & Storkey, 2017) , and we use RAdam optimizer (Liu et al., 2019) . For training our policy, we use PPO (Schulman et al., 2017; Kostrikov, 2018) with Adam optimizer (Kingma & Ba, 2015) . Performance is computed across 5 seeds trained from the same model.

One possible way to adapt a policy to an unseen set of actions is to spend time fine-tuning the learned policy for the new action space. This process keeps the earlier layers of the policy but fine-tunes the actual action selection in the final layer by re-initializing the output layer. This process is able to adapt to an unseen action set of any size. However, this fine-tuning process can be time consuming which is why we try to achieve generalization in this paper. In Figure 10 , we show the performance of such fine tuning on the three CREATE tasks. As seen from the figure, it takes millions of steps to fine-tune, proving prohibitive for adaptive policies. On the other hand our framework generalizes to the new action space without the need for expensive RL retraining. Observations: The observation space is 9x9x3 where three channels represent (1) object id (agent, wall, goal, lava), (2) color of an object, and (3) zero in each cell. The agent gets a flattened vector of this 9x9x3 matrix.

Actions: An action of the agent is 5 consecutive moves in 4 directions. Hence, 4 5 = 1, 024 actions are possible in total. Once the agent selects an action, it executes 5 sequential moves step-by-step.

During an action execution, if the agent hits the boundary, it will stay in the current cell. If the agent steps on lava, the game will be terminated. The whole action set is divided into 8:1:1 split of train, validation and test action sets.

Rewards: GRID WORLD provides a sparse reward, 1 -0.9 × (step/max_step), only when the agent reaches the goal. The reward is discounted based on the number of actions taken to encourage a shorter path to the goal.

Termination: Each game is terminated when the agent takes max_steps (64) actions or the agent reaches the goal or the lava wall.

Action Datasets: The action datasets are observations of an agent performing actions in a 80x80 grid with no obstacles. The observations constitute the trajectory of states the agent sees during application of a macro-action, starting from randomized initial states. A dataset of 1024 such trajectories is used to represent a single macro-action, and is encoded by the HVAE to an action representation which identifies the underlying action behavior. We consider three kinds of embeddings in our experiments:

• One-hot (default): State is represented by two 80-dimensional one-hot vectors of x and y coordinates of the agent's locations on the 80x80 grid. Autoencoder reconstruction is based on a softmax cross-entropy loss over the states in the trajectory. The learned action embeddings are 16-dimensional. Rewards: The agent gets +1 reward when the marker ball hits the target ball, and +10 reward when the target ball passes the goal. In addition to reward for success, the environment provides some intermediate rewards. For every time-step, +0.01 reward is given to encourage the agent to keep balls inside the screen. At the same time, an invalid action is penalized by giving -0.01 reward (e.g., placing a tool outside of the screen or placing a tool on top of other tools).

Termination: Each game is terminated when the agent takes 20 actions, or the marker ball goes out of the screen before it hits the target ball, or the target ball goes out of the screen before it passes the goal.

Action Dataset: The action dataset is constructed by testing the properties of each tool through scripted interactions with a probe ball. A ball is launched at a given tool from various angles and speeds and then interacts with the tool. The properties of the tool will determine the deflection path of the ball. Testing interaction from various angles and speeds helps to build a better understanding of the tool. The collected datasets consist of 1024 trajectories of length 10 of the ball's interaction with the tool. Each trajectory consists of observations in the form of either environment state (default) or 48x48 gray-scale images (Ours (image) in Fig. 3 ). We consider three kinds of embeddings in our experiments:

• State trajectory (default): Each environment state in a trajectory is represented by 2D coordinates and 2D velocity of the probe ball, concatenated with 2D tool's location. Autoencoder reconstruction is based on Gaussian log-likelihood over the states in the trajectory. The learned action embeddings are 128-dimensional.

• Video: Each state in a trajectory is represented by 48x48 gray-scale image, which makes the trajectory a video of probe ball's interaction with the tool. Autoencoder reconstruction is based on Gaussian log-likelihood over the images in the input video. The learned action embeddings are 128-dimensional. • Ground-truth: These are 32-dimensional embeddings containing the true knowledge of the tool in the form of a one-hot encoding of the tool type along with its properties such as angle, length, elasticity, etc.

C.3 SHAPE STACKING SHAPE STACKING is a mujoco simulation environment (Todorov et al., 2012) where the agent is given a set of objects of different shapes of varying sizes, for instance, cubes, rectangles, spheres, round cylinders, archs, etc. The objective is to stack a tower as high as possible by choose the appropriate objects given in a particular episode.

Observations: An observation for each time-step is an 84x84x3 image of the shapes laying on the table.

Actions: In total there are 900 distinct shapes in the environment. 675 are used for learning the policy. The remaining 225 are used for evaluating performance on unseen actions. The same types of polygons do not appear in both train and test. Like CREATE the action space includes making a discrete selection over the shape to drop and the x, y coordinates to drop the shape above the scene.

Rewards::

The agent receives a reward for the positive difference in height of the tower. A penalty of -0.25 is given for every repeated shape.

Termination: The game is terminated either after the tower of shapes exceeds 3 in height or after 10 shape placements.

Action Dataset: In Shape Stacking the functionality of each action is characterized by the physical appearance of the shape. For this reason, the action dataset consists of images of the shape from various angles and heights (viewpoints). 1,024 images of resolution 84x84 constitute the dataset to represent a single shape.

:



Figures and Tables:

Figure 1: Generalizing the knowledge of solving a task to a new set of actions. (a) CREATE is a sequential environment where the task is to help the green ball reach the goal (blue) by selecting tools and deciding where to place them. (b) In Shape Stacking the goal is to stack a tall tower by selecting the right shapes and their placements. Scenario A depicts the training scenario when the agent learns to utilize a given set of actions to solve the task. Scenario B presents an unseen set of actions to the agent which is expected to generalize to solve the task zero-shot.

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. (

Figure 2: Framework for generalization to unseen actions. (1) Action datasets for all training actions are used to train a Hierarchical VAE (HVAE) model. (2) The action encoder embeds each dataset to define the approximate posterior q a (c|D) over action latents c. (3) The instance encoder q s (z|x, c) encodes each data sample x, while conditioned on the action latent c, into a distribution over instance latents z. (4) The decoder p(x|z, c) reconstructs the action sample x based on the action embedding c and sample latent z. (5) The policy π takes current state s t and the inferred action embeddings c i for each of the given actions and produces a categorical distribution to represent the policy. Similar flow occurs at inference, when new actions and their datasets are given.

Figure 3: Quantitative results: displayed are 3 of the CREATE tasks, the Block Stacking task, the Recommender task and the Grid World task. The performance displayed is measured on generalization to the test set of actions across 3,200 episodes. All results are averaged across 6 seeds. The legend describes ablations of our method (shades of red), embedding baselines (shades of blue), policy architecture baselines (shades of green), and alternate modalities in learning embeddings (yellow).

Figure 4: Varying difficulty of test action space: (i) Each test action is at least a specific angle apart from all actions seen during training (ii) Each test action is at least a specific distance in embedding space apart from all actions seen during training (iii) Test set contains seen/unseen ratio

ANALYSIS OF TRAINING PROCEDURE AND MODELS B.1 VISUALIZATION OF HIERARCHICAL EMBEDDING SPACES (a) CREATE

Figure 6: T-SNE Visualization of learned embedding space for CREATE environment. Tools in CREATE are labeled by their properties which define them to be floor, trampolines, high-frictional, etc. The action embeddings clearly group similar actions together.

Figure 8: T-SNE Visualization of learned embedding space for Shape-stacking environment.

Figure 11: GRID WORLD environment. The agent is the red triangle and the goal is the green cell. The environment contains one row or column lava wall with a single opening. Each action consists of 6 consecutive moves in 4 directions.

Figure 12: CREATE environment. In CREATE, the green ball is falling into the scene, which must push the red target ball into the blue goal location. The top and bottom rows show actual evaluation results when our model is tested on CREATE UP and CREATE DOWN, respectively.

Set of known actions A K = {a 1 , . . . a N } , associated action datasets {D 1 , . . . D N }, number of clusters k and number of clusters to sample from S Initialize HVAE model parameters for epoch = 1, 2, . . . do Sample a i ∼ A K Compute loss in Eq. 1 for D i Update HVAE model with gradient of Eq. 1 end for Initialize empty set C K = {} Infer action representations c ai = q µ a (c|D i ), ∀a i ∈ A K {µ denotes mean of distribution} Store action representations C K ← {c ai . . . c a N } Initialize policy parameters θ Compute k-means clustering on C K : G k ← {g 1 , . . . , g k } t ← 0 Receive initial state s 0 from ENV and initialize replay buffer while not done do for step in rollout buffer do Sample clusters G ⊂ G K Sample action set A g from clusters G as in Eq. 7 Compute f θ (s t , c aj ), ∀a j ∈ A g Sample action a t ∼ π θ (s) using Eq. 2 s t+1 , r t ← ENV(s t , a t ) Store experience (s t , a t , s t+1 , r t ) in replay buffer Set of test actions A T = {a N +1 , . . . a N +M } and action datasets {D N +1 , . . . D N +M } Infer action representations c ai = q µ a (c|D i ), ∀a i ∈ A T while not done do Compute f θ (s t , c ai ), ∀a i ∈ A T Sample action a t ∼ π θ (s) using Eq. 2 s t+1 , r t ← ENV(s t , a t ) t ← t + 1 end while HVAE implementation is based on the PyTorch implementation of Neural Stastician

Table content: <table><tr><td>Algorithm 1 Training</td></tr><tr><td>Require: t ← t + 1</td></tr><tr><td>end for</td></tr><tr><td>Update θ based with PPO using rollout buffer</td></tr><tr><td>end while</td></tr><tr><td>Algorithm 2 Testing</td></tr><tr><td>Require:</td></tr></table>

Peer review (iclr2020):
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. In general, this paper is written very well, with some minor suggested tweaks in the following paragraphs. The key strategies used here (HVAE to identify contexts, ERM as the objective, entropy regularisation, etc) all make sense, and are shown to work well in the experiments carried out. While the experiments are sufficiently varied, it worries me that only 3 or 2 seeds were used. In some cases, such as NN and VAE in the CREATE experiments show large variances in performance. Perfects a few more seeds would have been nice to see. This is the key reason why I chose a *Weak Accept* instead of an *Accept*. Some of the results (the latent spaces) shown in the appendix are very interesting too, particularly since they show how similar actions spaces cluster together in most cases. Minor issues: 1) In Figure 3, I am not clear about what *im* and *gt* settings are. 2) In Figure 3, it would have been nice to have consistent colors for the different settings. 3) It would have been nice to see the pseudocode of the algorithm used.