No Representation, No Trust: Connecting Representation, Collapse, and Trust Issues in PPO
Authors: Skander Moalla
Andrea Miele
Daniil Pyatko
Razvan Pascanu
Caglar Gulcehre
Venue: NeurIPS 2024 poster (2024)
Keywords: proximal policy optimization
plasticity loss
trust region
feature rank collapse
regularization
Primary area: reinforcement_learning
TL;DR: We show that PPO suffers from feature rank degradation, that it harms its heuristic trust region, and that it is connected to performance collapse.

Abstract: Reinforcement learning (RL) is inherently rife with non-stationarity since the states and rewards the agent observes during training depend on its changing policy.
Therefore, networks in deep RL must be capable of adapting to new observations and fitting new targets.
However, previous works have observed that networks trained under non-stationarity exhibit an inability to continue learning, termed loss of plasticity, and eventually a collapse in performance.
For off-policy deep value-based RL methods, this phenomenon has been correlated with a decrease in representation rank and the ability to fit random targets, termed capacity loss.
Although this correlation has generally been attributed to neural network learning under non-stationarity, the connection to representation dynamics has not been carefully studied in on-policy policy optimization methods.
In this work, we empirically study representation dynamics in Proximal Policy Optimization (PPO) on the Atari and MuJoCo environments, revealing that PPO agents are also affected by feature rank deterioration and capacity loss.
We show that this is aggravated by stronger non-stationarity, ultimately driving the actor's performance to collapse, regardless of the performance of the critic.
We ask why the trust region, specific to methods like PPO, cannot alleviate or prevent the collapse and find a connection between representation collapse and the degradation of the trust region, one exacerbating the other.
Finally, we present Proximal Feature Optimization (PFO), a novel auxiliary loss that, along with other interventions, shows that regularizing the representation dynamics mitigates the performance collapse of PPO agents.
Code and run histories are available at https://github.com/CLAIRE-Labo/no-representation-no-trust.

No Representation, No Trust: Connecting Representation, Collapse, and Trust Issues in PPO
Skander Moalla 1* Andrea Miele'
Daniil Pyatko
Razvan Pascanu
Caglar Gulcehre'
CLAIRE, EPFL "Google DeepMind
Abstract
Reinforcement learning (RL) is inherently rife with non-stationarity since the states and rewards the agent observes during training depend on its changing policy. Therefore, networks in deep RL must be capable of adapting to new observations and fitting new targets. However, previous works have observed that networks trained under non-stationarity exhibit inability to continue learning, termed loss plasticity, and eventually a collapse in performance. For off-policy deep value-based RI. methods, this phenomenon has been correlated with decrease in representation rank and the ability to fit random targets, termed capacity loss. Although this correlation has generally been attributed neural network learn- ing under non-stationarity, the connection to representation dynamics has not been carefully studied in on-policy policy optimization methods. In this work, we empirically study representation dynamics in Proximal Policy Optimization (PPO) the Atari and MuloCo environments, revealing that PPO agents are also affected by feature rank deterioration and capacity loss. We show that this is aggravated by stronger non-stationarity. ultimately driving the actor's perfor- mance to collapse, regardless of the performance of the critic. We ask why the trust region, specific to methods like PPO, cannot alleviate or prevent the col- lapse and find a connection between representation collapse and the degradation of the trust region, one exacerbating the other. Finally, we present Proximal Feature Optimization (PFO), novel auxiliary loss that, along with other in- terventions, shows that regularizing the representation dynamics mitigates the performance collapse of PPO agents. Code and run histories are available at https://githur httpg//thub.comm/AAARee aat.AaR-iitboommrrprreeeeattttto
1 Introduction
Reinforcement learning (RL) agents are inherently subject to non-stationarity as the states and rewards they observe change during learning. Therefore, neural networks in deep RL. must be capable of adapting new inputs and fitting new targets. However, previous works have observed that networks trained under non-stationarity exhibit an inability to continue learning, termed loss of plasticity, and collapse in performance (Dohare et al., 2021; Abbas et al., 2023; Kumar et al., 2023; Dohare et al., 2023a,b). Kumar et al. (2021); Lyle et al. (2022) connect this phenomenon to representation dynamics and show that value networks in off-policy value-based RL algorithms exhibit decrease in the rank their representations, termed feature rank collapse, and a decrease in their ability regress arbitrary targets, called capacity loss. Although this deterioration in representation is more generally attributed to neural networks trained under non-stationarity (Lyle et al., 2023), the connection to representation dynamics has been overlooked in on-policy policy optimization methods. In particular, Proximal Policy Optimization (PPO) (Schulman et al., 2017), one of the most popular policy optimization methods, makes several minibatch updates over non-statee -tasstationay data, unlike vanilla policy gradient methods, and optimizes surrogate loss that depends on a moving old
Cooresponneecce Laanderrmaala@eefl.c.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

policy. This raises the question of how much PPO agents are impacted by the same representation degradation attributed to non-stationarity. Dohare et al. (2021, 2023a,b) have shown that PPO agents lose plasticity throughout training but have only diagnosed it as collapse in performance or as an Adam optimization issue. Igl et al. (2021) have shown that non-stationarity affects the generalization of PPO agents (learning speed when training episodes are very different otherwise performance at test time on novel episodes) but not necessarily training, and no connection was made with the feature rank and capacity measures used in the recent value-based works. One crucial outstanding question is why the trust region embedded in methods like PPO is unable to prevent the deterioration policy by constraining its update. To address these gaps, we present the following contributions:
1. We provide the first study of feature rank and capacity loss in on-policy policy optimization, revealing that PPO agents in the Arcade Learning Environment (Bellemare et al., 2013) and MuJoCo (Todorov et al., 2012) environments are subject to representation collapse.
2. We draw connections between representation collapse, performance collapse, and trust region issues in PPO, showing that PPO's clipping becomes ineffective under poor representations and fails to prevent performance collapse, which is irrecoverable due to loss of capacity. We further isolate the breakdown of the trust region a theoretical setting.
3. We corroborate these connections by performing interventions that regularize non- stationarity and representations and result in a better trust region and mitigation of per- formance collapse, incidentally giving insights sharing an actor-critic trunk.
4. We propose Proximal Feature Optimization (PFO), a new regularization on the representation of the policy that regularizes the change in pre-activations. PFO strengthens our analysis by addressing the representation issues and mitigating performance collapse.
5. We open source our code providing a comprehensive and reproducible codebase for studying representation dynamics in policy optimization and a large database of run histories with extensive logging for further investigation on this topic.
2 Background
Reinforcement Learning (Sutton & Barto, 2018) We formalize our RL. setting with the finite- horizon undiscounted Markov decision process, describing the interaction between agent and an environment with finite 2 sets of states S and actions A. and a reward function r x A x - R. An initial state So E 8 is sampled from the environment, then at each time step & E {0,... max 1}. the agent observes the state S. E 8, picks an action A₁ E A according to its policy π : S + (AA) with probability π (A,, (S,), observes the next state S4+1 S sampled from the environment and receives a reward Re+1 = r(,,, Ac, S(+1). We denote by G. = Luma" Rk+1 the return after the action at time step t. The goal of the agent is to maximize its expected return J(x) = E [,] over the induced random trajectories. We discuss the choice of this setting in Appendix A.1.
Ei=o mmmxx1 Re+1
Actor-Critic Agent We consider on-policy deep actor-critic agents which train a policy (or actor) network (-: O) also denoted no. and a value (or critic) network 0(: w) that approximates the return of Tg at every state. At every training stage, the agent collects a batch of samples, called rollout, with its current policy тe, and both networks are trained with gradient descent on this data. The critic is trained to minimize the Euclidean distance to an estimator of the returns (c.g., G.). We use X-returns computed with the Generalized Advantage Estimator (GAB) (Schulman et al., 2015b). The actor is trained with the Proximal Policy Optimization (PPO) (Schulman et al., 2017).
Proximal Policy Optimization PPO-Clip, the most popular variant of PPO algorithms (Schulman et al., 2017). optimizes the actor by repeatedly maximizing the objective in Equation at each rollout.
Timax -1
(A((S)) V., clip (A,aa(A,)S))
To (A, St) Total (A 1 + 1 1e. v. )
I.ILI Maad
(0) ETT
Σ
min
(1)
t=0
'The pixel-based environment with discrete actions used in our experiments and our simple theoretical example in Section 3.2.1 fit the finite state and action formalism but not our continuous action space environment. We refer the reader Szepesvári (2022) for formalism of RL in that setting.
'The time step t is included in the representation of S, to preserve the Markov property in finie-hhoizzn tasks as done by Pardo et al. (2018) and is analogous to considering dimendenendent policies in the classical formulation of Miite-horizon MDPs
2

---

The objective is defined for some small hyperparameter Told is the last of the previous optimiza- tion phase, used to collect a training batch (rollout) after each optimization phase; V, is an estimator of the advantage of Told (e.g., V G. 6(5; Wola)); we use the GAE in our experiments. An optimization phase consists of maximizing the objective with minibatch gradient steps over multiple epochs on the training batch. We refer to PPO-Clip as PPO and provide a pseudocode in Algorithm 1.
Vr defined by Kakade & Langford (2002), which participates in lower bound to the improvement of He given that it is close to Told (Schulman et al., 2015a, see Theorem 1). In this regard, a gradient step on (CL)((0) (0) would increase (resp. decrease) the probability of actions at states yiclding positive (resp. negative) advantage until the ratio between the policies for those actions reaches 1 + (resp. - e) at which point the gradient at those samples becomes null. This is a heuristic to ensure a trust region that keeps policies close to each other, resulting in policy improvement.
Intuitively, PPO aims to maximize the policy advantage ETTAS El=
Emmxe-1
Te (A,A. Trss (A.
Non-stationarity in deep RL and PPO The actor and the critic networks are both subject to non- stationarity in deep RL. As the agent improves, it visits different states, shifting the distribution of states which makes the networks' input distribution non-stationary. This also holds for the targets to fit the critic, which change as the returns of the policy change. Unlike vanilla policy gradient (Sutton et al., 1999), and A2C (Mnih al., 2016), PPO's objective is optimized by performing multiple epochs of minibatch gradient descent on the current collected batch, potentially making the networks more likely to be impacted by previous training rollouts. In this sense, increasing the number of epochs in PPO can cause the agent to "overfit" more to previous experience.
Feature rank As done in most works studying feature dynamics in deep RL (Lyle et al., 2022: Kumar et al., 2021), we refer to the activations of the last hidden layer of a network (the penultimate layer) as the features or representation learned by the network. On a batch of N samples, this gives a matrix of dimensions N x D denoted by L, where D < N is the width of the penultimate layer. Several measures of the rank of this matrix have been used to quantify the "quality" of the representation (Kumar et al., 2021; Gulcehre et al., 2022; Lyle et al., 2022; Andriushchenko et al., 2023). Their absolute values differ significantly, but their dynamics are often correlated. We track all of the different rank metrics in our experiments, compare them in Appendix E, and use the approximate rank in our main figures for its connection to principal component analysis (PCA). Given a threshold 8 E R and the singular values (op of & in decreasing order, the
Σ - > 11 which corresponds to the smallest dimension EE 0
approximate rank of & is ming
of the subspace recovering (1 5)% of the variance of Φ. We use 8 = 0.01 i.e. the reconstruction recovers 99% of the variance as done by Andriushchenko et al. (2023); Yang al. (2020). We refer to this metric as feature rank with reference to the rank of the feature matrix when there is no ambiguity.
Capacity loss Target-fitting capacity (Lyle et al., 2022) is computed on checkpoints of a network undergoing some training to measure the evolution of its ability to fit some chosen target independent from its training. It is concrete metric to evaluate plasticity. Given fixed target (distribution over inputs and outputs) and a fixed optimization budget, a checkpoint's capacity loss is the loss from fitting the checkpoint to the target at the end of the optimization budget. Usually, the capacity of deep RL agent is measured by its ability to fit the outputs of a model initialized randomly from the same distribution as the agent a fixed rollout collected by this target random model (Lyle et al., 2022: Nikishin et al., 2023). We follow this practice. The data would in expectation be from the same distribution as the agent's initial checkpoint. To fit the critic, use an 2 loss on the outputs of the models. To fit the actor, we use a KI. divergence between the target and the checkpoint (forward KI.).
3 Deteriorating representations, collapse, and loss of trust
It is well-known that non-stationarity in deep RL can be factor causing issues in representation learning. However, most of the observations have been made in value-based methods showing that value networks are prone to rank collapse, harming their expressivity, and in turn, the performance of the agent (Lyle et al.. 2022; Kumar et al., 2022). PPO's generalization Igl et al. (2021) and performance in the long run or in a continual learning setting (Dohare et al., 2021, 2023a,b), but no evidence of representation deterioration was shown. Our motivation is to reuse the tools that showed that value-based methods are prone to representation collapse but in policy optimization methods for the first time. We focus on PPO for its popularity and its con-staiionarity which is impacted and can be controlled by multi-epoch optimization.
stationart has been shown to impact
3

---

Furthermore, crucial question for PPO, compared to most value-based alternatives, is how the regularization implicit in PPO through its trust region interacts with representation and performance collapse. Intuitively it should prevent rapid degradation of the policy.
We begin our experiments by training PPO agents the Arcade Learning Environment (ALeyBellemaree et al., 2013) for pixel-based observations with discrete actions and on MuJoCo (Todorov et al., 2012) for continuous observations with continuous actions. To keep our experiments tractable, we choose the Atari-5 subset recommended by Aitchison et al. (2023) and add Gravitar to include at least one sparse-reward exploortooottoon game from the taxonomy presented by Bellemare et al. (2016). For MuJoCo, we train on Ants, Half- Cheetahss Humanoids, and Hoppers, which have varying complexity and observation and output sizes. We use the same model architectures and hyperparameters as popular implementations of PPO on ALE and MuJoCo (Raffin et al., 2021; Huang et al., 2022b); these are also the architectures and hyperparameters used by Schulman et al. (2017) in the original implementation of PPO; they do not include normalization layers. For MuloCo we further adopt parameterization of the output action distribution using a TanhNormal* with both its mean and variance depending on the state representation as done by Haarnoja et al. (2018); Andrychowicz et al. (2021). As we study the connection between performance and representation dynamics this is more natural choice than using the commonly implemented state-independent variance which would be independent of representation dynamics. The ALE models use ReLU activations (Nair& & Hinton, 2010) and the MuJoCo ones tanh; we also experiment with ReLU on MuJoCo. We use separate actor and critic models for both environments unless specified in Section 4. Details on the performance metries and tables of all environment parameters, model architectures, and algorithm hyperparameters are presented in Appendix B. Observing that the previous findings on the feature dynamics of value-based approaches (Gulcehre al., 2022; Lyle et al., 2022) apply the critic of PPO as well since the loss function is the same, we focus on studying the feature dynamics of the actor unless stated otherwise in the text or figures.
Experimental setup
We vary the number of epochs as a tool to control the effects of non-stationarity, which gives the agent a more significant number of optimization steps per rollout while not changing the optimal target it can reach due to clipping, as opposed to changing the value of e in the trust region for example. We keep the learning rate constant throughout training and use the same learning rate for all the epoch configurations. 6 To understand the feature dynamics, we measure different metrics that are proposed in the literature, including feature rank, number of dead neurons (Gulcehre al., 2022), capacity loss (Lyle et al., 2022), and penultimate layer pre-activation norm. Previous work has monitored feature norm values as well (Abbas al., 2023; Lyle et al., 2024); however, in our case, we found that the neurons in the policy network die, the feature norm might be stable while the pre-activation norm blows up. All the metrics are computed on on-policy rollouts except for the capacity loss.
We run five seeds per hyperparameter configuration and report mean curves with min/max shaded regions unless specified otherwise. All curves, except for capacity loss, are smoothed using an exponentially weighted moving average with a coefficient of 0.05.
3.1 PPO suffers from deteriorating representations
Deteriorating representation How do the representation metrics of a PPO agent such as the feature rank and the capacity loss, evolve during training? Are they subject to the same decline observed by Kumar et al. (2021); Lyle et al. (2022) in value-based methods? Does affect performance?
As illustrated in Figure 1 with ALE/Phoenix as an example, we observe a consistent increase in the norm of the pre- accivvaiiioins of the feature layer of the policy network. Learning curves for all the ALE games and MuJoCo tasks considered can be found in Appendix D. The increase in feature norm is present in all the games/tasks considered in both environments, that is, with the two different model architectures and activation functions in the case of MuJoCo. We associate the rapid growth in the norm of the pre-activations with an eventual decline in the policy network's feature rank. We observe a rank decline in five out of six ALE games and seven out of eight MuJoCo tasks (four with ReLU and three with tanh). The same observations about the increasing norm of the pre-activations can be made about the critic network. However, its rank varies more with the sparsity of the reward: in
'We also provide evidence of collapse with the Gaussian distribution parameterization in the appendix. 'ee show results with varying in Figure 34 of Appendix D.
6 "Although the environments we use are single-task task environments to ablate additional MDP non-stan staaiioataitt.. they are complex enough for the agents to keep improving when trained for longer than common benchmark limits without annealing the learning rate.

---

most environments. its rank experiences a significant deterioration after the policy's 's performance declines (not the policy's rank) and rewards become sparser, and in the sparse-reward game Gravitar, the critic's rank collapses before the policy. Furthermore, capacity loss is increasing for the critic, as observed in value-based plasticity studies (Lyle et al., 2022), and we also show that the case for the actor, for which it explodes around rank collapse.
Worse consequences How does increasing the number of epochs per rollout to vary non-stationarity affect PPO agent's representation? Does it degrade as observed in DQN and SAC agents when increasing the replay ratio (Nikishin et al.. 2022; Kumar al., 2022)?
Increasing the replay ratio DQN and SAC deteriorates the agent's representation and, in turn, its performance (Kumar et al., 2022: D' Oro et al., 2023). This is commonly attributed to "overfitting" to previous experience (Nikishin et al., 2022). Increasing the number of epochs in PPO is analogous, and natural hypothesis is that this would accelerate the deterioration of the policy's representation. Figure shows that increasing the number of epochs accelerates the increase of pre-activations norm and the decrease of the policy's 's feature rank." In some cases, the rank eventually collapses, coinciding with the policy's performance collapse. We observe the performance collapse in three of the six ALE games and three of the four MuJoCo tasks.
epachs
epocchs
Eepcet
70"
70x
lnn
Cuppe
Cuppo
Figure I: Deteriorating performance and representation metrics The policy network of PPO-Clip agent on ALE/Phoenix-v3 v5 is subject deteriorating representation. The norm of the pre activations of the penultimate layer consistently increases, and its rank eventually decreases. Performing more optimization epochs per rollout to increase the effects of non-stationarity accelerates the growth collapse of its rank. This ultimately leads to the collapse of the policy. This collapse is not driven by the value network, whose rank is still high. Both network's 's ability to fit arbitrary targets (capacity loss) is also worsening.
norm
pre-activations and the
Characterizing the collapse The collapse observe is distinct from the typical entropy col- lapse. Figure 2 shows that the policy reaches high entropy. A high overall entropy can come from an average of high-entropy states with dif- ferent action distributions or trivially from the same high-entropy distribution in all states. Our analysis reveals the latter, a zero policy variance collapsed
spoda
120077
100
Enterrat 8E3
across states. This corresponds representation where most neurons are inactive. The output thus relies solely on the bias term, as linear weights act on null feature vector, making actions near uniform across all states and collapsing performance on complex tasks.
Figure 2: Rank collapse gives a high but trivial en- tropy The rank collapse of the policy gives policy with high entropy but work outputs the same high-entropy action distribution all states, as all the neurons the feature layer are dead. Its output only depends on the constant bias torm.
variance
states. The
3.2 Collapsed representations create trust issues and unrecoverable loss
Intuitively, the heuristic trust region set by PPO-Clip should prevent sudden catastrophic changes and limit the rank collapse, which induces worse performance. However, empirically, it seems the trust region cannot mitigate the collapse. In this section, seek to understand the interaction between the rank collapse and the trust region. We argue that as rank collapses, the clipping constraint becomes unreliable and unable restrict learning. This is in line with previous works that have pointed out
'The collapse happens with all epochs configurations when trained for long enough seen in Figure 32, increasing the number of epochs a tool observe the collapse earlier rather than a condition for it happen. *We consider ReLU neuron as dead when its values are zero for all the samples in the batch and a tanh neuron dead when its standard deviation across samples is less than 0.001.
5

---

that probability ratios during training can go beyond the clipping limits with PPO-Clip (Engstrom al., 2020; Wang et al., 2020; Sun al., 2022). We believe, however, that this behavior systematic when rank collapses and does not merely happen occasionally.
Wang al. (2020, Theorem 2) state that when the gradients of the unclipped samples align with the gradients of clipped samples, the clipped samples' ratios will have their probabilities continue to go beyond the clip limit. They claim this condition would hold in practice because of "optimization tricks" or optimizer accumulated moments; however, there is no evidence that these factors induce the gradient alignment or that the alignment is present in practice. Our intuition is that representation degradation leads to alignment in the gradients and, therefore, breakdown of the trust region constraint. This can create a snowball effect, preventing PPO-Clip from preventing representation collapse. We summarize this in two observations:
Loss of trust is extreme around poor representations The average of probability ratios outside the clipping limits (below 1 - in Figure 3) significantly diverges from the clipping limit around the collapse of the agent's representation. This gives one more reason why the PPO trust region can be violated. We isolate this in a toy setting and analyze it formally in the next section. We further show in Figure scatter plots of the lowest average probability ratios in runs with their associated representation metrics.' 9 We observe no significant correlation in the regions where the representation rich (high rank, low pre-activation norm), but an apparent decrease of the average of probability ratios below 1 - < is observed as the representation reaches poor values. Note that we characterize the collapsing regime by an extremely low rank, however, is not straightforward to draw a line between low-rank representations beneficial for generalization and extremely low-rank representations causing aliasing as also acknowledged by Gulcehre al. (2022), but for environments like Atari, our figures seem to draw the line at single-digit ranks, which can be related to the action space of dimension 8+.
AlF./Mamee Saamers,, epochs
E.MmeeTiGGamm-a. 6 eppchs
F/MameeThaGammes apochs
Lso
s.0
21.00
091
DIVEENENT MessK2d
Eevicement 222
Envicement Stepx
Tae rwwwent Skeex
OeioUTe Mexiiia
Figure 3: Focusing on individual runs Individual training curves on ALE/NameThisGaame- with different epochs per batch. Extremely low ratios observed around the representation collapse of a PPO-Clip agent, implying that the heuristic trust region breaks down when representation power lacking. The last-minibatch value the PPO objective decreases towards around the representation collapse, implying a reduction in the ability to improve the policy and recover, which is corroborated by the increase in capacity loss. (Ratios are trivially above after collapse collapsed model does not change much to have values below 1 e.)
Spearment
Spewrmen
ALEPhseniovv
Deed Teurone polity
Testire cark policy (PCA)
Pozture proactraaoor rorm
Figure 4: Representation trust region Samples from ALE/Phoenix-v5 training curves. Each point maps an average the probability ratios below the clipping limit vs. its corresponding average representation metric (dead neurons, feature rank, feature norm). The average ratios are significantly lower around poor representations (high dead neurons, low policy rank, high feature norm) reflecting the failure of the trust region in this regime. Averages are over non-overlapping windows larger than episodes.
Loss of plasticity makes performance collapse unrecoverable The persistent decrease in perfor- mance overlaps with a monotonic decrease in policy variance and PPO objective. It appears that as
220 points per run, across windows of size 1% training progress, spanning at least the horizon of the environment, that points are well spaced in the run, with each point being the average of the window

---

the policy loses its ability to distinguish state, it can also ascend the PPO objective less and less at each batch (recall: after collecting a batch, the loss starts around zero with a normalized advantage. and through minibatch updates, the clipped policy advantage is ascended). Intuitively, this is implied by a loss of plasticity or a collapse entropy (no new actions to learn from). As seen in Section the entropy does not collapse, and measuring the capacity loss in Figure 3 shows that the decrease in objective gain is associated with a significant increase in capacity loss, implying loss of plasticity.
Connecting the dots Hence, around collapse, the representation of the policy is getting so poor that it is impacting its ability to distinguish and act differently across states; the trust region cannot prevent this catastrophic change as it also breaks down with a poor representation; finally, the policy's s plasticity is also becoming so poor that the agent cannot recover by optimizing the surrogate objective.
Implications and discussion The causal connection we draw between the representation dynamics, the trust region, and the performance primarily holds around the collapse regime and not necessarily throughout training. However, this does not mean that one should only be concerned about the link when performance is starting to deteriorate. The representations don't collapse all of a sudden; they deteriorate throughout training until they reach collapse. Thus, mitigating representation degradation should happen throughout training and not only when around the collapsing regime. In addition, the connection gives important insights into the failure mode of the popular PPO-Clip algorithm, whose trust region is highly dependent on the representation quality, and more generally about trust-region methods which only constrain the output probabilities.
3.2.1 A toy setting understand the effects of rank collapse on trust region
We present toy example that illustrates how collapsed representation bypasses the clipping set by PPO and cannot satisfy the trust region it seeks to set. PPO constructs a trust region around the policy Tee (-(s) of the agent evaluated at a given state 8, enforeing (in an approximate way) that the update computed on state can not move the policy (-(ss outside of the trust region. However, the constraint does not capture how updates computed on another state s' affect the policy's s probability distribution over the current state s. The underlying assumption is that updates computed on different states are, at least in expectation, approximately orthogonal to each other, and they do not interact. Therefore, restricting the update of the current state is sufficient to keep the policy within the region.
In our case, however, one can show that as the rank collapses or the neurons die, the representations corresponding to different states become more colinear." 10 Therefore, the gradients also become more colinear. the extreme case, when the rank collapses to 1, or there only one neuron alive, all representations are exactly colinear; therefore, all gradients are also. This means that even though clipping prevents the policy (-(s) on the current state s from changing due to the update of that state V (me (-Is). (-|ss will still change and move outside of the trust region due to the updates on other states 8'. Leading to the trust region constraint being ineffective and not constraining the learning process in any meaningful sense. This gives clear situation where the theorem of Wang et al. (2020) holds and can easily be analyzed as below without resorting to the theorem for an end-to-end proof or to get better intuition.
Formal statement of the toy setting Let us consider a batch containing two state-action pairs (Ir, a1) and (y, a1) with sampled probabilities Tota(a (x) and (a [y) and positive estimated advantages A(z, a1), A(y,aa)) > 0. Let
TyAA Vi -114
1000 1.15
-ic
1.5.
120.
2.30
1.05
1.00
1.05-
o(x), Φ(y) E R be fixed 1-dimensional representations of x, and y that can be seen as the output of the (frozen) penul- timate layer of a policy network with collapsed representation (all but one dead neuron), and let E RR. such that p(y)= = (x)(z). Let 0 = 02). be the last layer of the network, computing the logits of two actions, aj and ag, that are
1.00
0.5
1.5
1.0 Nirbbttch
3.5
Minibatch
x10
Figure 5: Simulation of the toy setting Left (a > 0): a gradient
on (æ, a1) takes the probability of (y. 41) up and vice versa. When one is above the threshold and should not increase, the other still pushes it. Right (a < 0): a gradient on (z, a1) takes the probability of (y, ai) down and vice versa. Both slow each down, with one forcing the other to be lower than its initial value.
fed into softmax to compute the probabilities. Le., To(a)ls) 21000
Consider PPO
The expected angle between representations shrinks to 0.
7

---

minibatch updates alternating between (æz, 01) and (v, 01). Ideally, the PPO loss increases (a (s) at gradients on (z, a1) until it reaches the clip ratio and similarly on a1). However, we show in Appendix C that a gradient step in (ar a1) also affects To (o((ll)) and depending on a will increase past its clipped ratio, or decrease it below its initial value. Hssentially, when a 2 0, a gradient on (I, 01) increases 0P therefore increasing both T (a1 |æ) and Ta The same holds for a gradient on (y. a1), causing one state to reach the clip limit first depending on a > but still have the other keep pushing its probability upwards. However, when a < 0, gradient on (z, a1) increases 0pe therefore increasing T [a) but decreasing T (ai (y). For a gradient on (y, a1) it is the opposite: 8 decreases therefore T T (a1 decreases and T (a1ly) ly) increases, causing each state to reduce the probability of the other, and depending on a < 1 one of the probabilities will dominate and push the other one down. Figure 5 shows the evolution of the probabilities when simulating the updates empirically.
Intervening to regularize representations and non-stationarity
Having observed that PPO is affected by frequent representation degradation that impacts its trust region heuristic and causes its performance to collapse, we turn to study interventions that aim at regularizing the representation of the policy network reducing the non-stationarity in the optimization. We investigate whether these interventions improve the representation metrics we track and if in turn, this affects performance. We choose simple interventions that do not apply modifications to the models during training (c.g., resetting or adding neurons) or require significantly more memory (e.g., maintaining separate copies of the models). We perform interventions on the games/tasks where the collapse is the most significant. We are interested in the state of the agent at the end of the training budget. We record the performance and representation metrics for each run as averages over the last 5% of training progress. We measure the excess ratio at a timestep as the average probability ratio above 1 + e divided by the average probability ratio below 1 at that timestep. This metric gives an idea of how much the policy exceeds the trust region. Its average value computed over the last 5% of training progress where the ratios are non-trivial, giving the same window at the end of training as the other metries when there is no collapse, otherwise window before total collapse covering 5% of training progress, as after collapse, the model does not change anymore and the ratios are trivially within the + c and 1 - c limits. We give additional details on the computation of these aggregate metrics and the interventions performed in Appendix B.
PFO: Regularizing features to mitigate trust-region issues The motivation for our first intervention and our proposed regularization method comes from our observation that the norm of the preactivation features is consistently increasing, which can be linked to the trust-region issues discussed in Section 3. We seek to mitigate this effect in way that is analogous trust region to the feature space. We apply an L loss on the difference between the pre-activated features of the optimized policy and the policy that collected the batch, as a way to keep the pre- activations of the network during an update within trust region. We apply this regularization to the pre-activations and not the activations, as dead neurons cannot propagate gradients, and even when they do, depending on the activation function, do so with a low magnitude. The regularization is an additional loss/penalty added to the overall loss. We term this loss the Proximal Feature Optimization (PFO) loss. With (s(s) as the pre-activation of the penultimate layer of the actor Te given a state
the PPO trust region, by extending the
Timaae
L PFO (0) = E
(S)) P (5:)
We apply two versions of PFO: one on only the penultimate layer's pre-activations and one on all the pre-activations until the penultimate layer. In the scope of this work. we do not tune the coefficient of PFO: we pick the closest power of 10 that sets the magnitude of this loss to a similar magnitude of the clipped PPO objective tracked on the experiments without intervention. This gives a coefficient of 1 for ALE, 1 for MuJoCo with tanh, and 10 with ReLU. The goal is not necessarily to obtain better performance but see PFO improves the representations learned by PPO and if, in turn, it affects its trust region and performance. As shown in Figure 6, the regularization of PFO effectively brings the norm of the preactivation down, the number of dead neurons down, the capacity loss down, and the rank up. This coincides with significant decrease in the excess probability ratio, especially in the upper tail. More importantly, we also see a significant increase in the lower tail of the returns where and ALE/Phocnix, with
no collapse in performance is observed anymore on ALE/Name
8

---

slight increase in the upper tail showing that PFO can increase performance. Among the interventions we have tried, PFO provided the most consistent improvements in representation and trust region.
Neat Madda
Dutles
mits
solly
pollyy (PC))
Capacity pallyy
Regularize last preactvvatinn Regularss is el vatins
O8
Raset idam
Acam equal bclax
0
200
235
5.0
107
0.5
x10'
x10
Ress Vnkin gutler
pricy
inververcion
Rogudariz lest preactvvtion
Regula be all preaction ons
Share actor criic features
Adam equal betas-
410
00
ManoThadamm
2.5
50
7/5
1%5
7.0
2.5
K10
Hess
Mid
valy
Doad
poliy
pray.
No Invervention
Regudartoe lest preactivation on-
Regularos all preactioations len--
-80
Share octor and crticc features- 60
Adom equal betas
-3
80
0.5
10
15
20 6.5
10 10#
0.0
25
5.0
+10
x10'
K10"
Figure 6: Effects of regularizing features and non-stationarity Top & Middle: ALF/Phoenix-v5 & ALE/NameThisGame-vS Regularizing the difference between the features consecutive policies with PFO results in better representations, a lower trust-region excess, and mitigates performance collapse. The same applies to sharing the actor-critic trunk. Bottom: ALE/Graviur: Sharing the feature trunk between the actor and the critic results in worse policy representation the value network is subject to rank collapse due to reward sparsity. A boxplot includes 15 runs with different epachs.
Sharing the actor-critic trunk In deep RL, the decision to use the same feature network trunk for both the actor and the critic is not trivial. Depending on the complexity of the environment, can significantly change the performance of a PPO agent (Andrychowicz et al., 2021; Huang et al., 2022a). We, therefore, attempt to draw a connection between sharing the feature trunk, the resulting representation, and its effects on the PPO objective. In this intervention, we make the actor and the critic share all the layers except their respective output layers and backpropagate the gradients from both the value and policy losses to the shared trunk. Figure 6 shows that the value loss acts as a regularizer, which decreases the feature rank and, depending on the reward's s sparsity, gives two distinct effects. In dense-reward environments such ALE/Phoenix and ALE/Name TaiaammmTeiam,,,, the ranks are concentrated at low but non-trivial values: the upper tail significantly decreases compared to the baselines while the lower tail increases. This coincides with a lower feature norm, lower excess probability ratio, and, in turn, a high tail for the returns. It also increases performance in some cases. However, the opposite is true in the sparse-reward environment Gravitar: the rank completely collapses, and the feature norms and excess ratios are very high, collapsing the model's performance. This is consistent with the observations made in the plasticity works studying value-based methods: they show that sparse rewards deteriorate the rank of the value network, and we show that when shared in an actor-critic architecture they, in turn, deteriorate the policy. is important to note that this distinction using the reward sparsity holds when comparing environments from the same family (e.g., ALE), but may not hold otherwise (e.g., comparing an ALE and a MuJoCo environment). We provide training curves showing the difference in the evolution of the feature rank when sharing the actor-critic trunk in Appendix D. To further strengthen this observation we run an intervention on ALE//hoenix (a dense reward environment), with a reward mask randomly masking a reward with 90% chance, comparing the effects of sharing the actor-critic trunk. As expected, while with dense rewards, sharing the trunk is beneficial in ALE/Phoenix (Appendix Figure 21), with the sparse reward, the opposite is true: sharing the trunk is detrimental (Appendix Figure 35).
Adapting Adam Asadi et al. (2023) argue that as the targets of the value function change with the changing policy rollouts, the old moments accumulated by Adam become harmful to fit the new targets and find that resetting the moments of Adam helps performance in DQN-like algorithms. As the PPO objective creates a dependency on the previous policy, and more generally, in the policy gradient, the advantages change with the policy, the same argument about Adam moments can be made for PPO. Furthermore, Dohare et al. (2023b); Lyle et al. (2023) advocate for decaying the second moment of Adam faster than its default decay of 0.999 when training under von-stationarity

---

and set it to mathh the decay of the first moment. Therefore. we experiment with both resetting Adam's moments after each batch collection (to avoid tuning its frequency) and setting the second moment to decay at the (smaller) default decay of the first moment for both the actor and the critie; the moments are thus only accumulated over the epochs on the same batch in the former and over shorter batch sequences in the latter. We observe in Figure 6 and Appendix D that these interventions reduce the feature norm and increase the feature rank on ALE, which also reduces the excess probability ratio and, in some cases, improves performance; however, they are not sufficient to prevent collapse and, like sharing the actor-critic trunk, result in poor performance on ALE/Gravitar.
5 Related Work
Our work is complementary to various other works studying the plasticity and representation dynamics of neural networks trained under non-stationarity. Kumar et al. (2023) provide a comprehensive comparison and categorization of methods used to mitigate plasticity loss in continual supervised learning tasks and their effects on representations. Our work provides insights into the transferability of some of these solutions to RL and tools to evaluate their impact on trust region methods. Sokar al. (2023) provide an alternative characterization of plasticity loss in RL using dormant neurons and observe an increase in dormant neurons for non-stationary objectives. Abbas et al. (2023) study representation metrics such as feature norms and observe decrease of the norm due to dying neurons. Like in the work of Lyle et al. (2022), both studies only include value-based methods. In this work, we study dead units and capacity loss as Lyle et al. (2022) and provide corroboration of the dying units phenomenon in policy optimization methods and, taking the dying neurons out of the equation, find that the norm of preactivations actually blows up.
Other feature regularizations similar to PFO have been studied in value- based offline RL. Kumar et al. (2022) propose DR3, which counteracts an implicit regularization in TD learning by minimizing the dot product between the features of the estimated and target states. Ma et al. (2023) propose Representation Distinction (RD) which tries to avoid unwanted generalization by minimizing the dot product between the features of state-action pairs sampled from the learned policy and those sampled from the dataset or an OOD policy. Both are related to PFO as the methods directly tackle an undesired feature learning dynamic, but there is no motivation for DR3 or RD in online RL. and PFO conceptually different. The implicit regularization that DR3 counteracts not present in on-policy RL as shown by Kumar et al. (2022) in the SARSA experiment, and PFO differs from DR3 as it extends a trust region rather than counteracts an implicit bias. Similarly, the overestimation studied by Ma et al. (2023) in the vicious backup-genee generalizzaainnn cycle is broken by on-policy data, and RD regularizes state features between the learned policy and the dataset policy, not consecutive policies.
6 Conclusion and Discussion
Conclusion In this work, we provide evidence that the representation deterioration under non- stationarity observed by previous work in value-based methods generalizes PPO agents in ALE and MuJoCo with their common model architectures and is connected to performance collapse. This brings a novel perspective to previous works that showed that PPO agents lose plasticity throughout training. We show that this is particularly concerning for the heuristic trust region set by PPO-Clip, which fails to prevent collapse as it becomes less effective when the agent's representation becomes poor. Finally, we present Proximal Feature Optimization (PFO), simple novel auxiliary loss based on regularizing the evolution of features that mitigates representation degradation and, along with other interventions, shows that controlling representation mitigates performance collapse.
L Limitations and open questions In this work, we study the common architecture and optimizer of PPO agents in ALE and MuloCo consisting of relatively small models without normalization layers, weight decay, or memory (e.g., not using Transformers and RNNs) Despite our best attempts, as with any other empirical machine learning work, the generalization of our results to other settings is not fully known. Still, this work should raise awareness about the representation collapse phenomenon observed in PPO and encourage future work to monitor representations when training PPO agents, as it can help diagnose performance collapse. We have focused on simple interventions that regularize non-stationarity and representations to highlight the effects of non-stationarity and the connection between representation, trust region, and collapse, but exploring interventions on plasticity is also valuable, as these may also influence the same dynamics. We believe further studies to analyze this problem, both empirically and particularly theoretically, to understand the reasons driving representation deterioration to be valuable. We hope that our study encourages work in this direction.
10

---

Acknowledgments and Disclosure of Funding
We extend our gratitude to the reviewers for their valuable insights, which significantly enhanced the clarity and rigor of this work. We particularly grateful the area chair for their guidance which shaped the final version of this paper. We thank the SCITAS team at EPFL for the access to the beta testing phase of their new cluster. We are also grateful Vincent Moens for his support with the TorchRL library. Finally, we thank nimble. ai for their generous gift to the CLAIRE Lab, which supported D.P.'s Master's project.
References
Zaheer Abbas, Rosie Zhao, Joseph Modayil, Adam White, and Marlos C. Machado. Loss of plasticity in continual deep reinforcement learning, In Sarath Chandar, Razvan Pascanu, Hanie Sedghi, and Doina Precup (eds.), Proceedings of The 2nd Conference on Lifelong Learning Agents, volume 232 of Proceedings of Machine Learning Research, pp. httops//procceeiings mlr www.ppreessssvee. v232// abbas23a. html.
PMLR, 22-25 Aug 2023. URL
Matthew Aitchison, Penny Sweetser, and Marcus Hutter. Atari-5: Distilling the arcade learn- ing environment down to five games. Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th Inter- national Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 421-438. PMLR, 23-29 Jul 2023. URL https: https: htttp://pprcceednngs mlr press// #202/aitchison28.. html.
Maksym Andriushchenko, Dara Bahri, Hossein Mobahi, and Nicolas Flammarion. Sharpness-aware minimization leads to low-rank features. In A. Oh, Neumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural In- formation Processing Systems, volume 36, pp. 47032-47051. Curran Associates, Inc., 2023. URL https: meeeoeeeeeddinnneeeeeessee... neurips. cc wwtpccppppppllaeeeeee www.s//pppe//////// 02accdaa9t330066ffffeee33000 Paper- Conference. pdf.
Marcin Andrychowicz. Anton Raichuk, Piotr Stañczyk, Manu Orsini, Sertan Girgin, Raphaäl Marinier, Leonard Hussenot, Matthieu Geist, Olivier Pietquin, Marein Michalski, Sylvain Gelly, and Olivier Bachem. What matters for on-policy deep actor-critic methods? a large-scale study. In International Conference on Learning Representations, 2021. URL https: http://ppeeeeeeeeeeeeeee net/ ffreetffformidd nIAxjaniDzz.
Kavosh Asadi, Rasool Fakoor, and Shoham Sabach. Resetting the optimizer in deep rl: An em- pirical study. In A. Oh, T. Neumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 72284-72324. Curran As- sociates, Inc., 2023. URL https: https://prooceedings neurips neurips.cc/paper files/ wwwes/ppperr/20222/ paper/2023/ Compperr-ooooeeeene...pe
Marc Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying count-based exploration and intrinsic motivation. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, vol- neurips. cct.pppe.cpppppr. ume 29, Curran Associates, Inc., 2016. URL https: Paper Paper.pdf. files/ httppp//ooo....aaaaaaaaaaaaaaaar
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Albert Bou, Mattco Bettini, Sebastian Dittert, Vikash Kumar, Shagun Sodhani, Xiaomeng Yang, Gianni De Fabritiis, and Vincent Moens. TorchRL: A data-driven decision-makir library for pytorch. In The Twelfth International Conference on Learning Representations, 2024. URL https: https:://pentes /otes////ppeeeiieeeee wwt.:nfeurmmddddeeooooo Aarrmmm?Add
Shibhansh Dohare, Richard S Sutton, and A Rupam Mahmood. Continual backprop: Stochastic gradient descent with persistent randomness. arXiv preprint 206:22008.06355, 2021.
11

---

Shibhansh Dohare, J Fernando Grmaaddz--Gariii Parash Rahman, A Rupam Mahmood, and Richard S Sutton. Maintaining plasticity in deep continual learning. arXiv preprint arXiv:2306. 138/2, 2023a.
Shibhansh Dohare, Qingfeng Lan, and A. Rupam Mahmood. Overcoming policy collapse in deep reinforcement learning. In Sixteenth European Workshop on Reinforcement Learning, 2023b. URL httts:://peenreveei hwtp:/iewweeeeeeeeeeeeeeeeeeeeee fent//ooummiinn mf dz4ymm.
Pierluca D'Oro, Max Schwarzer, Evgenii Nikishin, Pierre-Luc Bacon, Marc G Bellemare, and Aaron Courville. Sample-efficient reinforcement learning by breaking the replay ratio barrier. In The Eleventh International Conference on Learning Representations, 2023. URL https:
9aBBVJe.
Logan Engstrom, Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Firdaus Janoos, Larry Rudolph, and Aleksander Madry. Implementation matters in deep rl: A case study on ppo and trpo In International Conference on Learning Representations, 2020. URL https / http://pppereeee.eeeeee/e. net/ forum?id=r1etNirre2
Caglar Gulcehre, Srivatsan Srinivasan, Jakub Sygnowski, Georg Ostrovski, Mehrdad Farajtabar, Matthew Hoffman, Razvan Pascanu, and Arnaud Doucet. An empirical study of implicit regular- ization in deep offline RL. Transactions on Machine Learning Research, 2022. ISSN 2835-8856.
URL https: /tttps://pperrevieee net mntt:/forrumm33d--B.........
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint arXiv: 1812.0095 2018.
Shengyi Huang, Rousslan Fernand Julien Dossa, Antonin Raffin, Anssi Kanervisto, and Weixun Wang. The 37 implementation details of proximal policy optimization. The ICLR Blog Track 2023, 2022a.
Shengyi Huang, Rousslan Fernand Julien Dossa, Chang Ye, Jeff Braga, Dipam Chakraborty, Kinal Mehta, and João G.M. Araújo. Cleanrl: High-quality single-file implementations of deep reinforce- ment learning algorithms. Journal of Machine Learning Research, 23(274);1-18, 2022b. URL attp::/jmile
httg//aperr/v23//21 1342. html.
Minyoung Huh, Hossein Mobahi, Richard Zhang, Brian Cheung, Pulkit Agrawal, and Phillip Isola. The low-rank simplicity bias in deep networks. Transactions on Machine Learning Research, 2023.
ISSN 2835-8856. URL httpp://opennevview wtwpiiiwwwen//ffooreeree
Maximilian Igl, Gregory Farquhar, Jelena Luketina, Wendelin Boehmer, and Shimon Whiteson. Transient non-stationarity and generalisation in deep reinforcement learning. In International Conference on Learning Representations, 2021. URL https://openreviiii net/ forumidd Qun8fv4qSby.
Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. In Proceedings of the Nineteenth International Conference Machine Learning, ICML '02, pp. 267-274, San Francisco, CA, USA, 2002. Morgan Kaufmann Publishers Inc. ISBN 1558608737.
M. G. Kendall. A new measure of rank correlation. Biometrika, 30(1/2):81-93, 1938. ISSN 00063444. URL http::/wwww.. jstor org/stablle2333228.
Aviral Kumar, Rishabh Agarwal, Dibya Ghosh, and Sergey Levine. Implicit under-parameteriin inhibits data-cfficient deep reinforcement learning. In International Conference Learning Representations, 2021. URL
Aviral Kumar, Rishabh Agarwal, Tengyu Ma, Aaron Courville, George Tucker, and Sergey Levine.
DR3: Value-based deep reinforcement learning requires explicit regularization. In International Conference on Learning Representations, 2022. URL https: / bttps://openneview.neee. pppemreriew..ne//foumriiii POwMvLi91f.
Saurabh Kumar, Henrik Marklund, and Benjamin Van Roy. Maintaining plasticity in continual learning via regenerative regularization, 2023.
12

---

Clare Lyle, Mark Rowland, and Will Dabney. Understanding and preventing capacity loss in reinforcement learning. In International Conference on Learning Representations, 2022. URL htttps://openreview
Clare Lyle, Zeyu Zheng, Evgenii Nikishin, Bernardo Avila Pires, Razvan Pascanu, and Will Dabney. Understanding plasticity in neural networks. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning. volume 202 of Proceedings of Machine Learning Research, pp. 23190-23211. PMLR, 23-29 Jul 2023. httLttttt/ppcceeeeeeee v202/1yle23b html.
mlr.press.
Clare Lyle, Zeyu Zheng, Khimya Khetarpal, Hado Hasselt, Razvan Pascanu, James Martens, and Will Dabney, Disentangling the causes of plasticity loss neural networks. arXiv preprint arXiv:2402.1876. 2024.
Yi Ma, Hongyao Tang, Dong Li, and Zhaopeng Meng. Reining generalization in offline reinforcement learning via representation distinction. In A. Oh, T. Nau- mann, Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neu- Curran Associates, ral Information Processing Systems, volume 36, pp. laaarrffllls/eeeeeeeee willee/paperr222//Eiil. Inc., 2023. URL https://prooeeddnnss neurips.cc/paper 802a4350ca4f
1543- Paper- wwwprrr-oofeeeeeeeeeeeeeeeeeeeae Conferencee.cn.oeeeeeeeeeee pdf.
Marlos C Machado, Marc G Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61:523-562, 2018.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 5 518(7540):529-533 2015.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap. Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In Maria Florina Balcan and Kilian Q. Weinberger (eds.), Proceedings of The 33rd International Conference on Machine Learning, volume of Proceedings of Machine Learning Research, pp.
New York, New York, USA, 20 22 Jun 2016. PMLR. URL https: //proceedings. mlr press/vvl vt////8///mm mn hal.6........ html.
Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In Proceedings of the 27th international conference machine learning (ICML-10), pp. 807-814, 2010.
Evgenii Nikishin, Max Schwarzer, Pierluca D'Oro, Pierre-Luc Bacon, and Aaron Courville. The primacy bias in deep reinforcement learning. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 16828-16847. PMLR, 17-23 Jul 2022. URL httpp:://ppooteeee httpc//edoo/......../oo mlr, press/vv62/ nikishin22a. html.
Evgenii Nikishin, Junhyuk Oh, Georg Ostrovski, Clare Lyle, Razvan Pascanu, Will Dabney, and Andre Barreto. Deep reinforcement learning with plasticity injection. In Oh, T. Neumann, A. Globerson. K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 37142-37159. Curran Associates. Inc., 2023. URL http:htsts://occccediiiee
wwwuripp.ccc/.. www.cc/ppppeeeeeeeeeeeeeeeeeeeee files/ http:/pppa//////////// pdf.
Chris Nota and Philip S. Thomas. Is the policy gradient a gradient? In Proceedings of the 19th International Conference on Autonomous Agents and MultiAgent Systems, AAMAS '20, pp- 939-947, Richland, SC. 2020. International Foundation for Autonomous Agents and Multiagent Systems. ISBN 9781450375184.
Fabio Pardo, Arash Tavakoli, Vitaly Levdik, and Petar Kormushev. Time limits in reinforcement learning. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International
13

---

Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, Pp. 4045-4033 PMLR, 10-15 Jul 2018. URL URLhhtts:/pproeeeiing mlr www.s..ressv///// paat//aar....... html.
William H. Press, Saul A. Teukolsky, William T. Vetterling, and Brian P. Flannery. Numerical Recipes 3rd Edition: The Art of Scientific Computing. Cambridge University Press, USA, edition, 2007. ISBN 0521880688.
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1-8, 2021. URL http:// http://mmmrrrrrppperreeee
wwtg/pppprs/v2222000 1364. html..
Olivier Roy and Martin Vetterli. The effective rank: A measure of effective dimensionality. In 2007
15th European Signal Processing Conference, pp. 606-610 610, 2007.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 1889-1897, Lille, France, 07-09 Jul 2015a. PMLR. URL https: httpss:///ccceeedeeeeeeee mlr. press/ v37/schulman15 html.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional continuous control using generalized advantage estimation. arXiv preprint 1rXiiv:500.000 2015b.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Ghada Sokar, Rishabh Agarwal, Pablo Samuel Castro, and Utku Evei. The dormant neuron phe- nomenon in deep reinforcement learning. In International Conference Machine Learning, pp. 32145-32166 PMLR, 2023.
C. Spearman. The proof and measurement of association between two things. The American Journal of Psychology, 100(3/4):441471) 1987. ISSN 00029556. URL http:://wii jstor. org/stable/14226B5
Mingfei Sun, Vitaly Kurin, Guoqing Liu, Sam Devlin, Tao Qin, Katja Hofmann, and Shimon Whiteson. You may not need ratio clipping in ppo. arXiv preprint arXiv:2202.00079, 2022.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Richard Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradi- ent methods for reinforcement learning with function approximation. In S. Solla, T. Leen, and K. Müller (eds.), Advances in Neural Information Processing Systems, volume 12. MIT Press, 1999. URL https: https::/proccedddngg/ http://prr.ediii....r..rrr.....p nwwrins..nn.mm www.neeuw....ppppeeeeeeeeeeeeeee
Csaba Szepesvári. Algorithmms for reinforcemen learning. Springer Nature, 2022.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEFF/RSJ International Conference on Intelligent Robots and Systems, pp. 5026-5033. IEEE, 2012. doi: 10.1 1.1109/RROS.000 200.0I....22...66600 2.6386109
Mark Towers, Jordan K. Terry, Ariel Kwiatkowski, John U. Balis, Gianluca de Cola, Tristan Deleu, Manuel Goulão, Andreas Kallinteris, Arjun KG, Markus Krimmel, Rodrigo Vicez-Vieeet,, Andrea Pierré, Sander Schulhoff, Jun Jet Tai, Andrew Tan Jin Shen, and Omar G. Younis. Gymnasium, March 2023. URL https://zenodo
httg//ecord//11/000
Yuhui Wang, Hao He, and Xiaoyang Tan. Truly proximal policy optmizztiinn In Ryan P. Adams and Vibhav Gogate (eds.), Proceedings of The 35th Uncertainty in Artificial Intelligence Conference, volume 115 of Proceedings of Machine Learning Research, 113-122. PMLR, 22-25 Jul 2020. URL htttt:://ppoceeednnnnn http:/aeeddinnna.....rrrreee///e www.s/eevvvvveeee wwtp///ng////.....///// .ttml
Yuzhe Yang, Guo Zhang, Zhi Xu, and Dina Katabi. Harnessing structures for value-based planning and reinforcement learning. In International Conference on Learning Representations, 2020. URL.
14

---

A Additional background
A.1 Reinforcement Learning
The undiscounted formulation presented in the background (Section 2) has also been used by Schulman et al. (b015b) and does not limit the use of a discount factor discount future rewards; for that purpose, as we consider finite-horizon setting, we can assume that discounting would already be present in the reward which depends on time through the state. This allows to isolate the discount factor 7 for the purpose of the value estimation with GAE which serves as a trade-off between the bias and the variance in the estimator, in addition to À used for the X-returns that combine multiple n-step returns. More importantly, this also allows us to reuse the policy gradient and PPO losses without discount factors, as the deep RL community used them while avoiding their incorrect use in the discounted setting as pointed out by Nota & Thomas (2020). In any case, our results can also be translated to the discounted setting using a biased gradient estimator (missing a discount factor), being the typical setting considered in deep RL works.
B Experiment details
B.1 Code and run histories
Our codebase is publicly available at https://github. CoAAIIAEC/o//AAREEaLaa.o irepresentatioonno-tte
It includes the development environment distributed as Docker image for GPU-accelerated machines and a Conda environment for MPS-accelerated machines, the training code, scripts to run all the experiments, and the notebook that generated the plots. The codebase uses TorchRL (Bou et al., 2024) and provides comprehensive toolbox study representation dynamics in policy optimization. We also provide modified scripts of CleanRL (Huang et al., 2022b) to replicate the collapse observed in this work and ensure is not bug from our novel codebase.
The code repository contains links the Weights&BBases (W&B) project with all of a summary W&B report of the runs, and a W&B report with the replication with CleanRL.
histories,
Runs are fully reproducible on the same acceleration device on which they were run. In particular, we have reproduced our results on three different clusters with the same NVIDIA GPU device.
B.2 Additional details on our experimental setup
We conduct experiments on an environment with pixel-based observations and discrete actions and an environment with continuous observations and actions, each with different model architecture. For the discrete action case, we use the Arcade Learning Environment (ALE)BBellemrre et al., 2013) with the specification recommended by Machado et al. (2018) in v5 on Gymnasium (Towers et al., 2023). That is, with a sticky action probability of 0.25 as the only form of environment stochasticity. using only the game-over signal for termination (as opposed to end-of-life signals) with the default maximum of 108 x 103 environment frames per episode and reporting performance over training episodes (i.e., with sampling according to the policy distribution as opposed to taking the mode action). We train all models for 100 million environment frames. We use standard algorithmic choices to make setting and results relevant to previous work. This includes taking only the sign of rewards (clipping) and frame skipping. We use frame skip of 3, as opposed to the standard value of 4, due to limitations in the ALE-v5 environment, which does not implement frame pooling." !1 We use the standard architecture of Mnih et al. (2015) consisting of convolutional layers followed by linear layers, all with ReLU activations, and no normalization layers. We also use Mnih et al. (2015)'s standard observation transformations with a resizing to 84x84, grayscaling, and a frame stacking of 4.
For the continuous case, we use MuJoCo (Todorov et al., 2012) with v4 on Gymnasium (Towers al., 2023) with the default maximum of 1,000 environment frames to mark episode termination. Similarly to Atari, we report performance as the average episode return over training episodes. We
That is taking the max over the last two skipped and unskipped frames to capture elements that only appear in even or odd frames of the game (https: h/ap::////t /ggb. wtt.b...oo/.....ee com/Farama Foundation/ Arcade -ecanin-earning--
Lew.aniinnn-iiiromeeneteeee Eww.rnneentteeeeeeeeeeeaeeeeeeee //ssutt//s11/1 Using an odd frame skip value alleviates the issue.
15

---

train all models for 5 million environment frames. We standardize the observations (subtract mean and divide by standard deviation) according to initial rollout of 4,000 environment steps (at least four episodes). The standardization parameters are kept the same to avoid adding non-stationarity. We use the same architecture as Schulman al. (2017), with only linear layers, tanh activations, and no normalization layers. We also experiment with ReLU activations. The network outputs mean and standard deviation (with softplus), both conditioning on the observation independently for each action dimension, which are then used to create TanhNormal distribution, similarly to Haarnoja et al. (2018).
To measure the capacity loss of a checkpoint, we use the same optimization hyperparameters used to train the checkpoint, i.e. the same batch size and learning rate. The optimizer is reconstructed from its initial state (loading the optimizer state is also valid design choice). The dataset sizes and fitting budgets for capacity are listed in Tables 2 and
We provide high-level pseudocode for PPO in Algorithm and list all hyperparameters considered Tables and? 3.
Algorithm 1 High-level Pseudocode for PPO
N: number of environments in parallel.
Bens agent steps per environment to collect in a batch. K: number of optimization epochs per batch.
CLIP (0): PPO-Clip objective.
H entropy bonus/loss; cи¹ entropy bonus coefficient.
L LVF (w): critic loss (L² to GAE); pt critic loss coefficient.
1: while collected environment steps I total environment steps do Collect batch of interaction steps of size B N Beny and computes advantages. for actor = 1 to N do
2:
3: Run policy oold environment for Bew agent steps. Bum with GAE.
4: Compute advantage estimates Vpactor
5: end for
6: Minimize overall policy and value loss (-L -H H (0) +000)ooFVVV (W))) with autograd
on the collected batch over K epochs with minibatch size M B.
7: Told < Te
8: end while
Proximal Feature Regularization With coefficient CPFO, the PFO loss is added the overall loss (- Fai ( + CPPOL (0) CH + cy FV V (W))) optimized with autograd over multiple minibatch epochs.
B.3 Additional details on metries used in the figures
Training curves A point in the training curves in Figures 1, 2, 3, before aggregating seeds and smoothing, corresponds to an average value over the last batch collected at the time of logging for metrics available at every batch (feature rank, entropy, etc.) or the latest batch where the metric was available the time of logging for the episodic return (as it's only available when episodes finish, and it requires multiple batches to finish an episode). F.g., in Figure 1 on ALE, feature rank corresponds to the average feature rank over the states in the last batch collected at the time of logging and is logged every 0.1% of the batches (i.e. every 1.4 env steps); A return corresponds to the average return across all workers that had episodes finished in the latest batch containing finished episodes at the time of logging.
Figure 4 window of size 1% of training progress represents approximately 1 million training steps on ALE and 50,000 training steps on MuloCo We average the metrics per window and then take the run
20 windows with the lowest average probability ratios below e. The probability ratios can be trivially within the 1 - region after the model collapses, resulting in less than 20 points if the model collapses before 20% of the training progress. When all runs give points, we can observe 300 points in total per scatter plot.
16

---

Table I: Hyperparameters for the toy setting in Figure 5.
Environment
x(x)
Sampled from a Normal distribution aΦ(z)
o(g)
3 (overshoot), -1 (interfere)
A(y, o1)
Policy
Network
2 output neurons representing the 2 logits Softmax
Optimization
Clipping epsilon (PPO-Clip)
0.1
Optimizer
SGD
Learning rate
1.5
Minibatch size
Number of epochs
10
Number of steps
20 alternating between and y
Figure A window size 5% of training progress represents approximately five million training steps in ALE and captures at least five episodes per environment so in total at least 40 episodes. For MuJoCo this represents approximately 256,000 training steps and captures at least 128 episodes per environment so in total at least 256 episodes.
When model collapses, it typically doesn't change anymore its optimization trivially gives ratios within the clipping limits (no value above + and below 1 - is logged). In that case, we are more interested in the evolution of the excess ratio before the ratios become trivial. Therefore, the upper limit of the 5% of training progress is taken such that it is the latest timestep where there are least 10 non-trivial ratios, i.e. 10 logged excess ratios. This coincides with window before the collapse of the model capturing the values we are interested in. Note that when model collapses this window may coincide with the window used report the other metries such the average return, however, these other metrics typically do not change after collapse, so it is more robust to capture them at the end of training rather than looking for an arbitrary window after the collapse. We give training curves similar to Figure with the interventions performed.
In MuJoCo, with continuous action distributions the ratios diverge infinity and 0 before collapse therefore to get meaningful plots, we clip average probability ratios above + e and below 1 - to 10 and 10 respectively, before computing the average excess ratio.
We group the different epoch configurations of an intervention on the same environment, giving 15 runs per boxplot (three epochs with five seeds each). The right (resp. left) whiskers are determined by the highest (resp. lowest) observed datapoint below Q3 + 1.5 IQR (default of Matplotlib). The outliers are points outside of the whiskers.
B.4 Statistical significance
Stochasticity our experiments arises from network initialization, environment transitions (e.g., sticky actions in ALE), agent action sampling, and minibatch sampling for optimization. A seed fully controls the sequence of randomness a run with the same hyperparameter configuration. We repeat each configuration with five seeds using the same collection of seeds, resulting in the same initialization of the networks and environments for a given seed across configurations. This form of repeated measures allows us to compare the configurations with lower variance they share the same initial conditions, hence requiring a lower number seeds.
In Figures/ and 2, we aggregate the five runs of each experiment into mean curves with min/max shaded areas. The use of min/max error bars allows us to demonstrate full range of observed outcomes, although it may result in shaded areas that overlap more than with other types of error bars. Most of the claims we make based on those figures do not rely on non-overlapping shaded areas and are instead stronger when the max min boundaries consistent with the observation made
17

---

(min boundary of feature norm increasing, max boundary of feature rank decreasing). Otherwise, we made comparative claims when shaded areas did not overlap (feature rank decreasing faster with more epochs and more non-stationarity).
Figure displays individual seeds zoom single-run dynamics around collapse. It is used for an illustrative purpose to provide intuition and does not depend on the number of runs or statistical aggregation of results. The main claim made with the intuition (breakdown of the trust region) is backed by Figure 4, which includes 300 points per plot per environment, subsampled from 15 training curves per environment.
To evaluate the effects of the interventions in Figure 6, we show boxplots to give a complete idea of the distribution of the data which is formed by grouping the different configurations in the same environment. Each boxplot contains 15 runs. We make claims such as preventing collapse using the tails and medians and claims about lower excess ratio and higher rank using the interquartile range. Without clear intuition about the distribution of combined configurations per environment, we consider this approach appropriate for comparing interventions.
In summary, we believe our experimental design provides a balanced tradeoff between statistical significance and richness of claims. The computational cost of running more seeds may not yield proportionately valuable insights.
B.5 Hardware and runtime
The experiments in this project took a total of -11, ~11,300 GPU hours on NVIDIA V100 and A100 GPUs (ALE) and -25,500 CPU hours (MuJoCo). A run on ALE takes around 10 hours on A A100 and 16 hours on a V100. A run MuJoCo takes around 5 hours on 6 CPUs.
18

---

Table 2: Hyperparameters for ALE.
Environment
Repeat action probability (Sticky actions)
0.25
Frameskip
Max environment steps per episode
108,000
Noop reset steps
Observation transforms
Grayscale
True
Resize width ('resize_w) w")
84
Resize height ('resize_h")
84
Frame stack
Normalize observations
False
Reward transforms
Sign
True
Collector
Total environment steps
100,000,000
Num envs in parallel
8
Num envs in parallel capacity
Agent steps per batch
10,24 (128 per env)
Total agent steps capacity
36,000 (at least one full episode)
Models (actor and critic)
Activation
ReLU
Convolutional Layers
Filters
[32, 64, 64]
Kernel sizes
[8, 4, 3]
Strides
[4, 2, 1]
Linear Layers
Number of layers
Layer size
512 Optimization
Advantage estimator
Advantage estimator
GAE
Gamma
0.99
Lambda
0.95
Value loss
Value loss coefficient
0.5
Loss type
L2
Policy loss
Normalize advantages
minibatch normalization
Clipping epsilon
0.1
Entropy coefficient
0.01
Feature regularization coefficient
(last pre-activation), 10 (all pre-activations)
Optimizer (actor and critic)
Optimizer
Adam
Learning rate
0.00025
Betas
(0.9, 0,999), (0.9, 0,9) for the intervention
Max grad norm
0.5
Annealing linearly
False
Number of epochs
4, 6, 8
Number of epochs capacity fit
Minibatch size
256
Logging (% of the total number of batches)
Training Capacity
every 0.1% (-100,000 env steps)
every 2.5% (41 times in total)
19

---

Table 3: Hyperparamelers for MuloCo.
Environment
Frameskip
Max env steps per episode Noop reset steps
1,000
Observation transforms
Normalize observations
True (from initial steps collected by uniform policy)
Initial random steps for normalization
4000 (at least 4 episodes)
Collector
Total environment steps
5,000,000
Num envs in parallel
Num envs in parallel capacity
Agent steps per batch
2048 (1024 per env)
Total environment steps capacity
4,000 (at least full episodes)
Models (actor and critic)
Activation
Tanh, ReLU
Convolutional layers
Number of Layers
Linear layers
Number of layers
Layer size
64
Optimization
Advantage estimator
Advantage estimator
GAE
Gamma
0.99
Lambda
0.95
Value loss
Value coefficient
0.5
Loss type
L22
Policy loss
Normalize advantages
minibatch normalization
Clipping epsilon (PPO-Clip)
0.2
Entropy coefficient
0.0
Feature regularization coefficient Optimizer (actor and critic)
(tanh), 10 (ReLU)
Optimizer
Adam
Learning rate
0.0003
Betas
(0.9, 0,999) (0.9, 0,9) for the intervention 0.5
Max grad norm
Annealing linearly
False
Number of epochs
10, 15, 20
Number of epochs capacity
Minibatch size
64
Logging (% of the total number of batches)
Training Capacity
every 0.1% (6,144 env steps) every 2.5% (41 times in total)
20

---

C
Toy setting derivation details
The derivatives of the softmax probability Te(a2]s) with respect and 0g are follows:
@To(ai|s)
(0(()
clo e((((s)
(Bi((e) + CB20(03 = o(0)
(3)
001
00,
(el10(e) + cl2d(0))2
(mp(a))s)
θ (l)((s)
(l1((s) c²20(()
((s)
(4)
202
002
el10(+)) + eo20(s)
(E"1o(a) + e02o(t))²
The update rule for each parameter in with SGD is 6.ew 8, + 17 OL where 17 is the learning rate. 3e, Therefore, given the partial derivatives, the updated values for and 0g after taking gradient step are (if the probability is still inferior to I c, otherwise the gradient is 0)
A(s, a1)
(010(s) e020(0) (E + c
A(s,
(li((e) (220(0) (c + e220(1))22
(())) 81 +1 (aa(a
and 6.Th
02 1
Tosal (a:ls)))) ((())
Hence,
with 8, = A(s,a)) (as)(s (c010(0) + e220(0132
(d(s)
( (00((s)
0Te 01
0g 82 8,
Let 2 0 and without loss of generality, let's take 1. After gradient step has
3T 6(x)
Tgr (a (r)
)(x) ((x) 0Th ((x)
(01 181+88)000
e00 el02-da)o(e]
(1b(())
c91o(x) + c(02-226))0())
e10(x))
(010(2) + (0220())226012
(10(())
(since 26zØ(z) S 0)
elo((z) + el2d(x) (21lz)
e riss ad(x))
Tore (apn(aily)) 6 a(x)) + adla)
(ab(a)
e(0)tó,)ap(z) (0
(las(z)
(lad(a)) + ((02-26_)a0(())
(laa(())
(lind())) + = 26z00(z)
(lap(())
2 Blad(a)) + (laad(a)
(since 25 < 0)
(
21

---

And after a gradient step y:
e ((x)
TTgnc (a1 [z)
e ((x) + o(L)
((()+)))(())
+ e(02-8))0(2)
e1((z))
eP10(x) + e(02-26g)((3)
(²10(2)
ed10(z) + (66g(()))))((2)
el0(())
(since 28,0(2) s 0)
(l0(z)) + cl2b(1)
Tgge (a1lz) (To(()) + adlax)
el (((++))))
((0++)))0(()) + 6,(02- elap(z)
(lab())) + 2,(22226,)v((a) (lad(()
(lar((w) + (((2ad())-2 (la((())
(adl(a) + elaal(z) (since 28,a0(() S 0)
Ta11,a2 0)
(a1 ly)
Let a < O and without loss of generality, let's take a S 1, after a gradient step on x one has
e 2t s(x)
Tgr (a1 (
(r ((1)
0T Ō(x)
(o1+)))
el0, Jo(K)) (002=
(1((2)
(lo(x) + el02 ((02-26.)((e)
(1((2)
(l0(x) + ((020(())-220 26x0(z)
(10(x)
(since 26_ $ 0)
(P10(x) + el2d(a)
(T(2)
22

---

e ax(x)
Te (11, y)
e a0(2) + t ao(a)
el01t),)a0(z)
((0++)))00(2)) + ((02-)))ad(a) (l0((2)
(hiad(a) + (002-2)))a000))
(la0(())
(lao(()) + -22 (lov(())
(since 2))((((1) 2 0)
(abla) + el2ad(a)
To (11l))
And after a gradient step on y:
e ((x)
The (a1 [r] d(a) + 2 6(0)
(0)+))))(e))
e (10++))0(()) + e(02-6,))02)
(1(((2)
eP10(z) + (002-26,)00
(16(2)
e010(z) + (l20(1)) 26,
(li((u)
(since 28,Φ(z) 2 0)
(l1(() + c02d(a)
ma (x)
6 ax(x)
Tggir (all))))))))))))))))){))) (P000((2) + t no(a)
e(0(+d,)o0(z) + e(02-8,)ood clap(z)
((ao(a) + 2002-22,)a00cn
(lad(())
(lab(x)) + (2f20o(a)22,000) (0)))(z)
(la((()
(since -2,, S 0)
(aad())) + (laa((n)
(o(a)))
23

---

D
Main paper figures on all environments
epocts
epecch
Reppzta
220
Nor Procctodt 1il pole
Fulure with citte (PON
10x
101
Caoagtty
1.0
-5ppchs
Muat unn sekicp IPCA
Narm reatteroo www.ooopttt..... purg
Pout TuELANTIE 0102 (TCA)
coddy l000 poley
1.0
Eploage
101
0.5
0.5
10
0.5
0.5
1.0
spocts
6 epochs
- spocts
2.0 422
-22
pneethhette palks
3.0
fertue meek oritc (PC)
Crarctylos cfti:
80+
8.6
30t
300
0.0
0.0
0.0
255
On someet Mapa
Enttorentt Magaof
Indorrmntt Masy
Envrrsmen sezze
2.5
3.0
Sature Bank cikcet (PCR)
1.25
100
9.8
0.00
los
0.8
Toluce
0.75
Cypnn
0.25
03
05
03
0.5
0.0
spocts
epechs
spacts
10"
Capsechy oOe crte
2.0
201
0.6
2.5
10
107
0.
03
Capoo
epoche
Gepochs
Sepocts
6422
-32
Foo O paicy MCA
DRcIUm wek. cnte (PCA)
Catuebly tas pelog
1.00
2.0
0.50
30.5
dazr
Cri sage
rape
Figure 7: Figure on ALE. QBert is the only game where rank decline and collapse are not ohserved, apart from outlier run that collapsed initialization. The performance of the policy should be taken into consideration when comparing the capacity loss of the critic. F.g., for Phoenix, the capacity loss of the critic associated with the policy that doesn't collapse ends up higher than that of policies that do collapse.
24

---

112
6-5
110
spc
opacis
4.15
opecp
Figure 8: Figure on MuloCo with the tanh activation.
M00 CER POL
M LAN CANTS
spacts
Drandty Mertuniy
Dapact SAN ORC.
AA LOL PTM
E
Figure 9: Figure on MuloCo with the ReLU activation.
25

---

courta po trr
14
1.2
000
025
035
275
1:00
030
3.50
075
100
0.00
015
035
0.75
120
espodt
acodts
cpodss
x20
230"
1.5-
4.0
1.00
2.5
3.50
Enainnment staa x10
Envicrmmnt staps
K10
Ensrrnmmnnt staps K10
spocis
essccts
0.25 0.50
1.0
0700
421
330
6.00
035
077
Ensronment aiapo
Envronment
7.0
F-0007
0.4
2.2
0.2
2.0
030
025
0.50
375
1700
030
335
6%
180
2.0
525
075
120
4 4esodhs
1por's
Beaach
x00-2
2.5
2.5
1.3
0.0
0.2
Enstrnment
x
Enshanment
KI0
eppcl
scocis
appcho
x10
X307
1.75
0.6
0.4
070
630
0.23 0.55 0.75 1.25 E37
525
370
stace
Eostronment stage
steps
Figure 10: Figure 2 on ALE.
26

---

spack -20gat
20-
240
164
122
Depodt Bapee tecch
EwiiorPent 2003
XI2
Delennee Moo
Figure Figure on MuJoCo with the tanh acti- vation. With continuous action distribution, the policy variance can either drop or explode. Dead neurons for the tanh activation are hard to compute they are dependent arbitrary threshold.
Figure 12: Figure 2 on MuJoCo with the ReLU activation.
ALC.Namm Game-vd, A.cNamattnGGameed,
2101
X2211
65
Figure 13: Figure on ALE, (No other environments considered; same figure as Figure 3).
2.5
2.0
Joy wetteres
Low
60
Figure 14: Figure on MuloCo with the tanh activation. The PPO-Clip objective explodes in the negative direction after collapse so clip the y-axis of that plot to -1
27

---

Dead Teurona polay
Testure polccy (PCA)
Posture preactrastor rorm
Sandil:
Rendul
p0005
Ocad polcy
Posturc rark polcy
prscciiation
Fencal Spesman
Kendal
0.00
Kendall Spearsn -670
0..2
(CA)
preactratter rorm
Kencall Spasman: 0.47
Kendel Spearmar:
Kandat: 0. Spearman -0.33
ADGrevttees
13
20
Rondat 420 Spsaman: 0.35
Kerdal: 0 Spearmar 0.15
Sandil: 0.45, Spearmnn: 0.60
of
10 15
26
02
23
121
121
Zensal 3.R,, Spesman 0.26
Kendall 13,3, Speemani 6.18
gendall: o 35, Spearmon -0.56
ALL
Dead nexrane
Ao festure
(PCA)
Figure 15: Figure ALE. Qbert and Gravitar do not have runs with poor representation regions (dead neurons > 510) to exhibit the correlation around collapse. Qbert has one outlier where the agent collapsed at the very beginning of the training and kept high (but lower than 510) number of dead neurons and trivial rank, but excess ratio.
28

---

Figure 16: Figure on MuloCo with the tanh activation. Dead neurons for the tanh activation are hard to compute as they are dependent on an arbitrary threshold. In Humanoid the rank does not arrive at low values to exhibit the correlation around collapse.
C.l.
Urmuub Sprcman
200
Spranee A.3
235
Heatur cnarovitin
6.41
3.2
2.1 Opmamee
Desi 1r poice
Figure 17: Figure 4 on MuJoCo with the ReLU activation. In Ant, the rank does not arrive at low values to exhibit the correlation around collapse.
29

---

palty
palty
Requ last preactivation
Regulare al creamvatiinn
Seeat liam
Aoam
Phonsix
200
25
5.
16
x10
X10
X20
K30
Haa Mel
butles
pallyy
Rogudarize list preadivation Nagula uo all proacilsttons hose
1
Allm equal bolas
01
810
000
MamuThaGame
2,5
5.0
7.5 x30
1.5
2.0
25
10
10
x10"
x30
Hoan Med
Dutlers
rate
solicy
palayy
Inververtion
hogudarize list preadivation
Nagular uo all proaciloation hasel Allm equal botas
-80
Greatar
05
100
25 3.005
10 10
0.0
2.5
5.0
x30
Figure 18: Figure 6 on ALE. The tails of the capacity loss on Phoenix with interventions can be higher than without interventions on the runs where the models collapse too early without interventions, leading to the capacity loss of the non-collapsed models with interventions eventually becoming higher. This can be observed from the training curves with interventions. Nevertheless, their medians are lower.
Wis Malin
Silles
Regularize lest preactvvtion + o000
Regula se preacisonions ons-
Share Adam equal betas:
HD
Happer
101
10
2.5
5.8
10
10
Maas
Madian
Outlees
Eploude seturn
Faceor rati
Dead
Norm amactivation
Festure policy (PCN)
Capacty lons
invervenionn
Lertte test preactivation on-
+B0++
680
Regularos preactionl lon-
-B
Share octor crticc features-
Adam equal betas--
08
0.5
1.0
10
16
10P
10
10
Maas Madian
Outiees
Eplonde returr
Dead
Norm amactivaaiinn
Fasture mok (PCM)
Capacty lons
invervention
Regutarte test preactivation on
Regularos preactiont
Share octor crtic features-
Adam equal betas--8
8.00
00
100
0.5
25
3.
101
10
0.0
25
3.0
10'
Mian
Mellan Dutles
Eploods
Escos ratl
Doad neurors
Norm proectivtton policy
Fasture rank paliy (PCA)
Capac losa
No in
00
Rogularize lost preacivvtion
Nogula preacihutions ons
Share sctor and citic festures
Adam equal bocas
0
0.0
25
10F
10
107
5.0
Figure 19: Figure on MuJoCo with the tanh activation.
30

---

Pades Outlia
Tpleode retum
Docers reto
Dead seumns polte
Norn preactiadior sclicy
Teature rant polloy IPCA
Capacity loes pelicy
No Inverventonn
0
Requlartzn last preactrvation
Resst Adam
Share actor and crtlic features
8+
Acem aqual
Happer
x1ot
K12
X12
tyisode selum
EauDi ruflo
Dcuil reuttre polley
Marm inraasinttion polkr
Paaturn pallyy IPCAI
Cupaably lasie polcc
ivervention
[-20)
60
Reguariz preactivations
Passt Adam
+
Share actor
Hamonssi
x107
K10
Eplsods velurr
Esoos 1ol0
Dood nourors pollyy
Norm proxcination polcy
faalure palky (PCA)
Cupac loxo pallcy
voriion
Rogula insacivotions
Resee Azom
08
0.0
Share
HolChhetthh
x107
x10
Spisods votur
Esoos 1ate
Doad polley
Norm proaciiation policy
Paslure palky
Capas loso palccy
Rogularize anexcivvtion
Rogula ons
Resee Auom
Share tctor
Ant
2.5
10A
10
10
X30
X10
x10
Figure 20: Figure on MuJoCo with the ReI.U activation.
Orrall
2g"
0.B
Capocli boss
Eolodd
02
0.0
05
03
0.0
03
10
6.0
0.5
030
G3
0.0
05
10
Adam equal betas
No Intervention
Reaer Adam
Reguracze led prectivation
Regidarine al precttvaaiion
Share schrr andir features
212
(cout 2.41 SanTun
Num posthatton petccy
TWELUN MAR SIL (PCA
0.5
40*
0.0
Envecemen M0
Dwizsneettaeeee 10
Enviceeen Meawa
Adam equal betos
intervention
Rose Adam
Reguiaiaa last proactivation
Roqularizo proocivations
Share actor
(roul Saae mu 10TT
30"
12
2EV WAT.ETA----- Citl (PEA
Narm preattio wor.onooppppleoo
Eolose nturn
104
0.3
0.5
Figure 21: Figure on ALE/Phocnix-v5 with interventions.
31

---

prsccllation
Segidative
Neam prticheutan 22
T M SPA
2..0
1.3
1.5
2.8
03
05
10
0.0
05
0.5
10
Adam ezual betas
intervention
Rasset Adam
last preccivation
Regularicc al preectvvaiion
Share actor criicc features
10'
Captchk Lon pal
10
Edhood
101
0.
www.ereee. Cryy 210
Coppe
2010
Adam equal betas
Intervention
Raner Adam
Beguirra preactivaiion
Regularto preactvvation
Share sttor
2207
Fest-mm ea pallin 12Ck
159
enc UCA
Copaccty Ross crticc
120
Epsosa
0.50
0.25
R.S
05
17
ENVEDRTAT: Haae
Envinemen: Mage
Ervinement
Figure 22: Figure I on A.F.NameeThsSaameee with interventions.
preactvvttion
(POU)
wag poiny
Capar Do YouR catte
10
15
010
05
85
R55
Adam equal betos
intervention
Rose Ado
Reguran laz prcacivation
Reqularicc all proocivvtions
Share 200 of Rebuuncs
Nure inacciaaion on DEECy
gottey
1.5
Caradiy VOS
1.0
0.5
0.0
ecual betos
Rost Adam
Reguiaraa last proactivation
Roqularize prooctivation
Share actor
2302
1322
30"
1.0
Festurr mark oa tt: (PCA)
Crewn loe percey
Worm prw.ttaaiionnnnoooo
B.A
1.0
10
201
0.0
6.
03 Dnitsmenttoom
05
G3
0.5
1.0
Figure 23: Figure I on ALE/Name/ThiSaamenn with interventions.
32

---

Neslur FaR pocko PCC.
Norm www.tttttiionnnoooo.o De dg
(CCA)
204
Feetee mn
Adam esual betss
No interventioon on
Reset Adim
Regultiize last preactivation
Reg. Jariee preeetvvtioos
Share actor critic features
2.5
Peatue saR sokk (ros)
Rarm sraaad vation saloy
PICA)
Crpack tocpptt
2.0
1.55
1.
0.5
Foot
0.0
0.0
Wenonnee Cipyn
RESISRMNnT Weglati
sodoered Starett
Andeesee Oeyeste
PrwCOMMEN copectt
Adam equal betss
interventioon
Reset Mat
Regulorize last presctivvtion
Reg Aaiie preectvatioon
Share actor cnticc feeture
X422
.322
x32
Teat SAT patt (PCA)
Norm SANML: wartnn seleey
(RCA)
Croedhy effic of
oos
10
Festume
coppett
Figure 24: Figure I on MuJoCo Hopper with the tanh activation.
Peau WaR patte (NCA)
Nor ploadiaal on oolkz
10-4
(CCA)
0.00
604
0.4
0.0
Adam ecusl bates
To intervertioo on
Reeet Adem
Segilstte last preactivation
Reg Jarine preactivation:
Share sctor and critic fret
Poturr TwN co (PCA)
Freture park cito
1.00
0.75
tor
0.75
0.49
101
2,30
y
Narm
0.25
0.25
0.00
0.00
Unaistrent Sepec1
inferrent deas
Lnd of drpo
Adam equal
Infervention on
Beeet
Segularize preactivation
Reg preetivat
Share entiic
-12
Peour 2. koy PCN)
150
1.03
Tobose
0.75
0.50
Cpport
0.25
0.00
STAREET dean
Figure 25: Figure I on MuloCo Humanoid with the tanh activation.
33

---

spoidd
TNat (4R park PIN
Wern Ciclettion polee
Cipar By Ere Easgole
Noste ont onti
-01
Adam equat betts
No interventioon
Reset Adr
Regulatiz lest preactivation
Reg Jartee preeetvvtions
Share actor crttic features
Desture mee pakry FPC
madeyy
150
Capadiy us CIC
Entrade
101
0.02
Adam equat bette
Intervention
Reset Adtm
Regulatize lost presctivation
Regidarine presctvatiios
Share actor ontic feat.rre
Testti puRk IPCA
6,500
mm pip Sroeecioon voniin sellyy
1.50
S10
1.
210
X32l
2.3
Estocdn
10
Caroer
22
Fnd po
Figure 26: Figure I on MuJoCo HalfCheetah with the tanh activation.
og (PCA)
Nom proctio patce
Zgo
Nusuur Ac PCA)
Capar by tou pabyy
Nagure or
Adam aquae aquae betee
No Intervention
Reeat Adam
Regilarize lnd pracciiatiion
Regulattte preectivations:
Stare actor crtiic That.
Festure saT pelto (PCA)
Nim Drg coverrarooo on polcy
Capicta low ontley
255
0.00
Swnsec Cepoc
on drap
Zaa equal
Intervention
Reest
Regularize pracctivation
Regudarioe prective
Share enttic
.127
220
POTUR M TO (ROA)
festur conee (PCA
Capacity Iane
Detcod
20P
copor33!
caperto!
Figure 27: Figure I on MuloCo Ant with the tanh activation.
34

---

2.5
Nealowe nath pocke PCCA
Nar www.tttttttttooooo. pai
0.0
Lo
044
0.5
Fester
Adam esual betss
No intervention on
Reset Adim
Regulbdize last preactivation
Reg Jarize preeetivatioss
Share actor crticc features
3.0
Nattir calk polks (CCU
Borm caaa Nario sello
20"
Cipect Cipectytod toosppts polc
22.0
Caparty one
0.5
30
0.4
0.0
Weseeree City coefcrt
Eesiceeent Meglatt
sodeered Starett
Prwccemand coppcca
Adam equal betss
interventioon
Reset Mat
Regularize last presciivation
Reg Jaiie preectvvtioos
Share actor citiic features
Teatue petto (PCA)
X122
Norm p4Mat1 San veleey
1.2
A3p2
Teat cre (000
232
2527
22.0
31.0
1rpeohy pree ortic
0.0
201-
0.8
0.6
0.2
City
Figure 28: Figure 1 on MuloCo Hopper with the ReLU activation.
Reqularies
FORBU 20t gate (PCO
gaks
10"
3 1.25
30'
Capacit palcy
Quaecily LOSt reC
1.00
10"
30
0.50
D.2
Adam equal betes
No intervention on
Reset Adem
Segulstte last preactivation
Reg.dartze preactivation:
Share actor criiic Teet
Teature BaK Ba (PCN
Feer odei (PCA)
Capadit top polkcy
30%
Uwdeerer deper1
www.arse capecar
Smaloooad copocut
Adam equal
Beset
Segularize preactivation
Regulario prectivat
Share cntlic meet
Teguie ar kny PCN
pelicy
Eotorce rntum
3..0
0.75
0.75
101
STAVEE dean
Figure 29; Figure 1 on MuloCo Humanoid with the ReLU activation.
35

---

Testumm Typl pator (ICA.
1.33
Crouiiy poky
1.25
1.00
120
0.50
0.30
Adam equat equat betes
No Intervention
Reset Adtm
Regularize lest preactivation
Regulsriz preectivaiios
Share actor entie features
PeatU CaM socky (PCA
No mo preaet vatie potcy
12
31.50
1:25
Cacaeday crilee
8.
Eptodde
5.0
6.4
1p
0.
0.25
0.0
SeArccee Stpaett
conent Orpicit
Andeese Seyae
Adam equai betes
Intervention
Recet kdtm
Regulatize lud presctivation
Regidarine preectivation
Share actor and onticc feat ure.
-312
x32
6,50
FOx OO paRo (PCA)
Wm MRO aartaan setee
Teat Ce mert cre (PD
Estocdn
20"
0.4
Divisiment 23
Destomment sapertd
Figure 30: Figure 1 on MuloCo HalfCheetah with the ReLU activation.
FORtI ORN po Cy (POO
Nor m proactivt un polky
1.0
Caparityloe polky
Epiodd
107
0.
1.0
Adam aquat betas
No Interaation
Reeet Adam
Reguladte not practiivation
Rey Sartterte all preectivatioo one
Share actor ciitic Teat Lree
Fostur geN poki (PCU
Nom m pearclegroon on specyy
(CC
Spr
Eolocde
2gr.
0.6.
0.
Cerediy am
0.6
201-
0.0
wwicerer.cuoe d1ge11
SIsE Dipex12
camar copecar
Analarnont expocil
Adam equal
Reeet
Regularize prscctivation
Regularioe preectiva
erttic
2102
A12"
3.50 stol
232!
THALUIE L..k 2ekkg PCN
m pro coaattianinnn pelccy
Copacte Vos peloor
1.75
1.00
Estccde
1.pp.
0.75
Cspactt
D:25
ConDec Capol
SINESTE Dape
6:d4 mast
Figure 31: Figure I on MuJoCo Ant with the ReLU activation.
36

---

Baselines PFO
x10
x10
M.G.PhonniowS
Feature (PCA)
Norm preactivation policy
10
2.0
Episode retum
10a
1.8
2
10
Entropy
1.6
1.4
10:
2
1-
1.2
1.0
Environment steps x10
Environment steps X10
Environment steps X00
Environment X
Figure 32: PPO on ALF/Phoenix-v5 collapses with its standard tuned hyperprameters from (Schulman al., 2017) (4 epochs) when training for 200M steps. Regularizing with PFO mitigates the collapse (applied on the last pre-activation..
Seed
Seed 25
Seed
64
x10
x10
policy (PCA)
Norm prea rectiivation policy
10a
1.8
1.6
Episode return
10a
Entropy
10a
1.2
10"
1.0
0.8
0.6
Environment xtopsx10"
Environment steps x10*
Environment steps x10
Environment steps
ALE/Name ThisGame-v5 v5 showing that it also collapses with standard tuned Figure Three seeds ThinSameThssGaeee
hyperprameters from (Schulman et al., 2017) (4 epochs) when trained for long enough.
0.05
==
= 0.2
104
Norm preactivation policy
2.0
2.5
10a
Episode return
policy
1.5
102
Entropy
1.0
10a
2-
10a
0.5
0.0
0-
Environment stepsx10"
Environment
Figure 34: Varying cannot be used reliable tool to study collapse monotonic changes in yield
mon-monotonic collapse speeds, unlike when varying the number of epochs.
Separate networks
Shared
X10
x102
3.0
ALE Phoenss-v
Feature rank policy (PCA)
Feature rank cntic (PCA)
Episode return
2.0
1.8
2.0
2
2.5
Entropy
1.5
1.6
1.0
1.0
0.0
0.5
0.0
0.5
0.0
0.5
1.0
0.5 Envtronment stopsc10
Environment stapox10*
Environment stepsx10*
Enviranment stepsx
Figure 35: PPO with a shared actor-critic on ALE/Phoenix-Ve probability 0.9) with the standard hyperparameters (4 epochs). While with dense rewards, sharing the trunk was beneficial in ALE/Phocnix (Figure 21 Appendix), with the sparse reward, the opposite is true: sharing the trunk is detrimental.
sparse reward (random masking with
37

---

E
Measuring and comparing rank dynamics
Several matrix rank approximations have been used in the deep learning literature, and more specifi- cally the deep RL literature, measure the rank of the representation network. In complement to the background presented in section 2, we give here all the rank metrics we have tracked in this work and their correlations, showing that although their absolute values differ, their dynamics tend to describe the same evolution.
features learned by deep
E.1 Definitions of different rank metries
Essentially, the main difference between the rank metrics considered in the literature is whether they apply a relative thresholding of the singular values an absolute one. Their implementation can be found under src/ wecc/oooote po_dynami cs /tt/do/ol//////////eee///eeeeee/ www.eeeeeesiss./eeeeeeeeeeeeeeee in our codebase.
Referring by T the N D matrix of representations in Section 2, and letting 5 0.01 be the threshold, and (o( the singular values of + decreasing order, the different rank definitions follows.
Effective rank (Roy & Vetterli, 2007) relative measure of rank. Let H (pI..... pp) denote the Shannon entropy of a probability distribution over & events and ||o|]] be the sum of the singular
values. Let ($)
o.($) be the normalized singular values. The effective rank ||al|
(9)
This rank measure has also been used in deep learning by Huh al. (2023).
Approximate rank (PCA) relative measure of the rank. Intuitively this rank measures the number of PCA values that together explain 99% of the variance of the matrix. This can also be viewed as the lowest-rank reconstruction of the feature matrix with an error lower than 1%. also used in RL by Yang et al. (2020).
Σ -i=1
( j= (9)
(
min
>
8
srank (Kumar al., 2021) relative measure of the rank. This is relative thresholding of the singular values, similar to the approximate rank but with connection to low-rank reconstruction or variance of the feature matrix.
k E
(.( j= (2
min
- 8
Feature Rank (Lyle et al., 2022) An absolute measure of the rank. The number of singular values of the normalized & that are larger than a threshold 6.
( 0 VN
o
> 8 for iE {L. D
PyTorcch rank An absolute measure of the rank. This is the rank computed by torch. linalg....... wwammaatiirreee rank and torch. linalg ence possible between points of the data type of the singular values, i.c. for torch. float32 that is 1.19209e 7, This rank computed follows.
Let be the smallest differ-
{ o.(()) Ø1 x N
for (1, (1,..,, D3}
It also appears in Press et al. (2007) in the discussion of SVD solutions for linear least squares.
http:///ww.....oeeeaaa///acottte 4actures/12//ecture/2... http:////eee/eeeeeeeeooooooooooc htw.://teppaammttteeeee......e pdf
38

---

E.2
Correlations between the rank metrics
We compute various correlation coefficients and distance measures between the rank metries. To compute a correlation/distance on pair of rank metrics (X, Y), we take for each training run the set {(24,9)) ye)t E {0. .... } of coinciding values of the curves of the two rank metrics during the run that had 7 logged steps, compute the correlation/distance on this set, and average the correlation/distance values across considered runs. We also compute the worst correlation/distance between each rank metric pair for a worst-case analysis. We separate the average values and worst-case values by environment (ALE vs. MuloCo) for a more granular analysis. We consider all the runs without the interventions and exclude a few runs where the models collapse since the beginning of training, giving constant trivial ranks, as these result in undefined or trivial correlation coefficients.
We compute Kendall's coefficient (Kendall, 1938), Spearman's 's coefficient (Spearman, 1987), the
vTxL Pearson correlation coefficient, and normalized 1.2 computed VE,,, (=,- v.)² where L the width of the feature layer considered (i.e., 512 for ALE and 64 for MuJoCo).
Results We visualize the correlation/distance between the pairs of ranks as heatmaps annotated with averages and standard deviations. Overall, the metries are highly correlated with average correlation the coefficients varying between 0.99 and 0.51. Individually, no rank metric correlates significantly more on average with the other metries. Interestingly. from the average correlations, we clearly see two consistent clusters of stronger correlations between the relative rank metrics (approximate rank (PCA) and Effective rank (Roy & Vetterli, 2007)) and absolute rank metries (Feature Rank (Lyle et al., 2022) and PyTorch rank). The srank (Kumar et al., 2021) which is technically relative metric, but with a weak normalization rationale, correlates more with the relative metries on MuloCo with tanh but more with the absolute metrics on ALE and MuJoCo with ReLU.
0.39
400.00
041
3.30
646
0.00
0.34
MSdd
1.0
0:2
FM
LO0K
8.39
0.00
40.22
3.49
0.96
6.02
C.93
0.PL
0.54
3.30
640
5.30
Figure 36: Average correlation hetween rank metrics on MuloCo ALE.
39

---

1o00
2.50
3.70
3:50
2.77
3:67
0.22
9:90
Pyled
1:3
6.1
3.00
79%
1.00
1.00
0:00
41.00
40.10
45.00
$12.0
2.55
PSGad
3:50
3.70
69%
3:
0.58
1:00
0.00
2
Figure 37: Average correlation between rank metrics on MuloCo with the tanh activation.
1:00
2.40
1:52
0:44
P0
3:70
3:48
0.59
C:44
1.40
CA7
1:90
0.08
9:53
0.22
1:
6.
1:30
3:30
6:44
C4E
0.7
one
20.20
1:00
3:21
1.40
6.1%
1:48
6:14
9.1%
0.09
1:50
Figure 38: Average correlation between rank metrics on MuloCo with the ReLU activation.
20143
23.41
C4
C.49
0
3.00
6.64
9
317
11
1M
6:79
1.0
PStc
3.56
0/0
6.79
3.50
1.00
0:22
622
0.30
C.00
3.30
D.56
C73
0.56
6.1
Bure
3.37
2.80
0.75
4.30
2.7
3.50
GU14
6.2
1:00
6:23
0:54
450
0.00
1.00
0.35
1.39
0.10
40:08
640
0.77
0.05
3.55
Figure 39: Worst-case correlations between rank metrics on ALE.
40

---

1:20
1AB
1:T
1CB
3.50
3.80
1:00
6:54
0.22
6:36
3:30
1.40
3:90
0.22
9:84
600
1.4
PCA
9:11
L.C
$0.00
10.00
3:20
Figure 40: Worst-casc correlations between rank metrics on MuJoCo with the tanh activation.
L:N
6:24
4079
-272
0.31
30.00
4:74
2:30
4045
40.72
0.00
1:00
6:58
400
4.72
4073
3000
0054
3.34
0.08
PO
3.54
148
5.75
0.14
L.C
9.82
QUB
2.14
0.05
0.23
3:00
4012
40.15
+0.00
Figure 41: Worst-case correlations between rank metries on MuJoCo with the ReLU activation.
41

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Using the numbered contributions and claims in the introduction, we provide evidence for claims 1, 2, 3, 4, and in Sections 3.1,3.2 4, 4, and Appendix B.I, respectively.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
Limitations are discussed in Section 6.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
42

---

Answer: [Yes]
Justification: Proofs relative to Section 3.2.1 are presented in Appendix C.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We list all the experimental details in Appendix B and provide fully repro- ducible codebase.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
43

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes)
Justification:
We provide fully reproducible codebase with instructions in Appendix B.1.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: ctsps.cc/public//eeeee guwdee/eee/ee// /CodeSubmia sgionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We describe our experimental setup Section and give details in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification:
We discuss statistical significance in Appendix B.4.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
44

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We discuss computing resources in Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with NeurlPS Code of Ethics https pettps:/nneurpalr htt////eurii...pe cc/ppblic//chicssuudeeeee
Answer: [Yes]
Justification:
The research conducted in the paper conforms with the NeurIPS Code of
Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
The paper addresses fundamental algorithmic issue with no specific applica- tion. The impact of any application should be discussed specifically when introducing the application.
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
45

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The codebase implements already existing algorithms and training curves do not pose such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
All the assets used have been cited.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
46

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We introduce a codebase and run histories documented in Appendix B.1.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing or research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
47

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
48


Reviews:
[{"review_id": "OTp3AUdrHz", "reviewer": "Reviewer_qwmj", "summary": "This paper presents a series of experimental studies to diagnose the learning issues of PPO under non-stationarity in Atari-5 and MuJoCo tasks. Based on the results, this paper establishes a connection among feature rank/norm, plasticity loss and trust region violation and learning performance. To mitigate the issues in feature representation, this paper proposes Proximal Feature Optimization (PFO) and demonstrates its effectiveness in mitigating the learning issues under non-stationarity. Besides, the effects of sharing network parameters and adapting Adam are also evaluated in the same context.", "questions": "1. The epsiode return curves in Figure 24 and Figure 26 (i.e., for MuJoCo Hopper) look strange. Is there any explanation to the collapse? A regular implementation of PPO should work in Hopper when the epoch num is 10.\n2. The authors mentioned 4 MuJoCo tasks and 5 Atari games in the Experimental Setup paragraph. However, it seems only 3 Atari games (Phoenix, NameThisGame, Gravitar) and 2 MuJoCo tasks (i.e., Humanoid, Hopper) are used for the evaluation in Section 4. Are there more results on the remaining tasks/games?", "limitations": "The limitations are discussed in Section 5. However, one major limitation of this work is that the results in this paper do not provide sufficient support for the point that \"mitigating the learning issues in feature rank, plasticity loss of PPO can improve the learning performance in terms of episode return\".", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Most existing works on plasticity loss or learning under non-stationarity focus on value-based RL or off-policy AC methods. In contrast, this paper focuses on on-policy algorithms, mainly PPO, which helps the RL community gain a better understanding of this problem.\n- The experimental studies contains the results with multiple metrics, i.e., feature rank, feature norm, plasticity loss, dead neuron, excess ratio. These results will be a useful reference to the audiences.\n- This paper reveals the connection among factors like feature rank, plasticity loss, excess rate, and learning performance.", "weaknesses": "- The writing is almost clear. The empirical results and conclusions can be better organized and made prominent.\n- The proposed method PFO is closely related to DR3 (Kumar et. al, 2022) and RD [1], both of which propose regularizing the feature representation. More discussions are necessary.\n- Although PFO are demonstrated to mitigate the feature rank, plasticity loss, excess ratio in Section 4, it does not bring clear improvement in terms of learning performance like episode return (according to Figure 18 to Figure 27). This cripples the significance of PFO and also the potential causal connection between the learning issues investigated and the learning performance.\n\n---\n\nReference:\n\n[1] Reining Generalization in Offline Reinforcement Learning via Representation Distinction. NeurIPS 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">**Q1.** … PFO is closely related to DR3 and RD …. More discussions are necessary.\n\n**A1.** We thank the reviewer for pointing these out. Although **these regularizations emerge from value-based offline-RL challenges**, a discussion of their similarities with PFO can be valuable to our audience. We can add the following to our camera-ready version.\n\nOther interventions regularizing feature representations have been studied in value-based offline RL. Kumar et al. (2022) propose DR3, which counteracts an implicit regularization in TD learningby minimizing the dot product between the features of the estimated and target states.  \nMa et al. (2024) propose Representation Distinction (RD) which tries to avoid unwanted generalization by minimizing the dot product between the features of state-action pairs sampled from the learned policy and those sampled from the dataset or an OOD policy.  \n\n**Both are related to PFO as the methods directly tackle an undesired feature learning dynamic, but there is no motivation for DR3 or RD in online RL, and PFO is conceptually different.**\nThe implicit regularization that DR3 counteracts is not present in on-policy RL as shown by Kumar et al. (2022) in the SARSA experiment, and PFO differs from DR3 as it extends a trust region rather than counteracts an implicit bias. Whence the different implementations: PFO regularizes the state features between two consecutive policies, as opposed to consecutive ones of the same policy in DR3.  \nSimilarly, the overestimation studied by Ma et al. (2024) in the vicious backup-generalization cycle is broken by on-policy data. PFO’s motivation to bring the trust region to the feature space resembles RD’s motivation to bring the overestimation constraint to the feature space. However, they are again different in their implementation as RD regularizes state features between the learned policy and the dataset policy or an OOD policy.  \n**Finally, PFO is applied to the actor, while both DR3 and RD are applied to the critic.**\n\n>**Q2.** … do not provide sufficient support for the point that \"mitigating the learning issues in feature rank, plasticity loss of PPO can improve the learning performance in terms of episode return\". \n>PFO … does not bring clear improvement in terms of learning performance like episode return ... This cripples the significance of PFO and also the potential causal connection …\n\n**A2.** First, we claim that **the causal connection holds around the collapse regime**, not necessarily throughout training, and that discovering and **describing such a relation only around collapse is nontrivial.** We have clarified these arguments in GA1 and GA2 in the global rebuttal.\n\nSecond, we distinguish between consistently improving the best run performance (i.e. claiming to be “better” than PPO) and improving the aggregate performance across multiple runs and hyperparameters, e.g., robustness to collapse in our case. The limitation the reviewer raises is about the former, however this work argues about the latter.\n\nIn this sense, the **significance of PFO is realized by a) successfully mitigating collapse** as observed in Figure 6 with a significantly higher median performance indicated by the black lines, and b) improving the representation metrics and obtaining a better trust region, therefore **strengthening the relation we’ve drawn** between representation, trust region, and performance around collapse.\n\nUltimately, **this improvement can result in best performance improvement as well but less consistently as it requires running for long enough** to observe a collapse with the standard (tuned) hyperprameters. This is the case for NameThisGame in Figures 18 and 22, Humanoid in Figures 19 and 25, and Hopper in Figures 20 and 26, where PFO performs better than the baseline PPO. This is not a central claim of the paper, as it requires a larger computational budget to show for all environments.  \n**To strengthen these points, we have added Figure 34 to the rebuttal PDF** where we show that on ALE/Phoenix with the tuned standard hyperparameters, the agent collapses when training for longer and this is mitigated with PFO, which attains the best performance in that case.\n\n>**Q3.** The epsiode return curves … for MuJoCo Hopper look strange.  Is there any explanation to the collapse? A regular implementation of PPO should work in Hopper when the epoch num is 10.\n\n**A3.** We use the same hyperparameters as the original PPO implementation, which are also the default ones in popular codebases (as noted in lines 140-143).  \nOur setting differs from a “regular” implementation in the following main points, which can help explain why a collapse is not typically observed in previous work.\n1.  We train for longer and collapse always happens after the default 1M steps. Dohare et al. (2023) also observe collapse when training for longer.\n2.  We do not decay the learning rate (line 154) we are interested in agents that can ideally train indefinitely.\n3.  we parameterize the action space by a TanhNormal distribution with both state-dependent mean and variance following the implementation of Haarnoja et al. (2018). (lines 614-616).\n    \n>**Q4.** … 4 MuJoCo tasks and 5 Atari games in the Experimental Setup ... However … only 3 Atari games … and 2 MuJoCo tasks … used in Section 4. Are there more results on the remaining tasks/games?\n\n**A4**. Indeed, as noted in lines 305-306, while we used 4 MuJoCo tasks and 5 Atari games to demonstrate the collapse phenomenon with enough evidence, we did not use all of the environments for the evaluation.\nWe selected the environments where the collapse was the most consistent to test multiple interventions while maintaining a reasonable compute budget (this is common practice like in Gulcehre et al. 2023 and Kumar et al. 2021). We believe this to be a right tradeoff between claiming that interventions are necessary to prevent collapse and providing enough insights about several types of interventions.", "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": "glWBnQT6F8", "reviewer": "Reviewer_YzGo", "summary": "The paper investigates the phenomenon of loss of trainability and performance collapse with PPO agents in standard RL benchmarks. The same phenonema from other settings is found to hold with PPO agents and, moreover, increasing the number of training epochs exacerbates the effect. \nFurther investigation finds that PPO's trust-region constraint is not maintained when the feature rank decreases and a mechanism for this finding is explained. \nThe authors propose a regularizer, PFO, which softly constrains the change in features is shown to mitigate trainability loss throughout the learning process.", "questions": "I would be willing to revise my score based on answers to the following.\n\n_Clarification questions_\n- Line 308. Could you clarify how the average probability ratio is computed?\n\tAs written, my interpretation is that, given some window of updates, you consider all ratios above $1+\\epsilon$ and take their mean. Then, you do the same for the ratios less than $1-\\epsilon$. Finally you divide the first mean ratio by the second mean ratio. Is this correct?\n\tIf so, this seems potentially misleading because the means are computed conditionally, considering only ratios that are above 9or below) the clipping ratios. Then, we cannot tell what the mean ratio is overall, and whether there are many violations of the trust-region occurring. Moreover, ratios can be misleading since they may overemphasize small probabilities in the denominator. \n\n\tI would suggest also reporting how often the clipping criteria is violated. E.g. compute on collected states which proportion of them have a ratio outside the clipping region after a round of updates on a batch of trajectories.\n\tPerhaps it would also be useful to consider ratios less than 1 or greater than 1 seperately instead of computing a ratio.\n\n\n- Line 164 Using the term \"plasticity loss\" as a metric is confusing because it also refers to the overall phenomena of neural networks being unable to learn additional tasks after training on other tasks. \nDoes this term refer to the capacity to fit random labels? In that case, I would use a term that is more descriptive e.g. random label trainability (RLT). \nAlso, how are these random labels generated in this context? In the original paper, a regression task was made since a value-based algorithm was used. With a policy-based algorithm, how is the plasticity evaluation task set up?\n\n\n- Fig.1 Why does the preactivation norm plateau at 10^4? \n\n- Which data is used to produce the box plots? The caption for Fig. 6 mentions \"A boxplot includes 15 runs with the different epochs\". Does this mean the box plot conains data across training times and epochs? \n\n- In the box plots, how are outliers determined? We can see they lie outside the interquartile range indicated by the \"whiskers\" but what additional criteria is there? \n\n\n- How is the PFO regularizer related to adding a regularizer in parameter space? Say $\\ell_2$ regularization between the current parameters and the previous parameters? Basically, this seems similar in spirit to adding a target network to value-based methods. \n\n\n- Line 94-99 (and elsewhere) \"PPO introduces additional non-stationarity...\": I do not think it would be fair to see PPO has more nonstationarity than other policy gradient methods. REINFORCE and PPO both have nonstationarity due to the changing policies. REINFORCE can also be interpreted as a trust-region-like method. Since REINFORCE is essentially gradient ascent on the expected return objective, we can view it as minimizing a linear approximation with a quadratic constraint (see paragraph containing eq. 5.5 in [4]). In this light, by doing a single optimization step before changing the surrogate objective, we could say REINFORCE has _more_ nonstationarity.\n\n\n- Line 257. \"The underlying assumption...approximately orthogonal...\" Could you give some more support to this claim?\n\tAs mentioned above, with state aggregation, this clipping strategy could be effective without having orthogonal features. \n\n\n- Line 8-9: \"...policy optimization methods which are often thought capable of training indefinitely\". I would disagree with this claim. Deep RL algorithms have often be thought to be high-variance and unstable so I do not think it would be common to believe they can train indefinitely. \nI suggest replacing this phrase and simply mention the phenonmenon has been underexplored in the context of on-policy policy optimization.\n\n\n- The parameter norm has been linked to plasticity/trainability loss in neural networks. Is there a reason why the paper measures feature norms instead? \n\n\n\n\n\n_Broader questions_\n- The performance degradation from training more epochs could be attributed to difficulties of offline RL. As the number of epochs is increased, the more the procedure ressembles offlne RL---a setting in which standard RL algorithms are not suitable for. What do you think of this?\n\n\n- I find the proposed solution (Proximal Feature Optimization) to be slightly unsatisfying, Generally, I can see how using regularization on the policy and features can mitigate the performance collapse but would this be simply delaying the problem? If we train for longer, would we still expect to see the same performance degradation? \n\n\n- Interestingly, the value network is not the problem. It is the policy network that suffers from representation collapse. Do you have any thoughts about this? I think this could be an interesting avenue of investigation.\n\n\n_Suggestions_\n\n- I would suggest removing or summarizing the paragraph at line 193 into a single sentence. \n\tIt is clear that the policy would be the same for all states if the features are constant. This would imply the policy is uniform. Currently the paragraph seems to explain this in a roundabout manner when a concise explanation suffices. \n\n- A suggestion would be to try the Atari environemnts identified in [1] to benefit from plasticity injection. On many Atari games, little or no benefit was found, so it may be more meaningful to focus on those where there is more evidence for trainability loss.  \n\n\n- I would be careful about mentioning that low feature rank necessarily means that the representation is poor. Successful training of deep neural networks often involves a reduction in feature rank and the relationship between rank and performance can be complex [3].\n\n\n- As a sidenote: even if there is a low dimensional representation, this does not necessarily imply the clipping criteria will be violated. For example, if we use state aggregation, then multiple states are clustered into one aggregate state. Then, the policy will be exactly the same for all these states, so if the policy exits the clipping region for one those states, there will no longer be any incentive to update the policy any further on any of them. \n\n\n- A minor suggestion is to use stable rank [2] as a measure of the rank instead of effective rank since it is a continuous quantity in the singular values. \n\n\n\n\n[1] \"Deep Reinforcement Learning with Plasticity Injection\" Nikishin et al. \n\n[2] https://nickhar.wordpress.com/2012/02/29/lecture-15-low-rank-approximation-of-matrices/\n\n[3] \"An Empirical Study of Implicit Regularization in Deep Offline RL\" Gulcehre et al. \n\n[4] https://www.stat.cmu.edu/~ryantibs/convexopt-S15/scribes/05-grad-descent-scribed.pdf", "limitations": "These are discussed.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper investigates the topic of trainability loss in a less-explored context, on-policy policy optimization, and focusing on a popular algorithm, PPO. Better understanding this algorithm and how it interacts with neural networks could lead to widespread impact as the research direction is developed.\n\nA wide variety of experiments are conducted with some interesting findings. To me, the most intriguing insight is the connection between low feature ranks and increased violations of PPO's trust region. The investigation is fairly thorough, examining various metrics in tandem throughout the training process. Experiments seem to be conducted and reported in a reasonable fashion.\nI appreciate the simple, clear example to demonstrate how a poor representation may lead to the clipping constraint being ineffective (line 270-...).", "weaknesses": "Currently, the paper feels a bit disorganized and it is a bit difficult to follow the train of thought. There are many different experiments done and placing more focus on the most important could help streamline the paper. For example, the section at line 193 feels a bit out of place since this line of investigation is not mentioned earlier. Another example is that the research questions Q1. and Q2. do not mention the trust-region although it is a large part of the paper.\n\nAdditionally, various phrases (or claims) in the paper could use more support or evidence. These are discussed more specifically in the Questions section. \n\nI am also concerned about the link between feature rank and trust-region violation, a central point of the paper. \nIn Fig.4, the correlation between the prob ratios and dead neurons or prob ratios and feature rank seem weak. There is only a substantial dip in the prob ratios once the feature rank is near zero or the number of dead units is very high. In other words, only when the agent is doing very poorly can we see a strong relationship, which may limit the applicability of the finding. Ideally, one would avoid this regime since performance is poor to start with.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We address the key points in the rebuttal and the remaining ones in a comment.\n\n>**Q1.** line 193 feels a bit out of place … investigation is not mentioned earlier … suggest removing or summarizing the paragraph …\n\n**A1. This information is relevant to practitioners in diagnosing the kind of collapse they face.** One common collapse or stagnation is due to entropy collapse and lack of exploration. In this work, we highlight that when collapse is associated with a high entropy, it may likely be due to collapsed representations. **We agree with the suggestion of the reviewer and will summarize this paragraph in a sentence.**\n\n> **Q2.** … research questions Q1. and Q2. do not mention the trust-region although it is a large part of the paper.\n\n**A2.** Q1 and Q2 serve as motivation for section 3.1, where they are answered. These are not the only questions addressed in the work. The trust-region exposition in section 3.2 is only motivated after answering Q1 and Q2 and observing collapse.  \nWe will move Q1 and Q2 inside section 3.1 to clarify their scope and reformulate section 3.2 with two questions, Q3 and Q4, to mirror 3.1 and ensure the question format covers all the main points of our work.\n\n>**Q3.** I would be careful about mentioning that low feature rank necessarily means that the representation is poor …  I am also concerned about the link between feature rank and trust-region violation …\n\n**A3.** Reviewer XWGd has also raised this concern. We have clarified this in GA1 and GA2 of our global rebuttal.\n\n> **Q4.a** how the average probability ratio is computed? …\n\n**A4.a** Yes, this is correct, and more details can be found in the Appendix lines 635-640.\n\n> **Q4.b** If so, this seems potentially misleading …\n\n**A4.b** The mean ratio without conditioning does not give information about learning or the trust region. The PPO-Clip objective increases the ratios of actions with positive advantages and decreases those with negative advantages until they reach the clip limit. Therefore, to have a signal for learning, we have to at least condition on the sign of the advantage. Still, a takeaway from the toy example is that taking the mean ratio of actions with, say, a positive advantage under bad representations would mix the ratios that suffer from interference (Figure 5 right) and would give a misleading average.\n\nTherefore, we found that the right way to quantify the trust-region violation is to condition on the violation itself (e.g., above $1+\\epsilon$). We took the mean rather than the count because fully optimizing the PPO objective should push the actions until the clip limit, so it's not clear whether a high clip count would be worse than a lower one.\n\n>**Q4.c** they may overemphasize small probabilities in the denominator.\n\n**A4.c** This is true but has not been a critical issue for our analysis. In Figure 4 with ALE the smallest probability ratio is around 0.4, and in Figure 6, the largest excess ratios are between 1.5 and 2.5.\n\n>**Q4.d** I would suggest also reporting how often the clipping criteria are violated …\n\n**A4.d** Yes, we track this. As discussed above (A4.b) it’s not clear what the count of the clip fraction should be and it is not an interesting quantity for our study.\n\n>**Q4.e** consider ratios less than 1 or greater than 1 seperately\n\n**A4.e** We do separate them in Figure 4, where we look at the ratios below $1-\\epsilon$ and isolate the causal relation around poor representations. For Figure 6, aggregating them into a ratio provides a better summary with larger plots.\n\n>**Q8.** How is the PFO regularizer related to adding a regularizer in parameter space? …\n\n**A8.** We can construct a spectrum ranging from regularizing the parameters to the outputs. The lower end regularizes the parameters like an L2 weight difference. The higher end regularizes the network's output like PPO clipping and almost like a target network in value-based methods. **PFO sits in the middle of this spectrum, where the regularization allows both the network's weights and final outputs to change without explicit constraints while maintaining a regularization of the feature space.**\nAlso, PFO should mostly be seen combined with PPO, to extend the trust region to the feature space.\n\n>**Q12.** Is there a reason why the paper measures feature norms instead?\n\n**A12.** Other work we cite, like Lyle et al. (2024), has also looked at feature norms in addition to weight norms. We noticed that the model weights were consistently and steadily increasing regardless of collapse; however, the feature norm showed a sudden jump around collapse. **We found this to be a more apparent symptom that researchers and practitioners would want to investigate further.**\n\n>**Q13.** degradation from training more epochs could be attributed to difficulties of offline RL\n\n**A13.** Off-policiness could be a good characterization here as the algorithm has access to the action probabilities of the collected data and the policy that collected it more generally.\nIn PPO the more epochs performed, the more the on-policy approximation becomes off-policy. However, the PPO-Clip objective is supposed to be robust to slight changes in the number of epochs, as it only depends on epsilon. Therefore, **we view observing a drastic collapse by moving from four to six epochs in Figure 1 not as a failure of standard RL algorithms in the off-policy setting but as a failure of the learning dynamics of PPO-Clip.**\n\n> **Q14.** I find the proposed solution PFO to be slightly unsatisfying, …\n\n**A14.** The first purpose of PFO in this work, as the reviewer mentions, is to show that explicitly regularizing the feature dynamics builds robustness to trust-region and performance collapse.\nNow as with most empirical studies and regularizations using a coefficient, there is no guarantee on how long PFO will stand the test of time, but we hope that our paper encourages further theoretical studies and only then will we have provably mitigated the problem.", "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": "wOCj8qU0b9", "reviewer": "Reviewer_XWGd", "summary": "This work provides an empirical study of the feature rank deterioration and loss of plasticity of the Proximal Policy Optimization (PPO) algorithm on Atari and Mujoco tasks. Then links the deterioration of the performance to representation collapse and hence the break of the trust region. From there, the authors propose an auxiliary loss to PPO to maintain the representation rank.", "questions": "Could the authors address my questions in the weakness section?\n\nFrom Figure 6, does it mean that the proposed method does not consistently improve the loss of plasticity?\n\nIf the main goal is to prevent the preactivation norm from going larger, why not just regularize it to be smaller?", "limitations": "Ues", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The study is very interesting and brings some new insight into how PPO works in practice. The experiments are conducted in clear logic and the analysis and observations are all novel and interesting.", "weaknesses": "Figures can be plotted with better quality, don't overlap the labels (in Figure 3), and maybe set the titles to the quantity of interest. \n\nI understand that several works are measuring the feature rank at the last policy/critic layer. But why that is sufficient evidence that the policy is losing rank? For instance, if my action space is binary, then the policy's last layer could learn to be rank-2. More information might be stored in previous layers. From another perspective, isn't it expected that the policy learns to lose rank, by compressing state information to only correlate with actions?\n\nThe proposed Proximal Feature Optimization objective is a bit confusing. First it should be a $\\ell_2$ norm? What does it mean by $(\\phi - \\phi_\\text{old})^2$. Secondly, from a high-level, isn't this just a regularization loss as in continual learning that prevents forgetting? But based on my knowledge it will exacerbate the the loss of plasticity. If you want to preserve rank, should not you add some reconstruction loss?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">**Q1.a** Figures can be plotted with better quality, don't overlap the labels (in Figure 3),\n\n**A1.a** Indeed. We will correct this in the camera-ready version.\n\n>**Q1.b** and maybe set the titles to the quantity of interest.\n\n**A1b.** Can the reviewer please elaborate on this or give an example? What is meant by title and quantity of interest?\n\n>**Q2.a** isn't it expected that the policy learns to lose rank, by compressing state information to only correlate with actions?\n\n**A2.a** **Indeed, however, we distinguish between a low rank that’s beneficial and an extremely low rank that is detrimental. We have clarified this in answer GA1 in the global rebuttal.**\n\n>**Q2.b** I understand that several works are measuring the feature rank at the last policy/critic layer. But why that is sufficient evidence that the policy is losing rank? For instance, if my action space is binary, then the policy's last layer could learn to be rank-2. More information might be stored in previous layers.\n\n**A2.b** **The penultimate layer serves as a bottleneck** for decoding the actions or the value, so when this layer reaches an extremely low rank (as described in the previous answer, to distinguish for a beneficial low rank), this becomes irrecoverable. This also strongly correlates with plasticity loss, as this layer serves as a bottleneck for reconstructing information from the input state.\n\n>**Q3.a** The proposed Proximal Feature Optimization objective is a bit confusing. First it should be a l2 norm? What does it mean by (phi(s) - phi_old(s))^2.\n\n**A3.a** Yes, we thank for pointing this out. $\\phi(s)$ is a vector containing the pre-activations. \nThe square of the difference should be more clearly written as the squared norm of the difference between the vectors: $\\|\\| \\phi_\\theta(S_t) - \\phi_\\text{old}(S_t)\\|\\|_2^2$.\n\n>**Q3.b** Secondly, from a high-level, isn't this just a regularization loss as in continual learning that prevents forgetting? But based on my knowledge it will exacerbate the the loss of plasticity. If you want to preserve rank, should not you add some reconstruction loss?\n\n**A3.b** **PFO regularizes the change in features at every batch, unlike methods that prevent forgetting, which regularize the parameters towards a fixed optimum of a previous task.**  \nIn this sense, PFO is sought to be an extension of the clipping trust region applied in the output space to a trust region in the feature space. It comes from our observation that the clipping is not enough to satisfy the trust region.  \nIt also allows us to address the undesired symptoms we observed (high feature norm and aliased features).  \nOther auxiliary losses, such as a reconstruction loss, can also be beneficial in this case but do not necessarily constrain or regularize the step size.  \nFinally, to support our claim that regularizing representations would prevent trust region and performance collapse, we found that using a regularization that directly targets one of the undesired symptoms we observed would make the connection more straightforward.\n\n>**Q4.** From Figure 6, does it mean that the proposed method does not consistently improve the loss of plasticity?\n\n**A4.** **It does consistently improve the plasticity loss.** This can not show in the the figure when the non-regularized model collapses too early. This is described in the caption of Figure 18: \"the tails of the plasticity loss on Phoenix with interventions can be higher than without interventions on the runs where the models collapse too early without interventions, leading to the plasticity loss of the non-collapsed models with interventions eventually becoming higher.\"\n\n>**Q5.** If the main goal is to prevent the preactivation norm from going larger, why not just regularize it to be smaller?\n\n**A5.** The increasing preactivation norm phenomenon can be seen as an effect of a bigger issue: the failure of the clipping trust region imposed on the output to ensure a trust region on the rest of the network. Therefore,  as mentioned in A3.b, PFO is sought to extend this trust region. Analogously to constraining ratios between the previous and current policy, it regularizes the difference in features between the previous and current policy.  \n**This allows PFO to address both the feature norm and the trust region issue. This is described in lines (319-322), which we will rephrase to highlight this point more.**", "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": "zZhz6iznwI", "reviewer": "Reviewer_6zAY", "summary": "The paper addresses non-stationarity in RL and its impact on deep learning networks, focusing on PPO. It identifies that networks in PPO, like those in off-policy methods, suffer from representation rank deterioration, leading to performance collapse. The authors propose Proximal Feature Optimization (PFO), a new regularization on the policy representation that regularizes preactivation changes, which can mitigate performance collapse and improve the agent’s performance.", "questions": "- Please provide a more detailed explanation on Eq.2. I don't understand the meaning of $\\phi$, and should the formula be $s_t$ instead of $S_t$?\n- In section 4, the authors introduce an intervention that allows the actor and critic to share all the layers. The experiment results illustrate that this can bring some improvements but not significantly. I am more curious about whether the reason for this is the sparsity of reward or the environment itself. I think the latter is more likely, and we can verify this by setting up a control experiment with different reward functions in the same environment.", "limitations": "N/A", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The work exhibits a coherent structure, and the author adeptly elucidates the underlying motivation. The author adeptly elucidates the shortcomings of PPO and offers enough reasons. This paper is well-written.\n- Section 3's Q1 and Q2 effectively describe the potential problems of PPO, and I concur with the underlying motivation.\n- The paper not only identifies problems but also suggests practical interventions, such as PFO and sharing the actor-critic trunk, to regularize representations and reduce non-stationarity, showing promising improvements in performance.\n- The authors provide details of the experiment design, and it's clear.", "weaknesses": "The authors claim to open data and code; however, I could not locate them. Therefore, I apologize if I overlooked their presence. Refer to the question section for other weaknesses.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">**Q1.** The authors claim to open data and code; however, I could not locate them. Therefore, I apologize if I overlooked their presence.\n\n**A1.** **Yes, these are available on an anonymous GitHub repo mentioned in Appendix line 578.** The GitHub repo contains further links to the Weights&Biases project with all the interactive training curves and links to download the raw logs.\n\n>**Q2.** Please provide a more detailed explanation on Eq.2. I don't understand the meaning of phi, and should the formula be s_t instead of S_t?\n\n**A2.** $\\phi(S_t)$ refers to the vector of the preactivations of the penultimate layer of the network on the state $S_t$. As PFO can also take all the pre-activations in the network, $\\phi(S_t)$ could also be the concatenation of the preactivations of all the layers.  \nWe will add these details before the equation in the camera-ready version.\n(This is analogous to $\\Phi$ defined in the background (lines 100-104) as the feature matrix, consisting of the activations of the last hidden layer.)\n$S_t$ refers to the state that we capitalize as it is a random variable in the expectation over trajectories, like all other random variables in the background (lines 58 - 66) This is the notation from Sutton and Barto (2018).\n\nTherefore, the equation is essentially an L2 norm between the features of the same state at the old policy and the learned one: $\\|\\| \\phi_\\theta(S_t) - \\phi_\\text{old}(S_t)\\|\\|_2^2$.\n\n>**Q3.** In section 4, the authors introduce an intervention that allows the actor and critic to share all the layers. The experiment results illustrate that this can bring some improvements but not significantly. I am more curious about whether the reason for this is the sparsity of reward or the environment itself. I think the latter is more likely, and we can verify this by setting up a control experiment with different reward functions in the same environment.\n\n**A3.** This is a good observation! The representation of the critic plays a key role in the reasoning. When the critic, if trained separately, is subject to a degraded representation (or collapse) with a very low rank and many dead units, we observed that sharing its representation with the policy makes the policy collapse faster.\n**When the environments are similar** – we consider two ALE environments, e.g., Phoenix and Gravitar, to be similar because of the observation space, etc – **the main difference in making a critic collapse or not is sparsity of the reward** as observed by Lyle et al. (2022). **Therefore in this case, it is fair to say that the reason for the collapse is mainly due to reward sparsity.**\n\nNevertheless, one needs to be careful when comparing environments when more than one degree of freedom changes; for example, this would make comparisons of experimental results obtained on Mujoco and ALE environments harder to attribute to sparsity only.\n\n**Per the reviewer's request we ran an experiment on ALE/Phoenix (a dense reward environment), with a reward mask** randomly masking a reward with 90% chance, and compared the effects of sharing the actor-critic trunk. The results are in the additional PDF of the general rebuttal. **As expected, while with dense rewards, sharing the trunk was beneficial in ALE/Phoenix (Figure 21 Appendix), with the sparse reward, the opposite is true: sharing the trunk is detrimental (Figure 37 Rebuttal).** This confirms our conclusion. We thank the reviewer for suggesting the experiment, we will add it to the Appendix of the paper. We hope this strengthens the reviewer's support for the paper.", "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": "kiLWkRKb96", "reviewer": "Reviewer_v7XX", "summary": "The paper examines the loss of plasticity and its connection representation collapse in policy networks in online reinforcement learning (as opposed to previously studied value networks in offline reinforcement learning). The paper establishes the problem in the Atari game domain including the growth in the norm of feature pre-activations of the policy network. The analysis isolates and illustrates the problem of increasing representation state collinearity and representation collapse in the trust region of a policy in a toy domain. To mitigate representation collapse the paper proposes a feature regularization loss term based on these pre-activation features of the policy network.\n\nThe evaluation compares methods to mitigate policy network representation collapse including the new loss, sharing the actor and critic feature trunk, and modifying the Adam optimizer (by resetting moments or faster decay of second moments) in the ALE and MuJoCo environments. The results show the loss has does not impact episode returns but does mitigate loss of plasticity.", "questions": "- line 143: \"do not include normalization layers\"\n\t- Does this help plasticity preservation? (even if only to cite prior work on the effects)\n- lines 151-153: \"We vary the number of epochs as a way to control the effect of non-stationarity, which gives the agent a more significant number of optimization steps per rollout while not changing the optimal target it can reach due to clipping, as opposed to changing the value of ε in the trust region for example.\"\n\t- It might be interesting to contrast the effects of varying $\\epsilon$ as well.\n- line 280 (minor): $\\pi_\\theta(a_i | s)$\n\t- $a_i$ should be $a_1$ in this case\n- The toy setting was good demonstration of the problems with trust region constraints. It may help to introduce this sooner as I was confused about the trust region claims being made until that point.\n- Figure 6 (minor)\n\t- The caption only mentions \"Top\" and \"Bottom\", omitting the middle row of \"NameThisGame\".", "limitations": "Yes. The paper acknowledges remaining opportunities to relate the findings on representation collapse to other research on plasticity loss. And to incorporate more complex model architectures (like memory).\n\nThere is no need for additional comments on potential societal impacts.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "# originality\nGood\n- Investigating plasticity loss in offline RL is established, but has not been extended to online RL.\n- The analysis provides new insights on trust region optimization.\n\n# quality\nGood\n- The toy domain isolates the effect of interest to strengthen the empirical results and motivate the new regularization term.\n- the experiments show modest evidence for improvements and report variation in results over runs.\n\n# clarity\nMostly good\n- The exposition on trust region weaknesses and collapse could be clearer and make the main points sooner. The rest of the text was good.\n\n# significance\nGood\n- Loss of plasticity is an important problem, particularly in the case of online RL.\n- The paper will be of interest to researchers in continual learning and RL communities; both theoretical (for the trust region points) and practical (for the regularization term) researchers.", "weaknesses": "The paper would benefit from outlining the trust region insight sooner as it is a core idea that only gets explained in detail on page 7 of the paper.\n\nThe results on the regularization term are a bit confusing and mixed. The narrative would benefit from clarifying the aspects that remain unresolved or ambiguous that merit further attention. The Gravitar analysis of sparse reward environments was helpful. But I was not really sure what to make of the effects on episode returns compared to representation collapse. Some way of teasing these effects out more cleanly would help.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">**Q1.** The toy setting was good demonstration of the problems with trust region constraints. It may help to introduce this sooner … The paper would benefit from outlining the trust region insight sooner as it is a core idea that only gets explained in detail on page 7 of the paper.\n\n**A1.** We thank the reviewer for the recommendation. This is indeed a key point of interest to the readers. **We will update the introduction section to describe the insights sooner and point to the theoretical derivation for interested readers.**\n\n>**Q2.** The results on the regularization term are a bit confusing and mixed. The narrative would benefit from clarifying the aspects that remain unresolved or ambiguous that merit further attention. The Gravitar analysis of sparse reward environments was helpful. But I was not really sure what to make of the effects on episode returns compared to representation collapse. Some way of teasing these effects out more cleanly would help.\n\n**A2.** Figure 6 presents summaries of runs at the end of training aggregating runs with different numbers of epochs. It, therefore, shows the distribution of returns and representation metrics across different hyperparameters.  \nWith this, when a method like PFO presents an improvement of the lower tail of episodic returns and, at the same time, an improvement of the representation metrics (no extremely poor ranks or fully dead networks), we can say that **the method improves robustness to representation collapse and performance collapse, and strengthens the link we’ve drawn between the two in the earlier sections.   In this sense, the effects on episode returns and representation collapse should be appreciated in tandem.**  \nIt should also be noted that the strong relation between them mainly holds around the collapsing regime, as described in Figure 4 and explained in our global rebuttal answer GA2.\n\n>**Q3.** line 143: \"do not include normalization layers\". Does this help plasticity preservation? (even if only to cite prior work on the effects)\n\n**A3.** Lyle et al. (2023) conducted a study with non-stationary supervised learning and value-based methods and found that normalization layers do provide improvements to plasticity. In the scope of this work, we did not look at interventions that change the “degree” of non-stationarity in the input or output of the network, such as layer normalization, batch normalization, and running observation or reward normalization, as these would require a different analysis. As stated in our discussion and limitation plan to address these in future work.\n\n>**Q4.**  lines 151-153: \"We vary the number of epochs as a way to control the effect of non-stationarity, which gives the agent a more significant number of optimization steps per rollout while not changing the optimal target it can reach due to clipping, as opposed to changing the value of ε in the trust region for example.\" It might be interesting to contrast the effects of varying epsilon as well.\n\n**A4.** As mentioned, varying epsilon would change the optimum of the surrogate objective so it’s not exactly clear if we can use this as a tool to investigate the relation between non-stationarity and feature collapse, like we did the number of epochs.  \n**Nevertheless, per the reviewer’s request, we rerun Figure 1 with ALE/Phoenix with different values of epsilon** and the baseline number of epochs of 4. As expected, and as shown in Figure 36 of the additional rebuttal PDF, **we observe no apparent correlation with the time of collapse as both doubling (epsilon =0.2) and halving (epsilon =0.05) the baseline epsilon (epsilon =0.1) can yield training curves with a delayed representation and performance collapse.**\n\n>**Q5.** line 280 (minor): a_i should be a_1 in this case\n\n**A5.** Indeed. We thank the reviewer for the detailed feedback!\n\n>**Q6.** Figure 6 (minor). The caption only mentions \"Top\" and \"Bottom\", omitting the middle row of \"NameThisGame\".\n\n**A6.** Top and middle share the same properties. We will add this to the caption.", "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": "dxmCMF48JC", "reviewer": "Reviewer_4jUo", "summary": "The paper investigates the phenomenon of feature collapse in the popular on-policy RL algorithm PPO. It demonstrates that increasing the sample reuse (number of epochs) in PPO deteriorates the feature rank and plasticity. It also finds that the clipping operation for the conservative update in PPO does not prevent this feature collapse. Furthermore, it shows that feature collapse leads to the degradation of the trust region, and vice versa. Finally, the paper proposes several techniques to prevent feature collapse, including feature regularization, network sharing, and Adam hyperparameter adjustments, though these techniques do not fully resolve the issue.", "questions": "- In Section 3.1 and Figures 7, 8, and 9, there appears to be no clear correlation between the number of epochs and the plasticity loss of the critic network, which is counterintuitive since it contrasts with findings from value-based plasticity studies. Could you explain more about this?\n- In Section 3.2, you state that the clipped sample’s ratio will have their probabilities continue to go beyond the clip limit in PPO. How about replacing the clipping operation with a KL penalty? This technique was proposed in the original PPO paper [1], and it has shown no significant difference in performance compared to the clipping operation in some environments [2]. What impact do you think replacing the clipping operation with a KL penalty would have on feature collapse?\n- Regarding PFO, instead of applying regularization in the feature space to ensure the current representation does not deviate from previous representations, how about applying regularization in the weight space using EWMA update?\n- It would be beneficial to investigate Phasic Policy Gradient (PPG) [2]. It trains the actor network to predict the value of the old policy, which is similar to PFO. It would be interesting to analyze feature collapse in PPG to see if similar patterns emerge as observed in PFO.\n\n[1] John Schulman et al., Proximal Policy Optimization Algorithms, arXiv 2017 \\\n[2] Karl Cobbe et al., Phasic Policy Gradient, ICML 2020", "limitations": "N/A", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- To the best of my knowledge, this is the first work to investigate the phenomenon of feature collapse in on-policy algorithms.\n- The experimental design is strong. They appropriately select benchmarks: Atari for discrete action spaces and MuJoCo for continuous action spaces. Additionally, the experiments are extensive and well-conducted.\n- The connection between feature collapse and trust region is intriguing and distinguishes this work from previous studies on feature collapse in value-based methods.", "weaknesses": "- They demonstrated that training policy and value networks for more epochs leads to feature collapse, but this is not typical in PPO settings. It is uncommon to see PPO run for 6 or 8 epochs (Figure 1, 2, and 3).\n- While they attempt to explain feature collapse in the critic networks, the explanation is insufficient and requires further elaboration  (Section 3.1).\n- While they emphasize that the focus of the paper is not on improving the performance of PPO, they do not provide a clear solution to mitigate feature collapse in PPO (Figure 6).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": ">**Q1.** … leads to feature collapse, but … it is uncommon to see PPO run for 6 or 8 epochs (Figure 1, 2, and 3).\n\n**A1.** **The collapse phenomenon we show is certainly not exclusive to running with more epochs. The phenomenon does happen with the “standard” hyperparameters of the environments we study.**\nThis is the case on MuJoCo, and on ALE in Figures 1, 2, and 3, mentioned by the reviewer, we see the start of the collapse at 100M steps, which motivates us to increase the number of epochs to accelerate the collapse, and allows us to show that the phenomenon is due to non-stationarity and the inability of the policy to adapt to new targets.  \nIn this sense, **increasing the number of epochs is a tool in our analysis and is not a condition for collapse.**  \nTo observe collapse with the “standard” hyperparameters in Figures 1, 2, and 3 it suffices to train for longer. **We show this on ALE/Phoenix and ALE/NameThisGame with 200M steps in Figures 34 and 35 in the PDF of the rebuttal.**\n\n>**Q2.** While they attempt to explain feature collapse in the critic, the explanation is insufficient\n\n**A2.** **We do not attempt to explain the collapse of the critic network (see lines 147-149)** which has been studied in previous works, e.g., Kumar et al., 2022 through implicit regularization and Lyle et al., 2022 who argue that sparser targets are more detrimental.  \n**We make a few observations** about the representation dynamics of the critic network **in relation to the sparsity of the environment** (lines 178-184), which then allow us to draw conclusions about sharing an actor-critic trunk in section 4 (lines 339-355).\n\nIf the reviewer meant the actor network, then we do describe its features and dynamics and connect them to non-stationarity. By showing that this collapse also holds in on-policy policy optimization in addition to offline, off-policy, and value-based RL as shown by previous work, we strengthen the hypothesis of non-stationarity being a major cause of the phenomenon.\n\nUltimately, discovering deeper causes beyond non-stationarity remains an open question as discussed in our conclusion (389-391).\n\n>**Q3.** While they emphasize that the focus of the paper is not on improving the performance of PPO, they do not provide a clear solution to mitigate feature collapse in PPO (Figure 6).\n\n**A3.** In this work, we study interventions that have been shown in other settings to help and present their effects on the representation dynamics of PPO. This serves two purposes. 1) **showing which ones can help** and under which conditions (which environment features), and 2) when they do help and improve the trust region at the same time, **strengthen the relation we have drawn between representation collapse and trust-region collapse.**\nWe have shown that sharing the actor-critic network when the environment has a rich reward signal does mitigate collapse (Figure 6, top and middle), and applying PFO also mitigates collapse when sharing the actor-critic trunk does not (Figure 6, 19, 20).\n\nHowever, like any other empirical study, and as mentioned in our limitations, there is no guarantee on how well our results would generalize to other settings. **We hope that our paper encourages further theoretical studies, and only then will we have provably or “clearly” mitigated the problem.**\n\n>**Q4.** … no clear correlation between the number of epochs and the plasticity loss of the critic network … it contrasts with findings from value-based plasticity studies.\n\n**A4.** **Runs with higher epochs tend to have higher plasticity loss which can be explained by a stronger overfitting to previous targets and is consistent with the findings from value-based methods.**\n\n**There are some exceptions,** which have probably raised the reviewer’s attention, which can be explained by the behavior of the policy. When the policy collapses earlier in a run with a higher number of epochs, the critic's plasticity loss is prematurely stopped (there are no more changing objectives) and ends up lower than in a run with a lower number of epochs which has continued training until the end.\n\n**This is the case of Figure 7 and is explained in its caption.**\n\n>**Q5.** How about replacing the clipping operation with a KL penalty? …\n\n**A5.** We have also performed early experiments with PPO-KL. However, the adaptive coefficient of the KL trust region makes the algorithm extremely high variance and often collapses due to the coefficient exploding or collapsing independently of feature dynamics.\n\nNevertheless, for the runs where the coefficient was stable, **we did indeed observe the same behavior as PPO-Clip**, specifically that the features became more aliased until they collapsed, followed by a trust region collapse characterized by a blow-up of the KL divergence.\n\n>**Q6.** ...how about applying regularization in the weight space using EWMA update?\n\n**A6.** PFO doesn’t stem from the sole motivation of regularizing the representations, it also seeks to do so by extending the trust region set by the clipping in the output space to an additional regularizer in the feature space.\n\nThe idea is to not impose an explicit constraint or a regularization on the weights directly to allow them to adapt to the moving targets but to regularize their dynamics through the features and output trust regions.\n\n>**Q7.** It would be beneficial to investigate Phasic Policy Gradient (PPG) [2]\n\n**A7.** This is a great suggestion. We would put PPG in the same box of algorithms that add an auxiliary representation loss to the policy network while not interfering too much with its objective, unlike sharing the actor-critic trunk, which we show can be detrimental when its rank is low to start with.\n\nIn the scope of this work, we did not investigate such algorithms, as a thorough analysis would require assessing a comprehensive sample of these algorithms such as PPG, InFer (Lyle et al., 2022), and CURL (Laskin et al., 2020). We reserve this study for future work.", "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:
We thank the reviewers for their thorough and insightful reviews. We are glad that the reviewers appreciate the novelty and impact of assessing loss of plasticity in on-policy optimization and its connection to PPO’s trust region and acknowledge our thorough experimental setup.

We have **addressed all the concerns and questions of the reviewers** in the individual rebuttals.
We hope that this helps the reviewers **increase their support for the paper otherwise, kindly ask the reviewers to point out any remaining issues.**  

In addition, we use this section of the rebuttal to group and address **some common questions and concerns raised by reviewers. We reference those in our individual answers to the reviewers who raised them.**

**GQ1. The distinction between low-rank and poor representations**

>**YzGO Q3.** I would be careful about mentioning that low feature rank necessarily means that the representation is poor. Successful training of deep neural networks often involves a reduction in feature rank and the relationship between rank and performance can be complex Gulcehre et al. (2022).
**XWGd Q2a.** isn't it expected that the policy learns to lose rank, by compressing state information to only correlate with actions?

**GA1.** Indeed, we reference the work of Gulcehre et al. (2022) in our work and highlight in lines (236-238) that the relation we draw between the representation dynamics, the trust region, and the performance **primarily holds around the poor representation regime which we characterize by an extremely low rank, and not necessarily throughout training.** (lines 236-238: “We observe no significant correlation in the regions where the representation is rich … but an apparent decrease of the average probability ratios below $1 − \epsilon$ is observed as the representation reaches poor values”).

It may **not be straightforward to draw a line between low-rank representations beneficial for generalization and extremely low-rank representations causing aliasing**, as also acknowledged by Gulcehre et al. (2022) at least in offline value-based RL (“Unfortunately, reasoning about what it means for the rank to be too low is hard in general”), but for environments like Atari, **our figures seem to draw the line at single-digit ranks, which can be related to the action space of dimension 8+.**

We will further clarify this distinction in the same paragraph (lines 236-238) and in the captions.

**GQ2. The scope of the connection between representations and the trust region, and the significance of this scope.**

>**YzGO Q3.** I am also concerned about the link between feature rank and trust-region violation, a central point of the paper. In Fig.4, the correlation between the prob ratios and dead neurons or prob ratios and feature rank seem weak. There is only a substantial dip in the prob ratios once the feature rank is near zero or the number of dead units is very high. In other words, only when the agent is doing very poorly can we see a strong relationship, which may limit the applicability of the finding. Ideally, one would avoid this regime since performance is poor to start with.
> **Qwmj Q2**  ... This cripples the significance of PFO and also the potential causal connection …

**GA2.** As mentioned in the previous answer we claim that the causal connection we draw between the representation dynamics, the trust region, and the performance primarily holds around the collapse regime and not necessarily throughout training.

**Yet discovering and describing such a relation only around collapse is nontrivial.** First, this gives **evidence that this regime is often attained and should be avoided.** Second, it **gives important insights into the failure mode of the popular PPO algorithm**, whose trust region is highly dependent on the representation quality, and **more generally about current trust-region methods, which only constrain the output probabilities.**  
The discovery of this link can further drive research on training deep networks in non-stationary settings and influence the design of future trust-region methods (e.g., PFO forms a trust region in the representation space as well).

Meta-review: This paper is on rank collapse and loss of plasticity in reinforcement learning policies, which is claimed to have not been studied for on-policy policy gradient methods such as PPO. Thorough experiments are conducted to demonstrate and investigate this collapse phenomenon in PPO, revealing its connection to the degradation of the trust region, which the reviewers found to be novel, insightful, and helpful for the community.

Overall, this is a valuable addition to the recent investigative papers on the collapse and loss of plasticity phenomenon. Some reviewers showed concern about the work not leading to clear performance improvement, as mentioned by the rejecting reviewer. The authors need to address this concern by clearly describing the scope of the paper. They are also suggested to include results on the usual set of Mujoco tasks, which they produced during the rebuttal phase. 

Further, the area chair discussed with the authors whether loss of plasticity and collapse had previously been shown for PPO. The area chair provided references where such evidence already exists. The authors agreed to modify their claim incorporating such existing literature, which strengthens the paper by situating it well within the literature. Due to this shared understanding and overall agreement during the discussion, this paper will be a valuable contribution to this venue.
Decision: Accept (poster)