General Munchausen Reinforcement Learning with Tsallis Kullback-Leibler Divergence

Tsallis entropy was first proposed by Tsallis (1988) and is defined by the $q$-logarithm. The $q$-logarithm and its unique inverse function, the $q$-exponential, are defined as:

where $[\cdot]_{+}:=\max\{\cdot,0\}$. We define $\ln_{1}=\ln\,,\exp_{1}=\exp$, as in the limit $q\rightarrow 1$, the formulas in Eq. (2) approach these functions. Tsallis entropy can be defined by $S_{q}(\pi):=p\left\langle-\pi^{q},\ln_{q}\!\pi\right\rangle,p\in\mathbb{R}$ (Suyari and Tsukada, 2005). We visualize the $q$-logarithm ​, $q$-exponential and Tsallis entropy for different $q$ in Figure 1. As $q$ gets larger, $q$-logarithm (and hence Tsallis entropy) becomes more flat and $q$-exponential more steep¹¹1 The $q$-logarithm defined here is consistent with the physics literature and different from prior RL works (Lee et al., 2020), where a change of variable $q^{*}=2-q$ is made. We analyze both cases in Appendix A. . Note that $\exp_{q}\!$ is only invertible for $x>\frac{-1}{1-q}$.

Tsallis policies have a similar form to softmax, but using the $q$-exponential instead. Let us provide some intuition for these policies. When $p=\frac{1}{2},q=2$, $S_{2}(\pi)=\frac{1}{2}\left\langle\pi,1-\pi\right\rangle$, the optimization problem $\operatorname*{arg\,max}_{\pi\in\Delta(\mathcal{A})}\left\langle\pi,Q\right\rangle+S_{2}(\pi)=\operatorname*{arg\,min}_{\pi\in\Delta(\mathcal{A})}\left|\!\left|\pi-Q\right|\!\right|_{2}^{2}$ is known to be the Euclidean projection onto the probability simplex. Its solution $\left[Q-\psi\right]_{+}$ is called the sparsemax (Martins and Astudillo, 2016; Lee et al., 2018) and has sparse support (Duchi et al., 2008; Condat, 2016; Blondel et al., 2020). $\psi:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{S}$ is the unique function satisfying $\left\langle\mathbf{1},\left[Q-\psi\right]_{+}\right\rangle=1$.

Theorem 1 characterizes the role played by $q$: controlling the degree of truncation. We show the truncation effect when $q=2$ and $q=50$ in Figure 2, confirming that Tsallis policies tend to truncate more as $q$ gets larger. The theorem also highlights that we can set $q=2$ and still get more or less truncation using different $\tau$, helping to explain why in our experiments $q=2$ is a generally effective choice.

Lee et al. (2020) also used $\exp_{q}$ to represent policies but they consider the continuous action setting and do not give any computable threshold. By contrast, Theorem 2 presents an easily computable $\hat{\psi}_{q}$ for all $q\notin\{1,\infty\}$.

The Tsallis KL divergence is defined as $D^{q}_{\!K\!L}\!\left(\pi\left|\!\right|\mu\right):=\left\langle\pi,-\ln_{q}\frac{\mu}{\pi}\right\rangle$ (Furuichi et al., 2004). It is a member of $f$-divergence and can be recovered by choosing $f(t)=-\ln_{q}t$. As a divergence penalty, it is required that $q>0$ since $f(t)$ should be convex. We further assume that $q>1$ to align with standard divergences; i.e. penalize large value of $\frac{\pi}{\mu}$, since for $0<q<1$ the regularization would penalize $\frac{\mu}{\pi}$ instead. In practice, we find that $0<q<1$ tend to perform poorly. In contrast to KL, Tsallis KL is more mass-covering; i.e. its value is proportional to the $q$-th power of the ratio $\frac{\pi}{\mu}$. When $q$ is big, large values of $\frac{\pi}{\mu}$ are strongly penalized (Wang et al., 2018). This behavior of Tsallis KL divergence can also be found in other well-known divergences: the $\alpha$-divergence (Wang et al., 2018; Belousov and Peters, 2019) coincides with Tsallis KL when $\alpha=2$; Rényi’s divergence also penalizes large policy ratio by raising it to the power $q$, but inside the logarithm, which is therefore an additive extension of KL (Li and Turner, 2016). In the limit of $q\rightarrow 1$, Tsallis entropy recovers Shannon entropy and the Tsallis KL divergence recovers the KL divergence. We plot the Tsallis KL divergence behavior in Figure 2.

Now let us turn to formalizing when value iteration under Tsallis regularization converges. The $q$-logarithm has the following properties: Convexity: $\ln_{q}\!\pi$ is convex for $q\leq 0$, concave for $q>0$. When $q=0$, both $\ln_{q}\!\,\,,\exp_{q}\!$ become linear. Monotonicity: $\ln_{q}\pi$ is monotonically increasing with respect to $\pi$. These two properties can be simply verified by checking the first and second order derivative. We prove in Appendix A the following similarity between Shannon entropy (reps. KL) and Tsallis entropy (resp. Tsallis KL). Bounded entropy: we have $0\leq\mathcal{H}\left(\pi\right)\leq\ln|\mathcal{A}|$; and $\forall q,\,0\leq S_{q}(\pi)\leq\ln_{q}{|\mathcal{A}|}$. Generalized KL property: $\forall q$, $D^{q}_{\!K\!L}\!\left(\pi\left|\!\right|\mu\right)\geq 0$. $D^{q}_{\!K\!L}\!\left(\pi\left|\!\right|\mu\right)=0$ if and only if $\pi=\mu$ almost everywhere, and $D^{q}_{\!K\!L}\!\left(\pi\left|\!\right|\mu\right)\rightarrow\infty$ whenever $\pi(a|s)>0$ and $\mu(a|s)=0$.

However, despite their similarity, a crucial difference is that $\ln_{q}$ is non-extensive, which means it is not additive (Tsallis, 1988). In fact, $\ln_{q}$ is only pseudo-additive:

Pseudo-additivity complicates obtaining convergence results for Eq. (1) with $q$-logarithm regularizers, since the techniques used for Shannon entropy and KL divergence are generally not applicable to their $\ln_{q}$ counterparts. Moreover, deriving the optimal policy may be nontrivial. Convergence results have only been established for Tsallis entropy (Lee et al., 2018; Chow et al., 2018).

We next show that the optimal regularized policy under Tsallis KL regularization does more than uniform averaging. It can be seen as performing a weighted average where the degree of weighting is controlled by $q$. Consider the recursion

where we dropped the regularization coefficient $\tau$ for convenience.

The form of this policy is harder to intuit, but we can try to understand each component. The first component actually corresponds to a weighted averaging by the property of the $\exp_{q}\!\,$:

Eq. (8) is a possible way to expand the summation: the left-hand side of the equation is what one might expect from conventional KL regularization; while the right-hand side shows a weighted scheme such that any estimate $Q_{j}$ is weighted by the summation of estimates before $Q_{j}$ times $1-q$ (Note that we can exchange 1 and $q$, see Appendix A). Weighting down numerator by the sum of components in the demoninator has been analyzed before in the literature of weighted average by robust divergences, e.g., the $\gamma$-divergence (Futami et al., 2018, Table 1). Therefore, we conjecture this functional form helps weighting down the magnitude of excessively large $Q_{k}$, which can also be controlled by choosing $q$. In fact, obtaining a weighted average has been an important topic in RL, where many proposed heuristics coincide with weighted averaging (Grau-Moya et al., 2019; Haarnoja et al., 2018; Kitamura et al., 2021).

Now let us consider the second term with $q=2$, therefore the leading $(q-1)^{j}$ vanishes. The action-value cross-product term can be intuitively understood as further increasing the probability for any actions that have had consistently larger values across iterations. This observation agrees with the mode-covering property of Tsallis KL. However, there is no concrete evidence yet how the average inside $q$-exponential and the cross-product action values may work jointly to benefit the policy, and their benefits may depend on the task and environments, requiring further categorization and discussion. Empirically, we find that the nonlinearity of Tsallis KL policies bring superior performance to the uniform averaging KL policies on the testbed considered.

Even for the standard KL, it is difficult to implement KL-regularized value iteration with function approximation. The difficulty arises from the fact that we cannot exactly obtain $\pi_{k+1}\propto\pi_{k}\exp\left(Q_{k}\right)$. This policy might not be representable by our function approximator. For $q=1$, one needs to store all past $Q_{k}$ which is computationally infeasible.

An alternative direction has been to construct a different value function iteration scheme, which is equivalent to the original KL regularized value iteration (Azar et al., 2012; Kozuno et al., 2019). A recent method of this family is Munchausen VI (MVI) (Vieillard et al., 2020b). MVI implicitly enforces KL regularization using the recursion

We see that Eq. (9) is Eq. (1) with $\Omega(\pi)=-\mathcal{H}\left(\pi\right)$ (blue) plus an additional red Munchausen term, with coefficient $\alpha$. Vieillard et al. (2020b) showed that implicit KL regularization was performed under the hood, even though we still have tractable $\pi_{k+1}\propto\exp\left(\tau^{-1}Q_{k}\right)$:

where $Q^{\prime}_{k+1}\!:=\!Q_{k+1}-\alpha\tau\ln{\pi_{k+1}}$ is the generalized action value function.

The implementation of this idea uses the fact that $\alpha\tau\ln\pi_{k+1}=\alpha(Q_{k}-\mathcal{M}_{\tau}Q_{k})$, where $\mathcal{M}_{\tau}Q_{k}:=\frac{1}{Z_{k}}\left\langle\exp\left(\tau^{-1}Q_{k}\right),Q_{k}\right\rangle,Z_{k}=\left\langle\boldsymbol{1},\exp\left(\tau^{-1}Q_{k}\right)\right\rangle$ is the Boltzmann softmax operator.²²2Using $\mathcal{M}_{\tau}Q$ is equivalent to the log-sum-exp operator up to a constant shift (Azar et al., 2012). In the original work, computing this advantage term was found to be more stable than directly using the log of the policy. In our extension, we use the same form.

The MVI(q) algorithm is a simple extension of MVI: it replaces the standard exponential in the definition of the advantage with the $q$-exponential. We can express this action gap as $Q_{k}-\mathcal{M}_{q,\tau}{Q_{k}}$, where $\mathcal{M}_{q,\tau}{Q_{k}}=\left\langle\exp_{q}\left(\frac{Q_{k}}{\tau}-\psi_{q}\left(\frac{Q_{k}}{\tau}\right)\right),Q_{k}\right\rangle$. When $q=1$, it recovers $Q_{k}-\mathcal{M}_{\tau}Q_{k}$. We summarize this MVI($q$) algorithm in Algorithm B in the Appendix. When $q=1$, we recover MVI. For $q=\infty$, we get that $\mathcal{M}_{\infty,\tau}{Q_{k}}$ is $\max_{a}Q_{k}(s,a)$—no regularization—and we recover advantage learning (Baird and Moore, 1999). Similar to the original MVI algorithm, MVI($q$) enjoys tractable policy expression with $\pi_{k+1}\propto\exp_{q}\left(\tau^{-1}Q_{k}\right)$.

Unlike MVI, however, MVI($q$) no longer exactly implements the implicit regularization shown in Eq. (10). Below, we go through a similar derivation as MVI, show why there is an approximation and motivate why the above advantage term is a reasonable approximation. In addition to this reasoning, our primary motivation for this extension of MVI to use $q>1$ was to inherit the same simple form as MVI as well as because empirically we found it to be effective.

Let us similarly define a generalized action value function $Q_{k+1}^{\prime}=Q_{k+1}-\alpha\tau\ln_{q}\!{\pi_{k+1}}$. Using the relationship $\ln_{q}\!{\pi_{k}}=\ln_{q}\!\frac{\pi_{k}}{\pi_{k+1}}-\ln_{q}\!\frac{1}{\pi_{k+1}}-(1-q)\ln_{q}\!\pi_{k}\ln_{q}\!\frac{1}{\pi_{k+1}}$, we get

where we leveraged the fact that $-\alpha\tau\left\langle\pi_{k+1},\ln_{q}\!\frac{1}{\pi_{k+1}}\right\rangle=-\alpha\tau S_{q}(\pi_{k+1})$ and defined the residual term $R_{q}(\pi_{k+1},\pi_{k}):=(1-q)\ln_{q}\!\frac{1}{\pi_{k+1}}\ln_{q}\!\pi_{k}$. When $q=2$, it is expected that the residual term remains negligible, but can become larger as $q$ increases. We visualize the trend of the residual $R_{q}(\pi_{k+1},\pi_{k})$ for $q=2,3,4,5$ on the CartPole-v1 environment (Brockman et al., 2016) in Figure 3. Learning consists of $2.5\times 10^{5}$ steps, evaluated every $2500$ steps (one iteration), averaged over 50 independent runs. It is visible that the magnitude of residual jumps from $q=4$ to $5$, while $q=2$ remains negligible throughout.

A reasonable approximation, therefore, is to use $\ln_{q}\pi_{k+1}$ and omit this residual term. Even this approximation, however, has an issue. When the actions are in the support, $\ln_{q}\!\,$ is the unique inverse function of $\exp_{q}$ and $\ln_{q}\pi_{k+1}$ yields $\frac{Q_{k}}{\tau}-\psi_{q}\left(\frac{Q_{k}}{\tau}\right)$. However, for actions outside the support, we cannot get the inverse, because many inputs to $\exp_{q}$ can result in zero. We could still use $\frac{Q_{k}}{\tau}-\psi_{q}\left(\frac{Q_{k}}{\tau}\right)$ as a sensible choice, and it appropriately does use negative values for the Munchausen term for these zero-probability actions. Empirically, however, we found this to be less effective than using the action gap.

Though the action gap is yet another approximation, there are clear similarities between using $\frac{Q_{k}}{\tau}-\psi_{q}\left(\frac{Q_{k}}{\tau}\right)$ and the action gap $Q_{k}-\mathcal{M}_{q,\tau}{Q_{k}}$. The primary difference is in how the values are centered. We can see $\psi_{q}$ as using a uniform average value of the actions in the support, as characterized in Theorem 2. $\mathcal{M}_{q,\tau}{Q_{k}}$, on the other hand, is a weighted average of action-values.

We plot the difference between $Q_{k}-\mathcal{M}_{q,\tau}{Q_{k}}$ and $\ln_{q}\pi_{k+1}$ in Figure 3, again in Cartpole. The difference stabilizes around -0.5 for most of learning—in other words primarily just shifting by a constant—but in early learning $\ln_{q}\pi_{k+1}$ is larger, across all $q$. This difference in magnitude might explain why using the action gap results in more stable learning, though more investigation is needed to truly understand the difference. For the purposes of this initial work, we pursue the use of the action gap, both as itself a natural extension of the current implementation of MVI and from our own experiments suggesting improved stability with this form.

We provide the overall performance of MVI versus MVI($q=2$) in Figure 5. Using $q=2$ provides a large improvement in about 5 games, about double the performance in the next 5 games, comparable performance in the next 7 games and then slightly worse performance in 3 games (PrivateEye, Chopper and Seaquest). Both PrivateEye and Seaquest are considered harder exploration games, which might explain this discrepancy. The Tsallis policy with $q=2$ reduces the support on actions, truncating some probabilities to zero. In general, with a higher $q$, the resulting policy is greedier, with $q=\infty$ corresponding to exactly the greedy policy. It is possible that for these harder exploration games, the higher stochasticity in the softmax policy from MVI whre $q=1$ promoted more exploration. A natural next step is to consider incorporating more directed exploration approaches, into MVI($q=2$), to benefit from the fact that lower-value actions are removed (avoiding taking poor actions) while exploring in a more directed way when needed.

We examine the learning curves for the games where MVI($q$) had the most significant improvement, in Figure 4. Particularly notable is how much more quickly MVI($q$) learned with $q=2$, in addition to plateauing at a higher point. In Hero, MVI($q$) learned a stably across the runs, whereas standard MVI with $q=1$ clearly has some failures.

These results are quite surprising. The algorithms are otherwise very similar, with the seemingly small change of using Munchausen term $Q_{k}(s,a)-\mathcal{M}_{q=2,\tau}{Q_{k}}$ instead of $Q_{k}(s,a)-\mathcal{M}_{q=1,\tau}{Q_{k}}$ and using the $q$-logarithm and $q$-exponential for the entropy regularization and policy parameterization. Previous work using $q=2$ to get the sparsemax with entropy regularization generally harmed performance (Lee et al., 2018, 2020). It seems that to get the benefits of the generalization to $q>1$, the addition of the KL regularization might be key. We validate this in the next section.

In the policy evaluation step of Eq. (11), if we set $\alpha=0$ then we recover Tsallis-VI which uses regularization $\Omega(\pi)=-S_{q}(\pi)$ in Eq. (1). In other words, we recover the algorithm that incorporates entropy regularization using the $q$-logarithm and the resulting sparsemax policy. Unlike MVI, Tsallis-VI has not been comprehensively evaluated on Atari games, so we include results for the larger benchmark set comprising 35 Atari games. We plot the percentage improvement of MVI($q$) over Tsallis-VI in Figure 5.

The improvement from including the Munchausen term ($\alpha>0$) is stark. For more than half of the games, MVI($q$) resulted in more than 100% improvement. For the remaining games it was comparable. For 10 games, it provided more than 400% improvement. Looking more specifically at which games there was notable improvement, it seems that exploration may again have played a role. MVI($q$) performs much better on Seaquest and PrivateEye. Both MVI($q$) and Tsallis-VI have policy parameterizations that truncate action support, setting probabilities to zero for some actions. The KL regularization term, however, likely slows this down. It is possible the Tsallis-VI is concentrating too quickly, resulting in insufficient exploration.