SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
f (c) \approx f (0) + f ^ {\prime} (0) (c - 0) + \frac {f ^ {\prime \prime} (0) (c - 0) ^ {2}}{2} = \frac {c ^ {2}}{2} \left\langle \delta , \nabla^ {2} \phi^ {*} (\nabla \phi (\pi_ {t} ^ {s})) \delta \right\rangle
$$
Let us first consider the softmax case with the log-sum-exp mirror map, where  $\pi^s = z(s,\cdot)$  and  $\phi (z(s,\cdot)) = \log (\sum_{a}\exp (z(s,a)))$ ,  $\phi^{*}(p^{\pi}(\cdot |s)) = \sum_{a}p^{\pi}(a|s)\log (p^{\pi}(a|s))$ . Since the negative entropy and log-sum-exp are Fenchel conjugates (see Lemma 9),  $\nabla \phi (z_t(s,\cdot)) = p^{\pi_t}(\cdot |s)$ . Hence, we need to compute  $\nabla^2\phi^* (p^{\pi_t}(\cdot |s))$
$$
\nabla \phi^ {*} (p ^ {\pi_ {t}} (\cdot | s)) = 1 + \log (p ^ {\pi_ {t}} (\cdot | s)) ; \quad \nabla^ {2} \phi^ {*} (p ^ {\pi_ {t}} (\cdot | s)) = \operatorname {d i a g} (1 / p ^ {\pi_ {t}} (\cdot | s))
$$
For the softmax representation, using the policy gradient theorem,  $[\delta]_a = p^{\pi_t}(a|s)[A(s,a) - \hat{A} (s,a)]$  and hence,
$$
\langle \delta , \nabla^ {2} \phi^ {*} (\nabla \phi (\pi_ {t})) \delta \rangle = \sum_ {a} p ^ {\pi_ {t}} (a | s) [ A (s, a) - \hat {A} (s, a) ] ^ {2}.
$$
Hence, for the softmax representation, the second-order Taylor series expansion around  $c = 0$  is equal to,
Now let us consider the direct case, where  $\pi^s = p^\pi (\cdot |s)$ ,  $\phi (p^{\pi}(\cdot |s)) = \sum_{a}p^{\pi}(a|s)\log (p^{\pi}(a|s))$ ,  $\phi^{*}(z(s,\cdot)) = \log (\sum_{a}\exp (z(s,a)))$ . Since the negative entropy and log-sum-exp are Fenchel conjugates (see Lemma 9),  $\nabla \phi (p^{\pi_t}(\cdot |s)) =$
$z_{t}(s,\cdot)$ . Hence, we need to compute  $\nabla^2\phi^* (z_t(s,\cdot))$
$$
\begin{array}{l} \left[ \nabla \phi^ {*} (z _ {t} (s, \cdot)) \right] _ {a} = \frac {\exp z _ {t} (s , a)}{\sum_ {a ^ {\prime}} \exp (z _ {t} (s , a ^ {\prime}))} = p ^ {\pi_ {t}} (a | s) \quad ; \quad \left[ \nabla^ {2} \phi^ {*} (z _ {t} (s, \cdot)) \right] _ {a, a} = p ^ {\pi_ {t}} (a | s) - \left[ p ^ {\pi_ {t}} (a | s) \right] ^ {2} \\ \left[ \nabla^ {2} \phi^ {*} (z _ {t} (s, \cdot)) \right] _ {a, a ^ {\prime}} = - p ^ {\pi_ {t}} (a | s) p ^ {\pi_ {t}} (a ^ {\prime} | s) \Rightarrow \nabla^ {2} \phi^ {*} (z _ {t} (s, \cdot)) = \operatorname {d i a g} (p ^ {\pi_ {t}} (\cdot | s)) - p ^ {\pi_ {t}} (\cdot | s) [ p ^ {\pi_ {t}} (\cdot | s) ] ^ {\intercal}. \\ \end{array}
$$
For the direct representation, using the policy gradient theorem,  $[\delta]_a = Q^{\pi_t}(s,a) - \hat{Q}^{\pi_t}(s,a)$  and hence,
$$
\begin{array}{l} \langle \delta , \nabla^ {2} \phi^ {*} (\nabla \phi (\pi_ {t})) \delta \rangle = [ Q ^ {\pi_ {t}} (s, \cdot) - \hat {Q} ^ {\pi_ {t}} (s, \cdot) ] ^ {\intercal} [ \operatorname {d i a g} (p ^ {\pi_ {t}} (\cdot | s)) - p ^ {\pi_ {t}} (\cdot | s) [ p ^ {\pi_ {t}} (\cdot | s) ] ^ {\intercal} ] [ Q ^ {\pi_ {t}} (s, \cdot) - \hat {Q} ^ {\pi_ {t}} (s, \cdot) ] \\ = \sum_ {a} p ^ {\pi_ {t}} (a | s) [ Q ^ {\pi_ {t}} (s, a) - \hat {Q} ^ {\pi_ {t}} (s, a) ] ^ {2} - \left[ \langle p ^ {\pi_ {t}} (a | s), Q ^ {\pi_ {t}} (s, \cdot) - \hat {Q} ^ {\pi_ {t}} (s, \cdot) \rangle \right] ^ {2} \\ = \sum_ {a} p ^ {\pi_ {t}} (a | s) [ Q ^ {\pi_ {t}} (s, a) - \hat {Q} ^ {\pi_ {t}} (s, a) ] ^ {2} - [ J _ {s} (\pi_ {t}) - \hat {J} _ {s} (\pi_ {t}) ] ^ {2} \\ \end{array}
$$
(where  $\hat{J}_s(\pi_t)$  is the estimated value function for starting state  $s$ )
Hence, for the direct representation, the second-order Taylor series expansion around  $c = 0$  is equal to,
$$
\begin{array}{l} f (c) \approx \frac {c ^ {2}}{2} \left[ \sum_ {a} p ^ {\pi_ {t}} (a | s) \left[ Q ^ {\pi_ {t}} (s, a) - \hat {Q} ^ {\pi_ {t}} (s, a) \right] ^ {2} - \left[ J _ {s} (\pi_ {t}) - \hat {J} _ {s} (\pi_ {t}) \right] ^ {2} \right] \\ = \frac {c ^ {2}}{2} \left[ \sum_ {a} p ^ {\pi_ {t}} (a | s) \left[ Q ^ {\pi_ {t}} (s, a) - \hat {Q} ^ {\pi_ {t}} (s, a) \right] ^ {2} - 2 \left[ J _ {s} (\pi_ {t}) - \hat {J} _ {s} (\pi_ {t}) \right] ^ {2} \sum_ {a} p ^ {\pi_ {t}} (a | s) + \left[ J _ {s} (\pi_ {t}) - \hat {J} _ {s} (\pi_ {t}) \right] ^ {2} \sum_ {a} p ^ {\pi_ {t}} (a | s) \right] \\ = \frac {c ^ {2}}{2} \left[ \sum_ {a} p ^ {\pi_ {t}} (a | s) [ Q ^ {\pi_ {t}} (s, a) - \hat {Q} ^ {\pi_ {t}} (s, a) ] ^ {2} - 2 [ J _ {s} (\pi_ {t}) - \hat {J} _ {s} (\pi_ {t}) ] \sum_ {a} p ^ {\pi_ {t}} (a | s) [ Q ^ {\pi_ {t}} (s, a) - \hat {Q} ^ {\pi_ {t}} (s, a) ] \right. \\ + \left[ J _ {s} \left(\pi_ {t}\right) - \hat {J} _ {s} \left(\pi_ {t}\right) \right] ^ {2} \sum_ {a} p ^ {\pi_ {t}} (a | s) \Bigg ] \\ = \frac {c ^ {2}}{2} \left[ \sum_ {a} p ^ {\pi_ {t}} (a | s) \left[ \left[ Q ^ {\pi_ {t}} (s, a) - \hat {Q} ^ {\pi_ {t}} (s, a) \right] ^ {2} - 2 \left[ J _ {s} (\pi_ {t}) - \hat {J} _ {s} (\pi_ {t}) \right] \left[ Q ^ {\pi_ {t}} (s, a) - \hat {Q} ^ {\pi_ {t}} (s, a) \right] + \left[ J _ {s} (\pi_ {t}) - \hat {J} _ {s} (\pi_ {t}) \right] ^ {2} \right] \right] \\ = \frac {c ^ {2}}{2} \left[ \sum_ {a} p ^ {\pi_ {t}} (a | s) \left(\left[ Q ^ {\pi_ {t}} (s, a) - \hat {Q} ^ {\pi_ {t}} (s, a) \right] - \left[ J _ {s} (\pi_ {t}) - \hat {J} _ {s} (\pi_ {t}) \right]\right) ^ {2} \right] \\ = \frac {c ^ {2}}{2} \left[ \sum_ {a} p ^ {\pi_ {t}} (a | s) \left[ A ^ {\pi_ {t}} (s, a) - \hat {A} ^ {\pi_ {t}} (s, a) \right] ^ {2} \right] \\ \end{array}
$$
E.1 Bandit examples to demonstrate the benefit of the decision-aware loss
Proposition 15 (Detailed version of Prop. 5). Consider a two-armed bandit example with deterministic rewards where arm 1 is optimal and has a reward  $r_1 = Q_1 = 2$  whereas arm 2 has reward  $r_2 = Q_2 = 1$ . Using a linear parameterization for the critic,  $Q$  function is estimated as:  $\hat{Q} = x\omega$  where  $\omega$  is the parameter to be learned and  $x$  is the feature of the corresponding arm. Let  $x_1 = -2$  and  $x_2 = 1$  implying that  $\hat{Q}_1(\omega) = -2\omega$  and  $\hat{Q}_2(\omega) = \omega$ . Let  $p_t$  be the probability of pulling the optimal arm at iteration  $t$ , and consider minimizing two alternative objectives to estimate  $\omega$ :
(1) Squared loss:  $\omega_{t}^{(1)} := \arg \min \left\{\frac{p_{t}}{2}\left[\hat{Q}_{1}(\omega) - Q_{1}\right]^{2} + \frac{1 - p_{t}}{2}\left[\hat{Q}_{2}(\omega) - Q_{2}\right]^{2}\right\}$ .
(2) Decision-aware critic loss:  $\omega_t^{(2)} := \arg \min \mathcal{L}_t(\omega) := p_t [Q_1 - \hat{Q}_1(\omega)] + (1 - p_t) [Q_2 - \hat{Q}_2(\omega)] + \frac{1}{c} \log \left(p_t \exp \left(-c [Q_1 - \hat{Q}_1(\omega)] + (1 - p_t) \exp \left(-c [Q_2 - \hat{Q}_2(\omega)]\right)\right)\right]$ .
Using the tabular parameterization for the actor, the policy update at iteration  $t$  is given by:  $p_{t+1} = \frac{p_t \exp(\eta \hat{Q}_1)}{p_t \exp(\eta \hat{Q}_1) + (1 - p_t) \exp(\eta \hat{Q}_2)}$ , where  $\eta$  is the functional step-size for the actor. For  $p_0 < \frac{2}{5}$ , minimizing the squared loss results in convergence to the sub-optimal action, while minimizing the decision-aware loss (for  $c, p_0 > 0$ ) results in convergence to the optimal action.
Proof. Note that  $\hat{Q}_1(\omega) - Q_1 = -2(\omega + 1)$  and  $\hat{Q}_2(\omega) - Q_2 = \omega - 1$ . Calculating  $\omega^{(1)}$  for a general policy s.t. the probability of pulling the optimal arm equal to  $p$ ,
$$
\begin{array}{l} \operatorname {M S E} (\omega) = \frac {p}{2} \left[ \hat {Q} _ {1} (\omega) - Q _ {1} \right] ^ {2} + \frac {1 - p}{2} \left[ \hat {Q} _ {2} (\omega) - Q _ {2} \right] ^ {2} = \frac {1}{2} \left[ 4 p (\omega + 1) ^ {2} + (1 - p) (\omega - 1) ^ {2} \right] \\ \Longrightarrow \nabla_ {\omega} \mathrm {M S E} (\omega) = 4 p (\omega + 1) + (1 - p) (\omega - 1) \\ \end{array}
$$
Setting the gradient to zero,
$$
\Rightarrow \omega^ {(1)} = \frac {1 - 5 p}{3 p + 1}
$$
Calculating  $\omega^{(2)}$  for a general policy s.t. the probability of pulling the optimal arm equal to  $p$ ,
$$
\begin{array}{l} L _ {t} (\omega) = 2 p (\omega + 1) - (1 - p) (\omega - 1) + \frac {1}{c} \log (p \exp (- 2 c (\omega + 1)) + (1 - p) \exp (c (\omega - 1))) \\ \Rightarrow \nabla_ {\omega} L _ {t} (\omega) = (3 p - 1) + \frac {1}{c} \nabla_ {\omega} [ \log (p \exp (- 2 c (\omega + 1)) + (1 - p) \exp (c (\omega - 1))) ] \\ \end{array}
$$
$$
\Rightarrow \nabla_ {\omega} [ \log (p \exp (- 2 c (\omega + 1)) + (1 - p) \exp (c (\omega - 1))) ] = (1 - 3 p) c
$$
Define  $A\coloneqq \exp (-2c(\omega +1))$  and  $B\coloneqq \exp (c(\omega -1)))$
$$
\Rightarrow \frac {- 2 p c A + (1 - p) c B}{p A + (1 - p) B} = (1 - 3 p) c \Rightarrow \frac {A}{p A + (1 - p) B} = 1 \Rightarrow \omega^ {(2)} = \frac {- 1}{3}.
$$
Now, let us consider the actor update,
$$
p _ {t + 1} = \frac {p _ {t} \exp (\eta \hat {Q} _ {1})}{p _ {t} \exp (\eta \hat {Q} _ {1}) + (1 - p _ {t}) \exp (\eta \hat {Q} _ {2})} \Rightarrow \frac {p _ {t + 1}}{p _ {t}} = \frac {1}{p _ {t} + (1 - p _ {t}) \exp (\eta (\hat {Q} _ {2} - \hat {Q} _ {1}))}
$$
Since arm 1 is optimal, if  $\frac{p_{t+1}}{p_t} < 1$  for all  $t$ , the algorithm will converge to the sub-optimal arm. This happens when  $\frac{1}{p_t + (1 - p_t)\exp(\eta(\hat{Q}_2 - \hat{Q}_1))} < 1 \Rightarrow \hat{Q}_2 - \hat{Q}_1 > 0 \Rightarrow \omega > 0$ . Hence, for any  $\eta$  and any iteration  $t$ , if  $\omega_t > 0$ ,  $p_{t+1} < p_t$ .
For the decision-aware critic loss,  $\omega_{t}^{(2)} = -\frac{1}{3}$  for all  $t$ , implying that  $p_{t + 1} > p_t$  and hence the algorithm will converge to the optimal policy for any  $\eta$  and any initialization  $p_0 > 0$ . However, for the squared MSE loss,  $\omega_{t}^{(2)} = \frac{1 - 5p_{t}}{3p_{t} + 1}$ ,  $\omega_{t}^{(2)} > 0$  if  $p_t < \frac{1}{5}$ . Hence, if  $p_0 < \frac{1}{5}$ ,  $p_1 < p_0 < \frac{1}{5}$ . Using the same reasoning,  $p_2 < p_1 < p_0 < 1/5$ , and hence the policy will converge to the sub-optimal arm.
Proposition 16. Consider two-armed bandit problem with deterministic rewards - arm 1 has a reward  $r_1 = Q_1$  whereas arm 2 has a reward  $r_2 = Q_2$  such that arm 1 is the optimal, i.e.  $Q_1 \geq Q_2$ . Using a linear parameterization for the critic,  $Q$  function is estimated as:  $\hat{Q} = x\omega$  where  $\omega$  is the parameter to be learned and  $x$  is the feature of the corresponding arm. Let  $p_t$  be the probability of pulling the optimal arm at iteration  $t$ , and consider minimizing the decision-aware critic loss to estimate  $\omega$ :  $\omega_t := \arg \min_{t \leq t_0} \mathcal{L}_t(\omega) := p_t [Q_1 - \hat{Q}_1(\omega)] + (1 - p_t) [Q_2 - \hat{Q}_2(\omega)] + \frac{1}{c} \log \left(p_t \exp \left(-c [Q_1 - \hat{Q}_1(\omega)] + (1 - p_t) \exp \left(-c [Q_2 - \hat{Q}_2(\omega)]\right)\right)\right]$ . Using the tabular parameterization for the actor, the policy update at iteration  $t$  is given by:  $p_t = \frac{p_t \exp(\eta \hat{Q}_1)}{p_t \exp(\eta \hat{Q}_1) + (1 - p_t) \exp(\eta \hat{Q}_2)}$ , where  $\eta$  is the functional step-size for the actor. For the above problem, minimizing the decision-aware loss (for  $c, p_0 > 0$ ) results in convergence to the optimal action, and  $\mathcal{L}_t(\omega_t) = 0$  for any iteration  $t$ .
Proof. Define  $A \coloneqq \exp(-c[Q_1 - \hat{Q}_1(\omega)])$  and  $B \coloneqq \exp(-c[Q_2 - \hat{Q}_2(\omega)])$ . Calculating the gradient of  $\mathcal{L}_t$  w.r.t  $\omega$  and setting it to zero,
$$
\begin{array}{l} \nabla_ {\omega} \mathcal {L} _ {t} (\omega) = p _ {t} x _ {1} + (1 - p _ {t}) x _ {2} - \frac {p _ {t} x _ {1} A + (1 - p _ {t}) x _ {2} B}{p _ {t} A + (1 - p _ {t}) B} = 0 \\ \Rightarrow p _ {t} \left(1 - p _ {t}\right) A \left(x _ {1} - x _ {2}\right) = p _ {t} \left(1 - p _ {t}\right) B \left(x _ {1} - x _ {2}\right) \\ \Rightarrow Q _ {1} - x _ {1} \omega_ {t} = Q _ {2} - x _ {2} \omega_ {t} \Rightarrow \omega_ {t} = \frac {Q _ {1} - Q _ {2}}{x _ {1} - x _ {2}}. \\ \end{array}
$$
Observe that  $Q_{1} - \hat{Q}_{1}(\omega_{t}) = Q_{2} - \hat{Q}_{2}(\omega_{t})$  and thus  $\mathcal{L}_t(\omega_t) = 0$  for all  $t$ . Writing the actor update,
$$
\begin{array}{l} p _ {t + 1} = \frac {p _ {t} \exp (\eta x _ {1} \omega_ {t})}{p _ {t} \exp (\eta x _ {1} \omega_ {t}) + (1 - p _ {t}) \exp (\eta x _ {2} \omega_ {t})} \\ \Rightarrow \frac {p _ {t + 1}}{p _ {t}} = \frac {1}{p _ {t} + \left(1 - p _ {t}\right) \exp \left(\eta \left(x _ {2} - x _ {1}\right) \omega_ {t}\right)} = \frac {1}{p _ {t} + \left(1 - p _ {t}\right) \exp \left(\eta \left(Q _ {2} - Q _ {1}\right)\right)} \geq 1 \\ \end{array}
$$
Proposition 17 (Detailed version of Prop. 7). Consider a two-armed bandit example and define  $p \in [0,1]$  as the probability of pulling arm 1. Given  $p$ , let the advantage of arm 1 be equal to  $A_1 := \frac{1}{2} > 0$ , while that of arm 2 is  $A_2 := -\frac{p}{2(1-p)} < 0$  implying that arm 1 is optimal. For the critic, consider approximating the advantage of the two arms using a discrete hypothesis class with two hypotheses that depend on  $p$  for:  $\mathcal{H}_0: \hat{A}_1 = \frac{1}{2} + \varepsilon, \hat{A}_2 = -\frac{p}{1-p} \left( \frac{1}{2} + \varepsilon \right)$  and  $\mathcal{H}_1: \hat{A}_1 = \frac{1}{2} - \varepsilon \operatorname{sgn} \left( \frac{1}{2} - p \right), \hat{A}_2 = -\frac{p}{1-p} \left( \frac{1}{2} - \varepsilon \operatorname{sgn} \left( \frac{1}{2} - p \right) \right)$  where  $\operatorname{sgn}$  is the signum function and  $\varepsilon \in \left( \frac{1}{2}, 1 \right)$ . If  $p_t$  is the probability of pulling arm 1 at iteration  $t$ , consider minimizing two alternative loss functions to choose the hypothesis  $\mathcal{H}_t$ :
(1) Squared (MSE) loss:  $\mathcal{H}_t = \arg \min_{\{\mathcal{H}_0,\mathcal{H}_1\}}\left\{\frac{p_t}{2} [A_1 - \hat{A}_1]^2 +\frac{1 - p_t}{2} [A_2 - \hat{A}_2]^2\right\}$ .
(2) Decision-aware critic loss (DA) with  $c = 1:\mathcal{H}_t = \arg \min_{\{\mathcal{H}_0,\mathcal{H}_1\}}$
$$
\left\{p _ {t} \left(1 - \left[ A _ {1} - \hat {A} _ {1} \right]\right) \log \left(1 - \left[ A _ {1} - \hat {A} _ {1} \right]\right) + \left(1 - p _ {t}\right) \left(1 - \left[ A _ {2} - \hat {A} _ {2} \right]\right) \log \left(1 - \left[ A _ {2} - \hat {A} _ {2} \right]\right) \right\}.
$$
Using the tabular parameterization for the actor, the policy update at iteration  $t$  is given by:  $p_{t + 1} = \frac{p_t(1 + \eta\hat{A}_1)}{p_t(1 + \eta\hat{A}_1) + (1 - p_t)(1 + \eta\hat{A}_2)}$ . For  $p_0 \leq \frac{1}{2}$ , the squared loss cannot distinguish between  $\mathcal{H}_0$  and  $\mathcal{H}_1$ , and depending on how ties are broken, minimizing it can result in convergence to the sub-optimal action. On the other hand, minimizing the divergence loss (for any  $p_0 > 0$ ) results in convergence to the optimal arm.
Proof. First note that when  $p > \frac{1}{2}$ ,  $\mathcal{H}_0$  and  $\mathcal{H}_1$  are identical, ensure that  $\hat{A}_1 > \hat{A}_2$  and the algorithm will converge to the optimal arm no matter which hypothesis is chosen. The regime of interest is therefore when  $p \leq \frac{1}{2}$  and we focus on this case. Let us calculate the MSE and decision-aware (DA) losses for  $\mathcal{H}_0$ .
$$
A _ {1} - \hat {A} _ {1} = \frac {1}{2} - \left(\frac {1}{2} + \varepsilon\right) = - \varepsilon ; A _ {2} - \hat {A} _ {2} = - \frac {p}{1 - p} [ 1 - (1 + \varepsilon) ] = \frac {p}{1 - p} \varepsilon
$$
$$
\operatorname {M S E} \left(\hat {A} _ {1}, \hat {A} _ {2}\right) = p \varepsilon^ {2} + (1 - p) \left(\frac {p}{1 - p} \varepsilon\right) ^ {2} = p \varepsilon^ {2} + \frac {\varepsilon^ {2} p ^ {2}}{1 - p}
$$
$$
\mathrm {D A} (\hat {A} _ {1}, \hat {A} _ {2}) = p (1 + \varepsilon) \log (1 + \varepsilon) + (1 - p) \left(1 - \frac {\varepsilon p}{1 - p}\right) \log \left(1 - \frac {\varepsilon p}{1 - p}\right)
$$
Similarly, we can calculate the MSE and decision-aware losses for  $\mathcal{H}_1$ .
$$
A _ {1} - \hat {A} _ {1} = \frac {1}{2} - \left(\frac {1}{2} - \varepsilon\right) = \varepsilon ; A _ {2} - \hat {A} _ {2} = - \frac {p}{1 - p} \left[ \frac {1}{2} - \left(\frac {1}{2} - \varepsilon\right) \right] = - \frac {p}{1 - p} \varepsilon
$$
$$
\mathrm {D A} (\hat {A} _ {1}, \hat {A} _ {2}) = p (1 - \varepsilon) \log (1 - \varepsilon) + (1 - p) \left(1 + \frac {\varepsilon p}{1 - p}\right) \log \left(1 + \frac {\varepsilon p}{1 - p}\right)
$$
For both  $\mathcal{H}_0$  and  $\mathcal{H}_1$ , the MSE loss is equal to  $p\varepsilon^2 + \frac{\varepsilon^2p^2}{1 - p}$  and hence it cannot distinguish between the two hypotheses. Writing the actor update,
$$
p _ {t + 1} = \frac {p _ {t} (1 + \eta \hat {A} _ {1})}{p _ {t} (1 + \eta \hat {A} _ {1}) + (1 - p _ {t}) (1 + \eta \hat {A} _ {2})} \Rightarrow \frac {p _ {t + 1}}{p _ {t}} = \frac {1}{p _ {t} + (1 - p _ {t}) \frac {1 + \eta \hat {A} _ {2}}{1 + \eta \hat {A} _ {1}}}
$$
Hence, in order to ensure that  $p_{t+1} > p_t$  and eventual convergence to the optimal arm, we want that  $\hat{A}_2 < \hat{A}_1$ . For  $\varepsilon \in \left(\frac{1}{2}, 1\right)$ , for  $\mathcal{H}_0, \hat{A}_1 > 0$  while  $\hat{A}_2 < 0$ . On the other hand, for  $\mathcal{H}_1, \hat{A}_1 < 0$  and  $\hat{A}_2 > 0$ . This implies that the algorithm should choose  $\mathcal{H}_0$  in order to approximate the advantage. Since the MSE loss is the same for both hypotheses, convergence to the optimal arm depends on how the algorithm breaks ties. Next, we prove that for the decision-aware loss and any iteration such that  $p_t < 0.5$ , the loss for  $\mathcal{H}_0$  is smaller than that for  $\mathcal{H}_1$ , and hence the algorithm chooses the correct hypothesis and pulls the optimal arm. For this, we define  $f(p)$  as follows,
$$
\begin{array}{l} f (p) := \left[ p \left(1 + \varepsilon\right) \log (1 + \varepsilon) + (1 - p) \left(1 - \frac {\varepsilon p}{1 - p}\right) \log \left(1 - \frac {\varepsilon p}{1 - p}\right) \right] \\ - \left[ p (1 - \varepsilon) \log (1 - \varepsilon) + (1 - p) \left(1 + \frac {\varepsilon p}{1 - p}\right) \log \left(1 + \frac {\varepsilon p}{1 - p}\right) \right] \\ \end{array}
$$
For  $f(p)$  to be well-defined, we want that,  $1 - \epsilon > 0 \Rightarrow \epsilon < 1$  and  $1 - \frac{\epsilon p}{1 - p} > 0 \Rightarrow p < \frac{1}{1 + \epsilon}$ . Since  $\epsilon \in (1/2, 1)$ ,  $p < \frac{1}{2}$ . In order to prove that the algorithm will always choose  $\mathcal{H}_0$ , we will show that  $f(p) \leq 0$  for all  $p \in [0, 1/2]$  next. First note that,
$$
f (0) = 0 \quad ; \quad f ^ {(1 / 2)} = \frac {1 + \varepsilon}{2} \log (1 + \varepsilon) + \frac {(1 - \varepsilon)}{2} \log (1 - \varepsilon) - \frac {1 - \varepsilon}{2} \log (1 - \varepsilon) - \frac {1 + \varepsilon}{2} \log (1 + \varepsilon) = 0
$$
Next, we will prove that  $f(p)$  is convex. This combined with the fact  $f(0) = f(1/2) = 0$  implies that  $f(p) < 0$  for all  $p \in (0, 1/2)$ . For this, we write  $f(p) = g(p) + h_1(p) - h_2(p)$  where,
$$
g (p) = p (1 + \varepsilon) \log (1 + \varepsilon) - p (1 - \varepsilon) \log (1 - \varepsilon)
$$
$$
h _ {1} (p) = (1 - p) \left(1 - \frac {\varepsilon p}{1 - p}\right) \log \left(1 - \frac {\varepsilon p}{1 - p}\right) = (1 - \epsilon^ {\prime} p) \log \left(\frac {1 - \epsilon^ {\prime} p}{1 - p}\right) \quad (\epsilon^ {\prime} = 1 + \epsilon)
$$
$$
h _ {2} (p) = (1 - p) \left(1 + \frac {\varepsilon p}{1 - p}\right) \log \left(1 + \frac {\varepsilon p}{1 - p}\right) = (1 - \epsilon^ {\prime \prime} p) \log \left(\frac {1 - \epsilon^ {\prime \prime} p}{1 - p}\right) \quad (\epsilon^ {\prime \prime} = 1 - \epsilon)
$$
Differentiating the above terms,
$$
g ^ {\prime} (p) = (1 + \varepsilon) \log (1 + \varepsilon) - (1 - \varepsilon) \log (1 - \varepsilon); g ^ {\prime \prime} (p) = 0
$$
$$
h _ {1} ^ {\prime} (p) = - \epsilon^ {\prime} \log \left(\frac {1 - \epsilon^ {\prime} p}{1 - p}\right) + \frac {1 - \epsilon^ {\prime}}{1 - p}
$$
$$
h _ {1} ^ {\prime \prime} (p) = - \frac {\epsilon^ {\prime}}{1 - \epsilon^ {\prime} p} \frac {1 - \epsilon^ {\prime}}{1 - p} - \frac {(1 - \epsilon^ {\prime}) ^ {2}}{(1 - p) ^ {2}} = \frac {(\epsilon^ {\prime} - 1) p \left[ (\epsilon^ {\prime} - 1) ^ {2} + \left(\frac {1}{p} - 1\right) \right]}{(1 - \epsilon^ {\prime} p) (1 - p) ^ {2}} > 0
$$
$$
h _ {2} ^ {\prime \prime} (p) = \frac {(\epsilon^ {\prime \prime} - 1) p \left[ (\epsilon^ {\prime \prime} - 1) ^ {2} + \left(\frac {1}{p} - 1\right) \right]}{(1 - \epsilon^ {\prime \prime} p) (1 - p) ^ {2}} <   0
$$
Combining the above terms,  $f''(p) = g''(p) + h_1''(p) - h_2''(p) > 0$  for all  $p \in (0,1/2)$  and hence  $f(p)$  is convex. Hence, for all  $p < \frac{1}{2}$ , minimizing the divergence loss results in choosing  $\mathcal{H}_0$  and the actor pulling the optimal arm. Once the probability of
pulling the optimal arm is larger than 0.5, both hypotheses are identical and the algorithm will converge to the optimal arm regardless of the hypothesis chosen.
E.2 Lemmas
Lemma 7. For a probability distribution  $p \in \mathbb{R}^A$ , the negative entropy mirror map  $\phi(p) = \sum_i p_i \log(p_i)$ ,  $\delta \in \mathbb{R}^A$ ,  $c > 0$ ,
$$
D _ {\phi^ {*}} \left(\nabla \phi (p) - c \delta , \nabla \phi (p)\right) = c \langle p, \delta \rangle + \log \left(\sum_ {j} p _ {j} \exp (- c \delta_ {j})\right).
$$
Proof. In this case,  $[\nabla \phi(p)]_i = 1 + \log(p_i)$ . Hence, we need to compute  $D_{\phi^*}(z', z)$  where  $z_i' := 1 + \log(p_i) - c\delta_i$  and  $z_i := 1 + \log(p_i)$ . If  $\phi(p) = \sum_i p_i \log(p_i)$ , using Lemma 9,  $\phi^*(z) = \log\left(\sum_i \exp(z_i)\right)$  where  $z_i - \log\left(\sum_i \exp(z_i)\right) = \log(p_i)$ . Define distribution  $q$  such that  $q_i := \frac{\exp(1 + \log(p_i) - c\delta_i)}{\sum_j \exp(1 + \log(p_j) - c\delta_j)}$ . Using Lemma 11,
$$
D _ {\phi^ {*}} \left(z ^ {\prime}, z\right) = \mathrm {K L} \left(p \| q\right) = \sum_ {i} p _ {i} \log \left(\frac {p _ {i}}{q _ {i}}\right)
$$
Simplifying  $q$
$$
\begin{array}{l} q _ {i} = \frac {\exp (1 + \log (p _ {i}) - c \delta_ {i})}{\exp (\sum_ {j} (1 + \log (p _ {j}) - c \delta_ {j}))} = \frac {p _ {i} \exp (- c \delta_ {i})}{\sum_ {j} p _ {j} \exp (- c \delta_ {j})} \\ \Rightarrow D _ {\phi^ {*}} (z ^ {\prime}, z) = \sum_ {i} p _ {i} \log \left(\frac {p _ {i} \sum_ {j} p _ {j} \exp (- c \delta_ {j})}{p _ {i} \exp (- c \delta_ {i})}\right) = \sum_ {i} p _ {i} \log \left(\exp (c \delta_ {i}) \sum_ {j} p _ {j} \exp (- c \delta_ {j})\right) \\ = c \sum_ {i} p _ {i} \delta_ {i} + \sum_ {i} p _ {i} \log \left(\sum_ {j} p _ {j} \exp (- c \delta_ {j})\right) = c \langle p, \delta \rangle + \log \left(\sum_ {j} p _ {j} \exp (- c \delta_ {j})\right) \\ \end{array}
$$
Lemma 8. For  $z\in \mathbb{R}^A$  , the log-sum-exp mirror map  $\phi (z) = \log (\sum_{i}\exp (z_{i}))$ $\delta \in \mathbb{R}^{A}$  s.t.  $\begin{array}{r}\sum_{i}\delta_{i} = 0,c > 0, \end{array}$
$$
D _ {\phi^ {*}} \Bigg (\nabla \phi (z) - c \delta , \nabla \phi (z) \Bigg) = \sum_ {i} \left(p _ {i} - c \delta_ {i}\right) \log \left(\frac {p _ {i} - c \delta_ {i}}{p _ {i}}\right),
$$
where  $p_i = \frac{\exp(z_i)}{\sum_j\exp(z_j)}$
Proof. In this case,  $[\nabla \phi (z)]_i = \frac{\exp(z_i)}{\sum_j\exp(z_j)} = p_i$ . Define distribution  $q$  s.t.  $q_{i}\coloneqq p_{i} - c\delta_{i}$ . Note that since  $\sum_{i}\delta_{i} = 0$ ,  $\sum_{i}q_{i} = \sum_{i}p_{i} = 1$  and hence,  $q$  is a valid distribution. We thus need to compute  $D_{\phi^{*}}(q,p)$ . Using Lemma 9,  $\phi^{*}(p) = \sum_{i}p_{i}\log (p_{i})$  where  $p_i = \frac{\exp(z_i)}{\sum_j\exp(z_j)}$ . Using Lemma 10,
$$
D _ {\phi^ {*}} (q, p) = \mathrm {K L} (q | | p) = \sum_ {i} \left(p _ {i} - c \delta_ {i}\right) \log \left(\frac {p _ {i} - c \delta_ {i}}{p _ {i}}\right)
$$
Lemma 9. The log-sum-exp mirror map on the logits and the negative entropy mirror map on the corresponding probability distribution are Fenchel duals. In particular for  $z \in \mathbb{R}^d$ , if  $\phi(z) := \log \left( \sum_{i} \exp(z_i) \right)$ , then  $\phi^*(p) = \sum_{i} p_i \log(p_i)$  where  $p_i = \frac{\exp(z_i)}{\sum_j \exp(z_j)}$ . Similarly, if  $\phi(p) = \sum_{i} p_i \log(p_i)$ , then  $\phi^*(z) = \log \left( \sum_{i} \exp(z_i) \right)$  where  $z_i - \log \left( \sum_{i} \exp(z_i) \right) = \log(p_i)$ .
Proof. If  $\phi(z) \coloneqq \log\left(\sum_{i} \exp(z_i)\right)$ ,
$$
\phi^ {*} (p) := \sup _ {z} [ \langle p, z \rangle - \phi (z) ] = \sup _ {z} \left[ \sum_ {i} p _ {i} z _ {i} - \log \left(\sum_ {i} \exp (z _ {i})\right) \right]
$$