Learning Pseudometric-based Action Representations for Offline Reinforcement Learning

1 INTRODUCTION . Reinforcement learning ( RL ) approaches have been applied successfully in many decision-making tasks . In a conventional setting , RL agents learn policies through an online fashion , where they collect trial-and-error experiences directly from the environment to improve the current policy . This can be done when the environment can be easily established , and the cost of deploying new policies online is low ( e.g. , computer games ) . However , many real-world scenarios allow only offline access to the environment due to cost or safety concerns ( Thomas , 2015 ) , meaning that the algorithm can not access the environment during the training phase . This setting is known as offline RL , or batch RL ( Levine et al. , 2020 ) . Despite the success that offline RL has achieved in robotic control tasks ( Fujimoto et al. , 2019b ) and some Atari games ( Gulcehre et al. , 2021 ) , the applicability of offline RL is still limited in many practical scenarios where the action space is large and discrete , including recommender systems and dialogue systems . Prior works on offline RL mainly focus on the setting that the action space is continuous or only consists of a few discrete actions . To address the issue of overestimating the values of o.o.d . actions , they usually constrain the learned policy to stay close to the data-generating policies ( Fujimoto et al. , 2019b ; Kumar et al. , 2019 ; Wu et al. , 2019 ; Kumar et al. , 2020 ; Kostrikov et al. , 2021 ; Zhou et al. , 2020 ) . However , this family of algorithms might suffer from poor performance when there are a large number of discrete actions . Firstly , the value function hardly generalizes over the entire action space without proper action representations because the actions are isolated in nature . By contrast , actions can be naturally represented in continuous control tasks by their physical semantics ( e.g. , directions , forces ) ; thus , the value functions could generalize well . Secondly , with the size of the action space increasing , the state-action pairs are more sparse to the entire state-action space , resulting in a large proportion of o.o.d . actions . Consequently , the learned policies would be overly restrictive – they are constrained to only select actions within the support of in-distribution actions , without considering a large number of other o.o.d . actions that might contain the optimal action ( Zhou et al. , 2020 ; Kumar et al. , 2019 ) . In this paper , we propose a novel framework for learning Behavioral Metric of Actions ( BMA ) to accelerate the offline RL tasks with large discrete action spaces . In BMA , the policy of an agent is trained with a latent , state-conditional action space , where the action representations are learned under a pseudometric that reflects their behavioral and data-distributional similarities . The behavioral metric , inspired by the bisimulation metric defined on states ( Zhang et al. , 2021 ) , aiming to explicitly quantify how similar that two actions ’ effects are on the environment . We also define a data-distributional metric , aiming to quantify how an action deviates from those actions in the dataset . Then , we propose an architecture of action encoder to learn action representations in a self-supervised way , where ` 1 distances between representations correspond to the defined metric . The learned action representations can be combined with any offline RL algorithms and improve them in two ways . On the one hand , benefiting from the behavioral similarity between actions , the Q-function would be more smooth on the action space and , therefore , easier to converge . On the other hand , the quantified distributional distance enables the algorithm to balance between penalizing and exploring the o.o.d . actions . We theoretically prove the continuity and the bound of the Q-values using our proposed action representations . We also empirically evaluate our methods on two simulated tasks and two real-world applications with large action sets . Experimental results show that policies trained under the BMA framework significantly outperform those trained using original actions . Our pseudometric-based action representations are substantially more effective to offline RL tasks when compared to prior action representations widely used in online RL tasks . 2 RELATED WORKS . Offline RL Offline RL aims to learn policies from logged transition data without any active data collection ( Levine et al. , 2020 ) . Recently , many works have been proposed in this area ( Kumar et al. , 2019 ; Wu et al. , 2019 ; Kumar et al. , 2020 ; Kostrikov et al. , 2021 ; Fujimoto et al. , 2019b ) . They empirically and theoretically reveal that the critical issue in offline RL is the overestimation error induced by o.o.d . actions . As a result , they propose a variety of behavioral regularizations in RL training that compel the learned policy to stay close to the offline data . These regularizations consist of incorporating some divergence regularization into the critic ( Kumar et al. , 2020 ; Nachum et al. , 2019 ) , policy divergence penalties ( Kumar et al. , 2019 ; Wu et al. , 2019 ) , and appropriate network initializations ( Matsushima et al. , 2021 ) . These works show effective results in simulated tasks , like Mujoco ( Todorov et al. , 2012 ) and Atari games ( Gulcehre et al. , 2021 ) . However , we empirically show that they fail in the real-world applications with large discrete actions , likely due to the difficulty of generalization over large action sets or the excessive restriction induced by a large proportion of o.o.d . actions . Action representations in RL In online RL , continuous action representations are usually used to exploit underlying structures of large discrete action spaces , thereby accelerating policy optimization in large-scale real-world tasks , like recommender systems ( Ie et al. , 2019 ) and dialogue systems ( Lee et al. , 2019 ) . In prior works , Lee et al . ( 2018 ) uses predefined action embeddings to represent discrete actions and utilize continuous policy gradients for policy optimization . Chandak et al . ( 2019 ) avoid predefined embeddings by linking action representations to state transitions . Tennenholtz and Mannor ( 2019 ) regard action trajectories as natural languages and thus learn action representations from trajectories of expert demonstrations . Wang et al . ( 2021 ) learn action representations that focus on accurate reconstruction of rewards and next observations . Though these works show good results in online RL tasks , we empirically show that their performance is extremely unstable when combined with offline RL algorithms . The reason might be that they build statistical-based representations from raw data and thus struggle in the data-poor problem in the offline settings , especially when the action spaces are large . Furthermore , they lack an explicit scheme to regularize distances between action representations , which we find are important to the behavioral regularizations of offline RL algorithms . Our method utilizes a relation network ( Santoro et al. , 2017 ) to enforce the distances between action representations to reflect their behavioral and data-distributional relations . This provides extra information for building action representations when the offline data is sparse in the entire state-action space and helps offline RL algorithms to derive more proper behavioral regularizations . Some other works use action representations as pre-training behavioral priors or primitives ( Zhou et al. , 2020 ; Ajay et al. , 2021 ; Singh et al. , 2021 ) . They train policies on latent action spaces which are learned from past successful trials and thus avoid a serious distributional shift in offline RL . However , these works seriously rely on the offline dataset ’ s quality ; experience data with poor performance heavily limits the policy improvement . On the other hand , they focus on continuous control tasks without any specific design to address large discrete action spaces . Metrics in RL A crucial principle to generalization in reinforcement learning is to assign similar predictions to similar states . A standard implementation is to use the similarity in an adaptive fashion and group states into clusters while preserving some desired properties . The fundamental assumption behind it is the existence of a metric characterizing the real-valued distance between states ( Le Lan et al. , 2021 ) . A related concept is bisimulation metrics that measure how “ behaviorally similar '' states are Ferns et al . ( 2004 ) . Recently , Zhang et al . ( 2021 ) extend this concept by proposing a gradient-based method for learning a representation space with the properties of bisimulation metrics . In the field of offline RL , Dadashi et al . ( 2021 ) propose a pseudometric to measure the closeness of state-action pairs to the support of logged transitions , but does so directly , without learning a representation . These works mainly focus on metrics on states or state-action pairs but ignore relations between actions . We believe that our work is the first to define a behavioral and data-distributional metric on actions and show that action representations based on this metric are beneficial to offline policy improvements . 3 PRELIMINARIES . We start by introducing notation in this work . We consider the underlying environment as a Markov decision process ( MDP ) with a discrete action space , represented by the tupleM = ( S , A , P , R , γ ) . Here , S is the state space . A is a finite set of actions , called the action set , and |A| denotes the size of the action set . P andR are the transition function and the reward function , respectively , indicating that when the agent takes the action a ∈ A under the state s ∈ S , the probability of transitioning from state s to state s′ ∈ S is P ( s′|s , a ) , and there is a environmental rewardR ( s , a ) ∈ R. The goal of the agent is to learn a good policy a ∼ π ( s ) that maximizes the expected cumulative discounted rewards : EP [ ∑∞ t=1 [ γ tR ( st , π ( st ) ] ] . In online RL , the agent usually learns policy from its interacting experience with the current environment . By contrast , our work focus on the offline setting , in which the agent can not collect new experience data and learns policy from a static dataset D = { ( s , a , s′ , r ) } generated by some other policy . We call the policy that generates D the behavioral policy and denote it as πβ ( a|s ) . 4 OFFLINE RL WITH BMA . In this section , we elaborate on BMA , our proposed method for leveraging pseudometric-based action representations to address offline RL tasks with large discrete action spaces . We begin by describing the paradigm of how to train and execute policies according to the learned action representations . Next , we discuss which properties of action representations are crucial to the performance of this framework and define a pseudometric function to measure these properties . Then , we propose an action encoder structure to learn action representations complying with the desired properties from the offline dataset . Finally , we give a theoretical analysis of the generalization ability and the bound of the value function based on the learned action representation space . Overview . We first introduce a two-phase paradigm for solving such tasks . It first learns an action representation space from the experience dataset by a self-supervised learning framework , and then , train an internal policy πi on the action representation space . This internal policy can build upon arbitrary offline RL algorithms with continuous control . In detail , we first train an action eocoder φ for generating action representations and then convert the original dataset to a new datasetDe = { ( s , e , s′ , r ) } , where e = φ ( a ; s ) is the state-conditional action representation and e ∈ E . Note that we introduce how to train φ in next subsections . Then , the applied offline RL algorithm would learn an interal policy πi ( ê|s ) from De . It provides a latent action ê ∈ E for a given state , but ê would likely not be a valid action , i.e. , it does not equal any action representation e of a ∈ A . Therefore , we need to map from ê to an element in A . Here , we adopt a simple nearest neighbor lookup g ( ê ) introduced in Lee et al . ( 2018 ) : g ( ê ) = argmin a∈A ‖ê− φ ( a ; s ) ‖1 ( 1 ) where g is a mapping function from the continuous representation space to the discrete action set . It returns the original action whose representation is the closest to ê by ` 1 distance . Therefore , the overall policy πo = g ( πi ( e|s ) ) . Its scheme is also described by Fig.1 ( a ) and Alg.1 in the appendix .

Summary: This paper studies the action representation learning for offline RL with large discrete action spaces. It proposes a new metric to capture the similarity between actions, called Behavioral Metrics of Actions. The new metric has two components, one is a behavioral metric that groups actions that share the similar MDP property (i.e., reward and transition model), and the other offline specific one is distributional metric, which uses a binary variable to capture whether the two actions coming from the same data distribution. Empirical studies show superior performance over the baselines, and the ablation study demonstrates the necessity of all of the components in the metric.  RL with large discrete action spaces is of interest in some real-world applications, such as recommender systems and dialogue systems, this paper gives a way to do the action abstraction to further reduce the action space and hence increase the effective sample size in any downstream learning task.