SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Setting the gradient to zero, we get that  $p_i = \frac{\exp(z_i^*)}{\sum_j \exp(z_j^*)}$  for  $z^* \in \mathcal{Z}^*$  where  $\mathcal{Z}^*$  is the set of maxima related by a shift (i.e. if  $z^* \in \mathcal{Z}^*$ ,  $z^* + C \in \mathcal{Z}^*$  for a constant  $C$ ). Using the optimality condition, we know that  $\sum_i p_i = 1$  and
$$
\log (p _ {i}) = z _ {i} ^ {*} - \log \left(\sum_ {j} \exp (z _ {j} ^ {*})\right) \Rightarrow z _ {i} ^ {*} = \log (p _ {i}) + \phi (z ^ {*})
$$
Using this relation,
$$
\begin{array}{l} \phi^ {*} (p) = \left[ \sum_ {i} p _ {i} z _ {i} ^ {*} - \log \left(\sum_ {i} \exp \left(z _ {i} ^ {*}\right)\right) \right] = \left[ \sum_ {i} p _ {i} \log \left(p _ {i}\right) + \phi \left(z ^ {*}\right) \sum_ {i} p _ {i} - \phi \left(z ^ {*}\right) \right] \\ \Rightarrow \phi^ {*} (p) = \sum_ {i} p _ {i} \log \left(p _ {i}\right) \\ \end{array}
$$
The second statement follows since the  $\phi^{*}(\phi^{*}) = \phi$
Lemma 10. For probability distributions,  $p$  and  $p'$ , if  $\phi(p) = \sum_{i} p \log(p_i)$ , then  $D_{\phi}(p, p') = KL(p||p')$ .
Proof. Note that  $[\nabla \phi(p)]_i = 1 + \log(p_i)$ . Using the definition of the Bregman divergence,
$$
\begin{array}{l} D _ {\phi} (p, p ^ {\prime}) := \phi (p) - \phi (p ^ {\prime}) - \langle \nabla \phi (p ^ {\prime}), p - p ^ {\prime} \rangle \\ = \sum_ {i} \left[ p _ {i} \log \left(p _ {i}\right) - p _ {i} ^ {\prime} \log \left(p _ {i} ^ {\prime}\right) - \left(1 + \log \left(p _ {i} ^ {\prime}\right)\right) \left(p _ {i} - p _ {i} ^ {\prime}\right) \right] \\ = \sum_ {i} \left[ p _ {i} \log \left(\frac {p _ {i}}{p _ {i} ^ {\prime}}\right) \right] - \sum_ {i} p _ {i} + \sum_ {i} p _ {i} ^ {\prime} \\ \end{array}
$$
Since  $p$  and  $p'$  are valid probability distributions,  $\sum_{i} p_{i} = \sum_{i} p_{i}' = 1$ , and hence,  $D_{\phi}(p, p') = \mathrm{KL}(p||p')$ .
Lemma 11. If  $\phi(z) = \log(\sum_{i} \exp(z_i))$ , then  $D_{\phi}(z, z') = KL(p' || p)$ , where  $p_i := \frac{\exp(z_i)}{\sum_j \exp(z_j)}$  and  $p_i' := \frac{\exp(z_i')}{\sum_j \exp(z_j')}$ .
Proof. Note that  $\left[\nabla \phi (z)\right]_i = \frac{\exp(z_i)}{\sum_j\exp(z_j)} = p_i$  where  $p_i\coloneqq \frac{\exp(z_i)}{\sum_j\exp(z_j)}$ . Using the definition of the Bregman divergence,
$$
\begin{array}{l} D _ {\phi} (z, z ^ {\prime}) := \phi (z) - \phi (z ^ {\prime}) - \langle \nabla \phi (z ^ {\prime}), z - z ^ {\prime} \rangle \\ = \log \left(\sum_ {j} \exp \left(z _ {j}\right)\right) - \log \left(\sum_ {j} \exp \left(z _ {j} ^ {\prime}\right)\right) - \sum_ {i} \left[ \frac {\exp \left(z _ {i} ^ {\prime}\right)}{\sum_ {j} \exp \left(z _ {j} ^ {\prime}\right)} \left(z _ {i} - z _ {i} ^ {\prime}\right) \right] \\ = \sum_ {i} p _ {i} ^ {\prime} \left[ \log \left(\sum_ {j} \exp (z _ {j})\right) - \log \left(\sum_ {j} \exp \left(z _ {j} ^ {\prime}\right)\right) - z _ {i} + z _ {i} ^ {\prime} \right] \quad (\text {S i n c e} \sum_ {i} p _ {i} ^ {\prime} = 1) \\ = \sum_ {i} p _ {i} ^ {\prime} \left[ \log \left(\sum_ {j} \exp (z _ {j})\right) - \log \left(\sum_ {j} \exp (z _ {j} ^ {\prime})\right) - \log (\exp (z _ {i})) + \log (\exp (z _ {i} ^ {\prime})) \right] \\ = \sum_ {i} p _ {i} ^ {\prime} \left[ \log \left(\frac {\exp \left(z _ {i} ^ {\prime}\right)}{\sum_ {j} \exp \left(z _ {j} ^ {\prime}\right)}\right) - \log \left(\frac {\exp \left(z _ {i}\right)}{\sum_ {j} \exp \left(z _ {j}\right)}\right) \right] \\ = \sum_ {i} p _ {i} ^ {\prime} \left[ \log \left(p _ {i} ^ {\prime}\right) - \log \left(p _ {i} ^ {\prime}\right) \right] = \sum_ {i} p _ {i} ^ {\prime} \left[ \log \left(\frac {p _ {i} ^ {\prime}}{p _ {i}}\right) \right] = \mathrm {K L} (p ^ {\prime} | | p) \\ \end{array}
$$
F.1 Heuristic to estimate  $c$
We estimate  $c$  to maximize the lower-bound on  $J(\pi)$ . In particular, using Prop. 1,
$$
J (\pi) \geq J (\pi_ {t}) + \hat {g} (\pi_ {t}) ^ {\top} (\pi - \pi_ {t}) - \left(\frac {1}{\eta} + \frac {1}{c}\right) D _ {\Phi} (\pi , \pi_ {t}) - \frac {1}{c} D _ {\Phi^ {*}} \left(\nabla \Phi (\pi_ {t}) - c [ \nabla J (\pi_ {t}) - \hat {g} (\pi_ {t}) ], \nabla \Phi (\pi_ {t})\right)
$$
For a fixed  $\hat{g} (\pi_t)$ , we need to maximize the RHS w.r.t  $\pi$  and  $c$ , i.e.
$$
\max  _ {c > 0} \max  _ {\pi \in \Pi} J (\pi_ {t}) + \hat {g} (\pi_ {t}) ^ {\top} (\pi - \pi_ {t}) - \left(\frac {1}{\eta} + \frac {1}{c}\right) D _ {\Phi} (\pi , \pi_ {t}) - \frac {1}{c} D _ {\Phi^ {*}} \left(\nabla \Phi (\pi_ {t}) - c [ \nabla J (\pi_ {t}) - \hat {g} (\pi_ {t}) ], \nabla \Phi (\pi_ {t})\right) \tag {6}
$$
Instead of maximizing w.r.t  $\pi$  and  $c$ , we will next aim to find an upper-bound on the RHS that is independent of  $\pi$  and aim to maximize it w.r.t  $c$ . Using Lemma 1 with  $y' = \pi$ ,  $y = \pi_t$ ,  $x = -\hat{g}(\pi_t)$  and define  $c'$  such that  $\frac{1}{c'} = \frac{1}{\eta} + \frac{1}{c}$ .
$$
\begin{array}{l} \langle - \hat {g} (\pi_ {t}), \pi - \pi_ {t} \rangle \geq - \frac {1}{c ^ {\prime}} \left[ D _ {\Phi} (\pi , \pi_ {t}) + D _ {\Phi} ^ {*} (\nabla \Phi (\pi_ {t}) + c ^ {\prime} \hat {g} (\pi_ {t}), \nabla \Phi (\pi_ {t})) \right] \\ \Longrightarrow J (\pi_ {t}) + \langle \hat {g} (\pi_ {t}), \pi - \pi_ {t} \rangle - \left(\frac {1}{\eta^ {\prime}} + \frac {1}{c}\right) D _ {\Phi} (\pi , \pi_ {t}) \leq J (\pi_ {t}) + \frac {1}{c ^ {\prime}} D _ {\Phi} ^ {*} (\nabla \Phi (\pi_ {t}) + c ^ {\prime} \hat {g} (\pi_ {t}), \nabla \Phi (\pi_ {t})) \\ \end{array}
$$
Using the above upper-bound in Eq. (6),
$$
\max _ {c > 0} \left[ J (\pi_ {t}) + \frac {1}{c ^ {\prime}} D _ {\Phi} ^ {*} (\nabla \Phi (\pi_ {t}) + c ^ {\prime} \hat {g} (\pi_ {t}), \nabla \Phi (\pi_ {t})) - \frac {1}{c} D _ {\Phi^ {*}} \Bigl (\nabla \Phi (\pi_ {t}) - c [ \nabla J (\pi_ {t}) - \hat {g} (\pi_ {t}) ], \nabla \Phi (\pi_ {t}) \Bigr) \right]
$$
This implies that the estimate  $\hat{c}$  can be calculated as:
$$
\hat {c} = \underset {c > 0} {\arg \max} \left\{\left(\frac {1}{\eta} + \frac {1}{c}\right) D _ {\Phi} ^ {*} \left(\nabla \Phi (\pi_ {t}) + \frac {1}{\left(\frac {1}{\eta} + \frac {1}{c}\right)} \hat {g} (\pi_ {t}), \nabla \Phi (\pi_ {t})\right) - \frac {1}{c} D _ {\Phi^ {*}} \left(\nabla \Phi (\pi_ {t}) - c [ \nabla J (\pi_ {t}) - \hat {g} (\pi_ {t}) ], \nabla \Phi (\pi_ {t})\right) \right\}
$$
In order to gain some intuition, let us consider the case where  $D_{\Phi}(u,v) = \frac{1}{2}\| u - v\| _2^2$ . In this case,
$$
\hat {c} = \underset {c > 0} {\arg \max} \left\{\frac {\| \hat {g} (\pi_ {t}) \| _ {2} ^ {2}}{2 \left(\frac {1}{\eta} + \frac {1}{c}\right)} - \frac {c}{2} \| \hat {g} (\pi_ {t}) - \nabla J (\pi_ {t}) \| _ {2} ^ {2} \right\}
$$
If  $\| \hat{g} (\pi_t) - \nabla J(\pi_t)\| _2^2\to 0$
$$
\hat {c} = \underset {c > 0} {\arg \max } \left\{\frac {\| \hat {g} (\pi_ {t}) \| _ {2} ^ {2}}{2 \left(\frac {1}{\eta} + \frac {1}{c}\right)} \right\} \Rightarrow c \rightarrow \infty
$$
If  $\| \hat{g} (\pi_t) - \nabla J(\pi_t)\| _2^2\to \infty$
$$
\hat {c} = \underset {c > 0} {\arg \max } \left\{- \frac {c}{2} \right\} \Rightarrow c \rightarrow 0
$$
F.2 Environments and constructing features
**Cliff World:** We consider a modified version of the CliffWorld environment [53, Example 6.6]. The environment is deterministic and consists of 21 states and 4 actions. The objective is to reach the Goal state as quickly as possible. If the agent falls into a Cliff, it yields reward of  $-100$ , and is then returned to the Start state. Reaching the Goal state yields a reward of  $+1$  and the agent will stay in this terminal state. All other transitions are associated with a zero reward, and the discount factor is set to  $\gamma = 0.9$ .
Frozen Lake: We Consider the Frozen Lake v.1 environment from gym framework [6]. The environment is stochastic and consists of 16 states and 4 actions. The agent starts from the Start state and according to the next action (chosen by the policy) and the stochastic dynamics moves to the next state and yields a reward. The objective is to reach the Goal state as quickly as possible without entering the Hole States. All the Hole states and the Goal are terminal states. Reaching the goal state yields  $+1$  reward and all other rewards are zero, and the discount factor is set to  $\gamma = 0.9$ .
Sampling: We employ the Monte-Carlo method to sample from both environments and we use the expected return to estimate the action-value function  $Q$ . Specifically, we iteratively start from a randomly chosen state-action pair  $(s,a)$ , run a roll-out with a specified length starting from that pair, and collect the expected return to estimate  $Q(s,a)$ .
Constructing features: Also, in order to use function approximation on the above environments, we use tile-coded features [53]. Specifically, tilde-coded featurization needs three parameters to be set:  $(i)$  hash table size (equivalent to the feature dimension)  $d$ ,  $(ii)$  number of tiles  $N$  and  $(iii)$  size of tiles  $s$ . For Cliff world environment, we consider following pairs to construct features:  $\{(d = 40, N = 5, s = 1), (d = 50, N = 6, s = 1), (d = 60, N = 4, s = 3), (d = 80, N = 5, s = 3), (d = 100, N = 6, s = 3)\}$ . This means whenever we use  $d = 40$ , the number of tiles is  $N = 5$  and the tiling size is  $s = 1$ . The reported number of tiles and tiling size parameters are tuned and have achieved the best performance for all algorithms. Similarly for Frozen Lake environment, we use the following pairs to construct features:  $\{(d = 40, N = 3, s = 3), (d = 50, N = 4, s = 13), (d = 60, N = 5, s = 3), (d = 100, N = 8, s = 3)\}$ .
F.3 Critic optimization
We explain implementation of MSE, Adv-MSE and decision-aware critic loss functions. We use tile-coded features  $\mathbf{X}(s,a)$  and linear function approximation to estimate action-value function  $Q$ , implying that  $\hat{Q}(s,a) = \omega^T\mathbf{X}(s,a)$  where  $\omega, \mathbf{X}(s,a) \in \mathbb{R}^d$ .
Baselines: For policy  $\pi$ , the MSE objective is to return the  $\omega$  that minimizes the squared norm error of the action-value function  $Q^{\pi}$  across all state-action weights weighted by the state-action occupancy measure  $\mu^{\pi}(s,a)$ .
$$
\omega^{\mathrm{MSE}} = \operatorname *{arg  min}_{\omega \in \mathbb{R}^{d}}\mathbb{E}_{(s,a)\sim \mu^{\pi}(s,a)}[Q^{\pi}(s,a) - \omega^{T}\mathbf{X}(s,a)]^{2}
$$
Taking the derivative with respect to  $\omega$  and setting it to zero:
$$
\begin{array}{l} \mathbb {E} _ {(s, a) \sim \mu^ {\pi} (s, a)} \left[ \left(Q ^ {\pi} (s, a) - \omega^ {T} \mathbf {X} (s, a)\right) \mathbf {X} (s, a) ^ {T} \right] = 0 \\ \Longrightarrow \underbrace {\sum_ {s , a} \mu^ {\pi} (s , a) Q ^ {\pi} (s , a) \mathbf {X} (s , a) ^ {T}} _ {:= y} = \underbrace {\left[ \sum_ {s , a} \mu^ {\pi} (s , a) \mathbf {X} (s , a) \mathbf {X} (s , a) ^ {T} \right] \omega} _ {:= K} \\ \end{array}
$$
Given features  $\mathbf{X}$ , the true action-value function  $Q^{\pi}$  and state-action occupancy measure  $\mu^{\pi}$ , we can compute  $K$ ,  $y$  and solve  $\omega^{\mathrm{MSE}} = K^{-1}y$ .
Similarly for policy  $\pi$ , the advantage-MSE objective is to return  $\omega$  that minimizes the squared error of the advantage function  $A^{\pi}$  across all state-action weights weighted by the state-action occupancy measure  $\mu^{\pi}$ .
$$
\omega^ {\mathrm {A d v - M S E}} = \underset {\omega \in \mathbb {R} ^ {d}} {\arg \min } \mathbb {E} _ {(s, a) \sim \mu^ {\pi} (s, a)} \left[ A ^ {\pi} (s, a) - \omega^ {T} \left(\mathbf {X} (s, a) - \sum_ {a ^ {\prime}} \mathbf {X} (s, a ^ {\prime})\right) \right] ^ {2}
$$
$$
\begin{array}{l} \mathbb {E} _ {(s, a) \sim \mu^ {\pi} (s, a)} \left[ \left[ A ^ {\pi} (s, a) - \omega^ {T} \left(\mathbf {X} (s, a) - \sum_ {a ^ {\prime}} \mathbf {X} (s, a ^ {\prime})\right) \right] \left[ \mathbf {X} (s, a) - \sum_ {a ^ {\prime}} \mathbf {X} (s, a ^ {\prime}) \right] ^ {T} = 0 \right. \\ \Longrightarrow \underbrace {\sum_ {s , a} \mu^ {\pi} (s , a) A ^ {\pi} (s , a) \left[ \mathbf {X} (s , a) - \sum_ {a ^ {\prime}} \mathbf {X} (s , a ^ {\prime}) \right] ^ {T}} _ {:= y} = \underbrace {\sum_ {s , a} \mu^ {\pi} (s , a) \left[ \mathbf {X} (s , a) - \sum_ {a ^ {\prime}} \mathbf {X} (s , a ^ {\prime}) \right] \left[ \mathbf {X} (s , a) - \sum_ {a ^ {\prime}} \mathbf {X} (s , a ^ {\prime}) \right] ^ {T}} _ {:= K} \omega \\ \end{array}
$$
Given features  $\mathbf{X}$ , the true advantage function  $A^{\pi}$  and state-action occupancy measure  $\mu^{\pi}$ , we can compute  $K$  and  $y$  and solve  $w^{\mathrm{Adv - MSE}} = K^{-1}y$ .
Decision-aware critic in direct representation: Recall that for policy  $\pi$ , the decision-aware critic loss in direct representation is the blue term in Prop. 4, which after linear parameterization on  $\hat{Q}^{\pi}$  would be as follows:
$$
\mathbb {E} _ {s \sim d ^ {\pi}} \left[ \mathbb {E} _ {a \sim p ^ {\pi} (\cdot | s)} \left[ Q ^ {\pi} (s, a) - \omega^ {T} \mathbf {X} (s, a) \right] + \frac {1}{c} \log \left(\mathbb {E} _ {a \sim p ^ {\pi} (\cdot | s)} \left[ \exp \left(- c \left[ Q ^ {\pi} (s, a) - \omega^ {T} \mathbf {X} (s, a) \right]\right) \right]\right) \right]
$$
The above term is a convex function of  $\omega$  for any  $c > 0$ . We minimize the term using gradient descent, where the gradient with respect to  $\omega$  is:
$$
- \mathbb {E} _ {s \sim d ^ {\pi}} \left[ \mathbb {E} _ {a \sim p ^ {\pi} (\cdot | s)} \mathbf {X} (s, a) - \frac {\mathbb {E} _ {a \sim p ^ {\pi} (\cdot | s)} \left[ \exp \left(- c \left[ Q ^ {\pi} (s , a) - \omega^ {T} \mathbf {X} (s , a) \right]\right) \mathbf {X} (s , a) \right]}{\mathbb {E} _ {a \sim p ^ {\pi} (\cdot | s)} \left[ \exp \left(- c \left[ Q ^ {\pi} (s , a) - \omega^ {T} \mathbf {X} (s , a) \right]\right) \right]} \right]
$$
The step-size of gradient ascent is determined using Armijo line-search [3] where the maximum step size is set to 1000 and it decays with the rate  $\beta = 0.9$ . The number of iteration for critic inner-loop,  $m_{c}$  in Algorithm 1, is set to 10000, and if the gradient norm becomes smaller than  $10^{-6}$  we terminate the inner loop.
Decision-aware critic in softmax representation: Recall that for policy  $\pi$ , the decision-aware critic loss in softmax representation is the blue term in Prop. 6, which after linear parameterization on  $\hat{Q}^{\pi}$  and substituting  $\hat{A}^{\pi}(s,a)$  with  $\omega^{T}(\mathbf{X}(s,a) - \sum_{a'}\mathbf{X}(s,a'))$  would be as follows:
$$
\frac {1}{c} \mathbb {E} _ {s \sim d ^ {\pi_ {t}}} \mathbb {E} _ {a \sim p ^ {\pi_ {t}} (\cdot | s)} \left[ \right.\left(1 - c \left[ A ^ {\pi_ {t}} (s, a) - \omega^ {T} \left(\mathbf {X} (s, a) - \sum_ {a ^ {\prime}} \mathbf {X} (s, a ^ {\prime})\right)\right]\right) \log \left( \right.1 - c \left[ \right. A ^ {\pi_ {t}} (s, a) - \omega^ {T} \left(\mathbf {X} (s, a) - \sum_ {a ^ {\prime}} \mathbf {X} (s, a ^ {\prime})\right)\left. \right)\left. \right]
$$
Similarly, the above term is convex with respect to  $\omega$  and we minimize it using gradient descent. The step-size is determined using Armijo line-search with the same parameters as mentioned in direct case. The number of iterations in inner loop is set to 10000 and we terminate the loop if the gradient norm becomes smaller than  $10^{-8}$ . The gradient with respect to  $\omega$ :
$$
E _ {s \sim d ^ {\pi_ {t}}} \mathbb {E} _ {a \sim p ^ {\pi_ {t}} (\cdot | s)} \left[ \right.\left( \right.1 + \log \left( \right.1 - c \left[ \right. A ^ {\pi_ {t}} (s, a) - \omega^ {T} (\mathbf {X} (s, a) - \sum_ {a ^ {\prime}} \mathbf {X} (s, a ^ {\prime}))\left. \right)\left. \right)\left(\mathbf {X} (s, a) - \sum_ {a ^ {\prime}} \mathbf {X} (s, a ^ {\prime})\right)\left. \right]
$$
F.4 Actor optimization
Direct representation: For all actor-critic algorithms, we maximize the green term in Prop. 4 known as MDPO [56].
$$
\mathbb {E} _ {s \sim d ^ {\pi_ {t}}} \left[ \mathbb {E} _ {a \sim p ^ {\pi_ {t}} (\cdot | s)} \left[ \frac {p ^ {\pi} (a | s)}{p ^ {\pi_ {t}} (a | s)} \left(\hat {Q} ^ {\pi_ {t}} (s, a) - \left(\frac {1}{\eta} + \frac {1}{c}\right) \log \left(\frac {p ^ {\pi} (a | s)}{p ^ {\pi_ {t}} (a | s)}\right)\right) \right] \right]
$$
In tabular parameterization of the actor,  $\theta_{s,a} = p^{\pi}(s,a)$ , the actor update is exactly natural policy gradient [25] and can be solved in closed-form. We refer the reader to Appendix F.2 of [57] for explicit derivation. At iteration  $t$ , given policy  $\pi_t$ , the estimated action-value function from the critic  $\hat{Q}^{\pi_t}$  and  $\eta$  as the functional step-size, the update at iteration  $t$  is:
$$
p ^ {\pi_ {t + 1}} (a | s) = \frac {p ^ {\pi_ {t}} (a | s) \exp (\eta \hat {Q} ^ {\pi_ {t}} (s , a))}{\sum_ {a ^ {\prime}} p ^ {\pi_ {t}} (a ^ {\prime} | s) \exp (\eta \hat {Q} ^ {\pi_ {t}} (s , a ^ {\prime}))} \Rightarrow \theta_ {s, a} = \frac {\theta_ {s , a} \exp (\eta \hat {Q} ^ {\pi_ {t}} (s , a))}{\sum_ {a ^ {\prime}} \theta_ {s , a ^ {\prime}} \exp (\eta \hat {Q} ^ {\pi_ {t}} (s , a ^ {\prime}))}
$$
When we linearly parameterize the policy, implying that for policy  $\pi$ ,  $p^{\pi}(a|s) = \frac{\exp(\theta^{T}\mathbf{X}(s,a))}{\sum_{a'}\exp(\theta^{T}\mathbf{X}(s,a'))}$  where  $\theta$ ,  $\mathbf{X}(s,a)\in \mathbb{R}^n$  and  $n$  is the actor expressivity, we use the off-policy update loop (Lines 10-13 in Algorithm 1) and we iteratively update the parameters using gradient ascent. The MDPO objective with linear parameterization will be:
$$
\mathbb {E} _ {s \sim d ^ {\pi_ {t}}} \left[ \mathbb {E} _ {a \sim p ^ {\pi_ {t}} (\cdot | s)} \left[ \frac {\exp \left(\theta^ {T} \mathbf {X} (s , a)\right)}{p ^ {\pi_ {t}} (a | s) \sum_ {a ^ {\prime}} \exp \left(\theta^ {T} \mathbf {X} (s , a ^ {\prime})\right)} \left(\hat {Q} ^ {\pi_ {t}} (s, a) - \left(\frac {1}{\eta} + \frac {1}{c}\right) \log \left(\frac {\exp \left(\theta^ {T} \mathbf {X} (s , a)\right)}{p ^ {\pi_ {t}} (a | s) \sum_ {a ^ {\prime}} \exp \left(\theta^ {T} \mathbf {X} (s , a ^ {\prime})\right)}\right)\right) \right] \right]
$$
And the gradient of objective with respect to  $\theta$  is:
$$
\mathbb {E} _ {s \sim d ^ {\pi_ {t}}} \left[ \mathbb {E} _ {a \sim p ^ {\pi_ {t}} (\cdot | s)} \left[ \frac {p ^ {\pi} (a | s)}{p ^ {\pi_ {t}} (a | s)} \left(\mathbf {X} (s, a) - \frac {\sum_ {a ^ {\prime}} \exp (\theta^ {T} \mathbf {X} (s , a)) \mathbf {X} (s , a ^ {\prime})}{\sum_ {a ^ {\prime}} \exp (\theta^ {T} \mathbf {X} (s , a ^ {\prime}))}\right) \left(\hat {Q} ^ {\pi_ {t}} (s, a) - \left(\frac {1}{\eta} + \frac {1}{c}\right) \big (1 + \log (\frac {p ^ {\pi} (a | s)}{p ^ {\pi_ {t}} (a | s)}) \big)\right) \right] \right]
$$
Softmax representation: For all actor-critic algorithms, we maximize the green term in Prop. 6 known as sMDPO [57].
$$
\mathbb {E} _ {s \sim d ^ {\pi_ {t}}} \mathbb {E} _ {a \sim p ^ {\pi_ {t}} (\cdot | s)} \left[ \left(\hat {A} ^ {\pi_ {t}} (s, a) + \frac {1}{\eta} + \frac {1}{c}\right) \log \left(\frac {p ^ {\pi} (a | s)}{p ^ {\pi_ {t}} (a | s)}\right) \right]
$$
In tabular parameterization of the actor,  $\theta_{s,a} = p^{\pi}(s,a)$ , at iteration  $t$  given the policy  $\pi_t$ , the estimated advantage function from the critic  $\hat{A}^{\pi_t}(s,a) = \hat{Q}^{\pi_t}(s,a) - \sum_{a'} p^{\pi_t}(a|s) \hat{Q}^{\pi_t}(s,a')$ , and functional step-size  $\eta$ , the actor update can be solved in closed-form and is as follows:
$$
p ^ {\pi_ {t + 1}} (a | s) = \frac {p ^ {\pi_ {t}} (a | s) \max (1 + \eta A ^ {\pi_ {t}} (s , a) , 0)}{\sum_ {a ^ {\prime}} p ^ {\pi_ {t}} (a ^ {\prime} | s) \max (1 + \eta A ^ {\pi_ {t}} (s , a ^ {\prime}) , 0)} \Rightarrow \theta_ {s, a} = \frac {\theta_ {s , a} \max (1 + \eta A ^ {\pi_ {t}} (s , a) , 0)}{\sum_ {a ^ {\prime}} \theta_ {s , a ^ {\prime}} \max (1 + \eta A ^ {\pi_ {t}} (s , a ^ {\prime}) , 0)}
$$
We refer the reader to Appendix F.1 of [57] for explicit derivation. When we linearly parameterize the policy, implying that for policy  $\pi$ ,  $p^{\pi}(a|s) = \frac{\exp(\theta^{T}\mathbf{X}(s,a))}{\sum_{a'}\exp(\theta^{T}\mathbf{X}(s,a'))}$ , we need to maximize the following with respect to  $\theta$ :
$$
\mathbb {E} _ {s \sim d ^ {\pi_ {t}}} \mathbb {E} _ {a \sim p ^ {\pi_ {t}} (\cdot | s)} \left[ \left(\hat {A} ^ {\pi_ {t}} (s, a) + \frac {1}{\eta} + \frac {1}{c}\right) \log \left(\frac {\exp \left(\theta^ {T} \mathbf {X} (s , a)\right)}{p ^ {\pi_ {t}} (a | s) \sum_ {a ^ {\prime}} \exp \left(\theta^ {T} \mathbf {X} (s , a ^ {\prime})\right)}\right) \right]
$$
Similar to direct representation, we use the off-policy update loop and we iteratively update the parameters using gradient ascent. The gradient with respect to  $\theta$  is:
$$
\mathbb {E} _ {s \sim d ^ {\pi_ {t}}} \mathbb {E} _ {a \sim p ^ {\pi_ {t}} (\cdot | s)} \left[ \left(\hat {A} ^ {\pi_ {t}} (s, a) + \frac {1}{\eta} + \frac {1}{c}\right) \left(\mathbf {X} (s, a) - \frac {\sum_ {a ^ {\prime}} \exp (\theta^ {T} \mathbf {X} (s , a)) \mathbf {X} (s , a ^ {\prime})}{\sum_ {a ^ {\prime}} \exp (\theta^ {T} \mathbf {X} (s , a))}\right) \right]
$$
F.5 Parameter Tuning
Parameter	Value/Range
Sampling	# of samples	{1000,5000}
length of episode	{20,50}
Actor	Gradient termination criterion	{10-3,10-4}
ma	{1000,10000}
Armijo max step-size	1000
Armijo step-size decay β	0.9
Policy initialization (linear)	N(0,0.1)
Policy initialization (tabular)	Random
Linear Critic	Gradient termination criterion (direct)	{10-6,10-8}
Gradient termination criterion (softmax)	{10-8,10-10}
mc	{1000,10000}
Armijo max step-size	1000
Armijo step-size decay β	0.9
Others	η in direct	{0.001,0.005,0.01,0.1,1}
c in direct	{0.001,0.01,0.1,1}
η in softmax	{0.001,0.005,0.01,0.1,1}
c in softmax	{0.001,0.01,0.1}
d	{40,50,60,80,100}
Table 1: Parameters for the Cliff World environment
Parameter	Value/Range
Sampling	# of samples	{1000, 10000}
length of episode	{20, 50}
Actor	Gradient termination criterion	{10-4, 10-5}
ma	{100, 1000}
Armijo max step-size	1000
Armijo step-size decay β	0.9
Policy initialization (linear)	N(0, 0.1)
Policy initialization (tabular)	Random
Linear Critic	Gradient termination criterion (direct)	{10-6, 10-8}
Gradient termination criterion (softmax)	{10-6, 10-8}
mc	{10000, 1000000}
Armijo max step-size	1000
Armijo step-size decay β	0.9
Others	η in direct	{0.01, 0.1, 1, 10}
c in direct	{0.01, 0.1, 1}
η in softmax	{0.01, 0.1, 1, 10}
c in softmax	{0.01, 0.1}
d	{40, 50, 60, 100}
Table 2: Parameters for the Frozen Lake environment
G Additional Experiments
Figure 2: Cliff World - Linear actor and Linear critic with exact  $Q$  computation Assessing the impact of  $c$  (trade-off parameter in decision-aware framework) on the performance. We perform the experiment on the same setting as Fig. 1, linear actor and linear (with four different dimensions) critic with known MDP on Cliff World environment. We consider two values of functional step-size  $\eta \in \{0.01, 0.1\}$  and three values of  $c \in \{0.01, 0.1, 1\}$  for direct and  $c \in \{0.001, 0.01, 0.1\}$  for softmax representations, and compare the performance of 6 combinations. Overall, among different critic capacities and step-sizes, the value of  $c = 0.01$  demonstrates superior performance in both policy representations.
(a) Linear policy parameterization
(b) Tabular policy parameterization
Figure 3: Cliff World - Linear/Tabular actor and Linear critic with exact  $Q$  computation: Comparison of decision-aware, Adv-MSE, and MSE loss functions using a linear actor Fig. 3a and Fig. 3b coupled with a linear (with four different dimensions) critic in the Cliff World environment for direct and softmax policy representations with known MDP. For  $d = 100$  (corresponding to an expressive critic) in both actor parameterizations and  $d = 80$  in linear parameterization, all algorithms have almost the same performance. In other scenarios, minimizing MSE loss function with any functional step-size leads to a sub-optimal policy. In contrast, minimizing Adv-MSE and decision-aware loss functions always result in reaching the optimal policy even in the less expressive critic  $d = 40$ . Additionally, decision-aware convergence is faster than Adv-MSE particularly when the critic has limited capacity (e.g. In  $d = 40$  for direct and softmax representations and for both actor parameterizations, decision-aware reaches the optimal policy faster.)