Original: G LOBAL OPTIMALITY OF SOFTMAX POLICY GRADIENT WITH SINGLE HIDDEN LAYER NEURAL NETWORKS IN THE MEAN - FIELD REGIME
Revised: G LOBAL OPTIMALITY OF SOFTMAX POLICY GRADIENT WITH SINGLE HIDDEN LAYER NEURAL NETWORKS IN THE MEAN - FIELD REGIME

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: We study the problem of policy optimization for inﬁnite-horizon discounted Markov Decision
Revised: We study the problem of policy optimization for inﬁnite-horizon discounted Markov Decision

Original: Processes with softmax policy and nonlinear function approximation trained with policy gradient algorithms.
Revised: Processes with softmax policy and nonlinear function approximation trained with policy gradient algorithms.

Original: We concentrate on the training dynamics in the mean-ﬁeld regime, modeling e.g. , the behavior of wide single hidden layer neural networks, when exploration is encouraged through entropy regularization.
Revised: We concentrate on the training dynamics in the mean-ﬁeld regime, modeling e g. , the behavior of wide single hidden layer neural networks, when exploration is encouraged through entropy regularization.

Original: The dynamics of these models is established as a Wasserstein gradient ﬂow of distributions in parameter space.
Revised: The dynamics of these models is established as a Wasserstein gradient ﬂow of distributions in parameter space.

Original: We further prove global optimality of the ﬁxed points of this dynamics under mild conditions on their initialization
Revised: We further prove global optimality of the ﬁxed points of this dynamics under mild conditions on their initialization

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: In recent years, deep reinforcement learning has revolutionized the world of Artiﬁcial Intelligence by outperforming humans in a multitude of highly complex tasks and achieving breakthroughs that were deemed unthinkable at least for the next decade.
Revised: In recent years, deep reinforcement learning has revolutionized the world of Artiﬁcial Intelligence by outperforming humans in a multitude of highly complex tasks and achieving breakthroughs that were deemed unthinkable at least for the next decade.

Original: Spectacular examples of such revolutionary potential have appeared over the last few years, with reinforcement learning algorithms mastering games and tasks of increasing complexity, from learning to walk to the games of Go and Starcraft (Mnih et al., 2013; 2015; Silver et al., 2016; 2017; 2018;
Revised: Spectacular examples of such revolutionary potential have appeared over the last few years, with reinforcement learning algorithms mastering games and tasks of increasing complexity, from learning to walk to the games of Go and Starcraft (Mnih et al., 2013; 2015; Silver et al., 2016; 2017; 2018;

Original: Haarnoja et al., 2018a; Vinyals et al., 2019).
Revised: Haarnoja et al., 2018a; Vinyals et al., 2019).

Original: In most cases, the main workhorse allowing artiﬁcial intelligence to pass such unprecedented milestones was a variation of a fundamental method to train reinforcement learning models: policy gradient (PG) algorithms (Sutton et al., 2000).
Revised: In most cases, the main workhorse allowing artiﬁcial intelligence to pass such unprecedented milestones was a variation of a fundamental method to train reinforcement learning models: policy gradient (PG) algorithms (Sutton et al., 2000).

Original: This algorithm has a disarmingly simple approach to the optimization problem at hand: given a parametrization of the policy, it updates the parameters in the direction of steepest ascent of the associated integrated value function.
Revised: This algorithm has a disarmingly simple approach to the optimization problem at hand: given a parametrization of the policy, it updates the parameters in the direction of steepest ascent of the associated integrated value function.

Original: Impressive progress has been made recently in the understanding of the convergence and optimization properties of this class of algorithms in the tabular setting (Agarwal et al., 2019; Cen et al., 2020; Bhandari & Russo, 2019), in particular leveraging the natural tradeoff between exploration and exploitation offered for entropy-regularized rewards by softmax policies (Haarnoja et al., 2018b; Mei et al., 2020).
Revised: Impressive progress has been made recently in the understanding of the convergence and optimization properties of this class of algorithms in the tabular setting (Agarwal et al., 2019; Cen et al., 2020; Bhandari & Russo, 2019), in particular leveraging the natural tradeoff between exploration and exploitation offered for entropy-regularized rewards by softmax policies (Haarnoja et al., 2018b; Mei et al., 2020).

Original: However, this simple algorithm alone is not sufﬁcient to explain the multitude of recent breakthroughs in this ﬁeld: in application domains such as Starcraft, robotics or movement planning, the space of possible states and actions are exceedingly large – or even continuous – and can therefore not be represented efﬁciently by tabular policies (Haarnoja et al., 2018a).
Revised: However, this simple algorithm alone is not sufﬁcient to explain the multitude of recent breakthroughs in this ﬁeld: in application domains such as Starcraft, robotics or movement planning, the space of possible states and actions are exceedingly large – or even continuous – and can therefore not be represented efﬁciently by tabular policies (Haarnoja et al., 2018a).

Original: Consequently, the recent impressive successes of artiﬁcial intelligence would be impossible without the natural choice of neural networks to approximate value functions and / or policy functions in reinforcement learning algorithms (Mnih et al., 2015; Sutton et al., 2000).
Revised: Consequently, the recent impressive successes of artiﬁcial intelligence would be impossible without the natural choice of neural networks to approximate value functions and / or policy functions in reinforcement learning algorithms (Mnih et al., 2015; Sutton et al., 2000).

Original: While neural networks, in particular deep neural networks, provide a powerful and versatile tool to approximate high dimensional functions on continuous spaces (Cybenko, 1989; Hornik, 1991; Barron, 1993), their intrinsic nonlinearity poses signiﬁcant obstacles to the theoretical understanding of their training and optimization properties.
Revised: While neural networks, in particular deep neural networks, provide a powerful and versatile tool to approximate high dimensional functions on continuous spaces (Cybenko, 1989; Hornik, 1991; Barron, 1993), their intrinsic nonlinearity poses signiﬁcant obstacles to the theoretical understanding of their training and optimization properties.

Original: For instance, it is known that the optimization landscape of these models is highly nonconvex, preventing the use of most theoretical tools from classical optimization theory.
Revised: For instance, it is known that the optimization landscape of these models is highly nonconvex, preventing the use of most theoretical tools from classical optimization theory.

Original: For this reason, the unprecedented success of neural networks in artiﬁcial intelligence stands in contrast with the poor understanding of these methods from a theoretical perspective.
Revised: For this reason, the unprecedented success of neural networks in artiﬁcial intelligence stands in contrast with the poor understanding of these methods from a theoretical perspective.

Original: Indeed, even in the supervised setting, which can be viewed as a special case of reinforcement learning, deep neural networks are still far from being understood despite having been an important and fashionable research focus in recent years.
Revised: Indeed, even in the supervised setting, which can be viewed as a special case of reinforcement learning, deep neural networks are still far from being understood despite having been an important and fashionable research focus in recent years.

Original: Only recently, a theory of neural network learning has started to emerge, including recent works on mean-ﬁeld point of view of training dynamics (Mei et al.,2018; Rotskoff & Vanden-Eijnden, 2018;
Revised: Only recently, a theory of neural network learning has started to emerge, including recent works on mean-ﬁeld point of view of training dynamics (Mei et al.,

Original:  Rotskoff et al., 2019; Wei et al., 2018; Chizat & Bach, 2018) and on linearized dynamics in the over-parametrized regime (Jacot et al., 2018; Allen-Zhu et al., 2018; Du et al., 2018; 2019; Zou et al., 2018; Allen-Zhu et al., 2019; Chizat et al., 2019; Oymak & Soltanolkotabi, 2020; Ghorbani et al., 2019; Lee et al., 2019).
Revised: ; Rotskoff & Vanden-Eijnden, 2018; Rotskoff et al., 2019; Wei et al., 2018; Chizat & Bach, 2018) and on linearized dynamics in the over-parametrized regime (Jacot et al., 2018; Allen-Zhu et al., 2018; Du et al., 2018; 2019; Zou et al., 2018; Allen-Zhu et al., 2019; Chizat et al., 2019; Oymak & Soltanolkotabi, 2020; Ghorbani et al., 2019; Lee et al., 2019).

Original: More speciﬁcally to the context of reinforcement learning, some works focusing on value-based learning (Agazzi & Lu, 2019; Cai et al., 2019; Zhang et al., 2020), and others exploring the dynamics of policy gradient algorithms (Zhang et al., 2019) have recently appeared.
Revised: More speciﬁcally to the context of reinforcement learning, some works focusing on value-based learning (Agazzi & Lu, 2019; Cai et al., 2019; Zhang et al., 2020), and others exploring the dynamics of policy gradient algorithms (Zhang et al., 2019) have recently appeared.

Original: Despite this progress, the theoretical understanding of deep reinforcement learning still poses a signiﬁcant challenge to the theoretical machine learning community, and it is of crucial importance to understand the convergence and optimization properties of such algorithms to bridge the gap between theory and practice.
Revised: Despite this progress, the theoretical understanding of deep reinforcement learning still poses a signiﬁcant challenge to the theoretical machine learning community, and it is of crucial importance to understand the convergence and optimization properties of such algorithms to bridge the gap between theory and practice.

Original: C ONTRIBUTIONS .
Revised: C ONTRIBUTIONS .

Original: The main goal of this work is to investigate entropy-regularized policy gradient dynamics for wide, single hidden layer neural networks.
Revised: The main goal of this work is to investigate entropy-regularized policy gradient dynamics for wide, single hidden layer neural networks.

Original: In particular, we give the following contributions:
Revised: In particular, we give the following contributions:

Original: • We give a mean-ﬁeld formulation of policy gradient dynamics in parameter space, describing the evolution of neural network parameters in the form of a transport partial differential equation ( PDE ).
Revised: • We give a mean-ﬁeld formulation of policy gradient dynamics in parameter space, describing the evolution of neural network parameters in the form of a transport partial differential equation ( PDE ).

Original: We prove convergence of the particle dynamics to their mean-ﬁeld counterpart.
Revised: We prove convergence of the particle dynamics to their mean-ﬁeld counterpart.

Original: We further explore the structure of this problem by showing that such PDE is a gradient ﬂow in the Wasserstein space for the appropriate energy functional.
Revised: We further explore the structure of this problem by showing that such PDE is a gradient ﬂow in the Wasserstein space for the appropriate energy functional.

Original: • We investigate the convergence properties of the above dynamics in the space of measures.
Revised: • We investigate the convergence properties of the above dynamics in the space of measures.

Original: In particular, we prove that under some mild assumptions on the initialization of the neural network parameters and on the approximating power of the nonlinearity, all ﬁxed points of the dynamics are global optima,
Revised: In particular, we prove that under some mild assumptions on the initialization of the neural network parameters and on the approximating power of the nonlinearity, all ﬁxed points of the dynamics are global optima,

Original: i
Revised: i

Original: e. , the approximate policy learned by the neural network is optimal, R ELATED W ORKS .
Revised: e. , the approximate policy learned by the neural network is optimal, R ELATED W ORKS .

Original: Recent progress in the understanding of the parametric dynamics of simple neural networks trained with gradient descent in the supervised setting has been made in (Mei et al., 2018; Rotskoff & Vanden-Eijnden, 2018;
Revised: Recent progress in the understanding of the parametric dynamics of simple neural networks trained with gradient descent in the supervised setting has been made in (Mei et al., 2018; Rotskoff & Vanden-Eijnden, 2018;

Original: Wei et al., 2018; Chizat, 2019; Chizat & Bach, 2020).
Revised: Wei et al., 2018; Chizat, 2019; Chizat & Bach, 2020).

Original: These results have further been extended to the multilayer setting in (Nguyen & Pham, 2020).
Revised: These results have further been extended to the multilayer setting in (Nguyen & Pham, 2020).

Original: In particular, the paper (Chizat & Bach, 2018) proves optimality of ﬁxed points for wide single layer neural networks leveraging a Wasserstein gradient ﬂow structure and the strong convexity of the loss functional WRT the predictor.
Revised: In particular, the paper (Chizat & Bach, 2018) proves optimality of ﬁxed points for wide single layer neural networks leveraging a Wasserstein gradient ﬂow structure and the strong convexity of the loss functional WRT the predictor.

Original: We extend these results to the reinforcement learning framework, where the convexity that is heavily leveraged in (Chizat & Bach, 2018) is lost.
Revised: We extend these results to the reinforcement learning framework, where the convexity that is heavily leveraged in (Chizat & Bach, 2018) is lost.

Original: We bypass this issue by requiring a sufﬁcient expressivity of the used nonlinear representation, allowing to characterize global minimizer as optimal approximators.
Revised: We bypass this issue by requiring a sufﬁcient expressivity of the used nonlinear representation, allowing to characterize global minimizer as optimal approximators.

Original: The convergence and optimality of policy gradient algorithms (including in the entropy-regularized setting) is investigated in the recent papers (Bhandari & Russo, 2019; Mei et al., 2020; Cen et al., 2020; Agarwal et al., 2019).
Revised: The convergence and optimality of policy gradient algorithms (including in the entropy-regularized setting) is investigated in the recent papers (Bhandari & Russo, 2019; Mei et al., 2020; Cen et al., 2020; Agarwal et al., 2019).

Original: These references establish convergence estimates through gradient domination bounds.
Revised: These references establish convergence estimates through gradient domination bounds.

Original: In (Mei et al., 2020; Cen et al., 2020) such results are limited to the tabular case, while (Agarwal et al., 2019; 2020) also discuss neural softmax policy classes, but under a different algorithmic update and assuming certain well-conditioning assumptions along training.
Revised: In (Mei et al., 2020; Cen et al., 2020) such results are limited to the tabular case, while (Agarwal et al., 2019; 2020) also discuss neural softmax policy classes, but under a different algorithmic update and assuming certain well-conditioning assumptions along training.

Original: Furthermore, all these results heavily leverage the ﬁniteness of action space.
Revised: Furthermore, all these results heavily leverage the ﬁniteness of action space.

Original: In contrast, this paper focuses on the continuous space and action setting with nonlinear function approximation.
Revised: In contrast, this paper focuses on the continuous space and action setting with nonlinear function approximation.

Original: Further recent works discussing convergence properties of reinforcement learning algorithms with function approximation via neural networks include (Zhang et al., 2019;
Revised: Further recent works discussing convergence properties of reinforcement learning algorithms with function approximation via neural networks include (Zhang et al., 2019;

Original: Cai et al., 2019).
Revised: Cai et al., 2019).

Original: These results only hold for ﬁnite action spaces, and are obtained in the regime where the network behaves essentially like a linear model (known as the neural or lazy training regime), in contrast to the results of this paper, which considers training in a nonlinear regime.
Revised: These results only hold for ﬁnite action spaces, and are obtained in the regime where the network behaves essentially like a linear model (known as the neural or lazy training regime), in contrast to the results of this paper, which considers training in a nonlinear regime.

Original: We also note the work (Wang et al., 2019) where the action space is continuous but the training is again in an approximately linear regime.
Revised: We also note the work (Wang et al., 2019) where the action space is continuous but the training is again in an approximately linear regime.

Original:  M ARKOV D ECISION P ROCESSES AND POLICY GRADIENTS
Revised:  M ARKOV D ECISION P ROCESSES AND POLICY GRADIENTS

Original: We denote a Markov Decision Process ( MDP ) by the 5 -tuple ( S , A , P, r, γ ) , where S is the state space, A is the action space, P = P ( s, a, s (cid:48) )
Revised: We denote a Markov Decision Process ( MDP ) by the 5 -tuple ( S , A , P, r, γ ) , where S is the state space, A is the action space, P = P ( s, a, s (cid:48) )

Original: s,s (cid:48) ∈S ,a ∈A a Markov transition kernel, r ( s, a, s (cid:48) )
Revised: s,s (cid:48) ∈S ,a ∈A a Markov transition kernel, r ( s, a, s (cid:48) )

Original: s,s (cid:48) ∈S ,a ∈A is the realvalued, bounded and continuous immediate reward function and γ ∈ (0 , 1) is a discount factor.
Revised: s,s (cid:48) ∈S ,a ∈A is the realvalued, bounded and continuous immediate reward function and γ ∈ (0 , 1) is a discount factor.

Original: We will consider a probabilistic policy , mapping a state to a probability distribution on the action space, so that π : S → M 1+ ( A ) , where M 1+ ( A ) denotes the space of probability measures on A , and denote for any s ∈ S the corresponding density π ( s, · ) :
Revised: We will consider a probabilistic policy , mapping a state to a probability distribution on the action space, so that π : S → M 1+ ( A ) , where M 1+ ( A ) denotes the space of probability measures on A , and denote for any s ∈ S the corresponding density π ( s, · ) :

Original: A → R + .
Revised: A → R + .

Original: The policy deﬁnes a state-to-state transition operator
Revised: The policy deﬁnes a state-to-state transition operator

Original: P π ( s, d s (cid:48) )
Revised: P π ( s, d s (cid:48) )

Original: = (cid:82) A P ( s, a, d s (cid:48) )
Revised: = (cid:82) A P ( s, a, d s (cid:48) )

Original: π ( s, d a ) , and we assume that P π is Lipschitz continuous as an operator M 

Revised: π ( s, d a ) , and we assume that P π is Lipschitz continuous as an operator M

Original: + ( S ) →
Revised: + ( S ) →

Original: M 1+ ( S ) wrt the policy.
Revised: M 1+ ( S ) wrt the policy.

Original: We further encourage exploration by deﬁning some (relative) entropy-regularized rewards [Equation] Williams & 
Revised: We further encourage exploration by deﬁning some (relative) entropy-regularized rewards (Williams & Peng, 1991)

Original: [Equation]
Revised: [Equation]

Original: where D KL denotes the relative entropy, ¯ π is a reference measure and τ indicates the strength of regularization.
Revised: where D KL denotes the relative entropy, ¯ π is a reference measure and τ indicates the strength of regularization.

Original: Throughout, we choose ¯ π to be the Lebesgue measure on A , which we assume, like S , to be a compact subset of the Euclidean space.
Revised: Throughout, we choose ¯ π to be the Lebesgue measure on A , which we assume, like S , to be a compact subset of the Euclidean space.

Original: This regularization encourages exploration and absolute continuity of the policy WRT
Revised: This regularization encourages exploration and absolute continuity of the policy WRT

Original: Lebesgue measure.
Revised: Lebesgue measure.

Original: Consequently, with some abuse of notation, we use throughout the same notation for a distribution and its density in phase space.
Revised: Consequently, with some abuse of notation, we use throughout the same notation for a distribution and its density in phase space.

Original: Note that the original, unregularized MDP can be recovered in the limit τ → 0 .
Revised: Note that the original, unregularized MDP can be recovered in the limit τ → 0 .

Original: In this context, given a policy π the associated value function V π : S → R maps each state to the inﬁnite-horizon expected discounted reward obtained by following the policy π and the Markov process deﬁned by P :
Revised: In this context, given a policy π the associated value function V π : S → R maps each state to the inﬁnite-horizon expected discounted reward obtained by following the policy π and the Markov process deﬁned by P :

Original: [Equation]
Revised: [Equation]

Original: where E π [ · | s 

Revised: where E π [ · | s

Original: = s ] denotes the expectation of the stochastic process s t starting at s 

Revised: = s ] denotes the expectation of the stochastic process s t starting at s

Original: = s and following the (stochastic) dynamics deﬁned recursively by the transition operator P π ( s, d s (cid:48) )
Revised: = s and following the (stochastic) dynamics deﬁned recursively by the transition operator P π ( s, d s (cid:48) )

Original: = (cid:82) P ( s, a, d s (cid:48) )
Revised: = (cid:82) P ( s, a, d s (cid:48) )

Original: π ( s, d a ) .
Revised: π ( s, d a ) .

Original: Correspondingly, we deﬁne the Q -function Q π : S ×
Revised: Correspondingly, we deﬁne the Q -function Q π : S ×

Original: A → R as
Revised: A → R as

Original: [Equation]
Revised: [Equation]

Original: where ¯ r ( s, a ) =
Revised: where ¯ r ( s, a ) =

Original: E [ r ( s, a, s (cid:48) )] is the average reward from ( s, a ) .
Revised: E [ r ( s, a, s (cid:48) )] is the average reward from ( s, a ) .

Original: Conversely, from the deﬁnition, we have the identity for V π ( s ) =
Revised: Conversely, from the deﬁnition, we have the identity for V π ( s ) =

Original: E π
Revised: E π

Original: [ Q π ( s [Equation] [Equation] | s 0 = s ] − τD KL ( π ( s
Revised: [ Q π ( s 0 , a 0 ) | s 0 = s ] − τD KL ( π ( s, · ); ¯ π ( · )) .

Original: [Equation] [Equation] ; ¯ π ( · )) . We are interested in learning the optimal policy π ∗ of a given MDP ( S , A , P, r, γ ) , which satisﬁes for all s ∈ S
Revised: We are interested in learning the optimal policy π ∗ of a given MDP ( S , A , P, r, γ ) , which satisﬁes for all s ∈ S

Original: [Equation]
Revised: [Equation]

Original: More speciﬁcally we would like to estimate this function through a family of approximators π w : S → M 1+ ( A ) parametrized by a vector w ∈ W :
Revised: More speciﬁcally we would like to estimate this function through a family of approximators π w : S → M 1+ ( A ) parametrized by a vector w ∈ W :

Original: = R p .
Revised: = R p .

Original: Note that since we consider entropy-regularized rewards, the optimal policy will be a probabilistic policy (given as a Boltzmann distribution) instead of a deterministic one.
Revised: Note that since we consider entropy-regularized rewards, the optimal policy will be a probabilistic policy (given as a Boltzmann distribution) instead of a deterministic one.

Original: A popular algorithm to solve this problem is given by policy gradient algorithm (Sutton & Barto, 2018).
Revised: A popular algorithm to solve this problem is given by policy gradient algorithm (Sutton & Barto, 2018).

Original: Starting from an initial condition w (0) ∈ W , this algorithm updates the parameters w of the predictor in the direction of steepest ascent of the average reward
Revised: Starting from an initial condition w (0) ∈ W , this algorithm updates the parameters w of the predictor in the direction of steepest ascent of the average reward

Original: [Equation]
Revised: [Equation]

Original: for a ﬁxed absolutely continuous initial distribution of initial states (cid:37) 0 ∈ M 1+ ( S ) and sequence of time steps { β t } t .
Revised: for a ﬁxed absolutely continuous initial distribution of initial states (cid:37) 0 ∈ M 1+ ( S ) and sequence of time steps { β t } t .

Original: Here ˜ E [ · ] denotes an approximation of the expected value operator.
Revised: Here ˜ E [ · ] denotes an approximation of the expected value operator.

Original: This work investigates the regime of asymptotically small constant step-sizes β t → 0 .
Revised: This work investigates the regime of asymptotically small constant step-sizes β t → 0 .

Original: In this adiabatic limit, the stochastic component of the dynamics is averaged before the parameters of the model can change signiﬁcantly.
Revised: In this adiabatic limit, the stochastic component of the dynamics is averaged before the parameters of the model can change signiﬁcantly.

Original: This allows to consider the parametric update as a deterministic dynamical system emerging from the averaging of the underlying stochastic algorithm corresponding to the limit of inﬁnite sample sizes.
Revised: This allows to consider the parametric update as a deterministic dynamical system emerging from the averaging of the underlying stochastic algorithm corresponding to the limit of inﬁnite sample sizes.

Original: This is known as the ODE method (Borkar, 2009) for analyzing stochastic approximation.
Revised: This is known as the ODE method (Borkar, 2009) for analyzing stochastic approximation.

Original: We focus on the analysis of this deterministic system to highlight the core dynamical properties of policy gradients with nonlinear function approximation.
Revised: We focus on the analysis of this deterministic system to highlight the core dynamical properties of policy gradients with nonlinear function approximation.

Original: The averaged, deterministic dynamics is given by the set of ODE s
Revised: The averaged, deterministic dynamics is given by the set of ODE s

Original: [Equation]
Revised: [Equation]

Original: where in the second equality we have applied the policy gradient theorem (Sutton et al., 2000; Sutton & Barto, 2018), deﬁning for a ﬁxed (cid:37) 0 ∈ M 1+ ( S )
Revised: where in the second equality we have applied the policy gradient theorem (Sutton et al., 2000; Sutton & Barto, 2018), deﬁning for a ﬁxed (cid:37) 0 ∈ M 1+ ( S )

Original: [Equation]
Revised: [Equation]

Original: (6) as the (improper) discounted empirical measure.
Revised: ) as the (improper) discounted empirical measure.

Original: For completeness, we include a derivation of (5) in Appendix A.
Revised: For completeness, we include a derivation of (5) in Appendix A.

Original: S OFTMAX POLICIES IN THE MEAN - FIELD REGIME
Revised: S OFTMAX POLICIES IN THE MEAN - FIELD REGIME

Original: We choose to represent our policy as a softmax policy:
Revised: We choose to represent our policy as a softmax policy:

Original: [Equation]
Revised: [Equation]

Original: and parametrize the energy f as a two-layer neural network in the mean-ﬁeld regime, i.e. ,
Revised: and parametrize the energy f as a two-layer neural network in the mean-ﬁeld regime, i.e. ,

Original: [Equation]
Revised: [Equation]

Original: for a ﬁxed, usually nonlinear function ψ : S × A × Ω
Revised: for a ﬁxed, usually nonlinear function ψ : S × A × Ω

Original: → R , where we have separated w ∈ W into N identical components w i ∈ Ω so that W = Ω N .
Revised: → R , where we have separated w ∈ W into N identical components w i ∈ Ω so that W = Ω N .

Original: We can rewrite the above expression in terms of an empirical measure:
Revised: We can rewrite the above expression in terms of an empirical measure:

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: This empirical measure representation removes the symmetry of the approximating functions under permutations of parameters w i .
Revised: This empirical measure representation removes the symmetry of the approximating functions under permutations of parameters w i .

Original: It also facilitates the limit N → ∞ , when ν ( N ) → ν weakly, so that f ν ( N ) → f ν .
Revised: It also facilitates the limit N → ∞ , when ν ( N ) → ν weakly, so that f ν ( N ) → f ν .

Original: Then, for a general distribution ν ∈ M 1+ (Ω) the softmax mean-ﬁeld policy reads:
Revised: Then, for a general distribution ν ∈ M 1+ (Ω) the softmax mean-ﬁeld policy reads:

Original: [Equation]
Revised: [Equation]

Original: Note that by our choice of softmax policy and mean-ﬁeld parametrization (7) we have
Revised: Note that by our choice of softmax policy and mean-ﬁeld parametrization (7) we have

Original: [Equation]
Revised: [Equation]

Revised: Thus the training dynamics (5)

Original: Thus the training dynamics (5), after an appropriate rescaling of time ( t (cid:55)→ t/N , which is due to the mean-ﬁeld parametrization for f w )
Revised: [Equation] after an appropriate rescaling of time ( t (cid:55)→ t/N [Equation] which is due to the mean-ﬁeld parametrization for f w )

Original: [Equation]
Revised: [Equation]

Original: can be rewritten as
Revised: can be rewritten as

Original: [Equation]
Revised: [Equation]

Original: The training dynamics can be more compactly represented by the evolution of the measure ν ∈ M 1+ (Ω) in parameter space, given by a mean-ﬁeld transport partial differential equation of the Vlasov type as
Revised: The training dynamics can be more compactly represented by the evolution of the measure ν ∈ M 1+ (Ω) in parameter space, given by a mean-ﬁeld transport partial differential equation of the Vlasov type as

Original: [Equation]
Revised: [Equation]

Original: where ω ∈ Ω and we have introduced the shorthand C π [ f, g ]( s ) to denote the covariance operator WRT the probability measure
Revised: where ω ∈ Ω and we have introduced the shorthand C π [ f, g ]( s ) to denote the covariance operator WRT the probability measure

Original: π ( s, d a ) .
Revised: π ( s, d a ) .

Original: Note that the above partial differential equation also captures the dynamics of the ﬁnite-width system, i.e. , of the empirical measure ν ( N ) where each w i follows (9).
Revised: Note that the above partial differential equation also captures the dynamics of the ﬁnite-width system, i.e. , of the empirical measure ν ( N ) where each w i follows (9).

Original: We further note that the dynamics introduced above have a gradient ﬂow structure in the probability space
Revised: We further note that the dynamics introduced above have a gradient ﬂow structure in the probability space

Original: M 1 + (Ω) : deﬁning the expected value function
Revised: M 1+ (Ω) : deﬁning the expected value function

Original: [Equation]
Revised: [Equation]

Original: (11) the dynamics (10) is a gradient ﬂow for E in the Wasserstein space (see e.g., (Santambrogio, 2017) for an introduction), as we prove in the appendix:
Revised: 11) the dynamics (10) is a gradient ﬂow for E in the Wasserstein space (see e.g., (Santambrogio, 2017) for an introduction), as we prove in the appendix:

Original: Proposition 2.1.
Revised: Proposition 2.1.

Original: For a ﬁxed initial distribution (cid:37) 0 ∈ M 1+ ( S ) , the dynamics (10) is the Wasserstein gradient ﬂow of the energy functional (11).
Revised: For a ﬁxed initial distribution (cid:37) 0 ∈ M 1+ ( S ) , the dynamics (10) is the Wasserstein gradient ﬂow of the energy functional (11).

Original: Analogous dynamics equation for evolution of parameter space measure in the supervised learning case has been derived in (Mei et al., 2018; Rotskoff & Vanden-Eijnden, 2018; Chizat & Bach, 2018) and in the TD learning case in (Agazzi & Lu, 2019; Zhang et al., 2020).
Revised: Analogous dynamics equation for evolution of parameter space measure in the supervised learning case has been derived in (Mei et al., 2018; Rotskoff & Vanden-Eijnden, 2018; Chizat & Bach, 2018) and in the TD learning case in (Agazzi & Lu, 2019; Zhang et al., 2020).

Original: In particular, in the case of supervised learning, the resulting dynamics is a Wasserstein gradient ﬂow, the structure of which is used to obtain the convergence of the particle system to the mean-ﬁeld dynamics.
Revised: In particular, in the case of supervised learning, the resulting dynamics is a Wasserstein gradient ﬂow, the structure of which is used to obtain the convergence of the particle system to the mean-ﬁeld dynamics.

Original: In our case, however, the energy functional is not convex WRT the policy and moreover the softmax parametrization destroys the convexity of the approximator of the policy with respect to ν t .
Revised: In our case, however, the energy functional is not convex WRT the policy and moreover the softmax parametrization destroys the convexity of the approximator of the policy with respect to ν t .

Original: Thus showing convergence of the dynamics becomes much more challenging.
Revised: Thus showing convergence of the dynamics becomes much more challenging.

Original:  S IMPLIFIED SETTING : THE BANDIT PROBLEM
Revised:  S IMPLIFIED SETTING : THE BANDIT PROBLEM

Original: We now introduce our results in the simple bandit setting, where state space S is one point (and will be henceforth suppressed in the notation) and without loss of generality action space A is continuous.
Revised: We now introduce our results in the simple bandit setting, where state space S is one point (and will be henceforth suppressed in the notation) and without loss of generality action space A is continuous.

Original: In this case, for a reward function r and a softmax policy
Revised: In this case, for a reward function r and a softmax policy

Original: [Equation]
Revised: [Equation]

Original: we have that the value function for the regularized problem reads (we denote V ν
Revised: we have that the value function for the regularized problem reads (we denote V ν

Original: =
Revised: =

Original: V π ν to simplify notation)
Revised: V π ν to simplify notation)

Original: [Equation]
Revised: [Equation]

Original: while the Q function is simply Q ( a ) =
Revised: while the Q function is simply Q ( a ) =

Original: r ( a ) .
Revised: r ( a ) .

Original: We further note that the optimal policy in the regularized case reads:
Revised: We further note that the optimal policy in the regularized case reads:

Original: [Equation]
Revised: [Equation]

Original: Recalling the deﬁnition of the covariance operator C π [ f, g ]( s ) from (10), the expression for the policy gradient vector ﬁeld in the latter case simpliﬁes to
Revised: Recalling the deﬁnition of the covariance operator C π [ f, g ]( s ) from (10), the expression for the policy gradient vector ﬁeld in the latter case simpliﬁes to

Original: [Equation]
Revised: [Equation]

Original: where Dπ ν , DV ν denote the Fr ´ ech ´ et derivative of π ν and V ν WRT ν and π respectively.
Revised: where Dπ ν , DV ν denote the Fr´ech´et derivative of π ν and V ν WRT ν and π respectively.

Original: Note that by the structure of the covariance operator C π , adding a constant to the function f ν ( · ) does not affect the dynamics.
Revised: Note that by the structure of the covariance operator C π , adding a constant to the function f ν ( · ) does not affect the dynamics.

Original: This reﬂects the fact that the softmax policy is normalized by deﬁnition.
Revised: This reﬂects the fact that the softmax policy is normalized by deﬁnition.

Original: 3. G LOBAL OPTIMALITY OF SOFTMAX POLICY GRADIENT
Revised: 3. G LOBAL OPTIMALITY OF SOFTMAX POLICY GRADIENT

Original: We now sketch the main steps in proving that the mean-ﬁeld policy gradient dynamics converge, under appropriate assumptions, to global optimizers.
Revised: We now sketch the main steps in proving that the mean-ﬁeld policy gradient dynamics converge, under appropriate assumptions, to global optimizers.

Original: The proof in this simpler setting is much more transparent than the general case to be discussed in the next section, and will provide some intuition for the latter.
Revised: The proof in this simpler setting is much more transparent than the general case to be discussed in the next section, and will provide some intuition for the latter.

Original: The ﬁrst part of the proof concerns the properties of ﬁxed points of the dynamics (10), while the second part concerns the training dynamics.
Revised: The ﬁrst part of the proof concerns the properties of ﬁxed points of the dynamics (10), while the second part concerns the training dynamics.

Original: S TATICS
Revised: S TATICS

Original: We ﬁrst informally prove global optimality of any ﬁxed point ν ∗ of the transport equation dd t ν t = − div
Revised: We ﬁrst informally prove global optimality of any ﬁxed point ν ∗ of the transport equation dd t ν t = − div

Original: ( ν t F t ) with F t from (12) such that
Revised: ( ν t F t ) with F t from (12) such that

Original: a) ν ∗ has full support in Ω , b) the nonlinearity is 1-homogeneous in the ﬁrst component of its parameters, i.e. , that writing ω= ( ω 0 , ¯ ω ) ∈
Revised: a) ν ∗ has full support in Ω , b) the nonlinearity is 1-homogeneous in the ﬁrst component of its parameters, i.e. , that writing ω

Original:  R × Θ one has ψ ( a ; ω ) :
Revised: = ( ω 0 , ¯ ω ) ∈ R × Θ one has ψ ( a ; ω ) :

Original: = ω 

Revised: = ω

Original: φ ( a ; ¯ ω ) for a regular enough φ : A × Θ → R , c) the span of { φ ( a ; ¯ ω ) } ¯ ω ∈ Θ is dense in L 2 ( A ) , so that i.e. , π ν ∗ ( a ) =
Revised: φ ( a ; ¯ ω ) for a regular enough φ : A × Θ → R , c) the span of { φ ( a ; ¯ ω ) } ¯ ω ∈ Θ is dense in L 2 ( A ) , so that i.e. , π ν ∗ ( a ) =

Original: π ∗ ( a ) =
Revised: π ∗ ( a ) =

Original: Z − 1 exp (cid:2)
Revised: Z − 1 exp (cid:2)

Original: τ − 1 r ( a ) (cid:3) .
Revised: τ − 1 r ( a ) (cid:3) .

Original: Weaker assumptions and the general statement are given in the next section (Proposition 4.3) while the general proof appears in the appendix.
Revised: Weaker assumptions and the general statement are given in the next section, while the general proof appears in the appendix.

Original: First, we note that by assumption a), div ( ν ∗ F ( · ; ν ∗ ))
Revised: First, we note that by assumption a), div ( ν ∗ F ( · ; ν ∗ ))

Original: = 0 directly implies that for  all ω ∈ Ω
Revised: = 0 directly implies that for almost all ω ∈ Ω

Original: [Equation]
Revised: [Equation]

Original: In particular, by homogeneity assumption b), the ﬁrst component of the above vector ﬁeld must vanish on Θ .
Revised: In particular, by homogeneity assumption b), the ﬁrst component of the above vector ﬁeld must vanish on Θ .

Original: [Equation]
Revised: [Equation]

Original: By assumption c) that span of φ is dense in L 2 ( A ) the above implies that
Revised: By assumption c) that span of φ is dense in L 2 ( A ) the above implies that

Original: [Equation]
Revised: [Equation]

Original: Finally, recalling that by the softmax parametrization and by the boundedness of φ , π ν ∗ ( a ) > 0 , we must have
Revised: Finally, recalling that by the softmax parametrization and by the boundedness of φ , π ν ∗ ( a ) > 0 , we must have

Original: [Equation]
Revised: [Equation]

Original: which directly implies the optimality of the policy.
Revised: which directly implies the optimality of the policy.

Original: D YNAMICS
Revised: D YNAMICS

Original: While it is clear that assumption b) and c) about the structure and approximating power of the nonlinearity ψ hold independently of t , we want to show that assumption a) also holds uniformly in time.
Revised: While it is clear that assumption b) and c) about the structure and approximating power of the nonlinearity ψ hold independently of t , we want to show that assumption a) also holds uniformly in time.

Original: In this sense, the continuity of the vector ﬁeld (12) will preserve the full support properties of the measure ν t for all t > 0 , as we will prove in a more general framework in Lemma C.1.
Revised: In this sense, the continuity of the vector ﬁeld (12) will preserve the full support properties of the measure ν t for all t > 0 , as we will prove in a more general framework in Lemma C.2.

Original: Consequently, any measure ν respecting assumption a) at initialization will do so for any ﬁnite positive time t > 0 .
Revised: Consequently, any measure ν respecting assumption a) at initialization will do so for any ﬁnite positive time t > 0 .

Original: However, the question remains of whether this property still holds at t = ∞ .
Revised: However, the question remains of whether this property still holds at t = ∞ .

Original: This is the object of Lemma C.2, where we prove that whenever the gradient approaches a ﬁxed point in parameter space, if this ﬁxed point is not a global minimizer, it must be avoided by the dynamics, and thus, the only possible ﬁxed points of the dynamics are global minimizers.
Revised: This is the object of Lemma C.3, where we prove that whenever the gradient approaches a ﬁxed point in parameter space, if this ﬁxed point is not a global minimizer, it must be avoided by the dynamics, and thus, the only possible ﬁxed points of the dynamics are global minimizers.

Original:  R ESULTS IN THE GENERAL SETTING
Revised:  R ESULTS IN THE GENERAL SETTING

Original: We now come back to the general MDP framework introduced in Section 2.
Revised: We now come back to the general MDP framework introduced in Section 2.

Original: 4. A SSUMPTIONS
Revised: 4. A SSUMPTIONS

Original: To state the main result of this section, the optimality of ﬁxed points of (10) we need the following
Revised: To state the main result of this section, the optimality of ﬁxed points of (10) we need the following

Original: Assumption 1.
Revised: Assumption 1.

Original: Assume that ω
Revised: Assume that ω

Original: = ( ω 0 , ¯ ω ) ∈ R × Θ for Θ =
Revised: = ( ω 0 , ¯ ω ) ∈ R × Θ for Θ =

Original: R m − 1 and ψ ( s, a ; ω ) =ω 0 φ ( s, a ; ¯ ω ) with a Regularity
Revised: R m − 1 and ψ ( s, a ; ω ) =

Original: of φ

Original:  : φ is bounded, differentiable and Dφ ω is Lipschitz.
Revised: ω 0 φ ( s, a ; ¯ ω ) with a) Regularity of φ : φ is bounded, differentiable and Dφ ω is Lipschitz.

Original: Also, for all f ∈ L  S 
Revised: Also, for all f ∈ L 2 ( S × A ) the regular values of the map ¯ ω

Original: ×

Original: A

Original: the

Original: regular

Original: values of the

Original: map ¯ ω (cid:55)→ g f (¯ ω ) :
Revised: (cid:55)→ g f (¯ ω ) :

Original: = (cid:82) f
Revised: = (cid:82) f

Original: ( s, a ) φ ( s, a ; ¯ ω ) are dense in its range, and g f ( r ¯ ω ) converges in C 1 ( { ¯ ω ∈ Θ : (cid:107) ¯ ω (cid:107)2 = 1 } ) as r → ∞ to a map ¯ g f (¯ ω ) whose regular values are dense in its range.
Revised: ( s, a ) φ ( s, a ; ¯ ω ) are dense in its range, and g f ( r ¯ ω ) converges in C 1 ( { ¯ ω ∈ Θ : (cid:107) ¯ ω (cid:107)

Revised: = 1 }

Revised: as

Original:  b) Universal approximation : the span of { φ ( · , ¯ ω ) : ¯ ω ∈ Θ } is dense in L 2 ( S × A ) ; c) Support of the measure : There existsr > 0 s.t.
Revised: r → ∞ to a map ¯ g f (¯ ω ) whose regular values are dense in its range. b) Universal approximation : the span of { φ ( · , ¯ ω ) : ¯ ω ∈ Θ } is dense in L 2 ( S × A ) ; c) Support of the measure : There exists

Revised: r

Original:  the support of the initial condition ν 0 is contained in Q r :=
Revised: > s.t. the support of the initial condition ν 0 is contained in Q r :=

Original: [ − r, r ] × Θ and separates {− r } × Θ from { r } × Θ , i 
Revised: [ − r, r ] × Θ and separates {− r } × Θ from { r } × Θ , i e

Original: e., any continuous path connecting {− r} × Θ to { r } × Θ intersects the support of ν 0 .
Revised: any continuous path connecting {− r × Θ to { r } × Θ intersects the support of ν 0 .

Original: Assumption 1 a) is a common, technical regularity assumption ensuring that (10) is well behaved and controlling the growth, variation and regularity of φ .
Revised: Assumption 1 a) is a common, technical regularity assumption ensuring that (10) is well behaved and controlling the growth, variation and regularity of φ .

Original: Alternative assumptions on the case Θ (cid:54) =
Revised: Alternative assumptions on the case Θ (cid:54) =

Original: R m − 1 are given in the appendix.
Revised: R m − 1 are given in the appendix.

Original: Assumption 1 b) speaks to the approximating power of the nonlinearity, assumed to be expressive enough to approximate any function in L 2 ( S × A ) .
Revised: Assumption 1 b) speaks to the approximating power of the nonlinearity, assumed to be expressive enough to approximate any function in L 2 ( S × A ) .

Original: This condition replaces the convexity assumption from
Revised: This condition replaces the convexity assumption from

Original: Chizat & Bach (2018), as the lack of convex structure in our setting prevents us from identifying the local and global minimization properties of a ﬁxed point.
Revised: Chizat & Bach (2018), as the lack of convex structure in our setting prevents us from identifying the local and global minimization properties of a ﬁxed point.

Original: Indeed, despite the one-point convexity of E (cid:37) [ V π ( s )] as a functional of π (Kakade & Langford, 2002) which can be leveraged in the tabular case, such property will be lost, in general, when restricting to policies through nonlinear function approximation.
Revised: Indeed, despite the one-point convexity of E (cid:37) [ V π ( s )] as a functional of π (Kakade & Langford, 2002) which can be leveraged in the tabular case, such property will be lost, in general, when restricting to policies through nonlinear function approximation.

Original: We bypass this issue by requiring sufﬁcient expressivity of the approximating function class, guaranteeing that the optimal policy can be represented with arbitrary precision.
Revised: We bypass this issue by requiring sufﬁcient expressivity of the approximating function class, guaranteeing that the optimal policy can be represented with arbitrary precision.

Original: Similar assumption on approximability of neural network representation was made in recent analysis of natural policy gradient algorithm (Agarwal et al., 2019).
Revised: Similar assumption on approximability of neural network representation was made in recent analysis of natural policy gradient algorithm (Agarwal et al., 2019).

Original: We note that this assumption is easily satisﬁed by widely used nonlinearities by the universal approximation theorem (Cybenko, 1989; Barron, 1993).
Revised: We note that this assumption is easily satisﬁed by widely used nonlinearities by the universal approximation theorem (Cybenko, 1989; Barron, 1993).

Original: Examples of activation functions satisfying Assumption 1 a)-b) include sigmoid, tanh and Gaussian radial function nonlinearities.
Revised: Examples of activation functions satisfying Assumption 1 a)-b) include sigmoid, tanh and Gaussian radial function nonlinearities.

Original: Extension to analogous results in the ReLU case was discussed in
Revised: Extension to analogous results in the ReLU case was discussed in

Original: Wojtowytsch (2020) for the supervised learning.
Revised: Wojtowytsch (2020) for the supervised learning.

Original: Finally, Assumption 1 c) guarantees that the initial condition is such that the expressivity from b) can actually be exploited.
Revised: Finally, Assumption 1 c) guarantees that the initial condition is such that the expressivity from b) can actually be exploited.

Original: This condition is satisﬁed for example by the product of a uniform distribution on any bounded set A ⊂ R with the normal distribution on Θ or, if Θ is compact, with the uniform distribution on Θ .
Revised: This condition is satisﬁed for example by the product of a uniform distribution on any bounded set A ⊂ R with the normal distribution on Θ or, if Θ is compact, with the uniform distribution on Θ .

Original: 4. C ONVERGENCE OF THE MANY - PARTICLE LIMIT
Revised: 4. C ONVERGENCE OF THE MANY - PARTICLE LIMIT

Original: Before discussing the optimality properties of the dynamics (10), we show that this PDE accurately describes the policy gradient dynamics of a sufﬁciently wide, single layer neural network.
Revised: Before discussing the optimality properties of the dynamics (10), we show that this PDE accurately describes the policy gradient dynamics of a sufﬁciently wide, single layer neural network.

Original: To this aim, we let P 2 (Ω) be the space of probability distributions on Ω with ﬁnite second moment.
Revised: To this aim, we let P 2 (Ω) be the space of probability distributions on Ω with ﬁnite second moment.

Original: Theorem 4.1.
Revised: Theorem 4.1.

Original: Let Assumption 1 hold and let w ( N ) t be a solution of (5) with initial condition w ( N )∈ W = Ω N .
Revised: Let Assumption 1 hold and let w ( N ) t be a solution of (5) with initial condition w ( N )∈ W = Ω N .

Original: If ν ( N )converges to ν 0 ∈ P 2 (Ω) in Wasserstein distance W 2 then ν ( N ) t converges, for every t > 0 , to the unique solution ν t of (10).
Revised: If ν ( N )converges to ν 0 ∈ P 2 (Ω) in Wasserstein distance W 2 then ν ( N ) t converges, for every t > 0 , to the unique solution ν t of (10).

Original: We note that by the law of large numbers for empirical distributions, the condition of convergence of ν ( N )to νis e.g. , satisﬁed when w ( i ) 0 are drawn independently at random from ν 0 .
Revised: We note that by the law of large numbers for empirical distributions, the condition of convergence of ν ( N )to νis e.g. , satisﬁed when w ( i ) 0 are drawn independently at random from ν 0 .

Original: The proof of this result is largely standard, under the given assumptions, and is provided in the appendix for completeness.
Revised: The proof of this result is largely standard, under the given assumptions, and is provided in the appendix for completeness.

Original: The idea of the proof is a canonical propagation of chaos argument (Sznitman, 1991).
Revised: The idea of the proof is a canonical propagation of chaos argument (Sznitman, 1991).

Original: In a nutshell the ﬁrst step of the argument establishes sufﬁcient regularity of the gradient dynamics, allowing to guarantee existence and uniqueness of the solution to (10).
Revised: In a nutshell the ﬁrst step of the argument establishes sufﬁcient regularity of the gradient dynamics, allowing to guarantee existence and uniqueness of the solution to (10).

Original: Then, one bounds the difference in differential updates for the particle system and the mean-ﬁeld dynamics by comparing them with the evolution of the particle system according to a linear, time-inhomogeneous PDE using the drift term of the mean-ﬁeld model.
Revised: Then, one bounds the difference in differential updates for the particle system and the mean-ﬁeld dynamics by comparing them with the evolution of the particle system according to a linear, time-inhomogeneous PDE using the drift term of the mean-ﬁeld model.

Original: The proof is ﬁnally concluded by application of Gronwall inequality.
Revised: The proof is ﬁnally concluded by application of Gronwall inequality.

Original: The main difﬁculty WRT similar results in the literature is to establish the needed Lipschitz continuity of the vector ﬁeld driving the transport PDE: while this is an immediate consequence of assumptions on the activation functions and on the risk functional in the supervised setting, proving this type of regularity requires more effort in the RL setting given the involved dependence of the vector ﬁeld on the measure ν t .
Revised: The main difﬁculty WRT similar results in the literature is to establish the needed Lipschitz continuity of the vector ﬁeld driving the transport PDE: while this is an immediate consequence of assumptions on the activation functions and on the risk functional in the supervised setting, proving this type of regularity requires more effort in the RL setting given the involved dependence of the vector ﬁeld on the measure ν t .

Original: 4. O PTIMALITY
Revised: 4. O PTIMALITY

Original: After discussing the connection between particle dynamics and mean-ﬁeld equations, we present the main convergence result of this paper:
Revised: After discussing the connection between particle dynamics and mean-ﬁeld equations, we present the main convergence result of this paper:

Original: Theorem 4.2.
Revised: Theorem 4.2.

Original: Let Assumption 1 hold and ν t given by (10) converge to ν ∗ , then π ν ∗ = π ∗ , the optimal policy for (3) .
Revised: Let Assumption 1 hold and ν t given by (10) converge to ν ∗ , then π ν ∗ = π ∗ , the optimal policy for (3) .

Original: Thus if the policy gradient dynamics (10) converges to a stationary point, that point must be a global minimizer.
Revised: Thus if the policy gradient dynamics (10) converges to a stationary point, that point must be a global minimizer.

Original: Again, we emphasize that in our regularized setting π ∗ is given by a probability distribution, and can thus be represented as a softmax policy.
Revised: Again, we emphasize that in our regularized setting π ∗ is given by a probability distribution, and can thus be represented as a softmax policy.

Original: To prove this result, weﬁrst connect the optimality of a stationary point with the support of the underlying measure in parameter space.
Revised: We prove this result in three steps. First, we connect the optimality of a stationary point with the support of the underlying measure in parameter space.

Original: Proposition 4.3.
Revised: More speciﬁcally, we show in

Original: Let ν ∗ be a stationary point of (10) such that  the full support property of Assumption 1 c) holds at ν ∗ ,then π ν ∗ = π ∗ . We prove this intermediate result in two steps: ﬁrst we show that, by homogeneity of ψ and Assumption 1 c), the velocity ﬁeld of  (10) must vanish a.e. in Ω = R × Θ . Then, by Assumption 1 b) we deduce that a stationary point ν ∗ cannot result from the orthogonality between Df ν and the error estimate δ ( ν ) = ( Q π ν ( s, a ) − τ log π ν ( s, a ) − V π ν ( s )) , implying δ ( ν ) = 0 .
Revised: Lemma C.1 that by the expressivity of φ , the transport vector ﬁeld of suboptimal ﬁxed points of the dynamics (10) cannot vanish everywhere in parameter space, so that a measure with sufﬁcient support cannot correspond to a suboptimal ﬁxed point.

Original: Following Proposition 4.3, one important technical step in the proof of Theorem 4.2 is to show that the mean-ﬁeld policy gradient dynamics (10)preserves Assumption 1 c) on the support of the measure throughout training, as shown in Lemma C.1.
Revised: We then show in Lemma C.2 that such sufﬁcient notion of support (Assumption 1 c) ) is preserved by the mean-ﬁeld policy gradient dynamics (10) throughout training.

Original: This result holds for any ﬁnite time by topological arguments, as the separation property of the measure cannot be altered by the action of a continuous ﬂow such as (10).
Revised: For any ﬁnite time, this is true by topological arguments: the separation property of the measure cannot be altered by the action of a continuous vector ﬁeld such as (10).

Original: To extend to inﬁnite time (and thus to the stationary point) leveraging the above partial result, we ﬁnally prove  in Lemma C.2, that spurious ﬁxed points are avoided by the policy gradient dynamics (10) when initialized properly.
Revised: We note in particular that we do not prove that assumption a) in Section 3 holds in this case. Finally, in Lemma C.3 we combine the above results and prove that spurious ﬁxed points are avoided by the policy gradient dynamics (10) when initialized properly.

Original: To establish this we argue by contradiction: assuming  that we are approaching such a spurious ﬁxed point ˜ν
Revised: To establish this we argue by contradiction: assuming [Figure] that we are approaching such a spurious ﬁxed point ˜

Original:  at time t 0 , we show in Lemma C.4 that the velocity ﬁeld will change little for any t > t 0 .
Revised: ν at time t 0 , we show in Lemma C.5 that the velocity ﬁeld will change little for any t > t 0 .

Original: In particular, it follows that in this regime the dynamics of (10) can be approximated by the gradient descent dynamics (in particle space) of an approximately [Figure]
Revised: In particular, it follows that in this regime the dynamics of (10) can be approximated by the gradient descent dynamics (in particle space) of an approximately ﬁxed potential.

Original: ﬁxed potential. On the other hand, by Assumption 1 c) and by the homogeneity of ψ , we are able to show that  a positive amount of measure ˜ ν will fall in a forward invariant region where its ω 0 component will grow linearly in t , thereby eventually contradicting the assumption that ˜ ν is a ﬁxed point of (10).
Revised: On the other hand, by Assumption 1 c) and by the homogeneity of ψ , we are able to show that by Lemma C.2 a positive amount of measure ˜ ν will fall in a forward invariant region where its ω 0 component will grow linearly in t (which exists by Lemma C.1), thereby eventually contradicting the assumption that ˜ ν is a ﬁxed point of (10).

Original: There are two main conceptual differences between the proof outlined above and the one carried out in the supervised learning setting.
Revised: There are two main conceptual differences between the proof outlined above and the one carried out in the supervised learning setting.

Original: On one hand, a necessary step in our proof is to establish Lipschitz continuity of the vector ﬁeld deﬁning the transport equation (10), also needed for convergence of the particle dynamics as discussed above.
Revised: On one hand, a necessary step in our proof is to establish Lipschitz continuity of the vector ﬁeld deﬁning the transport equation (10), also needed for convergence of the particle dynamics as discussed above.

Original: On the other hand, the landscape of the objective function for policy gradients does not enjoy the convexity ( WRT the predictor) typically assumed in the supervised case.
Revised: On the other hand, the landscape of the objective function for policy gradients does not enjoy the convexity ( WRT the predictor) typically assumed in the supervised case.

Original: To exclude the existence of local minima we assume sufﬁcient expressivity of the activation functions, absent in the supervised analysis.
Revised: To exclude the existence of local minima we assume sufﬁcient expressivity of the activation functions, absent in the supervised analysis.

Original: This assumption is key to deduce optimality of ﬁxed points of (10) in Proposition 4.3 in our less regular setting.
Revised: This assumption is key to deduce optimality of ﬁxed points of (10) in Theorem 4.2 in our less regular setting.

Original:  N UMERICAL EXAMPLES
Revised:  N UMERICAL EXAMPLES

Original: To test our theoretical results in a simple setting we train a wide, single hidden layer neural network with policy gradients to learn the optimal softmax policy (8) for entropy-regularized rewards with parametrization (7) and regularization parameter τ = 0 .
Revised: To test our theoretical results in a simple setting we train a wide, single hidden layer neural network with policy gradients to learn the optimal softmax policy (8) for entropy-regularized rewards with parametrization (7) and regularization parameter τ = 0 .

Original: We do so in two separate settings:
Revised: We do so in two separate settings:

Original: (a) S = { 0 } , A = [0 , 1] .
Revised: (a) S = { 0 } , A = [0 , 1] .

Original: This setting corresponds to bandits framework discussed in Section
Revised: This setting corresponds to bandits framework discussed in Section

Original: (b) S ×
Revised: (b) S ×

Original: A is a grid of size 100 × 100 in the set [0 , 1] 2 .
Revised: A is a grid of size 100 × 100 in the set [0 , 1] 2 .

Original: In this case, we have chosen a discount factor of γ = 0 .
Revised: In this case, we have chosen a discount factor of γ = 0 .

Original: 7 and a transition process given by P ( s, a, s (cid:48) )
Revised: 7 and a transition process given by P ( s, a, s (cid:48) )

Original: = 0 .
Revised: = 0 .

Original: 9 δ ( s (cid:48) − a ) + 0 .
Revised: 9 δ ( s (cid:48) − a ) + 0 .

Original: 1 / 100 ( i.e. , an action a leads to the corresponding state s (cid:48)
Revised: 1 / 100 ( i.e. , an action a leads to the corresponding state s (cid:48)

Original: =
Revised: =

Original: a with probability 0 .
Revised: a with probability 0 .

Original: and is uniformly distributed with probability 0 . 1 ).
Revised: and is uniformly distributed with probability 0 . 1 ).

Original: At each iteration we have computed the exact distribution (cid:37) π by computing the resolvent of the (weighted) transition matrix.
Revised: At each iteration we have computed the exact distribution (cid:37) π by computing the resolvent of the (weighted) transition matrix

Original: In both cases, we deﬁned the optimal Q function as Q ∗ ( s, a ) = τf ∗ w ( s, a ) , where f ∗ w ( s, a ) is given by a single hidden layer neural network of width n = 5 , ReLU nonlinearites and weights w drawn independently and identically distributed from a centered, normal distribution with variance σ 2 = 4 , i.e. , 
Revised: In both cases, we deﬁned the optimal Q function as Q ∗ ( s, a ) = τf ∗ w ( s, a ) , where f ∗ w ( s, a ) is given by a single hidden layer neural network of width n = 5 , ReLU nonlinearites and weights w drawn independently and identically distributed from a centered, normal distribution with variance σ 2 = 4 , i.e. , [Equation] for w

Original: [Equation]
Revised: [Equation]

Original:  We learn the optimal policy for the problem deﬁned above using a N = 800 -neurons wide single hidden layer neural network with ReLU nonlinearities in the mean-ﬁeld regime (7) used as energy for a softmax policy (8).
Revised: N [Equation] We learn the optimal policy for the problem deﬁned above using a N = 800 -neurons wide single hidden layer neural network with ReLU nonlinearities in the mean-ﬁeld regime (7) used as energy for a softmax policy (8).

Original: The initialization of the student network is as follows: ﬁrst-layer weights are initialized at random drawn independently from a centered, normal distribution with variance σ 2 = 4 , while output weights are initialized at 0 .
Revised: The initialization of the student network is as follows: ﬁrst-layer weights are initialized at random drawn independently from a centered, normal distribution with variance σ 2 = 4 , while output weights are initialized at 0 .

Original: The model is trained according to (4) with ﬁxed step-size β t = 10 − 3 .
Revised: The model is trained according to (4) with ﬁxed step-size β t = 10 − 3 .

Original: We report the results of this training procedure in Fig.
Revised: We report the results of this training procedure in Fig.

Original: 1, where we notice that all the paths monotonically decrease the error E ( ν ∗ ) − E ( ν t ) , as predicted by our results.
Revised: 1, where we notice that all the paths monotonically decrease the error E ( ν ∗ ) − E ( ν t ) , as predicted by our results.

Original: Note that the convergence rate of the model varies across experiments, consistently with the purely qualitative nature of the convergence result we proved.
Revised: Note that the convergence rate of the model varies across experiments, consistently with the purely qualitative nature of the convergence result we proved.

Original:  C ONCLUSIONS AND FUTURE WORK
Revised:  C ONCLUSIONS AND FUTURE WORK

Original: This work addresses the problem of optimality of policy gradient algorithms, a workhorse of deep reinforcement learning, when combined with mean-ﬁeld models such as neural networks.
Revised: This work addresses the problem of optimality of policy gradient algorithms, a workhorse of deep reinforcement learning, when combined with mean-ﬁeld models such as neural networks.

Original: More speciﬁcally, we provide a mean-ﬁeld formulation of the parametric dynamics of policy gradient algorithms for entropy-regularized MDP s and prove that, under mild assumptions, all ﬁxed points of such dynamics are optimal.
Revised: More speciﬁcally, we provide a mean-ﬁeld formulation of the parametric dynamics of policy gradient algorithms for entropy-regularized MDP s and prove that, under mild assumptions, all ﬁxed points of such dynamics are optimal.

Original: This extends similar results obtained in the “neural” or “lazy” regime to the mean-ﬁeld one, which is known to be much more expressive (E et al., 2019; Ghorbani et al., 2020), but also highly nonlinear.
Revised: This extends similar results obtained in the “neural” or “lazy” regime to the mean-ﬁeld one, which is known to be much more expressive (E et al., 2019; Ghorbani et al., 2020), but also highly nonlinear.

Original: The latter feature prevents, at present, from obtaining convergence results of these models, except in very speciﬁc settings (Chizat, 2019; Javanmard et al., 2019).
Revised: The latter feature prevents, at present, from obtaining convergence results of these models, except in very speciﬁc settings (Chizat, 2019; Javanmard et al., 2019).

Original: Interesting avenues or future research include the relaxing the adiabaticity assumption, i.e. , considering the stochastic approximation problem resulting from the ﬁnite number of samples and the ﬁnite gradient step-size, as well as establishing quantitative bounds for models with a large, but ﬁnite, number of parameters.
Revised: Interesting avenues or future research include the relaxing the adiabaticity assumption, i.e. , considering the stochastic approximation problem resulting from the ﬁnite number of samples and the ﬁnite gradient step-size, as well as establishing quantitative bounds for models with a large, but ﬁnite, number of parameters.

Original: Probably the most important open question, however, concerns establishing quantitative convergence of mean-ﬁeld dynamics of neural networks: even in the supervised setting, despite recent results in speciﬁc settings (Chizat, 2019;
Revised: Probably the most important open question, however, concerns establishing quantitative convergence of mean-ﬁeld dynamics of neural networks: even in the supervised setting, despite recent results in speciﬁc settings (Chizat, 2019;

Original: Javanmard et al., 2019), these guarantees remain mainly out of reach.
Revised: Javanmard et al., 2019), these guarantees remain mainly out of reach.

Revised: A CKNOWLEDGMENTS .

Revised: AA acknowledges the support of the Swiss National Science Foundation through the grant P2GEP2-17501 and by the NSF grant DMS-1613337.

Revised: The work of JL is in part supported by the US National Science Foundation via grants CCF-1910571 (Duke TRIPODS) and DMS-2012286.

Original: A D ERIVATION OF SOFTMAX POLICY GRADIENT DYNAMICS
Revised: A D ERIVATION OF SOFTMAX POLICY GRADIENT DYNAMICS

Original: Lemma A
Revised: Lemma A.1.

Original: [Equation]

Original: The gradient of the entropy-regularized value function can be written as
Revised: The gradient of the entropy-regularized value function can be written as

Original: [Equation]
Revised: [Equation]

Original: and thus the policy gradient dynamics (5).
Revised: and thus the policy gradient dynamics (5)

Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We choose throughout ¯ π as the Lebesgue measure
Revised: We choose throughout ¯ π as the Lebesgue measure

Original: [Equation]
Revised: [Equation]

Original: and use that π w is absolutely continuous wrt ¯ π .
Revised: and use that π w is absolutely continuous wrt ¯ π .

Original: Taking the gradient of (1) using a parametric policy π w we obtain
Revised: Taking the gradient of (1) using a parametric policy π w we obtain

Original: [Equation]
Revised: [Equation]

Original: Now, since (cid:82) S P ( s, a, d s (cid:48) )
Revised: Now, since (cid:82) S P ( s, a, d s (cid:48) )

Original: = 1 and (cid:82)
Revised: = 1 and (cid:82)

Original: A π w ( s, d a (cid:48) )
Revised: A π w ( s, d a (cid:48) )

Original: = 1 for all s, a, w we have for the ﬁrst term in brackets in the last line
Revised: = 1 for all s, a, w we have for the ﬁrst term in brackets in the last line

Original: [Equation]
Revised: [Equation]

Original: On the other hand, we can rewrite the second term in brackets as E π w [ γ ∇ V π w ( s 1 ) | s 0 = s ] , and recognize the
Revised: On the other hand, we can rewrite the second term in brackets as E π w [ γ ∇ V π w ( s 1 ) | s 0 = s ] , and recognize the

Original: LHS of (A.1) evaluated at the next state s 1 in the expectation.
Revised: LHS of (A.1) evaluated at the next state s 1 in the expectation.

Original: Therefore, we can sequentially repeat the same computation as above, and recalling the deﬁnition of (cid:37) π w ( · ) and Q π ( s, a ) in (6) and (2) we obtain
Revised: Therefore, we can sequentially repeat the same computation as above, and recalling the deﬁnition of (cid:37) π w ( · ) and Q π ( s, a ) in (6) and (2) we obtain

Original: [Equation]
Revised: [Equation]

Original: where in the second line we have used that, if π w > 0 on A , π w ( s t , d a t ) ∇ w log π w ( s t , a t )
Revised: where in the second line we have used that, if π w > 0 on A , π w ( s t , d a t ) ∇ w log π w ( s t , a t )

Original: = ∇ w π w ( s t , d a t ) .
Revised: = ∇ w π w ( s t , d a t ) .

Original: Proposition 2.1.
Revised: Proposition 2.1.

Original: For a ﬁxed initial distribution (cid:37) 0 , the dynamics (10) is the Wasserstein gradient ﬂow of the energy functional 
Revised: For a ﬁxed initial distribution (cid:37) 0 , the dynamics (10) is the Wasserstein gradient ﬂow of the energy functional E

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: We ﬁnd the potential of the gradient ﬂow by functional differentiation of E :
Revised: We ﬁnd the potential of the gradient ﬂow by functional differentiation of E :

Original: [Equation]
Revised: [Equation]

Original: and consider the two terms in the integral separately, starting from the second:
Revised: and consider the two terms in the integral separately, starting from the second:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: For the ﬁrst term in the integrand of (A.2)
Revised: For the ﬁrst term in the integrand of (A.2)

Original: [Equation]
Revised: [Equation]

Original: we use π ( s, a ) as a density and obtain
Revised: we use π ( s, a ) as a density and obtain

Original: [Equation]
Revised: [Equation]

Original: A.5) where in the last line we have used that δδπ [ π log π ]( s, a ) =
Revised: A.5) where in the last line we have used that δδπ [ π log π ]( s, a ) =

Original: log π ( s, a ) + 1 .
Revised: log π ( s, a ) + 1 .

Original: We evaluate the variational derivative of (cid:37) π as δ(cid:37)
Revised: We evaluate the variational derivative of (cid:37) π as δ(cid:37)

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: We further recognize the same derivative on the RHS of the above expression, allowing to write
Revised: We further recognize the same derivative on the RHS of the above expression, allowing to write

Original: [Equation]
Revised: [Equation]

Original: Finally
Revised: Finally

Original: [Equation]
Revised: [Equation]

Original: we notice that the last term in (A.4) is constant in a and therefore vanishes when integrated against (A.3)
Revised: we notice that the last term in (A.4) is constant in a and therefore vanishes when integrated against (A.3), so we only consider

Original: [Equation]
Revised: [Equation]

Original: so we only consider

Original: [Equation]
Revised: [Equation]

Original: We conclude by noting that combining (A.3) and (A.6) we obtain C π [ ψ ( ω ) , Q π − τ log π ]( s ) , and the Wasserstein gradient ﬂow corresponding to this potential is (10).
Revised: We conclude by noting that combining (A.3) and (A.6) we obtain C π [ ψ ( ω ) , Q π − τ log π ]( s ) , and the Wasserstein gradient ﬂow corresponding to this potential is (10).

Original: B P ROOFS OF THE MANY - PARTICLE LIMIT
Revised: B P ROOFS OF THE MANY - PARTICLE LIMIT

Original: Theorem 4.1.
Revised: Theorem 4.1.

Original: Let Assumption 1 hold and let w ( N ) t be a solution of (5) with initial condition w ( N )∈ W = Ω N .
Revised: Let Assumption 1 hold and let w ( N ) t be a solution of (5) with initial condition w ( N )∈ W = Ω N .

Original: If ν ( N )converges to ν 0 ∈ P 2 (Ω) in Wasserstein distance W 2 then ν ( N ) t converges, for every t > 0 , to the unique solution ν t of (10).
Revised: If ν ( N )converges to ν 0 ∈ P 2 (Ω) in Wasserstein distance W 2 then ν ( N ) t converges, for every t > 0 , to the unique solution ν t of (10).

Original: Proof.
Revised: Proof.

Original: As anticipated in the main text, the proof is divided in two parts:
Revised: As anticipated in the main text, the proof is divided in two parts:

Original: We prove sufﬁcient regularity of the dynamics (10), allowing to establish existence and uniqueness of its solution.
Revised: We prove sufﬁcient regularity of the dynamics (10), allowing to establish existence and uniqueness of its solution.

Original: We leverage the regularity proven above to establish a propagation of chaos result, showing that the system of interacting particles behaves asymptotically as its mean-ﬁeld limit.
Revised: We leverage the regularity proven above to establish a propagation of chaos result, showing that the system of interacting particles behaves asymptotically as its mean-ﬁeld limit.

Original: While carrying out this proof is needed in our context since the dependence of E ( ν ) on ν is more involved than in e.g. , Mei et al.
Revised: While carrying out this proof is needed in our context since the dependence of E ( ν ) on ν is more involved than in e.g. , Mei et al.

Original: ; Rotskoff & Vanden-Eijnden (2018); Chizat & Bach (2018), the steps of this derivation are mainly standard, see e.g., (Sznitman, 1991).
Revised: ; Rotskoff & Vanden-Eijnden (2018); Chizat & Bach (2018), the steps of this derivation are mainly standard, see e.g., (Sznitman, 1991).

Original: B. R EGULARITY
Revised: B. R EGULARITY

Original: We prove existence and uniqueness of the gradient ﬂow dynamics (10) through standard arguments from the optimal transportation literature (see e.g. , (Ambrosio et al., 2008)).
Revised: We prove existence and uniqueness of the gradient ﬂow dynamics (10) through standard arguments from the optimal transportation literature (see e.g. , (Ambrosio et al., 2008)).

Original: More speciﬁcally, recalling that π = π
Revised: More speciﬁcally, recalling that π = π

Original: ν we leverage the Lipschitz continuity of the vector ﬁeld
Revised: ν we leverage the Lipschitz continuity of the vector ﬁeld

Original: [Equation]
Revised: [Equation]

Original: with respect to ν .
Revised: with respect to ν .

Original: To prove such regularity result, decompose
Revised: To prove such regularity result, decompose

Original: [Equation]
Revised: [Equation]

Original: for

Original: [Equation]

Original: and S : ( S → M 1+ ( A ))
Revised: and S : ( S → M 1+ ( A ))

Original: → R maps µ (cid:55)→ E µ [ V µ ( s ) | s ∼ (cid:37) 0 ] and π : L 2 ( S × A ) → ( S → M 1+ ( A )) is the softmax policy parametrization (8) of its argument.
Revised: → R maps µ (cid:55)→ E µ [ V µ ( s ) | s ∼ (cid:37) 0 ] and π : L 2 ( S × A ) → ( S → M 1+ ( A )) is the softmax policy parametrization (8) of its argument.

Original: Recalling the deﬁnition of Q r from Assumption 1 and denoting F r = { (cid:82) ψν : supp ν ∈ Q r } , we further deﬁne the norms and constants needed in the following proof as:
Revised: Recalling the deﬁnition of Q r from Assumption 1 and denoting F r = { (cid:82) ψν : supp ν ∈ Q r } , we further deﬁne the norms and constants needed in the following proof as:

Original: [Equation]
Revised: [Equation]

Original: where (cid:107) · (cid:107) denotes the operator norm.
Revised: where (cid:107) · (cid:107) denotes the operator norm.

Original: While for any r > 0 the boundedness of (cid:107) Dψ (cid:107) r, ∞ , L Dψ results directly from Assumption 1, more work is needed to prove that (cid:107) DR (cid:107) r, ∞ < ∞ , L DR < ∞ .
Revised: While for any r > 0 the boundedness of (cid:107) Dψ (cid:107) r, ∞ , L Dψ results directly from Assumption 1, more work is needed to prove that (cid:107) DR (cid:107) r, ∞ < ∞ , L DR < ∞ .

Original: We prove this in
Revised: We prove this in

Original: Lemma C.4 below, and proceed with the proof of convergence of the particle dynamics.
Revised: Lemma C.5 below, and proceed with the proof of convergence of the particle dynamics.

Original: For any r and corresponding Q r from Assumption 1 we deﬁne the set of localized functionals
Revised: For any r and corresponding Q r from Assumption 1 we deﬁne the set of localized functionals

Original: [Equation]
Revised: [Equation]

Original:  Furthermore, we say that a coupling γ ∈ M 1+ (Ω × Ω) is an admissible transport plan if both its marginals have support in Q r and ﬁnite second moments.
Revised: else Furthermore, we say that a coupling γ ∈ M 1 + (Ω × Ω) is an admissible transport plan if both its marginals have support in Q r and ﬁnite second moments.

Original: To every admissible transport plan, for p ≥ 1 we associate a transportation costC p
Revised: To every admissible transport plan, for p ≥ 1 we associate a transportation cost

Original: [Equation]
Revised: [Equation]

Original: γ ) =

Original: [Equation]

Original: We prove the following results: for every r > 0 we have
Revised: We prove the following results: for every r > 0 we have

Original: There exists λ r > 0 such that for any admissible transport plan γ , deﬁning the interpolation map ν γt :=
Revised: There exists λ r > 0 such that for any admissible transport plan γ , deﬁning the interpolation map ν γt :=

Original: ( t Π 0 + (1 − t )Π 1 ) # γ , the function t (cid:55)→ E ( ν γt ) is differentiable with Lipschitz continuous derivative with constant λ r C 22 ( γ ) .
Revised: ( t Π 0 + (1 − t )Π 1 ) # γ , the function t (cid:55)→ E ( ν γt ) is differentiable with Lipschitz continuous derivative with constant λ r C 22 ( γ ) .

Original: Let ν 0 have support in Q r .
Revised: Let ν 0 have support in Q r .

Original: Then for any given transport plan γ with ﬁrst marginal given by ν 0 , a velocity ﬁeld F satisﬁes
Revised: Then for any given transport plan γ with ﬁrst marginal given by ν 0 , a velocity ﬁeld F satisﬁes

Original: [Equation]
Revised: [Equation]

Original: if and only if F ( u ) is in the subdifferential of D E ν ( u ) :
Revised: if and only if F ( u ) is in the subdifferential of D E ν ( u ) :

Original: = C π [ ψ, Q π − log π ]( u ) for g ∈ L 2 ( S × A ) (projected in the interior of Q r when u ∈ ∂ Q r ) for ν 0 almost every u ∈ Ω .
Revised: = C π [ ψ, Q π − log π ]( u ) for g ∈ L 2 ( S × A ) (projected in the interior of Q r when u ∈ ∂ Q r ) for ν 0 almost every u ∈ Ω .

Original: The proof of the two points above corresponds to (Chizat & Bach, 2018, Lemma B.2).
Revised: The proof of the two points above corresponds to (Chizat & Bach, 2018, Lemma B.2).

Original: We sketch the proof of these two points below, referring to the original reference for the details.
Revised: We sketch the proof of these two points below, referring to the original reference for the details.

Original: By the Lipschitz continuity of ψ : Ω → L 2 ( S × A ) and R (cid:48)
Revised: By the Lipschitz continuity of ψ : Ω → L 2 ( S × A ) and R (cid:48)

Original: ( f ) =
Revised: ( f ) =

Original: DR f = C π [ · , Q π − τ log π ] in Q r , the energy E ( r ) ( ν γt ) transported along an interpolating path ν γt is differentiable and we can write its derivative as
Revised: DR f = C π [ · , Q π − τ log π ] in Q r , the energy E ( r ) ( ν γt ) transported along an interpolating path ν γt is differentiable and we can write its derivative as

Original: [Equation]
Revised: [Equation]

Original: Then, again by the Lipschitz continuity of Dψ and DR we have for 0 ≤ t (cid:48) <
Revised: Then, again by the Lipschitz continuity of Dψ and DR we have for 0 ≤ t (cid:48) <

Original: t (cid:48)(cid:48) <
Revised: t (cid:48)(cid:48) <

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]
Revised: [Equation]

Original: for a λ r large enough, where in the last inequality we have used the uniform bounds on DR in Q r , that
Revised: for a λ r large enough, where in the last inequality we have used the uniform bounds on DR in Q r , that

Original: [Equation]
Revised: [Equation]

Original: ω [Equation] | and we applied H ¨ older’s inequality to bound C 21 ( γ ) ≤ C 22 ( γ ) .
Revised: | and we applied H ¨ older’s inequality to bound C 21 ( γ ) ≤ C 22 ( γ ) .

Original: The proof of this result leverages an expansion of the functionals R, ψ to the second order in their arguments:
Revised: The proof of this result leverages an expansion of the functionals R, ψ to the second order in their arguments:

Original: [Equation]
Revised: [Equation]

Original: Recalling the Lipschitz bounds on the remainders R ψ ( ω, ω(cid:48) )
Revised: Recalling the Lipschitz bounds on the remainders R ψ ( ω, ω

Original: <
Revised: [Equation]

Original:  12 L Dψ | ω − ω (cid:48) | 2 , R R ( f, g ) <
Revised: ) < 12 L Dψ | ω − ω (cid:48) | 2 , R R ( f, g ) <

Original: 12 L DR (cid:107) f
Revised: 12 L DR (cid:107) f

Original: − g (cid:107) 22 and combining the two expansions above we have, for a transport plan γ with marginals ν 0 =
Revised: − g (cid:107) 22 and combining the two expansions above we have, for a transport plan γ with marginals ν 0 =

Original: (Π 0 )
Revised: (Π 0 )

Original: [Equation]
Revised: [Equation]

Revised: γ ,

Original: [Equation]
Revised: [Equation]

Original: for a remainder term R .
Revised: for a remainder term R .

Original: We can bound such remainder term, again by the Lipschitz regularity of
Revised: We can bound such remainder term, again by the Lipschitz regularity of

Original: Dψ and DR and by the boundedness of Dψ, DR in Q r , by C 2 ( γ ) 2 and C 1 ( γ ) 2 ≤ C 2 ( γ ) 2 , thereby obtaining that
Revised: Dψ and DR and by the boundedness of Dψ, DR in Q r , by C 2 ( γ ) 2 and C 1 ( γ ) 2 ≤ C 2 ( γ ) 2 , thereby obtaining that

Original: [Equation]
Revised: [Equation]

Original: Noting that the integrand against the coupling is the gradient ﬂow vector ﬁeld, the above uniquely characterizes the velocity ﬁeld satisfying (B.2).
Revised: Noting that the integrand against the coupling is the gradient ﬂow vector ﬁeld, the above uniquely characterizes the velocity ﬁeld satisfying (B.2).

Original: We note that point 1) above immediately implies that E ( · ) is λ r -semiconvex along geodesics, while by 2) E ( r ) ( · ) admits strong Wasserstein subdifferentials on its domain (Ambrosio et al., 2008, Deﬁnition 10.3.1).
Revised: We note that point 1) above immediately implies that E ( · ) is λ r -semiconvex along geodesics, while by 2) E ( r ) ( · ) admits strong Wasserstein subdifferentials on its domain (Ambrosio et al., 2008, Deﬁnition 10.3.1).

Original: Combining the two results one obtains existence and uniqueness of the solutions of the Wasserstein gradient ﬂow through (Ambrosio et al., 2008, Theorem 11.2.1).
Revised: Combining the two results one obtains existence and uniqueness of the solutions of the Wasserstein gradient ﬂow through (Ambrosio et al., 2008, Theorem 11.2.1).

Original: B. P ROPAGATION OF CHAOS
Revised: B. P ROPAGATION OF CHAOS

Original: By the Lipschitz continuity of the transport ﬁeld in (10) in ν 0 with supp ν 0 ⊂ Q r , there exists a time t r >such that, supp ν ( N ) s ⊂ Q r for all s ∈ [0 , t r ] , N ∈ N .
Revised: By the Lipschitz continuity of the transport ﬁeld in (10) in ν 0 with supp ν 0 ⊂ Q r , there exists a time t r >such that, supp ν ( N ) s ⊂ Q r for all s ∈ [0 , t r ] , N ∈ N .

Original: Consider now two times 0 ≤ t 1 < t 2 ≤ t r .
Revised: Consider now two times 0 ≤ t 1 < t 2 ≤ t r .

Original: To prove the existence of the limiting curve ( ν t ) t , we show that the curves ν ( N ) t are uniformly in N equicontinuous in Wand as such, possess a converging subsequence by Arzela-Ascoli theorem.
Revised: To prove the existence of the limiting curve ( ν t ) t , we show that the curves ν ( N ) t are uniformly in N equicontinuous in Wand as such, possess a converging subsequence by Arzela-Ascoli theorem.

Original: To show equicontinuity, we bound the the W 2 distance between distributions by coupling positions of the same particles at different times and using Cauchy-Schwartz inequality:
Revised: To show equicontinuity, we bound the the W 2 distance between distributions by coupling positions of the same particles at different times and using Cauchy-Schwartz inequality:

Original: [Equation]
Revised: [Equation]

Original: Combining the above with the identity
Revised: Combining the above with the identity

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: [Equation]
Revised: [Equation]

Original: where we recall that F r = { ν ∈ M 1+ (Ω) :
Revised: where we recall that F r = { ν ∈ M 1+ (Ω) :

Original: supp ν ⊂ Q r } .
Revised: supp ν ⊂ Q r } .

Original: In particular, the above continuity bound is independent on N , proving equicontinuity of ν ( N ) in W 2 .
Revised: In particular, the above continuity bound is independent on N , proving equicontinuity of ν ( N ) in W 2 .

Original: We now prove that the limiting point of the converging subsequence whose existence was identiﬁed above must solve (10).
Revised: We now prove that the limiting point of the converging subsequence whose existence was identiﬁed above must solve (10).

Original: To do so we compare both the differential for the mean-ﬁeld and particle dynamics to the one of a linear inhomogneous PDE : for any bounded and continuous f : R × R m → R m , denoting by E =
Revised: To do so we compare both the differential for the mean-ﬁeld and particle dynamics to the one of a linear inhomogneous PDE : for any bounded and continuous f : R × R m → R m , denoting by E =

Original: ν t F t d t and
Revised: ν t F t d t and

Original: E N = ν ( N ) t
Revised: E N = ν ( N ) t

Original: F ( N ) t d t where F t , F ( N ) t are the vector ﬁelds of the mean-ﬁeld and particle system respectively, we write
Revised: F ( N ) t d t where F t , F ( N ) t are the vector ﬁelds of the mean-ﬁeld and particle system respectively, we write

Original: [Equation]
Revised: [Equation]

Original: By boundedness of fF t over Q r , the second term converges by our choice of subsequence.
Revised: By boundedness of fF t over Q r , the second term converges by our choice of subsequence.

Original: For the ﬁrst term, denoting throughout by (cid:107) · (cid:107) BL the bounded Lipschitz norm, we leverage the Lipschitz continuity of the vector ﬁeld F WRT the underlying parametric measure:
Revised: For the ﬁrst term, denoting throughout by (cid:107) · (cid:107) BL the bounded Lipschitz norm, we leverage the Lipschitz continuity of the vector ﬁeld F WRT the underlying parametric measure:

Original: [Equation]
Revised: [Equation]

Original: for C r > 0 large enough, again obtaining convergence by our choice of subsequence.
Revised: for C r > 0 large enough, again obtaining convergence by our choice of subsequence.

Original: This proves convergence of the particle model to the mean-ﬁeld equation (10) on [0 , t r ] .
Revised: This proves convergence of the particle model to the mean-ﬁeld equation (10) on [0 , t r ] .

Original: To extend the time interval on which we prove convergence, we use that E ( ν ) (and thus R ) decays along trajectories of (10).
Revised: To extend the time interval on which we prove convergence, we use that E ( ν ) (and thus R ) decays along trajectories of (10).

Original: Consequently, by the boundedness of the differential DR on sublevel sets of R the Lipschitz constant of E ( ν t ) is uniformly bounded.
Revised: Consequently, by the boundedness of the differential DR on sublevel sets of R the Lipschitz constant of E ( ν t ) is uniformly bounded.

Original: Using again the Lipschitz continuity of F we can show that sup u ∈Q r (cid:107) F (cid:107) <
Revised: Using again the Lipschitz continuity of F we can show that sup u ∈Q r (cid:107) F (cid:107) <

Original: A + Br , i.e. , that particle velocities can grow at most linearly in r , and an application of Gronwall allows to ﬁnd, for every T > 0 that there exists r > 0 such that supp ν t ∈ Q r for all t ∈ [0 , T ] and propagation of chaos follows.
Revised: A + Br , i.e. , that particle velocities can grow at most linearly in r , and an application of Gronwall allows to ﬁnd, for every T > 0 that there exists r > 0 such that supp ν t ∈ Q r for all t ∈ [0 , T ] and propagation of chaos follows.

Original: C P ROOFS OF OPTIMALITY
Revised: C P ROOFS OF OPTIMALITY

Original: Theorem 4.2.
Revised: Theorem 4.2.

Original: Let Assumption 1 hold and ν t given by (10) converge to ν ∗ , then π ν ∗ =
Revised: Let Assumption 1 hold and ν t given by (10) converge to ν ∗ , then π ν ∗ =

Original: π ∗ S × A -a.e.
Revised: π ∗ S × A -a.e.

Original: We prove the above result as sketched in the bandit setting by ﬁrst establishing optimality of full support stationary measures in parametric space through Proposition 4.3, and then investigating how the dynamics preserves full support property for any t > 0 in Lemma C.1 and avoids spurious minima in Lemma C.2.

Original: Before starting this program we introduce the equivalent of greedy policies in the entropy-regularized setting.

Original: For a given Q ( s, a ) , the associated Boltzmann policy π B with respect to a reference measure ¯ π is given by

Original: [Equation]

Original: for

Original: [Equation]

Original: and satisﬁes

Original: [Equation]

Original: One can then deﬁne the Boltzmann backup or soft Bellman backup operator T τ that, for a given Q and the associated Boltzmann policy π B , gives the action-value function T τ Q associated to π B :

Original: [Equation]

Original: It is known (Haarnoja et al., 2018b, Theorem 3) that the above operator is a contraction, and that it therefore has a unique ﬁxed point, corresponding to the optimal policy π ∗ B = π ∗ of the entropy-regularized MDP.

Original: With the above information at hand we now proceed to prove Proposition 4.3:

Original: Proposition 4.3.

Original: Let ν ∗ be a ﬁxed point of (10) such that the full support property of Assumption 1 c) holds at ν ∗ , then π ν ∗ = π ∗ .

Original: Proof of Proposition 4.3.

Original: A ﬁxed point ν = ν ∗ of (10) must satisfy

Original: [Equation]

Original: for all ω (cid:48) ∈ supp ( ν ) .

Original: In particular, the ﬁrst component (in ω 0 direction) of the above gradient must be 0 :

Original: [Equation]

Original: and by homogeneity of ψ the above identity must hold for all ¯ ω (cid:48) in the projection of supp ( ν ) on Θ .

Original: Then, by the full support property of ν from Assumption 1 c), (C.3) must hold for all ¯ ω (cid:48) ∈ Θ .

Original: To deduce from this the desired invariance in Ω , we apply Assumption 1 b) on the universal approximation property of the nonlinearity φ and the full support property of π ν given by the chosen parametrization, which imply that

Original: [Equation]

Original: Finally, we rewrite the above condition in compact notation as the ﬁxed point equation

Original: [Equation]

Original: for the soft Q learning or Boltzmann backup operator T τ deﬁned in (C.1).

Original: Since T τ is a contraction in (cid:107) · (cid:107)for γ < 1 Haarnoja et al.

Original: (2018b), its ﬁxed point is unique and hence π ν ∗ =

Original: π B [ Q ∗ ] =

Original: π ∗ for almost every ( s, a ) ∈ S × A .

Original: Before proceeding to prove Theorem 4.2, we state the alternative form of Assumption 1 a) in the case where Θ (cid:54) =
Revised: Before proceeding to prove Theorem 4.2, we state the alternative form of Assumption 1 a) in the case where Θ (cid:54) =

Original: R m − 1 .
Revised: R m − 1 .

Original: Our proof of the theorem above can be easily generalized.
Revised: Our proof of the theorem above can be easily generalized to the setting of Assumption 2.

Original: Assumption 2.
Revised: Assumption 2.

Original: Assume that ω
Revised: Assume that ω

Original: = ( ω 0 , ¯ ω ) ∈ R × Θ for Θ ⊂ R m − 1 which is the closure of a bounded open convex set.
Revised: = ( ω 0 , ¯ ω ) ∈ R × Θ for Θ ⊂ R m − 1 which is the closure of a bounded open convex set.

Original: Furthermore ψ ( s, a ; ω ) =
Revised: Furthermore ψ ( s, a ; ω ) =

Original: ω 0 φ ( s, a ; ¯ ω ) where φ is bounded, differentiable and Dφ ω is Lipschitz.
Revised: ω 0 φ ( s, a ; ¯ ω ) where φ is bounded, differentiable and Dφ ω is Lipschitz.

Original: Also, for all f ∈ L 2 ( S × A ) the regular values of the map ¯ ω
Revised: Also, for all f ∈ L 2 ( S × A ) the regular values of the map ¯ ω

Original: (cid:55)→ g f (¯ ω ) :
Revised: (cid:55)→ g f (¯ ω ) :

Original: = (cid:82) f
Revised: = (cid:82) f

Original: ( s, a ) φ ( s, a ; ¯ ω )
Revised: ( s, a ) φ ( s, a ; ¯ ω )

Original: d a d s are dense in its range and g f (¯ ω ) satisﬁes Neumann boundary conditions (i.e., for all ¯ ω ∈ ∂ Θ we have dg f (¯ ω )( n ¯ ω )
Revised: d a d s are dense in its range and g f (¯ ω ) satisﬁes Neumann boundary conditions (i.e., for all ¯ ω ∈ ∂ Θ we have dg f (¯ ω )( n ¯ ω )

Original: =where n ¯ ω ∈ R m − 1 is the normal of ∂ Θ at ¯ ω ).
Revised: =where n ¯ ω ∈ R m − 1 is the normal of ∂ Θ at ¯ ω ).

Revised: We prove Theorem 4.2 as sketched in Section 4 by ﬁrst connecting the optimality and the support of stationary measures in parametric space through Lemma C.1, and then investigating how the dynamics preserves full support property for any t > 0 in Lemma C.2 and avoids spurious minima in Lemma C.3.

Revised: Before starting this program we introduce the equivalent of greedy policies in the entropy-regularized setting.

Original: It is instrumental for the following proof to note that in light of the above spurious minima necessarily correspond to distributions ν such that (cid:90)
Revised: For a given Q ( s, a ) , the associated Boltzmann policy π B with respect to a reference measure ¯ π is given by [Equation] [Equation] )) /τ ] for

Revised: [Equation]

Revised: , and satisﬁes

Revised: [Equation]

Revised: One can then deﬁne the Boltzmann backup or soft Bellman backup operator T τ that, for a given Q and the associated Boltzmann policy π B , gives the action-value function T τ Q associated to π B :

Original: [Equation]
Revised: [Equation]

Original: does not hold ¯ ω (cid:48) -almost everywhere in Θ .
Revised: It is known (Haarnoja et al., 2018b, Theorem 3) that the above operator is a contraction, and that it therefore has a unique ﬁxed point, corresponding to the optimal policy π ∗B = π ∗ of the entropy-regularized MDP. To state the ﬁrst partial result towards the proof of Theorem 4.2, we observe that the ω 0 -component of the transport vector ﬁeld in (10) can be written as [Equation]

Original: Indeed, if this is the case we must have that Q π ν (cid:54) =
Revised: where we recall that C π [ f, g ] is the covariance operator WRT the probability measure π ( s, d a ) introduced below (10). We note in particular that the above expression only depends on ¯ ω . With the above information at hand we now proceed to prove Lemma C.1 relating the value of (C.2) and the optimality of ﬁxed points of (10): Lemma C.1.

Original: T τ Q π ν on some subset of A × S with positive measure so thatultimately Q π ν(cid:54) =
Revised: Let Assumption 1 hold and let ν satisfy [Equation] ¯ ω -almost everywhere in Θ . Then we have that Q π ν =

Original: Q π ∗.
Revised: Q π ∗

Revised: [Equation]

Original: C. P ROOF OF T HEOREM 4.
Revised: π [Equation] in S × A . Proof of Lemma C.1.

Revised: Assuming that (C.3) holds Lebesgue-a.e. in Θ , by the assumed continuity of φ in ¯ ω combined with the expressivity of φ Assumption 1 b) we must have that

Revised: [Equation]

Revised: Finally, we rewrite the above condition in compact notation as the ﬁxed point equation

Revised: [Equation]

Revised: for the soft Q learning or Boltzmann backup operator T τ deﬁned in (C.1).

Revised: Since all ﬁxed points of T τ for γ <are optimal (Nachum et al., 2017, Theorem 3), we must have that π ν = π B [ Q ∗ ]

Revised: =

Original: {+π ∗ for (cid:82) P π (cid:37) 
+}
Revised: {+π ∗ for (cid:82) P π (cid:37) 
+}

Revised: -almost every ( s, a ) ∈ S × A .

Revised: Consequently, suboptimal ﬁxed points of the dynamics (10) cannot satisfy (C.3) Lebesgue-a.e. in Θ .

Revised: C. P ROOF OF T HEOREM 4.

Original: We prove below that spurious local minima that do not satisfy (C.4) Θ -a.e.
Revised: We prove below that spurious local minima that do not satisfy (C.3) Θ -a.e.

Original: are avoided by the dynamics.
Revised: are avoided by the dynamics.

Original: We do so by leveraging the approximate gradient structure of the policy
Revised: We do so by leveraging the approximate gradient structure of the policy

Original: gradient vector ﬁeld when ν t is close to one of such stationary points, as discussed in the main text.
Revised: gradient vector ﬁeld when ν t is close to one of such stationary points, as discussed in the main text.

Original: Combining this fact with the assumed convergence to ν ∗ proves Theorem 4.2.
Revised: Combining this fact with the assumed convergence to ν ∗ proves Theorem 4.2.

Original: We note that by the assumed homogeneity of ψ in its ﬁrst component, if ν , ν (cid:48) are such that
Revised: We note that by the assumed homogeneity of ψ in its ﬁrst component, if ν , ν (cid:48) are such that

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: then
Revised: then

Original: [Equation]
Revised: [Equation]

Original: so that in turn we have π ν = π ν (cid:48)
Revised: so that in turn we have π ν = π ν (cid:48)

Original: a.e..
Revised: a.e..

Original: In other words, the homogeneity of the chosen class of approximators results in a degeneracy of the map ψ : M 1+ (Ω) (cid:55)→ L 2 ( S × A ) .
Revised: In other words, the homogeneity of the chosen class of approximators results in a degeneracy of the map ψ : M 1+ (Ω) (cid:55)→ L 2 ( S × A ) .

Original: To remove this degeneracy in our analysis, we identify all the distributions ν, ν (cid:48) that are equivalent under (C.5) by deﬁning the signed measure
Revised: To remove this degeneracy in our analysis, we identify all the distributions ν, ν (cid:48) that are equivalent under (C.4) by deﬁning the signed measure

Original: [Equation]
Revised: [Equation]

Original: Leveraging this deﬁnition, we prove the desired result Theorem 4.2 in two key steps: we show that
Revised: Leveraging this deﬁnition, we prove the desired result Theorem 4.2 in two key steps: we show that

Original: the solution to (10) does not lose (projected) support for any ﬁnite time, thereby preserving the property from Assumption 1 c), 2.
Revised: the solution to (10) does not lose (projected) support for any ﬁnite time, thereby preserving the property from Assumption 1 c), 2.

Original: stationary points ˜
Revised: stationary points ˜

Original: ν with Q π ˜ ν
Revised: ν with Q π ˜ ν

Original: (cid:54) =
Revised: (cid:54) =

Original: Q π ∗ – which by Theorem 4.2 cannot have full support – are avoided by the dynamics.
Revised: Q π ∗ – which by Theorem 4.2 cannot have full support – are avoided by the dynamics.

Original: These facts are respectively summarized in the following lemmas:
Revised: These facts are respectively summarized in the following lemmas:

Original: Lemma C.1.
Revised: Lemma C.2.

Original: Let Assumption 1 a) hold and let ν 0 satisfy Assumption 1 c), then for every t > 0 , ν t solving (10) with initial condition ν 0 also satisﬁes Assumption 1 c).
Revised: Let Assumption 1 a) hold and let ν 0 satisfy Assumption 1 c), then for every t > 0 , ν t solving (10) with initial condition ν 0 also satisﬁes Assumption 1 c).

Original: Throughout, we let (cid:107) · (cid:107)
Revised: Throughout, we let (cid:107) · (cid:107)

Original: BL denote the bounded Lipschitz norm.
Revised: BL denote the bounded Lipschitz norm.

Original: Lemma C.2.
Revised: Lemma C.3.

Original: Let Assumption 1 hold and let ˜ ν be a ﬁxed point of (10) such that (C.4) does not hold a.e..
Revised: Let Assumption 1 hold and let ˜ ν be a ﬁxed point of (10) such that (C.3) does not hold a.e..

Original: Then there exists ε > 0 such that if (cid:107) h 1˜ ν −
Revised: Then there exists ε > 0 such that if (cid:107) h 1˜ ν −

Original: h 1 ν t 1 (cid:107) BL < ε for a t 1 > 0 there exists t 2 > t 1 such that (cid:107) h 1˜ ν
Revised: h 1 ν t 1 (cid:107) BL < ε for a t 1 > 0 there exists t 2 > t 1 such that (cid:107) h 1˜ ν

Original: − h 1 ν t 2 (cid:107) BL > ε .
Revised: − h 1 ν t 2 (cid:107) BL > ε .

Original: Proof of Lemma C.1.
Revised: Proof of Lemma C.2.

Original: Analogously to (Chizat & Bach, 2018, Lemma C.13), we aim to show that the separation property Assumption 1 c) is preserved by the evolution of ν 0 along the characteristic curves X ( t, u ) solving
Revised: Analogously to (Chizat & Bach, 2018, Lemma C.13), we aim to show that the separation property Assumption 1 c) is preserved by the evolution of ν 0 along the characteristic curves X ( t, u ) solving

Original: [Equation]
Revised: [Equation]

Original: where F t is the transport ﬁeld in (10).
Revised: where F t is the transport ﬁeld in (10).

Original: To reach this conclusion, the analogous result in Chizat & Bach (2018) only relies on the continuity of the map u
Revised: To reach this conclusion, the analogous result in Chizat & Bach (2018) only relies on the continuity of the map u

Original: (cid:55)→ X ( t, u ) , established in (Chizat & Bach, 2018, Lemma B.4) under
Revised: (cid:55)→ X ( t, u ) , established in (Chizat & Bach, 2018, Lemma B.4) under

Original: Assumption 1 a).
Revised: Assumption 1 a).

Original: Hence, it is sufﬁcient for our purposes to establish continuity of the map X ( t, · ) from (C.7).
Revised: Hence, it is sufﬁcient for our purposes to establish continuity of the map X ( t, · ) from (C.6).

Original: This property, however results immediately from the one-sided Lipschitz continuity of the vector ﬁeld F t on
Revised: This property, however results immediately from the one-sided Lipschitz continuity of the vector ﬁeld F t on

Original: Q r = [ − r, r ] × Θ uniformly on compact time intervals, which is in turn guaranteed by the Lipschitz continuity and Lipschitz smoothness of ψ from Assumption 1 and boundedness of r .
Revised: Q r = [ − r, r ] × Θ uniformly on compact time intervals, which is in turn guaranteed by the Lipschitz continuity and Lipschitz smoothness of ψ from Assumption 1 and boundedness of r .

Original: To simplify the notation in the following proof, we denote throughout
Revised: To simplify the notation in the following proof, we denote throughout

Original: [Equation]
Revised: [Equation]

Revised: π ν

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: Proof of Lemma C.2.
Revised: Proof of Lemma C.3.

Original: We ﬁrst claim that by Theorem 4.2, for any spurious ﬁxed point (such that Q π ˜ ν
Revised: We ﬁrst claim that by Theorem 4.2, for any spurious ﬁxed point (such that Q π ˜ ν

Original: (cid:54) =
Revised: (cid:54) =

Original: Q π ∗ a.e.)
Revised: Q π ∗ a.e.)

Original: , there must exist a subset of Θ with positive Lebesgue measure where ˜ ν loses support and such that
Revised: , there must exist a subset of Θ with positive Lebesgue measure where ˜ ν loses support and such that

Original: [Equation]
Revised: [Equation]

Original: ψ, δ (˜ ν ) (cid:105) π (cid:54) = 0 .
Revised: ψ, δ (˜ ν ) (cid:105) π (cid:54) = 0 .

Original: This is easily proven by contradiction: if (cid:104)∇ ψ, δ (˜ ν ) (cid:105) π = 0 a.e.
Revised: This is easily proven by contradiction: if (cid:104)∇ ψ, δ (˜ ν ) (cid:105) π = 0 a.e.

Original: then following the proof of{+by Lemma C.
+}
Revised: then {+by Lemma C.
+}

Original: Theorem 4.2 from after (C.2) we have that Q π ˜ ν
Revised: we have that Q π ˜ ν

Original: = Q π ∗ .
Revised: = Q π ∗ .

Original: This implies that the quantity
Revised: This implies that the quantity

Original: [Equation]
Revised: [Equation]

Original: cannot vanish a.e. on Θ .
Revised: cannot vanish a.e. on Θ .

Original: Then, by Assumption 1 on the regularity of g , there exists a nonzero regular value − η of g ˜ ν
Revised: Then, by Assumption 1 on the regularity of g , there exists a nonzero regular value − η of g ˜ ν

Original: (¯ ω ) .
Revised: (¯ ω ) .

Original: Assuming without loss of generality that this regular value is negative, so that η > 0 (else invert the signs of ω 0 in the remainder of the proof), we deﬁne the nonempty sublevel set G :
Revised: Assuming without loss of generality that this regular value is negative, so that η > 0 (else invert the signs of ω 0 in the remainder of the proof), we deﬁne the nonempty sublevel set G :

Original: = { ( ω 0 , ¯ ω ) ∈ Ω : g ˜ ν
Revised: = { ( ω 0 , ¯ ω ) ∈ Ω : g ˜ ν

Original: (¯ ω ) < − η } and
Revised: (¯ ω ) < − η } and

Original: [Equation]
Revised: [Equation]

Original: Further denoting by ¯ G ⊆
Revised: Further denoting by ¯ G ⊆

Original: Θ the projection of G onto Θ , we have by deﬁnition
Revised: Θ the projection of G onto Θ , we have by deﬁnition

Original: that the gradient ﬁeld of g ˜ ν
Revised: that the gradient ﬁeld of g ˜ ν

Original: (¯ ω ) is orthogonal to the level set ∂ ¯ G , the latter being an orientable manifold of dimension m − 2 .
Revised: (¯ ω ) is orthogonal to the level set ∂ ¯ G , the latter being an orientable manifold of dimension m − 2 .

Original: Denoting by n ¯ ω the normal unit vector to ∂ ¯ G in the outward direction, by continuity of ∇ g ˜ ν
Revised: Denoting by n ¯ ω the normal unit vector to ∂ ¯ G in the outward direction, by continuity of ∇ g ˜ ν

Original: (¯ ω ) when ¯ G is compact 1 we can bound the scalar product between the two away from 0, i.e. , there exists
Revised: (¯ ω ) when ¯ G is compact 1 we can bound the scalar product between the two away from 0, i.e. , there exists

Original: [Equation]
Revised: [Equation]

Original: We now prove that the stationarity assumption in a ε -neighborhood of a spurious ﬁxed point
Revised: We now prove that the stationarity assumption in a ε -neighborhood of a spurious ﬁxed point

Revised: [Equation]

Revised: all t > t

Original: [Equation]
Revised: [Equation]

Original: leads to a contradiction for ε small enough.
Revised: leads to a contradiction for ε small enough.

Original: To do so, by Lemma C.4 we set ε ( α, η, β ) small enough so that for all ν t such that (C.10) holds we have g ν t (¯ ω ) <
Revised: To do so, by Lemma C.5 we set ε ( α, η, β ) small enough so that for all ν t such that (C.9) holds we have g ν t (¯ ω ) <

Original: − η/ 2 on ¯ G and n ¯ ω · ∇ ¯ ω g ν t > β/ 2 on ∂ ¯ G .
Revised: − η/ 2 on ¯ G and n ¯ ω · ∇ ¯ ω g ν t > β/ 2 on ∂ ¯ G .

Original: Then, the two inequalities above combined with ∂ ω 0 ψ ( ω 0 , ¯ ω ) =
Revised: Then, the two inequalities above combined with ∂ ω 0 ψ ( ω 0 , ¯ ω ) =

Original: ψ (1 , ¯ ω ) imply that the set G + deﬁned above is forward invariant and therefore that ∂ t ν t ( G + ) ≥ 0 as long as (C.10) holds.
Revised: ψ (1 , ¯ ω ) imply that the set G + deﬁned above is forward invariant and therefore that ∂ t ν t ( G + ) ≥ 0 as long as (C.9) holds.

Original: Furthermore, by similar arguments we notice that characteristic trajectories cannot enter the set G \ G + after t 1 .
Revised: Furthermore, by similar arguments we notice that characteristic trajectories cannot enter the set G \ G + after t 1 .

Original: Now, we consider two cases: either (i) a positive amount of mass is present at t 1 in the forward invariant set G + ( ν t 1 ( G + ) > 0 ) or (ii)
Revised: Now, we consider two cases: either (i) a positive amount of mass is present at t 1 in the forward invariant set G + ( ν t 1 ( G + ) > 0 ) or (ii)

Original: ν t 1 ( G + ) = 0 .
Revised: ν t 1 ( G + ) = 0 .

Original: We discuss these two cases separately, along the lines of (Chizat & Bach, 2018, Lemmas C.4, C.18), respectively.
Revised: We discuss these two cases separately, along the lines of (Chizat & Bach, 2018, Lemmas C.4, C.18), respectively.

Original: (i) Assume that ν t 1 ( G + ) > 0 .
Revised: (i) Assume that ν t 1 ( G + ) > 0 .

Original: We note that under our assumptions the ﬁrst component of the velocity ﬁeld in
Revised: We note that under our assumptions the ﬁrst component of the velocity ﬁeld in

Original: G is lower bounded by η/ 2 , so that ω 0 ( t ) = ω 0 (0) + tη/ 2 bounds from below the ω 0 -component of the trajectory of a test mass with initial condition with ω (0) ∈ G , as long as ¯ ω ( t ) ∈ ¯ G .
Revised: G is lower bounded by η/ 2 , so that ω 0 ( t ) = ω 0 (0) + tη/ 2 bounds from below the ω 0 -component of the trajectory of a test mass with initial condition with ω (0) ∈ G , as long as ¯ ω ( t ) ∈ ¯ G .

Original: Combining this bound with the forward invariance of G + , we see that if ω (0) ∈ G
Revised: Combining this bound with the forward invariance of G + , we see that if ω (0) ∈ G

Original: + then ω 0 ( t ) > tη/ 2 .
Revised: + then ω 0 ( t ) > tη/ 2 .

Original: Consequently, assuming that supp ( ν t ) ⊂ ( − M, M ) × Θ for every t > t 1 we have
Revised: Consequently, assuming that supp ( ν t ) ⊂ ( − M, M ) × Θ for every t > t 1 we have

Original: [Equation]
Revised: [Equation]

Original: This implies linear growth of h 1 ν t ( ¯ G ) for t > t 

Revised: This implies linear growth of h 1 ν t ( ¯ G ) for t > t

Original: + 2 M/η , contradicting the original assumption that 
Revised: + 2 M/η , contradicting the original assumption that (cid:107)

Original: [Equation]
Revised: [Equation]

Original: BL < 
Revised: BL < ε

Original: [Equation]

Original: for all t > t 1 .
Revised: for all t > t 1 .

Original: (ii) Consider now the complementary case ν t 1 ( G + ) = 0 .
Revised: (ii) Consider now the complementary case ν t 1 ( G + ) = 0 .

Original: We proceed to show that there exists t 2 > t 1 such that ν t 2 ( G + ) > 0 , thereby reducing this case at time t 2 to part (i).
Revised: We proceed to show that there exists t 2 > t 1 such that ν t 2 ( G + ) > 0 , thereby reducing this case at time t 2 to part (i).

Original: To do so, we consider ω ∗ ∈ supp ( ν t 1 ) such that ¯ ω ∗ ∈ ¯ G is a local minimum of g ˜ ν , i.e. , for which ∇ g ˜ ν
Revised: To do so, we consider ω ∗ ∈ supp ( ν t 1 ) such that ¯ ω ∗ ∈ ¯ G is a local minimum of g ˜ ν , i.e. , for which ∇ g ˜ ν

Original: = 0 (which exists by the preservation of the support property Assumption 1 c) ).
Revised: = 0 (which exists by the preservation of the support property Assumption 1 c) ).

Original: Then, choosing ˜ ε such that B ˜ ε (¯ ω ∗ ) ⊂ ¯ G , and setting M large enough that supp ( ν t 1 ) ⊆
Revised: Then, choosing ˜ ε such that B ˜ ε (¯ ω ∗ ) ⊂ ¯ G , and setting M large enough that supp ( ν t 1 ) ⊆

Original: [ − M, M ] × Θ , we prove below in Lemma C.3 that there exists t 2 > t 1 for which the image at t 2 of ω ( t 1 ) :
Revised: [ − M, M ] × Θ , we prove below in Lemma C.4 that there exists t 2 > t 1 for which the image at t 2 of ω ( t 1 ) :

Original: = ω ∗ under the characteristic ﬂow (C.7) is contained in G + .
Revised: = ω ∗ under the characteristic ﬂow (C.6) is contained in G + .

Original: By continuity of the ﬂow map X ( · , t ) , this conclusion extends to a neighborhood
Revised: By continuity of the ﬂow map X ( · , t ) , this conclusion extends to a neighborhood

Original: of ω
Revised: of ω

Original: ∗ , with positive mass under ν t 1 .
Revised: ∗ , with positive mass under ν t 1 .

Original: We denote throughout by (cid:107) · (cid:107) C 1 the maximum of the supremum norm of a function and the supremum norm of its gradient and recall the structure of the policy gradient vector ﬁeld
Revised: We denote throughout by (cid:107) · (cid:107) C 1 the maximum of the supremum norm of a function and the supremum norm of its gradient and recall the structure of the policy gradient vector ﬁeld

Original: [Equation]
Revised: [Equation]

Original: where g is deﬁned in (C.8) and ν t solves (10).
Revised: where g is deﬁned in (C.7) and ν t solves (10).

Original: With these deﬁnitions, we proceed to prove that case (ii) in the analysis above will ultimately reduce to case (i) for t large enough.
Revised: With these deﬁnitions, we proceed to prove that case (ii) in the analysis above will ultimately reduce to case (i) for t large enough.

Original: Lemma C.3.
Revised: Lemma C.4.

Original: Let ˜ ν ∈ M 1+ (Ω) and ¯ ω ∗ satisfy |∇ g ˜ ν
Revised: Let ˜ ν ∈ M 1+ (Ω) and ¯ ω ∗ satisfy |∇ g ˜ ν

Original: (¯ ω ∗ ) | = 0 , g ˜ ν
Revised: (¯ ω ∗ ) | = 0 , g ˜ ν

Original: (¯ ω ∗ ) < − η < 0 for some η > 0 .
Revised: (¯ ω ∗ ) < − η < 0 for some η > 0 .

Original: Then for every ˜ ε, M > 0 there exists t 2 , ε > 0 such that if for all t ∈ (0 , t 2 ) we have (cid:107) g ˜ ν − g ν t (cid:107) C 1 < ε and ω ∗ 0 ∈ [ − M, 0] , then the point ω ∗ is mapped, under the ﬂow of the policy gradient vector ﬁeld (C.11) at time t 2 to a subset of
Revised: Then for every ˜ ε, M > 0 there exists t 2 , ε > 0 such that if for all t ∈ (0 , t 2 ) we have (cid:107) g ˜ ν − g ν t (cid:107) C 1 < ε and ω ∗ 0 ∈ [ − M, 0] , then the point ω ∗ is mapped, under the ﬂow of the policy gradient vector ﬁeld (C.10) at time t 2 to a subset of

Original: B ˜ ε ((1 , ¯ ω ∗ )) .
Revised: B ˜ ε ((1 , ¯ ω ∗ )) .

Original: Proof of Lemma C.3.
Revised: Proof of Lemma C.4.

Original: By homogeneity of the approximator, we can bound the ﬁrst component of the velocity of a particle ( ω 0 ( t ) , ¯ ω ( t )) under (C.11) with initial condition ω (0) =
Revised: By homogeneity of the approximator, we can bound the ﬁrst component of the velocity of a particle ( ω 0 ( t ) , ¯ ω ( t )) under (C.10) with initial condition ω (0) =

Original: ¯ ω ∗ as
Revised: ¯ ω ∗ as

Original: [Equation]
Revised: [Equation]

Original: In the other directions, deﬁning q ( t ) :
Revised: In the other directions, deﬁning q ( t ) :

Original: = (cid:107)
Revised: = (cid:107)

Original: ¯ ω
Revised: ¯ ω

Original: ( t ) − ¯ ω ∗ (cid:107) , we have
Revised: ( t ) − ¯ ω ∗ (cid:107) , we have

Original: [Equation]
Revised: [Equation]

Original: ] for all t ∈ [0 , ¯ τ ] where ¯ τ := inf { t : ω 0 ( t ) (cid:54)∈ [ − M, 1] } .
Revised: for all t ∈ [0 , ¯ τ ] where ¯ τ := inf { t : ω 0 ( t ) (cid:54)∈ [ − M, 1] } .

Original: Moreover, Lipschitz continuity of the potential g ˜ ν ( · ) and its Lipschitz smoothness imply the existence of a L > 0 such that max {| g ˜ ν
Revised: Moreover, Lipschitz continuity of the potential g ˜ ν ( · ) and its Lipschitz smoothness imply the existence of a L > 0 such that max {| g ˜ ν

Original: (¯ ω ) − g ˜ ν
Revised: (¯ ω ) − g ˜ ν

Original: (¯ ω ∗ ) | , (cid:107)∇ g ˜ ν
Revised: (¯ ω ∗ ) | , (cid:107)∇ g ˜ ν

Original: (¯ ω ) − ∇ g ˜ ν
Revised: (¯ ω ) − ∇ g ˜ ν

Original: (¯ ω ∗ ) (cid:107)} ≤ L (cid:107) ¯ ω − ¯ ω ∗ (cid:107) .
Revised: (¯ ω ∗ ) (cid:107)} ≤ L (cid:107) ¯ ω − ¯ ω ∗ (cid:107) .

Original: Combining this with the assumed convergence of ν t to ˜ ν , which implies (cid:107) g ˜ ν
Revised: Combining this with the assumed convergence of ν t to ˜ ν , which implies (cid:107) g ˜ ν

Original: − g ν t (cid:107) C 1 < ε , we can bound the evolution of ( ω 0 ( t ) , q ( t )) for t ∈ [0 , ¯ τ ] in the perturbative regime of interest as follows:[Equation]
Revised: − g ν t (cid:107) C 1 < ε , we can bound the evolution of ( ω 0 ( t ) , q ( t )) for t ∈ [0 , ¯ τ ] in the perturbative regime of interest as follows:

Original: [Equation]
Revised: [Equation]

Original: We now show that, choosing both ε and a neighborhood around ω ∗ =
Revised: We now show that, choosing both ε and a neighborhood around ω ∗ =

Original: ( ω ∗ 0 , ¯ ω ∗ ) to be small enough, the forward dynamics of ω ∗ will reach the set { ω 0 > 0 } before q ( t ) can increase too much.
Revised: ( ω ∗ 0 , ¯ ω ∗ ) to be small enough, the forward dynamics of ω ∗ will reach the set { ω 0 > 0 } before q ( t ) can increase too much.

Original: More precisely, by possibly increasing the value of L such that η/ 4 L < ˜ ε , and deﬁning τ q = inf { t : q ( t )
Revised: More precisely, by possibly increasing the value of L such that η/ 4 L < ˜ ε , and deﬁning τ q = inf { t : q ( t )

Original: > η/ 4 L } we prove that there exists ε ∈ (0 , η/ 4) such that τ q > ¯ τ , i.e. , that the trajectory ω ∗ reaches G + before q ( t ) >
Revised: > η/ 4 L } we prove that there exists ε ∈ (0 , η/ 4) such that τ q > ¯ τ , i.e. , that the trajectory ω ∗ reaches G + before q ( t ) >

Original: ˜ ε .
Revised: ˜ ε .

Original: Note that as long as t ∈ [0 , τ q ] and ε ∈ (0 , η/ 4) the negative terms on the RHS (C.12) can be bounded from below, and we have
Revised: Note that as long as t ∈ [0 , τ q ] and ε ∈ (0 , η/ 4) the negative terms on the RHS (C.11) can be bounded from below, and we have

Original: [Equation]
Revised: [Equation]

Original: so that ω 0 ( t ) > ω 0 (0) ≥ − M .
Revised: so that ω 0 ( t ) > ω 0 (0) ≥ − M .

Original: Consequently, for all t ∈ [0 , ¯ τ ∧ τ q ] we bound the RHS of (C.13) as dd t q ( t ) <
Revised: Consequently, for all t ∈ [0 , ¯ τ ∧ τ q ] we bound the RHS of (C.12) as dd t q ( t ) <

Original: Mε + LMq ( t ) .
Revised: Mε + LMq ( t ) .

Original: Using that q (0) = 

Revised: Using that q (0) =

Original: and Gr¨onwall inequality we can bound the total excursion in the ¯ ω component q ( t ) ≤ εMt exp [ LMt ] .
Revised: and Gr¨onwall inequality we can bound the total excursion in the ¯ ω component q ( t ) ≤ εMt exp [ LMt ] .

Original: Finally, setting τ 0 :
Revised: Finally, setting τ 0 :

Original: = 

Revised: =

Original: ( M + 1) /η ≥ − 2( ω 0 (0) − 1)
Revised: ( M + 1) /η ≥ − 2( ω 0 (0) − 1)

Original: /η >
Revised: /η >

Original: ¯ τ so that ω 0 ( τ 0 ) > 1 we are still free to set ε small enough such that τ q > τ 0 > ¯ τ .
Revised: ¯ τ so that ω 0 ( τ 0 ) > 1 we are still free to set ε small enough such that τ q > τ 0 > ¯ τ .

Original: Indeed, by monotonicity of the upper-bound on q ( t ) we have
Revised: Indeed, by monotonicity of the upper-bound on q ( t ) we have

Original: [Equation]
Revised: [Equation]

Original: so that setting ε ∈ (0 , η/ 4) concludes the proof.
Revised: so that setting ε ∈ (0 , η/ 4) concludes the proof.

Original: We now proceed to show the needed regularity of the potential g ν from (C.8) in terms of the signed measure h 1 ν deﬁned in (C.6).
Revised: We now proceed to show the needed regularity of the potential g ν from (C.7) in terms of the signed measure h 1 ν deﬁned in (C.5).

Original: In doing so, we also prove Lipschitz smoothness of the operator R deﬁned in (B.1):
Revised: In doing so, we also prove Lipschitz smoothness of the operator R deﬁned in (B.1):

Original: Lemma C.4.
Revised: Lemma C.5.

Original: For any r > 0 , the operator R on F r = { (cid:82) ψν : supp ν ∈ Q r } is Lipschitz smooth, and DR f is bounded in the supremum norm.
Revised: For any r > 0 , the operator R on F r = { (cid:82) ψν : supp ν ∈ Q r } is Lipschitz smooth, and DR f is bounded in the supremum norm.

Original: Furthermore, for all C 0 > 0 there exists α > 0 and ε > 0 such that for all ν, ν (cid:48) satisfying (cid:107) h1 ν (cid:107) BL , (cid:107) h 1 ν (cid:48) (cid:107) BL < C 0 , (cid:107) h
Revised: Furthermore, for all C 0 > 0 there exists α > 0 and ε > 0 such that for all ν, ν (cid:48) satisfying (cid:107) h

Original: [Equation]
Revised: [Equation]

Original: h
Revised: BL

Original: [Equation]
Revised: [Equation]

Revised: BL < C 0 ,

Original: (cid:48) (cid:107)
Revised: [Equation]

Original: BL < ε , one has
Revised: BL < ε , one has

Original: [Equation]
Revised: [Equation]

Original: To prove the above lemma, we ﬁrst bound some relevant quantities.
Revised: To prove the above lemma, we ﬁrst bound some relevant quantities.

Original: Throughout, by slight abuse of notation, we denote for any function f : S ×
Revised: Throughout, by slight abuse of notation, we denote for any function f : S ×

Original: A → R
Revised: A → R

Original: [Equation]
Revised: [Equation]

Original: Lemma C.5.
Revised: Lemma C.6.

Original: For all f, f (cid:48) ∈ F r there exists ε > 0 , C (cid:48) , C (cid:48)(cid:48) , C[-(cid:48)(cid:48)(cid:48) > 
-]
Revised: For all f, f (cid:48) ∈ F r there exists ε > 0 , C (cid:48) , C [Equation] , C[-(cid:48)(cid:48)(cid:48) > 
-]

Original:  such that if (cid:107) f − f (cid:48) (cid:107) 2 < ε one has (cid:107) π ν − π ν (cid:48) (cid:107) 2 ≤ C (cid:48) (cid:107) f − f (cid:48)
Revised: [Equation] > 0 such that if (cid:107) f − f (cid:48) (cid:107) 2 < ε one has [Equation]

Original: (cid:107)(C.15) (cid:107) (cid:37) ν − (cid:37) ν (cid:48) (cid:107) 1 ≤ C (cid:48)(cid:48)
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Proof of Lemma C.5.
Revised: Proof of Lemma C.6.

Original: Throughout this proof, for simplicity of notation, we will write π = π ν and π (cid:48)
Revised: Throughout this proof, for simplicity of notation, we will write π = π ν and π (cid:48)

Original: =
Revised: =

Original: π ν (cid:48) .
Revised: π ν (cid:48) .

Original: Furthermore, we use that for f ∈ F r there exists C 0 > 

Revised: Furthermore, we use that for f ∈ F r there exists C 0 >

Original: so that
Revised: so that

Original: [Equation]
Revised: [Equation]

Original: implying, together with the assumed absolute continuity of (cid:37) 

Revised: implying, together with the assumed absolute continuity of (cid:37)

Original: that (cid:107) Q π (cid:107) ∞ , (cid:107) π (cid:107) ∞ , (cid:107) (cid:37) π
Revised: that (cid:107) Q π (cid:107) ∞ , (cid:107) π (cid:107) ∞ , (cid:107) (cid:37) π

Original: (cid:107) ∞ < ∞ .
Revised: (cid:107) ∞ < ∞ .

Original: Setting throughout τ = 1 to simplify the notation and combining the above with the pointwise upper bound e x < 1 + K r | x | for | x | < e C [Equation] φ (cid:107) C [-
-]
Revised: Setting throughout τ = 1 to simplify the notation and combining the above with the pointwise upper bound e x < 1 + K r | x | for | x | < e C 0 (cid:107) φ (cid:107) C [-
-]1 we obtain

Original: we obtain

Original: [Equation]
Revised: [Equation]

Original: [Table][Equation] [Table] [Equation]
Revised: [Table]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]
Revised: [Equation]

Original: where we have denoted by |A| the Lebesgue measure of the action space A .
Revised: where we have denoted by |A| the Lebesgue measure of the action space A .

Original: We now proceed to establish the second bound in the statement of the lemma.
Revised: We now proceed to establish the second bound in the statement of the lemma.

Original: In this case, denoting the t -steps transition probability as
Revised: In this case, denoting the t -steps transition probability as

Original: Ptπ [Equation] s, d s t ) =
Revised: P

Original: [Equation]
Revised: [Equation]

Original: 1 , d s t ) we have[Equation] [Table] [Equation]
Revised: ) we have

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]
Revised: [Equation]

Original: Observing that for any smooth (cid:37) ∈ M 1+ ( S ) for the operator norm of the difference in the above sum we have
Revised: Observing that for any smooth (cid:37) ∈ M 1+ ( S ) for the operator norm of the difference in the above sum we have

Original: [Equation]
Revised: [Equation]

Original: for large enough C (cid:48)(cid:48) , where we used (C.19) in the last line and the Lipschitz continuity of P in its second argument, and (cid:107) P (cid:107) 1 is the operator norm of the transition operator (cid:82) A P ( s, a, d s (cid:48) )
Revised: for large enough C (cid:48)(cid:48) , where we used (C.18) in the last line and the Lipschitz continuity of P in its second argument, and (cid:107) P (cid:107) 1 is the operator norm of the transition operator (cid:82) A P ( s, a, d s (cid:48) )

Original: π (cid:48)(cid:48) ( s, d a ) :
Revised: π (cid:48)(cid:48) ( s, d a ) :

Original: M 1+ ( A ) →
Revised: M 1+ ( A ) →

Original: M 1+ ( A ) , which is equal to 1 .
Revised: M 1+ ( A ) , which is equal to 1 .

Original: From this we conclude
Revised: From this we conclude

Original: [Equation]
Revised: [Equation]

Original: Finally, deﬁning for notational convenience R (cid:48)
Revised: Finally, deﬁning for notational convenience R (cid:48)

Original: τ := ¯ r − τD KL ( π (cid:48) , ¯ π ) we write:
Revised: τ := ¯ r − τD KL ( π (cid:48) , ¯ π ) we write:

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Revised: [Equation]

Revised: ∞ (C.

Original: [Equation]
Revised: [Equation]

Original: and bound each term separately letting C (cid:48)(cid:48)(cid:48) 1 , C (cid:48)(cid:48)(cid:48) 2 , C (cid:48)(cid:48)(cid:48) 3 > 0 be large enough constants.
Revised: and bound each term separately letting C (cid:48)(cid:48)(cid:48) 1 , C (cid:48)(cid:48)(cid:48) 2 , C (cid:48)(cid:48)(cid:48) 3 > 0 be large enough constants.

Original: For the ﬁrst
Revised: For the ﬁrst, we have: (cid:13)(cid:13)(cid:13)(cid:13)(cid:90) ( R τ − R (cid:48) τ ) ( s (cid:48)(cid:48) , a

Original: [Equation]
Revised: [Equation]

Original: we have:
Revised: ) (cid:37) π ( s (cid:48) [Equation] d s (cid:48)(cid:48) ) π ( s (cid:48)(cid:48)

Original: [Equation]
Revised: [Equation]

Original: where we have bounded the log term as follows:
Revised: d

Original: [Equation]
Revised: [Equation]

Original: log
Revised: ) (cid:13)(cid:13)(cid:13)(cid:13) ∞ ≤

Original: [Equation]
Revised: [Equation]

Revised: where we have bounded the log term as follows:

Original: [Table] [Equation] 
Revised: [Equation] [Table]

Original: − [Table]
Revised: [Equation]

Original: π
Revised: C.

Original: [Table]
Revised: [Table]

Original: = γ [Equation] − [Equation]
Revised: For the second term in (C.19), using the boundedness of (cid:107) R (cid:107) 2 < C R and that (cid:37) π − (cid:37) π

Original: [Table]
Revised: [Table]

Original:  ) for a certain, (cid:37) 

Revised: = γ ( (cid:37) π − (cid:37) π (cid:48) ) for a certain, (cid:37)

Original: depending on s [Equation]
Revised: depending on s (cid:48)

Original: we write
Revised: we write

Original: [Equation]
Revised: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]
Revised: [Table]

Original: We ﬁnally bound the third term by writing
Revised: We ﬁnally bound the third term by writing

Original: [Table]
Revised: [Table]

Original: [Equation]

Original: C [Equation] [Equation] and obtain (C.17) by combining (C.21)-(C.24).
Revised: and obtain (C.16) by combining (C.20)-(C.23).

Original: Proof of Lemma C.4.
Revised: Proof of Lemma C.5.

Original: We ﬁrst establish the desired properties of the functional R .
Revised: We ﬁrst establish the desired properties of the functional R .

Original: To do so we differentiate (8) at f ∈ L 2 ( S × A )
Revised: To do so we differentiate (8) at f ∈ L 2 ( S × A )

Original: [Equation]
Revised: [Equation]

Original: Then, combining the above with (A.6) and H ¨ older inequality we obtain
Revised: Then, combining the above with (A.6) and H ¨ older inequality we obtain

Original: [Equation]
Revised: [Equation]

Original: where we have used that all the terms appearing in DR f are bounded for every choice of r > 0 .
Revised: where we have used that all the terms appearing in DR f are bounded for every choice of r > 0 .

Original: To establish Lipschitz smoothness of R , letting f, f
Revised: To establish Lipschitz smoothness of R , letting f, f

Original: (cid:48) ∈ F r and denoting, to simplify notation, π = π f and π (cid:48) =
Revised: (cid:48) ∈ F r and denoting, to simplify notation, π = π f and π (cid:48) =

Original: π f (cid:48) , we proceed to bound the operator norm by splitting the RHS as
Revised: π f (cid:48) , we proceed to bound the operator norm by splitting the RHS as

Original: [Equation]
Revised: [Equation]

Original: and considering the resulting terms separately
Revised: and considering the resulting terms separately.

Original: [Equation]

Original: First of all, deﬁning throughout by slight abuse of notation δ ( π ) :
Revised: First of all, deﬁning throughout by slight abuse of notation δ ( π ) :

Original: = Q π − τ log π , we have
Revised: = Q π − τ log π , we have

Original: [Equation]
Revised: [Equation]

Revised: C.

Original: [Equation]
Revised: [Equation]

Original: where the last line was obtained using (C.15) and boundedness from above and below of π [Equation] Q π in F r .
Revised: where the last line was obtained using (C.14) and boundedness from above and below of π, Q π in F r .

Original: We further write, for a K < ∞ large enough
Revised: We further write, for a K < ∞ large enough

Original: [Equation]
Revised: [Equation]

Original: where we have used Cauchy-Schwartz inequality together with (C.15) and (C.16).
Revised: where we have used Cauchy-Schwartz inequality together with (C.14) and (C.15).

Original: Finally [Equation] using (C.17) and (C.22)
Revised: Finally, using (C.16) and (C.21)

Original: [Equation]
Revised: [Equation]

Original: we bound for K < ∞ possibly larger than above
Revised: we bound for K < ∞ possibly larger than above

Original: [Equation]
Revised: [Equation]

Original: III

Original: [Equation]

Original: Combining (C.26), (C.27) and (C.28) we obtain that
Revised: Combining (C.25), (C.26) and (C.27) we obtain that

Original: [Equation]
Revised: [Equation]

Original: proving the Lipschitz smoothness claim. Proceeding to the proof of (C.14)
Revised: DR f − DR f (cid:48) (cid:107) (cid:107) f − f (cid:48) (cid:107)

Original: [Equation]
Revised: [Equation]

Revised: < ∞ proving the Lipschitz smoothness claim.

Original:  combining the Lipschitz smoothness of R on bounded sets [Equation] the identity ψ ( s, a ; (1 , ¯ ω ))
Revised: Proceeding to the proof of (C.13), combining the Lipschitz smoothness of R on bounded sets, the identity ψ ( s, a ; (1 , ¯ ω ))

Original: =
Revised: =

Original: φ ( s, a ; ¯ ω ) and the boundedness of the set { (cid:82) ψν : ν ∈ P 2 (Ω) [Equation] | h 1 ν | < C 0 } we obtain 
Revised: φ ( s, a ; ¯ ω ) and the boundedness of the set { (cid:82) ψν : ν ∈ P 2 (Ω) , | h 1 ν | < C 0 } we obtain (cid:107)

Original: [Equation]
Revised: [Equation]

Original: which combined with the (cid:107) φ (cid:107) C 1 -Lipschitz continuity of the map (cid:82) (cid:96) ( s, a ) φ ( s, a ; ¯ ω )d s d a and with the Lipschitz smoothness of R concludes the proof.
Revised: which combined with the (cid:107) φ (cid:107) C 1 -Lipschitz continuity of the map (cid:82) (cid:96) ( s, a ) φ ( s, a ; ¯ ω )d s d a and with the Lipschitz smoothness of R concludes the proof.