SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sharan Vaswani
vaswani.sharan@gmail.com
Reza Babanezhad
Samsung - SAIT AI Lab, Montreal
babanezhad@gmail.com
Amirreza Kazemi
aka208@sfu.ca
Nicolas Le Roux
Microsoft Research, Mila
nicolas@le-roux.name
Actor-critic (AC) methods are widely used in reinforcement learning (RL), and benefit from the flexibility of using any policy gradient method as the actor and value-based method as the critic. The critic is usually trained by minimizing the TD error, an objective that is potentially decorrelated with the true goal of achieving a high reward with the actor. We address this mismatch by designing a joint objective for training the actor and critic in a decision-aware fashion. We use the proposed objective to design a generic, AC algorithm that can easily handle any function approximation. We explicitly characterize the conditions under which the resulting algorithm guarantees monotonic policy improvement, regardless of the choice of the policy and critic parameterization. Instantiating the generic algorithm results in an actor that involves maximizing a sequence of surrogate functions (similar to TRPO, PPO), and a critic that involves minimizing a closely connected objective. Using simple bandit examples, we provably establish the benefit of the proposed critic objective over the standard squared error. Finally, we empirically demonstrate the benefit of our decision-aware actor-critic framework on simple RL problems.
Reinforcement learning (RL) is a framework for solving problems involving sequential decision-making under uncertainty, and has found applications in games [38, 50], robot manipulation tasks [55, 64] and clinical trials [45]. RL algorithms aim to learn a policy that maximizes the long-term return by interacting with the environment. Policy gradient (PG) methods [59, 54, 29, 25, 47] are an important class of algorithms that can easily handle function approximation and structured state-action spaces, making them widely used in practice. PG methods assume a differentiable parameterization of the policy and directly optimize the return with respect to the policy parameters. Typically, a policy's return is estimated by using Monte-Carlo samples obtained via environment interactions [59]. Since the environment is stochastic, this approach results in high variance in the estimated return, leading to higher sample-complexity (number of environment interactions required to learn a good policy). Actor-critic (AC) methods [29, 43, 5] alleviate this issue by using value-based approaches [52, 58] in conjunction with PG methods, and have been empirically successful [20, 23]. In AC algorithms, a value-based method ("critic") is used to approximate a policy's estimated value, and a PG method ("actor") uses this estimate to improve the policy towards obtaining higher returns.
Though AC methods have the flexibility of using any method to independently train the actor and critic, it is unclear how to train the two components jointly in order to learn good policies. For example, the critic is typically trained via temporal difference (TD) learning and its objective is to minimize the value estimation error across all states and actions. For large real-world Markov decision processes (MDPs), it is intractable to estimate the values across all states and actions, and
algorithms resort to function approximation schemes. In this setting, the critic should focus its limited model capacity to correctly estimate the state-action values that have the largest impact on improving the actor's policy. This idea of explicitly training each component of the RL system to help the agent take actions that result in higher returns is referred to as decision-aware RL. Decision-aware RL [17, 16, 1, 10, 13, 14, 32] has mainly focused on model-based approaches that aim to learn a model of the environment, for example, the rewards and transition dynamics in an MDP. In this setting, decision-aware RL aims to model relevant parts of the world that are important for inferring a good policy. This is achieved by (i) designing objectives that are aware of the current policy [1, 14] or its value [17, 16], (ii) differentiating through the transition dynamics to learn models that result in good action-value functions [13] or (iii) simultaneously learning value functions and models that are consistent [51, 40, 35]. In the model-free setting, decision-aware RL aims to train the actor and critic cooperatively in order to optimize the same objective that results in near-optimal policies. In particular, Dai et al. [10] use the linear programming formulation of MDPs and define a joint saddle-point objective (minimization w.r.t. the critic and maximization w.r.t. the actor). The use of function approximation makes the resulting optimization problem non-convex non-concave leading to training instabilities and necessitating the use of heuristics. Recently, Dong et al. [11] used stochastic gradient descent-ascent to optimize this saddle-point objective and, under certain assumptions on the problem, proved that the resulting policy converges to a stationary point of the value function. Similar to Dong et al. [11], we study a decision-aware AC method with function approximation and equipped with theoretical guarantees on its performance. In particular, we make the following contributions.
Joint objective for training the actor and critic: Following Vaswani et al. [57], we distinguish between a policy's functional representation (sufficient statistics that define a policy) and its parameterization (the specific model used to realize these sufficient statistics in practice). For example, a policy can be represented by its state-action occupancy measure, and we can use a neural network parameterization to model this measure in practice (refer to Sec. 2 for more examples). In Sec. 3.2, we exploit a smoothness property of the return and design a lower-bound (Prop. 1) on the return of an arbitrary policy. Importantly, the lower bound depends on both the actor and critic, and immediately implies a joint objective for training the two components (minimization w.r.t the critic and maximization w.r.t the actor). Unlike Dai et al. [10], Dong et al. [11], the proposed objective works for any policy representation – the policy could be represented as conditional distributions over actions for each state or a deterministic mapping from states to actions [21]. Another advantage of working in the functional space is that our lower bound does not depend on the parameterization of either the actor or the critic. Moreover, unlike Dai et al. [10], Dong et al. [11], our framework does not need to model the distribution over states, and hence results in a more efficient algorithm. We note that our framework can be used for other applications where gradient computation is expensive or has large variance [39], and hence requires a model of the gradient (e.g., variational inference).
Generic actor-critic algorithm: In Sec. 3.2, we use our joint objective to design a generic decision-aware AC algorithm. The resulting algorithm (Algorithm 1) can be instantiated with any functional representation of the policy, and can handle any policy or critic parameterization. Similar to Vaswani et al. [57], the actor update involves optimizing a surrogate function that depends on the current policy, and consequently supports off-policy updates, i.e. similar to common PG methods such as TRPO [46], PPO [48], the algorithm can update the policy without requiring additional interactions with the environment. This property coupled with the use of a critic makes the resulting algorithm sample-efficient in practice. In contrast with TRPO/PPO, both the off-policy actor updates and critic updates in Algorithm 1 are designed to maximize the same lower bound on the policy return.
Theoretical guarantees: In Sec. 4.1, we analyze the necessary and sufficient conditions in order to guarantee monotonic policy improvement, and hence convergence to a stationary point. We emphasize that these improvement guarantees hold regardless of the policy parameterization and the quality of the critic (up to a certain threshold that we explicitly characterize). This is in contrast to existing theoretical results that focus on the tabular or linear function approximation settings or rely on highly expressive critics to minimize the critic error and achieve good performance for the actor. By exploiting the connection to inexact mirror descent (MD), we prove that Algorithm 1 is guaranteed to converge to the neighbourhood of a stationary point where the neighbourhood term depends on the decision-aware critic loss (Sec. 4.2). Along the way, we improve the theoretical guarantees for MD on general smooth, non-convex functions [15, 12]. As an additional contribution, we demonstrate a way to use the framework of Vaswani et al. [57] to "lift" the existing convergence rates [61, 37, 24] for the tabular setting to use off-policy updates and function approximation (Appendix D.2 and D.3). This gives rise to a simple, black-box proof technique that might be of independent interest.
Instantiating the general AC framework: We instantiate the framework for two policy representations – in Sec. 5.1, we represent the policy by the set of conditional distributions over actions ("direct" representation), whereas in Sec. 5.2, we represent the policy by using the logits corresponding to a softmax representation of these conditional distributions ("softmax" representation). In both cases, we instantiate the generic lower-bound (Propositions 4, 6), completely specifying the actor and critic objectives in Algorithm 1. Importantly, unlike the standard critic objective that depends on the squared difference of the value functions, the proposed decision-aware critic loss (i) depends on the policy representation – it involves the state-action value functions for the direct representation and depends on the advantage functions for the softmax representation, and (ii) penalizes the under-estimation and over-estimation of these quantities in an asymmetric manner. For both representations, we consider simple bandit examples (Propositions 5, 7) which show that minimizing the decision-aware critic loss results in convergence to the optimal policy, whereas minimizing variants of the squared loss do not. In App. B, we consider a third policy representation involving stochastic value gradients [21] for continuous control, and instantiate our decision-aware actor-critic framework in this case.
Experimental evaluation: Finally, in Sec. 6, we consider simple RL environments and benchmark Algorithm 1 for both the direct and softmax representations with a linear policy and critic parameterization. We compare the actor performance when using the squared critic loss vs the proposed critic loss, and demonstrate the empirical benefit of our decision-aware actor-critic framework.
We consider an infinite-horizon discounted Markov decision process (MDP) [44] defined by the tuple  $\langle S, \mathcal{A}, \mathcal{P}, r, \rho, \gamma \rangle$  where  $\mathcal{S}$  is the set of states,  $\mathcal{A}$  is the action set,  $\mathcal{P}: S \times \mathcal{A} \to \Delta_S$  is the transition probability function,  $\rho \in \Delta_S$  is the initial distribution of states,  $r: S \times \mathcal{A} \to [0,1]$  is the reward function and  $\gamma \in [0,1)$  is the discount factor. For state  $s \in S$ , a policy  $\pi$  induces a distribution  $p^{\pi}(\cdot | s)$  over actions. It also induces a measure  $d^{\pi}$  over states such that  $d^{\pi}(s) = (1 - \gamma) \sum_{\tau=0}^{\infty} \gamma^{\tau} \mathcal{P}(s_{\tau} = s \mid s_0 \sim \rho, a_{\tau} \sim p^{\pi}(\cdot | s_{\tau}))$ . Similarly, we define  $\mu^{\pi}$  as the measure over state-action pairs induced by policy  $\pi$ , implying that  $\mu^{\pi}(s, a) = d^{\pi}(s) p^{\pi}(a|s)$  and  $d^{\pi}(s) = \sum_{a} \mu^{\pi}(s, a)$ . The action-value function corresponding to policy  $\pi$  is denoted by  $Q^{\pi}: S \times \mathcal{A} \to \mathbb{R}$  such that  $Q^{\pi}(s, a) := \mathbb{E}[\sum_{\tau=0}^{\infty} \gamma^{\tau} r(s_{\tau}, a_{\tau})]$  where  $s_0 = s, a_0 = a$  and for  $\tau \geq 0$ ,  $s_{\tau+1} \sim \mathcal{P}(\cdot | s_{\tau}, a_{\tau})$  and  $a_{\tau+1} \sim p^{\pi}(\cdot | s_{\tau+1})$ . The value function of a stationary policy  $\pi$  for the start state equal to  $s$  is defined as  $J_s(\pi) := \mathbb{E}_{a \sim p^{\pi}(\cdot | s)}[Q^{\pi}(s, a)]$  and we define  $J(\pi) := \mathbb{E}_{s \sim \rho}[J_s(\pi)]$ . For a state-action pair  $(s, a)$ , the advantage function corresponding to policy  $\pi$  is given by  $A^{\pi}(s, a) := Q^{\pi}(s, a) - J_s(\pi)$ . Given a set of feasible policies  $\Pi$ , the objective is to compute the policy that maximizes  $J(\pi)$ .
Functional representation vs Policy Parameterization: Similar to the policy optimization framework of Vaswani et al. [57], we differentiate between a policy's functional representation and its parameterization. The functional representation of a policy  $\pi$  defines its sufficient statistics, for example, we may represent a policy via the set of distributions  $p^{\pi}(\cdot |s) \in \Delta_A$  for state  $s \in S$ . We will refer to this as the direct representation. The same policy can have multiple functional representations, for example, since  $p^{\pi}(\cdot |s)$  is a probability distribution, one can write  $p^{\pi}(a|s) = \exp(z^{\pi}(s,a)) / \sum_{a'} \exp(z^{\pi}(s,a'))$ , and represent  $\pi$  by the set of logits  $z^{\pi}(s,a)$  for each  $(s,a)$  pair. We will refer to this as the softmax representation. On the other hand, the policy parameterization is determined by a model (with parameters  $\theta$ ) that realizes these statistics. For example, we could use a neural-network to parameterize the logits corresponding to the policy's softmax representation, rewriting  $z^{\pi}(s,a) = z^{\pi}(s,a|\theta)$  where the model is implicit in the  $z^{\pi}(s,a|\theta)$  notation. As another example, the tabular parameterization corresponds to having a parameter for each state-action pair [61, 37]. The policy parameterization thus defines the set  $\Pi$  of realizable policies that can be expressed with the parametric model at hand. Note that the policy parameterization can be chosen independently of its functional representation. Next, we recap the framework in Vaswani et al. [57] and generalize it to the actor-critic setting.
We describe functional mirror ascent in Sec. 3.1, and use it to design a general decision-aware actor-critic framework and corresponding algorithm in Sec. 3.2.
3.1 Functional Mirror Ascent for Policy Gradient (FMAPG) framework
For a given functional representation, Vaswani et al. [57] update the policy by functional mirror ascent and project the updated policy onto the set  $\Pi$  determined by the policy parameterization.
Functional mirror ascent is an iterative algorithm whose update at iteration  $t \in \{0,1,\dots,T - 1\}$  is given as:  $\pi_{t + 1} = \arg \max_{\pi \in \Pi}\left[\langle \pi ,\nabla_{\pi}J(\pi_{t})\rangle -\frac{1}{\eta} D_{\Phi}(\pi ,\pi_{t})\right]$  where  $\pi_t$  is the policy (expressed as its functional representation) at iteration  $t$ ,  $\eta$  is the step-size in the functional space and  $D_{\Phi}$  is the Bregman divergence (induced by the mirror map  $\Phi$ ) between the representation of policies  $\pi$  and  $\pi_t$ . The FMAPG framework casts the projection step onto  $\Pi$  as an unconstrained optimization w.r.t the parameters  $\theta \in \mathbb{R}^n$  of a surrogate function:  $\theta_{t + 1} = \arg \max_{\ell_t(\theta)}\coloneqq \langle \pi (\theta),\nabla_\pi J(\pi (\theta_t))\rangle -\frac{1}{\eta} D_{\Phi}(\pi (\theta),\pi (\theta_t))$ . Here,  $\pi (\theta)$  refers to the parametric form of the policy where the choice of the parametric model is implicit in the  $\pi (\theta)$  notation. The policy at iteration  $t$  is thus expressed as  $\pi (\theta_t)$  whereas the updated policy is given by  $\pi_{t + 1} = \pi (\theta_{t + 1})$ . The surrogate function is non-concave in general and can be approximately maximized using a gradient-based method, resulting in a nested loop algorithm. Importantly, the inner-loop (optimization of  $\ell_t(\theta)$ ) updates the policy parameters (and hence the policy), but does not involve recomputing  $\nabla_{\pi}J(\pi)$ . Consequently, these policy updates do not require interacting with the environment and are thus off-policy. This is a desirable trait for designing sample-efficient PG algorithms and is shared by methods such as TRPO [46] and PPO [48].
With the appropriate choice of  $\Phi$  and  $\eta$ , the FMAPG framework guarantees monotonic policy improvement for any number of inner-loops and policy parameterization. A shortcoming of this framework is that it requires access to the exact gradient  $\nabla_{\pi}J(\pi)$ . When using the direct or softmax representations, computing  $\nabla_{\pi}J(\pi)$  involves computing either the action-value  $Q^{\pi}$  or the advantage  $A^{\pi}$  function respectively. In complex real-world environments where the rewards and/or the transition dynamics are unknown, these quantities can only be estimated. For example,  $Q^{\pi}$  can be estimated using Monte-Carlo sampling by rolling out trajectories using policy  $\pi$  resulting in large variance, and consequently higher sample complexity. Moreover, for large MDPs, function approximation is typically used to estimate the  $Q$  function, and the resulting aliasing makes it impossible to compute it exactly in practice. This makes the FMAPG framework impractical in real-world scenarios. Next, we generalize FMAPG to handle inexact gradients and subsequently design an actor-critic framework.
3.2 Generalizing FMAPG to Actor-Critic
To generalize the FMAPG framework, we first prove the following proposition in App. C.
Proposition 1. For any policy representations  $\pi$  and  $\pi'$ , any strictly convex mirror map  $\Phi$ , and any gradient estimator  $\hat{g}$ , for  $c > 0$  and  $\eta$  such that  $J + \frac{1}{\eta} \Phi$  is convex in  $\pi$ ,
$$
J (\pi) \geq J \left(\pi^ {\prime}\right) + \langle \hat {g} \left(\pi^ {\prime}\right), \pi - \pi^ {\prime} \rangle - \left(\frac {1}{\eta} + \frac {1}{c}\right) D _ {\Phi} \left(\pi , \pi^ {\prime}\right) - \frac {1}{c} D _ {\Phi *} \left(\nabla \Phi \left(\pi^ {\prime}\right) - c [ \nabla J \left(\pi^ {\prime}\right) - \hat {g} \left(\pi^ {\prime}\right) ], \nabla \Phi \left(\pi^ {\prime}\right)\right)
$$
where  $\Phi^{*}$  is the Fenchel conjugate of  $\Phi$  and  $D_{\Phi^*}$  is the Bregman divergence induced by  $\Phi^{*}$ .
The above proposition is a statement about the relative smoothness [34] of  $J$  (w.r.t  $D_{\Phi}$ ) in the functional space. Here, the brown term is the linearization of  $J$  around  $\pi'$ , but involves  $\hat{g}(\pi')$  which can be any estimate of the gradient at  $\pi'$ . The red term quantifies the distance between the representations of policies  $\pi$  and  $\pi'$  in terms of  $D_{\Phi}(\pi, \pi')$ , whereas the blue term characterizes the penalty for an inaccurate estimate of  $\nabla_{\pi}J(\pi')$  and depends on  $\Phi$ . We emphasize that Prop. 1 can be used for any continuous optimization problem that requires a model of the gradient, e.g., in variational inference which uses an approximate posterior in lieu of the true one.
For policy optimization with FMAPG,  $\nabla_{\pi}J(\pi)$  involves the action-value or advantage function for the direct or softmax functional representations respectively (see Sec. 5 for details), and the gradient estimation error is equal to the error in these functions. Since these quantities are estimated by the critic, we refer to the blue term as the critic error. In order to use Prop. 1, at iteration  $t$  of FMAPG, we set  $\pi^{\prime} = \pi_{t}$  and include the policy parameterization, resulting in inequality (I):
$$
J (\pi) - J (\pi_ {t}) \geq \langle \hat {g} _ {t}, \pi (\theta) - \pi_ {t} \rangle - \left(\frac {1}{\eta} + \frac {1}{c}\right) D _ {\Phi} (\pi (\theta), \pi_ {t}) - \frac {1}{c} D _ {\Phi^ {*}} \left(\nabla \Phi (\pi_ {t}) - c [ \nabla J (\pi_ {t}) - \hat {g} _ {t} ], \nabla \Phi (\pi_ {t})\right),
$$
where  $\hat{g}_t \coloneqq \hat{g}(\pi_t)$ . We see that in order to obtain a policy  $\pi$  that maximizes the policy improvement  $J(\pi) - J(\pi_t)$  and hence the LHS, we should maximize the RHS i.e. (i) learn  $\hat{g}_t$  to minimize the blue term (equal to the critic objective) and (ii) compute  $\pi \in \Pi$  that maximizes the green term (equal to the functional mirror ascent update at iteration  $t$ ). Using a second-order Taylor series expansion of  $D_{\Phi^*}$  (Prop. 14), we see that as  $c$  decreases, the critic error decreases, whereas the  $\left(\frac{1}{\eta} + \frac{1}{c}\right) D_{\Phi}(\pi, \pi_t)$  term increases. Consequently, we interpret the scalar  $c$  as a trade-off parameter that relates the critic error to the permissible movement in the functional mirror ascent update.
Hence, both the actor and critic objectives are coupled through Prop. 1 and both components of the RL system should be jointly trained in order to maximize policy improvement. We refer to the resulting framework as decision-aware actor-critic and present its pseudo-code in Algorithm 1.
Algorithm 1: Generic actor-critic algorithm
1 Input:  $\pi$  (choice of functional representation),  $\theta_0$  (initial policy parameters),  $\omega_{(-1)}$  (initial critic parameters),  $T$  (AC iterations),  $m_a$  (actor inner-loops),  $m_c$  (critic inner-loops),  $\eta$  (functional step-size for actor),  $c$  (trade-off parameter),  $\alpha_a$  (parametric step-size for actor),  $\alpha_c$  (parametric step-size for critic)  
2 Initialization:  $\pi_0 = \pi (\theta_0)$   
3 for  $t\gets 0$  to  $T - 1$  do  
4 Estimate  $\widehat{\nabla}_{\pi}J(\pi_t)$  and form  $\mathcal{L}_t(\omega)\coloneqq \frac{1}{c} D_{\Phi^*}\Bigg(\nabla \Phi (\pi_t) - c[\widehat{\nabla}_{\pi}J(\pi_t) - \hat{g}_t(\omega)],\nabla \Phi (\pi_t)\Bigg)$   
5 Initialize inner-loop:  $v_{0} = \omega_{t - 1}$   
6 for  $k\gets 0$  to  $m_c - 1$  do  
7  $\begin{array}{r}\mid v_{k + 1} = v_k - \alpha_c\nabla_v\mathcal{L}_t(v_k) / \star \end{array}$  Critic Updates  $\star /$   
8  $\omega_{t} = v_{m_{c}}\quad ;\quad \hat{g}_{t} = \hat{g}_{t}(\omega_{t})$   
9 Form  $\ell_t(\theta)\coloneqq \langle \hat{g}_t,\pi (\theta) - \pi_t\rangle -\left(\frac{1}{\eta} +\frac{1}{c}\right)D_\Phi (\pi (\theta),\pi_t)$   
10 Initialize inner-loop:  $\nu_{0} = \theta_{t}$   
11 for  $k\gets 0$  to  $m_a - 1$  do  
12  $\begin{array}{r}\mid \nu_{k + 1} = \nu_k + \alpha_a\nabla_\nu \ell_t(\nu_k) / \star \end{array}$  Off-policy actor updates  $\star /$   
13  $\theta_{t + 1} = \nu_{m_a}\quad ;\quad \pi_{t + 1} = \pi (\theta_{t + 1})$   
14 Return  $\pi_T = \pi (\theta_T)$
Unlike Wu et al. [60], Konda and Tsitsiklis [29], Algorithm 1 does not update the actor and critic in a two time-scale setting (one environment interaction and update to the critic followed by an actor update), but rather performs multiple steps to update the critic, then uses the critic to perform multiple steps to update the actor [2, 61]. At iteration  $t$  of Algorithm 1,  $\hat{g}_t$  (the gradient estimate at  $\pi_t$ ) is parameterized by  $\omega$  and the parametric model for the critic is implicit in the  $\hat{g}_t(\omega)$  notation. The algorithm interacts with the environment, uses these interactions to form the estimate  $\widehat{\nabla_{\pi}} J(\pi_t)$  and construct the critic loss function  $\mathcal{L}_t(\omega)$ . For the direct or softmax representations,  $\widehat{\nabla_{\pi}} J(\pi_t)$  corresponds to the empirical estimates of the action-value or advantage functions respectively. In practice, these quantities can be estimated using Monte-Carlo rollouts or bootstrapping. Given these estimates, the critic is trained (using  $m_c$  inner-loops) to minimize  $\mathcal{L}_t(\omega)$  and obtain  $\hat{g}_t$  (Lines 5-8). Line 9 uses  $\hat{g}_t$  to construct the surrogate function  $\ell_t(\theta)$  for the actor and depends on the policy parameterization. The inner-loop (Lines 10 - 13) involves maximizing  $\ell_t(\theta)$  and corresponds to  $m_a$  off-policy updates. Next, we establish theoretical guarantees on the performance of Algorithm 1.
4 Theoretical Guarantees