Stabilizing Extreme Q-learning by Maclaurin Expansion

TL;DR

This paper introduces Maclaurin Expanded Extreme Q-learning, a method that stabilizes extreme Q-learning by applying Maclaurin expansion to the loss function, improving stability and performance in offline and online reinforcement learning tasks.

Contribution

The paper proposes a novel Maclaurin expansion technique for Extreme Q-learning, enhancing stability and allowing flexible error distribution modeling, which improves RL performance.

Findings

Significantly stabilizes learning in online RL tasks from DM Control.

Improves performance in offline RL tasks from D4RL.

Enables adjustment of error distribution assumptions from normal to Gumbel.

Abstract

In offline reinforcement learning, in-sample learning methods have been widely used to prevent performance degradation caused by evaluating out-of-distribution actions from the dataset. Extreme Q-learning (XQL) employs a loss function based on the assumption that Bellman error follows a Gumbel distribution, enabling it to model the soft optimal value function in an in-sample manner. It has demonstrated strong performance in both offline and online reinforcement learning settings. However, issues remain, such as the instability caused by the exponential term in the loss function and the risk of the error distribution deviating from the Gumbel distribution. Therefore, we propose Maclaurin Expanded Extreme Q-learning to enhance stability. In this method, applying Maclaurin expansion to the loss function in XQL enhances stability against large errors. This approach involves adjusting the modeled value function between the value function under the behavior policy and the soft optimal value function, thus achieving a trade-off between stability and optimality depending on the order of expansion. It also enables adjustment of the error distribution assumption from a normal distribution to a Gumbel distribution. Our method significantly stabilizes learning in online RL tasks from DM Control, where XQL was previously unstable. Additionally, it improves performance in several offline RL tasks from D4RL.