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 presents a novel framework for learning action metrics from offline data called Behavioral Metric of Actions (BMA).  The goal of this metric is to position the actions in a latent space that is pertinent relative to the task being learned.  This is particularly useful in environments with a large number of actions where exhaustive exploration is unfeasible, and therefore generalization over actions is necessary.  BMA has two design principles: behavioral similarity between actions, and data-distributional.  Behavioral reflects the fact that two actions that have a similar induced transition with a similar reward should be considered relatively equivalent.  The data-distributional relationship separates actions that appear within the offline dataset from ones that are absent (or at least absent for a particular state).  The authors show that by learning an embedding with these principles they are able to train off-the-shelf algorithms on offline RL tasks in cases where learning would be impossible without an action representation.  They also show that compared to other action representations their approach performs better.