Paper: Identifying Policy Gradient Subspaces (ICLR-2024)

Reviewer comment (Experiments):
The contribution of the paper is focused on the **empirical existence** of gradient subspaces. Although the authors, in Section 5, discuss how the existence of this sub-space could be leveraged in practice, it remains an open question to provide empirical evidence of these eventual benefits.

Author response:
> Although the authors, in Section 5, discuss how the existence of this sub-space could be leveraged in practice, it remains an open question to provide empirical evidence of these eventual benefits.

There are potentially many ways to exploit subspaces in the context of PG (parameters space exploration and second order optimization being only two of them). Creating methods that effectively exploit this knowledge is a significant effort that is independent of the analysis. The goal of this paper is to understand gradient subspaces in PG methods and share our findings with the RL community. We hope to spark interest in leveraging this phenomenon for the design of PG methods, similarly to how Gur-Ari (2018) encouraged research on gradient subspaces in supervised learning that resulted in a wealth of practical applications(see the methods mentioned in the related work section). For this reason, we hope for your understanding, that advancing PG algorithms with the findings of this paper is left for future work.

Reviewer comment (Experiments):
**Limited experimental campaign**. It has to be noticed that the contributions of this paper are only empirical. As a consequence, I would have expected more experiments to prove the empirical existence of these sub-spaces. Results are limited to 6 domains (i.e., Finger-spin, Ball_in_cup, Ant, HalfCheetah, Pendulum, Walker2D). I invite the authors to expand the set of domains considered.

Author response:
> I invite the authors to expand the set of domains considered.

We doubled the number of domains by applying our analysis to six more tasks:  BipedalWalker, Hopper, LunarLanderContinuous, Reacher, and Swimmer from OpenAI Gym (Brockman et al., 2016) and the FetchReach from Gym Robotics (Plappert et al., 2018).
Please note that the selection of tasks encompasses a wide variety of different characteristics of RL tasks. The tasks range from low-dimensional classical control tasks (Pendulum) to complex high-dimensional locomotion tasks for diverse embodiments (e.g., Ant, Swimmer) and robotics-inspired tasks (FetchReach). Furthermore, the selection includes sparse (Ball_in_cup and FetchReach) and dense reward (e.g., Ant, Walker2d), as well as single-goal (e.g., Ant, Walker2d) and goal-conditioned (FetchReach and Reacher) tasks.
We expanded the plots in Figure 2 to show results for four (instead of the previous two) tasks. Appendix C displays detailed results for all twelve tasks. Consistently across all tasks, a significant fraction of the gradient lies in the high-curvature subspace. Furthermore, on all tasks, there is a substantial overlap between the subspaces at different timesteps. The consistency of these results across such a diverse set of tasks suggests that stable subspaces are a general phenomenon in PGs rather than being task-specific.

Reviewer comment (Novelty):
One of the perspectives to make the paper stronger and more convincing is to show the unique conclusion and domain-specific insight for policy gradient learning, since most of the conclusions actually come from the gradient subspace paper under a supervised learning setup.

Author response:
> [S]how the unique conclusion and domain-specific insight for policy gradient learning

The paper already contains insights that are specific to the domain of policy gradient / actor-critic methods.
Sections 4.2 and 4.3 highlight the differences in the quality of the gradient subspace for the actor and the critic.
The sections further point out differences between on-policy and off-policy learning.
The potential application to parameter-space exploration that we discuss in the conclusion is also specific to reinforcement learning.
Additionally, we are currently running more experiments that directly compare on-policy and off-policy learning by considering an on-policy variant of SAC. These experiment investigate to which extent the differences in the analysis results of SAC and PPO stem from the fact that SAC reuses data and PPO does not.  The results will allow disentangling the effects of on-policy/off-policy learning from the algorithmic details of the two RL algorithms.

Reviewer comment (Theory):
One of the perspectives to make the paper stronger and more convincing is to show the unique conclusion of the policy gradient since this is an extension of the gradient subspace in the supervised learning setup. I could imagine more in-depth discussion in the paper could be helpful.

Author response:
> [S]how the unique conclusion and domain-specific insight for policy gradient learning

The paper already contains insights that are specific to the domain of policy gradient / actor-critic methods.
Sections 4.2 and 4.3 highlight the differences in the quality of the gradient subspace for the actor and the critic.
The sections further point out differences between on-policy and off-policy learning.
The potential application to parameter-space exploration that we discuss in the conclusion is also specific to reinforcement learning.
Additionally, we are currently running more experiments that directly compare on-policy and off-policy learning by considering an on-policy variant of SAC. These experiment investigate to which extent the differences in the analysis results of SAC and PPO stem from the fact that SAC reuses data and PPO does not.  The results will allow disentangling the effects of on-policy/off-policy learning from the algorithmic details of the two RL algorithms.

Reviewer comment (Evaluation):
Investigate and explain why the PPO/on-policy method has a much lower gradient fraction in Figure 2, and how much of this is due to the data distribution shift.

Author response:
> Investigate and explain why the PPO/on-policy method has a much lower gradient fraction in Figure 2, and how much of this is due to the data distribution shift

As mentioned in the response to the first question, we will add experiments that consider an on-policy variant of SAC. Like PPO, the on-policy SAC variant collects a small dataset of on-policy transitions at every update, which is discarded after the update. Since the rest of the algorithm stays the same, these experiments will shed light on whether the differences in the results between PPO and SAC are due to the stronger distribution shift in on-policy learning or due to other details of the algorithms.

Reviewer comment (Experiments):
Have more experiments/configurations in Figure 2 to make the conclusion and discussion more sound, for instance, whether the PPO low gradient fraction is consistent across more RL tasks or not.

Author response:
> Have more experiments/configurations in Figure 2

To increase the confidence that gradient subspaces are a general phenomenon in PG methods, we doubled the number of tasks on which we conducted our analysis by adding the tasks BipedalWalker, Hopper, LunarLanderContinuous, Reacher, and Swimmer from OpenAI Gym (Brockman et al., 2016) and the FetchReach task from Gym Robotics (Plappert et al., 2018). We extended Figure 2 with results for  Ant and LunarLanderContinuous while describing the remaining tasks in Appendix C due to space constraints. Further, we are currently conducting an evaluation of the influence of the hyperparameter configurations on the analysis results, which will be added to the appendix in the revised paper.

Reviewer comment (Novelty):
It is unclear to which extent the contribution of identifying gradient subspace comparing to existing works in RL that focus on identifying parameter subspace (e.g. Gaya et al 2023 in the references) is significant, since both approaches have the same goal of improving training efficiency of policy parameters.

Author response:
Both works consider different subspaces and, hence, identifying each subspace can have distinct, maybe even orthogonal, benefits. Gaya et al. generate a new policy through a convex combination of policy parameters that are trained to be distinct (Eq. (6) in Gaya et al. 2022) and maximize the RL objective (Eq. (4) in Gaya et al. 2022). These parameters (referred to as anchor parameters) span the subspace. We, on the other hand, analyze the subspace spanned by the dominant eigenvectors of the Hessian matrix (second derivative of the RL objective wrt. the policy parameters = derivative of the policy gradient). 
Thus, Gaya et al. focus more on generalizing policies, and we appreciate that this line of work also reports benefits of using subspaces. Our work aims to accelerate the vast field of policy gradients methods by using subspaces. In that sense, there might exist multiple ways to exploit subspaces for more efficient policy gradient training. Guiding parameter-space exploration or enabling second-order optimization are just two ways that we propose for future work. In fact, since the works by Gaya et al. from 2022 and 2023 use PPO (line 5 in Figure 2 in Gaya et al. 2022) and SAC (line 8 in Algorithm 1 in Appendix C.1 in Gaya et al. 2023) within their framework, our works could even be combined.

Reviewer comment (Experiments):
Is it possible to experimentally verify in a realistic example that the methods of Gaya et al. fail and the methods in current paper succeed?

Author response:
Gaya et al. consider policy adaptation (increase the performance of the policy on similar but different MDPs at test time (Gaya et al., 2022)) and continual RL (training agents on sequences of tasks (Gaya et al., 2023)) settings, whereas our work addresses efficiency in single-task learning. So, the objectives are not fully aligned. In fact, the approach by Gaya et al. builds upon learned single-task policies. For that reason, it is not obvious how to create a scenario where both methods can be compared fairly.

Reviewer comment (Evaluation):
There is a very weak link between the experimental results and the claims made in the paper, such as:  
(a) Claim in Section 1 (last paragraph): "(i) parameter-space directions with significantly larger curvature exist in PG" → why can we say that the curvature is significantly larger? Figure 1 is highly qualitative in nature, and Figure 2 is essentially qualitative as well. Also, larger than what?  
(b) Claim in Section 1 (last paragraph): "(iii) the subspace is sufficiently stable to be useful for training" → how do the experiments justify sufficiency? Figure 3 only shows that the magnitude of the largest eigenvector projected to the subspace is about 0.4 for most of the learning. What does 0.4 mean? How does that imply sufficient stability for learning based on methods that make updates in this subspace?  
(c) Claim in Section 1 (last paragraph): "observe that the value function subspace often exhibits less variability and retains a larger portion of its gradient compared to the PG subspace" → this is inconclusive, as Figure 2 shows very similar trends for actor and critic.

Author response:
(a)

> [W]hy can we say that the curvature is significantly larger?

The curvature is determined by the magnitude of the Hessian eigenvalues. In Figure 1, we show for the actor and critic losses that a handful Hessian eigenvalues exist with a significantly larger magnitude than the rest. Therefore, a small number of parameter-space directions dominate the curvature of the objective.

> And larger than what?

We added “compared to the other parameter-space directions” to make this point clearer.

> "(iii) the subspace is sufficiently stable to be useful for training" --> how do the experiments justify the sufficiency?

Our current experiments do not yet evaluate the sufficiency of the subspace for downstream applications. We changed the formulation to “(iii) the subspace remains relatively stable throughout the RL training.” Although it is yet to be shown that the stability of the subspace throughout the training is useful, we believe this result by itself is already interesting due to the high variance nature of PG algorithms.

> "observe that the value function subspace often exhibits less variability and retains a larger portion of its gradient compared to the PG subspace" --> inconclusive (Figure 2 shows very similar trends for actor and critic)

We added results for more tasks to Figure 2. Furthermore, we fixed a bug in our implementation that changed the analysis results for the Walker2D task (see the general response for details about this bugfix). The updated Figure 2 shows the trend that the critic gradient tends to lie better in the respective subspace more clearly. This effect is particularly strong for PPO, but the average gradient subspace fraction is also slightly higher for the critic in SAC. Please note that this trend is also noticeable in the results for all tasks (see Appendix C) and is not limited to the four tasks displayed in Figure 2.

Reviewer comment (Evaluation):
In Section 4.3 (paragraph 3), the statement "Similar to the gradient subspace fraction results from Section 4.2, the subspace overlap is more pronounced for the critic than the actor" lacks explanation. Why is this the case? Are the scales comparable for the actor and critic graphs in Figure 3? Is this "extra" overlap really useful for designing algorithms? Without explanations or additional experiments, this remains speculative.

Author response:
(c) 
> "Similar to the gradient subspace fraction results from Section 4.2, the subspace overlap is more pronounced for the critic than the actor." --> why? Are the scales comparable for the actor and critic graphs in Figure 3?

Yes, the scales are comparable. We applied the same criterion and used the same subspace dimensions for both the actor and the critic. The criterion is inherently normalized to the range [0, 1] (since we assume the eigenvectors to be norm 1). Furthermore, both networks have the same hidden layer size and, thus, contain roughly the same number of parameters (there is a small difference due to the different input/output dimensions).

The results in Figure 3 and the detailed results in Figures 7 and 8 of Appendix C show that the subspace overlap tends to be larger for the critic than for the actor.

> Is this "extra" overlap really useful for designing algorithms?

The design of algorithms that exploit the subspace phenomenon in the context of PG algorithms can take various forms and, hence, requires substantial experimental validation, which is left for future works. Because there is a significant overlap between the subspaces at different timesteps, the subspace does not need to be determined anew at every step during training. Downstream applications could exploit this fact to increase the computational efficiency of the method. The fact that the subspace overlap is on average higher for the critic than the actor could be exploited for further efficiency gains. The scale of these efficiency gains, however, depends on the specific algorithm.

Reviewer comment (Evaluation):
Gradient subspace fraction (Eq. 5) seems uninformative. Why is showing that the gradient norm in the projected subspace is similar to the original gradient norm helpful? From the Johnson-Lindenstrauss lemma, any random projection matrix will ensure that the projected gradient norm is not too far away from the actual gradient norm. Thus, Eq. 5’s metric doesn’t seem useful immediately. A possible improvement would be to establish a scale by comparing the norm of the projected gradient to the Hessian eigenvector subspace versus the norm of the gradient projected to random subspaces.

Author response:
(a)
> Why showing that the gradient norm in the projected subspace is similar to the original gradient norm helpful? For instance, from the Johnson-Lindenstrauss lemma we know that any random projection matrix will ensure that the projected gradient norm is not too far away from the actual gradient norm.

The criterion essentially measures how much of the original gradient is preserved by the projection into the low-dimensional subspace. We regret to have not made it explicit in the original submission that the criterion $\frac{||P_k g||^2}{||g||^2}$ is equivalent to $1 - \frac{||\tilde{g} - g||^2}{||g||^2}$, where $\tilde{g}$ is the gradient first mapped into the subspace and then back to the original dimensionality. The latter criterion measures how well the original gradient can be reconstructed from its subspace representation. This measure exactly represents what matters in our scenario. However, the equivalence above holds only for semi-orthogonal matrices (matrices where the columns are norm 1 and orthogonal to each other). We added this explanation to Section 4.2. of the paper and added the derivation of the equivalence in Appendix A.

The random matrices that the Johnson-Lindenstrauss lemma considers are not semi-orthonormal (see e.g. the matrix in  https://www.cs.cornell.edu/courses/cs4786/2020sp/lectures/lecnotes4.pdf). While their rows are approximately orthogonal due to the high dimensionality, they are not norm 1. In essence, this random projection still loses information while mapping to the subspace, it just stretches the vector back to the original length due to its row norm > 1. The equivalence of the criteria above does not hold for the random projection, and, importantly, the relative projection error will be high for such a projection (as it is unlikely that a random subspace contains the gradient by chance).

> Maybe one way to make it more informative would be to establish a scale (so compare the norm of the projected gradient to the Hessian eigenvector subspace and the norm of the gradient projected to random subspaces).

For random projection matrices, the value of the gradient subspace fraction is in expectation $k/n$, i.e., the ratio of the subspace dimension to the original parameter space dimension (so 0.02 for PPO and 0.0014 for SAC). While we currently do not have a formal proof of this statement, we implemented an empirical sanity check by mapping vectors to random subspaces. The results of this sanity check conformed to our intuition.

Reviewer comment (Theory):
Subspace overlap (Eq. 7) is complicated, and its relevance for showing the stability of the subspace and what that means for restricting training there is unclear. Additional references might help clarify this.

Author response:
(b)

As for the gradient subspace fraction criterion (eq. (5) (and (6) in the updated paper)), we measure to which extent a vector lies in the subspace. In contrast to the gradient subspace fraction criterion, for the subspace overlap, we do not project the gradient into the subspace, but rather the vectors that span the subspace (the Hessian eigenvectors) of the later timestep $t_2$. Since the subspace is spanned by multiple eigenvectors, we average the criterion over these vectors. The criterion measures how similar the two subspaces are. If the two subspaces are identical (i.e., the subspace did not change over time), the value of the criterion is 1. If there is no overlap between the subspaces, the criterion is 0.

Reviewer comment (Reproducibility):
Missing details about the implementations, for example:  
- Section 4.3 (equation for S_overlap): what is $k$ in the experiments? Does $k$ change with $t$ in the graph?  
- Section 4.3: "where $v_i$ is the ith largest eigenvector at timestep t" → how is the largest eigenvector determined? By the size of the corresponding eigenvalue?  
- Section 4 (paragraph 2): This section lacks sufficient details for reproduction; for example, SAC usually has multiple Q networks, so which are analyzed in the experiments?  
- Page 6, Figure 2: The difference between "Estimated/true" gradients and "estimated/true" Hessians is unclear; why does it matter? Presumably the Hessian is used to find the projection matrix $P_k$, but this should be explicitly stated. Why care about estimated gradients? Some justification is needed.  
- Page 6 (first paragraph): The use of 2%, 0.14%, and 0.07% appears arbitrary. How do trends change when these values vary?  
- Page 5: Cutoffs for Equation 6 seem arbitrary.

Author response:
R2 & R3: ## 3.
(a) Throughout the analysis, the subspace dimensionality k remains fixed. As in Section 4.2, we use k=100. We added this clarification to Section 4.3.

(b) Yes, we use “i-th largest eigenvector” as a shorthand for ”eigenvector that corresponds to the i-th largest eigenvalue”. We will add this clarification to the preliminaries.

(c) The Q-loss of SAC is typically implemented as the mean over the losses for the individual networks, so a single gradient is calculated for the two Q-functions. Hence, we analyze the gradient of both networks.

## 3 (cont.)
(d)
> I was not able to understand the difference between the three categories: "Estimated/true" gradient, "estimated/true" Hessian

The difference between the estimated and true gradient / Hessian (note that we renamed these to *mini-batch* and *precise* Hessian in the paper; see answer 4(l)) lies in the number of samples that we use to compute each of the quantities (10⁶ samples vs. a mini-batch). We investigate the “true” quantities to analyze whether the phenomenon occurs in the individual underlying optimization problems that PPO and SAC attempt to solve and the “estimated” quantities since these low-sample approximations are used in practice during training.

> Why do we care about estimated gradients?

The “estimated gradients” are mini-batch estimates of the true gradient. RL algorithms like PPO and SAC compute such mini-batch estimates of the gradient for updating the policy. For the analysis results to be useful in practice, it is crucial that the observations also transfer to the mini-batch gradients that are used in practice during training.

> Hessian is used to find the projection matrix $P_k$; if so, please explicitly mention it somewhere

We mention at the beginning of Section 4.2 that “the subspaces [are] spanned by the high-curvature directions” and that “$P_k$ denotes the matrix that projects [...] into the subspace”. We further mention in the caption of Figure 2 that the “subspace [is] spanned by the 100 largest Hessian eigenvectors”.

(e) The subspace dimension is k=100 in the experiments. Different network sizes cause the differing percentages. As mentioned in the paper, we chose k=100 because subspaces of this size can be identified relatively efficiently while  still capturing the gradients to a large degree. In any downstream application, the subspace size would be a hyperparameter, and, depending on the task, smaller values might already suffice. Appendix B (Appendix C in the revised paper) displays analysis results for different subspace sizes.

(f) We chose these cut-offs because the resulting splits into the different training phases conformed well with our intuition in which phase of the training the agent is at each training step. Since there is no universally accepted notion of when an agent is in an initial training phase / starting to converge, there is not really a way to define these thresholds more rigorously. For this reason, we frankly name this criterion heuristic.

Reviewer comment (Writing):
The abstract seems incomplete or misleading and overclaims the contributions. For example, it states "we demonstrate the existence of such gradient subspaces for policy gradient algorithms," but the paper only empirically shows these for a very limited class (specifically PPO and SAC). This could be rectified by including the line from the introduction: "This paper conducts a comprehensive empirical evaluation of gradient subspaces in the context of PG algorithms, assessing their properties across various simulated RL benchmarks." Additionally, the statement "Our findings reveal promising directions for more efficient reinforcement learning..." suggests that the paper introduces methods exploiting these subspaces, but it merely suggests future research directions.

Author response:
## 4.
(a) We modified the abstract to be more precise about our contributions.

Reviewer comment (Novelty):
Consider adding the "LoRA: Low-Rank Adaptation of Large Language Models" paper (https://arxiv.org/abs/2106.09685) to the related works section.

Author response:
(c) We will add the LoRA paper to the related work section of the revised paper.

Reviewer comment (Theory):
In Section 3.1: Can $\gamma$ be 1? Is the setting episodic?

Author response:
(d) γ is usually smaller than 1 in practice, but standard definitions often state 0 ≤ γ ≤ 1 (see e.g., chapter 3.3. in (Sutton & Barto, 2018)).

Reviewer comment (Writing):
In Section 3.1 (paragraph 2): Replace the phrase "expected cumulative reward" with "expected (discounted) cumulative reward."

Author response:
(e) As suggested, we added *(discounted)* as clarification.

Reviewer comment (Writing):
In Section 3.1 (paragraph 2): The phrase "advantage function ... and can also be defined as" is odd—this is the standard definition of the advantage function. Consider rephrasing.

Author response:
(f) We changed the formulation to the definition A(s, a) = Q(s, a) - V(s) to make this clearer.

Reviewer comment (Writing):
Section 3.2, first line: Include a reference for the objective function $J(\theta)$.

Author response:
(g) We added a reference to (Kakade & Langford, 2002) for the objective.

Reviewer comment (Theory):
Section 3.2, first line: The definition of the expectation is unclear; in its current form, $J = \mathbb{E}[ \pi(a_t | s_t) \hat A_t]$ appears to depend on timestep $t$, which should not be the case.

Author response:
(h) In this section, we employ the notation from (Schulman et al., 2017) to facilitate readers to check the details of the algorithms in the original paper. The value of the objective depends on the timestep through the states that the agent observes / the actions that it executes at that timestep.

Reviewer comment (Reproducibility):
Section 3.2, first paragraph: The "estimator of advantage function" used is not specified.

Author response:
(i) This part is still about policy gradients in general, so the estimator depends on the concrete algorithm. The advantage estimator for PPO is specified below equation (1).

Reviewer comment (Writing):
Section 3.2, first paragraph (above equation 2): What is the target value for the critic?

Author response:
(j) We added the definition of the target value (the discounted sum of future episode rewards).