Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper proposes a method (LCER) for agents to develop compact task representations from high-dimensional inputs (e.g., images) during unsupervised exploration that are effective for downstream tasks. The paper hypothesizes that encouraging an agent’s representation to focus on changes related to the agent’s actions will be more useful to downstream tasks than representations that try to naively model all changes. In order to allow some non-controllable elements in the representation the paper proposes a blended objective which trades off representing action controllable aspects of state with non-controllable elements of state using an application specific empirically tuned hyperparameter \beta. The paper is set in a fully observable world where a state at time t, S_{t}, is projected to an instantaneous latent state Z_{t}. The idea of controllable elements is implemented by computing mutual information between the latent state projection in the future, Z_{t+k} and the actions taken by the agent I(  Z_{t+k}; A_{t:t+k-1} ).  The uncontrollable information is uncontrollable elements is implemented by computing mutual information between current and future latent states I( Z_{t}; Z_{t+k} ).  Mutual information between latent state and actions uses the neural InfoNCE bound trick. For the second term, the paper instead optimizes Jensen-Shannon divergence between current and future latent representations instead of mutual information, again using a neural bound as the paper claims this is more stable. The paper then shows that this objective can be combined with a standard soft actor-critic architecture to learn representations without extrinsic rewards. In empirical studies on the PlaNet, the paper’s proposed LCER method finds representations that can be used by a policy learner to converge faster than competing methods. The proposed LCER method also compares well to an SAC agent that has access to underlying state representation of the domain (e.g., joint angles, velocities). A second set of experiments on the Distracting  Control Suite which has a lot of extraneous background noise show wider separation between LCER and competing methods. A study shows that increasing the look ahead K improves the quality of representations. 
main_review: The high-level idea seems sound and practical and the implementation strategy builds on popular ideas in the literature and seems reasonable. 

Does not cite DIYAN (Eysenbach 2018) https://arxiv.org/abs/1802.06070 which also uses mutual info between latent state projection Z and actions A, though not for representation learning.  

In equation 6, the information in the sequence is given as the sum of information at each horizon value k which is probably a sensible approximation, but I am guessing not technically correct as the states at each horizon value are not mutually independent. 

Page 5 middle paragraph starting “where B is the batch size” contains “is a latent defined” which I am guessing should be “is a latent variable defined”. 

Why is JSD more stable than MI? Is it possible to give an intuition? 

I didn’t find the section on augmentation directly relevant to the paper’s point and I think it could be omitted from the main text if space was tight. It might allow the paper to expand section 4.2 a bit more to explain the recursion better.  

Alternative: LACER: Learning Agent Controllable Elements Representations 
summary_of_the_review: I think the idea of explicitly trading off the action focused vs. state focused loss to generate representations from unsupervised experience is new but closely related to precedents. It speeds convergence on easier PlaNet tasks and seems empirically to make a difference in achievable expected reward for the harder tasks in the Distracting Control Suite.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: The paper proposes the LCER algorithm that trains representations to retain the controllable elements of the environment, which can reflect the action-related environment dynamics and thus are likely to be task-relevant. The main experiments use DMControl suite and includes several ablation experiments.

main_review: ## Strengths

1.  The paper is generally well written and well motivated. While section 4 required multiple passes (likely due to multiple approximations that are used), the high-level motivation and ideas are easy to grasp.

2. I like the idea - it is quite simple and elegant - learn representations that retain information about the controllable elements of the environment. 

3. I liked the experiments and ablations that the paper includes. While I have several issues with the claims/interpretation-of-results, I think the experiment design is well thought.

## Areas for improvement

1. Several claims are not well-substantiated. 

    * In some cases, it could just be a matter of rewording. For example "LCER is the first representation learning algorithm that enables the pixel-based SAC to outperform state-based SAC on the DMControl100K benchmark". From Table 1, LCER lags behind StateSAC on 5 (out of 6 environments). While its performance is impressive, the claim itself is not supported. Comparing the average performance (across 6 envs) is misleading because of 1 outlier environment. 

    * In other cases, the claims are not tested. For example, there is no experiment to show that the representations indeed "retain the controllable elements of the environment" (which is the key motivation behind the work).

    * In several cases, the results are quite close to each other (eg Figure 1 walker walk, reacher-easy, ball-in-cup, cartpole...) In absence of statistical significance test, it is difficult to conclude that the proposed approach is indeed better. Same applies to several results in the tables as well.

2. When evaluation robustness to distractors, better baselines (like Zhang et al 2020) should be used.

## Further questions

Note that I did not consider these points as weakness as I did not fully understand these aspects.

1. In equation 4, what is the meaning of "$I([\phi(S_t), A_{t:t+k-1}]$" ie what is the "[...]" operator doing?
2. In equation 3, is only A conditioned on $\phi(S_t)$ or both $\phi(S_{t+k})$ and $A$ are conditioned? 
3. In page 3, Section 4, the authors correctly point out "For example, in navigation tasks, the position of the goal
is task-relevant but not controllable by the agent" as a motivation for controlling the degree of compression. However, their experiments did not have any goal conditioned environments. In that case, would it make sense to directly use equation 3 (atleast in the context of experiments considered in the paper).
4. In page 4, the authors mentions "To tackle this problem, we substitute I..." (right before equation 5). Could they describe why is this approximation valid?
5. What is the runtime cost of LCER compared to the baseline? I expect it to be much higher than the baseline (given it is used k-step lookaheads to compute the losses. Reporting the runtime comparison will be useful for a holistic comparison. Similarly, I expect LCER to be using much more params and would like to see a comparison of those numbers.
6. What is the effect of using the target encoder? How severe is effect if target encoder isnt used?
7. In page 6, the paper mentions "All parameters of StateSAC.." do you mean all hyper-parameters of StateSAC... ?
8. In figure 2, cartpole results seems to be stopped too early. The two baselines may be able to overtake the red curve.
9. In page 7, the paper mentions "In LCER, we also randomly shift the observations by [−4, 4] pixels". Is this aimed to provide translation invariance? If yes, isnt a conv net already providing that? Is this augmentation applied to the baselines as well?
summary_of_the_review: The paper proposes a simple and elegant idea and shows that the approach "works" in practice. However, it is difficult to evaluate the effectiveness/significance of the approach. I look forward to the author's response and interacting with them to understand their approach better.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: This paper proposes a novel representation learning flamework, LCER, which enables the agent to extract task-relevant representations from image-based observations. It optimizes the surrogate loss of conditional mutual information (CMI) between future-state embedding and action sequences given past-state embedding. The key contribution seems the balancing between InfoNCE objective (similar to CURL) and JSD-based CMI. The experimental comparison shows that LCER performs best among some image-based RL algorithms (PI-SAC, SLAC. CURL, PlaNet); in 100k DM Control benchmark (6 tasks), and distracted-background settings.
main_review: ## Strengths
- This paper is well structured and easy to follow.
- The empirical results in 100k DM Control benchmark (6 tasks), and even distracted-background settings, show LCER performs best among some image-based RL algorithms, which successfully reveals that learned representation of LCER captures the controllable features in the image observations.

## Weaknesses
- I think the derivation of LCER seems a bit misleading. The objective that LCER actually optimizes may not be introduced from the CMI directly; it starts with CMI between future-state embedding and action sequences given past-state embedding (Eq. 3), introduces the bottleneck parameter $\beta$ (similar to Lee et al. 2020) (Eq. 4), and deformulates the objective again (Eq. 5). In Eq. 5, $J_k^2$ is the same term as Eq. 3, which might be confusing. It can be simply argued that LCER introduces the CMI term into CURL objective ($J_k^1$) with some ratio $\beta>0$, and also introduces the surrogate form of CMI based on JSD (Eq. 10). In addition, balancing $\beta$ seems very important hyper-parameters for LCER (I guess even in standard DM Control, where $\beta=0.1$ is used. If LCER learns controllable elements from CMI, $\beta$ should be 1, but isn't).
- I am also a bit confused about the connection and difference between LCER and PI-SAC, since both of them optimize the CMI objective as an auxiliary loss for RL. While the authors just say, "PI-SAC tries to encode predictive information, while LCER focuses on encoding controllable elements", further clarification about this is required to emphasize the contribution of this paper.
- For PI-SAC baselines, the authors can include the "gradient\_step=2, 4" results for a fair comparison, since it improves the performance according to its codebase (https://github.com/google-research/pisac). In contrast, LCER can also utilize such extra gradient steps. Moreover, I find that while LCER uses $K=5$ sequence length for InfoNCE + CMI objective, PI-SAC only uses $K=3$ (from Appendix A of Lee et al. 2020), which might be "unfair" comparisons, considering the experiments excepting multiple gradient steps of PI-SAC.
- I think the authors may also include the results of Dreamer (Hafner et al. 2020) for 100k benchmark since it generally achieves better results than PlaNet; in addition, CURL and PI-SAC compared against it in their original papers (+ code is available online).
- In distracting-environments settings, the authors may also include Contrastive Variational  Reinforcement Learning (CVRL) as a baseline, proposed by Ma et al. (2020), which achieves better results compared to Dreamer, image-based SAC, and image-based D4PG (+ code is available online). In addition, Ma et al. (2020) also take a contrastive approach to handle distracting background, so it should be added as related work.


### Reference
Lee et al. Predictive Information Accelerates Learning in RL. Advances in Neural Information Processing Systems (2020).

Hafner et al. Dream to Control: Learning Behaviors by Latent Imagination. International Conference on Learning Representations (2020).

Ma et al. Contrastive Variational Reinforcement Learning for Complex Observations. Conference on Robot Learning (2020).

summary_of_the_review: As discussed in Main Review in detail, I think the derivation of LCER seems a bit misleading, and further clarification compared to PI-SAC is required (since both PI-SAC and LCER optimize the CMI objective). In addition, some important baselines (PI-SAC w/ gs=2 or gs=4, Dreamer, CVRL) seem missing in the current manuscript. Considering these aspects, I larn towards rejection.
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: There has been a school thought lately in RL that good representations for control should focus on abstracting out on only the task-relevant features. In high-dimensional tasks, this is especially useful in making optimal use of a limited neural network capacity and may even help the agent be robust to distractors in the environment. The authors motivate an approach to learning task-relevant representations by having the representations capture only the controllable elements of the environment. Their approach is based on a mutual information loss between actions sequences and final states, and they validate their approach on a series of DMControl task where both high-dimensional pixel observations as well as ground truth latent state representations are available. The idea is to compare the performance of agents that attempt to learn representations of latent states from pixels and compare against an agent that has direct access to the ground truth latent state. In their empirical results they show that:
- their new approach (called LCER) is able to outperform alternative approaches on a set of these tasks. On some tasks, it even outperforms an SAC agent with access to the ground truth states.
- LCER is able to outperform alternatives when distracting visual features are introduced into these tasks
- LCER's representations learned on a source task are able to transfer better to downstream tasks.
main_review: **Strengths:**
- paper was clearly written and well-motivated
- approach can be applied to any task
- clear mathematical exposition of the loss function used
- demonstration of better sample complexity against other baselines including CURL, PlaNet, SLAC, and PI-SAC.
- demonstration of robustness to distractors compared to a PI-SLAC baseline
- the authors have explored the robustness of the algorithms to the two new hyperparameters introduced, $\beta$ and $K$


**Comments:**
I do like the idea of using controllability to inform compression of state representations -- after all, what is important to the agent is what it can control about the environment. Yet there may be situations where this approach can be too compressive. Perhaps the most important situation is where there are non-controllable task-informative features. The authors do point out one example, the goal location. But there are in principle many others. For instance, the directions for solving a task (e.g. a map showing the solution to a maze) may be painted on a wall in the environment -- the agent could still solve the task by completely ignoring the directions, but it would do much better if it could exploit that information. Another limitation with this approach is if the environment dynamics change. For instance, suppose an agent wanted to navigate a drone to a goal. Controllability will differ depending on wind conditions. So if the agent learned a representation under e.g. no wind, that representation may be less effective if the agent had to navigate in the presence of strong and volatile wind. These limitations are by no means reasons to disqualify the paper, but it would be worth pointing this out in their discussion of their algorithm.


**Weaknesses:**
- the idea of using mutual info to guide representation learning has been considered previously in various guises. Perhaps the closest version to the current approach is by Lee et al, 2020, which uses mutual information between the initial and final states of a trajectory to guide state compression. The authors' main modification is really to introduce action sequences into the picture, motivated by the idea of using controllability to guide compression, and in this reviewer's view, this contribution seems borderline incremental. Indeed, the improvement over PI-SAC (Table 1) seems to be within statistical error for several of the tasks (cartpole swing-up, reacher easy, finger spin, and ball-in-cup), which makes me wonder to what extent is the improvement attributable to implementational differences instead.
- mutual information between two random variables depends on their probability distributions. One of the variables in their objective is the action sequence, but this depends, in particular, on the policy learnt (training data are sampled from the replay buffer). This begs the question to what extent are the representations learnt optimised for the policy (and hence the particular task). It may be helpful to freeze the encoder and evaluate the representations on different tasks in the same environment (for instance, by changing the reward function to encourage different behaviour). In this setting, what is controllable about the environment would remain the same, but the optimal policy for the task will have to differ. This would also shed light on how transferable these representations would be to different tasks in the same environment. Moreover, it would be helpful to compare against a ground truth baseline where a LCER agent learns representations directly on the target task and is then made to "transfer" to the same task -- the idea is to see how well the agent that transfers representations compares to one that has access to the "ground truth" representation.
- Section 5.4 presenting the evaluation of representation transfer was especially weak. The difference between the source and target task was not made clear, making it difficult to interpret what transfer is taking place between tasks. Moreover, why did the authors only compare against SAC baselines in this section? Clearly SAC was not designed to extract out task irrelevant features, so this seems like a weak baseline to compare against.
- as mentioned above, to what extent are these representations robust to perturbations in the state dynamics. 
- a limitation of the approach is the introduction of extra hyperparameters $\beta$ and $K$, which need to be tuned for each task, particularly the weighting hyperparameter $\beta$.
- minor typographical errors including spurious capitalizations and places where the authors wrote "dose" where they meant "does" (e.g. in the middle of the 1st paragraph of the "Predictive representations in RL" subsection and in the middle of the captions for Figs 3 and 4)

summary_of_the_review: While I liked the motivation of their approach, my main hesitation for recommending this paper lies both on the strength of their contribution and issues with their empirical results. On the strength of their contribution, to me it seems borderline incremental, so I would be curious to hear the opinion of the other reviewers. Moreover, the results seem marginally better than PI-SAC, which makes me wonder to what extent the better results can the better results be conclusively attributed to the better objective and not to implementational factors, like the choice of $h()$ function architecture or a better approximation of the mutual information objective.

On the empirical side, their explanation of transfer in Section 5.4 could have been clearer, and they could have compared against stronger baselines. I would also have appreciated insights into how much the representations are fine-tuned to the policy being used. Finally, it would strengthen the paper if the authors could provide insight into how well do these representations transfer to different objectives as well as to perturbations of dynamics in the same environment. The main motivation behind the authors' work is to have an agent learn state representations that capture controllable elements of the environment only. Evaluating transfer on different goals in the same environment would be a test of this, since an agent that truly captured the controllable features should not be adversely affected by a change in goal or reward function.
review_title: null

Meta-review: This paper introduces an objective for representation learning that captures "controllable elements" in the environment (i.e., things that are affected by the agent's actions). In their reviews and discussion, the reviewers agreed this idea was intuitive, well-motivated, and the paper well written. However, multiple reviewers raised concerns about the evaluation and the extent to which LCER is truly an improvement over PI-SAC. Although many of the reviewer's concerns were addressed in the rebuttal period, at the end of the discussion they were still unconvinced or confused about how much LCER really helps over PI-SAC. Based on this, my assessment is that this paper is a promising piece of work, and that with some more controlled comparisons (see suggestion below) it would be a useful contribution to the literature. However, given that the claims are not fully supported as it currently stands, I recommend rejection.

Specific suggestion to improve the paper: based on reading the paper and the discussion, it seems to me (as per the authors' own statement in response to Reviewer uWv6) that the most valid/controlled comparison between LCER and PI-SAC is in Figure 4, where LCER w/ $\beta=0.1$ "can be seen as a variant of PI-SAC with the same embedding choices as LCER" (author's words). However, when taking into account the error bars of the training curves, other values of $\beta$ are only clearly better than $\beta=0.1$ in 1/3 environments (D.walker-walk). This does not make for a particularly convincing result that LCER is better than PI-SAC. To improve the paper, I'd encourage the authors to run further well-controlled comparisons such as this in a larger number of environments. If they can show via such controlled comparisons that LCER is generally better than PI-SAC (i.e. LCER w/ $\beta=0.1$) then that would be a much more compelling demonstration of LCER's superiority.
Decision: Reject