Paper: Bridging State and History Representations: Understanding Self-Predictive RL (ICLR-2024)

Reviewer comment (Evaluation):
The experiment in Sec. 5.1 is inconclusive and requires unbiased and fair discussion (also see questions).

Author response:
> The experiments showing different variants of the minimalistic algorithm (Figure 3) are somewhat inconclusive or require further discussion (e.g., claims about better/similar performance compared to ALM(3)).  
> Please elaborate on "This suggests that the primary advantage ALM(3) brings to model-free RL, lies in state representation rather than policy optimization" from Sec. 5.1.  

Please see our general response on these two points.  

> No justification for superior performance of ALM on Humanoid.  

We have recognized this gap and have included a justification in our revision for ALM(3)'s superior performance in the Humanoid task. We suggest that this is likely attributed to the effectiveness of SVG policy optimization in leveraging first-order gradient information from latent dynamics, a feature that seems particularly advantageous in high-dimensional tasks like Humanoid.

Reviewer comment (Evaluation):
In Sec. 5.2, you write: "Surprisingly, model-free RL (ϕ_{Q∗}) performs worse than ϕ_O." Isn't that expected?

Author response:
> In Sec. 5.2, you write: "Surprisingly, model-free RL ($\phi_{Q∗}$) performs worse than $\phi_O$." Isn't that expected?  

This finding was indeed unexpected. The rationale is that model-free RL, not requiring observation prediction, would presumably be less susceptible to distractions. However, our results indicate otherwise. We have also included the explanation to the new version.

Reviewer comment (Experiments):
Please elaborate on "This suggests that the primary advantage ALM(3) brings to model-free RL, lies in state representation rather than policy optimization" from Sec. 5.1. Isn't your model based on ALM? How does the stripped-down version of ALM (that you use as a basis for your augmentations) compare to ϕ_Q∗ (TD3)? Is it the same?

Author response:
> Isn't your model based on ALM? How does the stripped-down version of ALM (that you use as a basis for your augmentations) compare to ϕ_Q∗ (TD3)? Is it the same?  

Our minimalist algorithm shares the goal of learning self-predictive representations with ALM, but diverges in the approach to policy optimization. The key differences in policy optimization between our algorithm and various ALM versions are detailed in Appendix D.2 and can be summarized as follows:

- ALM(3): uses SVG with 3-step horizon planning and additional intrinsic rewards
- ALM-no-model: uses SVG with 1-step horizon planning 
- ALM(0): uses model-free TD3 in the latent space
- Our minimalist algorithm: uses model-free TD3 in the latent space, akin to ALM(0), but with distinct encoder objectives regarding detaching latent states (see Appendix D.2 for details)

Regarding your question on the comparison to $\phi_{Q^*}$ (TD3), the stripped-down version of ALM (assumed to be ALM-no-model or ALM(0)) differs from [TD3 in the original state space ($\phi_{Q^*}$)] as it operates in the latent state space and aims to learn $\phi_L$.

Reviewer comment (Writing):
Define FKL / RKL abbreviations – a good place could be on p. 5 after "includes forward and reverse KL".

Author response:
> define FKL / RKL abbr. - good place good be on p. 5 after "includes forward and reverse KL".  

Thanks and we've added it.

Reviewer comment (Writing):
On p. 5: clarify "detached from the computation graph" [OR] "using a copy ...".

Author response:
> p. 5: detached from the computation graph [OR] using a copy ...  

Here we emphasized that the copy of $\phi$ is also detached from the computation graph (i.e., it is a `deepcopy` in python), so we think "and" is more appropriate.

Reviewer comment (Presentation):
The legend in Figure 5 is slightly different from the graphs.

Author response:
Thanks for your positive feedback and pointing the issue out! Upon reviewing the code, we found that the absence of borders around markers in legends seems a default behavior of the seaborn library. It is a bit hard to change this behavior. Therefore, to maintain consistency, we opted to remove the edges from the markers in the plots, which also makes the plots clearer. This change has been applied to the updated figure in the latest PDF. Please have a look!

Reviewer comment (Presentation):
While the main message of the paper seems clear, the core message of the empirical section is difficult to understand.

Author response:
Thank you for your positive feedback! In response to the main concerns about clarity, we've extensively revised our paper. This includes adding a new section on recommendations (Section 6) and a detailed discussion on motivations (Appendix G). We've also carefully addressed other writing concerns through targeted revisions and detailed responses below.

Reviewer comment (Novelty):
Is minimalist better than $\phi_{Q^{*}}$? This seems somewhat obvious, which is the reason we have all these works on representation learning (not a criticism, just a remark).

Author response:
Our experiments demonstrate that the minimalist algorithm indeed outperforms $\phi_{Q^*}$ in the tasks we evaluated, which include standard MuJoCo, distracting MuJoCo, and MiniGrid. However, it's important to note that these findings may not generalize to every RL task. The performance of deep RL algorithms can vary significantly depending on the specific task structure.

Reviewer comment (Novelty):
Is the purpose broadly to categorize the entire zoo of representation learning algorithms into their essential components i.e. $\phi_{Q^*}$, $\phi_L$, $\phi_O$ and just study how they perform empirically?

Author response:
Yes, we think it is indeed a good future direction for empirical studies.  Alongside this, we also see significant potential in developing new theoretical frameworks. These could include understanding the interplay between representation learning and policy optimization, as well as establishing formal guarantees about the relative sample efficiency of learning different representations, such as $\phi_L$ over $\phi_O$ in distracting tasks, and $\phi_O$ over $\phi_{Q^*}$ in tasks with sparse rewards.

Reviewer comment (Writing):
It would be very useful to include why it is important to test the questions posed in Section 5, and what the final algorithm recommendation should be.

Author response:
Thank you for highlighting these points! To address your queries, we have detailed the motivation behind each hypothesis in **Appendix G ("Motivating Our Hypotheses")**. This section aims to clarify the rationale behind our research questions, including the specific inquiry about the performance of $\phi_O$ in the presence of distractors. We invite you to review this dedicated section and would appreciate your feedback on its clarity.

Reviewer comment (Evaluation):
The sample efficiency claims in Section 5.1 are unclear. It appears that ALM(3) outperforms $\phi_L$ in all cases. It is true that $\phi_L$ does better than $\phi_{Q^*}$, but the former does not seem to be. It is suspected that ALM(3) will do better because it explicitly learns a reward model, whereas $\phi_L$ suggests that the reward model/representations are implicitly learned based on their implication graph.

Author response:
Our data shows varied results across different tasks. In Walker2d and HalfCheetah, $\phi_L$ outperforms ALM(3), while in the Ant task, they demonstrate comparable sample efficiency *at 500k steps*. However, in the Humanoid task, ALM(3) significantly outperforms $\phi_L$. For further clarification, please refer to our general response on this topic.


We want to clarify that although ALM(3) explicitly learns a reward model, its encoder objective does *not* include reward prediction, as evidenced by [their code](https://github.com/RajGhugare19/alm/blob/main/agents/alm.py#L209) uses `z_dist.sample()` that detaches the latent states when learning the reward model. Therefore, it's unlikely that explicit reward prediction is a significant factor in learning an effective encoder. Rather, ALM(3)'s enhanced performance can more plausibly be attributed to its use of SVG planning for policy optimization. This clarification has been added to the new version of our paper.

Reviewer comment (Writing):
Is there intuition for why $\phi_O$ may struggle with distractors? This question is posed in the empirical section, but it is unclear what motivates it.

Author response:
Thank you for highlighting these points! To address your queries, we have detailed the motivation behind each hypothesis in **Appendix G ("Motivating Our Hypotheses")**. This section aims to clarify the rationale behind our research questions, including the specific inquiry about the performance of $\phi_O$ in the presence of distractors. We invite you to review this dedicated section and would appreciate your feedback on its clarity.

Reviewer comment (Writing):
I'm not familiar with using the phrase "distracted MDPs" to mean MDPs with distracting elements.

Author response:
We appreciate this observation and have made adjustments to clarify our terminology. We've changed "distracted MDPs" to "distracting MDPs" in our paper. Additionally, we've included a footnote stating, "Distracting MDPs refer to MDPs with distracting observations irrelevant to optimal control in this work," to prevent any potential confusion.

Reviewer comment (Novelty):
Maybe I'm wrong, but I feel like we already know that ZP + $\phi_{Q*}$ implies RP.

Author response:
You raise a valid point. While the connection between ZP + $\phi_{Q^*}$ and RP might seem intuitive, our work contributes to formalizing this relationship. The inverse Bellman equation in our proof, also used in imitation learning literature [Garg et al., IQ-Learn: Inverse soft-Q Learning for Imitation], provides a novel perspective specific to state or history representations for RL. To the best of our knowledge, this precise articulation is new in both MDPs and POMDPs.

Reviewer comment (Presentation):
Fig 3. Kind of hard to distinguish lines.

Author response:
Acknowledging the difficulty in distinguishing lines in Figure 3, we also want to point out that the overlap of curves mainly illustrates similar performance trends. Here we offer a qualitative comparison at 500k steps using symbols: $<$ for worse, and $\ll$ for much worse. Hope this will help you read the figure. 

- Walker2d: ALM(3) and TD3 (they are quite similar) $\ll$ OP and ZP methods (they are quite similar).
- HalfCheetah: TD3 $<$ ALM(3) $<$ OP methods $<$ ZP methods.
- Ant: TD3 $\ll$ OP-FKL $<$ the other methods. Note that ALM(3) initially leads but later declines.
- Humanoid: TD3 and ZP methods $<$ OP methods $\ll$ ALM(3).

Reviewer comment (Presentation):
Fig 5. A little difficult to read. See if maybe log return would help?

Author response:
We tried to plot the log returns but found it did not help clarity much. This is due to the scale of episode returns, which are mainly around $10^3$, with the exception of the Ant. In the Ant with $\ge 2^7$ distractors, all methods are close to $0$, indicating they are equally bad. For the other environments with large number of distractors, the results clearly show that TD3 is the worst performer. Following this, OP-$\ell_2$, OP-FKL, and ZP-FKL exhibit similar performance levels. ZP-$\ell_2$ and ZP-RKL stand out as the best performers.

Reviewer comment (Experiments):
Top of p9, minigrid experiment. This feels like a minor bait-and-switch. Normally minigrid uses pixels, no? And why does detached rank drop on minigrid when theory suggests it should remain high?

Author response:
In our MiniGrid experiment, we used vector-based (symbolic) observations, which is a standard option in MiniGrid, rather than the pixel-based representation. This approach aligns with common practices in recent research. For instance:

- Parker-Holder et al., Evolving Curricula with Regret-Based Environment Design, NeurIPS 2022. (They used "147-dimensional observation".)
- Campero et al., Learning with AMIGo: Adversarially Motivated Intrinsic Goals, ICLR 2021. (They mentioned that the observations are "symbolic".)
- Zhang et al., NovelD: A Simple yet Effective Exploration Criterion, NeurIPS 2021. 
- Jiang et al., Replay-Guided Adversarial Environment Design, NeurIPS 2021.
- Seo et al., State Entropy Maximization with Random Encoders for Efficient Exploration, ICML 2021. ([Their code](https://github.com/younggyoseo/RE3/blob/master/a2c_re3/rl-starter-files/rl-starter-files/utils/env.py) indicates they use symbolic observations.)
- Dennis et al., Emergent Complexity and Zero-Shot Transfer via Unsupervised Environment Design, NeurIPS 2020. 
- Goyal et al., InfoBot: Transfer and Exploration via the Information Bottleneck, ICLR 2019. (They mentioned that the observations are "not RGB images".)

Hope our responses have resolved your concerns!

Reviewer comment (Presentation):
"Validation of end-to-end hypothesis" (and throughout). It's a bit hard to keep track of which is which. Text uses $\phi_o$, $\phi_L$, but fig uses OP, ZP, RP+OP, RP+ZP.

Author response:
To address this, we have revised the text in the new version of our paper to enhance readability and ensure consistency between the text and figures.