Studying the Interplay Between the Actor and Critic Representations in Reinforcement Learning

We thank the reviewer for their insightful review. In regards to weaknesses 1 and 2, we kindly point the reviewer to the comment titled [proposed changes to manuscript] (at the top).

In regards to weakness 3, we agree with the reviewer that keeping track of the different metrics can be difficult. In light of this, we are reworking Figure 1 (right) as a standalone table to contain compact definitions for each of the metrics, alongside the existing information reporting the changes observed in $\phi_A$ and $\phi_C$ once the representations are decoupled. A preview for this table can be accessed at https://ibb.co/vwHnfxS

We hope the new table can act as a "cheatsheet" for the reader throughout the paper, reminding them how to intuitively understand each metric, and what they say about the representation.

We thank the reviewer for their insightful review. We were glad to learn that the reviewer found the contributions made in this article interesting. We believe our comment titled [proposed changes to manuscript] (at the top) covers the weaknesses pointed out by the reviewer. We will therefore use this comment to answer the reviewer's question.

For the mutual information metrics, e.g., Figure 4, how much will they change if we simply change some of the hyperparameters of one algorithm?

While we lack the computational budget to conduct a sensitivity analysis on all the possible hyperparameters, the additional experiments we are currently running should shed some light on the effect of changing the model capacity (i.e. the number of learned parameters) and the learning rate.

We occasionally had to tune the weighting coefficient of the representation learning objectives tested to improve learning stability. The impact of this weighting coefficient on the mutual information metrics was as expected: If objective $X$ with coefficient $c_X$ increases information metric $I_A$ and decreases $I_B$, then increasing $c_X$ further increases $I_A$ and decreases $I_B$.

We thank the reviewer for their time and comments, and appreciate the reviewer’s suggestions. We believe incorporating these suggestions and addressing these question will improve the clarity of the presentation. We will take these suggestions into account in the new manuscript revision. We point the reviewer to the comment titled [proposed changes to manuscript] (at the top) for a summary of the changes we plan to make for the next revision. We hope these clarifications will make the reviewer consider improving their score.

We will use this thread to address the reviewer’s questions.

Line 154: Why measuring how easy it is to infer the identity of the current level from Z is a good measure of overfitting?

Being able to infer the identity of the training level from $Z$ implies the learned latent features are level-specific, and that $I(Z;L)$ is high.

To help build an intuition on the relationship between $I(Z_A;L)$ and overfitting, we link to this graphics depicting an example in which the RL agent would be trained to play Super Mario Bros: https://ibb.co/Sst7L3V

Formally, the theoretical relationship between $I(Z_A;L)$ and the generalisation error is stated in Theorem 3.1: reducing $I(Z_A;L)$ will minimise an upper bound on the generalisation error. We will be sure to make this connection clearer in the final version.

Line 158: The purpose of Theorem 3.1 is not explained. Also, is it a result from the other paper or proved first time in this paper?

Theorem 3.1 follows directly from Lemma 3.1 in [1]. We will include a sketch proof in the appendix, as the theorem was not explicitly stated in this form in [1].

Line 165: Why does a high I(Z_C, V) help optimizing the value loss?

A high $I(Z_C; V)$ means that the critic’s learned latent features $Z_C$ are highly informative of the underlying state values $V$. Hence a high $I(Z_C; V)$ helps minimising $\ell_V$.

Line 166: What do you mean by increasing I(Z_A, V) and what's the relationship between value distillation and increasing I(Z_A, V)?

The standard approach to value distillation, which we employ in this paper, is to learn to predict value targets from $\hat{V}^{aux}=v^{aux} \circ \phi_{A}$, $v^{aux}$ being parameterized as an additional head of the actor network.

Since distillation gradients are propagated through $\phi_{A}$, performing distillation will increase $I(Z_A; V)$. We point the reviewer to l114-123 for a more detailed description of value distillation.

Line 175: What do you mean by "whether it is possible to identify latents (ϕ(o), ϕ(o′))) obtained from consecutive observations from pairs stemming from non-consecutive observations?" Is it consecutive or not?

We meant here that a high $I(Z;Z’)$ implies that it is easier for a binary classifier to predict whether a pair of latents $(z_1,z_2)$ corresponds to two consecutive observations $(o, o’)$, or to two non-consecutive observations $(o_x, o_y)$. We will clarify this in the main text.

Line 194: What's the role of linear probing here?

Linear probing means freezing $\phi^*$ and optimising the learning objective using a linear layer conditioned on $\phi^*(o)$. We will rework the definition of $\phi^*$ in the next revision as we have come to the conclusion that mentioning linear probing adds unnecessary complexity.

Equation 6: What are $z^*_{A_0}$ and $z^*_{A_1}$?

$z^*_{A_0}$ and $z^*_{A_1}$ are the two possible outputs of $\phi^*_A$ for the assembly line environment. Can the reviewer clarify what didn't seem clear here?

Why I((Z^_A, Z^'_A); A) is guaranteed to be maximized when following the optimal policy?

Our definition for $\phi^*_A$ in Equation (6) states that we can always identify the optimal action from the value taken by $z^*_A$. It directly follows that $I(Z^*_A;A)$ is always maximised, and, by extension, $I((Z^*_A;Z^{*'}_A);A)$ is maximised as well.

[1] Garcin, S., Doran, J., Guo, S., Lucas, C. G., & Albrecht, S. V. (2024). DRED: Zero-Shot Transfer in Reinforcement Learning via Data-Regularised Environment Design. In Forty-first International Conference on Machine Learning.

We thank the reviewer for their insightful review. We will address the reviewer’s comments below.

it seems the rest of the paper operate on the assumption that it is indeed maximized

This was not our intention. The theoretical and experimental results presented In this work cover the more general setting in which an observation may not provide all the information about the underlying state. The toy example presented in Figure 2 depicts a fully observable environment, but this was for convenience and to make the environment easy to understand.

If there are other parts of the manuscript that suggest we are assuming full observability, could the reviewer point to them for us to make amendments? We are happy to revise this text to ensure our setting is clear.

Evaluation is lacking in breath, it would be great to include a general set of tasks

The new revision will include experiments conducted in a continuous control benchmark in several Brax environments with video distractors, similar to [1]. We point the reviewer to the [proposed changes to manuscript] comment for details on our planned experiments.

Why use the idea of context, if I'm not mistaken, the entire work can be done with a normal MDP and MDPs with contexts can be seen as a encompassing MDP with different starting states.

As pointed out by the reviewer, it is possible to reframe any CMDP as a (partially observable) MDP. However, by explicitly separating the environment properties that must remain constant over an episode (i.e. the context), the CMDP provides several advantages:

• It allows splitting environment instances (levels) into training & test sets. This minimises leakage between the train and test distributions, leading to a more rigorous experimental setup where we can enhance our study of generalisation.
• Explicitly defining the context lets us derive theoretical insights about the overfitting behaviour of RL agents (see for example theorem 3.1). These insights can be leveraged into metrics (e.g. $I(Z;L)$) to quantify how much the agent’s representation overfits to its training data.

We point to [2] for additional analysis and discussion on the advantages of employing the CMDP framing.

The increased compute and parameter size of decoupled approach may also play a role in learning, is there any intuition or validation on the topic

We are currently conducting additional experiments in which we match the number of parameters between the coupled and decoupled architectures. We will include our results in the next revision.

I disagree that "an optimal or near-optimal representation can never be reached under a shared architecture"; from the reduced MDP point of view, I agree that a shared architecture may not be as concise, but still can be optimal (for both policy and value) given that it is a shared architecture.

We believe this disagreement stems from whether one considers compactness a necessary property for an optimal representation. Currently, the notion that the optimal representation yields the most compact embedding from which the learning objective can be optimised is made implicit by our choice of $Z^*$ in l196. However it was not explicitly stated when we define the optimal representation in l194. We will amend the definition in the next revision to reflect this notion.

[1] Stone, A., Ramirez, O., Konolige, K., & Jonschkowski, R. (2021). The Distracting Control Suite--A Challenging Benchmark for Reinforcement Learning from Pixels. arXiv preprint arXiv:2101.02722.

[2] Kirk, R., Zhang, A., Grefenstette, E., & Rocktäschel, T. (2023). A survey of zero-shot generalisation in deep reinforcement learning. Journal of Artificial Intelligence Research, 76, 201-264.

We thank the reviewer for their thorough and insightful review. We agree with the reviewer that the clarity and presentation of the manuscript can be improved, and will incorporate changes to the manuscript in light of their suggestions. We point the reviewer to the comment titled [proposed changes to manuscript] at the top for a summary of the changes we plan to make for the next revision. We will use this thread to address the reviewer’s questions.

What does the delta percentage in Figure 1 represent? It would be helpful if the authors added a brief explanation to the caption.

We have reworked this table (Figure 1, right) as a standalone table with more information. See this link for a preview of the new table that will be included in the next revision: https://ibb.co/vwHnfxS

How is mutual information measured in this study? Currently, it’s detailed only in the appendix, but a high-level explanation in the background section would help readers understand the methodology better.

Could the authors add a performance plot to the figures in the main paper? This would clarify how the information-theoretic metrics relate to final performance outcomes.

Could Figures 7 and 8 be moved from the appendix into the main paper? The authors might need to shorten other parts of the text rather than placing key results in the appendix while referring back to them in the main sections.

Restructuring section 5 and 6 (see [proposed changes to manuscript] comment) will free up space in the main text and allow us to include this information. Namely:

• How mutual information is measured (currently in appendix)
• Train and test scores achieved by the agents (currently in appendix)

In Section 6.1, it would be clearer if the authors defined the dynamics prediction objective from Moon et al. with a brief equation or sentence rather than just citing the source. Similarly, in Section 6.2, could the authors introduce the MICo objective with an equation first?

We will provide a detailed description of each representation learning objective tested in the new revision. Whether this information is included within the main text or within the appendix will be contingent on the available space.

Is it true that $I((Z,Z'); A) + I(Z,Z')$ in general is a good indicator of the algorithm's performance? (Looks like it's the case from Figure 3 and Figure 7.) The authors seem to argue that while actor representation (decoupled) tends to have larger $I((Z,Z'); A)$ while smaller $I(Z,Z')$ and the reverse holds for critic representation, combing them in a coupled representation does not combine the best of both world and resulting in smaller $I((Z,Z'); A) + I(Z,Z')$ and thus worse performance?

We do not observe a clear connection between $I(Z,Z’;A) + I(Z;Z)$ and agent performance. For example:

• DCPG[sh] significantly outperforms PPO[sh], but they have similar $I(Z,Z’;A) + I(Z;Z)$.
• DCPG outperforms PPO, but they have similar $I(Z,Z’;A) + I(Z;Z)$ for both the actor (unhashed bar in figure 3) and critic (hashed bar in figure 3).
• PPG and DCPG perform relatively similarly, but PPG has greater $I(Z,Z’;A) + I(Z;Z)$ than DCPG for both the actor and critic.

In Section 6.1, lines 466–468, the distinction between explanations (1) and (2) is unclear; both seem to attribute reduced overfitting to the larger batch size used in value distillation. If the authors hypothesize that batch size is the main factor here, could they test this by running PPO with larger batch sizes, then measuring the relevant metrics to assess any changes in overfitting?

A broader concern is how the study's conclusions about decoupled versus coupled objectives hold across varying model sizes. For instance, in Section 6.2, the authors note that adding a dynamic prediction objective to the critic impacts certain metrics but harms performance due to an unintended specialization. Similarly, the study attributes overfitting to the coupled representation. Could these conclusions change with larger model capacity (and also training with larger batch sizes as the author claims that this results in more overfitting)? Intuitively, if the shared representation had sufficient capacity (and were trained on larger batches as discussed), could a shared representation still achieve strong performance across the information-theoretic metrics, potentially improving performance without conflicting specializations?

We thank the reviewer for suggesting these experiments, we see how they could help substantiate some of the hypotheses we made in Section 6. We point the reviewer to the [proposed changes to manuscript] comment for details on our planned experiments.

ADDITIONAL EXPERIMENTS

We thank the reviewers for suggesting several additional experiments. We are currently running the following experiments. We will include results in the next revision.

Additional experiments in a continuous control domain. The new revision will feature additional experiments conducted in a continuous control benchmark formed of several Brax environments with video distractors, similar to [1]. The new experiments match the scope of our Procgen experiments. They confirm that the key insights we identified in the original study also apply in a different domain.

Additional representation learning objective. We add Data-Regularised Actor Critic (DrAC) [2] as an additional representation learning objective baseline. DrAC explicitly encourages learning invariances to different augmentations of the same observation. The objectives we had tested so far focused on adding new information (i.e. through prediction tasks) and the removal of superfluous information (i.e. regularisation) occurred implicitly. In contrast, DrAC performs regularisation explicitly, and therefore it fills a gap in our study.

Investigation on the effect of scaling the number of parameters. Following suggestions made by the reviewers, we are conducting experiments with scaled-up architectures of PPO[sh] and PPO (decoupled). We aim to answer the following questions:

• Do we observe the same actor and critic specialisation as representation capacity increases?
• Is additional information being extracted?
• In our initial experiments, PPO (decoupled), fits ~2x as many parameters as PPO[sh], since it has separate representations. How does PPO[sh] compare to decoupled PPO at equal parameter counts?

Investigation on the importance of the batch size & data diversity for representation learning. In Section 6.1, we claim that value distillation improves performance in PPG, but not in PPO, because the batch of data used each time PPG performs value distillation (i.e. for the auxiliary phase) is 32 times larger. We hypothesise that a large batch size is a necessary condition to promote encoding information shared across training levels over information specific to individual levels, and improve performance for PPG. To substantiate this claim, we follow the approach taken in [3]. We are conducting an experiment where we gradually reduce the auxiliary phase batch size in PPG from 32x PPO down to 1x PPO's, while keeping the total number of distillation updates constant.

Lastly, we note that multiple reviewers have suggested conducting a similar analysis for off-policy algorithms. We concur with the reviewers that this analysis would represent a valuable contribution to the field. However, we have decided to leave this study to future work. We do so for two reasons:

1. We currently do not have access to the additional compute that would be required for conducting this study.
2. This study would likely highlight nuances and complexities that are unique to the off-policy setting. We believe this would make the scope of this work too broad for a conference paper. The 10 page limit would be ill-suited for imparting the additional information to the reader, and it would adversely affect the clarity of the current paper contents.

We are working towards a new version of the manuscript, which we will upload before the end of the discussion period. Whenever possible, we will post here partial results for any additional experiment that did not finish by the time we upload the new revision.

[1] Stone, A., Ramirez, O., Konolige, K., & Jonschkowski, R. (2021). The Distracting Control Suite--A Challenging Benchmark for Reinforcement Learning from Pixels. arXiv preprint arXiv:2101.02722.

[2] Raileanu, R., Goldstein, M., Yarats, D., Kostrikov, I., & Fergus, R. (2021). Automatic data augmentation for generalization in reinforcement learning. Advances in Neural Information Processing Systems, 34, 5402-5415.

[3] Wang, K., Zhou, D., Feng, J., & Mannor, S. (2023). PPG reloaded: an empirical study on what matters in phasic policy gradient. Proceedings of the 40th International Conference on Machine Learning (ICML'23), Vol. 202. JMLR.org, Article 1526, 36694–36713.