Jointly-Trained State-Action Embedding for Efficient Reinforcement Learning

1 INTRODUCTION . Reinforcement learning ( RL ) has been successfully applied to a range of tasks , including challenging gaming scenarios ( Mnih et al. , 2015 ) . However , the application of RL in many real-world domains is often hindered by the large number of possible states and actions these settings present . For instance , resource management in computing clusters ( Mao et al. , 2016 ; Evans & Gao , 2016 ) , portfolio management ( Jiang et al. , 2017 ) , and recommender systems ( Lei & Li , 2019 ; Liu et al. , 2018 ) all suffer from extremely large state/action spaces , thus challenging to be tackled by RL . In this work , we investigate efficient training of reinforcement learning agents in the presence of large state-action spaces , aiming to improve the applicability of RL to real-world domains . Previous work attempting to address this challenge has explored the idea of learning representations ( embeddings ) for states or actions . Specifically , for state embeddings , using machine learning to obtain meaningful features from raw state representations is a common practice in RL , e.g . through the use of convolutional neural networks for image input ( Mnih et al. , 2013 ) . Previous works such as by Ha & Schmidhuber ( 2018b ) have explored the use of environment models , termed world models , to learn abstract state representations , and several pieces of literature explore state aggregation using bisimulation metrics ( Castro , 2020 ) . While for action embeddings , the most recent works by Tennenholtz & Mannor ( 2019 ) and Chandak et al . ( 2019 ) propose methods for learning embeddings for discrete actions that can be directly used by an RL agent and improve generalization over actions . However , these works consider the state representation and action representation as isolated tasks , which ignore the underlying relationships between them . In this regard , we take a different approach and propose to jointly learn embeddings for states and actions , aiming for better generalization over both states and actions in their respective embedding spaces . To this end , we propose an architecture consisting of two models : a model of the environment that is used to generate state and action representations and a model-free RL agent that learns a policy using the embedded states and actions . By using these two models , our approach combines aspects of model-based and model-free reinforcement learning and effectively bridges the gap between both approaches . In contrast to model-based RL , however , we do not use the environment model for planning , but to learn state and action embeddings . One key benefit of this approach is that state and action representations can be learned in a supervised manner , which greatly improves sampling efficiency and potentially enables their use for transfer learning . In sum , our key contributions are : • We formulate an embedding model for states and actions , along with an internal policy πi that leverages the learned state/action embeddings , as well as the corresponding overall policy πo . We show the existence of an overall policy πo that achieves optimality in the original problem domain . • We further prove the equivalence between updates to the internal policy πi acting in embedding space and updates to the overall policy πo . • We present a supervised learning algorithm for the proposed embedding model that can be combined with any policy gradient based RL algorithms . • We evaluate our methodology on some game-based as well as real-world tasks and find that it outperforms state-of-the-art models in both discrete/continuous domains . The remainder of this paper is structured as follows : In Section 2 , we provide some background on RL . We then give an overview of related work in Section 3 , before presenting our proposed methodology in Section 4 , which we evaluate in Section 5 . 2 BACKGROUND . We consider an agent interacting with its environment over discrete time steps , where the environment is modelled as a discrete-time Markov decision process ( MDP ) , defined by a tuple ( S , A , T , R , γ ) . Specifically , S and A are the sets of all possible states and actions , referred to as the state space and action space , respectively . In this work , we consider both discrete and continuous state and action spaces . The transition function from one state to another , given an action , is T : S × A 7→ S , which may be deterministic or stochastic . The agent receives a reward at each time step defined by R : S × A 7→ R. γ ∈ [ 0 , 1 ] denotes the reward discounting parameter . The state , action , and reward at time t ∈ { 0 , 1 , 2 ... } are denoted by the random variables St , At , and Rt . The initial state of the environment comes from an initial state distribution d0 . Thereafter , the agent follows a policy π , defined as a conditional distribution over actions given states , i.e. , π ( a|s ) = P ( At = a|St = s ) . The goal of the reinforcement learning agent is to find an optimal policy π∗ that maximizes the expected sum of discounted accumulated future rewards for a given environment , i.e. , π∗ ∈ arg maxπ E [ ∑∞ t=0 γ tRt|π ] . For any policy , we also define the state value function vπ ( s ) = E [ ∑∞ k=0 γ kRt+k|π , St = s ] and the state-action value function Qπ ( s , a ) = E [ ∑∞ k=0 γ kRt+k|π , St = s , At = a ] . 3 RELATED WORK . For the application of state embeddings in reinforcement learning , there are two dominant strands of research , namely world models and state aggregation using bisimulation metrics . World model approaches train an environment model in a supervised fashion from experience collected in the environment , which is then used to generate compressed state representations ( Ha & Schmidhuber , 2018a ) or to train an agent using the learned world model ( Ha & Schmidhuber , 2018b ; Schmidhuber , 2015 ) . Further applications of world models , e.g . for Atari 2000 domains , show that abstract state representations learned via world models can substantially improve sample efficiency ( Kaiser et al. , 2019 ; Hafner et al. , 2020 ) . Similar to this idea , Munk et al . ( 2016 ) pre-train an environment model and use it to provide state representations for an RL agent . Furthermore , de Bruin et al . ( 2018 ) investigate additional learning objectives to learn state representations , and Francois-Lavet et al . ( 2019 ) propose the use of an environment model to generate abstract state representations ; their learned state representations are then used by a Q-learning agent . By using a learned model of the environment to generate abstract states , these approaches capture structure in the state space and reduce the dimensionality of its representation – an idea similar to our proposed embedding model . In contrast to world models , we considers both states and actions , rather than just states . Additionally , our approach trains an RL agent on the original environment using its embedded representation , rather than learning a policy using a surrogate world model . Bisimulation , on the other hand , is a method for aggregating states that are “ behaviorally equivalent ” ( Li et al. , 2006 ) and can improve convergence speed by grouping similar states into abstract states . Several works on bisimulationbased state aggregation , e.g . Ferns et al . ( 2004 ) ; Givan et al . ( 2003 ) , present different metrics to measure state similarity . Furthermore , Zhang et al . ( 2020 ) and Castro ( 2020 ) propose deep learning methods for generating bisimulation-based state aggregations that scale beyond the tabular methods proposed in several earlier works ( Castro , 2020 ) . While there are parallels between bisimulation and our approach , we do not propose the aggregation of states . Instead , our embedding technique projects states into a continuous state embedding space , similar to Zhang et al . ( 2020 ) , where their behavioral similarity is captured in their proximity in embedding space . Furthermore , our method embeds both states and actions and does not employ an explicit similarity metric such as a bisimulation metric , but instead learns the relationships among different states and actions via an environment model . State representations are also used in RL-based NLP tasks , such as Narasimhan et al . ( 2015 ) , who jointly train an LSTM-based state representation module and a DQN agent , and Ammanabrolu & Riedl ( 2018 ) , who propose the use of a knowledge graph based state representation . In addition to state representations , previous work has explored the use of additional models to learn meaningful action representations . In particular , Van Hasselt & Wiering ( 2009 ) investigate the use of a continuous actor in a policy gradient based approach to solve discrete action MDPs , where the policy is learned in continuous space and actions are discretized before execution in the environment . Dulac-Arnold et al . ( 2015 ) propose a similar methodology , where a policy is learned in continuous action embedding space and then mapped to discrete actions in the original problem domain . Both Van Hasselt & Wiering ( 2009 ) and Dulac-Arnold et al . ( 2015 ) only consider actions and assume that embeddings are known a priori . Tennenholtz & Mannor ( 2019 ) propose a methodology called Act2Vec , where they introduce an embedding model similar to the Skip-Gram model ( Mikolov et al. , 2013 ) that is trained using data from expert demonstrations and then combined with a DQN agent . One significant drawback of this approach is that information on the semantics of actions has to be injected via expert demonstration and is not learned automatically . In contrast , Chandak et al . ( 2019 ) propose a method that enables the self-supervised learning of action representations from experience collected by the RL agent , using an embedding model that resembles CBOW . The approach presented in ( Chandak et al. , 2019 ) is most closely related to our method in the way of integrating embeddings with policy gradient based RL algorithms . However , they only consider action embeddings , while we jointly embed states and actions and improve the performance significantly . As an alternative to embedding methods , Pazis & Parr ( 2011 ) and Sallans & Hinton ( 2004 ) represent actions in binary format , where they learn a value function associated with each bit , and Sharma et al . ( 2017 ) propose the use of factored actions , where actions are expressed using underlying primitive actions . Again , these three methods rely on handcrafted action decomposition . Previous embedding approaches consider either states or actions in isolation , embedding only one of the two or using separate models for each . By contrast , in addition to capturing structure in the state or action space respectively , we leverage interdependencies between the two by jointly learning embeddings for them , and thus improve the quality of embeddings . 4 METHODOLOGY . We propose an embedding model to jointly learn state and action embeddings and present a generic framework that uses the learned embeddings in conjunction with any policy gradient algorithms .

Summary: The paper proposes a method to jointly learn: (a) a latent state embedding; (b) a latent action embedding; (c) a state transition model; and (d) an RL policy.  The latent models should allow for better generalization over states and actions, and therefore result in improved learning, particularly for discrete action domains. The method shows improved performance over vanilla policy gradient on a grid-world task, a slot machine task, a recommender system, and half-cheetah locomotion.