Learning Pseudometric-based Action Representations for Offline Reinforcement Learning
Authors: Pengjie Gu
Mengchen Zhao
Chen Chen
Dong Li
Jianye Hao
Bo An
Venue: ICLR 2022 Submitted (2022)
Keywords: offline reinforcement learning，representation learning， metric learning

Abstract: Offline reinforcement learning is a promising approach for practical applications since it does not require interactions with real-world environments. However, existing offline RL methods only work well in environments with continuous or small discrete action spaces. In environments with large and discrete action spaces, such as recommender systems and dialogue systems, the performance of existing methods decreases drastically because they suffer from inaccurate value estimation for a large proportion of o.o.d. actions. While recent works have demonstrated that online RL benefits from incorporating semantic information in action representations, unfortunately, they fail to learn reasonable relative distances between action representations, which is key to offline RL to reduce the influence of out-of-distribution (o.o.d.) actions. This paper proposes an action representation learning framework for offline RL based on a pseudometric, which measures both the behavioral relation and the data-distributional relation between actions.  We provide theoretical analysis on the continuity and the bounds of the expected Q-values using the learned action representations. Experimental results show that our methods significantly improve the performance of two typical offline RL methods in environments with large and discrete action spaces. 

Under review as conference paper at ICLR 2022
LEARNING PSEUDOMETRIC-BASEN ACTION REPRESENTA- TIONS FOR OFFLINE REINFORCEMENT LEARNING
Anonymous authors Paper under double-biee blind review
ABSTRACT
Offline reinforcement learning promising approach for practical applications since does not require interactions with real-world environments. However, existing offline RI. methods only work well environments with continuous or small discrete action spaces. In environments with large and discrete action spaces, such recommender systems and dialogue systems, the performance of existing methods decreases drastically because they suffer from inaceurate value estimation for large proportion of actions. While recent works have demonstrated that online RL benefits from incorporating semantic information in action representations, unfortunately, they fail learn reasonable relative distances between action representations, which is key to offiin RI. to reduce the influence of o.o.d. actions. This paper proposes an action representation learning framework for offline RL based on a pseudometric, which measures both the behavioral relation and the iata-diitributional relation between actions. We provide theoretical analysis on the continuity and the bounds of the expected Q-values using the learned action representations. Experimental results show that our methods significantly improve the performance of two typical offline RL. methods in environments with large and discrete action spaces.
(o.o.d.)
INTRODUCTION
Reinforcement learning (RL) approaches have been applied successfully in many decision-maie tasks. In a conventional setting, RL agents learn policies through an online fashion, where they collect trial-and-error experiences directly from the environment 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 (c.g., computer games). However, many real-world scenarios allow only offline access the environment due to cost safety concerns (Thomas) 2015), meaning that the algorithm cannot access the environment during the training phase. This setting is known as offline RI, batch RL (Levine et al. 2020). Despite the success that offline RI has achieved in
robotic control tasks (Fujimoto et al., 2011b) 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 dat--geneg policies (Fujimoto et al. 20196] Kumar al. 2019: Wu et al. 2019: Kumar et al. [2000; Kostrikov et al.) 2021 Zhou et al. 2020). However, this family of algorithms might suffer from poor performance when there are 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 semanties (e.g., directions, forces); thus, the value functions could generalize well. Secondly, with the size of the action space increasing, the state- aiat-acttion 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 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 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 latent, state-conditional action space, where the action representations are learned under a pseudometric that reflects

---

Under review as conference paper at ICLR 2022
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 iitta-distibuttional metric, aiming to quantify how action deviates from those actions in the dataset. Then, we propose an architecture of action encoder to learn action representations in ceff-supeviiee way, where & 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, casier 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-bar effective to offline RL tasks when compared to prior action representations widely used in online RL tasks
action representations are substantially more
2 RELATED WORKS
Offline RL Offline RI. aims to learn policies from logged transition data without any active data collection
(Levine 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 the overestimation error induced by o.o.d. actions. As 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.) 202T). These works show effective results simulated tasks, like Mujoco (Todorov et al. 2012) and Atari games (Gulcehre et al.) 202T). 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, [ee 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 natural languages and thus learn action representations from trajectories of expert demonstrations. |Wang et al. (202T) Icarn 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-bast 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., (017)) to enforce the distances between action representations reflect their behavioral and data-distrilutionaa relations. This provides extra information for building action representations when the offline data 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. 202 Singh et al. 202T). They train policies on latent action spaces which are learned from past successful trials and thus avoid serious distributional shift in offline RL. However, these works seriously rely on the offline dataset's 's quality; experience data with poor performance heavily limits the policy improvement. On the other hand, they focus 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. 202T). A related concept is
2

---

Under review as conference paper at ICLR 2022
bisimulation metrics that measure how "behaviorally similar" states are Ferns ct al. 2004). Recently, Zhang et al. (202T) extend this concept by proposing a gradient-bass based method for learning a representation space with the properties of bisimulation metrics. In the field of offline RL, Dadashi 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 metrics states or state-action pairs but ignore relations between actions. We believe that our work is the first to define a behavioral and data- cataadistributional metric on actions and show that action representations based on this metric are beneficial 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 discrete action space, represented by the tuple . AS, P, P. R) Here, is the state space. A is finite set of actions, called the action set, and |A|] denotes the size of the action set. P and R are the transition function and the reward function, respectively, indicating that when the agent takes the action a E A under the state 8 E S, the probability of transitioning from state s to state s' E S is P(s'|s, environmental reward [s, a) E R. The goal of the agent is to learn good policy a ~ π (T) that maximizes the expected cumulative discounted rewards: Ep Sp R((, (s:))]].
(ss a), and there is a
(())))(()))...
In online RL, the agent usually learns policy from its interacting experience with the current environment. By contrast, our work focus the offline setting, in which the agent cannot collect new experience data and learns policy from a static dataset D = {(s, a, s)rr)} generated by some other policy. We call the policy that generates D the behavioral policy and denote it TT (a(s)..
4 OFFLINE RL WITH BMA
In this section, we elaborate on BMA, our proposed method for leveraging pseudometric-base bbsed action representa- tions 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 the performance of this framework and define pseudometric function 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 self-supervised learning framework, and then, train an internal policy n; 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 new dataset De- = {(8, e, 8', rr], where e = o(a;s) s) is the state-conditional action representation and E. Note that introduce how to train á in next subsections. Then, the applied offline RL algorithm would learn an interal policy n, (éls) from De. It provides latent action è E E for a given state, but è would likely not be a valid action, i.e., it does not equal any action representation of E A. Therefore, we need to map from è to an element in A. Here, we adopt a simple nearest neighbor lookup g(ê) introduced in Lee et )2018):
g(ē) arg min ||ē p(a; 8)]] aEd
(1)
where y is mapping function from the continuous representation space to the discrete action set. It returns the original action whose representation is the closest to by l distance. Therefore, the overall policy To (eT.(e)))) Its scheme is also described by Fig I (a) and Alg. I in the appendix.
4.1 A PSEUDOMETRIC FUNCTION FOR MEASURING RELATIONS BETWREN ACTIONS
In this section, we discuss which properties are important action representations and introduce pseudo- metric function to measure these properties.
How formulate the structure of action representations crucial the performance of our framework. The agent acting in the representation space would generalize the information of an action to other actions with
3

---

Under review conference paper at ICLR 2022
H6 ())
Ue 2..
Action
Pt xay
Transition
Representation space
Gloseit
Dutnset
dotacce J
S
Reward Model
Reabtanc
Encoding Decoding Latent action
Action
A
P(xe)))
Model
a.a.d penalit
Discrete uction set
Sample
Forward
Forward wle gradients
Summation
Figure BMA framework. (a) The overall policy scheme. (b) The architecture of pseudometric-basel action representation learning. The action encoder aims to generate action representations, whose & distances between each other equals the expected pscudometric distance between their corresponding original actions. The generated action embeddings e and eg are both conditioned on the same state s. During training, since it is hard to find two transitions that have the same state but different actions from the dataset, (s, aj, Rg:) ) is sampled from the dataset, and aj is uniformly sampled from the finite action set A. The reward of (s, az) and the transitions of (s, a,) and (s, 0z) are estimated by corresponding models. f is function estimate whether aj an in-distribution action and thus determine whether there is o.o.d. penalty the pseudometric distance. The objective J is the mean square error between le;, eallı and the summation the reward distance, the discounted transition distribution distance, and the o.o.d. penalty.
similar representations. However, if actions with similar representations have different envirommental effects, this generalization might result in poor performance. Furthermore, existing offline RL algorithms usually adopt various approaches for behavioral regularizations to avoid serious overestimation errors; they enforce the learned policy to choose actions close to the in-distribution actions. Therefore, different action spaces induce distinct results of the behavioral regularizations, and an inappropriate representation space would destroy the effect of behavioral regularizations in our paradigm. For example. If the behaviorally different actions, which induce distinct transitions and rewards, are close in the representation space, false result of behavioral regularization would happen. By contrast, the representation space where behaviorally similar actions are clustered helps the agent choose the optimal safe actions. To this end, expect that the learned action representations' relative distances reflect two major relations between any two discrete actions:
(i) The behavioral relation: The distance between any actions in the representation space should reflect the difference between their induced transitions and rewards. This helps the behavioral regularized policy select actions close the in-distribution actions in terms of their effects on the MDPs.
(ii) The data-distributional relation: The distance between any two actions the representation space should reflect whether they are in the same distribution of the experience dataset not. If they not in the distribution, there is a penalty distance between them. This helps offline RI. algorithms tend to choose in- distribution actions, avoiding serious overestimation errors.
To derive the first relation, we first define a distance function d: A A + R>0 measuring the behavioral similarity between two actions.
Definition 1 (Behavioral Metric of Actions) DGiven MDP M, behavioral metric of actions is function d x + R20 such that:
daas Uz)) Eses [R R2 + W2(Po, P.))
where Wg is the 2nd Wasserstein distance between two distributions.
The Wasserstein metric is widely used in prior works (Zhang al., 2021 Ferns al., 2004) for measuring the distance between two probability distributions of states, indicating the cost of transporting mass from one
Note that is pseudometric, meaning that it allows the distance between two different actions to be 0.
4

---

Under review conference paper at ICLR 2022
distribution to another. Here, we adopt 2-Wasserstein metric W2 in Eq as it has a convenient formulation: (2,
W,), (,, E,)) Ips 1yyl2 + E - 12 where
|||= the Frobenius norm.
Exactly computing the behavioral distance between twO actions according Eq generally impractical in the offline setting since the data points of states in the offline dataset are difficult to cover the entire state space, especially when the action space large. Estimating the results of the max expectation over might induce a large variance. Therefore, we extend the definition to conditional version as follows:
Definition 2 (State-conditional Behavioral Metric of Actions) Given an MDP M and a state a state- conditional behavioral metric of actions a function
S x A H Rzo such that:
d(a;, 0g ag1s)
[2 R +
W2 (
This definition avoids estimating the expectation over S and brings a higher representation capability: It allows the actions' relations to be changed under different states.
Data-distribution detection. The major issue in offline RL is the overestimation problem caused by o.o.d. actions. Since o.o.d. actions are data-poor in the dataset, the estimation about relations of them might be inaccurate. This mismatch enables them to be close to those in-distribution actions which have distinct behavioral effects from them. So also expect that there is penalty distance between the o.o.d. actions and the in- distribution actions, forcing the behavioral regularized policy tends to choose in-distribution actions. Thus, we modify the pscudometric function (3 to reflect the citaaddsttrbuutinaal relation between actions:
d(a;, (s)s)) s) = |R R.d | W2(P".. P.))) + P- Lp(a,, asls)
(4)
where p >> is penalty coefficient, [a az (s) is a function indicating the relation of actions: If ag and aj are from two data distributions, I (a,, 0z (s) O. otherwise, Ig (a,, ayls) = 1.
4.2 LEARNING PSEUDOMETRIC BASED ACTION REPRESENTATIONS
This section introduces how to learn action representations complying with the pseudometric d (Eq 4 and gives a theoretical analysis about the generalization ability and the bound of the value function using our action representations.
In detail, we expect to train an action encoder o S x A A E for generating action representations that possess the desired property d(a,, sy]s) ;= [[o(a;; s) o(q,;s)ll It means that the & distance between any two actions in the learned representation space equals their pseudometric distance. Concretely, the learning objective can be written as follows:
( - d(a,, asls))²
(5)
where et = -(a; =f Mayz;) state but different actions a; and aj- However, it is challenging to search for sufficient such pairs from the offline dataset. Thus, we modify the objective by using the estimated samples:
s). 3) In practice this objective requires sampling state-action pairs with the same
(()) F. R, ~D.yy (llet edllı - d(a.,,yls)) asls))
(6)
where
(aa, = R R(s, 2l)l + W2 (22PP(111 et), (-1,, + +p- is(a5ls)
(7)
Specifically, dia;, asls) is an estimation of d(a; as)s). denotes s(a;s) with stop gradients, π and P is the reward model and transition model, which have their training steps. īs(ayls) is also trainable model predict whether aj an o.o.d. action. If a, is out of distribution, fa [s) = 1, otherwise fg(ayls) = In practice, we follow the prior work to derive [a(ss) (Fujimoto et al. 2019a). We train model G(als) Rs (a)ss) predict the probabilities of every action in the dataset and then scale all probabilities by the maximum probability. a, would be predicted as the o.o.d. action if the relative probability is below a threshold constant T.
This procedure can be summarized as fg(a(s)
max G(ā)s) < T
C(als)
5

---

Under review as conference paper at ICLR 2022
During training, we sample the tuple (s, ag, Rg:) from D> and directly sample the other action a, from the finite action set A. Then, other lacked data would be estimated by the trained models and the modified objective can be computed. To a certain extent, this estimation operation would reduce the accuracy of the learned action representations, but our empirical study shows that it is sufficient to derive a latent space with effective representative ability in offline RL tasks. Further details of training the action encoder and all models are described in Fig T and Alg. 2 in the appendix.
Theoretical Analysis. We first theoretically analyze the generalization ability of the value function using our action representations. We prove that the value function of any given policy is Lipschitz with respect to our proposed pseudometric function
Theorem 1 (Q" is Lipschiz with respect to d) Given policy n, let Q" be the value function for given discount factor n. Q" is Lipschitz continuous with respect to d with Lipschitz constant 1-y
1
1 - d(a,, az1s) 1-7
(8)
Proofs in appendix. This theorem means that, the closer two actions are in terms of d, the more likely they are share the similar value. If explicitly force the distance between action representations to comply with d, the value function of the learned policy would be Lipschitz continuous in the action representation space. This continuity brings an effective generalization capability the value function, and thus, intuitively, reduces the estimation errors of o.o.d. actions (Le Lan et al.) 2021]. On the other hand, when the policy is regularized to select actions close to the un-distribution actions in our representation space, it would choose similar actions in terms of the long-term return, avoiding selecting actions with distinct effects. Further, This conculsion can be further represented as: one of o; and aj is the o.o.d. action but the other is not, (QT (s) Q" (s. a;)l S 1-7 (|R(s, az) - a;) + 7 W2(P(ssls: a;). P(s"(s. az) + p), otherwise, T (R(s a.) R(s, + W2(P(s'|s, a.), P(s' (ssss, (s,,))))))))),);,,,))) This indicates that if two actions are from two distributions, the difference between their values tend to have a looser upper bound depending on the constant p. and plays a role in balancing between penalizing and exploring the o.o.d. actions. If p is big, the difference between two values tends to be large, the offline algorithms would ignore most o.o.d. actions. If p is proper and the two actions' behavioral relation is close, the values o.o.d. actions and the values of in-distribution actions would be similar but avoid to be the same. This can encourage the offline policy cautiously explore the effect of o.o.d. actions.
Now, we would show how action representations complying the pseudemetric d improve dowenstream offline RL. Assuming that the BMA action encoder ő(a; s) has a learning error e: If the pseudometric distance between any two action is smaller than e, they would be aggregated together. Building on the policy theoretical analysis from (Ajay et al. 2021) and (Kumar et al. 2020), we now bound the performance of the policy obtained when offline RL is performed with BMA.
Theorem 2 (Performance bound in offline RL.) Let oद (els) be the policy obtained by CQL performing with BMA in the constructed MDP M and Ti.g (als) refer to the overall policy when тद (els) is used together with nearest lookup function g- Let J(r,M) refer to the expected return of a in M and a(a; s) the BMA action encoder, which has a learning error e. Let Ty refer to the behavioral policy genearting D and [[((() = e = :(a; s), a ~ re(als))Then J ( M) Z (TT, k where
1
k 0
7)2
E
snd
(o) M
v VEIDOOL (T) (s) +1
(9)
E
37 (n) +
E
1
I-
Proofs in appendix. This bound suggests that the lower bound over the performance of the learned overall policy depends three factors: 1) the divergence between the learned latent-space policy and the latent behavioral policy (Doql (m' (s)) which is controlled by the applied offline algorithm. 2) The size of the latent space [g]. According to (6) in (Kumar et al., 2020), note that if directly learn an offline policy from the original action space A, the lower bound of offline policy performance would depend on [A]. This indicates that a latent
6

---

Under review as conference paper at ICLR 2022
action space would induce an improvement the worst-case deterioration over the learned policy, since we can set [&] < [.A,, especially when |A] is enormous. 3) The learning error of BMA action encoder c. This term is induced by self-suuervised loss function. Compared with errors induced by offline RL algorithms, this term is much easier to be controlled. In conclusion, BMA provides a therotically gurrantee of policy improvement for offline RL task with large discrete action spaces.
5 EXPERIMENTS
In this section, we empirically show that BMA could be used as a drop-in extension for improving the policy performance in existing offline RL algorithms for problems with large discrete action spaces. Combing BMA with two standard offline RL algorithms, we evaluate our method in two simulated tasks and two real-world problems. All experiments in this paper are carried out with 5 different random seeds, and results are shown with a 95% confidence interval of standard deviation. We believe that the reported results can be further improved by using our framework with other offline RL algorithms; we leave this for future work. Details of architectures and hyperparameters can be found in the appendix.
Evaluation Environments. We consider four environments with large discrete action spaces: maze environ- ment introduced in Chandak et al. (2019). It has 4096 discrete actions. A multi-step maze environment that requires select actions from the 6-step planning decisions. It is modified from the maze environment and also has 4096 actions. A video recommender system that sets 1000 videos as the action set (Ie et al. 2019). A dialogue system that sets 1500 dialogues as the action set. The experience trajectories are policies with sub-optimal performance. They are all trained by online methods. Note that more details of environments and datasets can be found in the appendix.
5.1 EXPERIMENTAL RESULTS
We empirically reveal that optimizing policies on the learned action representation spaces leads to better performance than directly training policies on the original action spaces. We consider two widely-used offline RL algorithms: CQL (Kumar et al., 2020) and BCQ (Fujimoto et al. 2019h). We first learn our pseudometric-base action representations from the given datasets. Then, we utilize our framework to train policies using the two offline RL algorithms. The developed policies are annotated by BMA-CQL and BMA-CQL, respectively. The baselines are the discrete versions of CQL and BCQ and behavior cloning (BC).
The average performance curves are illustrated in Fig (a). We also visualize the average performance of the behavioral policy. Overall, we find that BMA-CQL and BMA-BCQ show faster convergence and better performance than their discrete versions in the four euvironments. Furthermore, discrete CQL and discrete BCQ cannot outperform the behavior policies and show unstable learning curves in some situations. This indicates that in offline RL with large action sets, directly training policies the original action spaces might suffer from serious generalization or extrapolation errors. In addition, the poor performance of BC in most environments indicates that simple imitation learning cannot address this problem well.
Comparison with Other Action Encoders. A natural question is whether training policies on other kinds of action representation spaces can lead similar results? To reply to this question, compare BMA-CQL. against baselines using the same architecture but exchange our action representations with other representations. We first consider two random action representations to examine whether simply projecting discrete actions into continuous action spaces can improve the performance. The first scheme generates random vectors ranging from -1 as action representations (Rand-1). The other scheme also generates random vectors but forces the & distance between any two action embeddings to be larger than a constant. This scheme avoids actions overlapping together (Rand-2). Then, we consider three schemes for learning action representations widely used online tasks: (1) Transition (Chandak et al. 2019): An encoder that captures information about corresponding state transitions into action representations. (2) Reconstruction: An encoder-decoder scheme in which the decoder takes the state and the action representation generated by the encoder as input, and its objective is to predict the next state and reward. (3) External (Lee et al. 2018): Action representations are given by external information. In the maze environment, the action embeddings set as the movement vectors. In the multi-step environment, the actions are represented as the concatenation of the base action at each step of the plan. In the recommendation system, the action embedding the concatenation of video type (expressed as one-hot vector), video quality, and video length. Note that there is no proper external information to represent dialogue actions, so we did not
7

---

Under review conference paper at ICLR 2022
Maze
Multi-step maze
Recommendation system
Dialogue system
100
Total watch time (min)
100
80
0.6
Rvg. retur
60
0.4
40
40
60
20
20
40
0.22
0.0
Gradient
Grad cnt stops
20003
40000
20000 Gradient steps
Gradient steps
20000
40000
BMA-CQL.
Discrete
Discrete
Rehavioral
b
Maze
Multi-step maze
Recommendation system
Dialogue system
100
100
Tocal watch time (min)
1000
75
80
0.0
Aug. return
Avg. return
success rate
50
2.4
60
25
0.2
20
0.0
Gradient
Gradient steps
20000
40000
Gradient steps
20000
40000
Gradienn steps
20000
40000
B.M.A.
Rand-1
Rand-2
Transition
Reconstruction
External
PLAS
Behavloral
Multi-step maze
Recommendation system
100
Total watch time (min)
100
70
80
60
40
50
20
60
40
20000
20000
40000
20000 40000
Gradient sieps
Gradient
Gradient stops
BMA-CQL
BMA-BCQ
BMA-CQL
BMA-BCQ
BMA 0.1
BMA
0,3
BMA
0.5
BMA w/op CYAE
Figure 2: Experimental results. (a) Comparing BMA equipped with BCQ and CQL against directly training offline RL algorithms (Discrete BCQ, Discrete CQL, and BC) on the original action spaces in environments with large action spaces. (b) Comparing the performance of BMA against other widely used action representations (transition-basse representation, reconstruction-base. representation, two kinds of random representations, and external representation in environments with large action spaces. (c) Ablations. We perform two offline algorithms (CQL and BCQ) with BMA two environments (Multi-step maze and Recommendation system). We consider BMA with different penalty distances (BMA with 0.1, 0.3, 0.5), removing the penalty distance from the learning objective (BMA w/o p), and removing the distance learning objective (Eq G) from the learning procedure (CVAE).
consider external action representations in this task. Finally, we consider a similar work for offline RL (PLAS (Zhou et al. 2020)) since it also trains policies on the latent action spaces.
The average performance curves in four environments are illustrated in Fig Two random action representa- tions both fail in the four environments, indicating that simply converting discrete action spaces to continuous action spaces cannot address this task. By comparison, transition and reconstruction representations show better performance in certain situations, revealing that generalization over actions with similar effects can improve the policy performance. However, these two representations also show poor performance in most environments. In contrast, our methods show superior events among all environments, revealing that action representations com- plying with the proposed pseudometric PLAS both show poor performance in all tasks. The reason might be that external representations do not capture the information of dynamics MDPs, resulting in misleading generalization effects. The objective of PLAS is to constrain policies more naturally instead of improving the generalization ability of action representations, so it
suitable to offline RI. In addition, external representations and
8

---

Under review conference paper at ICLR 2022
Online embedding
Rand embedding
Reconstruction embedding
Ours
State
State 2
0.2
0.4
0.6
0.8
Figure Visualization of different representation spaces. Since we set the dimension of representations as 2 in the multi-step maze environment, we directly plot the action representations. The position of each dot in these plots equals the corresponding action embeddings, and the color corresponds to the ground-truth action value predicted by model (Chandak al. (019) with online training. Two rows are in terms of two different states. Each column from left to right corresponds to the representations learned by the online model, the rand-1 1 representations, the reconstruction representations, and ours.
cannot address the tasks with large discrete action spaces well. Fig illustrates the action representations. We sample states from the multi-step maze environment and plot the different action representations. We see that, in offline settings, the rand action representations cannot generalize well since actions with different values are mixed together. The reconstruction-basee representations and our pseudometric-basee representations both show generalization over actions: The changes of action value are continuous in both representation spaces to a certain extent. Furthermore, the reconstruction- based representations are scattered throughout the space while our representations with similar values are arranged as more compact structure. Intuitively, we believe that this is the reason why our method shows faster convergence performance than reconstruction-bant based representations.
A Allaiions We perform several ablations on the multi-step maze environment and recommendation system to try and determine the importance of each component in our framework. Fig (c) illustrate the results. We first consider the effects of different scales of the penalty constant. We set p 0.1, 0.3, 0.5 respectively. The results show that a proper penalty constant would bring better performance, but large penalty constant might make the performance collapse in generalization structure of action representations. Then, we consider removing the penalty constant (denoted BMA w/o p). The results of this ablation show that without the distance penalty mechanism ensure the distribution detection, there performance degradation in our framework. Finally, we remove the objective function (Eq for controlling the distances between actions. As result, the architecture of this scenario is similar to conditional variational auto-encoder that takes current states and actions as input and aims to reconstruct next states and rewards (denoted CVAE) from the action embeddings. This ablation shows worse performance than BMA w/o p in many situations, indicating that controlling the pseudometric distances between actions provides extra information for action representations, leading to better policy performance.
situations. The reason might be that large penalty distances destroy the
5.2 CONCLUSIONS
This paper proposes Behavioral Metrics of Actions: a new action representation learning framework for offline RL tasks with large discrete action spaces. In the learned action representation spaces, distances between representations of discrete actions reflect their behavioral and data-distributional relations. We derive theoretically statements about the benefits of learning policies based on our action representations. Experimental results show that our methods significantly outperform prior works in a variety of environments.

---

Under review conference paper at ICLR 2022
REFERENCES
A. Ajay, A. Kumar, P. Agrawal, S. Levine, and O. Nachum. OPAL: offline primitive discovery for acceler- ating offline reinforcement learning. In ICLR, 2021. URL ottea:///openeevieevveee/rrroro V69LGwJ01IN
Y. Chandak, G. Theocharous, J. Kostas, S. M. Jordan, and P. S. Thomas. Learning action representations for reinforcement learning. In ICML, pages 941-950, 2019.
Dadashi, S. Rezaeifar, N. Vieillard, L. Hussenot, O. Pietquin, and M. Geist. Offline reinforcement learning with pseudometric learning. In ICML, pages 2307-2318, 2021.
N. Ferns, P. Panangaden, and D. Precup. Metrics for finite markov decision processes. In UAI, pages 2052-2062, 2004.
S. Fujimoto, E. Conti, M. Ghavamzadeh, and J. Pineau. Benchmarking batch deep reinforcement learning algorithms. arXiv preprint 1910.0/708, 2019a.
S. Fujimoto, D. Meger, and D. Precup. Off-policy deep reinforcement learning without exploration. In ICML, pages 2052-2062, 2019b.
C. Gulcehre, Z. Wang, A. Novikov, T. L. Paine, S. G. Colmenarejo, K. Zolna, R. Agarwal, J. Merel, D. Mankowitz, C. Paduraru, G. Dulac-Arnold, J. Li, M. Norouzi, M. Hoffman, O. Nachum, G. Tucker, N. Heess, and N. de Freitas. Rl unplugged: A suite of benchmarks for offline reinforcement learning. In NeurIPS, pages 7248-7259, 2021.
E. le, C. wei Hsu, M. Mladenov, V. Jain, S. Narvekar, J. Wang, R. Wu, and C. Boutilier. Recsim: A configurable simulation platform for recommender systems. arXiv preprint 1909, 000470 2019.
I. Kostrikov, Tompson, R. Fergus, and O. Nachum. Offline reinforcement learning with fisher divergence critic regularization. In ICML, pages 5774-5783, 2021.
A. Kumar, J. Fu, G. Tucker, and S. Levine. Stabilizing off- Q-Learning via bootstrapping error reduction. NeurIPS, pages 11761-1177,, 2019.
A. Kumar, A. Zhou, G. Tucker, and S. Levine. Conservative Q-Learning for offline reinforcement learning. In NeurlPS, pages 1179-1191, 2020.
Ce Lan, M. G. Bellemare, and P. S. Castro. Metrics and continuity in reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence, pages 8261-8269, 2021.
K. Lee, S.-A. Kim, J. Choi, and S.-W. Lee. Deep reinforcement learning in continuous action spaces: case study in the game of simulated curling. In ICML, volume 80, pages 2937-2946, 2018.
S. Lee, Q. Zhu, R. Takanobu, X. Li, Y. Zhang, Z. Zhang, J. Li, B. Peng. X. Li, M. Huang, and J. Gao. Convlab: Multi-domain end-to-end dialog system platform. In ACI, pages 64 69, 2019.
S. Levine, A. Kumar, G. Tucker, and Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint 0164,01643, 2020.
T. Matsushima, H. Furuta, Y. Matsuo, O. Nachum, and S. Gu. Deployment-efficcentt reinorrcmet learning via masel-based offline optimization. In ICLR, 2021. URL BRRLAtppss//peeeeeeeeeeee///ooo 3hGNop 4WS
O. Nachum, B. Dai, I. Kostrikov, Y. Chow, L. Li, and D. Schuurmans. Algaedice: Policy gradient from arbitrary experience. arXiv preprint 1912.02074, 2019.
A. Santoro, D. Raposo, D. G. Barrett, M. Malinowski, R. Pascanu, P. Battaglia, and T. Lillicrap. A simple neural network module for relational reasoning. In NewrIPS, 2017. URL WtRLhhtpss:///// //ttps///rooeedingg neurips cc/paper/2017/ (tt.e//ecccccboeeeeeeeeeeeeeeee (0w.66666666008815555.3 Faaaa1153aaalff-aaaeeee Paper.ppe
A. Singh, H. Liu, G. Zhou, A. Yu, N. Rhinehart, and S. Levine. Parrot: Data-driven behavioral priors for reinforcement leamning In ICLR, 2021.
10

---

Under review conference paper at ICLR 2022
G. Tennenholtz and S. Mannor. The natural language of actions. In ICML, pages 6196- 6205,, 2019.
P. S. Thomas. Safe reinforcement learning. Ph.D. diss., University of Massachusetts Libraries, 2015.
E. Todorov, T. Erez, and Y. Tassa. Mujoco: A physics engine for model-based control. In IROS, pages 5026-5033, 2012.
T. Wang, T. Gupta, A. Mahajan, B. Peng, S. Whiteson, and C. Zhang. RODE: Learning roles to decompose multi-agent tasks. In ICLR, 2021.
Y. Wu, G. Tucker, and O. Nachum. Behavior regularized offline reinforcement learning. arXiv preprint 1911.11361, 2019.
A. Zhang, R. McAllister, R. Calandra, Y. Gal, and S. Levine. Learning invariant representattionss for reinforce- ment learning without reconstruction. ICLR, 2021. URL htl.tpsss/peenneeiivveeee/oooooo 2FCWDXRREu
W. Zhou, S. Bajracharya, and D. Held. Plas: Latent action space for offline reinforcement learning. In CoRI, pages 1719-1735, 2020.
A APPENDIX
A.I ADDITIONAL THEOREMS AND PROOFS
Theorem ( is Lipsehiz with respect to d) Given policy =(als), let Q be the value function for given discount factor n, Q" is Lipschitz continuous with respect to d with Lipschitz constant
1 (("s.a))Q((ssaa)) 1- d(a;, az1s)
(10)
Proofs. follows by expanding each Q, rearranging terms and then simplifying the expression.
- ((s),aa)(((ssss)))) (s, = [R(s, + P(s'ls, ((s) (
$ R(s, + Σ
(11)
We assume that R < 1, ge
(T( Σ r" R. 5
2
1
(12)
1-7
Combing the above equations, we have 1Q* az) ("((, R(s,a;) R(s,a;)) +7 (P((s'ss,a)) -P(sssss)))()))) V=(s)
R(s,a;) - R(s,a;)) 1-7
? (P(('ss,, - -
(1 (1-7) (R(s, -((s,)))((ss,))++++++ -1P(:1as10))(((111..))))
1-7
(R(s,a)) R(s, y. (P(s'ls, P(s')s -P(61ssa))P(((5....)))))
1-7
(13)
II practice, there might be some environments in which R E [- Rmaz, Rmaz] and Rmaz > 1. To address this issue, we normalize the reward signals by setting R + R/Rmax
11

---

Under review conference paper at ICLR 2022
Since in Eq. 4. 0 and (g ajls) E {0,1}, get
(s*(s,a)--""(s,. - 1-7 (R(s,a)) - R(s, a(s,aj)]+ E a(s -(ssssssssssss ((s) Is ayls)
$ d(a;,
(14)
According to theorem we can derive the bound between value functions the original MDP and the MDP constructed by the learned BMA action representations. Assuming that the BMA action encoder p(a; s) has a learning error e: the pseudometric distance between any two action smaller than c, they would be aggregated together.
Lemma 1 (Value bound based on Φ) Given the BMA action encoder ((a;8) which has a learning error and maps from actions in the original MDP M to actions in the MDP M constructed by the action representations. The value functions for a given policy n the orginal MDP and its converted version w. the constructed MDP are bounded as:
o(a; 8)
(15)
1-
Proofs. First, us give some notation and assumptions concemiin the form of the constructed MDP M. We assume M = (S,A,, Ps. RR, 7), where S is the same state space, Aş. is the action representation space where the l distances between any two representations corresponds our defined pseudometric distance, and there exists learning error of Φ): If the pseudometric distance between any two action is smaller than c, they would be aggregated together. Recall that the nearest neighbor lookup g is as follow:
g(e) arg min ||e p(a; 8)]l aEA
(16)
As a result, we can get the relationships between the transition probabilities and rewards in the two MDPs, i.e.
Po(s'ss,) e) = P(''|s, Ro(s,e) = R(s,g(e))
(17)
By this way, we can get
QM((, (a
R)( + Σ P(sls,.))VM(() (Ro(s,
+ +7, E(db6la
= +r P((s's, -R( g((aa(((( E (18)
< (R(s.a)-R(s.dlle s))) + I Σ 16d:11.006d(:::))1))
According to the proofs of theorem I we get
(x)()))) - ((, < 1- da M(Ha; s)]{s)
(19)
We have argued that there exists some near-optimal value functions based the constructed MDP M, if of is sufficiently learned. Now, we would show how action representations complying the pseudemetric d improve dowenstream offline RL. Building on the policy theoretical analysis from (Ajay et al. (202T) and (Kumar et al. 2020), we now bound the performance of the policy obtained when offline RL is performed with BMA.
Theorem (Peorrmm(eerfrrmaaac bound in offline RL.) Let गÝ (els) be the policy obtained by COL performing with BMA in the constructed MDP M and कं (als) refer to the overall policy when गद (els) is used together
12

---

Under review conference paper at ICLR 2022
with nearest lookup function 9- Let J M) refer to the expected return of in M and á(a;s) s) is the BMA action encoder, which has a learning error c. Let ag refer to the behavioral policy genearting and x Φ(a; 5), (p(als).Then [Tigg 2 (T k where
k: 0
E (1 7)-
v Vel (B) +1
(20)
soud sndd (a) DDQL [T) I-7
Proofs. First, let's break [J(m M) (TB into
gJ M) - -J (TT M)] $ [ MM) (T;
+ [ M) M)
(21)
+ ( J[T,,
where Th is the converted behavioral policy executed on the constructed MDP M mentioned in the proofs of thereom ??. It can be defined as e ~ F = e = o(a; 8), a ~ Ta(aas). According to theorem 3.6 in (umarr et al. 2020), we apply it to M, we get
J - J MTT
2
Crd
γ
(1-
Pu.
E
and
(o)
3l [D()]
(f ()) +
(22)
7
E
snd
t(t))
(DQQ (o)
k
Then, we try to bound [ g. M) J ( Since the decisions of x and migg are in one-to-one correspon- dence, we get [m M) ; 0.
Now, we try to bound [( - M(T
[J(T, J(TR,MM) = Σ do(s) p (8) - Σ do(s)
8
Σ do(s)
( (8,e) E do(s) Σ (ss.)s
(23)
Σ do(s) Σ Tt(a|s) (s,e) "Edoos do(s) Σ
M
S
S
Σ
do(8)
Ta(a|s)
1 (8, e) - (x)
S
and
According to lemma 1, (s, eg) - QM (8,a))=
-0 M) J(3 1 = k2. Finally, we
get k = k1 + l2. We apply O to get the notation in the theorem.
A.2 ALGORITHMS
A.3 ARCHITECTURE, HYPERPARAMETERS. AND INFRASTRUCTURE
In implementations, the nearest neighbor lookup function g can be regarded as finding the minimum distance between and all actions' embeddings. So, we directly adopted the torch. farchh.min () function in the Py PyTocc python package. Its time complexity is O(N), there would not be serious scale problem when the action set is enormous.
For [s(a)s). We train a model G(als) (als)) predict the probabilities of every action under a given state and then scale all probabilities by the maximum probability, aj would be predicted as the o.o.d. action
13

---

Under review conference paper at ICLR 2022
Algorithm 1: Train policy
Data: Offline Dataset D, the action encoder Φ Any offline RI. algorithm Y, and number of iterations T.
Result: Policy (T),(é(s) for = 0 to I" do
Sample data (s, a, s', ,'r) D;
e t Φ s);
Update policy (é)(( with (8, (,,,,,',7 sr) using the given algorithm Y
end
Algorithm Pseudometric-basee representation learning
Data: Offline Dataset D and number of iterations T. Result: Action encoder
for ==
0 to T do
Sample data (s, 0, R(s, a;)) D;
Sample the other action (j A:
Compute representations e, = o(a; s) and e o(aj;s)
Estimate the transition distributions: P(.(a, et) and P( e!);
Estimate the reward of (s, a.):: (R(s,e)))
Estimate whether aj is o.o.d. action: io(ay Is);
Compute the estimated pseudometric distance: dia,, ayls) Eq
Train encoder: J(Ø) (Ileg egll ddl;, a,1s))
Train transition model: [[P, Φ) = =(P(-1s,e)) - 2¹)2
Train reward model: J(R, o) (R(s, R(s, a; 2:
Train prediction model: J(G) = CrossEntropy G (8), az):
end
14

---

Under review conference paper at ICLR 2022
J6
[(]] s
J0H r
FC, [E] X |Ajurits
FC, [S] units
FC, 1 units
MSE Loss
MSE Loss
5
Tanh
FC, 200 units
FC, 200 units
200 units
FC, [S] units
FC, 1 units
units
FC, [S] +|F] units
FC. +|F] units
FC, 200 units
FC, 200 units
$
(s,e)
(s,e)
BMA action encoder
Transition model
Reward model
FC, 50+[F] units
CrossEntropy
J(G)
Loss
FC, units
FC, [F] units
FC, [A] units
FC, [E] units
FC, |E] units
FC. 200 units
FC,200 units
FC, 200 units
FC, 200 units
FC, 200 units
FC, [S] units
FC. |A] units
FC, |5] units
FC, 2|5| units
FC, |A] units
Reconstruction-ban based action encoder
$
(s,s')
a
Reward model
Transition- baseid action encoder
Figure 4: Details of the network architectures used in our experiments.
if its relative probability below threshold constant т. This procedure can be summarized as īs(a|s) G(als)
G(āls) I T. We have searched over {0, 0.1, 0.3, 0.5, 0.7, 0.9, 1} and found that the best setting is АЕd
T= 0.3
Furthermore, the transition model P. the reward model R, and the distribution detection model G are all implemented by neural networks. Their learning objective can be represented follow;
J(P) ((()s,e)) s'22
J(R) = (R(8, R(s, a:))
(24)
J(G) CrossEEnroopp ((()ss,,
During training, all data are randomly sampled from the offline dataset. We set the batch size as 128 and set the training gradient steps for all models 10000. We control the scale of the learning objective function in all models by controlling the optimization procedure. It is conducted using Adam with a learning rate of 10 -2 and with momentum weight decay. We set the dimension of the action representations = 10, and the penalty coefficient = 0.01, 0.1, 0.1, 0.3 respectively in the maze environment, the multi-step maze environment, the recommender system, and the dialogue system respectively. Details of the neural network architectures used in our experiments are provided in Fig.4 4 in the newest version.
15

---

Under review conference paper at ICLR 2022
Agent
1
Goal
Start position
Experience trajectory
Optimal action
Behavior constraint
Discrete
Representation space
Representation space
Figure 5: An example of the influence of different action spaces on the behavioral regulariziiis of offline RL algorithms.
Online Performance in Multi-step maze
100
75
Avg. Return
50
25
BMA Reconstruction Transition
0
-25
10000 20000 30000 40000 50000 Steps
Figure 6: Online performance of BMA, the transition-based action representations, and the reconstruction-basedd representations in the Multi-step maze environment.
For the applied offline RL algorithms: BCQ (Fujimoto al., 2019b). CQL/Kumar al. 2020) and their discrete versions, we all adopt their open-source implementations released by the authors. For the discrete CQL, We have searched over the crucial hyperparameter a (0.1, 0.3, 0.5, 0.7, 0.9). which determines the extent of conservative estimation of value functions. The best settings for four environments (Maze, Multi-step maze, Recommendation system, Dialogue system) are 0.5, 0.5, 0.5, 0.3, respectively. For the discrete BCQ, we have searched over 0.1, 0.3, 0.5, 0.7, 0.9 for the hyperparameter of the threshold which determines the range of the candidate actions. The best settings for four environments (Maze, Multi-step maze, Recommendation system, Dialogue system) are 0.3, 0.3, 0.3, 0.3, respectively.
A.4 ENVIRONMENTS AND DATASETS
Maze:
We first adopt the maze environment introduced in Chandak al. (2019). The state is the coordinates of the agent's current location, and the agent has n equally spaced actuators (each actuator corresponds to a unit movement in one direction). It choose whether each actuator should be off. The effect of action is the vectorial summation of the movements associated with the selected actuators. Thus, the size of the action space |A| 2n, We n 12, there are 4096 actions in total. There are three reward types: a reward of 10 is given if the agent visits a rewarding region, a reward of -20 is given for visiting a puddle region, and reward
16

---

Under review as conference paper at ICLR 2022
of 100 is given with the episode ending if it is on goal region. The noise was added to the action 10% of the time, and the maximum episode length was 60 steps.
Multi-step maze: Selecting one decision from all possible m-step plans is general problem with large action set. If an environment has actions available at each step, the number of m-step plans would increase exponentially, In this setting: |A| = m'. We implement version of this task on the maze environment mentioned above. The reward settings are the same in this task, but the agent only has 4 base actions: down, right, upward, and left. It needs choose 6-step plan per time steps. Thus, the actual action size in this environment is 4 4096.
Recommender environment: We consider a real-world domain with large action set. The agent needs recommend different videos to the simulated users, and its goal is to extend the user's 's viewing time. We utilize the open-source platform Reesim (Ie et al.) 2019) and adopt the provided scenario called interest t__eeolluu uu in as our simulated environment There are 1000 candidate videos in the action set. All videos are split into 20 types, and each video has its quality and video length. The sampled user has different initial interest values ranging from -1 to on different video types at each episode. The user would decide whether to view the recommended video according to the interest value of the video's ' type. The higher the interest value, the more likely the user is to view the video. The user has a latent attribute called time budget to control whether to continue to view another video. The episode ends when the time budget is smaller than 0. The initial time budget is 100, and it would be deducted by the video length if the user has viewed video. In addition, the viewed video's quality would also influence the user's interest value and time budget: a high-quality video brings an increase, but a low-quality video leads to a decrease. To improve the diversity of the recommended videos, we deduct the time budget by 2.5 if two videos of the same type are recommended successively. The observed state at each time step is the users' interest values and the history of its viewed vidcos' types. A reward equaling the video length is given once one video is viewed.
Dialogue environment: Building dialogue systems for providing information (Young 2006) or improving engagement (Li et al.) 2016) is a real-world challenge. One solution is to employ RL to optimize dialog strategies in multi-turn dialog models. We adopt an open-source platform called convlab (Lee et al. 2019) simulate a dialogue environment in which each action is dialogue. We focus on the MultiWOZ domain (Eric et al., 2019). Its main task is to help a tourist various situations involving multiple sub-domains, such as requesting basic information about attractions and booking hotel room. Specifically, there are 7 sub-domains Attraction, Hospital, Police, Hotel, Restaurant, Taxi, Train. We adopt dialogues in the top 1500 using frequencies as our action set. In detail, we utilize the provided scenario called UserPo wwteroolicyA.. wiiyyyaeedaaaattt..........aaaae tiwo to simulate the tourists' policies and use Mult t Prozstare eettteeeennneee to encode the histories of dialogs into vectorized states.
Since there are no open-source datasets for the offline RL tasks with large discrete action spaces, we collected logged experience trajectories generated from online RL policies. These policies were all developed by the open-source implementation of the work (Chandak et al. (020)), which can be used to address tasks with large discrete action spaces by setting the hyperparameter of the action change number as 1. We trained policies until they achieved sub-optimal performance. Then, policies of checkpoints were used to collect transition data, and the noise was added to the action 50% of the time. In this way, we collected 100000 pieces of transition data in each environment.
Data collection:
A.5 EXAMPLE
To better understand the influence of action spaces to the behavioral regularization of offline algorithms, we give a toy example iluustrated in 5ig.5. The task is to reach the goal from the starting position as fast as possible, and collisions with the wall induce negative reward. The agent needs to learn a policy from the logged suboptimal (blue) trajectory. Assuming that a simple behavioral regularization scenario is applied to the target policy: the agent can only select actions within a particular range near the in-distribution action. At the agent's 's current position, we can find that different action spaces induce different results of the behavioral regularization. The target policy is limited to selecting the sub-optimal in mndistribution action (move upward) in the discrete action space since other actions are isolated from it. By contrast, the policy is permitted to choose actions near the in-distribution action in the other representation spaces where there could be some other actions in the range of behavioral regularization. This is crucial to the policy improvement since the agent has the opportunity to choose the optimal action (move right). However, actions' relative distances also largely influence the results of the behavioral regularization. In representation space 1, behaviorally different actions (the transitions and
17

---

Under review conference paper at ICLR 2022
rewards induced by them are different) grouped together, so the agent is more likely to choose "bad" actions. By contrast, if the relative distances between behaviorally similar actions are close (like in representation space 2), the agent tends to select the optimal other 'safe' actions.
A.6 ONLINE PERFORMANCE BMA
Actually, BMA action representations can be directly applied on the online setting. Utilizzngg the framework mentioned in Chandak et al. (2019) which can incorporating action representaions and online RL algorithms. We get the online performance of BMA, the transion-based representations, and the reconstruction-basee representations in the Multi-step maze environment. The results are illustrated in Fig.6. We can find that, the online setting, three action representations can all achieve fast convergence, while only BMA get good performance the offline setting (b)). This indicates that the BMA more effective in offline RL tasks with large discrete action spaces.
18


Reviews:
[{"review_id": "kJ2ugxE1UUV", "reviewer": "Reviewer_bL2q", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n- I like the setting this paper studies, how to learn efficient action representation, which seems a necessary task when we face large discrete action spaces, which is also the case for some real-world applications. \n- The paper is well-written and it is easy to follow.\n- The experiments are well done. I do like the ablation part, which demonstrates the effectiveness and necessarity of different components.\n\nWeakness:\n- The novelty of this paper seems very limited. For the behavior metric, it just seems like doing some bisimulation metric under the action space, while there are already tons of literature studies on this over the state space. This is a simple and trivial extension to the action space.\n- Regarding the behavior metric, it would be great if you can also the effectiveness in the online setting since this is not a component specific to offline RL.\n- Specific to offline RL, it adds a distributional penalty. Though some theoretical results are given, I do not gain much insight on it regarding how this distributional penalty helps there. Theorem 1 and 2 seem also to hold without this distributional term in the metric. How this helps theoretically, compared with not adding it. \n- Also for Figure 7, is there any explanation that the metric will give more compact clustering than the reconstruction-based method, and why is this the case?\n\n", "summary_of_the_paper": "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.\n\nRL 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.", "summary_of_the_review": "This paper is well-written and the experiments are well-done to illustrate the effectiveness of the proposed metric. However, the novelty of this paper is pretty low, as a bisimulation type metric is well known. Though adding a new distributional component in the offline setting helps, it does not well-justified at all. Given this, I recommend a weak reject for the paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "LC1Um4Hhuw-", "reviewer": "Reviewer_pJJJ", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n  - This paper is well-written to follow. It clearly explains the problem, the motivation, and the proposed method.\n  - It provides theorems on the continuity and the generalization of Q-function.\n  - The proposed algorithm on environments with large discrete action spaces can improve other offline algorithms' performance, and the algorithm using the learned representation space outperforms other representations.\n\nWeakness:\n  - The proposed BMA is a simple extension of the bisimulation metric [1, 2] for offline RL. The pseudo-metric is a bisimulation metric for actions plus an indicator function that outputs one when the two actions are drawn from the identical distributions and 0 otherwise. \n  The minor thing is that the explanation of Theorem 2 is too similar to that of Theorem 2 of [2]. I think that the explanation should be rewritten.\n   \nTypo:\n  - In eq (8), the reward difference term\n  - In the 5.1 Experimental results section on page 8, \"annotated by BMA-CQL and BMA-CQL, respectively.\"\n\n[1] N. Ferns, P. Panangaden, and D. Precup. Metrics for finite Markov decision processes. In UAI, 2004.\n\n[2] A. Zhang, R. McAllister, R. Calandra, Y. Gal, and S. Levine. Learning invariant representations for reinforcement learning without reconstruction.\n", "summary_of_the_paper": "This paper proposed a pseudo-metric between actions, called Behavioral Metric of Actions (BMA), for offline RL in environments with large discrete action spaces. This metric considers both the data-distributional relation and the behavior relation. The distance between two actions drawn from similar distributions with similar rewards and transition probabilities is small. This paper also provides theorems about the continuity and generalization of Q-function using the proposed pseudo-metric. The experimental results show that training a policy in the proposed action representation space has significantly improved performance for offline RL algorithms in environments with large discrete action spaces.", "summary_of_the_review": "This paper has interesting experimental results, but the proposed method and the theory are not novel enough to accept. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "JrIigOM5en2", "reviewer": "Reviewer_hprr", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n\nI think the paper is well-written and easy to follow, the motivation of learning a representation specifically for offline RL learning is very clear. The empirical evaluations show good improvement of the proposed method than other existing ones. \n\nWeakness:\n\nI have the following concerns/questions\n1. Since [1] is also working on using pseudometric for representation as mentioned in Section 2, it would be a more proper baseline to compare with. Besides, I'm particularly interested in this comparison because I think one of the difference between [1] and BMA is that the distance defined in BMA is conditional on the states, while the distance in [1] is defined jointly by states and actions. There has been discussions about the conditional distribution and joint distribution in other areas such as imitation learning [2], so I'm interested to see the comparisons. \n2. The hyparameters of the alternatives in Figure 5 are all default parameters from their prior settings, which make the comparisons less convincing since the tasks are very different. It would be more convincing if a grid search over the hyperparameters as indicated in their original papers can be applied. \n3. It would be also interesting to see how the performances will be changed with respect to the value of gamma since the transition distributions can be challenging to estimate when the action dimension is large. \n\nMinor point: \n\n1. Eq(6) looks confusing because $||e_s^i - e_s^j||$ and $d(a_i,a_j|s)$ are the same by definition.\n2. Will the complexity of eq(1) be a problem when the action space is very large? (e.g., on recommendation system situation, the action space can be million scale)\n\n[1]: R. Dadashi, S. Rezaeifar, N. Vieillard, L. Hussenot, O. Pietquin, and M. Geist. Offline reinforcement learning\nwith pseudometric learning, 2021.\n[2] Ghasemipour, Seyed Kamyar Seyed, Richard Zemel, and Shixiang Gu. \"A divergence minimization perspective on imitation learning methods.\" Conference on Robot Learning. PMLR, 2020.", "summary_of_the_paper": "This paper proposes a framework for learning Behavioral Metrics of Actions that combines both behavioral relation and data-distribution relation between actions to help the learning of offline RL algorithms on large discrete-action tasks. \n", "summary_of_the_review": "In summary, I think the proposed method BMA is interesting and the evaluation is promising. It would be great to add a few experiments as mentioned above to make more fair comparisons and obtain a more convincing conclusion. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "7nGIlVPFtf", "reviewer": "Reviewer_gR99", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Goals\nThis paper sets out to define a new action metric for large action spaces which aids in learning on both on-line and off-line scenarios.  They provide a reasonable approach that is demonstrated to perform empirically.\n\n## Strengths\n* Generalization over states has been one of the main powerful aspects of function-approximated RL, however as actions are often without any direct feature representation they have often been treated as a an arbitrary discrete set.  Providing representations over actions is key to scaling RL up to large systems, and is therefore an important problem to tackle.\n\n* The proposed approach makes good intuitive sense, and the conclusion that the embedding line up to correlate with the value function is an important one, as this has been my personal intuition for a while: as far as the policy is concerned with optimizing the return, two actions are similar iff they have the same expected value.\n\n* The ablations are clear, especially around the value of p which is nice to see.\n\n* The embedding visualizations are nice on the maze, however it's a bit surprising that the embedding from Chanak don't look better as they have figures in their paper which seem more well-distributed for a maze environment.\n\n* Overall the bounds make sense and are intuitive, and looking over the proof the general idea seems good, I am however not the right person to give them a true seal of approval.\n\n## Weaknesses\nThe paper has a couple important issues, more in the form than the function:\n* Many figures are small, with very small text for the legends, pleas try and make the text bigger and consume more surrounding whitespace with your figures as much as possible.\n\n* There is a lot of space spent on presenting the problem, not leaving enough for clear experimental results and discussion.  Perhaps it's because I'm already convinced of the importance of this problem, but I think for example that Figure 2. doesn't bring much to the table and actually confused me more than anything at first.  \n\n* Implementation details are sorely lacking: what is your nearest-neighbour lookup? How do do you actually set up your embedding network? I'm not entirely sure I'd be able to replicate this, please provide more concrete information on the actual configuration of the network and any tricks to get the loss to actually converge.  For example in Algorithm 2 in the appendix you mention estimating the transition distributions, but then you train them afterwards, I'm sure this makes sense somehow but it's not quite clear to me and I would appreciate it being clearer :)\n\n\n## Paper Readability\nOverall the approach is legible but there are some concerns that I would like addressed:\n* There are no descriptions for Figure 5 and Figure 6, please provide a descriptive legend.\n* All figures have very small text and have a lot of surrounding whitespace, they are overall hard to read and I had to zoom in a lot to try and understand them.\n* Overall I would suggest perhaps restructuring Sec. 4.2 some to be easier to follow.\n* One small NIT is that you talk about data-distribution separability, but there is really only one distribution (the expert's), and then you talk about OOD detection.  I believe what you're trying to do is OOD detection through some mechanism, whereas distrubition separability seems like there is perhaps two sources of training data.  Is there a reason not to just call this OOD detection throughout? This would reduce some confusion I had. \n* For the OOD detection can you elaborate in just one sentence what Fujimoto et al.'s approcah is for I_B?  I didn't have time to go through that paper to see what they did and a one-liner would have made things clearer.\n\nThanks again for this work and hopefully we can clear up some of the form-related nits!\n", "summary_of_the_paper": "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.", "summary_of_the_review": "I believe this paper should be accepted on the merit of having:\n* A pertinent problem: action embeddings for large action spaces.\n* A well thought-out approach to tackle this problem: Using a set of reasonable priors on action similarity to provide a general embedding for actions given an existing dataset of behavior on the environment.\n* Convincing experimental results: Both in the on-line and off-line case the approach works well, and is shown to work better than other action embedding schemes.\n\nHowever it has issues in its form:\n* Clearly missing description on Figures 5 and 6.\n* Hard-to-read figures for results.\n* Not always easy to follow explanation of the method.\n\nIf the authors can manage to clear up at least most of the form issues I will support strong acceptance.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a new pseudometric for action representations.  The reviewers generally liked the work and the rebuttal helped to clarity may concerns.  However, the degree of novelty of the approach remains a concern.  In addition, a technical error was discovered by a reviewer in the revised paper.  Hence the paper is not ready for publication.
Decision: Reject