Paper: Adaptive Regularization of Representation Rank as an Implicit Constraint of Bellman Equation (ICLR-2024)

Reviewer comment (Theory):
The authors intuitively explained the intrinsic similarity between the cosine similarity and the representation rank with examples. I am curious about if there is some formal statement or proof on this topic, since it serves as a very important property of your paper. There is some parameter in the definition of the representation rank (e.g. $\epsilon$) in Definition 1, but I didn't see it in the following discussion.

Author response:
> Q1.1: The authors intuitively explained the intrinsic similarity between the cosine similarity and the representation rank with examples. I am curious about if there is some formal statement or proof on this topic since it serves as a very important property of your paper. 

**A1.1**: 
- To the best of our knowledge, there is currently no established formal statement or proof directly correlating cosine similarity with representation rank in the DRL literature. Our work is theoretically motivated. 
- Theoretically, rank cannot strictly bound the cosine similarity without making more further information (e.g. Bellman equation and neural network in our paper). 
- Empirically, researchers [1,2,3] could use cosine similarity to affect the rank. In our work, we explore the influence of cosine similarity on the rank, as detailed in Appendix D. 
- On the one hand, several practices and algorithms in the field [1,2,3] manipulate representation vectors to affect the rank, which is aligned with our approach. For example, [1] maximizes representation rank by aligning learned representation vectors with their initial counterparts.
- On the other hand, existing theoretical research [4,5,6] about representation rank cannot inherently bound the cosine similarity of the representation vectors.


> **Q1.2**: There is some parameter in the definition of the representation rank (e.g. $\epsilon$) in Definition 1, but I didn't see it in the following discussion.

**A1.2**: The parameter $\epsilon$ in Definition 1 aims to filter out small singular values, thereby focusing on vectors that mainly contribute to the rank. We selected $\epsilon=0.01$ for our experiments, as detailed in Tables 2, 3, and 4. This value aligns with the default setting in prior literature (referenced in section C.2 of [1]), providing a standardized basis for comparison and analysis in our work.

Reviewer comment (Theory):
When the authors transform the constraint in Eqn. (11) to the penalty regularizer with ReLU in Eqn.(12), is there any theoretical support for this transformation? And also is there any theoretical analysis on how to choose the value of $\beta$? I check the ablation study on this hyperparameter on Appendix, but I feel its choice is still quite important but unclear in practice.

Author response:
> Q2.1: When the authors transform the constraint in Eqn. (11) to the penalty regularizer with ReLU in Eqn. (12), is there any theoretical support for this transformation?

**A2.1**:
- The primal-dual method, typically used for some convex optimization problems, is not feasible here due to the non-convex nature of the original problem in Eqn. (11).  The original problem in Equation 11 involves optimizing $\phi$ and $w$ separately inside. The primal-dual method is still not guaranteed to convergence. 
- Our work is theoretically motivated. Our use of the ReLU function in the regularizer ensures the optimizer does not maximize the representation rank once the upper bound holds. Specifically, when the left-hand side of Equation (11) is less than its right-hand side, the gradients of our regularizer remain consistently zero due to the properties of $\nabla_x \text{ReLU}(x)=0$ for $x<0$.
- Empirically, according to our analysis in Figure 6 (Appendix E), the $C$ values of BEER empirically ensure the constraint is satisfied. Our method aligns with general practices in DRL literature [1,2,3], where adding a regularizer to the original optimization objective is a common method.


> Q2.2: And also is there any theoretical analysis on how to choose the value of $\beta$? I check the ablation study on this hyperparameter on Appendix, but I feel its choice is still quite important but unclear in practice.

**A2.2**:
- Theoretical insights: The choice of $\beta$ determines the relative impact of the two losses (primary optimization and regularization) on parameter gradients. To avoid undermining the primary optimization objective, the regularizer's loss should be considerably smaller than the primary loss. 
- Empirically, our experiments in Appendix F indicate that a relatively small $\beta$ works well. $\beta$ varying over a small domain of values ([0.0005,0.001] in our experiments) results in equally good performance.
- P.S. One of the reasons why we selected a consistent, fixed $\beta$ across all the experiments is to fairly evaluate our algorithm considering the recent crisis in the model evaluation [7].

Reviewer comment (Theory):
The argument that the cosine similarity of the **adjustment** state-action pairs $(s, a)$ and $(s', a')$ does not rigorously lead to the rank. This is because one can only control the similarity between limited state-action pairs, instead of all possible state-action pairs. Think of an extreme case that the representation in adjustment state-action pairs is iterating between $(0, 1, 0, \cdots, 0)$ and $(1, 0, 0, \cdots, 0)$, the implicit regularization will fail.

Author response:
**A1.** We thank you for pointing out this critical aspect! Our design indeed aims to access the entire spectrum of possible state-action pairs. As highlighted in [2], "In RL, the reachability of a state plays a fundamental role. For example, representation of states that are not reachable by any policy are irrelevant". The unreachable state-action pairs are irrelvant for the numerical representation rank (in Definition 1) because they cannot be sampled. We humbly hope to draw your attention to our sampling process as detailed in Algorithm 1, lines 7 to 10.

- **Sampling mechanism allows access to all possible (s,a) pairs.**: BEER employs a uniform sampling strategy from the replay buffer, which theoretically allows access to the entire spectrum of state-action pairs. This strategy is a fundamental aspect of our methodology, ensuring a comprehensive representation of the state-action pairs. Additionally, the exploration tricks in BEER, which introduce noise to actions, further ensure the agent covers as many state-action pairs as possible.


- **Practical Implementation**: We implement this with substantial sampling iterations – one million – with each batch comprising 256 samples. This extensive sampling framework significantly enhances the likelihood of covering all possible state-action pairings, thereby mitigating concerns about limited coverage.

We have added more details in our revised manuscript (Appendix J) to clarify these points and highlight the robustness of our sampling process.

Reviewer comment (Writing):
On page 5, what is the exact definition of $\mathbb{SG}$?

Author response:
**A2**. The term $\mathbb{SG}$, denoting 'Stop Gradient', is a concept from the semi-gradient nature of the Bellman backup, as discussed in Section 9.3 of [1].

To provide a more precise understanding, $\mathbb{SG}$ can be defined as follows:

$\textbf{Definition 2: Stop Gradient.}$ The $\mathbb{SG}$ operator is applied to a differentiable function $f: \mathbb{R}^n \rightarrow \mathbb{R}^m$ within the context of gradient-based optimization. When this operator is applied, it generates a new function $\mathbb{SG}(f)$, such that for any input vector $\mathbf{x} \in \mathbb{R}^n$, the output during the forward computation remains the same as that of $f$, i.e., $\mathbb{SG}(f)(\mathbf{x}) = f(\mathbf{x})$. However, during the backward computation phase, where gradients are typically calculated through backpropagation, the gradient of $\mathbb{SG}(f)$ with respect to $\mathbf{x}$ is explicitly set to zero. This is irrespective of the actual gradient of $f$ at $\mathbf{x}$. Formally, this concept is encapsulated by the following equation

$$
\forall \mathbf{x} \in \mathbb{R}^n, \quad \mathbb{SG}(f)(\mathbf{x}) = f(\mathbf{x}) \quad \text{and} \quad \nabla_{\mathbf{x}} \mathbb{SG}(f)(\mathbf{x}) = \mathbf{0}.
$$

We have included this analysis in Appendix I of our revised manuscript.

Reviewer comment (Theory):
I do not quite get why the authors need to design the regularization term as in (12). It seems more natural to me to set the regularization term as in (19), which avoids a calculation of the gradient of an inverse term $1/|\phi(s,a)|$ which might hurt the stability of the optimization process.

Author response:
**A3**. We apologize for the unintentional missing of a $\mathbb{SG}$ symbol in Eq 12. The corrected equation should read as follows: 

$$
\mathcal{R}(\theta) = \text{ReLU}\Big( \cos(\phi(s,a),\mathbb{SG}\overline{\phi(s',a')})  -  \mathbb{SG}\big((\| \phi(s,a) \|^2 + \gamma^2\| \overline{\phi}(s',a') \|^2 - \frac{\|r\|^2}{\|w\|^2} )\frac{1}{2\gamma \| \phi(s,a) \| \| \overline{\phi(s',a')} \| } \big)\Big),
$$

where the $\mathbb{SG}$ operation is applied to the derived upper bound. Note that our experimental results are based on this corrected equation (demonstrated by the submitted code) rather than the original Eq. (12). As a result of the $\mathbb{SG}$ operation, the gradient of $1/\phi(s,a)$ is not computed, thereby preserving the stability of the optimization process.

Reviewer comment (Evaluation):
The logic behind the superior performance of BEER is not clear to me. The authors tried to claim that 'representation rank affects the model performance (approximation error), while BEER explicitly controls the rank, thus BEER outperforms other baseline methods'. However, according to figure 2, InFeR has a similar representation rank as BEER, while it performs worse than BEER. Is there any explanation why such a phenomenon happens? Meanwhile, in simpler tasks (Grid World), the rank of BEER is higher than other baselines, while in complex tasks (Lunan radar), the rank of BEER is lower than other baselines. Due to the inconsistency, it is doubtful whether the performance gain is due to the better control of representation rank. If so, it indeed deserves more explanation.

Author response:
**A4.1 (Logic).**  BEER's effectiveness lies in its adaptive regulation of representation rank. Different tasks require different representation ranks. Overly high ranks may lead to overfitting, while too low ranks can cause underfitting. BEER's principle is based on the Bellman equation and adaptively balances this rank to optimize the performance.


**A4.2 (Experiments).** We highlight the fact that different tasks demand different representation ranks. Consistently high or consistently low representation rank is **not optimal** for all the tasks.



Determining whether a rank is good or bad is challenging. An effective guideline is that a rank is preferred if it reduces the approximation error. We thus utilize approximation error as our primary criterion for rank and model evaluation. This metric directly and precisely assesses the quality of ranks and models. 

Our experiments show that BEER consistently achieves lower approximation errors, leading to superior performance. BEER's representation rank, as depicted in Figures 1(b) and 2(b), demonstrates its efficacy in maintaining ranks adaptive to different tasks, which is critical to its performance.


Regarding Figure 2, the ranks of BEER and InFeR are statistically different, as indicated by their non-overlapping confidence intervals.

We hope these clarifications address your concerns. We remain open to any additional queries and are committed to continually improving our work. Thank you for considering our response.

Reviewer comment (Evaluation):
It is recommended to thoroughly investigate the computational overhead of the proposed method.

Author response:
Dear Reviewer ykJX,

Thank you very much for your detailed review and insightful comments on our manuscript. Your feedback has been instrumental in enhancing the quality of our work.  We appreciate your recognition of the foundational aspects of our method and its capability in controlling the representation rank in DRL.

In response to your concerns about the computational overhead of BEER, we would like to offer the following clarification. 

> **Q1.** It is recommended to thoroughly investigate the computational overhead of the proposed method.


**A1.** BEER introduces an additional loss term but maintains the original network architecture. This significantly mitigates the increase in computational burden. To provide a more detailed evaluation, we have conducted two key analyses:


1. **The number of learnable parameters**. BEER operates as a regularizer without introducing any new learnable parameters to the base algorithm. This contrasts with methods like InFeR, which incorporates multiple additional heads to the value networks and significantly increases the number of learnable parameters.

2. **Runtime analysis.** BEER’s runtime is competitive with other algorithms and, in some cases, even more efficient. We have carefully compared the runtime of BEER with other algorithms focusing on representation rank, including SAC, using the same standard evaluation protocol. The experiments were executed on a GPU server equipped with an Intel(R) Xeon(R) Gold 6240 CPU @ 2.60GHz and NVIDIA GeForce RTX 2080 Ti GPUs. Each task was run sequentially to ensure a fair comparison. The results are as follows (H=hours). 



Table 1. Runtime comparison of BEER and other algorithms. The runtime of BEER is competitive with other algorithms.

| Env                | BEER  | INFER | TD3   | SAC   | DR3   |
|--------------------|-------|-------|-------|-------|-------|
| pendulum-swingup   | 1.43H | 1.79H | 1.20H | 1.68H | 1.34H |
| acrobot-swingup    | 1.46H | 1.83H | 1.20H | 1.71H | 1.36H |
| hopper-hop         | 1.48H | 1.83H | 1.18H | 1.74H | 1.30H |
| finger-spin        | 1.48H | 1.81H | 1.25H | 1.85H | 1.29H |
| **Average Time**   | 1.47H | 1.81H | 1.21H | 1.74H | 1.32H |


These results show that BEER's runtime is competitive with other algorithms.

We have included this analysis in Appendix H of our revised manuscript for a comprehensive understanding.

We hope this response addresses your concerns and further supports the feasibility of our approach. We remain open to any additional queries and are committed to continually improving our work. Thank you for considering our response.


Sincerely,
Authors of Submission 1504.

References

[1] Lyle C, Rowland M, Dabney W. Understanding and Preventing Capacity Loss in Reinforcement Learning[C]//International Conference on Learning Representations. 2021.