Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 6: Marginally above acceptance threshold
review_content: 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. 

The 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.

Overall, 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. 

First, 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? 

Second, 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. 

For 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.

I think also reporting the best hyperparameters found for each agent and experiment would be helpful as well.
review_title: Interesting work

final_score: rating: 3: Clear rejection
review_content: ### Summary
The 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.

### Strong points
- important problem: learned latent models of the world are key to making progress in RL
- evaluated on several tasks, with improvements on some
- it is surprising that the all of the components can be trained, in a stable fashion, without further regularization or conditioning

### Weak points
- 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.
- 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
- the method is only compared against vanilla policy gradient for a number of the results

### Recommendations
Currently 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.

### Questions
Q1. 4.3.1 component (iii): function g:  How is g well-posed, given that it is inverting a possibly many-to-one mapping?

Q2. 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.

Q3. Algorithm 1:  Where are the models for g and T updated?

Q4: 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?

Q5:  Figure 2e has comparisons with PPO and SAC. How would these algorithms fare for the other problems? 

### Additional feedback

This 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.

Figure 1b: label the transition model, T

4.2 Assumption 1: Given an _action_ embedding

Concluding sentence for the paper:  this makes a very strong statement that is not really supported by the results.
review_title: possibly has merit, but limited impact with the current evaluations

final_score: rating: 4: Ok but not good enough - rejection
review_content: 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. 

The 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. 

Second, 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. 

Finally, 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. 

Comments:  
1. 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. 

2. 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.
 
3. If the space permits, it is good to include an example where a joint state-action representation will help in learning the policy. 

4. 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^*.

5. Looking at Figure 2(a,b,d) JSA seems to have worse asymptotic performance. This could be discussed.
 
6. 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)

7. It's good to see many of the experimental details mentioned. I would recommend including the following details:  
	a. Size of batch and mini-batch used.
	b. Epsilon value, for importance sample clipping in PPO.

Questions:
1. 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? 

2. 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. 

3. 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. 

4. 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. 

5. How are the state and actions fed into the transition model (e.g., are they concatenated, etc.)?

References 
1. Chandak, Y., Theocharous, G., Kostas, J., Jordan, S., & Thomas, P. S. (2019). Learning Action Representations for Reinforcement Learning. http://arxiv.org/abs/1902.00183

2. Henderson, P., Islam, R., Bachman, P., Pineau, J., Precup, D., & Meger, D. (2019). Deep Reinforcement Learning that Matters. http://arxiv.org/abs/1709.06560

3. 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

----- Update

Thank 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. 

The 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.

Lemma 1 is overly complex.

Alternative proof:

Assume pr(a | s) is pi(a | s) (i.e., action probabilities given s are defined under pi).

vpi(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)

Lemma 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.

I 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.

I also appreciate that Figure 2 was added. But, it is a bit hard to interpret. More explanation is needed there.
review_title: Official Review for Reviewer 4

final_score: rating: 5: Marginally below acceptance threshold
review_content: 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.

The main issues for me are:
1. **There are a lot of relevant references missing.** With regards to planning over latent spaces here are some:
  - "Recurrent World Models Facilitate Policy Evolution", David Ha and Jürgen Schmidhuber, NeurIPS 2018.
  - "Model-Based Reinforcement Learning for Atari", Lukasz Kaiser et al. ICLR 2020
  - "Learning Latent Dynamics for Planning from Pixels", D Hafner, T Lillicrap, I Fischer, R Villegas, D Ha, H Lee, J Davidson, ICML 2019
  - "Mastering Atari with Discrete World Models", D Hafner, T Lillicrap, M Norouzi, J Ba, arXiv 2020
  
  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:
  - 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.
  - 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.
  - Li, L.; Walsh, T. J.; and Littman, M. L. 2006. "Towards a Unified Theory of State Abstraction for MDPs. In ISAIM.
  - Taylor, J.; Precup, D.; and Panagaden, P. 2009. "Bounding performance loss in approximate MDP homomorphisms". In Advances in Neural Information Processing Systems, 1649–1656.
  - 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).
  - Castro, P. S. 2020. "Scalable methods for computing state similarity in deterministic Markov Decision Processes". In Proceedings of the AAAI Conference on Artificial Intelligence.
  - 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 .

2. **Lemma 1 is not correct.**
  - 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$.
  - 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')$.
  - 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.
  - (Minor) The jump from $P(s, a)$ to $\pi(a | s)P(s)$ is probably ok, but requires a little more justification.

3. **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).

Medium issues:
1. 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..
2.  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.
3. 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.

Minor issues:
1. 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.
2. 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)$

review_title: Official Blind Review

final_score: rating: 5: Marginally below acceptance threshold
review_content: 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. 

Novelty. The idea is not new but the particular method of defining a policy in the embedding space seems to be new. 

Significance. The work may be of general interest to the reinforcement learning research community.  

Clarity. The paper presents the key idea/method clearly. 

Quality. There are a few theoretical and empirical issues in this work. 

First, 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. 

Second, I do not see why the assumptions can get satisfied in practice. 

Third, 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. 

Motivation. 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. 

Experiments. 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. 
review_title: The paper presents a somewhat novel way to learn and to use joint state-action representation for policy gradient methods in a deep learning setting. There are some concerns about the soundness of the method.

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