Studying the Interplay Between the Actor and Critic Representations in Reinforcement Learning
Authors: Samuel Garcin
Trevor McInroe
Pablo Samuel Castro
Christopher G. Lucas
David Abel
Prakash Panangaden
Stefano V Albrecht
Venue: ICLR 2025 Poster (2025)
Keywords: Reinforcement Learning
Representation Learning
Transfer Learning
Information Theory
Primary area: reinforcement learning
TL;DR: When separated, the actor and critic specialise in the information extracted from the environment, but still influence one-another in online RL. Some representation learning techniques further this specialisation and improve performance.

Abstract: Extracting relevant information from a stream of high-dimensional observations is a central challenge for deep reinforcement learning agents. Actor-critic algorithms add further complexity to this challenge, as it is often unclear whether the same information will be relevant to both the actor and the critic. To this end, we here explore the principles that underlie effective representations for the actor and for the critic in on-policy algorithms. We focus our study on understanding whether the actor and critic will benefit from separate, rather than shared, representations. Our primary finding is that when separated, the representations for the actor and critic systematically specialise in extracting different types of information from the environment---the actor's representation tends to focus on action-relevant information, while the critic's representation specialises in encoding value and dynamics information. We conduct a rigourous empirical study to understand how different representation learning approaches affect the actor and critic's specialisations and their downstream performance, in terms of sample efficiency and generation capabilities. Finally, we discover that a separated critic plays an important role in exploration and data collection during training. Our code, trained models and data are accessible at https://github.com/francelico/deac-rep.

Published as a conference paper at ICLR 2025
STUDYING THE INTERPLAY BETWEEN THE ACTOR AND CRITIC REPRESENTATIONS IN REINFORCEMENT LEARNING
Samuel Garcin' University of Edinburgh
Trevor McInroe" University of Edinburgh
Pablo Samuel Castro Google DeepMind, Mila
Prakash Panangaden McGill, Mila
Christopher G. Lucas University of Edinburgh
David Abel
Stefano V. Albrecht University Edinburgh
Google DeepMind, University of Edinburgh
ABSTRACT
Extracting relevant information from stream of high- sigh-dimensional observations is central challenge for deep reinforcement learning agents. Actor-critic algorithms add further complexity to this challenge, formation will be relevant to both the actor and the critic. To this end, we here ex- plore the principles that underlie effective representations for the actor and for the critic in on-policy algorithms. We focus our study on understanding whether the actor and critic will benefit from separate, rather than shared, representations, Our primary finding is that when separated, the representations for the actor and critic systematically specialise in extracting different types of information from the environment-thee actor's representation tends to focus on action-relevant infor- mation, while the critic's 's representation specialises in encoding value and dynam- ics information. We conduct rigourous empirical study to understand how differ- ent representation learning approaches affect the actor and critic's specialisations and their downstream performance, in terms of sample efficiency and generation capabilities. Finally, we discover that a separated critic plays an important role exploration and data collection during training. Our code, trained models and data are accessible at https: : //githbb www/francelicoeeeocooocccococcoe reep
is often unclear whether the same in-
INTRODUCTION
In recent years, auxiliary representation learning objectives have become increasingly prominent in deep reinforcement learning (RL) agents (aarats et al. Dunion al. 2023a). These objectives facilitate extracting relevant features from high dimensional observations, and can help improve the sam- ple efficiency and generalisation capabilities of both value- based (Anand et al. 2019 Schwarzer et al.] (021) and actor- critic methods (Yarats et al. 2022bb Zhang al. 2021 McIn- roe et al. 2023). However, knowing whether a particular repre- sentation learning objective will work and understanding why it works is often difficult due to the interplay between the different components of modern RL algorithms.
T(alo)
(TT
T(loo)
(s)
i
Φ
Φ
Figure Models with shared (left) and decoupled representa- tions (right).
Online actor-critic algorithms like PPO (Schulman et al. 2017) jointly optimise policy improve- ment and value estimation objectives. When parametrised by deep neural networks, the actor (in charge of improving the policy) and the critic (in charge of estimating the value of the current pol- icy) often share the same learned representation o, which maps observations to latent features z.
"Equal contribution. Correspondence ss.garcin, t. t.m.ei.nenn mecinroe)eeee Sed ac.uk

---

Published as a conference paper at ICLR 2025
Cobbe et al. (2021) and Raileanu & Fergus (2021) report that fully separating the actor and critic networks, L.c. decoupling the two (Figure right), improves sample efficiency and generalisation over shared architecture (Figure ! left). We hypothesise that decoupling is effective because it en- courages information specialisation in Φ and dc. which in turn improves performance. To test our hypothesis, introduce metrics to quantify specialisation, and conduct extensive empirical study of three on-policy actor-critic algorithms (Schulman et al. [0117] Cobbe et al. 2021 Moon al. 2022) across discrete and continuous control benchmarks, and under various representation
learning approaches (Raileanu & Fergus, 202T, Moon et al. 2022 Raileanu et al. 2021 Castro
al., 2021) applied the actor, to the critic, or both. We supplement this empirical study by
theoretical characterisation of the information extracted by the actor and critic's 's respective optimal representations.
Once decoupled, the actor and critic representations Φa and dc specialise in capturing different information from the environment. Reported values correspond to a PPO agent trained in Procgen (Cobbe et al., 2020). See 2 and | for formal definitions of the quantities quoted.
Table 1:
% change from using shared representation
I.. is high,
it is possible to
Φ Φ
I(Z;L)
overfit to training levels (enviromment instances).
-20%
+35%
I(Z;V)
use to predict state values.
+37%
+41%
and z' obtained from consecutive timesteps 1, 1 to identify the action taken at timestep /.
((2,2));
+23%
-48%
differentiate between latent pairs obtained from consecutive and non-consecutive timesteps.
1(2:2)
-96%
+324%
Our main findings are summarised below.
Decoupled actor and critic representations extract different information about the environ- ment. This information specialisation, described and quantified in Table systematically occurs the on-policy algorithms and benchmarks covered by our study, and is consistent with the actor's and critic's 's respective optimal representations.
The actor benefits from representation learning approaches that prioritise extracting level- invariant information over level-specific information. This bias for level-invariant infor- mation matters than what specific information quantity is targeted by the repre- sentation learning objective. Nevertheless, approaches antithetical to the actor's inherent information specialisation tend to perform poorly.
Through its role bascline in the actor's objective, decoupled critic will tend to bias policy updates to facilitate the optimisation of its own learning objective. The critic, there- fore, plays important role in exploration and data collection during training. Thus, find that care must be taken when selecting representation learning objective for the critic: certain objectives improve the critic's 's value predictions but may prevent conver- gence to the optimal policy because the objective induces significant bias.
BACKGROUND
RL framework. We follow the framework established by Kirk et al. (202)), which, given fixed timestep budget, lets us quantify the agent's performance in its original training environment (i.c. its sample efficiency) and its performance in held-out instances (i.c. its generalisation capabilities). We consider the episodic setting, and model the environment (S, A. C, S2 R, C, P(c), Pu. 7) with state, action and observation spaces S, A and O and discount factor y. In a CMDP, the reward R S C x & + R, and observation functions Ω $ x C + O (S). and initial state Po - (S) kernels C. with P(c) at the start of each episode. The CMDP is
Contextual-MDP (CMDP) M A
as well as the transition T S x x can change with the context
2

---

Published as a conference paper at ICLR 2025
therefore conceptually equivalent an MDP with state space X S x C. Each context maps to-one to a particular environment instance, or level, and thus represents the component the state that cannot change during the episode. The agent's policy π 0 - (A) maps observations to action distributions and induce a value function V* X + IR mapping states to expected future returns V*(z) (z) E, IEi y'rel, where fre fre) are possible sequences of rewards obtainable when following policy from and until the episode terminates. We define the optimal policy ' as the policy maximising expected returns Ec- Pic-P(e),zo-PP(()) (zo)]. During training, we assume access to a limited set of training levels L ~ P(c). We evaluate sample efficiency by measuring returns over L and generalisation by evaluating on an held-out Lues P(c).
Actor-critic architectures. On-policy actor-critic models consist of policy network Te, and value network Vec. with actor parameters 0 and critic parameters Og (we use // ase when referring to the actor/ critic in this work). When learning from high dimensional observations, such as pixels, representation @ O + Σ maps observations to latent features E Z. When coupled, the policy and value networks share a representation and split into actor and critic heads f and 0. That is, we have a fw On and Vee - fe o on with O - (w, v)) and Oc _ (5. 17). When decoupled, two representation functions 1.. ΦC with parameters (n,, nC) are learned.
PPO and PPG. In this work, we investigate the representation properties of PPO (Schulman et al,) 2017) and Phasic Policy Gradient (PPG) (Cobbe al., 202I), two actor-critic algorithms that have been reported to benefit from improved sample efficiency and transfer upon decoupling (Raileanu & 2021 In PPO, the actor maximises Cobbe et al.
Fergus 2021
Te (a)lor) Ae, clip(
To (aslo)) To (ar(o:))
Jx (0 Eg
min
1 c,1++)AA)) HH (arlo)))
ABna (ar)
where 0 Aold are the actor weights before starting a round of policy updates, a batch of trajecto- ries collected with Toa A, is an estimator for the advantage function at timestep t, H()) denotes the entropy and and 33 are hyperparameters controlling clipping and the entropy bonus. The critic minimises
(v()))
1 B
Σ (Vec (or) vi)2 B
where VÚ are value targets. Both À and V are computed using GAE (Schulman et al. 2016). PPG performs auxiliary phase after conducting PPO updates over N₂ policy phases. To prevent over- fitting, the auxiliary phase fine-tunes the critic and distills value information into the representation from much larger trajectory batches Baur UiEI.... B;, using the loss Ejoim ly + lauxs with
laux (Ax)
Σ
(Vgu (ot) V/)² + DKL (aelot)) (ar
(3)
Bausl
where B. controls the distortion of the policy. When decoupled, frauz A your Φ distills value information into representation parameters Va through an additional head yauz, When coupled, your = U. and stop-gradient operation on yy ensures 77 is updated by the critic during the auxiliary phase only.
Mutual information. We study the information embedded in features outputted by d. To do so, we propose metrics based on the mutual information I(X;;) Y), measuring the information shared between sets of random variables X and Y. defined as
I(X;)) = H(X) - H(Y H(X,Y)= = ΣΣ p(x, v) log p(r)p(u)
p(x, v)
(4)
where integrals replace sums for continuous quantities. I(X; Y) is symmetric, and quantifies how much information about Y obtained by observing X, and vice versa. Similarly, the conditional mutual information I(X;Y\Z) measures the information shared between X and Y that does not depend on Z. We measure mutual information using the k-nearest neighbors entropy estimator proposed by Kraskov et al. (004) and extended to pairings of continuous and discrete variables by Ross (2044). We briefly introduce notation for random variables used in following sections. L. P(e) denotes the set of training levels drawn from the CMDP context distribution. A, R, O,
3

---

Published as a conference paper at ICLR 2025
O², X'' and X' are sets constructed from transitions (ae, r2, Ot, 01+1, Lt, F++1) uniformly sampled from a batch of trajectories collected in L using policy т. and are latents features, with =()),2' 6(o'). We construct V using the rewards obtained from t until episode termination, with 21
3 CATEGORISING AND QUANTIFYING THE INFORMATION EXTRACTED BY LEARNED REPRESENTATIONS
To conduct analysis of the respective functions of the actor and critic representations, analyse the information being extracted from observations at agent convergence. We propose four mutual information metrics to measure information extracted about the identity of the current training level, the value function, and the inverse and forward dynamics of the environment. Each metric relates to quantities relevant to the actor and critic's respective learning objectives, and to the agent's general- isation performance. They are introduced in turn below.
Overfitting. Our first metric, [(; L), quantifies overfitting of the actor and critic representations to the set of training levels, as it measures how easy is to infer the identity of the current level from Z. We follow a similar reasoning as Garein al. (2024) to derive an upper bound for the generalisation error that is proportional to I[Z [
Theorem 3.1. The difference in returns achieved in train levels and under the full distribution, or generalisation error, has an upper bound that depends on A(Z L). with
2D [L]
Eeoet((()=zo (xo)] Pto- (zo)) S V
X L),
(5)
where e ~ U(L) indicates c is sampled uniformly over levels in L, D is a constant such that [V*())) S D/2, Va, and Z is the output space of the actor's learned representation.
Value information. The second metric quantifies I(2;V), the mutual information between Z and state values. While a high 1(2c: V) facilitates the minimisation of Ły (Equation 2 ). we wish to un- derstand whether increasing 1(Z4; V) is always desirable. In fact, we find an apparent contradiction on this matter in prior work: Cobbe et al. (2021) and Wang et al. (023) report that value distillation into the actor's is representation (which implies high I[ZA; VJ) improves sample efficiency and gen- eralisation of coupled and decoupled PPG agents, whereas Raileanu & Fergus) (2021) and Garcin al. (2024) report positive correlation between generalisation and a high Ey for coupled PPO agents (which implies a low I(ZA; V)).
Dynamics in the latent space. The remaining two metrics investigate the transition dynamics T. Z x A + (Z)) within the latent state space Z spanned by the representation. We will see in later sections that the reduced MDP (Z, A. Tz, Rg 7) spanned by the actor or critic's 's repre- sentation tend to have distinct T., which often markedly differ from the transition dynamics TT in the original environment. I(Z; Z') measures how easy obtained from consecutive observations from latent pair obtained from non-consecutive observa- tions. I((Z, Z'); A) quantifies how easy is to predict the action taken during transition (o, a, o') given the pair (o(o), é(o')). In Theorem 3.2 establish that T. maintains the Markov property of the original MDP when both of these metrics attain their thcoretical maximum
differentiate latent pair (ø(0), o(o())
Theorem 3.2. ifT x A + (X) satisfies the Markov property, and we have 1(X,X') A) [((Z, Z A) and I(X I( z') for any X,X', A, Z, Z' collected using policy then T, x & + (Z) satisfies the Markov property when following म. Ta always satisfies the Markov property if the above conditions hold for any π.
Given that & only induces T. for the current æ in the on-policy setting, we make the distinction between T. being Markov when following л and the more general notion of TT being Markov when following any policy. Crucially, Theorem 3.2 generalises the equivalence relations obtained by Allen
Proof's for the theoretical results presented this work are provided Appendix
'The Markov property is satisfied for a MDP (2, A, Tz, Rs, 7) if and only if "M (z441 21
Ts Te(++ (r, 2,) and R (#++1 (ar-i, 1 R₂ 2++ ac, 21), Va A, 2,2 E 7, 11.
4

---

Published as a conference paper at ICLR 2025
Redused MDP Ge 9
Reduced MDP for o::
Mags to stales ln kevel
dic.
Maps to in levele 2and
Figure 2: (Top) the initial observations and state spaces of three levels from the assembly line environment (Bottom) the reduced MDPs spanned by d and dz.
al. (2021) continuous metries. As such, I((Z, Z"): A) and Z") quantify how close any comes to have Tx satisfy the Markov property. They remain applicable in settings in which it isn't t practical (or even possible) for T. to satisfy the Markov property, e.g. when é has finite capacity and bottlenecks how much information can be extracted from raw observations, or when observations are not Markov
4 INFORMATION SPECIALISATION IN ACTOR AND CRITIC REPRESENTATIONS
Raileanu & Fergus (201]) Cobbe et (2021) have attributed the performance improvements ob- tained from decoupled architectures to the disappearance of gradient interference between the actor and critic, and to the critic tolerating higher degree of sample reuse than the actor before overfit- ting. We propose different interpretation: given their different learning objectives, the actor's 's and critic's optimal representations (defined below) prioritise different types of information from the environment. While not incompatible with prior interpretations, claim is stronger. We posit that an optimal (or near-optimal) representation for both the actor and critic will generally be impossible under shared architecture.
Definition 4.1. Given the model m & fw 6 and associated loss (m( Φ). an optimal representa- tion 6² O + Z" satisfies the conditions:
1. Optimality conservation. min em (w, (w, min @ (ws ())
2. Maximal compression. 6" E arg ming 1²), with & the set of all & satisfying condition 1.
We will provide theoretical insights on the respective specialisations and mutual incompatibility of 6a, and oc which we further highlight through a motivating example. In our example, depicted in Figure the agent inspects parts for defects levels drawn from P(c). A level is characterised by a particular combination of part specifications, number and ordering, each part having probability P of being defective. At each timestep. the agent observes the part specifications for the current level, which parts are on the assembly line and which part is up for inspection. The agent picks action E A {ao accept, ai reject) and moves to the next part. It receives reward R т+ when correctly accepting/rejectigg r. when it makes a mistake, with r+ r.. The episode terminates
assembly line. The agent is trained
set L of
good/defective part and early when the agent accepts defective part, otherwise it terminates after N, timesteps, where N. the number of parts level E
In contrast, Allen al. 2021 assume observations always Markov.
5

---

Published as a conference paper at ICLR 2025
4.1 THE ACTOR'S 'S OPTIMAL REPRESENTATION
The combinatorial explosion of possible specifications and part assortments means Φ should ide- ally map observations reduced MDP spanning much smaller state space that in the original environment. However, Φ should still provide the information necessary to select the optimal ac- tion each timestep of each level, including those not in the training set.
Dynamics of the reduced MDP. Under our definition, the mapping
[Ahor if a- 0g
(0)
(6)
if a1
satisfies the conditions for being an optimal representation, and spans the reduced MDP Figure 2 (bottom left). This reduced MDP describes the perceived environment dynamics when only observ- ing the latent states in 2 By construction, 1((Z), Z '): A) is guaranteed to be maximised when following the optimal policy. We note that, even if I(X;; X') is maximised under x² in the original environment, we have I i( - = 0 in the reduced MDP. other words, the current reduced state yields no information about the next reduced state, and vice versa.
Overfitting to training levels. In our assembly line example, an overfit ΦΛ with high I(Za; L) may leverage this information to first identify, and then solve certain levels. This overfit Φ may be optimal over I., but would be heavily biased to the training set, and fail to generalise to unseen levels that do not satisfy certain spurious correlations. For example: "if there then I must be in level 1, and, in level 1, should reject the triangle". In contrast, a representation that invariant individual levels, i.e. with A(ZA; 4 L) O, guarantees zero generalisation error under Nevertheless, achieving I(ZA5 L) = 0 is not achievable under certain conditions,
three objects on the line
Theorem 3.1
stated in Lemma 41
Lemma 4.1. L) > Of > 0 and Bzk Cj E 2 x L such that μlzz (e)) & u(zz).
In our example, I(O;,) T.) 0, since level can be identified from its observations. We can verify that the second condition holds for 6 by inspecting the stationary distributions in a particular level e and over all levels.
PF if AAo pF
{ [Pe if AAA - (pr. if = AA1
m22)
A(z|c)
= 2
where 1 PF and PF is the defect probability when in level We may have PF * P.. since individual levels do not all have the same distribution of defective parts. In other words, while it is useful to reduce I(Za; L) to guard against overfitting, the optimal representation dA may carry some irreducible information about level identities.
Value and dynamics distillation can induce overfitting. we can employ the chain rule of mutual information to decompose the information & captures about some arbitrary quantity Y as
(8) where ((Z:YL)) the level- cevel-invariant information encoded about Y, and I(Z; L. - I(Z; L|Y) is the level-speccficc information being encoded about Y. Thus, encouraging ΦA to capture extraneous information about state values or transition dynamies can cause I(Z I.) to increase and promote overfitting, In Lemma 4.2 we show that increasing I(ZA; V) or I(ZA; Z'x) (which is not necessary for obtaining an optimal Φ as discussed above) can cause I(ZA; L) to increase as well.
I( = I(Z;Y + I(2;
I(Z; L|Y),
Lemma 4.2. I(Z:L) monotonically increases with a) I[Z V) I(Z:VIL) and b) I(Z: Z') (')..
In fact, I(Z; I.) will increase when I(2;V)= > I(Z;VIZ) (or when (2; > I(2; (L)),,,, that is when the information encoded carries a level-specific component. our example, the higher state values only occur in subset of levels, since the optimal value for any given state depends how many parts are left to inspect. State values therefore inherently contain level-specific information, and may induce overfitting if this information is captured by Φb. This challenges the notion that value distillation would systematically improve the actor's 's representation.
The key implications of the above are 1) While achieving zero I) is not always possible, a high I(Z,; L) implies Φ is overfit; 2) Increasing I(Z or I(Z V) is not necessary for obtaining Φ and will increase I( L.) if the information captured is level-specific.

---

Published as a conference paper at ICLR 2025
((Z-2)
KIZZ%A
DCPG
2.0
2.5
3.0
3.5
0.30
0.45
0.60
0.75
0.0
0.2
0.4
0.6
0.16
0.24
0.32
PPO
PPG
DCPG
[O:L)
[O,V)
110:0)
((0,0%)
Figure 3: Mean and 95% confidence interval aggregates of I( MI( (top/bottom rows) in Procgen. Gray bars indicate ((Z;))/((O:)) for a shared db. Blue and orange bars indicate [(Z:-) measured for Φλ and &c when employing decoupled architecture. Pink bars indicate I(O;-) measured when using decoupled architecture. X-axes are shared across top and bottom. For all algorithms, decoupling induces specialisation consistent with $4
4.2 THE CRITIC'S OPTIMAL REPRESENTATION
6 has higher (2: L) than a The reduced MDP spanned by ot is depicted Figure 2 (bottom right). In order to ensure perfect value prediction, de maps each possible optimal state value to different element in Z and it maximises I(Z;; V) by construction. I(Z Zg') is also high due [R(r, a) + t= (V=())) This points to V being a quantity inherently more level specific than the optimal action for the current state, because V" encodes information pertaining to all possible future states of the current level. We should then expect that, in general, I(Z6; L) > I L), implying that decoupling the actor and critic helps with i(Zx; L) regularisation.
to the recurrence (*(x) = En [[(,,
6 is not compatible with x². Paradoxically, while 0z would necessitate trajectories collected using the optimal policy in order to be learnt, in our example is not possible to have an optimal policy that only depends on z The information contained in = is not sufficient for picking the optimal action in any given timestep, and therefore the best response is to always pick a in order prevent early termination. Therefore, in addition the information prioritised by de being in general irrelevant to π², employing shared ó with finite capacity may prevent extracting the necessary information to execute the optimal policy.
PPO[sh] Normalized Returns
1.8
4.3 CONFIRMING SPECIALISATION EMPIRICALLY
We conclude this section by studying the representations learned by PPO (Schulman et al. 2017). PPG (Cobbe et al. 2021] and DCPG (Moon al. 2022), close variant of PPG that employs delayed value targets to train the critic and for value distillation. We evaluate all algorithms with and without decoupling their representation. We conduct our experiments in Procgen (Cobbe et al., 2020), a benchmark of 16 games de- signed to measure generalisation in RL. We report our main observations in below, with extended results and details on our methodology included in Appendix C.2
1.4
(Test) PPCtsh] (Test)
PPCtch) (Tran)
0.6 1.2 Number Parameters 4,8
2.4
Figure 4: Effect of parame- ter scaling in coupled (blue) and decoupled (orange) PPO. Scores normalized by model per- formance at 0.6M parameters.
Specialisation is consistent with Φ and dec As no algorithm achieves optimal scores in all games, now consider the sub- optimal representations QA and de realistically obtainable by the end of training. Figure 3 we observe clear specialization upon decoupling consistent with the properties expect for Φ and 6- de has high I(Z; V), I(Z; Z') and ((Z L), while Φ specializes I((Z, A))
Decoupling is more parameter efficient. Since decoupled representations fit twice many param- eters, it is fair to wonder whether the performance improvements are mainly caused by the increased model capacity. To test this, measure performance scale model size in shared and de-

---

Published as a conference paper at ICLR 2025
coupled architecture in Figure
4 Surrrsingll, the decoupled model turns out to be more parameter efficient, and still outperforms a shared model with four times its parameter count.
On Markov representations. According to Theorem 3.2
representation is considered Markov when both I( Z'); and I(Z; Z') are maximized. However, our observations reveal an inter- esting pattern during decoupling: the actor representation shows increase in I((ZA, A)); ))::A] A) but a decrease in IZZ Z'A), while the critic representation shows the opposite effect- a decrease in I((Ze Zc);; A) and an increase in I(2c; Zc). This divergent behavior suggests that neither the actor nor the critic networks inherently benefit from maintaining a Markov representation. This finding aligns with theoretical expectations, as neither O nor Φc need be Markov to be optimal. Furthermore, we found no significant correlation between the sum of these mutual information terms (((ZZ A) + ((: and agent performance, as shown by comparing Figure and Figure 13
5 REPRESENTATION LEARNING FOR THE ACTOR
In this section, we study how different representation learning objectives affect O in PPO, PPG and DCPG, We consider advantage (Raileanu & Fergus 2021 and dynamics (Moon et al. 2022)
prediction, data augmentation (Raileanu et al. (2021) and MICo (Castro et al.] 2021). an objective explicitly shaping the latent space to embed differences in state values. We study these objectives in Proegen (Figure 5) and in four continuous control environments with video distractors (Mcinroe & Garein 7025)(Figure) LTp.
Representation learning impacts information specialisation. As expected, applying auxiliary tasks alters what information is extracted by the representation. Dynamics prediction generally enhances the specialization of Φ by increasing I((2, Z"); A) while reducing the three other quanti- ties. Conversely, MICo produces the opposite effect- in most cases, it increases I(Z; Z²). I(2; V), and I(Z; L) at the expense of I((Z, Z'); A). The effects of the last two objectives are not as clear- cut. Data augmentation produces little change in each quantity. while advantage prediction tends to reduce the measured mutual information, but is inconsistent in the quantities it affects. Performance- wise, data augmentation improves train and test scores for all algorithms; dynamics prediction tends to improve performance for PPG and DCPG; MICo generally decreases performance, and advan- tage prediction makes no noticeable impact. Based on these findings, we find advisable to not use objectives increasing A(ZA; L) or playing directly against the specialisation of dA..
On the importance the batch size and data diversity. We now turn our attention to an apparent contradiction in the relationship between value distillation and performance. Decoupling PPO, and thus completely forgoing value distillation, leads to improved train and test scores (Figure [3). How- ever, PPG and DCPG perform extensive value distillation (four times many distillation updates as coupled PPO in our experiments), and achieve significant performance improvements. Crucially, conducting value distillation every N, policy phases ensures the batch size Baux is N₂ times larger than the PPO batch size, greatly increasing its data diversity. |Wang et al. (2023) report that increas- ing diversity key driver of performance improvement at equal number of gradient updates. We reproduce their experiment in Figure I2 while also tracking the evolution of I AZ L) and I(ZA; V) as Baux increases. We find that the increase in I(ZA: L) in PPG is mainly caused by these addi- tional gradient updates, whereas agent performance and I(Za; V) only increase when performing value distillation over large, diverse training batches. Increasing data diversity should promote level- invariant state value information to be distilled into é.A and we conclude that it is this level invariant information plays a key role in improving agent performance in PPG and DCPG.
Our experiments further reveal that dynamics prediction yields significant performance improve- ments, but only when applied to algorithms with large batch sizes (specifically in PPG and DCPG, not in PPO). In comparison, data augmentation. which increases data diversity, demonstrates per- formance benefits across all three algorithms. These findings suggest an important hypothesis: the actor benefits more from encouraging da extract level-invariant information about any arbitrary quantity Y, rather than from specifically choosing what that quantity y should be (such as state values or dynamics).
8

---

Published as a conference paper at ICLR 2025
((Z:L)
(ZZZ)
((ZZ2))
PPO
PPO-MICaIA)
PPO+DIAU
PPO+DriAI
PPG
PPG-MICalA)
PPG+DIA)
PPG+DDrA)
DEPG
DCPG MiCa(A)
DCPG+DIA)
DCPG+Dr(A)
1.8
2.4
3.0
3.6
0.4
0.8
1.2
0.0
0.3
0.6
0.9
0.00
0.15
0.30
PPO-MICa(A)
PPO
+PO+Dr(A)
8
MRGaMICa(A)
PIG+DDA)
PPG+Dr(A)
DCPG
DOPG MICa(A)-
DCPG:
DCPG+Dr(A)
Jran
[(O)L]
(OOV)
1(000")
1((0,0')AA)
Figure 5: Mcan and 95% confidence intervals of ((Z;))/((;))) (top/bottom) for actor (blue) and critic (orange) representations in Proegen. Information measured from agent observations shown in pink. X-axes are shared across top and bottom. Auxiliary tasks shown are MICo, dynamics prediction (D), and data augmentation (Dr) applied to the actor (A).
THE CRITIC'S S OBJECTIVE INFLUENCES DATA COLLECTION
We now consider how the same set of representation learning objectives affect the critic's represen- tation and present our results in Figures and The effect of given objective on the information extracted by &c is consistent with how they would have affected Φ in the previous section. How- ever, report two surprising findings: a) Without conducting any value distillation, decoupled PPO has a 37% higher I(Zx; V) than shared PPO (Table [. and b) the information specialisation of Φo incurred by applying objective the critic is often observed in A albeit to a lesser extent. Given that the two representations are decoupled, how can an objective applied to Φc affect Φ
As maintain different optimisers for the actor and critic, their only remaining interaction in decoupled PPO through 2R (Equation [I): Å, being computed from the critic's value estimates. Therefore, at least one of the following hypothesis must hold:
1. Data collection bias. Through J updates, the critic biases to collect trajectories containing information relevant to its own learning objective. This information could then leak through O because more of this information is contained in its input. In this scenario, it is not necessary for Φ to become more proficient at extracting critic-relevant information.
2. Implicit knowledge transfer. The advantage targets in J induce information transfer between Φc and Φ when applying the gradients Vo, J Here, dA becomes proficient at extracting the same information &c extracts.
The first hypothesis broadly holds in our experiments: in most cases, applying MICo to the critic increases I(O; V) and O"), and applying dynamics prediction increases I((0,0'): A Fur- thermore, I(O; V) increases when PPO is decoupled (Figure B). Without the critic's 's influence, there would be no direct incentive for the actor to collect data that contains value information, since no value distillation is taking place.
To test the second hypothesis, we measure the compression efficiency, applicalle whenever I(O;-) > 0, and defined
C(Z:-) min (1(7:-) (I(O:-)' )
In contrast, I(O;I.) does not vary significantly, given that the policy does not control which level is played episode.

---

Published as a conference paper at ICLR 2025
For example, C(Z measures the fraction of available information in I(O; V) that is extracted by δ In Tables and 3 we report that A(Z V) does not signficantly change between PPO[sh] and decoupled PPO, or when MICo is applied to the critic in decoupled PPO. This result appears to disprove the second hypothesis, at least in PPO. We cannot formally confirm whether implicit knowl- edge transfer occurs for PPG and DCPG, explicit knowledge transfer already occurs through value distillation.
Finally, we highlight that this data collection bias generally leads worse performance (Figure 13 Interestingly, employing different representation learning objectives for the actor and the critic results in surprising interactions. Consider PPO in the Procgen setup: when applied separately, Adv(A) and MICo(C) have respectively a neutral and adverse effect on performance, the latter be- ing caused by the data collection bias caused by the critic. When applied together, they exhibit sharp performance gain, and bring I(O; V) down (Figure 10). This suggests that certain actor rep- resentation objectives could improve performance by cancelling the bias in data collection induced by the critic.
7 RELATED WORK
Representation learning in RL. Representation learning objectives have been used in RL for
variety of reasons such as sample efficiency (Jaderberg et al. 2017 Gelada et al. 2019 Laskin et al.,
2020a Lee et al. 2020 [askin al., 2020b), planning (Sekar et al. | [2000] McInroe et al. 2024). disentanglement (Dunion al. (023b), and generalisation (Higgins et al. 2017: Li et al. 202
Some works focus on designing metrics motivated by theoretical properties such as bisimulation
metrics, pseudometries. decompositions of MDP components, successor features (Ferns et al.
2004
Mahadevan & Maggioni 2007 Dayan 1993:
Castro 2020
Agarwal et al. 2021 Castro
al. 2023)
Analysing representations in RL. Despite the large body of research into representation learn- ing objectives in RL, relatively little work has gone into understanding the learned representations themselves (Wang et al., 2024). Several works use linear probing to determine how well learned representations relate environment or agent properties (Racah & Pall 2019 Guo al. 2018 Anand et al.] 2019)] Zhang et al.) (024). Other works analyse the learned representation functions via saliency maps which help visualise where an agent is "paying attention" (Rosynski et al. 2020
Atrey et al. 2020 Dunion et al. 2023a).
8 CONCLUSION
In this work, conducted an in-depth analysis of the representations learned by actor and critic networks in on-policy deep reinforcement learning. Our key findings revealed that when decoupled, actor and critic representations specialise in extracting different types of information from the en- vironment. We found that the actor benefits from representation learning objectives that promote extracting level-invariant information. Finally, discovered that the critic influences policy up- dates to collect data that is informative for its own learning objective.
We identify three important research directions to be tackled in future work: 1) Broadening this study's 's scope to cover different network architectures, representation learning objectives, or RL al- gorithms (such as off off-policy RI. (Haarnoja et al. 2018 ); 2) Employing these findings to design new representation learning objectives for the actor that target the critic's ' influence on data collection to devise new exploration strategies in online RL.
information; 3) Leveraging
REPRODUCIBILITY STATEMENT
Reproducibility can be challenging without access to the data generated during experiments. To assist with this, we will make all of our experimental data, including model checkpoints, logged
'By the data processing inequality must have I(O:.) 2 1(7;-), and C(Z:-) cannot be larger than 1. We enforce this upper bound as our estimator sometimes underestimates I(O;)) for high dimensional observations.
10

---

Published as a conference paper at ICLR 2025
data and the code for reproducing the figures in this paper openly available at/tttps///gthhb com/ wrranceeico/deac---E- rep
REFERENCES
Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Mare G. Bellemare. Con- trastive behavioral similarity embeddings for generalization in reinforcement learning, In 9th International Conference on Learning Representutions, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL hittps://ppenreee ew. net forum? di qda7-sVg84
Cameron Allen, Neev Parikh, Omer Gottesman, and George Konidaris. Learning markov state abstractions for deep reinforcement learning. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Ad- vances in Neural Information Processing Systems 34: Annual Conference Neural In- formation Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp- 8229 8241, 2021. URL (el.eecc6822279eeldd22aldaaaaa
neur httpsc/ppppprrrrepeeppppppp /2011 /hash/ -bstrractcbbattatc.ccc mansttract.tmll
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Alex-nlex Côté, and R. De- von Hjelm. Unsupervised state representation learning in atari. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché- Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Informa- tion Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 8766-8779, 2019, URL https: /tttps://ppoocee.ngss.comm www.:eeiww...................... CC httppppper229///aaii.. 2tt.996aaaee1116777-aaabaaaaaaa Abst html
Akanksha Atrey, Kaleigh Clary, and David D. Jensen. Exploratory not explanatory: Counterfac- tual analysis of saliency maps for deep reinforcement learning. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenRe- view.net, 2020. http:httts/:/oommme
net. www.fforrroooooomm 3m1 BFDB
Marc G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning en- vironment: An evaluation platform for general agents (extended abstract). In Qiang Yang and Michael J. Wooldridge (eds.). Proceedings of the Twenty- Fourth International Joint Conference on Artificial Intelligence, IJCAI Buenos Aires, Argentina, July 25-31, 2015. pp. 4148-4152. AAAI Press, 2015. URL http:http:/ttit .rg raasstrac111/// 585
Martín Bertrán, Natalia Martínez, Mariano Phielipp, and Guillermo Sapiro. Instance-based gen- eralization in reinforcement learning. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Had- sell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Process- ing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020, URL httpttpps//ppccceiiiiii wwwp:procce.eiiio..eeee.e....... neurips.cc/ paper/2020/har
62674fc2 2at.22/44422220000022236662----3
.html
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic markov decision processes. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAA/ 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY. USA, February 7-12, 2020, pp. 10069-10076. AAAI Press, 2020. URL https://aaall
org/ htts/js//iae index..ex...o phwp www.AAaarticeevviii.....
6564
Pablo Samuel Castro, Tyler Kastner, Prakash Panangaden, and Mark Rowland. Mico: Im- proved representations via sampling-based state similarity for markov decision processes. In Mare Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wort- Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Confer- ence on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp. 30112- 2021. URL htttps://proceeeeetse Ac- htmm
wwt.erii...cppppeeeeee cc/paper/
2021/hash 0d0 http:///bbbbeeeeeeeeeeeeeeeeeeee
Pablo Samuel Castro. Tyler Kastner, Prakash Panangaden, and Mark Rowland. A kernel perspective on behavioural metrics for markov decision processes. TMLRR 2023.
11

---

Published as a conference paper at ICLR 2025
Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural genera- tion to benchmark reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning. ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 2048-2056. PMLR, 2020. URI.ttpp:// www.httppp/prrcceeinnge
m w.p.rpes//ee. https//vv19////bbbbbe cttp://obbbbbbaebb html
Karl Cobbe, Jacob Hilton, Oleg Klimov, and John Schulman. Phasic policy gradient. In Ma- rina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Ma- chine Learning, ICMI. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Ma- chine Learning Research, pp. 2020-2027. PMLR, 2021. URLhttp: hnttp://procent ngs. mmlr. press/v139 ctte9//obbbbaaaaaaaa.. html
Peter Dayan. Improving generalization for temporal difference learning: The successor representa- tion. Neural Computation, 5:613 1993.
Mhairi Dunion, Trevor McInroe, Kevin Sebastian Luck, Josiah Hanna, and Stefano V. Albrecht. Conditional mutual information for disentangled representations in reinforcement learning. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko. Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neu- ral Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023a. URL http: http:/pprrrrrniss..cccppppeeerrr Sacpaperrflle//eeeeaeeeeee fww.s/ wttese/ppppeee22.... htm
Mhairi Dunion, Trevor McInroe, Kevin Sebastian Luck, Josiah P. Hanna, and Stefano V. Albrecht. Temporal disentanglement of representations for improved generalisation in reinforcement learn- ing. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023b. UtL/itttss://opeeeeeee http:/eewwwieeeee.....e..ee Tee=ssgP6aIInnn
Norman Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite markov decision pro- cesses. In Conference on Uncertainty in Artificial Intelligence, 2004.
Samuel Garein, James Doran, Shangmin Guo, Christopher G. Lucas, and Stefano V. Albrecht. DRED: zero-shot transfer in reinforcement learning via data-regularised environment design. In Forty-first International Conference on Machine Learning. ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https: http::////peeeeee..eoooo http:/wwwww.weeeeeeeeee
uuu 9 6RRO1
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofrr Nachum, and Mare G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 2170-2179 PMLR, 2019. http:: mlr. wwt.:ppsssss/eeeeeeeeee
9a. html
Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Bernardo A. Pires, and Rémi Munos. Neural predictive belief representations. ArXiv preprint, abs/1811.06410 2018. URL https:
//axxiv.oro/aas//1111.000
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stackholmsmüsssan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 1865. PMLR, 2018. htt.:htt://ppee htLL:tttpp//rrrceinggss mlr. www.prres/vvvi. hww.00//aaaaooo/a jalbb.
Bb. html
Irina Higgins, Arka Pal, Andrei A. Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick. Charles Blundell, and Alexander Lerchner. DARLA: improving zero-shot transfer in reinforcement learning. In Doina Precup and Yee Whye Teh (eds.). Proceedings of the 34th International Conference Machine Learning, ICML 2017, Sydney, NSW, Australia. 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research pp. 1480-1490. PMLR, 2017. URL U htttt:://procee/ii ngs. ml wwwppp/sww/..h...ppppppp ns 1ww.s117.... html
12

---

Published as a conference paper at ICLR 2025
Shengyi Huang, Rousslan Fernand Julien Dossa, Chang Ye, Jeff Braga, Dipam Chakraborty, Ki- nal Mehta, and João G.M. Araújo. Cleanrl: qugh-quuaity single-file implementations of deep reinforcement learning algorithms. Jourmal of Machine Learning Research, 23(274):1-18, 2022. URL http:://mmmmmm.../peeeeeeeeee
v2tp//23/2111///// 1342. html
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czamecki, Tom Schaul, Joel Z. Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. htttp:://oppereeeeee net net/forum? =
Minqi Jiang, Edward Grefenstette, and Tim Rocktäschel. Prioritized level replay. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 194.-440-444 4950. PMLR, 2021. URL hRL http::tpppoccce.ngggg...... mlr htd...pees//139. jiang2lb.brs
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot gener- alisation in deep reinforcement learning. Journal of Artificial Intelligence Rexearch, 76:201-2660 2023.
Alexander Kraskov, Harald Stögbauer, and Peter Grassberger. Estimating mutual information. Phys- ical Review E Statistical, Nonlinear, and Soft Matter Physics, (966))066138 2004.
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Rein- forcement learning with augmented data. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Had- sell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances Neural Information Process- ing Systems 33: Annual Conference Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020a. / paper/2020/har
nwwrips....... cc/ www.eccccaaaaaa... 681 (1d.11110ee122222220000000000000
.html
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: contrastive unsupervised representa- tions for reinforcement learning. In Proceedings of the 37th International Conference on Ma- chine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 5639-5650. PMLR, 2020b. URL htttp::/proceedinggg v 1ilvl996 19/ www.inakaana...... html
Kuang-Husd Huei Lee, Ian Fischer, Anthony Liu, Yijie Guo, Honglak Lee, John Canny, and Sergio Guadarrama. Predictive information accelerates learning in RL. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL
wwt httpppppp//222/aat... html
Bonnie Li, Vincent François-Lavet, Thang Doan, and Joelle Pineau. Domain adversarial reinforce- ment learning. ArXiv preprint, abs/2102.07097, 2021. URL https://arxiv. abs/ 2102. 07097
Sridhar Mahadevan and Mauro Maggioni. Proto-value functions: A laplacian framework for learn- ing representation and control in markov decision processes. JMLR, 2007.
Trevor Mclnroe and Samuel Garcin. Pixelbrax: Learning continuous control from pixels end-to-end on the gpu. ArXiv preprint, abs/2502.00001,, 2025. htt//tttpa///axxi.orr///bb//////
00021
Trevor McInroe, Lukas Schäfer, and Stefano V. Albrecht. Multi-horizon representations with hier- archical forward models for reinforeement learning. TMLR, 2023.
Trevor McInroe, Adam Jelley, Stefano V. Albrecht, and Amos Storkey. Planning to go out-of- distribution in offline-to-online reinforcement learning. In RIC., 2024.
13

---

Published as a conference paper at ICLR 2025
Seungyong Moon, JunYeong Lee, and Hyun Oh Song. Rethinking value function learn- ing for generalization in reinforcement learning. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural In- formation Processing Systems 35: Anmual Conference on Neural Information Process- ing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022. URL
wwwp..ccpp.pp........e.pppp wtwes//ppprr/22222/a Abstract Conference.
el (1laaboddeeeeeofffddleeeeeeee e115222deeeesfff6aeeeeeeeessseee oerersseeeeettttttcaoooeeeeeeeee html
Evan Racah and Christopher Pal. Supervise thyself: Examining self-supervised representations
in interactive environments. ArXiv preprint, abs/1906. 11901, 2019. URL https://arxiv..
org/ htg/abs// 1906. 11951
Roberta Raileanu and Rob Fergus. Decoupling value and policy for generalization in reinforce- ment learning. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th Interna- tional Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, vol- ume 139 of Proceedings of Machine Leurning Research, pp. 8787-8798. PMLR, 2021. URL html
.mr
Roberta Raileanu, Max Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Auto- matic data augmentation for generalization in reinforcement learning. In Mare' Aurelio Ran- 7ato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neu- ral Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 5402-5415, 2021. URL 2b38c2df 4 91w.499999999991
neurips.... http:///ia...oioosoeooo////////a cc/paperr2021/hastt 88886 Abstract .htll
Brian C Ross. Mutual information between discrete and continuous data sets. PloS one, 9(2): e87357, 2014.
Matthias Rosynski, Frank Kirchner, and Matias Valdenegro- Are gradient-based saliency maps useful in deep reinforcement learning? ArXiv preprint, abs/2012.01281, 2020. URL https:
// //arxiv.org/aaae wttpv//.wrrgaabs/22....111
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-
dimensional continuous control using generalized advantage estimation. In Yoshua Bengio and Yann LeCun (eds.), 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016. URL http:
//arxiv.org abs 1506. 02438
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv, 2017.
Max Schwarzer, Ankesh Anand, Rishab Goel, R. Devon Hjelm, Aaron C. Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL
net.iiewww.e.eeeeeeeee.eee.eeeee forum?id
uCOEP ZwRUul
Ramanan Sckar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In Proceedings of the 37th In- ternational Conference on Machine Learning. ICML 2020, 13-18 July 2020, Virtual Event, vol- ume 119 of Proceedings of Machine Learning Research, pp. 8583-8592. PMLR, 2020. URL .m.. wtr.ppres/vll http://1111s.aee...... setssskka2222e .htmm2maa..mlll..
Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite- challenging benchmark for reinforcemeett learning from pixels. ArXiv preprint, abs/2101.02722. 2021. URL BARLLattss//rrxxv.ooo///a....
02722
Han Wang, Erfan Miahi, Martha White. Marlos C Machado, Zaheer Abbas, Raksha Kumaraswamy, Vincent Liu, and Adam White. Investigating the properties of neural network representations in reinforcement learning. Artificial Intelligence, 330:104100, 2024.
14

---

Published as a conference paper at ICLR 2025
Kaixin Wang, Daquan Zhou, Jiashi Feng, and Shie Mannor. PPG reloaded: An empirical study on what matters phasic policy gradient. In Andreas Krause, Emma Brunskill, Kyunghyun
Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), International Conference
on Machine Learning, ICMI. 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of PMLR, 2023.
Proceedings of Machine Learning Research, pp. wang23aw. html
URL https:
7/proceedings m.r. www...t.........................
Jiayi Weng, Min Lin, Shengyi Huang, Bo Liu, Denys Makoviichuk, Viktor Makoviychuk, Zichen Liu, Yufan Song, Ting Luo, Yukun Jiang, Zhongwen Xu, and Shuicheng Yan. En- ypool: A highly parallel reinforcement learning environment execution engine. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022. New Orleans, LA, USA, November 28 Decem- ber 9, 2022, 2022. URL http: htttt::/ppaprs..niggt nips cc cwppppppppppppreeepppppppppppe.e _ttppppee_eeeeee/ppaaapppppeepp files/paper/200 and
hash 4 (abbaaa694 Benchmarks html
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In 9th International Conference on Learning Repre- sentations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021a. URL https://cpenrevieww net
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Im- proving sample efficiency in model-free reinforcement learning from images. In Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI 2021, Thirty-Third Conference on Innova- tive Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium on Educa- tional Advances Artificial Intelligence, EAAI 2021, Virtual Event, February 2-9, 2021, pp. 10674-1066 AAAI Press, 2021b. URL https: / wwai.aaa.aaa.....a.....o.o wtwrg/iinee.../moo php/ AAAI/ article/view/1722
Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL. httpphtts////peereeeee eww. 2ett:/formmmddd=2FFooa
Wancong Zhang, Anthony GX-Chen, Vlad Sobal, Yann LeCun, and Nicolas Carion. Light-weight probing of unsupervised representations for reinforcement learning. In RLC, 2024.
A THEORETICAL RESULTS
Theorem 3.1.
The difference in returns achieved in train levels and under the full distribution, or generalisation error, has an upper bound that depends on 1(ZA; L.), with
2D I(Za;
Feo_l(())zo-PPtell (xo)) Peo (2O)) S
[T]
where c U(L) indicates e is sampled uniformly over levels in L, D is a constant such that [V*(x) D/2, V2, and Z is the output space of the actor's learned representation.
Proof. This result directly follows from a result obtained by Bertrán et al. 2020 and reproduced below.
Theorem A.1. For any CMDP such that [ D/2, Vz, n, with D being a constant, then for any set of training levels L,, and policy π
2D [T.]
Beo2e(())an~<Boee)1 (VTU)) - (
x [[T; L).
(10)
15

---

Published as a conference paper at ICLR 2025
Then, _ dA. by the data processing inequality always have I(T;L) - a( L), and therefore,
2DD2
Pende Plo
x xI(T;; L)
[T]
2D2
< |4]
x XI(ZA:L)
Garcin et al. (2024) follow the same reasoning and obtain equivalent result, without restating the bound.
Theorem 3.2. ifT (X) satisfies the Markov property, and we have ((X,X'); A) I((Z, A) and -(;;" for any X,X', A, Z, Z' collected using policy then Tg x A + (Z) satisfies the Markov property when following a. Tz always satisfies the Markov property if the above conditions hold for any n.
Proof. This proof has two part. We first demonstrate that the Inverse Model condition of Theo-
rem A.2 from [Allen al. (2021) (reproduced below) is satisfied if and only if I((Z, Z');
((X,x) A). We then show that if (Z; Z') = I(;; X') then the Density Ratio condition is also satisfied.
Theorem A.2. á is a Markov representation if the following conditions hold for every timestep and any policy n.
1. Inverse Model. The inverse dynamic model, defined as (a]ss's) = T( [(s(aa.))((ass)) P (= where
p= (s'(ss)) Eaea T(s s)T(a should be equal in the original and reduced MDPs. That is have (alz 2) = PT (als, s'), Ya E A. s,s' E S.
2. Density Ratio. The original and abstract next-state density ratios are equal when con- p (=11z) p=(5(())
ditioned on the same abstract state:
Vx' E S, where p (s((szz) =
(=(2)
(F(s))
ESEs PT (((()))) and p(s)z) Eves (a)a))
1ocess (*) (=(((() is the probability of transi-
tioning to state s' and p(((22) is the probability of currently being in state when in latent state 2.
We begin with two observations that are useful for our derivation.
Observation Given that any 2 E Z is obtained from the mapping x 2 o 9 2, and that h é á f2 is deterministic (but not necessarily invertible) function, each element E X maps to single element E It directly follows that Va, ž1, 22 x Z x Z, we have
p(a, 21, = Σ
22)
21.2
and
2(21,22)
Σ
=h(=)) h(x2))
21.
Observation B: Let Pd (0, x,, x!) be the joint distribution elements in (A, X, X') collected under policy n, we have pa (a, > 0 if and only if a, æ1, 22 XA,
Observation C: Similarly to obs. B, we have pa (21,, x2) > 0 if and only if (1, 02 E (X,X")
1) Proving that the Inverse Model condition is satisfied if and only if [((7, Z'); A)
16

---

Published as a conference paper at ICLR 2025
The above is equivalent to showing that the Inverse Model condition is satisfied if and only if HA Z') = H( For H(A Z"), we have
H HA = Σ
pd (a, log PM (a)2,
A,
(from obs. A)
Σ
Σ
(= x1, 21,
log (=(a2,,))
Ax &
21 2X22
(from obs. B)
Σ P'
h(x), h(=), ) log P (alz,2)
Σ
AxXxX
2
Σ AxXxX
Pd I'
Pd ("| 2
22
Σ
Pd (" x ) P" (a log
pd (a 2') 1 M2),01)
AxXxX
2
Ex,
p (a log
PT (a (f)
It follows that
Pd (alrr21 x')
H(A H(AX, = Ex.x'
Σ A
pd (alx, log
Il Pa (alz, (ee,'=
= Ex.x (PIQ)
with P = P= x') and Q = E P" (alz, 2/Ja//. From Gibbs inequality we always have DRr. (pl.(pl (pl1q) 2 0, therefore 1((2, Z"); = X'); A) if and only if DKI. (Pll()))Q)) = Vx, E X, X', which the case if and only if Q almost eeeeeryheere.
From observation A, any J1,I2 T2 E X² maps to exactly one pair 21, E z and by construction of X,X', Z' for any pair x, X,X' must have Q (* with 2, z' being the corresponding pair in
2 Px (alz, =f 11 Z,Z'
Therefore (((z, = X((x,X');; A) if and only if (=(a2,,)') = X,X' Z,Z', and we recover the Inverse Model condition.
2,
Conversely, if the Inverse Model condition is not satisfied, then ,', a xx, X,''',,'''', A for which + Q. Then DkI > at v. and [(ZZ, Z'); A) (X,XX) A).
2) Proving that the Density Ratio condition is satisfied if I(Z; 2') = I(X;X'.
We first show that satisfying
Pd (2' (x)
(=(2112)
V' d' E X, X', Z, Z' (x(x")
(II)
P
is sufficient for satisfying the Density Ratio condition (x
(T()))))
We then show that the
(p
(=))
condition in Equation 11 holds if and only if 1(2; 7')
I(X; X').
i) Showing the Density Ratio condition holds when Equation is satisfied. First we notice that,
Vx', 2 E X', Z, we have
Ex [x"
REX
Then, supposing Equation
IT) holds, we must have
pm (=11== (T(22)
nd Ex (x' [=]==]] PM (x6'
V2',2,2'' E x', xx, 2.7,
and the Density Ratio condition holds.
ii) Proving Equation 11 holds if and only if 1(7;7') I(X;X').
17

---

Published as a conference paper at ICLR 2025
We have
(=1122)
I(ZiZ') Σ
p (z, log
pm (z')
(2((2)
(from obs. A)
Σ
pa (21 h(x1), h(a2) log (=)
2
21,
pd (=(222))
(from obs. C)
Σ
Pd (0.
1 1 h(x), (xx2)(x)))))) log
(=(22)
22a
p (=(2112)
1|2,5' 10),aa
Ex,x'
log
(T()))
Then,
I(X;X') I(Z:Z') Ex.x (DKL(P'II')
with
(=(('() and -=Il (=(()))) 1k,e' (x
P'
P=(²))
The remainder this part follows the same structure as for the first part of the proof.
X() I(Z; Z') if and only if Vx, E X', Q almost p-everywhere. Any 21, 22 X² maps to exactly one pair 21, 22 E z², and by construction of X, Z', for any pair x, x' E X, X', we must have
(=(21
1=.='M+)))(''' P=(2))
PM (='l=) PT(22)
II E,5
with 2 being the corresponding pair in Z, Z'
Therefore X() 2") if and only if Vx, a 2,,2 X,X', Z' we have p*(a")
(='(a())
(T P())
Finally, from i) being true, the Density ratio condition must hold.
Lemma 4.1. (Z; L) > 0 Of () > O and E2k, Cy E x L such that p(zxle)) + u(zx))
Proof. Given π is fixed while the batch O collected, for a single batch the causal interaction between L, O and Z is described by the Markov chain X O Z, where (s,c) E $ x L and isn't directly observed. By the data processing inequality, I(L; Z) I I(L;O), and as such [(L; O) 0 is necessary condition for I(L; Z) be positive.
Note that
I(L HH(L + H(Z) - H(L, = = + -(L,2)=H(L) + H(2),
that if and only if Z and L are independently distributed. Given the causal relationship between
I and Z, p(zlc) is well defined Vz, c E 7 x L.. If 3e,, Cj E Z x I. such that p(= (c)) # m(=x) then and cannot be independently distributed, and [[T: Z) > 0.
Lemma 4.2. I(;L) monotonically increases with a) I(Z; V) I(ZVVIL) and b) Il: Z') ()).
Proof. Proof for condition a) : By the causal structure V + X + 0 + Z and the chain rule of mutual information, we have
= -12 +12
LV) is the information encoded in Z about the training levels that does not depend on V. ((Z;V) L) represents the information encoded in Z about state values that is level-specific. If this term increases then 1(Z; L.) will also increase.
18

---

Published as a conference paper at ICLR 2025
Sni
FW
Figure 6: In the top row, left, we depict the causal graph of states, observation and latents obtained over an episode. On the same row we draw a simplified graph that focuses on the relationship between and No.N. and utilises the notion that the context remains the same throughout the episode. In the bottom row draw the resulting causal relationship between L, Z and Z'.
Proof for condition b) Consider episode of arbitrary length N collected with policy We depict the causal structure that exists between elements in the top row of Figure (elements may be repeated within each sequence). It naturally follows that we have the causal structure depicted in the bottom row when considering all levels in L. By the chain rule mutual information, we have
I(2; L) = (ZZ;
I(Z; L' = I(Z L|Z') + I(Z 2') I(2; Z
I(Z; T.). and it follows that I(Z; L) increases with I(Z; Z') 7')
B ADDITIONAL FIGURES AND TABLES
I((Z,'))A) I(2,2)
PPO
PPG
DCPG
0.2
0.4
0.6
0.8
Figure 7: I((Z Z'): + [(Z; Z') for shared (gray), actor (blue) and critic (orange) for PPO, PPG, and DCPG in Proegen.
C IMPLEMENTATION DETAILS
C.I MUTUAL INFORMATION ESTIMATION
We measure mutual information using the estimator proposed by Kraskov et al. (2004) and later extended to pairings of continuous and discrete variables by Ross (2014). These methods are based
19

---

Published as a conference paper at ICLR 2025
((ZL)
(Z:V)
(2:22)
2((221)
RPO
#0--DCC
wo-Dr(C
PRG+DrTC
DCPG
M
DCPG+D(C)
DCPG-DDrt) Orts)
1.8
2.4
3.0
3.6
0.4
0.8
1.2
0.0
0.3
0.6
0.9
0.16
0.24
0.32
0.40
PPO
PPLMMCu((C
PPD+Dr(C)
PPG
C
RCG
PG-DDriC
DGPG D(C)
1000
DCPG
I(0:L)
[(O.V]
(0000))
1((O,0')A)
Figure 8: Mutual information measurements for the actor (blue) and critic (orange) for auxiliary losses applied to the critic for PPO, PPG, and DCPG in Proegen. Top/bottom rows are (Z; (I;)) with a shared x-axis.
KZ:L)
((:V)
1(:22)
K2Z,2AA))
PPC
#O-MICOlA?
PPG-DIA)
PPO+DIAI-ICOIC
PRO+AovIA)
PPO DriA)
PPO+0rIC)
PGGMICo
PIG MICaAA)
PPG-DICH
PPG+DIA)- MICaCO
PPG)AdVIA)
(PG+AdVAA
PPG+ DriA
PPG:DiCC
DCPG MiCoiC3
MICG-M
DCPG-DIC3
DCPG+DIAM- MiCotCH
MICoCC
DCPGIDr1C
1.8
2.4
3,0
3.6
0.4
0.8
1.2
0.0
0.3
0.6
0.9
0.00
0.15
0.30
Figure 9: .(Z:.) measurements for the actor (blue) and critic (orange) for auxiliary losses for PPO, PPG, and DCPG in Procgen.
on performing entropy estimation using k-nearest neighbors distances. We use and determine nearest neighbors by measuring the Euclidian (L2) distance between points. We checked measure- ments obtained when using different and under different metric spaces, and we found that our measurements broadly invariant the choice of estimator parameters.
At the end of training we collect a batch of trajectories consisting of 26 timesteps (2 timesteps in Brax) from L. We construct (A, O'. Z. Z', L) from (a,, or, O.+1, 2t, 2441, D,, c.). Subsampling is necessary to compute mutual information estimates in reasonable time, while ensuring we sample states from most levels in I. and at various point of
4096 timesteps yielding
20

---

Published as a conference paper at ICLR 2025
ICOO
5POtsh
PPG
(PO+RICa(A)
PC+
PPG-A
PRC+DrC]
PPG
C
PPG DRGCDAA
PPG-AdK(A
DCPG
DCPS
DCPG-MICaCCI
DCPG+MICNAA
DCC
DCPG
3.42
3.48 3.54
3.60
0.25
0.50
0.75
0.2
0.4
0.6
0.8
0.05
0.10
0.15
0.20
Figure 10: I(O;,) measurements for the actor (blue) and critic (orange) for auxiliary losses for PPO, PPG, and DCPG in Procgen.
(Z;L)
(2,22)
I(Z.Z'AA)
PPO
MRCAA
PCO+DCC
MIC
PPG
PRG1RICA/A
PPG+
DCPG
MICCIM
DEPGRMICHA!
DCAG+DIA
DCCG+DICI
DCPG MICCCI
0.5
1.0
1.5
2.0
0.00
0.08
0.16
0.24
0.0
0.3
0.6
0.9
0.0
0.4
0.8
1.2
Figure 11: Mutual information measurements for the actor (blue) and critic (orange) for auxiliary losses for PPO, PPG, and DCPG in Brax.
the trajectories followed by the agent in each level. Timesteps sampled uniformly and without replacement from the batch, after having excluded:
Odd timesteps, to ensure and O' will not overlap (i.c. O contains only even timesteps, and O', being sampled at / contains only odd timesteps).
21

---

Published as a conference paper at ICLR 2025
PPO[sh] Normalized Returns
KZA; and ((Za; V)
PPG (TesU
PPG (Train)
2.0
PPO (Test)
2.5
(Train)
1.9
2.0
1.8
MZ VI
1.5
(Zz:L)
PPO 12
1.6
1.0
1.5
0.5
1.4
32x
16x
8x
1x
32x
16x
8x
1x
Auxiliary batch size relative to PPO batch size
Auxilarr batch size relative to PPO batch size
Figure 12: Procgen PPG returns (left) normalized by PPO[sh] performance and mutual information quantities [(ZA; (right) for varying auxiliary batch size levels.
Returns
PPO
PPO+MICO(C)
PPO+MICo(C)
PPO+MICO(A)
PPO+MICoLA)
PPO+D(A)
PPO
PRO
(PO+D(A-MICo(C)
PPO+D(C)
AddVA)
MMOCCC
MPO-D(A)-MICO(C)
PPO+
(PO+
PRO+Dr(C)
PPG[sh]
PPO+Dr(C)
PPG+MICo(A)
MIG-MICoIC)
PPG+D(CI)
PPG-MICoIA)
MICACC))
PPG+D(A)
PPG+AdWAI
MMCCCCC
PRG+D[C)
PPG+DNA)
PFG-D(AI-MICoIC)
PPG DDCCC
DCPG[shi
DCPG
DCPG
DCPG+MICo(C)
DCFG MICOC)
MiCaAA)
DCPG+MICoIA)
DCPG+D(A)
DDCC+
DCPG+DIA)
DCFG+DIA)+MICOCI)
DCCG
DCPG-Dr(A)
DCPG-Dr(C)
DCPG DIA) MICo(C)
0.8
1.6
2.4
3.2
Figure 13: Returns in Procgen (left) and Brax (right).
2. Timesteps corresponding episode terminations, to ensure the pair or, o, cannot orig- inate from different levels.
3. Timesteps from episodes that have terminated, ensure always compute tt.
22

---

Published as a conference paper at ICLR 2025
Table 2: Measurements of compression efficiency C(ZA V) (Equation ( with standard error in Procgen. Statistical significance bolded, determined by Welch's t-test. Results highlighted in red when decoupling decreases A O; V), and highlighted in green when decoupling increases CO V), otherwise yellow. Coupled architectures are denoted with algorithm name plus "[sh]".
Algorithm
C(ZA ((
PPO[sh]
89.3 22
65.2 =
PPO
90.1 4
52.3 =
PPG[sh]
8
70.0 =
PPG
94.1
75.5 =
DCPG[sh]
95.4
77.6 =
DCPG
92.3 7
76.4 2
Table 3: Measurements of compression efficiency C(ZA|O;:) (Equation ( of the actor's repre- sentation Φ Proegen. Results highlighted red when the auxiliary loss decreases the metric relative to the base algorithm, and highlighted in green when the auxiliary loss increases the metric relative to the base algorithm. Auxiliary losses are applied to the actor (A) and critic (C) the form of dynamics prediction (D), MICo, and advantage distillation (Adv).
Algorithm PPO PPO+MICO(C)
L(
C((Za|O,ZZ]0))
90.1 +
52.3
9
93.9 2
60.4 +
99.4 0
PPO+MICo(A)
98.6 = 1
84.7 2
87.5 5
PPO+D(A)
62.6 + 12
61.3 +
100.0 + 0
PPO+D(C)
86.8 7 7
52.8 3
100.0 + 0
PPO-D(A)+MICO(C)
76.3 +: 6
63.7 3
99.5 + 0
PPO+Adv(A)
96.9 + 2
53.2 3
100.0 + 0
(PO+AUN(A)+MICOC))
100.0 + 0
57.0 2
98.5 +
PPO+Dr(A)
89.1 + 6
54.3 +
100.0 + 0
PPO+Dr(C)
98.1 + 1
50.8 3
99.6 + 0
PPG
94.1
75.5 +
100.0
PPG+MICo(C)
98.0 +
76.4
100.0 + 0
PPG+MICo(A)
95.4 2
85.8 +
100.0 + 0
PPG+D(A)
89.5 4
71.2
100.0 + 0
PPG+D(C)
96.3 + 2
75.1 + 2
100.0 + 0
PPG+D(A))MICO(C)
85.9 + 7
71.6 +
100.0 + 0
PPG+Adv(A)
98.4 +
63.3 +
100.0 + 0
PPG+AdY(A)+MICOCC)
99.4 +
62.9 +2
100.0 + 0
PPG+Dr(A)
91.3 + 3
74.9 2
100.0 + 0
PPG+Dr(C)
93.1 + 6
75.6
100.0 + 0
DCPG
92.3 7
76.42
100.0 +
DCPG+MICO(C)
98.1 +
76.6 2
100.0 + 0
DCPG+MICO(A)
91.7 3
74.3 2
100.0 + 0
DCPG+D(A)
80.9 4
69.9 + 2
100.0 + 0
DCPG+D(C)
97.8 +
76.3 +
100.0 + 0
DCPG+D(A)+MICO(C)
83.4 5
69.6 2
100.0 + 0
DCPG+Dr(A)
97.5 +
76.7 +
100.0 + 0
C.2 PROCGEN
The Procgen Benchmark is a set of 16 diverse PCG environments that echoes the gameplay variety seen in the ALE benchmark Bellemare et al. (2015). The game levels, determined by a random seed,
23

---

Published as a conference paper at ICLR 2025
Table Measurements of compression efficiency C(ZolO;-) (Equation (9) of the actor's represen- tation de in Procgen. Results highlighted in red when the auxiliary loss decreases the metric relative the base algorithm, highlighted in green when the auxiliary loss increases the metric relative to the base algorithm, and highlighted in yellow otherwise. Auxiliary losses are applied to the actor (A) and critic (C) in the form of dynamics prediction (D), MICo, and advantage distillation (Adv).
Algorithm
C(ZclO:V)
L)
C((Zc O, ((
PPO PPO+MICO(C)
93.7 #3
88.4 2
85.6 + 4
100.0 + 0
90.3 +
82.7 + 3
PPO+MICO(A)
97.6 +
92.4 +
64.4 + 6
PPO+D(A)
99.7 =
90.2 =
87.4 + 3
PPO+D(C)
87.6
77.4 + 2
99.1 + 0
PPO+D(A)MMIOO(C)
91.0 6
88.1
84.4
PPO+Adv(A)
96.7 + 2
89.3 +
87.6 + 4
PPO+Ad((A++MCOOCC)
100.0 +
89.9
87.0 + 3
PPO+Dr(A)
98.0 +
90.0
87.9 =
PPO+Dr(C)
97.5 +
87.0 =
86.3 +
PPG
99.2 + 1
81.6 2
90.3 +
PPG+MICo(C)
93.0 + 6
90.9 + 2
91.8 + 2
PPG+MICo(A)
100.0
80.9
84.4
PPG+D(A)
100.0 = 0
79.2 = 2
91.3 + 3
PPG+D(C)
89.4 + 4
77.6 2
100.0 + 0
PPG+D(A))MICO(C)
93.3 +
89.1
87.1 4
PPG+Adv(A)
100.0 + 0
80.9 + 2
89.7 +
PPG+Ad((A+MMCOOC)
99.9 +
92.3
93.1 + 3
PPG+Dr(A)
98.7 +
81.7 2
90.2 + 3
PPG+Dr(C)
96.4 3
81.8 2
85.8 4
DCPG
99.5 =
81.7
92.1 =
DCPG+MICO(C)
98.9 +
88.6
93.5 2
DCPG+MICO(A)
99.8 + 0
81.4
87.2 + 4
DCPG+D(A)
100.0 = 0
80.7 2
92.6 +
DCPG+D(CC
91.3 + 3
81.2 +
100.0 + 0
DCPG+D(A)+MIC((O)
99.6 + 0
87.4 2
92.7 + 3
DCPG+Dr(A)
100.0 0
81.3 2
89.4 3
DCPG+Dr(C)
97.0 2
81.2
90.7 + 3
can differ in visual design, navigational structure, and the starting locations of entities. All Procgen environments use common discrete !5-dimensional action space and generate 64 x 64 x RGB ob- servations. A detailed description of each of the 16 environments is provided by Cobbe et al. (2020). RL algorithms such as PPO reveal significant differences between test and training performance in all games, making Procgen valuable tool for evaluating generalisation performance.
We conduct our experiment on the easy setting of Procgen, which employs 200 training levels and budget of 25M training steps, and evaluate the agent's 's scores the training levels and the full range of levels, excluding the training levels. We use the version of Procgen provided by Env Pool
Weng et al. 2022). Following prior work, (Raileanu al. 2021 Jiang et al. 2021 Moon al. 2022). for each game we normalise train/ test scores by the mean train test score achieved by PPO that game.
For PPO, we base our implementation on the CleanRL PPO implementation (Huang et al. 2022). which reimplements the PPO agent from the original Procgen publication in JAX. We use the same
24

---

Published as a conference paper at ICLR 2025
ResNet policy architecture and PPO hyperparameters (identical for all games) as Cobbe et al. (2020)
and reported Table 5
We remmplement PPG and DCPG in JAX, based the Pytorch implementations provided by Huang ct al. 22022) and Moon et al. 2202). We use the default recommended hyperparameters for each algorithms, which are reported in Table 6 We note that our PPG implementation ends up outper- forming the original implementation by about 10% on the test set, while our DCPG implementation underperforms scores reported by Moon et al. (2022 by about 10%. We attribute this discrep- ancy to minor differences between the JAX and Pytorch libraries, and decided to not investigate further.
We conduct our experiments on A100 and RTX8000 Nvidia GPUs and 6 CPU cores. One seed for one game completes in 2 12 hours, depending on the GPU, algorithm, and whether the architecture is coupled or decoupled (for example, PPG decoupled can be expected to run 4x to 6x slower than PPO coupled).
C.3 BRAX
For experiments in Brax, implement custom "video distractors" set of tasks, similar to those from (Stone et al. 2021). In this setup, a video plays in an overlay on the pixels the agent views. There is disjoint set of videos between the training and testing environments. The random seed determines the environment's ' initial physics and the video overlay at the beginning of training. The pixels themselves full-RGB 64 64 arrays, but agent input to 64 x 64 x 9 pixels.
framestacking bring each
Similar the algorithms used in the Procgen experiments, we implement our algorithms in JAX and base them on ClearnRL
We conduct our experiments on RTX A4500 Nvidia GPUs and CPU cores. One seed completes in 7.5-48 hours, depending on the environment and its physics backend as well as the algorithm.
25

---

Published as a conference paper at ICLR 2025
Table 5: Hyperparameters used for PPO in Procgen and Brax experiments. All runs employing
specific (or combination of) representation learning objective use the same hyperparameters.
Parameter
Procgen
Brax
PPO
?
0.999
0.999
AGAE
0.95
0.95
rollout length
256
128
minibatches per epoch
8
8
minibatch size
2048
512
Jr clip range
0.2
0.2
number of environments Adam learning rate
32
5e-4
5e-4
Adam
le-5
1e-8
max gradient norm
0.5
0.5
value clipping
no
no
return normalisation
no
value loss coefficient
0.5
0.5
entropy coefficient
0.01
0.01
PPO (coupled)
PPO epochs (actor and critic)
PPO (decoupled)
Actor epochs
Critic epochs
MICo objective
MICo cnefficient
Target network update coefficient
0.005
0.05
Dynamics objective
Dynamics loss coefficient
1.0
0.01
In-distribution transitions weighting
1.0
1.0
Out-of disttribbution states weighting
1.0
1.0
Out-of-disse -itto--distrbbution actions weighting
0.5
0.5
Adventage distillation objective Advantage prediction coefficient
0.25
26

---

Published as a conference paper at ICLR 2025
Table 6: Hyperparameters used for PPG and DCPG in Proegen experiments. Hyperparameters shared between methods are only reported if they change from the method above. All runs employ- ing specific (or combination of) representation learning objective use the same hyperparameters.
Parameter
Procgen
PPG
0.999
AGAE
0.95
rollout length
256
minibatches per epoch policy phase
minibatch size policy phase
2048
minibatches per epoch auxiliary phase
512
minibatch size auxiliary phase Jn clip range
1024
0.2
number of environments
64
Adam learning rate
5e-4
Adam
le-5
gradient
0.5
value clipping
return normalisation
yes
value loss coefficient policy phase
0.5
value loss coefficient auxiliary phase
1.0
entropy coefficient
0.01
policy phase epochs
auxiliary phase epochs
6
number of policy phases per auxiliary phase
32
policy regularisation coefficient
1.0
auxiliary value distillation coefficient
1.0
DCPG
value loss coefficient policy phase delayed value loss coefficient policy phase
1.0
MICo objective
MICo coefficient
0.5
Target network update coefficient
0.005
Dynamics objective
Dynamics loss coefficient
1.0
In-distribation transitions weighting
1.0
Out-of-distribun states weighting
1.0
Oui-of-distribution actions weighting
0.5
27


Reviews:
[{"review_id": "7bnno9nsp7", "reviewer": "Reviewer_Cooi", "summary": "The paper investigates representation learning objectives for actor and critic networks in deep reinforcement learning, specifically in the context of policy gradient algorithms. Using information-theoretic metrics, the study reveals that the actor and critic networks naturally specialize in capturing distinct information types, with the actor focusing on action-relevant features and the critic on value and environmental dynamics. This specialization is shown to enhance model architecture design and improve representation learning objectives, ultimately contributing to greater sample efficiency in reinforcement learning. In addition, this paper also how various representation learning objectives could help shaping the actor and critic's representation and improve the RL sample efficiency.", "questions": "- What does the delta percentage in Figure 1 represent? It would be helpful if the authors added a brief explanation to the caption.\n- 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.\n- 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.\n- 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?\n- 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.\n- 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?\n- 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?\n- 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "-  **Novel Perspective on Representation Specialization**: This paper provides a fresh and thorough analysis of actor and critic representation specialization using information-theoretic metrics. \n- **Convincing Empirical Evidence**: The authors provide substantial empirical support for their claims, with extensive experiments that effectively demonstrate the impact of decoupling actor and critic representations, as well as impacts of different representation learning objectives on the resulting representation.\n- **Insightful Conclusions**: The study draws interesting and impactful conclusions regarding how specialized representation learning objectives can improve sample efficiency and generalization in RL, a great contribution to the RL community, shedding light on the design of future deep RL algorithms.", "weaknesses": "- **Limited Scope of Algorithm Selection**: The paper exclusively focuses on online, on-policy, policy gradient algorithms (e.g., PPO, DCPG), leaving it unclear how well the conclusions apply to other RL frameworks, particularly off-policy algorithms like DDPG or SAC. Broader consideration of diverse algorithm types would strengthen the generalizability of the findings. Also, whether the same conclusion holds under different model sizes? It would be useful to add in the paper to show that this conclusion would still hold if we simply added capacity to the actor/critic network without changing any learning objectives.\n- **Presentation and Clarity**: The paper’s structure and presentation could be significantly improved to make it more easier to understand for readers. The writing is currently dense, making it challenging to grasp the main takeaways and hypotheses. For instance, restructuring sections with clearer headings, especially in Sections 6.1 and 6.2, would help readers navigate the study’s conclusions. Using bullet points to highlight key hypotheses and findings, as well as starting each paragraph with a bolded summary sentence, could enhance readability and ensure that readers capture the core messages more effectively.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "N1q8zrNxyU", "reviewer": "Reviewer_fkit", "summary": "This paper focus on the difference in representation graph between actor and critic in AC methods in RL. The authors aims to more formally validate an empirical finding where decoupling actor and critics may help with representation learning through", "questions": "- 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.\n- 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?\n- 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.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The way in which the topic is approached in the section is well considered, I find it readable\n- I appreciate the clear problem formulation and motivation of the paper\n- The results, while not counter-intuitive nor surprising, gives some form of closure to the decoupling findings in earlier works, and brings this issue to attention does help clarify the topic\n- Assumption seems reasonable", "weaknesses": "- Potentially need more ground work setup, such as in the problem definition where an observation function is present, which normally implies that the MI between state and observation is not maximized, but it seems the rest of the paper operate on the assumption that it is indeed maximized and so that the environments are fully observable.\n- Evaluation is lacking in breath, it would be great to include a general set of tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DQdmqc2Zxp", "reviewer": "Reviewer_qthx", "summary": "This paper analyzes how different choices of architecture and representation learning objectives affect overfitting on, transfer to, and specializing in different levels of games (ProcGen). The analysis is largely established on the mutual information between different quantities. The author concludes that actor and critic representation specialize in different information, that critic can influence exploration significantly, and that representation learning objectives (i.e., auxiliary loss) should be consistent with the specialization of actor and critic representation.", "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?\n- 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?\n- Line 165: Why does a high I(Z_C, V) help optimizing the value loss? \n- 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)?\n- 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?\n- Line 194: What's the role of linear probing here?\n- Equation 6: What are $z^*_{A_0}$ and  $z^*_{A_1}$? Why I((Z^*_A, Z^*'_A); A) is guaranteed to be maximized when following the optimal policy?\n- Line 256: I don't see why zero-shot transfer is evident in your example.\n- Figure 3: How should one interpret these values? The higher the better? The lower the better? Better summarize the takeaway in the caption.\n- Line 374: Why does maximizing J^\\pi bias \\pi to collect trajectories with high I(O; V)?\n- Line 431: Why does actor objective promote learning invariant quantity?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "The analysis to representation learning based on mutual information is interesting. The conclusion drawn from this analysis is helpful for deep RL researchers on designing representation learning objectives.", "weaknesses": "Writing is very unclear. See my questions for details. Overall, there are overwhelming notations, making the presentation unnecessarily complex. The reason of defining each notation is also not clear to me. Takeaway messages in each section is also not clear. For example, in section 4.1, the author starts with lots of new definitions without explaining why these are necessary. I would suggest the author explain why we need those definitions and then say the high-level messages before diving into details.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Gl7ghEfRSC", "reviewer": "Reviewer_5hRY", "summary": "This paper studies whether actors and critics benefit from a decoupled representation. The authors find that when decoupled, actors and critics learn representations that specialize in different types of information from the environment. Additionally, there is an interplay between these representations even when they are decoupled. The authors show how this insight helps select representation learning objectives for better performance.", "questions": "- 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? I am curious about factors other than coupled/decoupled learning and how much impact they can have on these metrics.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "**originality**\n- although the fact that decoupled actor and critic learn different representations is not really surprising, the other findings in the paper are very interesting and novel. \n-  the four mutual information metrics are interesting and helps to understand the effect of decoupled learning. \n- the paper gives some very novel and interesting analysis and results\n\n**quality**\n- it is good that authors plan to release all data and code for better reproducibility. \n- overall great quality\n\n**clarity**\n- writing is clear and easy to follow. \n\n**significance**\n- empirical and theoretical results are good contributions \n- the analysis and insights from the paper are especially interesting", "weaknesses": "I don't have major concerns, minor ones:\n- I think it can be nice to explicitly highlight, either in the beginning or conclusion, what are the best representation learning objectives and most important specialisations with quantitative measures. \n- The paper can be even stronger if similar analysis can be done for Q-learning based methods such as DDPG/SAC/TD3 type algorithms. But I understand that might be too much extra work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CtcgcIE8Dd", "reviewer": "Reviewer_r4WB", "summary": "Authors systematically investigate the impact of coupling and decoupling of the representations for actors and critics in on-policy actor-critic framework.", "questions": "My questions are based on the weaknesses\n\n1. Could authors please elaborate on weaknesses 1 and 2?\n2. According to weakness 3, I would recommend authors to add some additional information to the tables and plots about whether the metric is expected to be minimized or maximized.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Theoretically-grounded study with the demonstration that actor and critic should be decoupled as their objectives contradict each other.\n2. Good empirical study for confirming claims using various algorithms. When result contradict assumptions, authors test and confirm different hypothesis which could explain the observations. These assumptions reveal novel things for RL community, such as the potentially underestimated role of critic in exploration.\n3. Additional study on the effects of different representation learning approaches and when they benefit performance.", "weaknesses": "1. Study is only conducted using the on-policy algorithms. It would be interesting to also investigate off-policy methods\n2. If I understand correctly, study is performed only on discrete state and action spaces environments. Would the conclusions change if we switch to the continuous environments? It is not clear for me whether the proposed analysis can be applied to such tasks as the computation of mutual information becomes intractable.\n3. For me, it was easy to get lost in many different metrics and keeping in mind what we want to maximize and what we want to minimize for the optimal solutions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)