Bridging State and History Representations: Understanding Self-Predictive RL

Thank you for your positive feedback!

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.

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

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.

Given the scope of the article, why did you preferer conference vs. journal? (not answering this question can be understandable)

We chose to submit to ICLR due to its high relevance to our research topic (representation learning) and its high impact factor. Additionally, we plan to submit to a journal for more in-depth reviews.

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

Thanks and we've added it.

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.

in Fig. 6 caption, mention that RP+OP / RP+ZP are phased; it is not obvious what the difference is and why they should perform worse

The caption of Figure 6 has been updated to indicate that RP+OP / RP+ZP use a phased training approach, as detailed in Section 3.1.

Our initial hypothesis did not anticipate a significant performance difference between the phased and end-to-end approaches, given that both utilize the same R2D2 backbone and target similar representations. However, the end-to-end approach was found to be at least comparable, if not superior, to the phased one in our experiments. We speculate this could be due to the phased approach dealing with accumulated errors from three separate losses (reward prediction, OP or ZP, Q-learning loss) as opposed to the end-to-end approach's two losses (OP or ZP, Q-learning loss), potentially complicating the optimization process.

legend in Figure 5 is slightly different from the graphs

We have carefully reviewed the figure and its legend but did not identify any discrepancies. Could you please specify which part of the legend you find to be different from the graph?

Hope our responses have resolved your concerns!

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!

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.

It would be very useful to include why it is important to test the questions that are posed in Section 5. Is there intuition for why $\phi_O$ may struggle with distractors? The question is posed in the empirical section, but it’s a bit unclear what motivates this question.

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.

It would be very useful to include what the final algorithm recommendation should be.

This a very good point! Please see our general response on our recommendations.

Should we be using minimalist over everything else? If so, ALM outperforming minimalist in Figure 3 on multiple occasions would not support that.

Our intention is not to position the minimalist algorithm as superior to others, such as ALM, in all aspects. The primary goal of the minimalist algorithm is to provide a straightforward means to validate the hypotheses derived from our theoretical framework. Its simplicity aids in easy implementation and understanding. Additionally, it serves as a baseline for future research, enabling separate examination of the impacts of representation learning and policy optimization.

As an example, in Section 5.1, we discuss how ALM's performance advantage over TD3 at 500k steps in MuJoCo tasks is largely attributable to its approach to representation learning, illustrating the utility of the minimalist algorithm in isolating these factors for analysis.

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

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.

Is it 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?

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.

The sample efficiency claims in Section 5.1 are unclear to me. It appears that ALM(3) outperforms $\phi_L$ in all cases.

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.

I would suspect that ALM(3) will do better because it is explicitly learning a reward model, whereas $\phi_L$ is suggesting that the reward model/representations are implicitly learned based on their implication graph.

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

Hope our responses have resolved your concerns!

Thank you for your feedback! It seems the main concern is the limited contribution. We agree that many prior methods are based on a similar intuition of self-prediction. However, these methods come from different perspectives and have various theoretical properties with nuanced-yet-crucial implementation differences (see Table 1 and Appendix C). Our objective is to contextualize these varied approaches within a unified framework, aiming to distill fundamental principles of representation learning in RL. This endeavor seeks not just to deepen theoretical understanding and prevent overlapping research efforts, but also to offer practical guidance to practitioners (Section 6), as elaborated in our general response.

How does your paper help practitioners identify and use the best approach for their specific RL problem?

We have included recommendations for practitioners in Section 6 in the revision. For more detailed information, please refer to our general response on recommendations.

The results in the paper are either trivial or indecisive. The paper is built around an insight that different works do similar things trying to optimize self-prediction of certain features, but this is, in my opinion, trivial.

We respectfully disagree with the view that our work is trivial or unconvincing. Contrary to this opinion, the other three reviewers (5DUm, kS5Z, sXU8) have recognized the significance and value of our unified framework, which is much needed to the RL community. 5DUm noted its utility in enabling comparisons of various methods based on their core representation learning principles without being confounded by other components. sXU8 commented on its potential to prevent redundant research efforts.

The paper uses a lot of abbreviations.

To enhance readability, we have implemented clickable abbreviations that link directly to their definitions. Additionally, we occasionally spell out the full names of these abbreviations to aid recall.

The proofs a sketchy and uncommented, and verifying or even following the proofs takes tremendous effort.

Are there any specific proofs that you found challenging to follow? We would appreciate your pointing them out and offer further clarifications as needed.

The empirical evaluation is unconvincing. According to the plots presented in the paper, the proposed unified algorithm does not outperform (and does not always perform comparably) to algorithms from the literature.

The aim of the experiments and the paper as a whole is not to propose a SOTA or competitive method, but rather understand various self-predictive methods within the broader context of representation learning. The empirical evlauation is intended solely to validate our hypotheses suggested by our theory. In fact, our empirical results validate most of our hypotheses.

Notably, the findings related to stop-gradients were particularly convincing, as acknowledged by the other reviewers (5DUm, kS5Z, sXU8). For detailed clarification on the comparisons with ALM(3), please refer to our general response.

Looking at the algorithm pseudocode and implementation, this is not surprising, given that the 'minimalist algorithm' is more of a boilerplate, which, when filled with details, reduces to one of the earlier published algorithms. However, any practical implementation requires attending to details, which the unifying minimalist algorithms fails to achieve.

To the best of our knowledge, our algorithm, employing end-to-end training with a self-predictive auxiliary task, is novel within the POMDP literature. Is the reviewer aware of prior work that can be reduced to our algorithm in POMDPs?

In addition, we do care about practical implementation in our algorithm. For example, we tuned the coefficient of the auxiliary task for each benchmark. We also base our algorithm on existing, well-tuned implementations, rather than writting a new one from scratch.

Thank you for your valuable feedback! We have addressed the main concerns around some confusing terminologies and the seemingly incorrect proof in the appendix. Below are our responses to each point, with the changes in the updated PDF highlighted in magenta.

Page 2: "encoder of the value function" -> "encoding learned with/for value function"

We've fixed this term to be "the encoder learned for the value function" in the new revision to avoid confusion.

Page 2: next latent state prediction --- the distribution of the encodings rather than the state is predicted.

We've renamed this term to be "next latent state distribution prediction (ZP)" to avoid confusion.

Page 2: "encoder is capable of predicting" the encoder is defined only encodes the state.

We've changed this to "the encoder can be used to predict" for clearer understanding.

Page 3: "expected next latent state" --- the latent state is not defined on a codomain for which the mean can be defined (e.g. R^n or Z^n).

We've added this assumption into Appendix A.2 (see "Remark on the latent state distribution") to address your concern.

Page 4: "The bootstrapping effect ..." phi does not in "both sides of the next latent state prediction" explicitly, one must guess that z' = phi(h') .

We've added this explanation "since $z'$ also relies on $\phi(h')$" to make it clearer.

Page 4-5" This is not a stochastic encoder, but a probabilistic one. A probabilistic encoder returns a distribution \Delta(Z), a stochastic encoder returns a sample from a distribution. This is significantly different.

We've renamed all instances of "stochastic encoder" to "probabilistic encoder" to avoid confusion between these two concepts.

There is almost no way to train a 'deterministic encoder' using differentiable models (deep networks).

We believe we can, and actually have, trained a deterministic encoder with neural networks in our implementation. The deterministic encoder $f_\phi$ takes input a history $h$ and outputs the latent state $z$. We assume $f_\phi$ is differentiable w.r.t. $\phi$ for any history $h$, which is the standard assumption in deep RL.

Page 5: after expr (1), where D(., .) \in R compares two distribution. When (1) reaches minimum --- it does not have to, functions on R do not have to have a minimum.If it is on R+, then it still would be good to know what features D should possess.

We've changed $\mathbb R$ to $\mathbb R_{\ge 0}$ (the set of positive reals and zero) to ensure the existence of a minimum.

Page 5: Equations (2) and (3) expectations are on variable o' which does not appear in the expression under expectation.

We want to clarify that $o'$ indeed appears in the expression through $h' = (h, a, o')$. We define $h'$ in the "MDPs and POMDPs" in the background section.

Page 5: D (apparently divergence) uses different notation in (1) -- $D(., .)$ and in (3) $D(.||.)$.

We've unified the notation to $\mathbb D(\cdot\mid \mid \cdot)$ to avoid confusion. Since we allow both symmetric and asymmetric functions, we use this notation.

Proof of proposition 1 (appendix): second terms in eqs (90) and (91) are mutually cancelled, despite being different.

We believe our proof is correct. The two cancelled terms in (90) and (91) are equal due to the following elementary fact: Inner product is a bilinear form, i.e., $\langle a + b, c \rangle = \langle a, c \rangle + \langle b, c \rangle$. Therefore, it follows that for a constant vector $a$ and random vector $X$ (with PDF $P_X$),

Note that in our case, $g_\theta(f_\phi(h),a)$ does not depend on $o'$, thus it can be treated as a constant for the derivation.

We hope our responses and the revisions in the manuscript have resolved the concerns raised. Looking forward to your feedback!

Thank you for your positive feedback!

I'm not familiar with using the phrase "distracted MDPs" to mean MDPs with distracting elements.

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.

Maybe I'm wrong, but I feel like we already know that ZP + $\phi_{Q^*}$ implies RP.

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.

Top of p5: The use of $\mathbb P$ and $\mathbb Q$ is confusing, given that $P$ and $Q$ are often overloaded in RL.

We acknowledge this potential confusion. To address it, we have now replaced $\mathbb Q$ with $\mathbb P$ in our paper, eliminating the notational conflict and avoiding confusion with Q-values in RL.

How does Thm 3 predict low absolute cosine similarity? Because we start with full rank?

In this experiment, the latent state dimension is set to $2$ and we orthogonally initialize the $\phi$ matrix. This results in $\phi^\top \phi = [1, 0; 0, 1]$ initially. According to Thm 3, $\phi^\top \phi$ maintains its initial value in continuous-time. Thus, the cosine similarity (and the inner product) between the two columns of $\phi$ will remain close to $0$ in discrete-time.

Therefore, Thm 3 predicts low absolute cosine similarity because we initialize $\phi$ as an orthogonal matrix, which is a stronger condition of having full rank.

Fig 3. Kind of hard to distinguish lines.

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

Sec 5, first para: would be helpful to summarize the findings when introducing the hypotheses. Do the authors have a clear recommendation on when to use OP vs ZP vs RP+[*]? It would be nice to have some clear takeaways after all this analysis.

Thanks for your suggestion! We've added our findings to the paper and please the general response on recommendations.

p7, penultimate para: I don't know about similar sample efficiency in ant. It seems a lot worse. I also don't know if I agree with the conclusion that "the primary advantage ALM(3) brings to model-free RL, lies in the state representation rather than policy optimization, except for Humanoid." I feel like it's too soon to conclude something as sweeping as that.

Please see our clarification on this in our general response.

Fig 5. A little difficult to read. See if maybe log return would help?

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.

And why does detached rank drop on minigrid when theory suggests it should remain high?

We acknowledge that the observed drop in detached rank in the MiniGrid experiment is intriguing, especially since our theory does not predict a significant degradation in such cases. We plan to explore this further in our future research.

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

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.