No Representation, No Trust: Connecting Representation, Collapse, and Trust Issues in PPO

We formalize our RL setting with the finite-horizon undiscounted Markov decision process, describing the interaction between an agent and an environment with finite ¹¹1The 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 to Szepesvári (2022) for a formalism of RL in that setting. sets of states ${\cal S}$ and actions ${\cal A}$, and a reward function $r:{\cal S}\times{\cal A}\times{\cal S}\to\mathbb{R}$. An initial state $S_{0}\in{\cal S}$ is sampled from the environment, then at each time step $t\in\{0,\dots,t_{\max}-1\}$, the agent observes the state $S_{t}\in{\cal S}$, picks an action $A_{t}\in{\cal A}$ according to its policy $\pi:{\cal S}\to\Delta({\cal A})$ with probability $\pi(A_{t}|S_{t})$, ²²2The time step $t$ is included in the representation of $S_{t}$ to preserve the Markov property in finite-horizon tasks as done by Pardo et al. (2018) and is analogous to considering time-dependent policies in the classical formulation of finite-horizon MDPs. observes the next state $S_{t+1}\in S$ sampled from the environment and receives a reward $R_{t+1}\doteq r(S_{t},A_{t},S_{t+1})$. We denote by $G_{t}\doteq\sum_{k=t}^{t_{\max}-1}R_{k+1}$ the return after the action at time step $t$. The goal of the agent is to maximize its expected return $J(\pi)\doteq\mathbb{E}_{\pi\!\!}\left[\sum_{t=0}^{t_{\max}-1}R_{t+1}\right]=\mathbb{E}_{\pi\!\!}\left[G_{0}\right]$ over the induced random trajectories. We discuss the choice of this setting in Appendix A.1.

We consider on-policy deep actor-critic agents which train a policy network $\pi(\cdot;{\bm{\theta}})$, called actor and denoted $\pi_{\bm{\theta}}$, and a value network $\hat{v}(\cdot;{\bf w})$, called critic that approximates the return of $\pi_{\bm{\theta}}$ at every state (expected sum of rewards starting from a state). At every training stage, the agent collects a batch of samples, called rollout, with its current policy $\pi_{\bm{\theta}}$, 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 (e.g., $G_{t}$). We use $\lambda$-returns computed with the Generalized Advantage Estimator (GAE) (Schulman et al., 2015b). The actor is trained with the Proximal Policy Optimization (PPO) (Schulman et al., 2017).

PPO-Clip, the most popular variant of PPO algorithms proposed by Schulman et al. (2017), optimizes the actor by maximizing the surrogate objective in Equation 1 at each rollout.

The objective is defined for some small hyperparameter $\epsilon$, where $\pi_{\text{old}}$ is the last $\pi_{\bm{\theta}}$ of the previous optimization stage and is used to collect the training batch, and $\Psi_{t}$ is an estimator of the advantage of $\pi_{\text{old}}$ (e.g., ${\Psi_{t}=G_{t}-\hat{v}(S_{t};{\bf w}_{\text{old}})}$). We use the GAE in our experiments. The objective is optimized with minibatch gradient steps over multiple epochs on the rollout data (batch). We refer to PPO-Clip as PPO and provide a high-level pseudocode in Algorithm 1.

Intuitively, PPO aims to maximize the policy advantage $\mathbb{E}_{\pi_{\text{old}}\!\!}\left[\sum_{t=0}^{t_{\max}-1}\frac{\pi_{\bm{\theta}}(A_{t}|S_{t})}{\pi_{\text{old}}(A_{t}|S_{t})}\Psi_{t}\right]$ defined by Kakade & Langford (2002), which participates in a lower bound to the improvement of $\pi_{\bm{\theta}}$ given that it is close to $\pi_{\text{old}}$ (Schulman et al., 2015a, see Theorem 1). In this regard, a gradient step on $L_{\pi_{\text{old}}}^{CLIP}(\theta)$ would increase (decrease) the probability of actions at states yielding positive (negative) advantage until the ratio between the policies for those actions reaches $1+\epsilon$ ($1-\epsilon)$ 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.

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 policy returns change. This form of non-stationarity has been shown by previous work to hinder the ability of several deep RL agents (e.g., DQN (Mnih et al., 2015), SAC (Haarnoja et al., 2018)) to continue learning (Lyle et al., 2022; Nikishin et al., 2022). PPO introduces additional non-stationarity to the actor compared to policy gradient methods (such as vanilla policy gradient (Sutton et al., 1999), A2C Mnih et al. (2016)) by making its objective non-stationary as it depends on the previous policy. In addition, the practical benefit of the PPO objective is that it can be optimized by performing multiple epochs of mini-batch gradient descent on every new collected batch; however, this makes the networks more likely to be impacted by previous training rollouts as new rollouts are collected. Increasing the number of epochs can cause the agent to “overfit” more to previous experience, making the impact of non-stationarity stronger.

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 dimension $N\times D$ denoted by $\Phi$, 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.1 and E, and use the approximate rank in our main figures for its connection to principal component analysis (PCA). Given a threshold $\delta\in\mathbb{R}$ and the singular values of the feature matrix $\langle\sigma_{i}(\Phi),\dots,\sigma_{D}(\Phi)\rangle$ in decreasing order, the approximate rank of $\Phi$ is $\min_{k}\left\{\frac{\sum_{i=1}^{k}\sigma_{i}^{2}(\Phi)}{\sum_{j=1}^{D}\sigma_{j}^{2}(\Phi)}>1-\delta\right\}$ which corresponds to the smallest dimension of the subspace recovering $(1-\delta)\%$ of the variance of $\Phi$. We use $\delta=0.01$ i.e. the reconstruction recovers $99\%$ of the variance as done by Andriushchenko et al. (2023); Yang et al. (2020). We also refer to this metric as feature rank with reference to the rank of the feature matrix when there is no ambiguity.

Plasticity is computed on checkpoints of a network undergoing some training to measure the evolution of its ability to fit targets. Given a fixed target and a fixed optimization budget (number of gradient steps), a checkpoint’s plasticity loss is the loss from fitting the checkpoint to the target at the end of the optimization budget. Usually, the plasticity of a deep RL agent throughout training is measured by its ability to fit the outputs of a model initialized from the same distribution as the agent, on a fixed rollout collected by this target random model (Lyle et al., 2022; Nikishin et al., 2023). The data would in expectation be from the same distribution as the agent’s initial checkpoint. To fit the critic we use an $L^{2}$ loss on the outputs of the models. To fit the actor we use a KL divergence between the target and the checkpoint (forward KL).

We begin our experiments by training PPO agents on the Arcade Learning Environment (ALE)(Bellemare 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 hard-exploration game from the taxonomy presented by Bellemare et al. (2016). For MuJoCo, we train on Ants, Half-Cheetahs, 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. The ALE model uses ReLU activations (Nair & Hinton, 2010) and the MuJoCo one 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 metrics 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 et al., 2022; Lyle et al., 2022) apply to 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.

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 $\epsilon$ in the trust region for example. We keep the learning rate constant throughout training and use the same value for all the epoch configurations. To understand the feature dynamics, we measure different metrics that are proposed in the literature, including feature rank, number of dead neurons (Gulcehre et al., 2022), plasticity loss (Lyle et al., 2022), and penultimate layer pre-activation norm. Previous work has monitored feature norm values as well (Lyle et al., 2024); however, in our case, we found that as 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 plasticity loss.

We run five seeds per hyperparameter configuration and report mean curves with min/max shaded regions unless specified otherwise. All curves, except for plasticity loss, are smoothed using an exponentially weighted moving average of coefficient $0.05$. We seek to answer the following questions:

• •

  Q1. How do a PPO agent’s representation metrics, such as the feature rank and the plasticity 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 it affect performance?

• •

  Q2. How does increasing the number of epochs per rollout to vary non-stationarity affect a PPO agent’s representation? Does it degrade as observed in DQN and SAC agents when increasing the replay ratio (Nikishin et al., 2022; Kumar et al., 2022)?

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 issues discussed in Section 3. We seek to mitigate this effect in a way that is analogous to the PPO trust region between the optimized policy and the policy that collected the batch. We apply an $L^{2}$ loss on the difference between the pre-activated features of the optimized policy and the policy that collected the batch. This can also be thought of as a way to keep the pre-activations of the network during an update within a 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.

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 for MuJoCo with ReLU. The goal is not necessarily to obtain better performance but to see if PFO improves the representations learned by PPO and if, in turn, it affects its 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 plasticity loss down, and the rank up. This coincides with a 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 no collapse in performance is observed anymore on ALE/NameThisGame and ALE/Phoenix, with a slight increase in the upper tail showing that PFO can increase performance. Among the interventions we have tried, PFO provided the most consistent improvements across the tasks and metrics we monitored.

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, it can significantly vary 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 sparsity, gives two distinct effects; In dense-reward environments such as ALE/Phoenix and ALE/NameThisGame, the ranks are concentrated at low but non-zero 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 where sparse rewards deteriorate the rank of the value network. We provide training curves showing the difference in the evolution of the feature rank when sharing the actor-critic trunk in Appendix D.

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. Therefore, we experiment with resetting Adam’s moments after each batch collection (to avoid tuning its frequency) for both the actor and the critic; the moments are thus only accumulated over the epochs on the same batch. We observe that this intervention reduces the feature norm and increases the feature rank on ALE, which also reduces the excess probability ratio and, in some cases, improves performance; however, it is not as effective as the other interventions in preventing collapse and, like sharing the actor-critic trunk results in poor performance on ALE/Gravitar. Nevertheless, further tuning the frequency of the reset or coupling it with representation dynamics may yield different conclusions.

In this work, we have provided evidence that the representation deterioration under non-stationarity observed by previous work in value-based methods generalizes to PPO agents in ALE and MuJoCo with their common model architectures and can lead to performance collapse. We have shown 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 presented Proximal Feature Optimization (PFO), a simple novel auxiliary loss based on regularizing the features’ evolution, which mitigates representation degradation and along with other interventions shows that controlling the representation degradation improves performance.

Although we study pixel-based discrete action environments and continuous control environments, we only study the common architectures of each environment, which do not include normalization layers and are limited to relatively small architectures with no memory (i.e., exclude Transformers or RNNs). Therefore, the generalization of the results of this work to other settings is unknown. 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 failure to train indefinitely and sudden collapse may be due to representations. We have presented here methods that mitigate representation degradation, but none of them seem to be a fundamental solution and a deeper understanding of the reasons driving representation deterioration under non-stationarity is still needed. We plan to address these issues in future work.