L EARNING P SEUDOMETRIC - BASED A CTION R EPRESENTATIONS FOR O FFLINE R EINFORCEMENT L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Ofﬂine reinforcement learning is a promising approach for practical applications since it does not require interactions with real-world environments.
However, existing ofﬂine 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 out-of-distribution (o.o.d.) actions.
While recent works have demonstrated that online RL beneﬁts from incorporating semantic information in action representations, unfortunately, they fail to learn reasonable relative distances between action representations, which is key to ofﬂine RL to reduce the inﬂuence of o.o.d. actions.
This paper proposes an action representation learning framework for ofﬂine 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 signiﬁcantly improve the performance of two typical ofﬂine RL methods in environments with large and discrete action spaces.
I NTRODUCTION
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 ofﬂine access to the environment due to cost or safety concerns (Thomas, 2015), meaning that the algorithm cannot access the environment during the training phase.
This setting is known as ofﬂine RL, or batch RL (Levine et al., 2020).
Despite the success that ofﬂine RL has achieved in robotic control tasks (Fujimoto et al., 2019b) and some Atari games (Gulcehre et al., 2021), the applicability of ofﬂine 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 ofﬂine 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 ofﬂine 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 reﬂects their behavioral and data-distributional similarities.
The behavioral metric, inspired by the bisimulation metric deﬁned on states (Zhang et al., 2021), aiming to explicitly quantify how similar that two actions’ effects are on the environment.
We also deﬁne 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 (cid:96) 1 distances between representations correspond to the deﬁned metric.
The learned action representations can be combined with any ofﬂine RL algorithms and improve them in two ways.
On the one hand, beneﬁting 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 quantiﬁed 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 signiﬁcantly outperform those trained using original actions.
Our pseudometric-based action representations are substantially more effective to ofﬂine RL tasks when compared to prior action representations widely used in online RL tasks.
R ELATED WORKS
Ofﬂine RL Ofﬂine 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 ofﬂine 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 ofﬂine 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 difﬁculty 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 predeﬁned action embeddings to represent discrete actions and utilize continuous policy gradients for policy optimization.
Chandak et al.
(2019) avoid predeﬁned 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 ofﬂine 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 ofﬂine settings, especially when the action spaces are large.
Furthermore, they lack an explicit scheme to regularize distances between action representations, which we ﬁnd are important to the behavioral regularizations of ofﬂine RL algorithms.
Our method utilizes a relation network (Santoro et al., 2017) to enforce the distances between action representations to reﬂect their behavioral and data-distributional relations.
This provides extra information for building action representations when the ofﬂine data is sparse in the entire state-action space and helps ofﬂine 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 ofﬂine RL.
However, these works seriously rely on the ofﬂine 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 speciﬁc design to address large discrete action spaces.
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.
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 ﬁeld of ofﬂine 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 ﬁrst to deﬁne a behavioral and data-distributional metric on actions and show that action representations based on this metric are beneﬁcial to ofﬂine policy improvements.
P RELIMINARIES
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 tuple M = ( S , A , P , R , γ ) .
Here, S is the state space.
A is a ﬁnite 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 ∈ A under the state s ∈ S , the probability of transitioning from state s to state s (cid:48) ∈ S is P ( s (cid:48) | s, a ) , and there is a environmental reward R ( s, a ) ∈ R .
The goal of the agent is to learn a good policy a ∼ π ( s ) that maximizes the expected cumulative discounted rewards: E P [ (cid:80) ∞ t =
[ γ t R ( s t , π ( s t )]] .
In online RL, the agent usually learns policy from its interacting experience with the current environment.
By contrast, our work focus on the ofﬂine setting, in which the agent cannot collect new experience data and learns policy from a static dataset D = { ( s, a, s (cid:48) , r ) } generated by some other policy.
We call the policy that generates D the behavioral policy and denote it as π β ( a | s ) .
O FFLINE RL WITH BMA
In this section, we elaborate on BMA, our proposed method for leveraging pseudometric-based action representations to address ofﬂine 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 deﬁne 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 ofﬂine 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 ﬁrst introduce a two-phase paradigm for solving such tasks.
It ﬁrst 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 ofﬂine RL algorithms with continuous control.
In detail, we ﬁrst train an action eocoder φ for generating action representations and then convert the original dataset to a new dataset D e = { ( s, e, s (cid:48) , 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 ofﬂine RL algorithm would learn an interal policy π i (ˆ e | s ) from D e .
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.
[Equation]
(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 (cid:96) 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.
4. A PSEUDOMETRIC FUNCTION FOR MEASURING RELATIONS BETWEEN ACTIONS
In this section, we discuss which properties are important to our action representations and introduce a pseudometric function to measure these properties.
How to formulate the structure of action representations is crucial to the performance of our framework.
The agent acting in the representation space would generalize the information of an action to other actions with Dataset s i a j a Action Encoder
Action Encoder  i as  ˆ( | , ) js s e ⋅  ˆ( | , ) is s e ⋅  js e is e  1 distance Transition Model Reward Model Transition Model
W 2 distance
ˆ j i a as s −   o.o.d.
penalty 
ˆ I β
Sample
Forward
Forward w/o gradients
Summation i π
Discrete action set
Representation space
EncodingDecoding
C l o s e s t
Latent action {( s, e, s’, r )} {( s, a, s’, r )}
[Figure] similar representations. However, if actions with similar representations have different environmental effects, this generalization might result in poor performance. Furthermore, existing ofﬂine 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, a false result of behavioral regularization would happen.
By contrast, the representation space where behaviorally similar actions are clustered helps the agent choose the optimal or safe actions.
To this end, we expect that the learned action representations’ relative distances reﬂect two major relations between any two discrete actions:
(i)
The behavioral relation : The distance between any two actions in the representation space should reﬂect the difference between their induced transitions and rewards.
This helps the behavioral regularized policy select actions close to the in-distribution actions in terms of their effects on the MDPs.
(ii)
The data-distributional relation : The distance between any two actions in the representation space should reﬂect whether they are in the same distribution of the experience dataset or not.
If they are not in the same distribution, there is a penalty distance between them.
This helps ofﬂine RL algorithms tend to choose indistribution actions, avoiding serious overestimation errors
To derive the ﬁrst relation, we ﬁrst deﬁne a distance function d : A × A (cid:55)→ R ≥ 0 measuring the behavioral similarity between two actions.
Deﬁnition 1 (Behavioral Metric of Actions) 1 Given an MDP M , a behavioral metric of actions is a function d : A × A (cid:55)→ R ≥
such that:
[Equation]
where W 2 is the 2 nd Wasserstein distance between two distributions.
The Wasserstein metric is widely used in prior works (Zhang et al., 2021; Ferns et al., 2004) for measuring the distance between two probability distributions of states, indicating the cost of transporting mass from one distribution to another.
Here, we adopt 2-Wasserstein metric W 2 in Eq.2, as it has a convenient formulation:
{+W 2 ( N ( µ i , (cid:80) i ) , N ( µ j , (cid:80) j )) 
 = (cid:107) µ i − µ j (cid:107) 
 + (cid:107) (cid:80) 1 / 
+}
i − (cid:80) 1 / 2 j (cid:107) 2 F , where (cid:107) · (cid:107) F is the Frobenius norm.
Exactly computing the behavioral distance between two actions according to Eq.2 is generally impractical in the ofﬂine setting since the data points of states in the ofﬂine dataset are difﬁcult to cover the entire state space, especially when the action space is large.
Estimating the results of the max expectation over S might induce a large variance.
Therefore, we extend the deﬁnition 1 to a conditional version as follows:
Deﬁnition 2 (State-conditional Behavioral Metric of Actions)
Given an MDP M and a state s , a stateconditional behavioral metric of actions is a function d : S ×
A × A (cid:55)→ R ≥ 0 such that:
[Equation]
This deﬁnition avoids estimating the expectation over S and brings a higher representation capability: It allows the actions’ relations to be changed under different states, e.g., in the task illustrated in Fig. ?? , the actions “move left” and “move right” are behaviorally equivalent at the start position ( they both induce a collision) but different at the agent’s current position (“move left” induces a collision but the other induces a movement).
Data-distribution detection.
The major issue in ofﬂine 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 we also expect that there is a penalty distance between the o.o.d.
actions and the indistribution actions, forcing the behavioral regularized policy tends to choose in-distribution actions.
Thus, we modify the pseudometric function d (Eq.3) to reﬂect the data-distributional relation between actions:
[Equation]
where p > 0 is a penalty coefﬁcient,
I β
( a i , a j | s ) is a function indicating the relation of actions: If a i and a j are from two data distributions, I β
( a i , a j | s )
= 0 , otherwise,
I β ( a i , a j | s )
= 1 .
4. L EARNING 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 φ : S × A (cid:55)→ E for generating action representations that possess the desired property d ( a i , a j | s ) :
= (cid:107) φ ( a i ; s ) − φ ( a j ; s ) (cid:107) 1 .
It means that the (cid:96) 1 distance between any two actions in the learned representation space equals their pseudometric distance.
Concretely, the learning objective can be written as follows:
[Equation]
(5) where e is =
φ ( a i ; s ) , e js = φ ( a j ; s ) .
In practice, this objective requires sampling state-action pairs with the same state s but different actions a i and a j .
However, it is challenging to search for sufﬁcient such pairs from the ofﬂine dataset.
Thus, we modify the objective by using the estimated samples:
[Equation]
where
[Equation]
Speciﬁcally, ˆ d ( a i , a j | s ) is an estimation of d ( a i , a j | s ) .
e s denotes φ ( a ; s ) with stop gradients, ˆ R and ˆ P is the reward model and transition model, which have their own training steps.
ˆ I β ( a j | s ) is also a trainable model to predict whether a j is an o.o.d.
action.
If a j is out of distribution, ˆ I β ( a j | s ) = 1 , otherwise ˆ I β ( a j | s )
= 0 .
In practice, we follow the prior work to derive ˆ I β ( a | s ) (Fujimoto et al., 2019a).
We train a model G ( a | s )
≈
π β ( a | s ) to predict the probabilities of every action in the dataset and then scale all probabilities by the maximum probability.
a j would be predicted as the o.o.d.
action if the relative probability is below a threshold constant τ .
This procedure can be summarized as ˆ
[Equation]
During training, we sample the tuple ( s, a i , R
a i s )
from D and directly sample the other action a j from the ﬁnite action set A .
Then, other lacked data would be estimated by the trained models and the modiﬁed 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 sufﬁcient to derive a latent space with effective representative ability in ofﬂine RL tasks.
Further details of training the action encoder and all models are described in Fig.1 () and Alg. 2 in the appendix.
Theoretical Analysis.
We ﬁrst 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 d .
Theorem 1 ( Q π is Lipschiz with respect to d ) Given a policy π , let Q π be the value function for a given discount factor γ .
Q π is Lipschitz continuous with respect to d with a Lipschitz constant 11 − γ
[Equation]
Proofs in appendix.
This theorem means that, the closer two actions are in terms of d , the more likely they are to share the similar value.
If we 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 to 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 in-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: If one of a i and a j is the o.o.d. action but the other is not, | Q π ( s, a i ) − Q π ( s, a j ) | ≤ 11 − γ · ( |R ( s, a i ) − R ( s, a j )
| + γ · W 2 ( P ( s (cid:48) | s, a i ) , P ( s (cid:48) | s, a j ) + p ) , otherwise, | Q π ( s, a i ) − Q π ( s, a j ) | ≤ 11 − γ · ( |R ( s, a i ) − R ( s, a j ) | + γ · W 2 ( P ( s (cid:48) | s, a i ) , P ( s (cid:48) | s, a j )) .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 p 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, so the ofﬂine 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 ofﬂine policy to cautiously explore the effect of o.o.d. actions. Now, we would show how action representations complying the pseudemetric d improve dowenstream ofﬂine RL.
Assuming that the BMA action encoder φ ( a ; s ) has a learning error (cid:15) : If the pseudometric distance between any two action is smaller than (cid:15) , 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 ofﬂine RL is performed with BMA. Theorem 2 (Performance bound in ofﬂine RL.) Let π ∗
i ( e | s ) be the policy obtained by CQL performing with BMA in the constructed MDP M and π ∗ i,g ( a | s ) refer to the overall policy when π ∗ i ( e | s ) is used together with nearest lookup function g . Let J ( π, M ) refer to the expected return of π in M and φ ( a ; s ) is the
BMA action encoder, which has a learning error (cid:15) . Let π β refer to the behavioral policy genearting D and π β ( e | s ) ≡ e = φ ( a ; s ) , a ∼ π β ( a | s ) .Then , J ( π ∗ i,g , M ) ≥ J ( π β , M ) − k where
[Equation]
Proofs in appendix.
This bound suggests that the lower bound over the performance of the learned overall policy depends on three factors: 1) the divergence between the learned latent-space policy and the latent behavioral policy ( D CQL ( π ∗ i , π β )( s ) ), which is controlled by the applied ofﬂine algorithm.
The size of the latent space |E| . According to Eq.(6) in (Kumar et al., 2020), note that if we directly learn an ofﬂine policy from the original action space A , the lower bound of ofﬂine policy performance would depend on |A| . This indicates that a latent action space would induce an improvement on the worst-case deterioration over the learned policy, since we can set |E| (cid:28) |A| , especially when |A| is enormous.
The learning error of BMA action encoder (cid:15) .
This term is induced by a self-supervised loss function.
Compared with errors induced by ofﬂine RL algorithms, this term is much easier to be controlled. In conclusion, BMA provides a therotically gurrantee of policy improvement for ofﬂine RL task with large discrete action spaces.
E XPERIMENTS
In this section, we empirically show that BMA could be used as a drop-in extension for improving the policy performance in existing ofﬂine RL algorithms for problems with large discrete action spaces.
Combing BMA with two standard ofﬂine 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% conﬁdence interval of standard deviation.
We believe that the reported results can be further improved by using our framework with other ofﬂine 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: A maze environment introduced in Chandak et al.
It has 4096 discrete actions.
A multi-step maze environment that requires select actions from the 6 -step planning decisions.
It is modiﬁed 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. E XPERIMENTAL 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 ofﬂine RL algorithms: CQL (Kumar et al., 2020) and BCQ (Fujimoto et al., 2019b).
We ﬁrst learn our pseudometric-based action representations from the given datasets.
Then, we utilize our framework to train policies using the two ofﬂine 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.2 (a).
We also visualize the average performance of the behavioral policy.
Overall, we ﬁnd that BMA-CQL and BMA-BCQ show faster convergence and better performance than their discrete versions in the four environments.
Furthermore, discrete CQL and discrete BCQ cannot outperform the behavior policies and show unstable learning curves in some situations.
This indicates that in ofﬂine RL with large action sets, directly training policies on 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 to similar results?
To reply to this question, we compare BMA-CQL against baselines using the same architecture but exchange our action representations with other representations.
We ﬁrst consider two random action representations to examine whether simply projecting discrete actions into continuous action spaces can improve the performance.
The ﬁrst scheme generates random vectors ranging from -1 to 1 as action representations ( Rand-1 ).
The other scheme also generates random vectors but forces the (cid:96)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 in 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 are set as the movement vectors. In the multi-step environment, Multi-step maze BMA-CQL
BMA-BCQ
BMA-CQL
BMA-BCQ c Multi-step maze Recommendation system Dialogue system a
BMA-BCQ
BMA-CQL
Discrete BCQ
Discrete CQL
BC
Behavioral
Maze
Maze
Multi-step maze
Recommendation system Dialogue system b
BMA
RandRand-
Transition
Reconstruction
External
PLAS
Behavioral
BMA p = 0.
BMA p = 0.
BMA p = 0.
BMA w/o p
CVAE
Recommendation
[Figure]
We consider BMA with different penalty distances (BMA with p = 0 .
, removing the penalty distance from the learning objective (BMA w/o p
and removing the distance learning the actions are represented as the concatenation of the base action at each step of the plan.
In the recommendation system, the action embedding is the concatenation of video type (expressed as a one-hot vector), video quality, and video length.
Note that there is no proper external information to represent dialogue actions, so we did not consider external action representations in this task.
Finally, we consider a similar work for ofﬂine 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.2 (b).
Two random action representations 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 complying with the proposed pseudometric are more suitable to ofﬂine RL. In addition, external representations and Online embedding Rand embedding
Reconstruction embedding
Ours S t a t e S t a t e  0. PLAS both show poor performance in all tasks.
The reason might be that external representations do not capture the information of dynamics in 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 cannot address the tasks with large discrete action spaces well.
Fig.3 illustrates the action representations.
We sample two states from the multi-step maze environment and plot the different action representations.
We see that, in ofﬂine settings, the rand action representations cannot generalize well since actions with different values are mixed together.
The reconstruction-based representations and our pseudometric-based 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 a more compact structure.
Intuitively, we believe that this is the reason why our method shows faster convergence performance than reconstruction-based representations.
Ablations 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.2 (c) illustrate the results.
We ﬁrst consider the effects of different scales of the penalty constant.
We set p = 0 .
5 respectively.
The results show that a proper penalty constant would bring better performance, but a large penalty constant might make the performance collapse in some situations.
The reason might be that large penalty distances destroy the generalization structure of action representations.
Then, we consider removing the penalty constant (denoted as BMA w/o p ).
The results of this ablation show that without the distance penalty mechanism to ensure the distribution detection, there is performance degradation in our framework.
Finally, we remove the objective function (Eq.6) for controlling the distances between actions.
As a result, the architecture of this scenario is similar to a conditional variational auto-encoder that takes current states and actions as input and aims to reconstruct next states and rewards (denoted as 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.
5. C ONCLUSIONS
This paper proposes Behavioral Metrics of Actions: a new action representation learning framework for ofﬂine
RL tasks with large discrete action spaces.
In the learned action representation spaces, distances between representations of discrete actions reﬂect their behavioral and data-distributional relations.
We derive theoretically statements about the beneﬁts of learning policies based on our action representations.
Experimental results show that our methods signiﬁcantly outperform prior works in a variety of environments.
A A PPENDIX A. A DDITIONAL T HEOREMS AND P ROOFS
Theorem 1 ( Q π is Lipschiz with respect to d ) Given a policy π ( a | s ) , let Q π be the value function for a given discount factor γ .
Q π is Lipschitz continuous with respect to d with a Lipschitz constant 11 − γ
[Equation]
Proofs. follows by expanding each Q , rearranging terms and then simplifying the expression.
[Equation]
We assume that R ≤ 1 , so we get
[Equation] [Equation]
Combing the above equations
[Equation]
we have
[Equation]
Since in Eq.(4), p > 0 and I β ( a i , a j | s ) ∈ { 0 , 1 } , we get
[Equation]
According to theorem 1, we can derive the bound between value functions in the original MDP and the MDP constructed by the learned BMA action representations. Assuming that the BMA action encoder φ ( a ; s ) has a learning error (cid:15) : If the pseudometric distance between any two action is smaller than (cid:15) , they would be aggregated together. Lemma 1 (Value bound based on φ ) Given the BMA action encoder φ ( a ; s ) which has a learning error (cid:15) 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 π in the orginal MDP and its converted version π in the constructed MDP are bounded as:
[Equation]
Proofs.
First, let us give some notation and assumptions concerning the form of the constructed MDP M .
We assume M =
( S , A φ , P φ , R φ , γ ) , where S is the same state space, A φ is the action representation space where the (cid:96) 1 distances between any two representations corresponds to our deﬁned pseudometric distance, and there exists a learning error of φ : If the pseudometric distance between any two action is smaller than (cid:15) , they would be aggregated together. Recall that the nearest neighbor lookup g is as follow: [Equation]
[Equation]
As a result
[Equation]
we can get the relationships between the transition probabilities and rewards in the two MDPs, i
[Equation]
e.
[Equation]
By this way
[Equation]
we can get
[Equation]
According to the proofs of theorem 1, we get
[Equation]
We have argued that there exists some near-optimal value functions based on the constructed MDP M , if φ is sufﬁciently learned.
Now, we would show how action representations complying the pseudemetric d improve dowenstream ofﬂine RL.
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 ofﬂine RL is performed with BMA.
Theorem 2 (Performance bound in ofﬂine RL.)
Let π ∗ i
( e | s ) be the policy obtained by CQL performing with
BMA in the constructed MDP M and π ∗ i,g ( a | s ) refer to the overall policy when π ∗ i ( e | s ) is used together with nearest lookup function g .
Let J ( π, M ) refer to the expected return of π in M and φ ( a ; s ) is the
BMA action encoder, which has a learning error (cid:15) .
Let π β refer to the behavioral policy genearting D and π β ( e | s ) ≡ e = φ ( a ; s ) , a ∼ π β ( a | s ) .Then
, J ( π ∗ i,g , M ) ≥ J ( π β , M ) − k where
[Equation]
=
[Equation]
Proofs.
First, let’s break | J ( π ∗ i,g , M ) − J
( π β , M ) | into
[Equation]
where π β is the converted behavioral policy executed on the constructed MDP M mentioned in the proofs of thereom ?? .
It can be deﬁned as e ∼ π β ≡ e = φ ( a ; s ) , a ∼ π β ( a | s ) .
According to theorem 3.6 in (Kumar et al., 2020), we apply it to M , we get
[Equation] [Equation]
Then, we try to bound | J ( π ∗ i,g , M ) − J ( π ∗ i , M ) | . Since the decisions of π ∗ i and π ∗ i,g are in one-to-one correspondence, we get | J ( π ∗ i,g , M ) − J ( π ∗ i , M ) | = 0 . Now, we try to bound | J ( π β , M ) − J ( π β , M ) | : | J (
[Equation]
| (cid:88)
[Equation] [Equation] According to lemma 1, | Q π β M ( s, e as ) − Q π β M ( s, a ) | = (cid:15) 1 − γ , so | J ( π β , M ) − J ( π β , M ) | =
(cid:15) 1 − γ = k 2 . Finally, we get k = k 1 + k 2 .
We apply O to get the notation in the theorem. A. A LGORITHMS A. A RCHITECTURE , H YPERPARAMETERS , AND I NFRASTRUCTURE In our implementations, the nearest neighbor lookup function g can be regarded as ﬁnding the minimum distance between ˆ e and all actions’ embeddings. So, we directly adopted the torch.min() function in the PyTorch python package.
Its time complexity is O ( N ) , so there would not be a serious scale problem when the action set is enormous. For ˆ I β ( a | s ) . We train a model G ( a | s )
≈ π β ( a | s ) to predict the probabilities of every action under a given state s and then scale all probabilities by the maximum probability. a j would be predicted as the o.o.d.
action
Data: Ofﬂine Dataset D , the action encoder φ , Any ofﬂine RL algorithm Y , and number of iterations T .
Result:
Policy π i,θ (ˆ e | s ) for t = 0 to T do
Sample data ( s, a, s (cid:48) , r ) ∼ D ; e ← φ ( a ; s ) ; Update policy π i,θ (ˆ e | s ) with ( s, e, s (cid:48) , r ) using the given algorithm Y end Result: Action encoder φ for t = 0 to T do Sample data ( s, a i , R ( s, a i )) ∼ D ;
Sample the other action a j ∼ A ; Compute representations e is = φ ( a i ; s ) and e js = φ ( a j ; s ) ; Estimate the transition distributions: ˆ P ( ·| s, e is ) and ˆ P ( ·| s, e js ) ; Estimate the reward of ( s, a j ) :
ˆ R ( s, e js ) ; Estimate whether a j is an o.o.d. action: ˆ I β ( a j | s ) ;
Compute the estimated pseudometric distance: ˆ d ( a i , a j | s ) Eq.3; Train encoder: J ( φ ) = ( (cid:107) e is − e js (cid:107) 1 − ˆ d ( a i , a j | s )) 2 ;
Train transition model: J ( ˆ P , φ ) = ( ˆ P ( ·| s, e is ) − s (cid:48) ) 2 ; Train reward model: J ( ˆ R , φ ) = ( ˆ R ( s, e is ) − R
( s, a i ))
Train prediction model: J ( G ) =
CrossEntropy ( G ( ·| s ) , a i ) ; end
[Figure]
BMA action encoder
[Figure]
Transition-based action encoder
[Figure]
Reward model r
[Figure]
Reward model
[Figure]
Transition model
[Figure]
Transition model
[Figure]
s
Reconstruction [Figure] action encoder [Figure]
s
Reconstruction-based action encoder
[Figure]
if its relative probability is below a threshold constant τ .
This procedure can be summarized as ˆ I β ( a | s ) =
G ( a | s ) max ˆ a ∈A G (ˆ a | s ) ≤ τ .
We have searched over τ = { 0 , 0 . 1 , 0 . 3 , 0 . 5 , 0 .
and found that the best setting is τ = 0 .
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 as follow:
[Equation]
[Equation] During training, all data are randomly sampled from the ofﬂine dataset. We set the batch size as 128 and set the training gradient steps for all models as 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 no momentum or weight decay.
We set the dimension of the action representations |E| = 2 , 2 , 10 ,and the penalty coefﬁcient p = 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 in the newest version. [Figure]
[Figure] For the applied ofﬂine RL algorithms: BCQ (Fujimoto et al., 2019b), CQL(Kumar et 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 α = { 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 .
9 for the
the A. E NVIRONMENTS AND D ATASETS
Maze: We ﬁrst adopt the maze environment introduced in Chandak et al.
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 of movement in one direction).
It can choose whether each actuator should be on or off.
The effect of an action is the vectorial summation of the movements associated with the selected actuators.
Thus, the size of the action space |A| = 2 n .
We set n = 12 , so there are 4096 actions in total.
There are three reward types: a reward ofis given if the agent visits a rewarding region, a reward of -20 is given for visiting a puddle region, and a reward of 100 is given with the episode ending if it is on a 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 a general problem with a large action set.
If an environment has i actions available at each step, the number of m -step plans would increase exponentially.
In this setting, |A| =
m i .
We implement a 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 to choose a 6 -step plan per 6 time steps.
Thus, the actual action size in this environment is 4 6 = 4096 .
Recommender environment: We consider a real-world domain with a large action set.
The agent needs to recommend different videos to the simulated users, and its goal is to extend the user’s viewing time.
We utilize the open-source platform Recsim (Ie et al., 2019) and adopt the provided scenario called interest_evolution 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 1 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 a video.
In addition, the viewed video’s quality would also inﬂuence 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 videos’ 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) to simulate a dialogue environment in which each action is a dialogue.
We focus on the MultiWOZ domain (Eric et al., 2019).
Its main task is to help a tourist in various situations involving multiple sub-domains, such as requesting basic information about attractions and booking a hotel room.
Speciﬁcally, 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 UserPolicyAgendaMultiwoz to simulate the tourists’ policies and use MultiWozStateEncoder to encode the histories of dialogs into vectorized states.
Data collection: Since there are no open-source datasets for the ofﬂine 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., 2020), 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. A. E XAMPLE
To better understand the inﬂuence of action spaces to the behavioral regularization of ofﬂine algorithms, we give a toy example iluustrated in Fig.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 current position, we can ﬁnd that different action spaces induce different results of the behavioral regularization.
The target policy is limited to selecting the sub-optimal in-distribution 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 inﬂuence the results of the behavioral regularization.
In representation space 1, behaviorally different actions (the transitions and rewards induced by them are different) are 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 or other ‘safe’ actions.
A. O NLINE PERFORMANCE OF BMA
Actually, BMA action representations can be directly applied on the online setting.
Utilizing 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-based representations in the Multi-step maze environment.
The results are illustrated in Fig.6.
We can ﬁnd that, in the online setting, three action representations can all achieve fast convergence, while only BMA can get good performance in the ofﬂine setting (Fig.2 (b)).
This indicates that the BMA is more effective in ofﬂine RL tasks with large discrete action spaces.