Annealed Implicit Q-learning in Online Reinforcement Learning

[2] Michal Nauman, , Mateusz Ostaszewski, Krzysztof Jankowski, Piotr Mi\los, Marek Cygan. Bigger, Regularized, Optimistic: scaling for compute and sample efficient continuous control. The Thirty-eighth Annual Conference on Neural Information Processing Systems. 2024.

[3] Sebastian Thrun and Anton Schwartz. Issues in using function approximation for reinforcement learning. In Fourth Connectionist Models Summer School, 1993.

[4] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1179–1191. Curran Associates, Inc., 2020.

[5] Motoki Omura, Takayuki Osa, Yusuke Mukuta, and Tatsuya Harada. Stabilizing extreme q-learning by maclaurin expansion. In Reinforcement Learning Conference, 2024.

Thank you very much for your valuable feedback and comments! Below, we provide our responses.

Q1. The authors should comment on the choice of learning the Q-function directly instead of using a separate value function as done in IQL. This choice should normally harm the performances as the learned expectile is now influenced by the stochasticity of the environment. Evaluating the proposed algorithm in a highly stochastic MDP would help justify this choice.

We conducted experiments in stochastic environments to compare the performance of AIQL when using both V-function and Q-function versus using only Q-function, and we have added the results in Section F. The stochastic environments were created by adding Gaussian noise to the actions input into the DM Control tasks. Even when significant noise (standard deviation of 0.5) was added to the action with range [-1, 1], the Q-function-only AIQL achieved better performance. Additionally, the annealing mechanism effectively mitigated the performance degradation caused by the noise. These results support the validity of our implementation, which excludes the V-function. We will include experiments on other tasks in the camera-ready version.

Q2. AIQL is only applied to SAC. Applying it to TD3 and CrossQ, for example, would strengthen the authors' claim.

CrossQ is known to be sensitive to hyperparameters and environment settings [2], making it less generalizable. Therefore, we conducted additional experiments with TD3 and added the results in Section D. By simply replacing the critic's L2 loss in TD3's official implementation, we observed similar performance improvements to those seen with SAC in the hopper-hop and humanoid-run tasks. We will include results for other tasks.

Q3. Analyzing the influence of the hyperparameter T would be beneficial for understanding its influence on AIQL.

We have added results in Section E showing the effects of varying the annealing duration $T$. Interestingly, even with a small $T$ of 1M steps, AIQL still achieved significant performance improvements compared to SAC. Within the range of $T$ tested in these experiments, AIQL demonstrated a certain degree of robustness to changes in $T$.

Q4. The code is not provided, which limits the reliability of the experiments. Additionally, the baselines’ hyperparameters are not reported.

We have uploaded the code. Additionally, we have updated Section A.2 with details about the baseline hyperparameters.

Q5 Line 155, the authors claim that the scenario where all action values are equal has the largest overestimation bias. Can this be justified?

Here is an intuitive example. Consider two scalar values $a$ and $b$ as the true Q-values for two actions, where the noise $e_1$ and $e_2$ is added to each value. The $e_1$ and $e_2$ follow a uniform distribution $U(−1,1)$. In this case, the overestimation bias can be expressed as:

$Z = \mathbb{E}_{e_1,e_2}[\max(a+e_1,b+e_2)]-\max(a,b).$

When $a \geq b$, overestimation does not occur if only $a + e_1$ is chosen in $\max(a + e_1, b + e_2)$ because $Z = \mathbb{E}[a+e_1]-a=0$. However, due to the influence of noise $e_1$ and $e_2$, there are instances where $b + e_2$ can be larger than $a + e_1$, resulting in overestimation.

This means that when $a$ and $b$ are far apart, $\max(a + e_1, b + e_2)$ is more likely to be $a + e_1$. On the other hand, when $a$ and $b$ are close, $\max(a + e_1, b + e_2)$ is more likely to select $b + e_2$, increasing the chances of overestimation. Overestimation reaches its maximum when $a = b$.

For example, if $a = 2$ and $b = 0$, $Z = 2 - 2 = 0$, meaning no overestimation occurs. However, if $a = b = 2$, then:

$Z=\mathbb{E}_{e_1,e_2}[\max (2+e_1,2+e_2)]-2$

$= \mathbb{E}_{e_1,e_2}[\max (e_1, e_2)].$

In this case, overestimation occurs. This type of formulation was derived in [3], which may serve as a useful reference.

Q6. To mitigate the issue presented in Weakness 2, the authors could consider having a Q-network with several heads, where each head is responsible for learning one expectile, similar to Quantile-Regression Deep Q-Network [1]. The head used for training the policy would evolve during training, starting from a low target expectile value to finish with a high target expectile value as proposed by the authors. I would be interested in the authors' comments about this suggestion.

Thank you very much for the highly intriguing suggestion! While it may lead to an increase in computational cost, we agree that using a distributional critic to achieve a stationary loss is an effective approach. (We avoid discrete changes in the loss by not using max-backup [4] or MXQL [5], instead opting for expectile loss to ensure continuous changes.) Although our computational resources are limited and results are not yet available, we plan to include the results in the camera-ready version.

Q7. Remarks

Thank you for pointing out those mistakes! We have corrected them.

Thank you for your valuable comments! Here are our responses to them.

Q1. Section 2.2 (Line 115): Why is it stated that "In IQL, near-optimal values are estimated by using τ close to 1"? Could you clarify this choice?

The goal of IQL is to compute the optimal target value ($r + \max Q(s',a')$) for continuous action tasks. In IQL's expectile loss, setting $\tau = 1$ allows for handling this optimal target value directly. However, due to issues such as overestimation bias and stability concerns, values like $\tau = 0.7$ or $0.9$ are typically used to estimate a near-optimal value instead.

Q2. Sections 4.1 & 4.6: Could you expand your analysis to discuss Exp1 in Section 4.6? Based on the explanation in Section 4.1, it seemed the natural annealing strategy might align with Exp1 in Section 4.6, yet it performed relatively worse than the linear and sigmoid strategies.

An annealing method like Exp1, which reduces $\tau$ more significantly at the beginning, might potentially improve performance further. However, Exp1's initial decay is too aggressive, suggesting that an annealing strategy falling between Exp1 and Linear could be more effective. Despite this, we demonstrated that a straightforward linear annealing approach achieves strong performance without requiring detailed tuning of the decay method. Discovering the optimal annealing strategy remains a promising direction for future research.

Q3. Section 4.2: In the XQL (Garg et al., 2023), the extension of SAC (X-SAC) showed comparable performance to SAC, even outperforming it in Hopper-Hop. However, the results in Figure 4 don’t seem to align with this finding. Could you provide any reasoning or justification for this discrepancy?

In our study, to ensure a fair comparison, we set the batch size and the number of hidden units in the network to 256, following the configuration used in the SAC paper. In the XQL paper, as noted in Appendix D, the batch size and the number of hidden units were set to 1024, which may have influenced the results. For XQL in our study, we used the official implementation of XQL and kept all other hyperparameters consistent with those reported in the paper. We have added details about the baseline hyperparameters in Section A.2.

We have included new experiments in Appendix D, E, and F to provide stronger support for the claims in our paper. In addition, we have expanded the explanations in the Introduction and Section 3.3 to address any potential ambiguity regarding the motivation and claims. We would be grateful if you could review these updates and our earlier response and let us know your thoughts.

Thank you for your insightful comments. Below, we provide our responses to your remarks.

Q1. The proposed method mainly follows the ideas of both IQL and SAC.

In our study, we found that the simple application of IQL loss to online RL significantly improves performance. While the approach is straightforward, it represents a novel finding that has not yet been explored, which we believe is a major contribution. To demonstrate the generality of this approach, we have conducted additional experiments combining it not only with SAC but also with TD3 and added the results in Section D. Similar performance improvements were observed in the case of TD3 as well.

Q2. The claims regarding the overestimation bias are not sufficiently convincing. Equation (1) is derived following the strong assumption that Q-functions are independent from actions, which tends to be an over-simplification of the scenario. In addition, the experiment shown in Figure 2 mainly concerns a single training step, but does not explain why such bias can still accumulate after online exploration (apart from claiming such result by 'if the bias becomes too large').

The simplified formulation is designed to analytically derive the overestimation bias, and its simplicity does not imply a lack of relevance to actual learning. The same formulation has been employed in works such as [1, 2, 3], and in [2, 3], the bias adjustments derived from this formulation have been shown to be effective in actual learning scenarios. While it is true that bias can be somewhat mitigated during online exploration, as shown in Figure 7 of our study, skewness accumulates progressively during learning, and bias, after an initial decrease in the early stages of learning, gradually increases over time

Q3. In Section 3.3, the advantage of tuning τ is described, after which a linearly decreasing τ(t) is proposed. However, there is neither theoretical/empirical observation showcasing how different τ affects the training in different stages (as a support to the advantage part), nor justification of the choice of linear functions (as a support to the proposal part).

Figure 7 demonstrates that skewness and bias increase when $\tau$ is fixed, providing the experimental motivation for annealing. While it is possible to consider annealing schedules other than linear one, such as dynamically adjusting $\tau$ based on bias or skewness, we prioritized simplicity in our study. Thus, we adopted linear annealing, which has low computational cost and requires fewer hyperparameters. (In contrast, scheduling based on bias or skewness may introduce additional computation and require extra hyperparameters.)

Q4. In Section 2.1, d0 is not clearly defined by just claiming it is 'the probability distribution', which actually seems to be the initial state distribution. And the horizon T is not defined. In the review of IQL, the loss function for the parameters of Q-functions misses the dependence on Vψ.

Thank you for pointing out those mistakes! We have corrected them accordingly.

Q5. In Section 3.1 the paper emphasizes its advantage by estimating directly the optimal value compared to methods like SAC, while in actual implementation according to Section 4.1, the proposed method is just an expectile-loss version of SAC. More clarifications on distinctions between methods could be helpful.

Our proposed method is simple: we replace the L2 loss used in the critic learning of SAC and TD3 with expectile loss and introduce annealing for $\tau$. When using L2 loss in critic learning, the target value is based on the SARSA framework ($r + \gamma \mathbb{E}[Q(s',a')]$). By using expectile loss, we instead estimate the near-optimal value ($r + \gamma \max[Q(s',a')]$). This straightforward modification leads to performance improvements and represents a key contribution of our work. Moreover, it is a generalizable approach that can be combined not only with SAC but also with other methods.

Q6. In some experimental results in Figure 4, SAC reaches higher average returns earlier than AIQL.

In the cheetah-run and humanoid-stand tasks, SAC shows slightly better performance than AIQL during the early stages of training (e.g., at 1 million steps). However, in the other 8 tasks, AIQL outperforms SAC, as can also be observed from the average performance in Table 1, where AIQL demonstrates a clear advantage. For cheetah-run and humanoid-stand, it is possible that their sensitivity to bias causes the initial bias to slow down learning. Nevertheless, due to the annealing mechanism, AIQL achieves similar final scores to SAC in these tasks.

Q7. Why is there no comparison with some online version of IQL, if any?

In IQL, the policy is trained to avoid deviating significantly from the dataset policy, making it less suitable for direct use in online RL. As a result, an alternative training method is required. To address this, we implemented our approach by replacing only the critic's loss in SAC, which is widely used in online RL. In the additional experiments combining our method with TD3, where entropy maximization is absent, the critic's learning closely resembles that of IQL in offline RL.

References

[1] Sebastian Thrun and Anton Schwartz. Issues in using function approximation for reinforcement learning. In Fourth Connectionist Models Summer School, 1993.

[2] Qingfeng Lan, Yangchen Pan, Alona Fyshe, and Martha White. Maxmin q-learning: Controlling the estimation bias of q-learning. In International Conference on Learning Representations, 2020.

[3] Xinyue Chen, Che Wang, Zijian Zhou, and Keith W. Ross. Randomized ensembled double q-learning: Learning fast without a model. In International Conference on Learning Representations, 2021.

Thanks for the additional details.

Proof for accumulation of over-estimation bias

• For convenience of reference, let's refer to the equations in your last reply as 'Equation 1-3', and the last inequality following the sentence 'the target becomes' as 'the inequality'. Then according the reviewer's understanding, through the process of online learning, the policy is being changed. Then the estimation bias, in terms of Q-function, should be the difference between the estimated Q of one specific intermediate policy and its underlying Q-function. Now going back to the derivation in the reply: by Equation 3, one can argue that the estimation of Q, in the sense of empirical loss function, is performed with a possibly higher target. But this does not guarantee that Q can be over-estimated for all state-action pairs simultaneously. In addition, when talking about the estimation bias due to Q\^\\prime, it is with respect to the policy in the next iteration, which is different from the one concerned when talking about using $Q$ as the target. That means one need to compare $\\hat{Q}\^{\\pi_{i+1}} - Q\^{\\pi_{i+1}}$ and $\\hat{Q}\^{\\pi_{i}} - Q\^{\\pi_{i}}$, which is not determined by the inequality provided. If I misunderstood the problem setting, please let me know.

Comparison

• The clarification about the implementation-performance-wise difference between the method and SAC looks reasonable.

Thank you for your response!

Proof for accumulation of over-estimation bias

First, in value function learning where the max operator is applied, as in Q-learning, this process corresponds to value iteration rather than policy iteration. Therefore, it does not estimate the value function of a specific policy but instead performs Bellman optimality backups to approximate the optimal Q-value (Q*). Consequently, evaluating $\hat{Q}^{\pi_k} - {Q}^{\pi_k}$ is not feasible. Furthermore, it is challenging to accurately calculate and compare the bias at each update step in a general form. Given that the discount rate contributes to bias reduction, the use of the term "accumulation" may require reconsideration, and we will revise such expressions accordingly. However, the fact that max leaves residual bias is evident from the left plot in Figure 7, where the bias increases at the final steps as $\tau$ grows. Furthermore, due to the inherent randomness in training methods like stochastic gradient descent, the Q-function remains noisy over time. This noise, based on Jensen's inequality, continuously introduces bias. This suggests that convergence to an unbiased value is hindered, and the annealing of $\tau$ is effective in mitigating this issue. We will revise the relevant expressions in the paper for the camera-ready version to clarify these points.

We would greatly appreciate any additional feedback you might have, as the rebuttal period will be concluding soon.

Thank you for your valuable comments and questions! We would like to address the concerns raised under Weaknesses and Questions.

Q1. Lacks justification of expectile loss

it is still unclear why introduce the implicit Q-learning loss to begin with?

The primary motivation for introducing the IQL loss lies in enabling the calculation of maxQ(s′,a′) during the critic's update process. In actor-critic methods for continuous action tasks, such as SAC and TD3, policy improvement depends solely on the actor. Our study aims to enhance performance by enabling improvement through the critic as well, achieved by calculating the maximum value. This motivation is consistent with that stated in XQL, one of the comparison methods in our research, and is similarly rooted in prior work like Soft Q-Learning, which also computes a (soft) maximum. However, XQL suffers from instability, and Soft Q-Learning requires complex approximate inference procedures, prompting the development of SAC. In contrast, we used IQL loss as a stable alternative to XQL loss, which involves exponential terms, that implicitly calculates the maximum from the samples with simplicity.

By applying this loss to SAC, we aimed to demonstrate performance improvements by merely replacing the loss function of a commonly used method. The choice of SAC as the base algorithm is not obligatory. To further validate this, we added experiments combining AIQL with TD3 in Section D, where similar performance improvements were observed.

Then what is the root cause of AIQL outperforming SAC?

Therefore, in response to the question of why AIQL performs better than SAC, we believe this is due to more efficient policy improvement as well as the influence of bias. As you pointed out, SAC exhibits lower bias and skewness compared to AIQL. However, prior work such as [1] has experimentally demonstrated that positive bias can effectively encourage exploration.

From Figure 7, we observe that AIQL has a larger bias during the early stages of training, which gradually diminishes to levels comparable to SAC in the later stages. This behavior resembles an epsilon-greedy policy, where exploration is emphasized during the early stages of learning, while exploitation dominates in the later stages, facilitated by annealing.

We believe that this annealing mechanism and exploration driven by bias significantly contribute to the performance improvements seen in AIQL.

Q2. The annealing schedule of optimality is not sufficiently justified

I am still unable to understand why the bias should be kept high in the beginning (and lower in the end)?

As stated earlier in Section 3.2, overestimation bias arises from the calculation of max⁡Q(s′,a′) [2]. With the expectile loss, the closer $\tau$ is to 1, the closer it approximates the maximum, whereas at $\tau = 0.5$, it calculates the expectation.

This means that as $\tau$ approaches 1, overestimation bias becomes more likely to occur, while it becomes less likely as $\tau$ approaches 0.5. At the same time, when $\tau$ is close to 1, the loss approximates the optimal value through a maximum-like calculation, whereas at $\tau = 0.5$, it estimates the value function of the policy.

Based on this understanding, we claim that $\tau$ controls the trade-off between optimality and bias. We have added this explanation in Section 3.3 for clarity.

Q3. The scheduling of tau depends heavily on the environment

So, how does one define the annealing schedule in such cases?

The end value of $\tau$ could also be treated as a hyperparameter, but we recommend setting it to 0.5. By doing so, the algorithm transitions into SAC at the end of annealing, allowing the learning process to continue as SAC afterward. While the start value of $\tau$ and the annealing duration are hyperparameters, they introduce a continuous change rather than a discrete one to SAC. For example, setting the start value of $\tau$ to 0.5 and the annealing duration to 0 results in SAC. By gradually increasing the start value of $\tau$ and the annealing duration, the benefits of our method can be realized.

Table 3 demonstrates that the algorithm is robust to the specific values of $\tau$ when annealing is employed. Additionally, we have added experimental results regarding the duration of annealing in Section E. Even with an annealing duration as short as just 1 million steps, our method significantly outperforms SAC. Furthermore, the algorithm is not sensitive to the annealing duration, making it practical to use.

Q4. IQL loss on algorithms other than SAC

We have added results in Section D combining AIQL with TD3 for the tasks where AIQL with SAC had the most significant impact — hopper-hop and humanoid-run. Similar to the results of AIQL+SAC, we observed significant performance improvements. Experiments on other tasks will be included in the camera-ready version.

Q5. Interpretation of the values of τ annealing

What happens between 0.5 and 0.6 that causes the improvement?

Regarding the difference between Fixed (0.6) and SAC, we believe it stems from whether or not there is even a slight occurrence of policy improvement and bias during the critic's learning process, as discussed in response to A1. Improvements and biases accumulate as learning progresses, so even a small increase in $\tau$ from 0.5 should have an impact.

In Figure 7, it is evident that Fixed (0.6) exhibits greater skewness and bias compared to SAC. Similarly, in the IQL paper, although $\tau = 0.7$ (not so close to 1) is often used across various tasks, the performance significantly surpasses that of AWAC, which is similar to IQL with $\tau=0.5$.

Q6. Why is high optimality not necessary towards the end? If the value function is suboptimal, then the policy would also be suboptimal.

We had assumed that an optimal policy and value function could be learned; however, there is also the possibility that they may remain suboptimal. Our intention was to convey a state in which further policy improvement is no longer achievable. We have revised the relevant statement in Section 3.3 to make it more precise.

Q7. The paper categorizes hopper-hop and humanoid-run as "more challenging tasks". What makes hopper-hop more challenging than quadruped-run?

We assessed the difficulty of tasks based on the performance of prior studies. In the DM Control suite, the rewards are normalized to 1 per step, and each episode consists of 1000 steps, resulting in a maximum possible return of 1000. For the quadruped-run task, returns exceeding 800 have been achieved with methods like SAC and TD3. However, for tasks like hopper-hop and humanoid-run, returns are approximately 100-200, indicating that these are more challenging tasks.

Q8. Replacing SAC's value function update with expectile loss does mean that the entropy regularization is removed, correct?

We did not remove entropy regularization. Instead, we simply replaced the L2 loss in the critic's loss function with expectile loss, without altering the rest of the loss function.

Q9. What is the reason for AIQL to perform so well in hopper-hop? It seems this is the main reason for the cumulative results to be so good.

The improvement in AIQL's performance, as mentioned in response to Q1, is attributed to the efficiency of policy improvement and the exploration facilitated by bias. One way to enhance policy improvement is by increasing the number of updates (replay ratio) at the expense of computational cost. According to Appendix D.1 of [3], increasing the replay ratio significantly improves the performance in the hopper-hop task. This suggests that, in hopper-hop, policy improvement has a substantial impact on improving sample efficiency.

References

[1] Qingfeng Lan, Yangchen Pan, Alona Fyshe, and Martha White. Maxmin q-learning: Controlling the estimation bias of q-learning. In International Conference on Learning Representations, 2020.

[2] Sebastian Thrun and Anton Schwartz. Issues in using function approximation for reinforcement learning. In Fourth Connectionist Models Summer School, 1993.

[3] P. D’Oro, M. Schwarzer, E. Nikishin, P.-L. Bacon, M. G. Bellemare, and A. Courville. Sampleefficient reinforcement learning by breaking the replay ratio barrier. In International Conference on Learning Representations, 2023.

Thank you for the additional feedback. Below is our responses:

I believe the IQL paper already has experiments on online RL fine-tuning with IQL, so I don't see how that can be claimed as a novelty for this work. Is my understanding incorrect about IQL's results?

While fine-tuning is performed online after offline training in IQL paper, to the best of our knowledge, there is no existing work that trains an agent from scratch in an online setting using the IQL loss and compares it with widely-used online RL methods. We consider it a significant finding that a simple modification to the critic loss in existing methods can lead to performance improvements. Furthermore, the annealing of the $\tau$ parameter, implemented to suppress the increased bias introduced by using the IQL loss, is an approach uniquely feasible in the context of online RL. This approach enhances performance and makes the method more robust to the choice of $\tau$, which we regard as another key contribution of our work.

Furthermore, the performance of SAC and TD3 seem to be truncated and exceptionally low in these environments. Are these methods really fully hyperparameter-tuned to be state-of-the-art?

For the SAC implementation, we use the same hyperparameters as the original paper, while for TD3, we rely on the authors' implementation. We use the same hyperparameters when combining them with AIQL. Regarding the training steps, we extended the training to 5 million steps and added the results to Figure 13, which similarly show that the proposed method outperforms SAC. Moreover, the performance improvements in shorter training steps are also valuable from the perspective of improving sample efficiency.

The gains coming primarily from hopper and humanoid seems unconvincing to me and there is no direct evidence that incorporating the "max" step in IQL is the direct cause of this improvement in these environments.

In tasks other than Hopper and Humanoid, SAC's returns are close to 800, leaving limited room for improvement. However, even in tasks like Quadruped and Walker, AIQL demonstrates better performance than SAC, suggesting that AIQL is consistently effective. Furthermore, the utility of the "max" operation is evidenced by the experiments on max-backup in Figure 6 of our paper. In max-backup, the only modification to SAC is in the computation of the target value for the critic update, where multiple actions are sampled, and the maximum action is selected. Comparing this with the SAC returns shown in Figure 4, max-backup achieves higher returns. We believe this demonstrates the value of introducing the "max" step to accelerate policy improvement, which is a key element in RL.

There needs to be a better demonstration of why exactly IQL works, more than that "it happens to work" in these environments. Also, I think the authors should report results on a different set of benchmarks than dm control to really justify this difference empirically.

While we believe that evaluations on various locomotion tasks in DM Control demonstrate the effectiveness of our method, we are also conducting additional experiments on manipulation tasks in [1]. Due to the limited time and computational resources, we are uncertain whether we can share the results before the end of the rebuttal period. Nevertheless, we will share them as soon as they are completed.

[1] T. Yu, D. Quillen, Z. He, R. Julian, K. Hausman, C. Finn, and S. Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning (CoRL), 2020.