Jointly-Trained State-Action Embedding for Efficient Reinforcement Learning
Authors: Paul Julian Pritz
Liang Ma
Kin Leung
Venue:  (2021)
Keywords: reinforcement learning
embedding
representation learning
state-action embedding

Abstract: While reinforcement learning has achieved considerable successes in recent years, state-of-the-art models are often still limited by the size of state and action spaces. Model-free reinforcement learning approaches use some form of state representations and the latest work has explored embedding techniques for actions, both with the aim of achieving better generalization and applicability. However, these approaches consider only states or actions, ignoring the interaction between them when generating embedded representations. In this work, we propose a new approach for jointly learning embeddings for states and actions that combines aspects of model-free and model-based reinforcement learning, which can be applied in both discrete and continuous domains. Specifically, we use a model of the environment to obtain embeddings for states and actions and present a generic architecture that uses these to learn a policy. In this way, the embedded representations obtained via our approach enable better generalization over both states and actions by capturing similarities in the embedding spaces. Evaluations of our approach on several gaming, robotic control, and recommender systems show it significantly outperforms state-of-the-art models in both discrete/continuous domains with large state/action spaces, thus confirming its efficacy and the overall superior performance.

Under review conference paper ICLR 2021
JOINTLY-TRAINED STATE-ACTION EMBEDDING FOR EFFICIENT REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
While reinforcement learning has achieved considerable successes in recent years, state-of-the-art models are often still limited by the size of state and action spaces. Model-fre freee-free reinforcement learning approaches use some form of state represen- tations and the latest work has explored embedding techniques for actions, both with the aim of achieving better generalization and applicability. However, these approaches consider only states or actions, ignoring the interaction between them when generating embedded representations. In this work, we propose a new ap- proach for jointly learning embeddings for states and actions that combines as- pects of model-free and model-based reinforcement learning, which can be ap- plied both discrete and continuous domains. Specifically, we use a model the environment to obtain embeddings for states and actions and present generic architecture that uses these to learn a policy. In this way, the embedded repre- sentations obtained via our approach enable better generalization over both states and actions by capturing similarities in the embedding spaces, Evaluations of our approach on several gaming, robotic control, and recommender systems show it significantly outperforms state-of-the-ar art models in both discrete/continuous do- mains with large state/action spaces, thus confirming its efficacy and the overall superior performance.
INTRODUCTION
Reinforcement learning (RL) has been successfully applied to range of tasks, including challenging gaming scenarios (Mnih al. 2015). However, the application of RL in many real-world domains often hindered by the large number of possible states and actions these settings present. For instance,
resource management in computing clusters (Mao et al. 2016 Evans & Gao 2016). portfolio man-
agement (Jiang al. 2017), and recommender systems (Lei Lij 2019
al. 2018) all suffer
from extremely large state/action spaces, thus challenging to be tackled by RL.
In this work, we investigate efficient training of reinforcement learning agents in the presence of large state-action spaces, aiming improve the applicability of RL to real-world domains. Previous work attempting to address this challenge has explored the idea of learning representations (embed- dings) for states or actions. Specifically, for state embeddings, using machine learning to obtain meaningful features from raw state representations is common practice in RL, e.g. through the use of convolutional neural networks for image input (Mnih et al., 2013). Previous works such as by Ha Schmidhuber (20T8h) have explored the use of environment models, termed world models, to learn abstract representations. and several pieces of literature explore state aggregation using bisimulation metries (Castro] 2020). While for action embeddings, the most recent works by Ten- nenholtz & Mannor (019)) and Chandak et al. (2019) propose methods for learning embeddings for discrete actions that can be directly used by RL agent and improve generalization over actions. However, these works consider the state representation and action representation as isolated tasks, which ignore the underlying relationships between them. In this regard, we take a different approach and propose to jointly learn embeddings for states and actions, aiming for better generalization over both states and actions their respective embedding spaces.
To this end, we propose an architecture consisting of two models: a model of the environment that is used to generate state and action representations and model-free RL agent that lears policy using the embedded states and actions. By using these models, our approach combines aspects

---

Under review conference paper ICLR 2021
of model-based and model -ree reinforcement learning and effectively bridges the gap between both approaches. In contrast to model-based RL, however, we do not use the environment model for planning. but to learn state and action embeddings. One key benefit of this approach is that state and action representations can be learned in a supervised manner. which greatly improves sampling efficiency and potentially enables their use for transfer learning. In sum, our key contributions are:
We formulate embedding model for states and actions, along with internal policy Te that leverages the learned state/action embeddings, as well as the corresponding overall policy To. We show the existence of an overall policy To that achieves optimality in the original problem domain.
We further prove the equivalence between updates the internal policy m, acting in bedding space and updates to the overall policy To-
We present supervised learning algorithm for the proposed embedding model that can be combined with any policy gradient based RL algorithms.
We evaluate our methodology on some game-based as well as real-world tasks and find that it outperforms state-of-the-an models in both discrele/continuous domains.
The remainder of this paper is structured as follows: Section we provide some background on RL. We then give an overview of related work in Section 3 before presenting our proposed
methodology in Section which we evaluate in Section
2 BACKGROUND
We consider agent interacting with its environment over discrete time steps, where the en- vironment is modelled discrete-time Markov decision process (MDP), defined by tuple (S, A, R. Specifically, S and A are the sets of all possible states and actions, referred to as the state space and action space, respectively. In this work, we consider both discrete and con- tinuous state and action spaces. The transition function from one state another, given action, is T A + S, which may be deterministic or stochastic. The agent receives a reward at each time step defined by R S x A + R. E [0, denotes the reward discounting parameter. The state, action, and reward time & {0, 1..... denoted by the random variables S,, A,, and RG
The initial state of the environment comes from an initial state distribution dg- Thereafter, the agent follows a policy x, defined as a conditional distribution over actions given states, ie.. <(als) P(A, a.S. s). The goal of the reinforcement learning agent is to find an op- timal policy '² that maximizes the expected sum of discounted accumulated future rewards for a given environment, i.e., x=" E arg max, E| > R& /=]- For any policy, we also define the state value function (*)) E 7"th Retx S, x] and the state-action value function Q Q²(ssa) a) E|E * R
8, =a] a].
3 RELATED WORK
For the application of state embeddings in reinforcement learning. there are two dominant strands of research, namely world models and state aggregation using bisimulation metrics. World model approaches train an environment model in a supervised fashion from experience collected in the environment, which is then used to generate compressed state representations (Ha & Schmidhuber 2018a) or train an agent using the learned world model (Ha & Schmidhuber] [2088]] (Schmidhu- ber (2015). Further applications of world models, c.g. for Atari 2000 domains, show that abstract state representations learned via world models can substantially improve sample efficiency (Kaiser al. 2019) [affer et al. 2020). Similar this idea, Munk et al. (0066) pre-train environment model and use to provide state representations for an RL agent. Furthermore, de Bruin et all (2018) investigate additional learning objectives to learn state representations, and Francois- Lavee
al. (2019) propose the use of an environment model to generate abstract state representations: their learned state representations are then used by a Q-learning agent. By using a learned model the environment to generate abstract states, these approaches capture structure in the state space and reduce the dimensionality of its representation an idea similar proposed embedding model.

---

Under review as conference paper ICLR 2021
In contrast to world models, we considers both states and actions, rather than just states. Addition- ally, our approach trains an RL agent on the original environment using its embedded representation. rather than learning a policy using surrogate world model. Bisimulation, on the other hand, is a method for aggregating states that are "behaviorally equivalent" (Li al., 2006) and can improve convergence speed by grouping similar states into abstract states. Several works on bisimulation- based state aggregation, e.g. | et al. (2004); Givan et al. (2003), present different metries to measure state similarity. Furthermore, Zhang et al. (2020) and Castro (2020) propose deep learning methods for generating bisimulation-basee proposed in several earlier works (Castro) 2020). While there are parallels between bisimulation and our approach, we do not propose the aggregation of states. Instead, our embedding technique projects states into a continuous state embedding space, similar to Zhang et al. (2020), where their behavioral similarity is captured in their proximity in embedding space. Furthermore, our method embeds both states and actions and does not employ an explicit similarity metric such as a bisimula- tion metric, but instead learns the relationships among different states and actions via an environment model. State representations are also used in RL--based NI.P tasks, such as Narasimhan al. (2015). who jointly train LSTM-based state representation module and DQN agent, and Ammanabrolu & Riedl (2018), who propose the use of a knowledge graph based state representation.
state aggregations that scale beyond the tabular methods
In addition to state representations, previous work has explored the use of additional models to leam meaningful action representations. In partieular, Van Hasselt & Wiering (2009) investigate the use continuous actor policy gradient based approach to solve discrete action MDPs, where the policy is learned in continuous space and actions are discretized before execution in the environment. ADulac-Amold et al. (2015) propose a similar methodology, where a policy is Iearned continuous action embedding space and then mapped to discrete actions in the original problem domain. Both Van Hasselt & Wiering (2009) and Dulac-Arnold et al. (2013) only consider actions and assume that embeddings are known a priori. TTnnenholtz & Mannor (2019) propose method- ology called Vec2 where they introduce an embedding model similar to the Skip-Gram model (Mikolov et al. (203) that is trained using data from expert demonstrations and then combined with a DQN agent. One significant drawback of this approach is that information on the semantics of actions has to be injected via expert demonstration and is not learned automatically. In contrast, Chandak et al. (2019) propose method that enables the self-supervised learning of action repre- sentations from experience collected by the RL agent, using embedding model that resembles CBOW. The approach presented in (Chandak et al.) 2019) is most closely related to our method the way of integrating embeddings with policy gradient based RI. algorithms. However, they only consider action embeddings, while we jointly embed states and actions and improve the performance significantly. As an alternative to embedding methods, |Pazis & Parr (2011) and [Sallans & Hinton (2004) represent actions in binary format, where they learn a value function associated with each bit, and (Sharma al. (2017) propose the use of factored actions, where actions are expressed using underlying primitive actions. Again, these three methods rely on handcrafted action decomposition.
Previous embedding approaches consider either states or actions in isolation, embedding only one of the two or using separate models for each. By contrast, in addition to capturing structure in the state or action space respectively, we leverage interdependencies between the two by jointly learning embeddings for them, and thus improve the quality of embeddings.
4 METHODOLOOGY
We propose an embedding model to jointly learn state and action embeddings and present generic framework that uses the learned embeddings in conjunction with any policy gradient algorithms.
4.1 ARCHITECTURE
Our proposed approach has three components: (i) state-action embedding model, (ii) policy gradient based RL agent that takes in the embeded states and outputs actions in action embedding space, and (iii) action mapping function that maps embedded actions to actions in the original problem domain. We define two embedding spaces, i.c., a continuous state embedding space X E R. and a continuous action embedding space E R2; states are projected to the state embedding space via a function @ : S + X. On the other hand, we define a function f : E + A that maps points in the embedding space to actions in the original problem domain. With these embedding spaces, we

---

Under review conference paper ICLR 2021
VC
X
S
Environmen
A
V2
T Tronstion
2 E
X;
A
E
V.
(a)
(b)
Figure I: (a) Overall architecture, using the state embedding function Φ, the internal policy T, and the action mapping function f. (b) Embedding model for learning & and the auxiliary action em- bedding function (the red arrows denote the gradient flow of the supervised loss in Equation (
further define internal policy my that takes in the embedded state and outputs the embedded action, i.e., N, ++ E. Furthermore, the mapping functions & and f. together with the internal policy m form the overall policy T. : S + A, i.c., To is characterized by the following three components:
X, = (S)), E (T(1)x),, and A, = f(E),,
where random variables S. E S, E X. E. E E, and A; E A denote the state in the original problem domain, state in the embedding domain, action in the embedding domain, and action in the original domain time respectively. The overall model architecture is illustrated in Figure Ta In Figure Ial function ó is directly learned using an embedding model illustrated in Figure [b] (detaile in Section ). However, function f is not directly part of this embedding model, but uses the obtained action embeddings to learn mapping from e & to
(see Section 4.3.2 for details).
For an RL problem, the ultimate goal optimize the policy ι the original problem domain, i.e.. v in the S x A space. For the proposed architecture, the idea is to first establish efficient embeddings for both states and actions in S and A. Then we reduce the problem of finding the optimal x' in the original problem domain to the problem of internal policy n; optimization that acts purely in the embedding space. However, it is still unclear how the optimization of n. is related to n²" In addition, the overall policy T. relying establish the relationships among n. To and n; in the next section, which is the foundation to justify the validity of leveraging hoth state-action embeddings in RI.
and further complicates the policy learning. In this regard, we
4.2 THEORETICAL FOUNDATIONS: RELATIONSHIPS AMONG ग. AND π..
To understand the relationships among म, To. and nú we first derive how n. is related ग, after which we prove the existence of n. that achieves optimality in the original problem domain. For this goal, two further assumptions are required on the nature of the state embedding function é and the action mapping function f.
Given action embedding EG, Ae is deterministic and defined by function f E + A, i.e. there exists a such that P(A, alFa =e)======
Assumption 1.
Assumption 2. Given state S., X, is deterministic and defined by a function é S + X, i.e. there exists such that P(X, x
We validate Assumption 1 which defines f as a many-to-one mapping, empirically via all experi- ments conducted in Section and find that no two actions share exactly the same embedded repre- sentation, i.c., Assumption holds in practice. Note that we also assume the Markov property for our environment, which is standard assumption for reinforcement learning problems. With slight abuse of notation, denote the inverse mapping from action
A to its corresponding points
the embedding space (one-to-many mapping) by f (a) == {e E E f(e) = a).
Lemma 1. Under Assumptions 7 and 2 such that
for policy n in the original problem domain, there exists m
(3) - Σ a3A (cl-f ((a)
T.(e|x = d(s))

---

Under review conference paper ICLR 2021
Proof Sketch. Based on the Bellman equation for the value function e the original domain, we introduce the embedded state & = d(s) using Assumption [ By the law of total probability and Assumption [ we then introduce the embedded action From the Markovian property and the definition of our model, we can derive Claims|T on (conditional) independence in Appendix which then allow us to remove the original state and action from the expression of the policy, leaving us with n.. See Appendix A.2 for the complete proof.
By Assumptions 1 and policy ny as
Lemma
I alloos express the overall policy % in terms of the internal
To (alss) - </(e|ı
= ((s)) de,
(I)
(c)=f-1(a))
bridging the gap between the original problem domain and the policy in embedding spaces X' and Under x² in the original domain, we define '" == ex' and Q == Q for the ease of discussions.
Then using Lemma [| we now prove the existence of overall policy %o that is optimal.
Theorem 1. Under Assumptions S [nd there exists an overall policy No that optimal, such that UT
Proof. Under finite state and action sets, bounded rewards, and 7 E [0, 1), at least one optimal I we then have
policy n' exists. From Lemma
(s) Σ Σ
(s,
(2)
Thus, 3 d. f. and n. representing overall policy To, which is optimal, i.c. vd = v².
Theorem I Juggests that in order to get the optimal policy x² in the original domain, we can focus on the optimization of the overall policy To, which will discussed in the next section.
4.3 ARCHITECTURE EMBODIMENT AND TRAINING
In this section, we first present the implementation and the loint-training of the state-action embed- ding model (illustrated in Figure lb. Based on this method and according to Theorem [] we then propose a strategy to train the overall policy Ro. where functions é and f are iteratively updated.
4.3.1 JOINT TRAINING OF THE STATE-ACTION ACTION EMBEDDING
this section, we elaborate on our proposed embedding model for jointly learning state and action embeddings. Specifically, we require two components: (i) function á that projects S; into embed- ding space X' and (ii) function f that maps each point in embedding space & to an action A, in the original problem domain. Since these functions are not known advance, we train estimators using a model of the environment (see Figure lb). This environment model requires two further compo- nents: (iii) function 9 A + E and (iv) transition model T E X + S that predicts the next state from the concatenated embeddings of the current state and action. Note that g a one-to-one map- ping, since according to Assumption T no two actions have exactly the same embedding. Denote the estimators of components (i)-(iv) by g. 9. T: the target of the environment model is P(S (S, (,) ((. 1/5++|X4 IX, E,)g(E, A.) The difference between the true transition probabilities P(S.+ Si,, Ar) and the estimated proba- bilities 1 Ar) can be measured using the Kullback-Leibler (KI.) divergence, where the A.).
expectation is over the true distribution
DKL (kI(P) - ES In ((S
P(S++1|S. (.).)))))))). A
(4)
Since we observe the full tuple (Se11. St Ad), we can compute sample estimate of this using Equation (4). Equation (4), the denominator does not depend on H. 9. or î. Therefore, we define the loss function for the embedding model as
L(O. g, /)) EE P [I((P(S++ISS, A.)]]
(5)

---

Under review conference paper ICLR 2021
Note that the estimator f not directly included in the embedding model. Instead, we find f by minimizing the reconstruction error between the original action and the action reconstructed from the embedding. We can again define this using the KL. divergence as
[() = -A ((A)E)) In(f(A/)E))
(6)
Then, all components of our embedding model be learned by minimizing the loss functions C( 9. T) and C(f). Note that the embedding model is trained in two steps. firstly by minimizing the loss in Equation ( to update Φ, 0. and 7 and secondly by minimizing the loss in Equation (6) to update f. Here, we give the target of the embedding model for the case of discrete state and action spaces. However, the model is equally applicable to continuous domains. In this case, the loss functions [(è, g. T') and L(f) must be replaced with loss functions suitable for continuous domains. For continuous domains, we adopt mean squared error loss instead of the losses given in Equations (5) and (6). but other loss functions may also be applicable. For our experiments, we parameterize all components of the embedding model illustrated in Figure Tb as neural networks. More details on this can be found in Appendix A..4
4.3.2 STATE-ACTION EMBEDDING AND POLICY LEARNING
Theorem I oonly shows that optimizing To help achieve the optimality in the original domain. Then a natural question is whether we can optimize v. by directly optimizing the internal policy a using the state/action embeddings in Section
To answer this question, we first mathematically derive that updating и; is equivalent updating wo. We then proceed to present iterative algorithm for learning the state-action embeddings and the internal policy.
Suppose T; is parameterized by 0. Then with the objective of optimizing no by only updating m. we define the performance function of To a
Jo f)
Σ do(s)
Σ [=((als)(='ss.))) aEA
SES
Let the Cate-action-val value function for the internal policy be Q (2x, e) E Ret4/ x, E, = e]. We can then define the performance function of the internal policy as
(,(
{
dolx))
S (x(e)))
(8)
JEX
The parameters of the overall policy T, can then be learned by updating its parameters in the direction of (J,(ó, 0. f)/80, while the parameters of the internal policy can be learned by updating 0 in the direction of (J.(0))00. With the following assumption, we then have Lemma 2
Assumption 3. The state embedding function d(s) maps each state s a unique state embedding x = ((s), i.e., Vs, # sy P(ő(s:) = =(s,))
Note that Assumption 3, defining é as a one-to-one mapping, is not theoretically guaranteed by the definition of our embedding model. Nevertheless, test this empirically in Section 5aand find that no two states share the same embedded representation in any of our experiments, thus justifying Assumption in practical scenarios considered in this work.
3 for all deterministic functions I and ó, which map euch point A and to an embedded state x E X, and the internal policy 3; parameterized by 0, the gradient of the internal policy's performance function 0d,10) equals the
Lemma 2. Under Assumptions
E E and 8 to action
gradient of the overall policy's 's performance function 8Jo(6.8.f) L.e.
69
0J((0) 0.1. a0
08
Proof Sketch. Assumptions [ and 3 allow us to show the equivalence between the internal state- action-value function and the overall state-action-valle function. Using this equivalence, then remove functions and from 0.,((b 0, f)/00 and are left with the gradient of the internal policy T, only. The complete proof is deferred to Appendix A.3
'Equation is for the overall policy with discrete states and actions. Note that approach is also applicable continuous domains. E.ss and E.sA would then replaced by J. and f.. respectively.

---

Under review conference paper ICLR 2021
Algorithm 1: Joint Training of State-Action Embeddings and the Internal Policy Initialize the state and action embeddings (optional pre-training);
for Epoch 0, 1. do for / 1, 2, ... do
Generate state embedding X; = 6(S:)
Sample action embedding E, ~ (1)+)
Map embedded action A, = [E,);
Execute A, in the environment to observe St, Ry;
Update the and the critic using some policy gradient algorithms:;
Update Φ, 0, T, and f by minimizing the losses in Equations (5) and (
1.0
0.6
3.0
0.2
0.2
00.
0.0
-100
-1.
0.50
0.0
an
(a) Discrete gridworld
(b) State embeddings
(c) Action movements
(d) Action embeddings
Figure 2: State and action embeddings for a discrete state gridworld with 1, 600 states and 2 actions actuators). The embedding dimension for both states and actions is set for visualization purposes. (a) Original environment, i.c., a discrete gridworld as described in Section [.1 The blue area is the goal state and the red dot is the starting position of the agent. (b) Learned state embeddings 2-dimensional state embedding space. (c) Displacement in the Cartesian co-ordinates (ci and c2) caused by actions. Each action is colored according to this displacement with [R = Ae1, G Acz, B 0.5]. (d) Learned action embeddings in continuous action embedding space.
Lemma 2 shows that the updates the internal policy n; and the overall policy equivalent. This allows us to optimize the overall policy by making updates directly to the internal policy mi thereby avoiding the potentially intractable computation of the inverse functions f and ø Since there special restrictions the internal policy и;. designed for continuous control optimize the policy. addition, we can also iteratively update parameters for jointly learning d, f, and as shown in Algorithm
any policy gradient algorithm
EMPIRICAL EVALUATION
We evaluate our proposed architecture game-based applications, robotic control, and real-world recommender system, covering different combinations discrete and continuous state and action spaces. Our methodology is evaluated using Vanilla Policy Gradient (VPG) (Achiam (2018). Proxi- mal Policy Optimization (PPO) (Schulman al. 2017 and Soft Actor Critic (SAC) (Haarnoja et al., 2018) algorithms. We benchmark the performance against these algorithms without embeddings and with action embeddings generated by the method called policy gradients with Representations for Actions (RA) proposed by Chandak et al. (2019). We pre-train the embedding model on randomly collected samples for all experiments and enable continuous updates for the Ant-v2 and recom- mender system environments. Complete parameterizations are presented in Appendix(A.4
5.1 PROOF-OF-CONCEPY GRIDWORLD SLOTMACHINE
Gridworld: similar to that used by Chandak et al. (001)). States are given as a continuous coordinate, while actions are defined via n actuators, equally spaced around the agent, which move the agent the direction they are pointing towards. Then each combination of actuators forms an action, resulting in 2" unique actions. We run two sets of experiments in this environment: (i) we use the continuous state directly and (ii) we discretize the coordinate in 40 x 40 grid. The agent receives small step and collision penalties and large positive reward for reaching the goal state.

---

Under review conference paper ICLR 2021
10
VIG
32 VPG
05
512 VPG
32 VPG
2049 VPG
0.0
2048 VPG
22 VPG RA
-0.5
VPG-A
RA
VPG
Epecode
-10
32 VPGJSAE
-15
-1.0
512 VPG I5NE
-20
VRGYANE
-15
2048 VRG/SAE
2046 VPGUSAE
Samples (Trairing)
Number Semples (Tr Triiting)
(a) Gridworld with discrete states and actions for dif- ferent numbers of actions
(b) Gridworld with continuous states and discrete actions for different numbers of actions
VPS RA
300
VPE--AAE
retum
600
-0.5
200
Edoode
-1.
-15
-2.
VS JAEE
Tranater
0.0
2.5
5.0
10.0
12.5
Te5
1w5
(c) Slotmachine
(d) Half-Cheetah
(c) Transfer learning in the discrete- state gridworld
5.000
VPG
4000
SAC
PPO
14000
PPO
VPG RA
3000
VPG-NAA
12000
PPO-RA
PPO-SSAE SAC/SAE
Epzods
2000
10000
VP6-JSAE
P5O
8000
SAC/SAE
1000
Number of Samples (Training) 140
Number of Samples (Training) 1e5
(f) Ant-v2
(g) Recommender system
Figure 3: Performance of our approach (JSAE, Joint State Action Embedding) compared against benchmarks without embeddings and/or benchmarks with action embeddings by RA (all results are the average return over episodes, with mean and dev. over 10 random seeds).
Figure 2 illustrates the learned embeddings for this environment. The obtained embeddings for states and actions illustrate that our approach is indeed able to obtain meaningful representations. Actions are embedded according to the displacement in the Cartesian coordinates they represent (Figure 2d Figure 26 shows that states embedded according to the coordinate they represent and that the embeddings capture the L-shaped obstacle the original problem domain.
Slotmackine: consists several reels, each which can individually spun by the agent for a fraction of a full turn. Then, each unique combination of these fractions constitutes an action. The discrete state is given the numbers the front of the reels.
Both scenarios yield large discrete action spaces, rendering them well-suited for evaluating our ap- proach. The results for the gridworld and slotmachine environments are shown in Figure Bal [36] and 3c] Our approach outperforms both benchmarks on the discretized gridworld and the slotmachine environments and is comparable for the continuous-state gridworld. Such observation suggests that our approach is particularly well suited for discrete domains. Intuitively, this is because the re- lationship between different states or action in discrete domains is usually not apparent, e.g., in one-hot encoding. continuous domains, however, structure in the state and action space al- ready captured some extent, rendering harder to uncover additional structure by embeddings; nevertheless, when continuous problem becomes more complicated, our embedding method again shows effectiveness in capturing useful information (see Section (.2

---

Under review conference paper ICLR 2021
5.2 ROBOTIC CONTROL
We use the half-cheetah-12 and the Ant-v2 environments from the robotic control simulator Mujoco (Todorov et al. 2012). Here, the agent observes continuous states (describing the current position of the cheetah or ant respectively) and then decides on the force to apply to each of the joints of calculated based
the distance covered in episode. From the results presented in Figure [3]] we observe that our method does not outperform the benchmark on the half-chectah-v2 environment. On the Ant-v2 environment, however, our ap- proach outperforms the benchmarks. Similar the experiments
the robot (i.e., continuous actions). Rewards
continuous gridworld, it is less likely that embedding uncovers additional structure on lop of the structure inherent to the continuous state and action, explaining the comparable performance with and without embeddings on the half- cheetah-v2 domain. Nevertheless, in the more challenging Ant-v2 environment, the state and action representation more complex and the agent might thus benefit from the reduced dimensionality in embedding space. Additionally, the learned embeddings appear to be able to capture sufficient structure in the state and action space, despite using an embedding space with lower dimensionality.
5.3 RECOMMENDER SYSTEM
In addition, also test methodology on real-world application of recommender system. We use data on user behavior from an e-commerce store collected 2019 (Kechinov) 2019). The environment is constructed as an n-gram based model, following Shann et al. (2005). Based on the previously purchased items, user's purchasing likelihood is computed for each item in the store. Recommending an item then scales the likelihood for that item by pre-defined factor. States are represented as a concatenation of the last n purchased items and each item forms an action. In this environment we have 835 items (actions) and approx. 700, 000 states. The results obtained under various RL methods are reported in Figure 3g We find that our approach leads to significant improvements in both the convergence speed and linal performance compared to benchmarks. This result confirms that our method is particularly useful in the presence of large discrete state and action spaces. Interestingly. the PPO and VPG benchmarks ouiperform the benchmarks using the RA methodology. We conjecture that the action embeddings generated in RA on this environment in fact obfuscate the effect action has, and thus limit the agent's 's performance.
5.4 APPLICATION TRANSFER LEARNING
In addition to improving the convergence speed and the accumulated reward, test whether the policies learned one environment can transferred to a similar one. We consider two environ- ments (old and new) that differ by only the state space or the action space. To leverage the previous experience, if the environments have the same state (or action) space, the state (or action) bedding component in the new environment initialized with the weights learned from the old environment; we then train our model in this new environment. For evaluation, we use two grid- worlds, where only the numbers of actions differ, i..., 512 (old) and 2048 (new). Results in Figure Belshow that the model with the transferred knowledge outperforms the case where the entire model is trained from scratch, in terms of the convergence speed. Therefore, it confirms that the previously learned policy serves critical initialization for the new gridworld environment and sheds light on the potential of our approach in transfer learning; further evaluations are left for future work.
6 CONCLUSION
In this paper. we presented new architecture for jointly training states/action embeddings and combined this with common reinforcement learning algorithms. Our theoretical results confirm the validity of the proposed approach, i.e., the existence of an optimal policy in the embedding space. We empirically evaluated our method on several environments, where it outperforms ttate-of-the- art RL approaches in complex large-scale problems. Our approach is easily extensible as it can be combined with most existing RL algorithms and there are no special restrictions on the parameteri- zations of the embedding model. Consequently, we conclude that the proposed methodology works towards the applicability of reinforcement learning to real-world problems.

---

Under review conference paper ICLR 2021
REFERENCES
Joshua Achiam. Spinning Up in Deep Reinforcement Learning. 2018.
Prithviraj Ammanabrolu and Mark O Riedl. Playing text-adventure games with graph-based deep reinforcement learning. arXiv preprint 1rXiv:182..0100 2018.
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic markov decision processes. In Proceedings of the AAAI Conference Artificial Intelligence, volume 34, pp. 10069-10076, 2020.
Yash Chandak, Georgios Theocharous, James Kostas, Scott Jordan, and Philip Thomas. Learning Action Representations for Reinforcement Learning. 36th International Conference on Machine Learning (ICML), pp. 1565 1582, 2019.
Tim de Bruin, Jens Kober, Karl Tuyls, and Robert Babuška. Integrating state representation learning into deep reinforcement learning, JEEE Robotics and Automation Letters, pp. 1394-1401, 2018.
Gabriel Dulac-Arold, Richard Evans, Hado van Hasselt, Peter Sunehag, Timothy Lillicrap. Jonathan Hunt, Timothy Mann, Theophane Weber, Thomas Degris, and Ben Coppin. Deep rein- forcement learning in large discrete action spaces. arXiv preprint arXiv: 2015.
Richard Evans and Gao. Deepmind AI reduces google data centre cooling bill by 40%. DeepMind blog. 2016. URL holl//// deepmind-al-reducee ort.oooooeoaaaaaaa.oooooooooaaaa
Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite markov decision processes. In UAI, volume 4, pp. 162-169, 2004.
Vincent Francois-Laver Yoshua Bengio, Doina Precup, and Joelle Pineau. Combined Reinforce- ment Learning via Abstract Representations. pp. 3582-3589. Association for the Advancement of Artificial Intelligence (AAAI), 2019.
Robert Givan, Thomas Dean, and Matthew Greig. Equivalence notions and model minimization markov decision processes. Artificial Intelligence, 147(1-2): 163-223, 2003.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Advances in Neural Information Processing Systems, pp. 2450-2462. 2018a.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 10122, 2018b.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with stochastic actor. arXiv preprint arXiv: 2018.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. arXiv preprint 2rXt::2010.02793, 2020.
Zhengyao Jiang, Dixing Xu, and Jinjun Liang. A deep reinforcement learning framework for the financial portfolio management problem. arXiv preprint arXiv: 10059, 2017.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model-based reinforcement learning for atari. arXiv preprint (:XXvv19033003744 2019.
Michael store, 2019.
Kechinov.
ecommerce URL
behavior
data
from multi category
httpp: http://www.k.ggg..g wwggle. wtw.eee..ooommkkcceeecee
eccnne
Yu Lei and Wenjie Li. Interactive recommendation with user-specific deep reinforcement learning. ACM Transactions Knowledge Discovery from Data (TKDD), 13:1-15, December 2019.
Lihong Li, Thomas Walsh, and Michael L Littman. Towards unified theory of state abstraction for mdps. In TSAIM, 2006.
10

---

Under review conference paper ICLR 2021
Feng Liu, Ruiming Tang. Xutao Li, Weinan Zhang, Yunming Ye, Haokun Chen, Huifeng Guo, and Yuzhou Zhang. Deep reinforcement learning based recommendation with explicit user-item interactions modeling. arXiv preprint arXiv: 1810. 12027, 2018.
Hongzi Mao, Mohammad Alizadeh, Ishai Menache, and Srikanth Kandula. Resource management with deep reinforcement learning. In Proceedings of the 15th ACM Workshop on Hot Topics in Networks, pp. 50-56, 2016.
Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean, Efficient estimation of word represen- tations in vector space, arXiv preprint 1exXiv:3301.320 2013.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprini 1rXiv:1312.5000 2013.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland. Georg Ostrovski, al. Human-level control through deep reinforcement learning. Nature, 51::599-533, 2015.
Jelle Munk, Jens Kober, and Robert Babuška. Learning state representation for deep actor-critic control. In 2016 IEEE 55th Conference Decision and Control (CDC), pp. 4667-4673, 2016.
Karthik Narasimhan, Tejas Kulkarni, and Regina Barzilay. Language understanding for text-based games using deep reinforcement learning. arXiv preprint arXiv:1506.0000 2015.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury. Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Veto,o, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high performance deep learning library. In Advances in Neural Information Processing Systems 32 (NIPS), pp. 8024- 8035. 2019.
Jason Pazis and Ronald Parr. Generalized value functions for large action sets. In Proceedings of the 28th International Conference on International Conference on Muchine Learning (ICML), pp. 1992, 2011.
Brian Sallans and Geoffrey E Hinton.
Reinforcement learning with factored states and actions. Journal of Machine Learning Research, ((Aug):1063-108) 2004.
Jürgen Schmidhuber. On learning think: Algorithmic information theory for novel combina- tions of reinforcement learning controllers and recurrent neural world models. arXiv preprint arXiv:1511.092. 2015.
John Schulman. Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal policy optimization algorithms. arXiv preprint 170.v:1707.0637. 2017.
Guy Shani, David Heckerman, and Ronen I Brafman. An mdp- mdp-based recommender system. Journal of Machine Learning Research, (2Sep):1265-122: 2005.
Sahil Sharma, Aravind Suresh, Rahul Ramesh, and Balaraman Ravindran. Learning to factor poli- cies and action-value functions: Factored action space representations for deep reinforcement learning. arXiv preprint arXiv: 170.00726, 2017.
Guy Tennenholtz. and Shie Mannor. The natural language of actions. volume 97 Proceedings of Machine Learning Research, pp. 6196-6226 6205, 2019.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physies engine for model-based control. 2012 IEFE/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033, 2012.
Hado Van Hasselt and Marco A Wiering, Using continuous action spaces to solve discrete problems. 2009 International Joint Conference on Neural Newworks (IJCNN). pp. 1149-1156, 2009.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006. 10742, 2020.

---

Under review conference paper ICLR 2021
A
APPENDIX
A.1 CLAIMS
In addition to Assumptions 1 and we derive three claims conditional independence from the definition of our embedding model and the Markovian property of the environment, listed below.
Claim 1. Action A, is conditionally independent of St given state embedding X
P(A. St s|X, =|x, P(A, a = P(S
s|X, =))
Claim 2. Action A, is conditionally independent of St given state and action embeddings X, and Eg:
s|X, = = e). P(A, a, s|X, = = -. e(,
=
Claim 3. Next state Si+1 is conditionally independent of E, and Xt given A, and S
0, =
& a. =) = =/4, a, a, =)
Using Claims T
we can derive two further auxiliary claims that be used the proof
Lemma 1
Claim 4. The probability of next state Si+1 independent of E4 and X given state and action
AG, i.e.
P(S++1 = = - e, I, A, S. = PP(S++1 8 8 S. =)-
Proof. From Claim
we have
e 2l P(5'aa, P(e,
P(s'le, a, P(e, =P(slla, s) P(e,
2, ())
Claim 5. The probability of action A; is independent of S₁ and Xt. given action embedding Es, i.e.,
P(A, = a|S, X, E, P s|E, c).
Proof. From Claim 2 we have
((ass)/,))=P(dae.))P(p P(a 8, x,.) e))
Since action only depends the action embedding in model, this becomes
P(a|s, e) =P(ale).
A.2 PROOF OF LEMMA
Lemma 1. Under Assumptions and 3y such that
for policy in the original problem domain, there exists
=
Σ
(. = o(s)) de
(a)
12

---

Under review conference paper ICLR 2021
Proof. The Bellman equation for a MDP is given by
Σ x(als) Σ
P(
aEA
d'E
where G denotes the return, iee. [R(,,)) + (s')l, which arranging terms we get
function of 8, a, and 8. By re-
1 (5(s) Σ E ((dls)P(d|1s,0)G
P(s', G P(s,a)
ΣΣ (Tals)
"EANES
ΣΣ <(als) s(s',
P(s)
aEAVES
P(als',s)PCe G.
ΣΣ AEA +E8
P(s)
Since can be deterministically mapped to z via x = o(s). by Assumption 2 we have
P(x,als', s(P(s
20 (8)
P(s)
aEA''SS
s(c
Σ
((s)))
acA
Σ
Σ
G
P(2,s)
aEAd'SS
P(z,a,,',s)
Σ
P((|s)
P(r,8)
aeA
NEE
Σ
P(a(s)
P(r,3)
Σ P(a[s)
Σ
P(a,&/1,,)GG
aEA
#'S
Σ P(r[s)) Σ P(6la...s)P(aaa.c...
e'es
From Claim we know that )(((s))))P(((() Therefore,
=
0. (5))) = Σ P(a|s) Σ
aEA
SES
Since P(x|s) is deterministic by Assumption 2 and evaluates to for the representation of d(s) = x, we can rewrite the equation above to
(T(()=
Σ Σ 2is/
13

---

Under review conference paper
ICLR 2021
We now proceed establish the relationship with the action embedding. From above we have
10 (5)
Σ
Σ
P(s.))
a(,',,,,,8) G. p(a,2,8)
Σ aEARSS
Σ
P(a(z)
From Claim and conditionally independent given P(alz)P(sle))P(t) (x))
Therefore, P(a,x,8)
which allows us rewrite the above equation
(5))) = E Σ P(a|x) P(s',a,
P(a)o)((s,8)
P(s,
P(s',,,I,s)
E
sEA''ES P(I,s)
P(als', s, P(s)) G. P(x,s)
ΣΣ
aEA
By the law of total probability, we can now introduce the new variable e, which is the embedded action. Then
(*())
ΣΣ BEA'TES
P(a, cla,cls'sa x.s',,,,)))(
r. 8) de
P(x,8)
Pla, els a P($ P(e|z,s) P(
Σ
Σ
p)
Ge.
s'
Since is uniquely determined by drop in P(elæ, 8), and thus
(4,,,,) Gd p(c,x,s)
1d (=((s) = Σ E
P(elx)
Σ Σ
"EA
ΣΣ R
Using the previously derived Claims
4 and 5
the above equation can be simplified to
o (s)
Σ aEAs'ES
Pledds))
Since the function f, mapping a, is deterministic by Assumption D and only evaluates to for particular and ( elsewhere, we can rewrite this further as
("()) Σ ΣΣ
((4(a s)G de.
REA+ES (a)
Summarizing the terms, this becomes
(8) Σ
(a(e)((s)))"
1-((a)

---

Under review conference paper ICLR 2021
S2
Environment
A
X2 7
E2
To
Figure 4: Illustration of x3. in the combined model.
A.3 PROOF LEMMA
Lemma Under Assumptions for all deterministic functions f and ó which map each point E & and E S to action A and to embedded state E X, and the internal policy parameterized by 0, the gradient of the internal policy's 's performance function 81,00) equals the
89
gradient of the overall policy's performance function 0J.66.811) i...
(J:(0) 0Jo(0,0,
00
Proof. Recall from Lemma I that the overall policy is defined using the internal policy
J , ((a)
<a(ass)
(c(c(((s) de.
We can then define the performance function of the overall policy using the internal policy
Σ do(s) Σl REAA
J. 0,
((eld(s))( de.
The gradient of this performance function w.r.t. 0
@J. f f)
do(s)
Σ AEA"
(s..
(9)
00
a0
(())
SES
Now consider another policy n/ as illustrated in Figure ie.. policy that takes the raw state input
S, and outputs E, in embedding space. By Assumption (s, = Q" (s, as is determin-
istically mapped to However, Q" e) E.-. ((x) Q (s, Enn 1(z) (s, a)]. where
a d-1(-) (-) is the inverse mapping from æ to s. Nevertheless, by Assumption 3 E E (w))" 1QT (s. a)]
* a). Therefore, we can replace the overall (a, Equation (9) by the internal
Q (o(s),e) e) as follows
aJ.(0, f,
8
Σ
do(s)
Σ aEA (a)
(T((el(())) de.
00
00
SES
Then the deterministic mapping of to a by function ff further allows us to replace the integral over f (a) and the summation over A with an integral over e. Moreover, by Assumption 3 we have (s)s) do(a (())). Therefore, the above be rewritten
0J.(0.0.1)
do(x)) /
JeES =e(el=)0*) de
08
09
(.,(())
15

---

Under review conference paper ICLR 2021
A.4 PARAMETERIZATION
A.4.1 IMPLEMENTATION AND ENVIRONMENTS
implement evaluated models Python and use the PyTorch library neural network plementation (Paszke et al. 2019). The implementation for the benchmark models is adapted from OpenAl's SpinningUp package (Achiam 2018). In the following, elaborate on the different environments used for experimentation.
Gridworld The gridworld environments are constructed using a x grid to represent the un- derlying continuous state space. For the continuous state version of the environment, a continuous coordinate with the position of the agent in this grid is used the state. In the discrete state version, we discretize this continuous coordinate ated from n actuators, equally spaced around the agent, which move the agent in the direction they pointing towards. Then each combination of actuators forms unique action, resulting in 2 unique actions. An illustration this environment can be seen Figure [ We limit the maximum number of steps the agent can take in an episode to 150 and set the maximum step size to 0.2. The agent receives positive reward of 1 for reaching the goal step and receives small penalty of -0.01 for each step taken and an additional penalty of -0.05 for colliding with an obstacle or attempting to leave the grid.
40 40 grid, Actions these environments are gener-
The slotmachine environment is defined using 4 reels, with 6 values per reel, ran- domly sampled from the range between and 8. Note that sample these reels and then keep them fixed across all experiments to ensure that these run the same environment. The action space is defined by the fraction of a full turn for which each reel is spun. We set the maximum fraction to 0.8 and use increments of 0.1 set the fraction of full turn. Since each combination of fractions forms an action, we get 4, 096 actions total. The maximum number of steps per episode is set to 20. Rewards are calculated according to
Slotmachine
R
Σ 1,123
1s x ( 2 m).
where et denotes the value the th reel, n, the number of times e, occurred across all reels, m is the minimum number of times it has to occur to result in non-zero reward, and 1 (.) is an indicator function. In experiments,
the threshold
3.
Half-Cheetah and Ant-v2 For the experiments in a robotic control environ- ment, we use the physics engine Mujoco and use the Half- Cheeethha and Aur-v2 environments imple- mented therein (Tooorvv al.) 2012). We limit the maximum number of steps per episode to 250 for the Half-Cheetah environment and 1000 for the Ant-v2 environment.
Robotic Control:
Real-world Recommender System The environment model for the recommender system con- structed from publicly available data set containing data on approximately 43 million user inter- actions collected from multi-category e-commerce store (Kechinov) 2019). We only use purchase
Figure 5: Gridworld environment as used for evaluation The goal state is in blue, the agent is in red, and the obstacles shown in gray.
16

---

Under review as conference paper ICLR 2021
events to create our environment model, where we filter out items that were bought fewer than 100 times and users that made less than 3 purchases. We are then left with approximately 250, 000 data points and 835 items. The environment model is then created using an n-gram based model, with n set to The state is represented as the concatenation of the last items purchased and each item available for recommendation forms an action. We thus have 835 actions and 697, 225 possible states. In the environment model, we compute the purchasing probability for user for all items, given the current state, i.e. the last purchases that user made. We parameterize a categorical distri- bution with these transition probabilities to sample the next item. The RL agent choses which item to recommend. Recommending an item increases its transition probability by factorr = 3, after which the sum of all transition probabilities for the current state are normalized to sum to The transition function in this environment is stochastic as we have transition probabilities for each item and each state. For more details on the environment model used, please refer to Shani et al. (2005). whose methodology adopted. As the reward function, we simply use the revenue obtained each time step, i.e., the price of the item that the customer bought after making a certain recommendation.
A.4.2 HYPER PARAMETERS
Training for all experiments is conducted in epochs, where each epoch corresponds a set of sam- ples. At the end of an epoch, the actor and critic are trained for a number of iterations using their respective learning rates. For the VPG and PPO algorithms, the actor network is trained for 3 itera- tions and the critic network trained for 6 iterations the end of an epoch in all of the conducted experiments, using a single batch in each iteration, consisting of all samples collected during the epoch. The SAC updates are also run at the end of an epoch, where we train the actor and critic on i batches sampled from the replay buffer, where is set equal to the number of samples collected during the epoch. For all experiments, training is conducted samples are collected per epoch. Rather than updating the models after every timestep or episode, we run several iterations of gradient descent at the end of epoch (see the reference implementa- tions in Achiam (2018)). The batch size is searched over {256, 512) in all SAC experiments. The environment configuration is kept fixed for each environment across different models as well as the hyper parameter searches to ensure comparability.
epochs, where a fixed number
We conduct gridsearch over different parameter combinations for each experiment and report the average per episode return and standard deviation achieved over 10 runs with different random seeds, using the hyper parameters that perform the best in the gridsearch. The best hyper parameters apart from those we kept fixed - are reported in Tables I 7 The actor and critic learning rates are searched over {0.001, 0.0003, 0.0001} for all environments. We also search the initialization value for the standard deviation of the actor over {0.3, 0.6, 0.8. 1.0} in all experiments and keep y = 0.99 and 0.97 fixed. Note that À only used for general advantage estimation in the VPG algorithm. For all experiments apart from those using the recommender system and the Ant-v2 envi- ronment, we pre-train the embedding model using 30, 000 randomly collected samples and keep the embeddings fixed thereafter. For the recommender system, we pre-train the embedding model using 50, 000 randomly collected samples andalso allow continuous updating of the learned embeddings. Similarly, for the Ant-v2 environment, collected samples and then enable continuous updating. In all cases, we apply a Tanh non-linearity to the learned embeddings to bound their effective range. Throughout our experiments, we keep the learning rate of the embedding model fixed at 0.01. For continuous domains, we parameterize the mapping function f as a neural network with a single layer with 64 hidden units. In discrete domains, we use simpler nearest neighbor mapping function to parameterize f. i.e., we map the output of the internal policy n, to the action that is closest in embedding space E.
pre-train the embedding model using 100, 000 randomly
For experiments the gridworld environments, we fix the state and action embedding dimensions 8 and 2, respectively. The network size for the actor and critic is set to (64, 64), i.e., a two-layer network with 64 neurons in each layer. We then train the agent for 700 epochs, with 505 steps, i.., minimum of 10 episodes per epoch.
In the experiments on the slotmachine environment, the dimension for state and action embeddings is searched over (8, 16) and {2, 4), respectively. The network size for the actor and critic is kept fixed at (64. 64). The agent is trained for 2, 500 epochs, with 200 samples (a minimum of 10 episodes per epoch).
17

---

Under review conference paper ICLR 2021
Parameter
VPG discrete
(5)
VPG disc.
(9)
VPG disc.
(11)
VPG
(5)
VPG
(9)
VPG (II)
cont.
cont.
cont.
Critic
0.001
0.0003
0.001
0.001
0.001
0.0001
Learning
Rate
Actor
0.001
0.001
0.001
0.001
0.001
0.001
Learning
Rate
Actor Std.
0.8
0.6
0.8
Dev.
Table I; Best hyper parameters for the gridworld experiments without embedding.
Parameter
VPG-
VPG-
VPG-
VPG-
VPG-
VPG-
JSAE
(5)
JSAE
(9)
JSAE (1I)
JSAE (5)
JSAE
JSAE (11)
disc.
disc
disc.
cont.
cont.
cont.
Critic
0.0003
0.001
0.001
0.0003
0.001
0.001
Learning
Rate
Actor
0.0003
0,0001
0.0003
0.0001
0.0003
0.0003
Learning
Rate
Actor
Std.
0.8
1.0
1.0
0.3
1.0
1.0
Dev.
Table 2: Best hyper parameters for the gridworld experiments with embedding (ISAE).
experiments on Half-Cheetall environment. the state and action embedding dimen- sions are searched over {10,20} and {6, 10), respectively. We search the network size over {[64, 64), (128, 128)) and train the agent for 1, 000 epochs with 1, 500 steps per epoch, ie., least episodes per epoch.
dimensions of state and action embed- dings are searched over {16, 32, 64] and {6, 8, 10), respectively. We search the network size over [(64, 64), (128, 128)) and train the agent for 2, 000 epochs with 3, 000 steps per epoch, i.e., min- imum of episodes per epoch. The embeddings are pre-trained using 100, 000 randomly generated samples. Continuous updating for the leared embedding model is enabled for the Ant-v2 environ- ment.
the experiments on Ant-v2 environment,
On the recommender system environment, we evaluate Proximal Policy Optimisation (PPO) and Soft Actor Critic (SAC) algorithms in addition to the previously used Vanilla Policy Gradient (VPG). For the experiments using the SAC algorithm, we search the actor parameter which corresponds an peploration-epplitaaeee rithm, we set the gradient clipping ratio to 0.2 and keep this fixed throughout. Embedding di- mensions for states and actions are searched over [8, 16, 32 and 8, 12). respectively. In all experiments, the embedding functions are pre-trained using 200, 000 randomly generated samples. contrast to experiments on other environments, we enable continuous updating these Iearned embeddings for the recommender system. We search the size of the actor and critic network over [644, 64), (128, 128)) and train the agent for I, 000 epochs with I, 000 steps per epoch. The episode length is set to 20 steps.
trade-off over {0.1, 0.3). For experiments involving a PPO algo-
18

---

Under review conference paper ICLR 2021
Parameter
VPG-RA (5) disc.
VPG-RA (9) disc.
VPG-RA
VPG-RA
VPG-RA
VPG-RA
(11) disc.
(5) cont.
(9) cont.
(11) cont.
Critic
0.0003
0.001
0.0003
0.0001
0.0003
0.0001
Learning
Rate
Actor
0.0003
0.0003
0.0003
0.001
0.0003
0.0003
Learning
Rate
Actor Std.
0.8
0.6
0.6
0.8
0.8
1.0
Dev.
Table 3: Best hyper parameters the gridworld experiments with embedding (RA).
Parameter
VPG
VPG-ISAE
VPG-RA
Critic Learning Rate
0.001
0.0001
0.001
Actor Learning Rate
0.0003
0.0003
0.0003
State
Embedding
16
Dim.
Action Embedding
4
2
Dim.
Actor Std. Dev.
0.6
Table 4: Best hyper parameters for the slotmachine experiments.
Parameter
VPG
VPG-JSAE
Critic Learning Rate
0.0003
0.001
Actor Learning Rate
0.0003
0.0003
State Embedding Dim.
20
Action Embedding Dim.
10
Actor Std. Dev.
0.6
0.6
Actor Hidden Units
(I28. 128)
(I28, 128)
Critic Hidden Units
(128, 128)
(128, 128)
Table Best hyper parameters for Half-Cheetah experiments.
Parameter
VPG
PPO
SAC
VPG-
PPO-
SAC-
JSAE
JSAE
JSAE
Critic Learning Rate
0.001
0.0003
0.001
0.001
0.001
0.001
Actor Learning Rate
0.0001
0.0003
0.001
0.001
0.001
0.001
State Embedding Dim.
64
32
64
Action Embedding Dim.
8
8
Alpha
0.1
0.1
Actor Std. Dev.
0.6
0.8
0.6
0.6
0.8
0.8
Batch Size
256
Actor Hidden Dim.
(128,
(128,
(128.
(128,
(128,
(128,
128)
128)
128)
128)
128)
128)
Critie Hidden Dim.
(128,
(128.
(128.
(128.
(128.
(128,
128)
128)
128)
128)
128)
128)
Table 6: Best hyper parameters for the Ant-v2 experiments.
19

---

Under review conference paper ICLR 2021
Parameter
VPG
PPO
VPG-
PPО-
SAC-
VPG-
PPO-
SAC-
ISAE
ISAE
JSAE
RA
RA
RA
Critic Learning Rate
0.001
0.0001
0.0003
0.001
0.001
0.0001
0.0001
0.0001
Actor Learning Rate
0.001
0.001
0.001
0.001
0.001
0.001
0.001
0.0001
Alpha
0.1
0.1
State Embedding Dim.
Action Embedding Dim.
8
8
8
8
8
Actor Std. Dev.
0.8
0.8
0.8
0.6
0.6
0.6
Batch Size
256
256
Actor Hidden Dim.
(64,
(64,
(64.
(64.
(64,
(64,
(64.
(128.
64)
64)
64)
64)
64)
64)
64)
128)
Critic Hidden Dim.
(64.
(64.
(64.
(64.
(64,
(64,
(64.
(128.
64)
64)
64)
64)
64)
64)
64)
128)
Table Best hyper parameters for the recommender system experiments.
20


Reviews:
[{"review_id": "uQpA46375Zz", "reviewer": "AnonReviewer2", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Learning on environments with large state-action spaces can be difficult. This paper addresses this issue by learning a joint state-action embedding and learn an internal policy(\\pi_i) on this embedded state-action space instead of the original state-action space. There are three parts of learning, 1. learning the embedding model that learns mapping from state to state embedding, 2. learning the internal policy, and 3. learning the mapping from action embedding to action space. The authors justify this approach by showing that the overall policy (\\pi_o) can be expressed in terms of the internal policy (\\pi_i). Furthermore, there is equivalence between the internal state-action-value function and overall state-action-value function and the authors show that updating \\pi_i is equivalent to updating \\pi_o. \n\nThe benefit of learning a policy on these joint state-action embedded space is that any policy gradient algorithm for continuous control can be used regardless of whether the original state-action is discrete or continuous. The authors claim that learning on this joint state-action embedded space is especially helpful in large discrete state-action spaces because relationships between state and action are often not clear with discrete representations. Their algorithm is compared on several environments highlighting this benefit, and also on several other benchmark domains.\n\nOverall, the paper is well-written with an interesting approach to tackle the learning problem in large state-action space. They propose learning an internal policy backed by theoretical proof that show this is equivalent. Their results especially on the gridworld domain and slotmachine showcase the purported benefits of using their joint state-action space embedding (JSAE). I would overall recommend a weak accept, and I think strengthening the experiment part would make this a much better paper. \n\nFirst, most of the experiments focus on VPG compared to PPO and SAC. While this is okay, there isn’t any citation on VPG or explanation of its exact implementation to help the readers follow what was actually used. Furthermore I think a justification on using VPG would be very helpful. Was it because it is a simple algorithm without complex structure? No replay buffer is necessary? \n\nSecond, I think there could be more discussion and insights into the various hyperparameters for JSAE. Figure 2a and 2c really highlights the benefit of using JSAE for large discrete state-action space domains. The authors hypothesize that it’s because the joint embedded space helps learn the relationship between states and actions. But were there any effects due to reduced embedded space? Would having a larger embedded space help capture the state-action relationship better? The experiment fixes the state and action embedding dimensions to a small value for the grid environments, and also only sweep a small range for other environments. A discussion on how these sweep ranges were chosen (e.g. was having small embedded space always helpful?), and other observations that the authors found while running the experiments would improve the paper a lot. \n\nFor other comments (won’t affect my decision but only for suggestions for improvement), I don’t think the experiment on HalfCheetah comparing VPG and VPG-JSAE shows much other than the fact that JSAE does not harm performance. SAC is known to have great performance on these continuous mujoco domains and it would be better to use SAC as the baseline and compare SAC-JSAE. HalfCheetah is also known to be one of the easy environments in Mujoco; I would be curious to know whether SAC-JSAE shows large improvements in high dimensional domains like Humanoid. Perhaps JSAE may capture a better state-action relationship in its embedded space and learn faster.\n\nI think also reporting the best hyperparameters found for each agent and experiment would be helpful as well.", "extra_scores": {}, "extra_text": {}}, {"review_id": "cWZT7-yYXJv", "reviewer": "AnonReviewer3", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper presents a method to take advantage of joint state-action representation. It proposes to use an environment model to obtain embeddings for state and action and then such representation should enable better generalization across state-action space. Experiments on several gaming and recommendation systems are conducted to show the superior performance of the proposed method. \n\nNovelty. The idea is not new but the particular method of defining a policy in the embedding space seems to be new. \n\nSignificance. The work may be of general interest to the reinforcement learning research community.  \n\nClarity. The paper presents the key idea/method clearly. \n\nQuality. There are a few theoretical and empirical issues in this work. \n\nFirst, one action can correspond to many embeddings and one embedding can correspond to only one action. Shouldn’t it be the case that one embedding can correspond to multiple actions? Intuitively, one wants to enjoy generalization across different actions. And generalization ability is also one of the goals mentioned by the authors. \n\nSecond, I do not see why the assumptions can get satisfied in practice. \n\nThird, based on the two assumptions, lemma 1 and theorem 1 are not that interesting. Furthermore, the bottom line says that theorem 1 indicates one can focus on the optimization of the overall policy \\pi_0; however, theorem 1 indicates only the existence of \\pi_0 equal to the optimal policy in the original MDP, it is unclear if optimizing \\pi_0 according to Alg 1 can really lead to the optimal one. \n\nMotivation. I think this is another serious concern of this paper. Although I agree with the general direction of learning some sort of joint embedding of state-action pairs, I am not persuaded by the motivation of using the method proposed by the authors. Many works are indeed using state-action representation, especially on those continuous control problems. Whenever taking both state and action as input for the critic/value network, one can think of the final hidden layer as the learned joint state-action representation. Why such simple way cannot be used? At least, those should be compared in the experiments. \n\nExperiments. In general, the experimental results are not really strong. Only Fig 2(a)(c) show a clear advantage of using joint representation. Given the weak motivation of the proposed method and the proposed theory, I think the most important experiment should be designed to show the method is indeed superior to other intuitive baselines. As I mentioned, one baseline is to simply take both state-action as input. Another intuitive baseline is to simply learn state-action representation by using an environment model first and then based on the learned representation to learn a policy. To enable the separability of the state and action representations, one can use the pairwise product to generate the joint representation. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "QX-tXplM_K", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 3, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\nThe paper proposes a method to jointly learn: (a) a latent state embedding; (b) a latent action embedding; (c) a state transition model; and (d) an RL policy.  The latent models should allow for better generalization over states and actions, and therefore result in improved learning, particularly for discrete action domains. The method shows improved performance over vanilla policy gradient on a grid-world task, a slot machine task, a recommender system, and half-cheetah locomotion.\n\n### Strong points\n- important problem: learned latent models of the world are key to making progress in RL\n- evaluated on several tasks, with improvements on some\n- it is surprising that the all of the components can be trained, in a stable fashion, without further regularization or conditioning\n\n### Weak points\n- there is a large and growing literature on leveraging learned latent models of all kinds, particularly for continuous systems. E.g., world models and many variants. It is unclear to this reader how to situate this work in relation to those.\n- the evaluation is weak:  two of the main examples are toy proof-of-concept; half-cheetah shows no benefit; the recommender system needs to be considered in the expansive volume of recommender system algorithms\n- the method is only compared against vanilla policy gradient for a number of the results\n\n### Recommendations\nCurrently recommend to reject.  The approach needs to be discussed and evaluated in the context of other latent \"world models\", and to show benefits on more challenging problems. I also currently remain unclear regarding the potentially underdetermined nature of learning all the given components in the absence of further regularization or constraints.\n\n### Questions\nQ1. 4.3.1 component (iii): function g:  How is g well-posed, given that it is inverting a possibly many-to-one mapping?\n\nQ2. Eqn (5):  Given that this is the only loss function that involves three of the models, it is unclear to me how they are fully determined in the absence of further regularization, or additional loss terms.\n\nQ3. Algorithm 1:  Where are the models for g and T updated?\n\nQ4: uniqueness of state embedding:  Why would this emerge as a property? How has this been tested emperically? Doesn't this also mean that some of the benefits of the latent space are lost, if working from possibly-redundant state observations?\n\nQ5:  Figure 2e has comparisons with PPO and SAC. How would these algorithms fare for the other problems? \n\n### Additional feedback\n\nThis reader found the phrase \"joint action-state embedding\" to be ambiguous.  Upon first reading, I interpreted it as learning a unique embedding for the state-action space, e.g., for Q(s,a) for example. Instead, it refers to separate state and action embeddings, which are jointly trained during learning, along with the policy. \"Jointly-trained action and state embeddings\" would clarify this ambiguity.\n\nFigure 1b: label the transition model, T\n\n4.2 Assumption 1: Given an _action_ embedding\n\nConcluding sentence for the paper:  this makes a very strong statement that is not really supported by the results.", "extra_scores": {}, "extra_text": {}}, {"review_id": "FVilI55OMt", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper explores the idea of using state and action embeddings for more scalable learning. The idea has merit so I encourage the authors to continue working in this direction, but unfortunately there are a number of technical issues with the paper that I detail below.\n\nThe main issues for me are:\n1. **There are a lot of relevant references missing.** With regards to planning over latent spaces here are some:\n  - \"Recurrent World Models Facilitate Policy Evolution\", David Ha and Jürgen Schmidhuber, NeurIPS 2018.\n  - \"Model-Based Reinforcement Learning for Atari\", Lukasz Kaiser et al. ICLR 2020\n  - \"Learning Latent Dynamics for Planning from Pixels\", D Hafner, T Lillicrap, I Fischer, R Villegas, D Ha, H Lee, J Davidson, ICML 2019\n  - \"Mastering Atari with Discrete World Models\", D Hafner, T Lillicrap, M Norouzi, J Ba, arXiv 2020\n  \n  The authors seem to be learning equivalence relations over the states and actions in MDPs (this is basically the result of Assumption 2). There is a large body of work on this, and here are some important references:\n  - Givan, R.; Dean, T.; and Greig, M. 2003. \"Equivalence notions and model minimization in Markov decision processes\". Artificial I Intelligence 147(1-2): 163–223.\n  - Ferns, N.; Panangaden, P.; and Precup, D. 2004. \"Metrics for finite Markov decision processes\". In Proceedings of the 20th conference on Uncertainty in artificial intelligence, 162–169. AUAI Press.\n  - Li, L.; Walsh, T. J.; and Littman, M. L. 2006. \"Towards a Unified Theory of State Abstraction for MDPs. In ISAIM.\n  - Taylor, J.; Precup, D.; and Panagaden, P. 2009. \"Bounding performance loss in approximate MDP homomorphisms\". In Advances in Neural Information Processing Systems, 1649–1656.\n  - Castro, P. S.; Panangaden, P.; and Precup, D. 2009. \"Notions of state equivalence under partial observability\". In Proceedings of the 21st International Joint Conference on Artificial Intelligence (IJCAI-09).\n  - Castro, P. S. 2020. \"Scalable methods for computing state similarity in deterministic Markov Decision Processes\". In Proceedings of the AAAI Conference on Artificial Intelligence.\n  - Zhang, A.; McAllister, R.; Calandra, R.; Gal, Y.; and Levine, S. 2020. \"Learning Invariant Representations for Reinforcement Learning without Reconstruction. arXiv preprint arXiv:2006.10742 .\n\n2. **Lemma 1 is not correct.**\n  - Consider the following simple example: an MDP with two states ($s, t$) and two actions ($a,b$). Action $a$ always transitions deterministically to state $t$, while action $b$ always transitions deterministtically to state $s$. In state $s$ taking action $a$ gives a reward of 1 while taking action $b$ gives a reward of 0; in state $t$ taking action $a$ gives a reward of 0 while taking action $b$ gives a reward of 1. The optimal policy is to take action $a$ from state $s$ and action $b$ from state $t$, resulting in $V^*(s) = V^*(t) = \\frac{1}{1-\\gamma}$. Now, take $X = \\lbrace x_1, x_2\\rbrace$ where $\\phi(s) = x_1$ and $\\phi(t) = x_2$ (this satisfies Assumption 2). Now take $E = \\lbrace e\\rbrace$ where $f(e) = a$, so $f^{-1}(a) = \\lbrace e\\rbrace$ (this satisfies Assumption 1). Letting $g:A\\rightarrow E$ be the action embedding function, we can see that $g(a) = g(b) = e$, which means that any internal policy will be suboptimal, as it will choose the same latent action for both $x_1$ and $x_2$.\n  - Part of the issue may lie in the way the proof is structured. The authors start by using $G$ which they define as \"denotes the return, which is a function of $s$, $a$, and $s'$\". However, the way they've included it in the definition of $v^{\\pi}$ seems like it does _not_ depend on $s'$ and is really just $Q^{\\pi}(s, a)$. This seems to be the case as that is where the authors end up at the end of the proof. If it _is_ supposed to depend on $s'$, then $G$ likely needs to be decomposed into the one-step reward $R(s, a)$ and the expected value at the next state $\\mathbb{E}V^{\\pi}(s')$.\n  - The other bug I found in the proof is at the top of page 12. The authors went from $P(a | s, a')$ at the end of page 11 to $P(a, x | s, a')$ at the top of page 12, which is most certainly not a valid equality.\n  - (Minor) The jump from $P(s, a)$ to $\\pi(a | s)P(s)$ is probably ok, but requires a little more justification.\n\n3. **The losses are incorrect**. Below equation (4), the authors state \"the denominator does not depend on $\\phi$, $g$, or $T$, and so they get rid of the $P(S_{t+1}|S_t, A_t)$ term in their loss, but this means there is no target for the loss! This means that the loss in equation (5) is just trying to maximize the probabilities uniformly, independent of whatever the true probabilities really are. This problem is also present in the loss in equation (6).\n\nMedium issues:\n1. The authors make a number of claims about their method working for continuous spaces, but this requires more details than simply replacing integrals for sums. For example, in the background the transition function T needs to be defined with respect to Borel sets, the summation in Lemma 1 doesn't \"just work\" by switching to integrals, etc..\n2.  Given that the authors motivate their work by claiming that existing algorithms don't work well outside of \"simple\" tasks, they should include larger scale experiments than the ones they are currently including, which are rather small.\n3. Doesn't Assumption 3 defeat the purpose of dealing with large state spaces? It's basically just converting the original state space into an isomorphic one.\n\nMinor issues:\n1. In the Introduction the authors refer to the ALE as \"comparatively simple tasks\", but they are not simple and are in fact more difficult than any of the environments evaluated in the paper.\n2. In equation (3) it should be $\\hat{T}(X_{t+1} | X_t, E_t)$ instead of $\\hat{T}(S_{t+1} | X_t, E_t)$\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "BxhtWu1DHr", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper proposes a framework of jointly learning a state and action embedding using the model of the environment, eventually using those embeddings to learn a parameterized control policy using standard policy gradient (PG) methods. Joint learning of state and action embeddings allows us to capture the interactions between actions in different states. The framework proposes to learn an internal (embedding) policy, a state embedding, an inverse function on action embeddings, combining all the parts to form an overall policy. The paper theoretically shows that optimizing the internal policy leads to an optimal overall policy. \n\nThe idea presented in the paper is a natural extension to work done in Chandak et al. (2019). The idea of jointly optimizing for the state and action embedding, as opposed to just separately optimizing the action embedding as in Chandak, is reasonable and warrants investigation. As is, however, the paper needs a bit more work for three primary reasons. Firstly, the embedding for states and actions involves a joint optimization, but the embedding itself seems to be separate for the two. The relationship between how the state embeddings influence the action embedding is unclear. The inverse embedding function ‘f’ does not seem to account for the state when inverting the embeddings, and the f seems to be a global unembedding (decoder) across all states. The properties of the learned embedding could be better explained. \n\nSecond, the experiments do not clearly highlight why the joint state and action embedding might help learn a good policy. For example, Chandak et al. (2019) present experiments, where they visualize the benefit of using action embeddings as part of the internal policy. Including experiments which show meaningful relationships between states and their corresponding action embedding will add soundness to the benefits of the proposed framework. \n\nFinally, the empirical results are inconclusive about the benefits of JSA. Five random seeds are not adequate to draw relevant conclusions given there is extensive overlapping between standard errors.  As pointed out in Henderson et al. (2017), random seeds vastly influence the performance of methods. \n\nComments:  \n1. Page 1, Para 3 mentions the proposed method to bridge the gap between model and model-free learning. This needs clarification, as in model-based approaches, the model is employed for planning, whereas, in this case, it's used to learn embeddings. \n\n2. Section 2 (background) : The distributions in this work seem to be defined as functions to the set, rather than as a distribution. For example, the transition function is defined as probability of transition to the next state from current state and action, but the terminology defines it as a deterministic function i.e. \\mathcal{T} : \\mathcal{S} \\times \\mathcal{A} \\rightarrow \\mathcal{S}. Either define it as deterministic or use \\mathcal{T} : \\mathcal{S} \\times \\mathcal{A} \\times \\mathcal{S} \\rightarrow [0,1]. In Equation (3), these are then used as distributions, rather than deterministic functions.\n \n3. If the space permits, it is good to include an example where a joint state-action representation will help in learning the policy. \n\n4. In Theorem 1, the optimal value function is defined as a function of \\pi_i and Q^*. It would be more clear to use separate notation to represent the optimal set of internal policies, maybe \\pi_i^*.\n\n5. Looking at Figure 2(a,b,d) JSA seems to have worse asymptotic performance. This could be discussed.\n \n6. The following missing citation looks relevant: “Reinforcement Learning with Function-Valued Action Spaces for Partial Differential Equation Control”, Pan et al., 2018. (see reference 3 below)\n\n7. It's good to see many of the experimental details mentioned. I would recommend including the following details:  \n\ta. Size of batch and mini-batch used.\n\tb. Epsilon value, for importance sample clipping in PPO.\n\nQuestions:\n1. Algorithm 1 mentions online training of the transition model, but the Appendix notes the use of pretraining the model for embeddings. Which one is used? \n\n2. In the case of the recommender system experiments, it is unclear to me how it is an RL problem i.e. it seems that the agent is defining the probability transition function, which would be wrong. The setting seems more relevant to pose as a contextual bandit. Also, more details on the reward function would be helpful; it seems like the agent can earn a high reward, even for wrong predictions, given the user is purchasing expensive items. \n\n3. Gridworld experiments: There are some details which are missing from the environment. For example, what is the exact reward function? Looking at the plots from Chandak et al. (2019), the reward scaling is near 100, and in the current paper figures for Gridworld rewards are near 1. \n\n4. When you say you train agent for 700 epochs with 1500 steps, do the steps refer to gradient descent steps and does 700 epoch mean that there 700 * 10 = 7000 episodes at a minimum? I was unsure about the connection between epochs and the number of episodes. \n\n5. How are the state and actions fed into the transition model (e.g., are they concatenated, etc.)?\n\nReferences \n1. Chandak, Y., Theocharous, G., Kostas, J., Jordan, S., & Thomas, P. S. (2019). Learning Action Representations for Reinforcement Learning. http://arxiv.org/abs/1902.00183\n\n2. Henderson, P., Islam, R., Bachman, P., Pineau, J., Precup, D., & Meger, D. (2019). Deep Reinforcement Learning that Matters. http://arxiv.org/abs/1709.06560\n\n3. Pan, Y., Farahmand, A., White, M., Nabi, S., Grover, P., & Nikovski, D. (2018). Reinforcement Learning with Function-Valued Action Spaces for Partial Differential Equation Control. http://arxiv.org/abs/1806.06931\n\n----- Update\n\nThank you for the update and response. Unfortunately, some of my concerns remain. The plots are now run with 10 runs, rather than only 5 runs in Figure 5. But, they look almost identical (in some cases, maybe they are identical?). That is not possible, unless there is a potentially invalid choice in the experimental design. If nothing else, the standard error should change. \n\nThe theory itself has some utility, since it is shown that learning in embedding space is equivalent. This is not surprising, considering it is assumed that there is a one-to-one mappings, but it’s good to be thorough. Nonetheless, this could maybe be shown more simply, and I am not sure Lemma 2 is exactly correct.\n\nLemma 1 is overly complex.\n\nAlternative proof:\n\nAssume pr(a | s) is pi(a | s) (i.e., action probabilities given s are defined under pi).\n\nvpi(s) = sum_a pr(a | s) qpi(s,a) = sum_a int_e pr(a, e| s) de qpi(s,a) = sum_a int_e pr(a | e, s) pr(e | s) de qpi(s,a) = sum_a int_e pr(a | e) pr(e | s) de qpi(s,a) (also using Claim 5 like they do) = sum_a int_{f^{-1}(a)} pr(e | s) de qpi(s,a) = sum_a int_{f^{-1}(a)} pr(e | phi(s)) de qpi(s,a)\n\nLemma 2 claims to show that the gradients are equivalent, but instead it seems to show that the functions themselves are equivalent and so should maybe be stated that way. Gradients are just placed in front of everything. Further the last step replacing d_0(s) with d_0(x) seems incorrect, as s is from a discrete space and x from a continuous space. Maybe you are suggesting that d_0 is some kind of delta distribution, but then it might be better to just sum over the same set of s.\n\nI am also a bit unsure about any smoothness assumptions required. Is J_0 even differentiable in theta? The requirements on the one-to-one mappings between discrete state to continuous state make for a piecewise flat function that could be problematic for such gradients.\n\nI also appreciate that Figure 2 was added. But, it is a bit hard to interpret. More explanation is needed there.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Most reviewers believe that the paper is not ready for publication. Among their concerns are:
- whether the new experiment with 10 runs are conducted correctly,
- the significance of the theoretical part,
- correctness of Lemma 2,
- generalization claims may not follow from the theoretical results,
- comparison with Zhang et al. (2020).

Given these and the lack of support from reviewers, unfortunately I cannot recommend acceptance of this paper at this stage. I encourage the authors to improve their paper according to these concerns.

I copy-paste some of the comments that came after Nov. 24th. The authors might want to use them to improve their work.
---
First, given the assumptions, the theorems/lemmas are not sufficient to be a solid contribution. I think what important is, can the Alg 1 lead to the optimal policy? More specifically, I do not doubt lemma 2; I am concerned about if updating the representations in an online manner (it should be highly nonstationary) can result in the optimal policy. Some two-time scale analysis may address this question. As an additional note, since f can be a many-to-one mapping, policy pi_i is a multimodal distribution. I am unsure if the authors consider this during implementation.

Second, without the two-time scale analysis, I would give weak acceptance if the authors can persuade that the superior performances are indeed due to the proposed jointly embedding learning method. That's why I ask for a baseline that directly considers environment model learning as an auxiliary task.

In the abstract, the authors state that "In this work, we propose a new approach for jointly learning embeddings for states and actions ...," in fact, this is not new. Almost any deep RL algorithm can be thought of as the process of learning state and/or action representations. In the response, the authors say, "We do not believe ..., as ... are embedded into the same space." How to do embedding is more like an implementation issue; one can encode them into different spaces and learn them by learning an environment model. It is nothing fancy/novel.

I consider this paper's main novelty to learn the optimal policy in the embedded state and action spaces, and the embeddings are learned by environment model learning. Thus, the authors need to have strong evidence to persuade people: 1) using an environment model to learn the embeddings is really useful; 2) a separate process of learning the policy in the embedded spaces is essential. Such evidence is necessary to make this paper a solid contribution.

Learning an environment model has been used as an auxiliary task in deep RL. Using such a baseline is to validate that it is necessary to learn the policy in the embedded space separately. The authors should also actively design other baselines to substantiate their claims.
Decision: Reject