Title: Decision-Aware Actor-Critic with Function Approximation and Theoretical Guarantees
Venue: NeurIPS 2023
Accepted: yes
Keywords: Decision-aware reinforcement learning

Actor-Critic algorithm

Off-policy updates

General function approximation

Theoretical guarantees
Abstract: 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.

Decision-Aware Actor-Critic with Function Approximation and Theoretical Guarantees

Abstract

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.

1 Introduction

Reinforcement learning (RL) is a framework for solving problems involving sequential decisionmaking 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

37th Conference on Neural Information Processing Systems (NeurIPS 2023).

Amirreza Kazemi

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 decisionaware 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.

2 Problem Formulation

We consider an infinite-horizon discounted Markov decision process (MDP) [44] defined by the tuple ⟨S , A , P , r, ρ, γ ⟩ where S is the set of states, A is the action set, P : S × A → ∆ S is the transition probability function, ρ ∈ ∆ S is the initial distribution of states, r : S × A → [0 , 1] is the reward function and γ ∈ [0 , 1) is the discount factor. For state s ∈ S , a policy π induces a distribution p π ( ·| s ) over actions. It also induces a measure d π over states such that d π ( s ) = (1 -γ ) ∑ ∞ τ =0 γ τ P ( s τ = s | s 0 ∼ ρ, a τ ∼ p π ( ·| s τ )) . Similarly, we define µ π as the measure over stateaction pairs induced by policy π , implying that µ π ( s, a ) = d π ( s ) p π ( a | s ) and d π ( s ) = ∑ a µ π ( s, a ) . The action-value function corresponding to policy π is denoted by Q π : S × A → R such that Q π ( s, a ) := E [ ∑ ∞ τ =0 γ τ r ( s τ , a τ )] where s 0 = s, a 0 = a and for τ ≥ 0 , s τ +1 ∼ P ( ·| s τ , a τ ) and a τ +1 ∼ p π ( ·| s τ +1 ) . The value function of a stationary policy π for the start state equal to s is defined as J s ( π ) := E a ∼ p π ( ·| s ) [ Q π ( s, a )] and we define J ( π ) := E s ∼ ρ [ J s ( π )] . For a state-action pair ( s, a ) , the advantage function corresponding to policy π is given by A π ( s, a ) := Q π ( s, a ) -J s ( π ) . Given a set of feasible policies Π , the objective is to compute the policy that maximizes J ( π ) .

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 π defines its sufficient statistics, for example, we may represent a policy via the set of distributions p π ( ·| s ) ∈ ∆ A for state s ∈ S . We will refer to this as the direct representation . The same policy can have multiple functional representations, for example, since p π ( ·| s ) is a probability distribution, one can write p π ( a | s ) = exp( z π ( s,a )) / ∑ a ′ exp( z π ( s,a ′ )) , and represent π by the set of logits z π ( 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 θ ) 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 π ( s, a ) = z π ( s, a | θ ) where the model is implicit in the z π ( s, a | θ ) 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 Π 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.

3 Methodology

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 Π determined by the policy parameterization.

Functional mirror ascent is an iterative algorithm whose update at iteration t ∈ { 0 , 1 , . . . , T -1 } is given as: π t +1 = arg max π ∈ Π [ ⟨ π, ∇ π J ( π t ) ⟩ -1 η D Φ ( π, π t ) ] where π t is the policy (expressed as its functional representation) at iteration t , η is the step-size in the functional space and D Φ is the Bregman divergence (induced by the mirror map Φ ) between the representation of policies π and π t . The FMAPG framework casts the projection step onto Π as an unconstrained optimization w.r.t the parameters θ ∈ R n of a surrogate function : θ t +1 = arg max ℓ t ( θ ) := ⟨ π ( θ ) , ∇ π J ( π ( θ t )) ⟩ -1 η D Φ ( π ( θ ) , π ( θ t )) . Here, π ( θ ) refers to the parametric form of the policy where the choice of the parametric model is implicit in the π ( θ ) notation. The policy at iteration t is thus expressed as π ( θ t ) , whereas the updated policy is given by π t +1 = π ( θ 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 ℓ t ( θ ) ) updates the policy parameters (and hence the policy), but does not involve recomputing ∇ π J ( π ) . 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 Φ and η , 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 ∇ π J ( π ) . When using the direct or softmax representations, computing ∇ π J ( π ) involves computing either the action-value Q π or the advantage A π 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 π can be estimated using Monte-Carlo sampling by rolling out trajectories using policy π 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 π and π ′ , any strictly convex mirror map Φ , and any gradient estimator ˆ g , for c > 0 and η such that J + 1 η Φ is convex in π ,

where Φ ∗ is the Fenchel conjugate of Φ and D Φ ∗ is the Bregman divergence induced by Φ ∗ .

The above proposition is a statement about the relative smoothness [34] of J (w.r.t D Φ ) in the functional space. Here, the brown term is the linearization of J around π ′ , but involves ˆ g ( π ′ ) which can be any estimate of the gradient at π ′ . The red term quantifies the distance between the representations of policies π and π ′ in terms of D Φ ( π, π ′ ) , whereas the blue term characterizes the penalty for an inaccurate estimate of ∇ π J ( π ′ ) and depends on Φ . 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, ∇ π J ( π ) 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 π ′ = π t and include the policy parameterization, resulting in inequality (I) :

where ˆ g t := ˆ g ( π t ) . We see that in order to obtain a policy π that maximizes the policy improvement J ( π ) -J ( π t ) and hence the LHS, we should maximize the RHS i.e. (i) learn ˆ g t to minimize the blue term (equal to the critic objective) and (ii) compute π ∈ Π that maximizes the green term (equal to the functional mirror ascent update at iteration t ). Using a second-order Taylor series expansion of D Φ ∗

(Prop. 14), we see that as c decreases, the critic error decreases, whereas the ( 1 η + 1 c ) D Φ ( π, π 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

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, ˆ g t (the gradient estimate at π t ) is parameterized by ω and the parametric model for the critic is implicit in the ˆ g t ( ω ) notation. The algorithm interacts with the environment, uses these interactions to form the estimate ̂ ∇ π J ( π t ) and construct the critic loss function L t ( ω ) . For the direct or softmax representations, ̂ ∇ π J ( π 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 L t ( ω ) and obtain ˆ g t (Lines 5-8). Line 9 uses ˆ g t to construct the surrogate function ℓ t ( θ ) for the actor and depends on the policy parameterization. The inner-loop (Lines 10 - 13) involves maximizing ℓ t ( θ ) and corresponds to m a off-policy updates. Next, we establish theoretical guarantees on the performance of Algorithm 1.

4 Theoretical Guarantees

We first establish the necessary and sufficient conditions to guarantee monotonic policy improvement in the presence of critic error (Sec. 4.1). In Sec. 4.2, we prove that Algorithm 1 is guaranteed to converge to the neighbourhood (that depends on the critic error) of a stationary point.

4.1 Conditions for monotonic policy improvement

According to inequality (I) , to guarantee monotonic policy improvement at iteration t , one must find a ( θ, c ) pair to guarantee that the RHS of (I) is positive. In Prop. 2 (proved in App. D), we derive the conditions on the critic error to ensure that it possible to find such an ( θ, c ) pair.

Proposition 2. For any policy representation and any policy or critic parameterization, there exists a ( θ, c ) pair that makes the RHS of inequality (I) strictly positive, and hence guarantees monotonic policy improvement ( J ( π t +1 ) > J ( π t ) ), if and only if

where b t ∈ R n := ∑ s ∈S ∑ a ∈A [ˆ g t ] s,a ∇ θ [ π ( θ t )] s,a , ˜ H t ∈ R n × n := ∇ θ π ( θ t ) T ∇ 2 π Φ( π t ) ∇ θ π ( θ t ) and ˜ H † t denotes the pseudo-inverse of ˜ H t . For the special case of the tabular policy parameterization, the above condition becomes equal to,

For the Euclidean mirror map with the tabular policy parameterization, this condition becomes equal to ∥ ˆ g t ∥ 2 2 > ∥∇ J ( π t ) -ˆ g t ∥ 2 2 meaning that the relative error in estimating ∇ J ( π t ) needs to be less than 1 . For a general mirror map, the relative error is measured in a different norm induced by the mirror map. The above proposition also quantifies the scenario when the critic error is too large to guarantee policy improvement. In this case, the algorithm should either improve the critic by better optimization or by using a more expressive model, or resort to using sufficiently many (high-variance) Monte-Carlo samples as in REINFORCE [59]. Finally, we see that the impact of a smaller function class for the actor is a potentially lower value for ⟨ b t , ˜ H † t b t ⟩ , making it more difficult to satisfy the condition. The improvement guarantee in Prop. 2 holds regardless of the policy representation and parameterization of the policy or critic. This is in contrast to existing theoretical results [41, 28, 18] that focus on either the tabular or linear function approximation setting for the policy and/or critic, or rely on using expressive models to minimize the critic error and achieve good performance for the actor. Moreover, this result only depends on the magnitude of the critic loss (after updates), irrespective of the optimizer, step-size or other factors influencing the critic optimization. The actor and critic are coupled via the threshold (on the critic loss) required to guarantee policy improvement.

4.2 Convergence of Algorithm 1

Prop. 2 holds when the critic error is small. We now analyze the convergence of Algorithm 1 for an arbitrary critic error. Define ¯ θ t +1 := arg max θ ℓ t ( θ ) , ¯ π t +1 = π ( ¯ θ t +1 ) = arg max π ∈ Π { ⟨ ˆ g t , π -π t ⟩ -( 1 η + 1 c ) D Φ ( π, π t ) } . Note that ¯ π t +1 is the iterate obtained by using the inexact mirror ascent (MA) update (because it does not use the true gradient ∇ π J ( π t ) ) starting from π t , and that the inner-loop (Lines 10-13) of Algorithm 1 approximates this update. This connection allows us to prove the following guarantee (see App. D.1 for details) for Algorithm 1.

Proposition 3. For any policy representation and mirror map Φ such that (i) J + 1 η Φ is convex in π , any policy parameterization such that (ii) ℓ t ( θ ) is smooth w.r.t θ and satisfies the Polyak-Lojasiewicz (PL) condition, for c > 0 , after T iterations of Algorithm 1 we have that,

where δ t := ∇ J ( π t ) -ˆ g t , 1 ζ := 1 η + 1 c , R is a random variable chosen uniformly from { 0 , 1 , 2 , . . . T -1 } and e t ∈ O (exp ( -m a )) is the projection error (onto Π ) at iteration t .

Prop. 3 shows that Algorithm 1 converges to the neighbourhood of a stationary point of J for an arbitrary critic error. The LHS of the above expression is a measure of sub-optimality similar to the one used in the analysis of stochastic mirror descent [65]. For the Euclidean mirror map, the LHS becomes equal to ∥∇ π J ( π R ) ∥ 2 2 , the standard characterization of a stationary point. Term (i) on the RHS is the initial sub-optimality, whereas Term (ii) is equal to the critic error and can be further decomposed into variance and bias terms. The variance decreases as the number of samples used to train the critic (Line 4 in Algorithm 1) increases. The bias can be decomposed into an optimization error (that decreases as m c increases) and a function approximation error (that decreases as we use more expressive models for the critic). Finally, Term (iii) is the projection (onto Π ) error, is equal to zero for the tabular policy parameterization, and decreases as m a increases. Hence, the performance of Algorithm 1 improves as we increase both m a and m c .

In Sec. 5, we specify the step-size η such that Assumption (i) is satisfied for both the direct and softmax representations. Assumption (ii) is satisfied when using a linear and, in some cases, a neural network policy parameterization [33]. For the above proposition to hold, we require that step-sizes α c and α a in Algorithm 1 be set according to the smoothness of critic ( L t ( ω ) ) and actor ( ℓ t ( θ ) ) objectives respectively. These choice of step-sizes guarantee ascent for the actor objective, and descent for the critic objective (refer to the proof of Prop. 3 in App. D for details). In practice, we set both step-sizes using an Armijo line-search, and refer the reader to App. F for details. Since Algorithm 1 does not update the actor and critic in a two time-scale setting, unlike [29, 60], the relative scales of the step-sizes and the number of inner iterations ( m a , m c ) do not affect the algorithm's performance.

In contrast to Prop. 3, Dong et al. [11] prove that their proposed algorithm results in an O ( 1 / T ) convergence to the stationary point (not the neighbourhood). However, they make a strong unjustified

assumption that the minimization problem w.r.t the parameters modeling the policy and distribution over states is jointly PL. Compared to [2, 61, 36] that focus on proving convergence to the neighbourhood of the optimal value function, but bound the critic error in the ℓ 2 or ℓ ∞ norm, we focus on proving convergence to the (neighbourhood) of a stationary point, but define the critic loss in a decision-aware manner that depends on D Φ ∗ . Since Algorithm 1 is not a two time-scale algorithm, unlike Konda and Tsitsiklis [29], Wu et al. [60], the proof of Prop. 3 does not require analyzing coupled recursions between the actor and critic. Furthermore, the guarantees in Prop. 3 are independent of how the critic loss is minimized. We could use any policy evaluation method in order to estimate the value function. Hence, unlike the standard two time-scale analyses, we do not make assumptions about the mixing time of the underlying Markov chain.

Compared to the existing theoretical work on general (not decision-aware) AC methods [62, 60, 8, 28, 22, 30, 18, 41, 9] that prove convergence for the tabular or linear function approximation settings, (a) our theoretical results require fewer assumptions on the function approximation. For instance, Prop. 2 holds for any actor or critic parameterization (including complex neural networks), while the guarantees in Prop. 3 hold for any critic parameterization, but require that ℓ t , the surrogate function for the actor satisfy smoothness and gradient domination properties. (b) On the other hand, since our analysis does not explicitly model how the critic error is minimized, we can only converge to the neighbourhood of a stationary point. This is in contrast to the existing two time-scale analyses that jointly analyze the actor and critic, and show convergence to a stationary point [60]. (c) Finally, we note that the proposed algorithm supports off-policy updates i.e. the actor can re-use the value estimates from the critic to update the policy multiple times (corresponding to Lines 10-13 in Algorithm 1). This is in contrast to existing theoretically principled actor-critic methods that require interacting with the environment and gathering new data after each policy update. Hence, compared to the existing literature on AC methods, Algorithm 1 is more practical, has weaker theoretical guarantees but requires fewer assumptions on the function approximation.

5 Instantiating the generic actor-critic framework

We now instantiate Algorithm 1 for the direct (Sec. 5.1) and softmax (Sec. 5.2) representation.

5.1 Direct representation

Recall that for the direct functional representation, policy π is represented by the set of distributions p π ( ·| s ) over actions for each state s ∈ S . Using the policy gradient theorem [53], [ ∇ π J ( π )] s,a = d π ( s ) Q π ( s, a ) . Similar to [57, 61], we use a weighted (across states) negative entropy mirror map implying that D Φ ( p π , p π ′ ) = ∑ s ∈S d π t ( s ) D ϕ ( p π ( ·| s ) , p π ′ ( ·| s )) where ϕ ( p π ( ·| s )) = -∑ a p π ( a | s ) log( p π ( a | s )) and hence, D ϕ ( p π ( ·| s ) , p π ′ ( ·| s )) = KL ( p π ( ·| s ) || p π ′ ( ·| s )) . We now instantiate inequality (I) in Sec. 3.2 in the proposition below (see App. E for the derivation).

For incorporating policy (with parameters θ ) and critic (with parameters ω ) parameterization, we note that p π ( ·| s ) = p π ( ·| s, θ ) and ˆ Q π ( s, a ) = Q π ( s, a | ω ) where the model is implicit in the notation. Using the reasoning in Sec. 3.2 with Prop. 4 immediately gives us the actor and critic objectives ( ℓ t ( θ ) and L t ( ω ) respectively) at iteration t and completely instantiates Algorithm 1. Observe that the critic error is asymmetric and penalizes the under/over-estimation of the Q π function differently. This is different from the standard squared critic loss: E s ∼ d πt E a ∼ p πt ( ·| s ) [ Q π t ( s, a ) -Q π t ( s, a | ω )] 2 that does not take into account the sign of the misestimation.

To demonstrate the effectiveness of the proposed critic loss, we consider a two-armed bandit example in Prop. 5 (see App. E for details) with deterministic rewards (there is no variance due to sampling), use the direct representation and tabular parameterization for the policy, linear function approximation

for the critic and compare minimizing the standard squared loss vs the decision-aware loss in Prop. 4.

Proposition 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 . Consider using linear function approximation to estimate the Q function i.e. ˆ Q = xω where ω 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 ˆ Q 1 ( ω ) = -2 ω and ˆ Q 2 ( ω ) = ω . Let p t be the probability of pulling the optimal arm at iteration t and consider minimizing two alternative objectives to estimate ω :

For p 0 < 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.

Hence, minimizing the decision-aware critic loss results in a better, more well-informed estimate of ω which when coupled with the actor update results in convergence to the optimal arm. For this simple example, at every iteration t , L t ( ω (2) t ) = 0 , while the standard squared loss is non-zero at ω (1) t , though we use the same linear function approximation model in both cases. In Prop. 16, we prove that for a 2-arm bandit with deterministic rewards and linear critic parameterization, minimizing the decision-aware critic loss will always result in convergence to the optimal arm.

5.2 Softmax representation

Recall that for the softmax functional representation, policy π is represented by the logits z π ( s, a ) for each s ∈ S and a ∈ A such that p π ( a | s ) = exp( z π ( s,a )) ∑ a ′ exp( z π ( s,a ′ )) . Using the policy gradient theorem, [ ∇ π J ( π )] s,a = d π ( s ) A π ( s, a ) p π ( a | s ) where A π is the advantage function. Similar to Vaswani et al. [57], we use a weighted (across states) log-sum-exp mirror map implying that D Φ ( z, z ′ ) = ∑ s ∈S d π t ( s ) D ϕ ( z ( s, · ) , z ′ ( s, · )) where ϕ ( z ( s, · )) = log( ∑ a exp( z ( s, a ))) and hence, D ϕ ( z ( s, · ) , z ′ ( s, · )) = KL ( p π ′ ( ·| s ) , p π ( ·| s )) (see Lemma 11 for a derivation). We now instantiate inequality (I) in Sec. 3.2 in the proposition below (see App. E for the derivation).

Proposition 6. For the softmax representation and log-sum-exp mirror map, c > 0 , η ≤ 1 -γ ,

where ˆ A π t is the estimate of the advantage function for policy π t .

For incorporating policy (with parameters θ ) and critic (with parameters ω ) parameterization, we note that p π ( a | s ) = exp( z π ( s,a | θ )) ∑ a ′ exp( z π ( s,a ′ | θ )) and ˆ A π ( s, a ) = A π ( s, a | ω ) where the model is implicit in the notation. Using the reasoning in Sec. 3.2 with Prop. 6 immediately gives us the actor and critic objectives ( ℓ t ( θ ) and L t ( ω ) respectively) at iteration t and completely instantiates Algorithm 1. Similar to the direct representation, observe that L t is asymmetric and penalizes the under/overestimation of the advantage function differently.

To demonstrate the effectiveness of the proposed critic loss, we construct a two-armed bandit example in Prop. 7 below (see App. E for details), use the softmax representation and tabular parameterization for the policy and consider a discrete hypothesis class (with two hypotheses) as the model for the critic. We compare minimizing the squared loss on the advantage: E s ∼ d πt E a ∼ p πt ( ·| s ) [ A π t ( s, a ) -A π t ( s, a | ω )] 2 with minimizing the decision-aware loss. We see that minimizing the decision-aware critic loss can distinguish between the two hypotheses and choose the correct hypothesis resulting in convergence to the optimal action.

Proposition 7. Consider a two-armed bandit example and define p ∈ [0 , 1] as the probability of pulling arm 1. Given p , let the advantage of arm 1 be equal to A 1 := 1 2 > 0 , while that of arm 2 is A 2 := -p 2 (1 -p ) < 0 implying that arm 1 is optimal. For ε ∈ ( 1 2 , 1 ) , consider approximating the advantage of the two arms using a function approximation model with two hypotheses that

depend on p : H 0 : ˆ A 1 = 1 2 + ε , ˆ A 2 = -p 1 -p ( 1 2 + ε ) and H 1 : ˆ A 1 = 1 2 -ε sgn ( 1 2 -p ) , ˆ A 2 = -p 1 -p ( 1 2 -ε sgn ( 1 2 -p )) where sgn is the signum function. If p t is the probability of pulling arm 1 at iteration t , consider minimizing two alternative loss functions to choose the hypothesis H :

t (1) Squared loss: H t = arg min {H 0 , H 1 } { p t 2 [ A 1 -ˆ A 1 ] 2 + 1 -p t 2 [ A 2 -ˆ A 2 ] 2 } . (2) Decision-aware critic loss with c = 1 : H t = arg min {H 0 , H 1 }

For p 0 ≤ 1 2 , the squared loss cannot distinguish between H 0 and 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.

In Prop. 13 in App. E, we study the softmax representation with the Euclidean mirror map and instantiate inequality (I) for this case. Finally, in App. B, we instantiate our actor-critic framework to handle stochastic value gradients used for learning continuous control policies [21]. In the next section, we consider simple RL environments to empirically benchmark Algorithm 1.

6 Experiments

--MSE-eta=0.01

MSE-eta= 0.1

---AdvMSE-eta=0.01

AdvMSE-eta=0.1

Ours-eta= 1.0

0.6

d=80

09=p

d=40

0.5

0.4-

Direct

(u)r

0.3

0.2

0.1-

-00

0.1

10000

20000

30000

40000

0

10000

20000

30000

40000

0

10000

20000

30000

40000

0.6-

0.5-

0.4

Softmax

(n)r

-E'0

0.2-

0.1

0.0-

0.1

0

10000

20000

30000

40000

0

10000

20000

30000

40000

0

10000

20000

30000

40000

#Iteration

#Iteration

#Iteration

We demonstrate the benefit of the decision-aware framework over the standard AC algorithm where the critic is trained by minimizing the squared error. We instantiate Algorithm 1 for the direct and softmax representations, and evaluate the performance on two grid-world environments, namely Cliff World [53] and Frozen Lake [6] (see App. F for details). We compare the performance of three AC algorithms that have the same actor, but differ in the objective function used to train the critic.

Critic Optimization : For the direct and softmax representations, the critic's objective is to estimate the action-value ( Q ) and advantage ( A ) functions respectively. We use a linear parameterization for the Q function implying that for each policy π , Q π ( s, a | ω ) = ⟨ ω, X ( s, a ) ⟩ , where X ( s, a ) ∈ R d are features obtained via tile-coding [53, Ch. 9]. We vary the dimension d ∈ { 80 , 60 , 40 } of the tile-coding features to vary the expressivity of the critic. Given the knowledge of p π and the estimate Q π ( s, a | ω ) , the estimated advantage can be obtained as: A π ( s, a | ω ) = Q π ( s, a | ω ) -∑ a p π ( a | s ) Q π ( s, a | ω ) . We consider two ways to estimate the Q function for training the critic: (a) using the known MDP to exactly compute the Q values and (b) estimating the Q function using Monte-Carlo (MC) rollouts. There are three sources of error for an insufficient critic - the bias due to limited model capacity, the optimization error due to an insufficient number of inner iterations ( m c ) and the variance due to Monte-Carlo sampling. We use a large value of m c and sufficiently large

Monte-Carlo rollouts to control the optimization error and variance respectively (see App. F). This ensures that the bias dominates, and enables us to isolate the effect of the form of the critic loss.

We evaluate the performance of the decision-aware loss defined for the direct (Prop. 4) and softmax representations (Prop. 6). For both representations, we minimize the corresponding objective at each iteration t (Lines 6-8 in Algorithm 1) using gradient descent with the step-size α c determined by the Armijo line-search [4]. We use a grid-search to tune the trade-off parameter c , and propose an alternative albeit conservative method to estimate c in App. F. We compare against two baselines (see App. F for implementation details) - (i) the standard squared loss on the Q functions (referred to as MSE in the plots) defined in Prop. 5 and (ii) squared loss on A function (referred to as Adv-MSE in the plots) defined in Prop. 7. We note that the Adv-MSE loss corresponds to a second-order Taylor series expansion of the decision-aware loss (see Prop. 14 for details), and is similar to the loss in Pan et al. [42]. Recall that the critic error consists of the variance when using MC samples (equal to zero when we exactly compute the Q function) and the bias because of the critic optimization error (controlled since the critic objective is convex) and error due to the limited expressivity of the linear function approximation (decreases as d increases). Since our objective is to study the effect of the critic loss and its interaction with function approximation, we do not use bootstrapping to estimate the Q π since it would result in a confounding bias term.

Actor Optimization : For all algorithms, we use the same actor objective defined for the direct (Prop. 4) and softmax representations (Prop. 6). We consider both the tabular and linear policy paramterization for the actor. For the linear function approximation, we use the same tile-coded features and set n = 60 for both environments. We update the policy parameters at each iteration t in the off-policy inner-loop (Lines 11-13 in Algorithm 1) using Armijo line-search to set α a . For details about the derivatives and closed-form solutions for the actor objective, refer to [57] and App. F. We use a grid-search to tune η , and compare different values. Our experimental setup 1 enables us to isolate the effect of the critic loss, without non-convexity or optimization issues acting as confounders.

Results : For each environment, we conduct four experiments that depend on (a) whether we use MCsamples or the true dynamics to estimate the Q function, and (b) on the policy parameterization. We only show the plot corresponding to using the true dynamics for estimating the Q function and linear policy parameterization, and defer the remaining plots to App. G. For all experiments, we report the mean and 95% confidence interval of J ( π ) averaged across 5 runs. In the main paper, we only include 2 values of η ∈ { 0 . 01 , 0 . 1 } and vary d ∈ { 40 , 60 , 80 } , and defer the complete figure with a broader range of η and d to App. G. For this experiment, c is tuned to 0 . 01 and we include a sensitivity (of J ( π ) to c ) plot in App. G. From Fig. 1, we see that (i) with a sufficiently expressive critic ( d = 80 ), all algorithms reach the optimal policy at nearly the same rate. (ii) as we decrease the critic capacity, minimizing the MSE loss does not result in monotonic improvement and converges to a sub-optimal policy, (iii) minimizing the Adv-MSE usually results in convergence to the optimal policy, whereas (iv) minimizing the decision-aware loss results in convergence to better policies at a faster rate, and is more beneficial when using a less-expressive critic (corresponding to d = 40 ). We obtain similar results for the tabular policy parameterization or when using sampling to estimate the Q function (see App. G for additional results).

7 Discussion

We designed a generic decision-aware actor-critic framework where the actor and critic are trained cooperatively to optimize a joint objective. Our framework can be used with any policy representation and easily handle general policy and critic parameterization, while preserving theoretical guarantees. Instantiating the framework resulted in an actor that supports off-policy updates, and a corresponding critic loss that can be minimized using first-order optimization. We demonstrated the benefit of our framework both theoretically and empirically. We note that Algorithm 1 can be directly used with any complex actor/critic parameterization in order to generalize across states/actions. The theoretical guarantees of Prop. 2 would still hold. From a practical perspective, w.r.t tuning hyperparameters, η does not depend on the actor/critic parameterization. On the other hand, α a and α c are set adaptively using an Armijo line-search that only requires the smoothness of actor/critic objectives, and does not depend on their convexity. However, Algorithm 1 does require tuning hyper-parameter c , and we will aim to investigate automatic adaptive ways to set it. In the future, we aim to benchmark Algorithm 1 for complex deep RL environments. Finally, we aim to broaden the scope our framework to applications such as variational inference.

References

Romina Abachi. Policy-aware model learning for policy gradient methods . PhD thesis, University of Toronto (Canada), 2020.

Alekh Agarwal, Sham M. Kakade, Jason D. Lee, and Gaurav Mahajan. Optimality and approximation with policy gradient methods in Markov decision processes. In Conference on Learning Theory (COLT) , pages 64-66, 2020.

Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural Computation , 10(2): 251-276, February 1998.

Larry Armijo. Minimization of functions having lipschitz continuous first partial derivatives. Pacific Journal of mathematics , 16(1):1-3, 1966.

Shalabh Bhatnagar, Richard S Sutton, Mohammad Ghavamzadeh, and Mark Lee. Natural actor-critic algorithms. Automatica , 45(11):2471-2482, 2009.

Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym, 2016.

Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends® in Machine Learning , 8(3-4):231-357, 2015.

Semih Cayci, Niao He, and R Srikant. Finite-time analysis of entropy-regularized neural natural actor-critic algorithm. arXiv preprint arXiv:2206.00833 , 2022.

Tianyi Chen, Yuejiao Sun, and Wotao Yin. Tighter analysis of alternating stochastic gradient method for stochastic nested problems. arXiv preprint arXiv:2106.13781 , 2021.

Bo Dai, Albert Shaw, Niao He, Lihong Li, and Le Song. Boosting the actor with dual critic. arXiv preprint arXiv:1712.10282 , 2017.

Jing Dong, Li Shen, Yinggan Xu, and Baoxiang Wang. Provably efficient convergence of primaldual actor-critic with nonlinear function approximation. arXiv preprint arXiv:2202.13863 , 2022.

Ryan D'Orazio, Nicolas Loizou, Issam Laradji, and Ioannis Mitliagkas. Stochastic mirror descent: Convergence analysis and adaptive variants via the mirror stochastic polyak stepsize. arXiv preprint arXiv:2110.15412 , 2021.

Pierluca D'Oro and Wojciech Ja´ skowski. How to learn a useful critic? model-based actiongradient-estimator policy optimization. Advances in Neural Information Processing Systems , 33:313-324, 2020.

Pierluca D'Oro, Alberto Maria Metelli, Andrea Tirinzoni, Matteo Papini, and Marcello Restelli. Gradient-aware model-based policy search. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 34, pages 3801-3808, 2020.

Radu Alexandru Dragomir, Mathieu Even, and Hadrien Hendrikx. Fast stochastic bregman gradient methods: Sharp analysis and variance reduction. In International Conference on Machine Learning , pages 2815-2825. PMLR, 2021.

Amir-massoud Farahmand. Iterative value-aware model learning. Advances in Neural Information Processing Systems , 31, 2018.

Amir-massoud Farahmand, Andre Barreto, and Daniel Nikovski. Value-aware loss function for model-based reinforcement learning. In Artificial Intelligence and Statistics , pages 1486-1494. PMLR, 2017.

Zuyue Fu, Zhuoran Yang, and Zhaoran Wang. Single-timescale actor-critic provably finds globally optimal policy. arXiv preprint arXiv:2008.00483 , 2020.

Gianluca Gorni. Conjugation and second-order properties of convex functions. Journal of Mathematical Analysis and Applications , 158(2):293-315, 1991.

Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Offpolicy maximum entropy deep reinforcement learning with a stochastic actor. arXiv preprint arXiv:1801.01290 , 2018.

Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learning continuous control policies by stochastic value gradients. In Advances in Neural Information Processing Systems , pages 2944-2952, 2015.

Mingyi Hong, Hoi-To Wai, Zhaoran Wang, and Zhuoran Yang. A two-timescale stochastic algorithm framework for bilevel optimization: Complexity analysis and application to actorcritic. SIAM Journal on Optimization , 33(1):147-180, 2023.

Shariq Iqbal and Fei Sha. Actor-attention-critic for multi-agent reinforcement learning. In International conference on machine learning , pages 2961-2970. PMLR, 2019.

Emmeran Johnson, Ciara Pike-Burke, and Patrick Rebeschini. Optimal convergence rate for exact policy mirror descent in discounted markov decision processes. arXiv preprint arXiv:2302.11381 , 2023.

Sham Kakade. A natural policy gradient. In NIPS , volume 14, pages 1531-1538, 2001.

Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. In International Conference on Machine Learning (ICML) , pages 267-274, 2002.

Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximalgradient methods under the polyak-łojasiewicz condition. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2016, Riva del Garda, Italy, September 19-23, 2016, Proceedings, Part I 16 , pages 795-811. Springer, 2016.

Sajad Khodadadian, Thinh T Doan, Justin Romberg, and Siva Theja Maguluri. Finite sample analysis of two-time-scale natural actor-critic algorithm. IEEE Transactions on Automatic Control , 2022.

Vijay Konda and John Tsitsiklis. Actor-critic algorithms. Advances in neural information processing systems , 12, 1999.

Harshat Kumar, Alec Koppel, and Alejandro Ribeiro. On the sample complexity of actor-critic method for reinforcement learning with function approximation. Machine Learning , pages 1-35, 2023.

Jonathan Wilder Lavington, Sharan Vaswani, Reza Babanezhad, Mark Schmidt, and Nicolas Le Roux. Target-based surrogates for stochastic optimization. arXiv preprint arXiv:2302.02607 , 2023.

Chongchong Li, Yue Wang, Wei Chen, Yuting Liu, Zhi-Ming Ma, and Tie-Yan Liu. Gradient information matters in policy optimization by back-propagating through model. In International Conference on Learning Representations , 2021.

Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in overparameterized non-linear systems and neural networks. Applied and Computational Harmonic Analysis , 59:85-116, 2022.

Haihao Lu, Robert M Freund, and Yurii Nesterov. Relatively smooth convex optimization by first-order methods, and applications. SIAM Journal on Optimization , 28(1):333-354, 2018.

Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorithmic framework for model-based deep reinforcement learning with theoretical guarantees. arXiv preprint arXiv:1807.03858 , 2018.

Jincheng Mei, Chenjun Xiao, Ruitong Huang, Dale Schuurmans, and Martin Müller. On principled entropy exploration in policy optimization. In IJCAI , pages 3130-3136, 2019.

Jincheng Mei, Chenjun Xiao, Csaba Szepesvari, and Dale Schuurmans. On the global convergence rates of softmax policy gradient methods. arXiv preprint arXiv:2005.06392 , 2020.

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature , 518(7540):529-533, 2015.

Shakir Mohamed, Mihaela Rosca, Michael Figurnov, and Andriy Mnih. Monte carlo gradient estimation in machine learning. The Journal of Machine Learning Research , 21(1):5183-5244, 2020.

Junhyuk Oh, Satinder Singh, and Honglak Lee. Value prediction network. Advances in neural information processing systems , 30, 2017.

Alex Olshevsky and Bahman Gharesifard. A small gain analysis of single timescale actor critic. arXiv preprint arXiv:2203.02591 , 2022.

Hsiao-Ru Pan, Nico Gürtler, Alexander Neitz, and Bernhard Schölkopf. Direct advantage estimation. Advances in Neural Information Processing Systems , 35:11869-11880, 2022.

Jan Peters, Sethu Vijayakumar, and Stefan Schaal. Natural actor-critic. In Machine Learning: ECML 2005: 16th European Conference on Machine Learning, Porto, Portugal, October 3-7, 2005. Proceedings 16 , pages 280-291. Springer, 2005.

Martin L. Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming . John Wiley & Sons, Inc., USA, 1994.

Andrew J Schaefer, Matthew D Bailey, Steven M Shechter, and Mark S Roberts. Modeling medical treatment using Markov decision processes. In Operations research and health care , pages 593-612. Springer, 2005.

John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International Conference on Machine Learning (ICML) , pages 1889-1897, 2015.

John Schulman, Xi Chen, and Pieter Abbeel. Equivalence between policy gradients and soft q-learning. arXiv preprint arXiv:1704.06440 , 2017.

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. CoRR , abs/1707.06347, 2017.

David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. Journal of Machine Learning Research , 2014.

David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. Nature , 529(7587):484, 2016.

David Silver, Hado Hasselt, Matteo Hessel, Tom Schaul, Arthur Guez, Tim Harley, Gabriel Dulac-Arnold, David Reichert, Neil Rabinowitz, Andre Barreto, et al. The predictron: End-toend learning and planning. In International Conference on Machine Learning , pages 3191-3199. PMLR, 2017.

Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning , 3(1):9-44, 1988.

Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction . MIT Press, 2 edition, 2018.

Richard S Sutton, David A McAllester, Satinder P Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In Advances in Neural Information Processing Systems (NeurIPS) , pages 1057-1063, 2000.

Jie Tan, Tingnan Zhang, Erwin Coumans, Atil Iscen, Yunfei Bai, Danijar Hafner, Steven Bohez, and Vincent Vanhoucke. Sim-to-real: Learning agile locomotion for quadruped robots. arXiv preprint arXiv:1804.10332 , 2018.

Manan Tomar, Lior Shani, Yonathan Efroni, and Mohammad Ghavamzadeh. Mirror descent policy optimization. arXiv preprint arXiv:2005.09814 , 2020.

Sharan Vaswani, Olivier Bachem, Simone Totaro, Robert Müller, Shivam Garg, Matthieu Geist, Marlos C Machado, Pablo Samuel Castro, and Nicolas Le Roux. A general class of surrogate functions for stable and efficient reinforcement learning. arXiv preprint arXiv:2108.05828 , 2021.

Christopher John Cornish Hellaby Watkins. Learning from delayed rewards. 1989.

Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning , 8(3-4):229-256, 1992.

Yue Frank Wu, Weitong Zhang, Pan Xu, and Quanquan Gu. A finite-time analysis of two time-scale actor-critic methods. Advances in Neural Information Processing Systems , 33: 17617-17628, 2020.

Lin Xiao. On the convergence rates of policy gradient methods. Journal of Machine Learning Research , 23(282):1-36, 2022.

Tengyu Xu, Zhe Wang, and Yingbin Liang. Non-asymptotic convergence analysis of two time-scale (natural) actor-critic algorithms. arXiv preprint arXiv:2005.03557 , 2020.

Rui Yuan, Robert M Gower, and Alessandro Lazaric. A general sample complexity analysis of vanilla policy gradient. In International Conference on Artificial Intelligence and Statistics , pages 3332-3380. PMLR, 2022.

Andy Zeng, Shuran Song, Johnny Lee, Alberto Rodriguez, and Thomas Funkhouser. Tossingbot: Learning to throw arbitrary objects with residual physics. IEEE Transactions on Robotics , 36 (4):1307-1319, 2020.

Siqi Zhang and Niao He. On the convergence rate of stochastic mirror descent for nonsmooth nonconvex optimization. arXiv preprint arXiv:1806.04781 , 2018.

Organization of the Appendix

A Definitions

B Extension to stochastic value gradients

C Proofs for Sec. 3

D Proofs for Sec. 4

E Proofs for Sec. 5

F Implementation Details

G Additional Experiments

A Definitions

[Solution set] . We define the solution set X ∗ for a function f as X ∗ := { x ∗ | x ∗ ∈ arg min x ∈ dom ( f ) f ( x ) } .

[Convexity] . A differentiable function f is convex iff for all v and w in dom ( f )

[Lipschitz continuity] . A differentiable function f is G -Lipschitz continuous, meaning that for all v and w and constant G > ,

[Smoothness] . A differentiable function f is L -smooth, meaning that for all v and w and some constant L > 0

[Polyak-Lojasiewicz inequality] . A differentiable function f satisfies the Polyak-Lojasiewicz (PL) inequality if there exists a constant µ p > 0 s.t. for all v ,

where f ∗ is the optimal function value i.e. f ∗ := f ( x ∗ ) for x ∗ ∈ X ∗ .

[Restricted Secant Inequality] . A differentiable function f satisfies the Restricted Secant Inequality (RSI) inequality if there exists a constant µ r > 0 that for all v

Supplementary material

where v p is the projection of v onto X ∗ .

[Bregman divergence] . For a strictly-convex, differentiable function Φ , we define the Bregman divergence induced by Φ (known as the mirror map) as:

[Relative smoothness] . A function f is ρ -relatively smooth w.r.t. D Φ iff f + ρ Φ is convex. Furthermore, if f is ρ -relatively smooth w.r.t. Φ , then, | f ( w ) -f ( v ) -⟨∇ f ( v ) , w -v ⟩| ≤ ρD Φ ( w,v ) .

[Mirror Ascent] . Optimizing max x ∈X f ( x ) using mirror ascent (MA), if x t is the current iterate, then the update at iteration t ∈ { 0 , 1 , . . . , T -1 } with a step-size η t and mirror map Φ is given as:

The above update can be formulated into two steps Bubeck [7, Chapter 4] as follows:

B Extension to stochastic value gradients

In Sec. 5, we have seen alternative ways to represent a policy's conditional distributions over actions p π ( ·| s ) for each state s ∈ S . On the other hand, stochastic value gradients [21] represent a policy by a set of actions. Formally, if ε are random variables drawn from a fixed distribution χ , then policy π is a deterministic map from S × χ → A . This corresponds to the functional representation of the policy, and is particularly helpful for continuous control, i.e. when the action-space is continuous. The action a chosen by π in state s , when fixing the random variable ε = ϵ , is represented as π ( s, ϵ ) , and the value function for policy π is given as:

and Silver et al. [49] showed that ∂J ( π ) ∂π ( s, ϵ ) = d π ( s ) ∇ a Q π ( s, a ) ∣ ∣ a = π ( s,ϵ ) . In order to characterize the dependence on the policy parameterization, we note that π ( s, ϵ ) = π ( s, ϵ, θ ) where θ are the model parameters. For a fixed ϵ , we will use a Euclidean mirror map implying that D Φ ( π, π ′ ) = ∑ s ∈S d π t ( s ) D ϕ ( π ( s, ϵ ) , π ′ ( s, ϵ ) and choose ϕ ( π ( s, ϵ )) = 1 2 ∥ π ′ ( s, ϵ ) ∥ 2 2 implying that D ϕ ( π ( s, ϵ ) , π ′ ( s, ϵ ) = 1 2 [ π ( s, ϵ ) -π ′ ( s, ϵ )] 2 . In order to instantiate the generic lower bound in Prop. 1 at iteration t , we prove the following proposition in App. E.

Proposition 8. For the stochastic value gradient representation and Euclidean mirror map, c > 0 , η such that J + 1 η Φ is convex in π .

where C is a constant and ̂ ∇ a Q π t ( s, a ) ∣ ∣ a = π t ( s,ε ) is the estimate of the action-value gradients for policy π at state s and a = π t ( s, ϵ ) .

For incorporating policy (with parameters θ ) and critic (with parameters ω ) parameterization, we note that π ( s, ε ) = π ( s, ε | θ ) and ̂ ∇ a Q π t ( s, a ) a = π t ( s,ε ) = ∇ a Q π t ( s, a | ω ) a = π t ( s,ε,θ t ) where the model is implicit in the notation. Using the reasoning in Sec. 3.2 with Prop. 8 immediately gives us the actor and critic objectives ( ℓ t ( θ ) and L t ( ω ) respectively) at iteration t and completely instantiates Algorithm 1. The actor objective is similar to Eq (15) of Silver et al. [49], with the easier to compute Q π t instead of Q π , whereas the critic objective is similar to the one used in existing work on policy-aware model-based RL for continuous control [13].

C Proofs for Sec. 3

Proposition 1. For any policy representations π and π ′ , any strictly convex mirror map Φ , and any gradient estimator ˆ g , for c > 0 and η such that J + 1 η Φ is convex in π ,

where Φ ∗ is the Fenchel conjugate of Φ and D Φ ∗ is the Bregman divergence induced by Φ ∗ .

Proof. For any η such that J + 1 η Φ is convex, we use Lemma 2 to form the following lower-bound,

Defining δ := ∇ J ( π ′ ) -ˆ g ( π ′ ) , and assuming that c δ is small enough to satisfy the requirement for Lemma 1, we use Lemma 1 with x = δ , y = π and y ′ = π ′ .

Lemma 1 (Bregman Fenchel-Young) . Let x ∈ Y ∗ , y ∈ Y , y ′ ∈ Y . Then, for sufficiently small c > 0 and x s.t. ( ∇ ϕ ) -1 [ ∇ ϕ ( y ′ ) -c x ] ∈ Y , we have

For a fixed y ′ , this inequality is tight for y = arg min υ { ⟨ x, υ -y ′ ⟩ + 1 c D Φ ( υ, y ) } .

Proof. Define f ( y ) := ⟨ x, y -y ′ ⟩ + 1 c D Φ ( y, y ′ ) . If y ∗ = arg min f ( y ) , then,

Note that according to our assumption, y ∗ ∈ Y . For any y ,

In order to simplify D Φ ( y ∗ , y ′ ) , we will use the definition of ϕ ∗ ( z ) . In particular, for any y ,

(using Eq. (4) to simplify the first term)

Let us focus on the first term and simplify it,

Using the above relations,

Using the above simplification with Eq. (3),

′

Using the definition of f ( y ) ,

Lemma 2. If J + 1 η Φ is convex, then, J ( π ) is 1 η -relatively smooth w.r.t to D Φ , and satisfies the following inequality,

′

′

∗

′

′

′

′

∗

′

Proof. If J + 1 η Φ is convex,

D Proofs for Sec. 4

Proposition 2. For any policy representation and any policy or critic parameterization, there exists a ( θ, c ) pair that makes the RHS of inequality (I) strictly positive, and hence guarantees monotonic policy improvement ( J ( π t +1 ) > J ( π t ) ), if and only if

where b t ∈ R n := ∑ s ∈S ∑ a ∈A [ˆ g t ] s,a ∇ θ [ π ( θ t )] s,a , ˜ H t ∈ R n × n := ∇ θ π ( θ t ) T ∇ 2 π Φ( π t ) ∇ θ π ( θ t ) and ˜ H † t denotes the pseudo-inverse of ˜ H t . For the special case of the tabular policy parameterization, the above condition becomes equal to,

Proof. As a warmup, let us first consider the tabular parameterization where π ( θ ) = θ ∈ R SA . In this case, the lower-bound in Prop. 1 is equal to,

We shall do a second-order Taylor expansion of the critic objective (blue term) in c around 0 and a second-order Taylor expansion of the actor objective (green term) around θ = θ t . Defining δ := ∇ J ( θ t ) -ˆ g t ,

where o ( c ) and o ( ∥ θ -θ t ∥ 2 2 ) consist of the higher order terms in the Taylor series expansion. A necessary and sufficient condition for monotonic improvement is equivalent to finding a ( θ, c ) such that RHS is positive. As c tends to 0, the θ maximizing the RHS is

With this choice,

If ⟨ ˆ g t , [ ∇ 2 Φ( θ t ) ] -1 ˆ g t ⟩ > ⟨ δ, ∇ 2 ϕ ∗ ( ∇ ϕ ( π t )) δ ⟩ , i.e. there exists an ϵ > 0 s.t. ⟨ ˆ g t , [ ∇ 2 Φ( θ t ) ] -1 ˆ g t ⟩ = ⟨ δ, ∇ 2 ϕ ∗ ( ∇ ϕ ( π t )) δ ⟩ + ϵ , then, RHS = cϵ 2 + o ( c ) . For any fixed κ > 0 , since o ( c ) /c → 0 as c → 0 , there exists a neighbourhood (0 , c κ ) around zero such that for all c in this neighbourhood, o ( c ) /c > -κ and hence o ( c ) > -κc . Setting κ = ε 4 , there is a c such that

Hence, there exists a c ∈ (0 , min { η, c κ } ) such that the RHS is positive, and is hence sufficient to guarantee monotonic policy improvement.

On the other hand, if ⟨ ˆ g t , [ ∇ 2 Φ( θ t ) ] -1 ˆ g t ⟩ < ⟨ δ, ∇ 2 ϕ ∗ ( ∇ ϕ ( π t )) δ ⟩ , i.e. there exists an ϵ > 0 s.t. ⟨ ˆ g t , [ ∇ 2 Φ( θ t ) ] -1 ˆ g t ⟩ = ⟨ δ, ∇ 2 ϕ ∗ ( ∇ ϕ ( π t )) δ ⟩ -ϵ , then, RHS = -cϵ 2 + o ( c ) which can be negative and hence monotonic improvement can not be guaranteed. Hence, ⟨ ˆ g t , [ ∇ 2 Φ( θ t ) ] -1 ˆ g t ⟩ > ⟨ δ, ∇ 2 ϕ ∗ ( ∇ ϕ ( π t )) δ ⟩ is a necessary and sufficient condition for improvement.

Let us now consider the more general case, and define m = SA , ˆ g t ∈ R m × 1 , π ( θ ) ∈ R m × 1 is a function of θ ∈ R n × 1 . Rewriting Prop. 1,

As before, we shall do a second-order Taylor expansion of the critic objective (blue term) in c around 0 and a second-order Taylor expansion of the actor objective (green term) around θ = θ t . Defining δ := ∇ J ( θ t ) -ˆ g t . From Prop. 14, we know that,

In order to calculate the second-order Taylor series expansion of the actor objective, we define ∇ θ π ( θ t ) ∈ R m × n as the Jacobian of the θ : → π map, and use ∇ θ [ π ( θ t )] i ∈ R 1 × n for i ∈ [ m ] to refer to row i

where o ( ∥ θ -θ t ∥ 2 2 ) consist of the higher order terms in the Taylor series expansion. For expanding the divergence term, note that D Φ ( π ( θ ) , π ( θ t )) = ϕ ( π ( θ )) -ϕ ( π ( θ t )) -⟨∇ ϕ ( π ( θ t )) , π ( θ ) -π ( θ t ) ⟩

Putting everything together,

Defining b t := ∑ m i =1 [ˆ g t ] i ∇ θ [ π ( θ t )] i and H t := ∇ θ π ( θ t ) T ∇ 2 π ϕ ( π t ) ∇ θ π ( θ t ) -1 ( 1 η + 1 c ) ∑ m i =1 ( [ˆ g t ] i ∇ 2 θ [ π ( θ t )] i )

As a sanity check, it can be verified that if π ( θ ) = θ , H t = ( 1 η + 1 c ) ∇ 2 Φ( θ t ) and b t = ˆ g t , and we recover the tabular result above. Notice that 〈 ( θ -θ t ) , 1 ( 1 η + 1 c ) ∑ m i =1 ( [ˆ g t ] i ∇ 2 θ [ π ( θ t )] i ) ( θ -θ t ) 〉 is o ( c ) as c goes to zero. Subsuming this term in o ( c ) ,

where ˜ H t := ∇ θ π ( θ t ) T ∇ 2 π ϕ ( π t ) ∇ θ π ( θ t ) . As before, a necessary and sufficient condition for monotonic improvement is equivalent to finding a ( θ, c ) such that RHS is positive. As c tends to 0, the θ maximizing the RHS is

With this choice,

As in the tabular case, since 1 1 η + 1 c is o ( c ) , we can subsume it, and we get that,

Using the same reasoning as in the tabular case above, we can prove that

is a necessary and sufficient condition for monotonic policy improvement. Finally, we use Gorni [19, Eq (1.5)] which shows that ∇ 2 ϕ ∗ ( ∇ ϕ ( π t )) = [ ∇ 2 π Φ( π t )] -1 and complete the proof.

D.1 Proof of Prop. 3

The following proposition shows the convergence of inexact mirror ascent in the functional space.

Proposition 9. Assuming that (i) J + 1 η Φ is convex in π , for a constant c > 0 , after T iterations of mirror ascent with 1 η ′ = 2 η + 2 c we have

where ζ = η ′ / 2 and R is picked uniformly random from { 0 , 1 , 2 , . . . T -1 } .

Proof. We divide the mirror ascent (MA) update into two steps:

Recursing for T iterations and dividing by 1 /ζ , picking R uniformly random from { 0 , 1 , 2 , . . . T -1 } and taking expectation we get

Compared to Dragomir et al. [15], D'Orazio et al. [12] that analyze stochastic mirror ascent in the smooth, non-convex setting, our analysis ensures that the (i) sub-optimality gap (the LHS in the above proposition) is always positive, and (ii) uses a different notion of variance that depends on D Φ ∗ .

Similar to Vaswani et al. [57], we assume that each π ∈ Π is parameterized by θ . In Algorithm 1, we run gradient ascent (GA) on ℓ t ( θ ) to compute π t +1 = π ( θ t +1 ) and interpret the inner loop of Algorithm 1 as an approximation to the projection step in the mirror ascent update. We note that ℓ t ( θ ) does not have any additional randomness and is a deterministic function w.r.t θ . Note that ¯ π t +1 = where π t = π ( θ t ) . Assuming that ℓ t ( θ ) is smooth, and satisfies the PL condition [27], we get the following convergence guarantee for Algorithm 1.

Proposition 3. For any policy representation and mirror map Φ such that (i) J + 1 η Φ is convex in π , any policy parameterization such that (ii) ℓ t ( θ ) is smooth w.r.t θ and satisfies the Polyak-Lojasiewicz (PL) condition, for c > 0 , after T iterations of Algorithm 1 we have that,











where δ t := ∇ J ( π t ) -ˆ g t , 1 ζ := 1 η + 1 c , R is a random variable chosen uniformly from { 0 , 1 , 2 , . . . T -1 } and e t ∈ O (exp ( -m a )) is the projection error (onto Π ) at iteration t .

Proof. For this proof, we define the following notation:

(Iterate obtained after running 1 step of mirror ascent starting from π )

where GradientAscent ( ℓ t ( θ ) , θ t , m a ) means running GradientAscent for m a iterattions on ℓ t ( θ ) with the initialization equal to θ t . Since we assume that ℓ t satisfies the PL -condition w.r.t. θ for all t , based on the results from Karimi et al. [27], we get

where c 1 , c 2 are problem-dependent constants related to the smoothness and curvature of ℓ t , and e t is the approximation error diminishes as we increase the value of m a . Following the same steps as before,

Recusing for T iterations and dividing by 1 /ζ , picking R uniformly random from { 0 , 1 , 2 , . . . T -1 } and taking expectation we get

Note that it is possible to incorporate a sampling error (in the distribution d π across states) for the actor update in Algorithm 1. This corresponds to an additional error in calculating D Φ , and we can use the techniques from Lavington et al. [31] to characterize the convergence in this case.

D.2 Exact setting with Lifting (Direct representation)

Recall the mirror ascent update in the functional space.

For the direct representation, we define π s := p π ( ·| s ) , ˆ g ( π t )( s, · ) = d π t ( s ) Q π t ( s, · ) and D Φ ( π, π t ) = ∑ s d π t ( s ) D ϕ ( π s , π t s ) . Rewriting the MA update,

For each state s and Q π t ( s, . ) we define the set Π s t = { π s : π s ∈ arg max p s ∈ ∆ A ⟨ Q π t ( s, . ) , p s ⟩} i.e. a set of greedy policies w.r.t. Q π t ( s, . ) . Similar to Johnson et al. [24] we define η ′ t as follows

where c t > 0 is a constant. Now we consider the policy parameterization, π = π ( θ ) . We assume that the mapping from θ → π is L π Lipschitz continuous.

GradientAscent ( ℓ t ( θ ) , θ t , m ) means that we run gradient ascent for m iterations to maximize ℓ t with θ t as the initial value. We assume that ℓ t satisfies Restricted Secant Inequality (RSI) and is smooth w.r.t. θ . Based on the convergence property of Gradient Ascent for RSI and smooth functions [27], we have:

Furthermore, we assume that ∥ ¯ π t +1 -˜ π t +1 ∥ 2 2 ≤ b 2 t for all t which represents the bias because of the function approximation. Before stating the main proposition of this section, we restate Johnson et al. [24, Lemma 2].

Lemma 3 (Lemma 2 of Johnson et al. [24]) . For all ( s, a ) ∈ S × A we have

Now we state the main proposition of this part.

Proposition 10 (Convergence of tabular MDP with Lifting) . Assume that (i) ℓ t ( θ ) is smooth and satisfies RSI condition, (ii) π ( θ ) is L π -Lipschitz continuous, (iii) the bias is bounded for all t i.e. ∥ ¯ π t +1 -˜ π t +1 ∥ 2 2 ≤ b 2 t , (iv) ∥ Q π ( s, · ) ∥ ≤ q for all π and s . By setting η ′ t as in Eq. (5) and running Gradient Ascent for m iterations to maximize ℓ t we have

where π ∗ is the optimal policy, π ∗ s refers to the optimal action in state s . Here, e t = O (exp( -m )) is the approximation error.

Proof. This proof is mainly based on the proof of Theorem 3 of Johnson et al. [24]. Using Lemma 3 and the fact that π s ≥ 0 , we have ⟨ Q π t ( s, · ) , ¯ π s t +1 ⟩ ≤ ⟨ Q ¯ π t +1 ( s, · ) , ¯ π s t +1 ⟩ = J s (¯ π t +1 ) . Using this inequality we get,

The last inequality is from the definition of Q and J as follows. For any action a ,

(Using Lemma 5 with d = Q π t ( s, · ) , y = ¯ π s t +1 , x = p s t )

From the above inequality,

= ⇒ -γ ∥ J ( π t ) -J ( π ∗ ) ∥ ∞ + J s ( π ∗ ) -J s (¯ π t +1 ) ≤ min p s t ∈ Π s t D ϕ ( p s t , π t s ) η ′ t ≤ c t (Based on the definition of η ′ in Eq. (5)) ⇒ -∥ -∗ ∥ ′ ∗ -′ ≤ ′ -′

(Since s is an arbitrary state, changing s = s for convenience)

= γ J ( π t ) J ( π ) ∞ + J s ( π ) J s ( π t +1 ) c t + J s (¯ π t +1 ) J s ( π t +1 ) ′

(Using performance difference lemma 6 with the starting state equal to s ′ )

(Cauchy Schwartz)

Since the above equation is true for all s ′ we have:

Recursing for T iterations we get:

We can control the approximation error e t by using a larger m . The bias term b t can be small if our function approximation model is expressive enough. c t is an arbitrary value and if we set c t = γ t c for some constant c > 0 , then ∑ T t =1 γ -t ( c t ) = Tc and therefore γ T Tc can diminish linearly. The above analysis relied on the knowledge of the true Q functions, but can be easily extended to using inexact estimates of Q π by using the techniques developed in [61, 24].

D.3 Exact setting with lifting trick (Softmax representation)

In the softmax representation in the tabular MDP, we consider the case that π is parameterized with parameter θ ∈ R n . In this setting Φ is the Euclidean norm. Using Prop. 1, for η such that J + 1 η ϕ is convex we have for a given π t ,

If we maximize h ( π ) w.r.t. π we get

Mei et al. [37, Lemma 8] proves that J ( π ) satisfies a gradient domination condition w.r.t the softmax representation. In particular, if a ∗ ( s ) is the optimal action in state s and µ := min π min s p π ( a ∗ ( s ) | s ) √ ∥ ∗ ∥ , they prove that for all π ,

Consider optimization in the parameter space where ℓ t ( θ ) := J ( π t ) + ⟨∇ J ( π ( θ t ) , π ( θ ) -π ( θ t ) ⟩ -1 η D Φ ( π ( θ ) , π ( θ t )) .

GradientAscent ( ℓ t ( θ ) , θ t , m ) means that we run gradient ascent for m iterations to maximize ℓ t with θ t as the initial value. Assuming that ℓ t is Lipschitz smooth w.r.t. θ and satisfies the Polyak-Lojasiewicz (PL) condition, we use the gradient ascent property for PL functions [27] to obtain,

Proposition 11 (Convergence of softmax+tabular setting with Lifting) . Assume (i) J + 1 η ϕ is convex, (ii) J satisfies gradient domination property above with µ > 0 , (iii) ℓ t ( θ ) is Lipschitz smooth and satisfies PL condition, (iv) | h (¯ π t +1 ) -h (˜ π t +1 ) | ≤ b t for all t . Then after running Gradient Ascent for m iterations to maximize ℓ t we have

where α := η µ 2 2 and e t is the approximation error at iteration t and [ T -1] := { 0 , 1 , 2 , . . . T -1 } .

Proof. Since J + 1 η ϕ is convex,

Summing up for T iterations and dividing both sides by T

The above analysis relied on the knowledge of the exact gradient ∇ J ( π ) , but can be easily extended to using inexact estimates of the gradient by using the techniques developed in [63].

D.4 Helper Lemmas

Lemma 4 (3-Point Bregman Property) . For x, y, z ∈ X ,

Lemma 5 (3-Point Descent Lemma for Mirror Ascent) . For any z ∈ rint dom ϕ , and a vector d , let

Then y ∈ rint dom ϕ and for any x ∈ X

Lemma 6 (Performance Difference Lemma [26]) . For any π , π ′ ∈ Π ,

E Proofs for Sec. 5

Proposition 12 (State-wise lower bound) . For (i) any representation π that is separable across states i.e. there exists π s ∈ R A such that π s,a = [ π s ] a , (ii) any strictly convex mirror map Φ that induces a Bregman divergence that is separable across states i.e. D Φ ( π, π ′ ) = ∑ s d π ( s ) D ϕ ( π s , π ′ s ) , (iii) any η such that J + 1 η Φ is convex, if (iv) ∇ J ( π ) is separable across states i.e. [ ∇ J ( π )] s,a = d π ( s ) [ ∇ π s J ( π ) a where ∇ π s J ( π ) ∈ R A , then (v) for any separable (across states) gradient estimator ˆ g i.e. [ˆ g ( π )] s,a = d π ( s ) [ˆ g s ( π )] a where ˆ g s ( π ) ∈ R A , and c ∈ (0 , ∞ ) S ,

Proof. Using condition (iii) of the proposition with Lemma 2,

Using conditions (iv) and (v), we know that [ ∇ J ( π t )] s,a = d π t ( s ) [ ∇ π s J ( π t )] a and [ˆ g ( π t )] s,a = d π t ( s ) [ˆ g s ( π t )] a . Defining δ s t := ∇ π s J ( π t ) -ˆ g s ( π t ) ∈ R A . Using conditions (i) and (ii), we can rewrite the lower-bound as follows,

Using Lemma 1 with x = δ s t , y = π s and y ′ = π t s ,

Proposition 4. For the direct representation and negative entropy mirror map, c > 0 , η ≤ (1 -γ ) 3 2 γ | A | ,

where C is a constant and ˆ Q π t is the estimate of the action-value function for policy π t .

Proof. For the direct representation, π s,a = p π ( a | s ) . Using the policy gradient theorem, [ ∇ π J ( π )] s,a = d π ( s ) Q π ( s, a ) . We choose ˆ g ( π ) such that [ˆ g ( π )] s,a = d π ( s ) ˆ Q π ( s, a ) as the estimated gradient. Using Vaswani et al. [57, Proposition 2], J + 1 η Φ is convex for η ≤ (1 -γ ) 3 2 γ | A | . Defining δ s t := ∇ π s J ( π t ) -ˆ g s ( π t ) = Q π t ( s, · ) -ˆ Q π t ( s, · ) ∈ R A , and using Prop. 12 with c s = c for all s ,

Since ϕ ( π s ) = ϕ ( p π ( ·| s )) = ∑ a p π ( a | s ) log( p π ( a | s )) , using Lemma 10, D ϕ ( π s , π t s ) = KL ( p π ( ·| s ) || p π t ( ·| s )) . Hence,

Using Lemma 7 to simplify the last term,

Putting everything together,

Proposition 6. For the softmax representation and log-sum-exp mirror map, c > 0 , η ≤ 1 -γ ,

where ˆ A π t is the estimate of the advantage function for policy π t .

Proof. For the softmax representation, π s,a = z ( s, a ) s.t. p π ( a | s ) = exp( z ( s,a )) ∑ a ′ exp( z ( s,a ′ )) . Using the policy gradient theorem, [ ∇ π J ( π )] s,a = d π ( s ) p π ( a | s ) A π ( s, a ) . We choose ˆ g ( π ) such that [ˆ g ( π )] s,a = d π ( s ) p π ( a | s ) ˆ A π ( s, a ) as the estimated gradient. Using Vaswani et al. [57, Proposition 3], J + 1 η Φ is convex for η ≤ 1 -γ . Define δ s ∈ R A such that δ s t [ a ] := ∇ π s J ( π t ) -ˆ g s ( π t ) = p π t ( a | s ) [ A π t ( s, a ) -ˆ A π t ( s, a )] . Using Prop. 12 with c s = c for all s ,

Since ϕ ( π s ) = ϕ ( z ( s, · )) = log ( ∑ a exp( z ( s, a ))) , using Lemma 11, D ϕ ( π s , π t s ) = KL ( p π t ( ·| s ) || p π ( ·| s )) where p π ( a | s ) = exp( z ( s,a )) ∑ a ′ exp( z ( s,a ′ )) and p π t ( a | s ) = exp( z t ( s,a )) ∑ a ′ exp( z t ( s,a ′ )) . Hence, the above bound can be simplified as,

Using Lemma 8 to simplify the last term,

Putting everything together,

Let us focus on simplifying ∑ a [ p π t ( a | s ) ˆ A π t ( s, a ) [ z ( s, a ) -z t ( s, a )] ] for a fixed s . Note that ∑ a p π t ( a | s ) ˆ A π t ( s, a ) = 0 = ⇒ log ( ∑ a ′ exp( z ( s, a ′ ))) ∑ a p π t ( a | s ) ˆ A π t ( s, a ) = 0 .

Similarly, simplifying ∑ a [ p π t ( a | s ) ˆ A π t ( s, a ) z t ( s, a ) ]

Using the above relations,

Proposition 8. For the stochastic value gradient representation and Euclidean mirror map, c > 0 , η such that J + 1 η Φ is convex in π .

where C is a constant and ̂ ∇ a Q π t ( s, a ) ∣ ∣ a = π t ( s,ε ) is the estimate of the action-value gradients for policy π at state s and a = π t ( s, ϵ ) .

Proof. For stochastic value gradients with a fixed ε , ∂J ( π ) ∂π ( s, ϵ ) = d π ( s ) ∇ a Q π ( s, a ) ∣ ∣ a = π ( s,ϵ ) . We choose ˆ g ( π ) such that [ˆ g ( π )] s,a = d π ( s ) ̂ ∇ a Q π ( s, a ) ∣ ∣ a = π ( s,ϵ ) . Define δ s t ∈ R A such that δ s t [ a ] := ∇ a Q π t ( s, a ) ∣ ∣ a = π t ( s,ϵ ) -̂ ∇ a Q π t ( s, a ) ∣ ∣ a = π t ( s,ϵ ) . Using Prop. 12 with c s = c for all s ,

For a fixed ε , since ϕ ( π s ) = ϕ ( π ( s, ϵ )) = 1 2 [ π ( s, ϵ )] 2 , D ϕ ( π s , π t s ) = 1 2 [ π ( s, ϵ ) -π t ( s, ϵ )] 2 . Hence,

Simplifying the last term, since ϕ ( π ( s, ϵ )) = 1 2 [ π ( s, ϵ )] 2 ,

Putting everything together,

Proposition 13. For the softmax representation and Euclidean mirror map, c > 0 , η ≤ (1 -γ ) 3 8 then

where C is a constant and ˆ A π is the estimate of advantage function for policy π t .

Proof. For the softmax representation, π s,a = z ( s, a ) s.t. p π ( a | s ) = exp( z ( s,a )) ∑ a ′ exp( z ( s,a ′ )) . Using the policy gradient theorem, [ ∇ π J ( π )] s,a = d π ( s ) p π ( a | s ) A π ( s, a ) . We choose ˆ g ( π ) such that [ˆ g ( π )] s,a = d π ( s ) p π ( a | s ) ˆ A π ( s, a ) as the estimated gradient. Define δ s ∈ R A such that δ s t [ a ] := ∇ π s J ( π t ) -ˆ g s ( π t ) = p π t ( a | s ) [ A π t ( s, a ) -ˆ A π t ( s, a )] . Using Mei et al. [37, Lemma 7], J + 1 η Φ is convex for η ≤ 1 -γ . Using Prop. 12 with c s = c for all s ,

Since ϕ ( π s ) = ϕ ( z ( s, · )) = 1 2 ∑ a [ z s,a ] 2 , D ϕ ( π s , π t s ) = 1 2 ∥ z ( s, · ) -z t ( s, · ) ∥ 2 2 . Hence,

Simplifying the last term, since ϕ ( z ( · , a )) = 1 2 [ z ( s, a )] 2 ,

Putting everything together,

Proposition 14. For both the direct (with the negative-entropy mirror map) and softmax representations (with the logsum-exp mirror map), for a fixed state s , if δ ∈ R A := ∇ π s J ( π t ) -ˆ g s ( π t ) , the second-order Taylor expansion of f ( c ) = D ϕ ∗ ( ∇ ϕ ( π t s ) -cδ, ∇ ϕ ( π t s )) around c = 0 is equal to

Proof.

By the second-order Taylor series expansion of f ( c ) around c = 0 ,

Let us first consider the softmax case with the log-sum-exp mirror map, where π s = z ( s, · ) and ϕ ( z ( s, · )) = log( ∑ a exp( z ( s, a ))) , ϕ ∗ ( p π ( ·| s )) = ∑ a p π ( a | s ) log( p π ( a | s )) . Since the negative entropy and log-sum-exp are Fenchel conjugates (see Lemma 9), ∇ ϕ ( z t ( s, · )) = p π t ( ·| s ) . Hence, we need to compute ∇ 2 ϕ ∗ ( p π t ( ·| s )) .

For the softmax representation, using the policy gradient theorem, [ δ ] a = p π t ( a | s )[ A ( s, a ) -ˆ A ( s, a )] and hence,

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 π s = p π ( ·| s ) , ϕ ( p π ( ·| s )) = ∑ a p π ( a | s ) log( p π ( a | s )) , ϕ ∗ ( z ( s, · )) = log( ∑ a exp( z ( s, a ))) . Since the negative entropy and log-sum-exp are Fenchel conjugates (see Lemma 9), ∇ ϕ ( p π t ( ·| s )) =

z t ( s, · ) . Hence, we need to compute ∇ 2 ϕ ∗ ( z t ( s, · )) .

For the direct representation, using the policy gradient theorem, [ δ ] a = Q π t ( s, a ) -ˆ Q π t ( s, a ) and hence,

(where ˆ J s ( π 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,

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: ˆ Q = xω where ω 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 ˆ Q 1 ( ω ) = -2 ω and ˆ Q 2 ( ω ) = ω . Let p t be the probability of pulling the optimal arm at iteration t , and consider minimizing two alternative objectives to estimate ω :

Using the tabular parameterization for the actor, the policy update at iteration t is given by: p t +1 = p t exp( η ˆ Q 1 ) p t exp( η ˆ Q 1 )+(1 -p t ) exp( η ˆ Q 2 ) , where η is the functional step-size for the actor. For p 0 < 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 ˆ Q 1 ( ω ) -Q 1 = -2( ω +1) and ˆ Q 2 ( ω ) -Q 2 = ω -1 . Calculating ω (1) for a general policy s.t. the probability of pulling the optimal arm equal to p ,

Setting the gradient to zero,

Calculating ω (2) for a general policy s.t. the probability of pulling the optimal arm equal to p ,

Setting the gradient to zero,

Define A := exp( -2 c ( ω +1)) and B := exp( c ( ω -1)))

Now, let us consider the actor update,

p

Since arm 1 is optimal, if t +1 p t < 1 for all t , the algorithm will converge to the sub-optimal arm. This happens when 1 p t +(1 -p t ) exp( η ( ˆ Q 2 -ˆ Q 1 )) < 1 = ⇒ ˆ Q 2 -ˆ Q 1 > 0 = ⇒ ω > 0 . Hence, for any η and any iteration t , if ω t > 0 , p t +1 < p t . For the decision-aware critic loss, ω (2) t = -1 3 for all t , implying that p t +1 > p t and hence the algorithm will converge to the optimal policy for any η and any initialization p 0 > 0 . However, for the squared MSE loss, ω (2) t = 1 -5 p t 3 p t +1 , ω (2) t > 0 if p t < 1 5 Hence, if p 0 < 1 5 , p 1 < p 0 < 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 ≥ Q 2 . Using a linear parameterization for the critic, Q function is estimated as: ˆ Q = xω where ω 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 ω : ω t := arg min L t ( ω ) := p t [ Q 1 -ˆ Q 1 ( ω )]+(1 -p t ) [ Q 2 -ˆ Q 2 ( ω )]+ 1 c log ( p t exp ( -c [ Q 1 -ˆ Q 1 ( ω )] + (1 -p t ) exp ( -c [ Q 2 -ˆ Q 2 ( ω )] ))] . Using the tabular parameterization for the actor, the policy update at iteration t is given by: p t +1 = p t exp( η ˆ Q 1 ) p t exp( η ˆ Q 1 )+(1 -p t ) exp( η ˆ Q 2 ) , where η 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 L t ( ω t ) = 0 for any iteration t .

Proof. Define A := exp( -c [ Q 1 -ˆ Q 1 ( ω )]) and B := exp( -c [ Q 2 -ˆ Q 2 ( ω )]) . Calculating the gradient of L t w.r.t ω and setting it to zero,

Observe that Q 1 -ˆ Q 1 ( ω t ) = Q 2 -ˆ Q 2 ( ω t ) and thus L t ( ω t ) = 0 for all t . Writing the actor update,

.

Proposition 17 (Detailed version of Prop. 7) . Consider a two-armed bandit example and define p ∈ [0 , 1] as the probability of pulling arm 1. Given p , let the advantage of arm 1 be equal to A 1 := 1 2 > 0 , while that of arm 2 is A 2 := -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: H 0 : ˆ A 1 = 1 2 + ε , ˆ A 2 = -p 1 -p ( 1 2 + ε ) and H 1 : ˆ A 1 = 1 2 -ε sgn ( 1 2 -p ) , ˆ A 2 = -p 1 -p ( 1 2 -ε sgn ( 1 2 -p )) where sgn is the signum function and ε ∈ ( 1 2 , 1 ) . If p t is the probability of pulling arm 1 at iteration t , consider minimizing two alternative loss functions to choose the hypothesis H t : (1) Squared (MSE) loss: H t = arg min {H 0 , H 1 } { p t 2 [ A 1 -ˆ A 1 ] 2 + 1 -p t 2 [ A 2 -ˆ A 2 ] 2 } . (2) Decision-aware critic loss ( DA ) with c = 1 : H t = arg min {H 0 , H 1 } { p t (1 -[ A 1 -ˆ A 1 ]) log(1 -[ A 1 -ˆ A 1 ]) + (1 -p t ) (1 -[ A 2 -ˆ A 2 ]) log(1 -[ A 2 -ˆ A 2 ]) } . Using the tabular parameterization for the actor, the policy update at iteration t is given by: p t +1 = p t (1+ η ˆ A 1 ) p t (1+ η ˆ A 1 )+(1 -p t ) (1+ η ˆ A 2 ) For p 0 ≤ 1 2 , the squared loss cannot distinguish between H 0 and 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 > 1 2 , H 0 and H 1 are identical, ensure that ˆ A 1 > ˆ 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 ≤ 1 2 and we focus on this case. Let us calculate the MSE and decision-aware (DA) losses for H 0 .

Similarly, we can calculate the MSE and decision-aware losses for H 1 .

For both H 0 and H 1 , the MSE loss is equal to pε 2 + ε 2 p 2 1 -p and hence it cannot distinguish between the two hypotheses. Writing the actor update,

Hence, in order to ensure that p t +1 > p t and eventual convergence to the optimal arm, we want that ˆ A 2 < ˆ A 1 . For ε ∈ ( 1 2 , 1 ) , for H 0 , ˆ A 1 > 0 while ˆ A 2 < 0 . On the other hand, for H 1 , ˆ A 1 < 0 and ˆ A 2 > 0 . This implies that the algorithm should choose 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 H 0 is smaller than that for H 1 , and hence the algorithm chooses the correct hypothesis and pulls the optimal arm. For this, we define f ( p ) as follows,

For f ( p ) to be well-defined, we want that, 1 -ϵ > 0 = ⇒ ϵ < 1 and 1 -ϵ p 1 -p > 0 = ⇒ p < 1 1+ ϵ . Since ϵ ∈ ( 1 / 2 , 1) , p < 1 2 . In order to prove that the algorithm will always choose H 0 , we will show that f ( p ) ≤ 0 for all p ∈ [0 , 1 / 2 ] next. First note that,

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 ∈ (0 , 1 / 2 ) . For this, we write f ( p ) = g ( p ) + h 1 ( p ) -h 2 ( p ) where,

Differentiating the above terms,

Combining the above terms, f ′′ ( p ) = g ′′ ( p ) + h ′′ 1 ( p ) -h ′′ 2 ( p ) > 0 for all p ∈ (0 , 1 / 2 ) and hence f ( p ) is convex. Hence, for all p < 1 2 , minimizing the divergence loss results in choosing H 0 and the actor pulling the optimal arm. Once the probability of

Similarly,

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 ∈ R A , the negative entropy mirror map ϕ ( p ) = ∑ i p i log( p i ) , δ ∈ R A , c > 0 ,

Proof. In this case, [ ∇ ϕ ( p )] i = 1 + log( p i ) . Hence, we need to compute D ϕ ∗ ( z ′ , z ) where z ′ i := 1 + log( p i ) -cδ i and z i := 1+log( p i ) . If ϕ ( p ) = ∑ i p i log( p i ) , using Lemma 9, ϕ ∗ ( z ) = log ( ∑ i exp( z i )) where z i -log( ∑ i exp( z i )) = log( p i ) . Define distribution q such that q i := exp(1+log( p i ) -cδ i ) ∑ j exp(1+log( p j ) -cδ j ) . Using Lemma 11,

Simplifying q ,

Lemma 8. For z ∈ R A , the log-sum-exp mirror map ϕ ( z ) = log( ∑ i exp( z i )) , δ ∈ R A s.t. ∑ i δ i = 0 , c > 0 ,

(

)

∑

i

(

p

i

-

p

i

cδ

i

)

where p i = exp( z i ) ∑ j exp( z j ) .

Proof. In this case, [ ∇ ϕ ( z )] i = exp( z i ) ∑ j exp( z j ) = p i . Define distribution q s.t. q i := p i -cδ i . Note that since ∑ i δ i = 0 , ∑ i q i = ∑ i p i = 1 and hence, q is a valid distribution. We thus need to compute D ϕ ∗ ( q, p ) . Using Lemma 9, ϕ ∗ ( p ) = ∑ i p i log( p i ) where p i = exp( z i ) ∑ j exp( z j ) . Using Lemma 10,

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 ∈ R d , if ϕ ( z ) := log ( ∑ i exp( z i )) , then ϕ ∗ ( p ) = ∑ i p i log( p i ) where p i = exp( z i ) ∑ j exp( z j ) . Similarly, if ϕ ( p ) = ∑ i p i log( p i ) , then ϕ ∗ ( z ) = log ( ∑ i exp( z i )) where z i -log( ∑ i exp( z i )) = log( p i ) .

Proof. If ϕ ( z ) := log ( ∑ i exp( z i )) ,

D

ϕ

∗

∇

ϕ

(

z

)

-

c δ,

∇

ϕ

(

z

)

=

(

p

i

-

cδ

i

) log

,

Setting the gradient to zero, we get that p i = exp( z ∗ i ) ∑ j exp( z ∗ j ) for z ∗ ∈ Z ∗ where Z ∗ is the set of maxima related by a shift (i.e. if z ∗ ∈ Z ∗ , z ∗ + C ∈ Z ∗ for a constant C ). Using the optimality condition, we know that ∑ i p i = 1 and

Using this relation,

The second statement follows since the ϕ ∗ ( ϕ ∗ ) = ϕ .

Lemma 10. For probability distributions, p and p ′ , if ϕ ( p ) = ∑ i p log( p i ) , then D ϕ ( p, p ′ ) = KL ( p || p ′ ) .

Proof. Note that [ ∇ ϕ ( p )] i = 1 + log( p i ) . Using the definition of the Bregman divergence,

Since p and p ′ are valid probability distributions, ∑ i p i = ∑ i p ′ i = 1 , and hence, D ϕ ( p, p ′ ) = KL ( p || p ′ ) .

Lemma 11. If ϕ ( z ) = log( ∑ i exp( z i )) , then D ϕ ( z, z ′ ) = KL ( p ′ || p ) , where p i := exp( z i ) ∑ j exp( z j ) and p ′ i := exp( z ′ i ) ∑ j exp( z ′ j ) .

Proof. Note that [ ∇ ϕ ( z )] i = exp( z i ) ∑ j exp( z j ) = p i where p i := exp( z i ) ∑ j exp( z j ) . Using the definition of the Bregman divergence,

F Implementation Details

F.1 Heuristic to estimate c

We estimate c to maximize the lower-bound on J ( π ) . In particular, using Prop. 1,

For a fixed ˆ g ( π t ) , we need to maximize the RHS w.r.t π and c , i.e.

Instead of maximizing w.r.t π and c , we will next aim to find an upper-bound on the RHS that is independent of π and aim to maximize it w.r.t c . Using Lemma 1 with y ′ = π , y = π t , x = -ˆ g ( π t ) and define c ′ such that 1 c ′ = 1 η + 1 c .

Using the above upper-bound in Eq. (6),

This implies that the estimate ˆ c can be calculated as:

In order to gain some intuition, let us consider the case where D Φ ( u, v ) = 1 2 ∥ u -v ∥ 2 2 . In this case,

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 γ = 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 γ = 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 X ( s, a ) and linear function approximation to estimate action-value function Q , implying that ˆ Q ( s, a ) = ω T X ( s, a ) where ω, X ( s, a ) ∈ R d .

Baselines: For policy π , the MSE objective is to return the ω that minimizes the squared norm error of the action-value function Q π across all state-actions weighted by the state-action occupancy measure µ π ( s, a ) .

Taking the derivative with respect to ω and setting it to zero:

Given features X , the true action-value function Q π and state-action occupancy measure µ π , we can compute K , y and solve ω MSE = K -1 y .

Similarly for policy π , the advantage-MSE objective is to return ω that minimizes the squared error of the advantage function A π across all state-actions weighted by the state-action occupancy measure µ π .

Taking the derivative with respect to ω and setting it to zero:

Given features X , the true advantage function A π and state-action occupancy measure µ π , we ca compute K and y and solve w Adv-MSE = K -1 y .

Decision-aware critic in direct representation: Recall that for policy π , the decision-aware critic loss in direct representation is the blue term in Prop. 4, which after linear parameterization on ˆ Q π would be as follows:

The above term is a convex function of ω for any c > 0 . We minimize the term using gradient descent, where the gradient with respect to ω is:

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 β = 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 π , the decision-aware critic loss in softmax representation is the blue term in Prop. 6, which after linear parameterization on ˆ Q π and substituting ˆ A π ( s, a ) with ω T ( X ( s, a ) -∑ a ′ X ( s, a ′ )) would be as follows:

Similarly, the above term is convex with respect to ω 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 ω :

F.4 Actor optimization

Direct representation : For all actor-critic algorithms, we maximize the green term in Prop. 4 known as MDPO [56].

In tabular parameterization of the actor, θ s,a = p π ( 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 π t , the estimated action-value function from the critic ˆ Q π t and η as the functional step-size, the update at iteration t is:

When we linearly parameterize the policy, implying that for policy π , p π ( a | s ) = exp ( θ T X ( s,a )) ∑ a ′ exp ( θ T X ( s,a ′ )) where θ, X ( s, a ) ∈ 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:

And the gradient of objective with respect to θ is:

Softmax representation : For all actor-critic algorithms, we maximize the green term in Prop. 6 known as sMDPO [57].

In tabular parameterization of the actor, θ s,a = p π ( s, a ) , at iteration t given the policy π t , the estimated advantage function from the critic ˆ A π t ( s, a ) = ˆ Q π t ( s, a ) -∑ a ′ p π t ( a | s ) ˆ Q π t ( s, a ′ ) , and functional step-size η , the actor update can be solved in closed-form and is as follows:

We refer the reader to Appendix F.1 of [57] for explicit derivation. When we linearly parameterize the policy, implying that for policy π , p π ( a | s ) = exp ( θ T X ( s,a )) ∑ a ′ exp ( θ T X ( s,a ′ )) , we need to maximize the following with respect to θ :

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 θ is:

F.5 Parameter Tuning

Table 2: Parameters for the Frozen Lake environment

G Additional Experiments

d=100

d=80

09=p

d=40

0.6-

Direct

0.4-

(n=0.01,c=0.01)

(u)r

(n=0.01, c =0.1)

0.2-

(n=0.01, c=1.0)

0.0-

0

10000

20000

30000

40000

0

10000

20000

30000

40000

0

10000

20000

30000

40000

10000

20000

30000

40000

0.6-

Softmax

0.4

(n=0.01,c=0.001)

(u)f

(n=0.01,c=0.01)

(n=0.01, c=0.1)

0.2-

(n=0.1, c=0.01)

0.0-

0

10000

20000

30000

40000

0

10000

20000

30000

40000

0

10000

20000

30000

40000

0

10000

20000

30000

40000

#Iteration

#Iteration

#Iteration

#Iteration

(a) Linear policy parameterization

....MSE-eta=0.005

.-MSE-eta=0.1

....

AdvMSE-eta=0.005

AdvMSE-eta=0.1

......

Ours-eta=0.005

--- MSE-eta= 0.01

MSE-eta= 1.0

AdvMSE-eta=0.01

AdvMSE-eta=1.0

Ours-eta=0.01

Ours-eta= 1.0

d=100

d=80

d=60

d=40

0.6-

0.4-

Direct

0.2-

0.0-

0

10000

20000

30000

400000

100002000030000

400000

10000

20000

30000

400000

10000

20000

30000

40000

0.6-

Softmax

0.4-

(E)r

0.2-

0.0-

0

10000

20000

30000

400000

10000

20000

30000

400000

10000

20000

30000

400000

10000

20000

30000

40000

#Iteration

#Iteration

#Iteration

#Iteration

MSE-eta= 0.001

MSE-eta= 0.1

AdvMSE-eta=0.001

AdvMSE-eta=0.1

Ours-eta=0.001

MSE-eta= 0.01

MSE-eta= 1.0

AdvMSE-eta=0.01

AdvMSE-eta= 1.0

Ours-eta= 0.01

Ours-eta= 1.0

d=100

d=80

09=p

d=40

-90

0.4

Direct

0.2-

0.0-

25000

50000

75000

0

25000

50000

75000

0

25000

50000

75000

0

25000

50000

75000

-90

0.4

Softmax

()r

0.2-

0.0-

0

25000

50000

75000

0

25000

50000

75000

0

25000

50000

75000

0

25000

50000

75000

#Iteration

#Iteration

#Iteration

#Iteration

(b) Tabular policy parameterization

(a) Linear policy parameterization

-MSE-eta=0.01

MSE-eta=0.1

---AdvMSE-eta=0.01

AdvMSE-eta=0.1

Ours-eta=1.0

d=80

09=p

d=40

0.6

0.5

0.4

Direct

(E)r

0.3

0.2

0.1-

-0:0

-0.1-

0

10000

20000

30000

0

10000

20000

30000

0

10000

20000

30000

0.6-

0.5

0.4

Softmax

-'O

0.2-

0.1-

0.0-

0.1

0

10000

20000

30000

0

10000

20000

30000

0

10000

20000

30000

#Iteration

#Iteration

#Iteration

MSE-eta=0.01

MSE-eta=0.1

---AdvMSE-eta=0.01

AdvMSE-eta=0.1

一

Ours-eta=1.0

d=80

d=60

d=40

-9:0

0.4

Direct

0.2

0.0

0

10000

20000

30000

0

10000

20000

30000

0

10000

20000

30000

0.6

Softmax

0.4-

(E)r

0.2-

0.0-

0

10000

20000

30000

0

10000

20000

30000

0

10000

20000

30000

#Iteration

#lteration

#Iteration

(b) Tabular policy parameterization

(a) Linear policy parameterization

--AdvMSE-eta=0.1

AdvMSE-eta=1.0

MSE-eta= 0.1

MSE-eta= 1.0

-Ours-eta= 1.0

d=100

09=p

d=40

0.07-

0.05-

Direct

(=)r

0.03-

0.01-

0.00-

o

10000

30000

50000

0

10000

30000

50000

10000

30000

500000

10000

30000

50000

0.07-

Softmax

0.05-

J(π)

0.03-

0.01-

0.00-

10000

30000

500000

10000

30000

500000

10000

30000

500000

10000

30000

50000

#Iteration

#Iteration

#Iteration

#Iteration

-- AdvMSE-eta= 0.1

AdvMSE-eta=1.0

--- MSE-eta= 0.1

MSE-eta= 1.0

---- Ours-eta= 0.1

Ours-eta= 1.0

d=100

09=p

d=40

0.05-

Direct

(u)r

0.03

0.01-

0.00-

10000

20000

30000

10000

20000

30000

10000

20000

30000

10000

20000

30000

0.07-

Softmax

0.05

(三)r

-00

0.01-

0.00-

10000

20000

30000

10000

20000

30000

10000

20000

30000

10000

20000

30000

#Iteration

#Iteration

#Iteration

#Iteration

(b) Tabular policy parameterization

Frozen Lake - Linear/Tabular actor and Linear critic with estimated Q : For the Frozen Lake environment, when estimating the Q functions using Monte Carlo sampling (all other choices being the same as in Fig. 5), we found that the variance resulting from Monte Carlo sampling (even with ≥ 1000 samples) dominates the bias. As a result, the effect of the critic loss is minimal, and all algorithms result in similar performance.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors present a decision aware actor-critic algorithm and then try to analyse the same. Subsequently they also provide some numerical results.\n\nSTRENGTHS:\nThe authors try and provide an analysis of an actor-critic algorithm. However, there are a number of questions that are unclear to me that I write in detail below.\n\nWEAKNESSES:\nI am writing down the entire list of comments on the paper below. One can call these weaknesses as these points have not been explained well or there are issues in the explanations provided.\n\n1. The authors say in the introduction that the objective of the critic is to minimize the value estimation error across all states and actions. This is incorrect since a critic's job is to estimate the expected TD error for any given policy provided by the actor and the actor's job is to find the policy that minimizes this expected TD error over policies. The authors seem to mention this at multiple places.\n\n2. Weird notation: J_s(\\pi) is well defined but what is J_s(\\rho)? Note that \\rho is the initial state distribution and then you say J(\\pi) = E[J_s(\\rho)]! \n\n3. Note that d^\\pi and \\mu^\\pi are not valid probability distributions. How do you then value sample states or state-action tuples from these?\n\n4. What does it mean to say that since p^\\pi(s,a) is a probability, you can write it in the Boltzmann form as you have written. Is it always possible to express any probability distribution as some Boltzmann distribution? Later you introduce a parameterization \\theta to this distributional form. Why not directly give a parameterized Boltzmann probability if that is what you want for the ramdomized policy instead of doing this in such a round about manner.\n\n5. In the definition of \\pi_{t+1} in Section 3.1, obviously \\eta plays a major role. What values can it take?\n\n6. For the blue term on page 4, you mention that as c decreases, critic error decreases. I don't see how? As c decreases, the first term inside D_{\\Phi^*}(.,.) will decrease as c decreases but that term is multiplied by 1/c which will blow up as c decreases. \n\n7. Algorithm 1: How do you compute $\\nabla_v L_t (v_k) and \\hat{g}_t?\n\n8. Step-sizes and stability of the algorithm: You use \\alpha_c and \\alpha_a as the step-sizes for the critic and actor respectively. I don't see any conditions written on these step-sizes. The reason is that actor-critic algorithms are meant to track policy iteration. This requires that the critic recursion moves on a faster timescale as compared to the actor recursion. When diminishing step-sizes are used, one requires that $\\alpha_a go to zero at a rate faster than \\alpha_c in order to get this effect of asymptotic convergence. Furthermore, the authors do not mention anything about the stability of their procedure. How does one ensure that this algorithm is stable - a precondition to ensuring convergence.\n\n9. Since you do not assume anything about mixing times of the underlying Markov process under any policy, how do you ensure that Markov noise does not create any problems in convergence of the scheme, see for instance,Sajad Khodadadian, Zaiwei Chen, Siva Theja Maguluri, \"Finite-Sample Analysis of Off-Policy Natural Actor-Critic Algorithm,\" ICML 2021.\n\n10. In Proposition 2, what are the terms \\tilde{H}_t^\\dagger and [\\hat{g}_t]_{s,a}?\n\n11. How does one interpret Proposition 3? Not clear at what rate the second term decreases and whether it even does so? In the absence of a proof of decrease and the rate at which it happens, the result is meaningless. \n\n12. In Section 4, you give a monotonic policy improvement result. TRPO also has a similar result. Which one is better - the improvement provided by your algorithm or TRPO?\n\n13. Mistakes: Section 5.1 - you say \\nabla_\\pi J(\\pi) = d^\\pi(s) Q^\\pi(s,a). What is s and what is (s,a) on the RHS since the LHS does not have any s or (s,a) tuple? Also, in Section 5.2, you say \\nabla_\\pi(J(\\pi)) = d^\\pi(s) A^\\pi(s,a) p^\\pi(a|s)? What are s, a on RHS. There is clearly a mistake since LHS does not depend on these quantities. Also, how do you reconcile these two definitions for the same object?\n\n14. The bandit proposition 5 and also 7 suddenly come in and appear crude with such assumptions as deterministic rewards, deterministic Q-value updates etc. It is not clear why these results have been provided.\n\nQUESTIONS:\nI have given my questions in detail above.\n\nLIMITATIONS:\nI think the analysis of convergence is flawed since there are no assumptions or results  regarding (a) Lipschitz continuity of the objectives, (b) step-sizes used in the actor-critic scheme, (c) stability of the two coupled recursions, (d) fast mixing nature of the Markov noise, etc., have been made or shown. Moreover, from my comments mentioned above, the paper lacks on several fronts and will require significant revision and re-review in order to be publishable.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-25"}

comments: SUMMARY:
The paper addresses the issue of objective mismatch in Actor-Critic methods by designing a joint objective that enables training the actor and critic in a decision-aware manner. The proposed algorithm ensures monotonic policy improvement, irrespective of the chosen policy and critic parameterization.

STRENGTHS:
The paper presents an algorithm that guarantees monotonic policy improvement and has solid theoretical foundations.

WEAKNESSES:
The experimentation in the paper is somewhat limited, which may raise concerns about the algorithm's performance in more complex environments.

QUESTIONS:
In complex real-world environments, $\Delta_\pi J(\pi)$could only be estimated, e.g. using MC sampling. And for large MDPs, function approximation error of Q is unavoidable for large MDPs. The article provides experiments on two simple environments, Cliff World and Frozen Lake, using linear/tabular action and Linear cirtic. Thus, my questions are:

(1) Can the algorithm be applied to more challenging environments, such as the standard benchmark tasks in MuJoCo? If so, how well does it perform in those environments? If there are limitations or challenges in applying the algorithm to more difficult tasks, it would be helpful to elaborate on those limitations.

(2) Does the algorithm introduce significant extra time overhead compared to other methods? It would be useful to report the additional time required by the algorithm on the Cliff World and Frozen Lake environments.

LIMITATIONS:
Potential Challenges in Complex Environments: The applicability of the proposed algorithm to more complex environments, such as those found in the MuJoCo benchmark tasks, remains uncertain.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors develop a generic decision-award AC algorithm where both the actor and the critic take steps iteratively to optimize some “policy improvement lower bound” under the FMAPG framework. In essence, each step takes the gradient estimation error as the critic error, and characterizes the policy improvement in terms of the estimated gradient $\\hat g$. This result seems to come from the Taylor expansions of the Bregman divergence. The authors argue that since both the actor and the critic act in a cooperative manner, the critic therefore only focuses on parts of the state-action pairs that have the largest impact on the actor’s performance improvement. The authors provide the conditions for monotonic policy improvement and show that the algorithms converge to some stable points. Examples and experiments show that the proposed algorithm outperforms TD/AdvTD methods on certain tasks.\n\nSTRENGTHS:\nThe algorithm is interesting and well-motivated. Although the idea is not new to the RL community, the theoretical results in this work are sound and in general support the authors’ claims. Moreover, the paper is well written and I find most of the statements clearly presented.\n\nWEAKNESSES:\nIn order to perform GD/SGD, you at least need to rollout trajectories collected under $d^{\\pi_t}$. I didn’t find evidence supporting the sample efficiency of the proposed algorithm as you have claimed. Questions include:\n1. How many samples are used for each GD step in your experiment? \n2. How large is the variance of the proposed algorithm compared to TD/AdvTD when the sampling budget is limited, especially when the samples do not have sufficient coverage? \n4. How do the relative scales of $m_c$ and $m_a$ influence your results?\n\nQUESTIONS:\n1. In practice, buffers containing trajectories collected under history policies are commonly used. Is there any evidence showing that your algorithm can handle the distribution shift issue?\n\n2. Can you perform gradient steps for the actor/critic simultaneously? If so, what are the relative step sizes for the inner and the outer loops?\n\nLIMITATIONS:\nNo concerns here.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

comments: SUMMARY:
This research addresses the mismatched objectives in actor-critic (AC) methods used in reinforcement learning (RL). By introducing a joint objective for training the actor and critic in a decision-aware fashion, a generic AC algorithm is developed. The algorithm ensures monotonic policy improvement regardless of the policy and critic parameterization. The proposed approach offers advantages over traditional methods, as demonstrated through rigorous analysis and empirical evaluations.

STRENGTHS:
I think this work really pushes the RL community research efforts further by answering:

> can we design a generic actor critic algorithm with joint objective?

The main contribution of generic actor critic algorithm with joint objective is a really nice idea worthy for publication at NeurIPS.

WEAKNESSES:
I have only have one minor weakness for this work as follows:

> Using linear representations as "general function approximation" is a bit weak. I presume this is the reason why only simple RL problems have been demonstrated in this work.

I am open to discussions with the authors and reviewers to increase my score. All the best for future decisions!

QUESTIONS:
na

LIMITATIONS:
na
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-05