Wasserstein Flow Meets Replicator Dynamics: A Mean-Field Analysis of Representation Learning in Actor-Critic

We consider a Markov decision process (MDP) given by $({\mathcal{S}},\mathcal{A},\gamma,P,r,\mathcal{D}_{0})$, where ${\mathcal{S}}\subseteq\mathbb{R}^{d_{1}}$ is the state, $\mathcal{A}\subseteq\mathbb{R}^{d_{2}}$ is the action space, $\gamma\in(0,1)$ is the discount factor, $P:{\mathcal{S}}\times\mathcal{A}\rightarrow\mathscr{P}({\mathcal{S}})$ is the transition kernel, $r:{\mathcal{S}}\times\mathcal{A}\rightarrow\mathbb{R}_{+}$ is the reward function, and $\mathcal{D}_{0}\in\mathscr{P}({\mathcal{S}})$ is the initial state distribution. Without loss of generality, we assume that ${\mathcal{S}}\times\mathcal{A}\subseteq\mathbb{R}^{d}$ and $\|(s,a)\|_{2}\leq 1$, where $d=d_{1}+d_{2}$. We remark that as long as the state-action space is bounded, we can normalize the space to be within the unit sphere. Given a policy $\pi:{\mathcal{S}}\times\mathcal{A}\rightarrow\mathscr{P}({\mathcal{S}})$, at the $m$th step, the agent takes an action $a_{m}$ at state $s_{m}$ according to $\pi(\cdot{\,|\,}s_{m})$ and observes a reward $r_{m}=r(s_{m},a_{m})$. The environment then transits to the next state $s_{m+1}$ according to the transition kernel $P(\cdot{\,|\,}s_{m},a_{m})$. Note that the policy $\pi$ induces Markov chains on both ${\mathcal{S}}$ and ${\mathcal{S}}\times\mathcal{A}$. Considering the Markov chain on ${\mathcal{S}}$, we denote the induced Markov transition kernel by $P^{\pi}:{\mathcal{S}}\rightarrow\mathscr{P}({\mathcal{S}})$, which is defined as $P^{\pi}(s^{\prime}{\,|\,}s)=\int_{\mathcal{A}}P(s^{\prime}{\,|\,}s,a)\pi({\mathrm{d}}a{\,|\,}s)$. Likewise, we denote the Markov transition kernel on ${\mathcal{S}}\times\mathcal{A}$ by $\widetilde{P}^{\pi}:{\mathcal{S}}\times\mathcal{A}\rightarrow\mathscr{P}({\mathcal{S}}\times\mathcal{A})$, which is defined as $\widetilde{P}^{\pi}(s^{\prime},a^{\prime}{\,|\,}s,a)=\pi(a^{\prime}{\,|\,}s^{\prime})P(s^{\prime}{\,|\,}s,a)$. Let ${\widetilde{\mathcal{D}}}$ be a probability measure on ${\mathcal{S}}\times\mathcal{A}$. We then define the visitation measure induced by policy $\pi$ and starting from ${\widetilde{\mathcal{D}}}$ as

$\displaystyle{\widetilde{\mathcal{E}}}^{\pi}_{{\widetilde{\mathcal{D}}}}\big{(}{\mathrm{d}}(s,a)\big{)}=(1-\gamma)\cdot\sum_{m\geq 0}\gamma^{m}\cdot\mathbb{P}\big{(}(s_{m},a_{m})\in{\mathrm{d}}(s,a){\,|\,}(s_{0},a_{0})\sim{\widetilde{\mathcal{D}}}\big{)},$

(2.1)

where $(s_{m},a_{m})$ is the trajectory generated by starting from $(s_{0},a_{0})\sim{\widetilde{\mathcal{D}}}$ and following policy $\pi$ thereafter. If ${\widetilde{\mathcal{D}}}=\mathcal{D}\otimes\pi$ holds, we then denote such a visitation measure by $\widetilde{\mathcal{E}}_{\mathcal{D}}^{\pi}$. Furthermore, we denote by $\mathcal{E}({\mathrm{d}}s)=\int_{\mathcal{A}}{\widetilde{\mathcal{E}}}({\mathrm{d}}s,{\mathrm{d}}a)$ the marginal distribution of visitation measure ${\widetilde{\mathcal{E}}}$ with respect to ${\mathcal{S}}$. In particular, when $(s_{0},a_{0})\sim\mathcal{D}\otimes\pi$ holds in (2.1), it follows that ${\widetilde{\mathcal{E}}}_{\mathcal{D}}^{\pi}=\mathcal{E}_{\mathcal{D}}^{\pi}\otimes\pi$. In policy optimization, we aim to maximize the expected total reward $J(\pi)$ defined as follows,

$\displaystyle J(\pi)=\mathbb{E}^{\pi}\biggl{[}\sum_{m\geq 0}\gamma^{m}\cdot r(s_{m},a_{m}){\,\big{|}\,}s_{0}\sim\mathcal{D}_{0}\biggr{]},$

where we denote by $\mathbb{E}^{\pi}$ the expectation with respect to $a_{m}\sim\pi(\cdot{\,|\,}s_{m})$ and $s_{m+1}\sim P(\cdot{\,|\,}s_{m},a_{m})$ for $m\geq 0$. We define the action value function $Q^{\pi}:{\mathcal{S}}\times\mathcal{A}\rightarrow\mathbb{R}$ and the state value function $V^{\pi}:{\mathcal{S}}\rightarrow\mathbb{R}$ as follows,

$\displaystyle Q^{\pi}(s,a)=\mathbb{E}^{\pi}\biggl{[}\sum_{m\geq 0}\gamma^{m}\cdot r(s_{m},a_{m}){\,\big{|}\,}s_{0}=s,a_{0}=a\biggr{]},\quad V^{\pi}(s)=\big{\langle}Q^{\pi}(s,\cdot),\pi(\cdot{\,|\,}s)\big{\rangle}_{\mathcal{A}},$

(2.2)

where we denote by $\langle\cdot,\cdot\rangle_{\mathcal{A}}$ the inner product on the action space $\mathcal{A}$. Correspondingly, the advantage function $A^{\pi}:{\mathcal{S}}\times\mathcal{A}\rightarrow\mathbb{R}$ is defined as

$\displaystyle A^{\pi}(s,a)=Q^{\pi}(s,a)-V^{\pi}(s).$

The action value function $Q^{\pi}$ corresponds to the minimizer of the following mean-squared Bellman error (MSBE),

$\displaystyle\mathrm{MSBE}(Q;\pi)=\frac{1}{2}\mathbb{E}_{(s,a)\sim\widetilde{\Phi}^{\pi}}\Bigl{[}\bigl{(}Q(s,a)-r(s,a)-\gamma\mathbb{E}_{(s^{\prime},a^{\prime})\sim\widetilde{P}^{\pi}(\cdot{\,|\,}s,a)}[Q(s^{\prime},a^{\prime})]\bigr{)}^{2}\Bigr{]},$

(2.3)

where $\widetilde{\Phi}^{\pi}$ is the sampling distribution depending on policy $\pi$, which will be defined by (4.8) in §4.2. Note that when $\widetilde{\Phi}^{\pi}$ is with full support, i.e., $\mathrm{supp}(\widetilde{\Phi}^{\pi})={\mathcal{S}}\times\mathcal{A}$, $Q^{\pi}$ is the unique global minimizer to the MSBE. Consequently, the policy optimization problem can be written as the following bilevel optimization problem,

$\displaystyle\max_{\pi}J(\pi)=\mathbb{E}_{s\sim\mathcal{D}_{0}}\Bigl{[}\big{\langle}Q^{\pi}(s,\cdot),\pi(\cdot{\,|\,}s)\big{\rangle}_{\mathcal{A}}\Bigr{]},\quad\text{subject to }Q^{\pi}=\mathop{\mathrm{argmin}}_{Q}\text{MSBE}(Q;\pi).$

(2.4)

The inner problem in (2.4) is known as a policy evaluation subproblem, while the outer problem is the policy improvement subproblem. One of the most popular way to solve the policy optimization problem is actor-critic (AC) methods (Sutton and Barto, 2018), where the job of the critic is to evaluate current policy and then the actor updates its policy according to the critic’s evaluation.

Let $\Theta\subseteq\mathbb{R}^{D}$ be a Polish space. We denote by $\mathscr{P}_{2}(\Theta)\subseteq\mathscr{P}(\Theta)$ the set of probability measures with finite second moments. Then, the Wasserstein-2 distance between $\mu,\nu\in\mathscr{P}_{2}(\Theta)$ is defined as follows,

$\displaystyle W_{2}(\mu,\nu)=\inf\Bigl{\{}\mathbb{E}\bigl{[}\|X-Y\|^{2}\bigr{]}^{1/2}{\,\Big{|}\,}\mathop{\mathrm{law}}(X)=\mu,\mathop{\mathrm{law}}(Y)=\nu\Bigr{\}},$

where the infimum is taken over the random variables $X$ and $Y$ on $\Theta$ and we denote by ${\rm law}(X)$ the distribution of a random variable $X$. We call $\mathcal{M}=(\mathscr{P}_{2}(\Theta),W_{2})$ the Wasserstein ($W_{2}$) space, which is an infinite-dimensional manifold (Villani, 2008). See §A.1 for more details.

For the MF-PPO update in (3.7), we establish the following theorem on global optimality and convergence rate.

The concentrability coefficient commonly appears in the reinforcement learning literature (Szepesvári and Munos, 2005; Munos and Szepesvári, 2008; Antos et al., 2008; Farahmand et al., 2010; Scherrer et al., 2015; Farahmand et al., 2016; Lazaric et al., 2016; Liu et al., 2019; Wang et al., 2019). In contrast to a more standard concentrability coefficient form, note that $\kappa$ is irrelevant to the update of the algorithm. To show the convergence of the MF-PPO, our condition here is much weaker since we only need to specify a base distribution $\widetilde{\phi}_{0}$ such that $\kappa<\infty$.

Theorem 4.1 shows that the MF-PPO converges to the globally optimal policy at a rate of $O(T^{-1})$ up to the policy evaluation error. Such a theorem implies the global optimality and convergence of a double-loop AC algorithm, where the critic $Q_{t}$ is solved to high precision and the policy evaluation error is sufficiently small. In the sequel, we consider a more challenging setting, where the critic $Q_{t}$ is updated simultaneously along with the update of the actor’s policy $\pi_{t}$.

In this section, we aim to provide an upper bound on the policy evaluation error when the critic and the actor are updated simultaneously. Specifically, the actor is updated via MF-PPO in (3.7) and the critic $Q_{t}=Q(\cdot;\rho_{t})$ is updated via the MF-TD in (3.5). The smooth function $\sigma$ in the parameterization of the Q function in (3.4) is taken to be the following two-layer neural network,

$\displaystyle\sigma(s,a;\theta)=B_{\beta}\cdot\beta(b)\cdot\widetilde{\sigma}\bigl{(}w^{\top}(s,a,1)\bigr{)},$

(4.2)

where $\widetilde{\sigma}:\mathbb{R}\rightarrow\mathbb{R}$ is the activation function, $\theta=(b,w)$ is the parameter, and $\beta:\mathbb{R}\rightarrow(-1,1)$ is an odd and invertible function with scaling hyper-parameter $B_{\beta}>0$. It then holds that $D=d+2$, where $d$ and $D$ are the dimensions of $(s,a)$ and $\theta$, respectively. It is worth noting that the function class of $\int\sigma(s,a;\theta)\rho({\mathrm{d}}\theta)$ for $\rho\in\mathscr{P}_{2}(\mathbb{R}^{D})$ is the same as

$\displaystyle\mathcal{F}=\Bigl{\{}\int\beta^{\prime}\cdot\widetilde{\sigma}\bigl{(}w^{\top}(s,a,1)\bigr{)}\nu({\mathrm{d}}\beta^{\prime},{\mathrm{d}}w){\,\Big{|}\,}\nu\in\mathscr{P}_{2}\bigl{(}(-B_{\beta},B_{\beta})\times\mathbb{R}^{d+1}\bigr{)}\Bigr{\}},$

(4.3)

which captures a vast function class because of the universal function approximation theorem (Barron, 1993; Pinkus, 1999). We remark that we introduce the rescaling function $\beta$ in (4.2) to avoid the study of the space of probability measures over $(-B_{\beta},B_{\beta})\times\mathbb{R}^{d+1}$ in (4.3), which has a boundary and thus lacks the regularity in the study of optimal transport. Furthermore, note that we introduce a hyper-parameter $\alpha>1$ in the Q function in (3.4). Thus, we are using $\alpha\cdot\mathcal{F}$ to represent $\mathcal{F}$, which causes an “over-representation” when $\alpha>1$. Such over-representation appears to be essential for our analysis. To anticipate briefly, we note that $\alpha$ controls the gap in the average total reward over time when the relative time-scale $\eta$ is properly selected according to Theorem 4.6. Furthermore, such an influence is imposed through Lemma 4.4, which shows that the Wasserstein distance between $\rho_{0}$ and $\rho_{\pi_{t}}$ is upper bounded by $O(1/\alpha)$. In what follows, we consider the initialization of the TD update to be $\rho_{0}=N(0,I_{D})$, which implies that $Q(s,a;\rho_{0})=0$. We next impose a regularity assumption on the two-layer neural network $\sigma$.

We remark that Assumption 4.2 is not restrictive and is satisfied by a large family of neural networks, e.g., $\widetilde{\sigma}(x)=\tanh(x)$ and $\beta(b)=\tanh(b)$. Noting that $\|(s,a)\|_{2}\leq 1$, Assumption 4.2 implies that the function $\sigma(s,a;\theta)$ in (4.2) is odd with respect to $w$ and $b$ and is also bounded, Lipschitz continuous, and smooth in the parameter domain, that is,

$\displaystyle|\nabla_{\theta}\sigma(s,a;\theta)|<B_{1},\quad|\nabla^{2}_{\theta\theta}\sigma(s,a;\theta)|<B_{2}.$

(4.4)

We then impose the following assumption on the MDP.

We remark that by assuming $\psi$ to be a probability measure and that $\varphi(s^{\prime})\geq 0$ in (4.6), the representation of the transition kernel does not lose generality. Specifically, the function class of (4.6) is the same as

$\displaystyle\mathcal{P}=\Big{\{}\int\widetilde{\sigma}((s,a,1)^{\top}w)\widetilde{\psi}(s^{\prime};{\mathrm{d}}w)\,\Big{|}\,\widetilde{\psi}(s^{\prime};\cdot)\text{ is a signed measure for any $s^{\prime}\in{\mathcal{S}}$}\Big{\}}.$

See §C.1 for a detailed proof. Assumption 4.3 generalizes the linear MDP in Yang and Wang (2019b, a); Cai et al. (2019a) and Jin et al. (2020). In contrast, our representation of the reward function and the transition kernel benefits from the universal function approximation theorem and is thus not as restrictive as the original linear MDP assumption. Note that the infinite-width neural network has a two-layer structure by $\eqref{eq:nn}$. We establish the following lemma on the regularity of the representation of the action value function $Q^{\pi}$ by such a neural network.

Property (i) of Lemma 4.4 shows that the action value function $Q^{\pi}$ can be parameterized with the infinite-width two-layer neural network $Q(\cdot;\rho_{\pi})$ in (3.4). Note that a larger $B_{\beta}$ captures a larger function class in (4.3). Without loss of generality, we assume that $B_{\beta}\geq 2(B_{r}+\gamma(1-\gamma)^{-1}B_{r}M_{1,\varphi})$ holds in the sequel. Hence, by Property (ii), we have that $\widetilde{W}_{2}(\rho_{\pi},\rho_{0})\leq O(1)$ for any policy $\pi$. In particular, it holds by Property (i) of Lemma 4.4 that $\|Q_{t}-Q^{\pi_{t}}\|_{2,\widetilde{\phi}^{\pi_{t}}}=\|Q(\cdot;\rho_{t})-Q(\cdot;\rho_{\pi_{t}})\|_{2,\widetilde{\phi}^{\pi_{t}}}$ and we have the following theorem to characterize such an error with regard to the $W_{2}$ space.

Here we give a nonrigorous discussion on how to upper bound $\Delta_{t}$ in (4.7). If $\widetilde{W}_{2}(\rho_{t},\rho_{0})\leq O(1)$ holds for any $t\in[0,T]$, by $\widetilde{W}_{2}(\rho_{\pi_{t}},\rho_{0})\leq O(1)$ in Lemma 4.4 and the triangle inequality of $W_{2}$ distance (Villani, 2008), it follows that $\widetilde{W}_{2}(\rho_{t},\rho_{\pi_{t}})\leq O(1)$ and $\Delta_{t}\leq O(\alpha^{1/2}\eta^{-1}+\alpha^{-1})$. Taking a time average of integration on both sides of (4.7), the policy evaluation error $\frac{1}{T}\int_{0}^{T}\|Q_{t}-Q^{\pi_{t}}\|_{2,\widetilde{\phi}^{\pi_{t}}}{\mathrm{d}}t$ is then upper bounded by $O(\eta^{-1}T^{-1}+\alpha^{1/2}\eta^{-1}+\alpha^{-1})$. Inspired by such a fact, we introduce the following restarting mechanism to ensure $\widetilde{W}_{2}(\rho_{t},\rho_{0})\leq O(1)$. Restarting Mechanism. Let $\widetilde{W}_{0}=\lambda\bar{D}$ be a threshold, where $\bar{D}$ is the upper bound for $\widetilde{W}_{2}(\rho_{\pi},\rho_{0})$ by Lemma 4.4, $\lambda\geq 3$ is a constant scaling parameter for the restarting threshold, $\rho_{t}$ is the distribution of the parameters in the neural network at time $t$, and $\rho_{0}$ is the initial distribution. Whenever we detect that $\widetilde{W}_{2}(\rho_{t},\rho_{0})$ reaches $\widetilde{W}_{0}$ in the update, we pause and reset $\rho_{t}$ to $\rho_{0}$ by resampling the parameters from $\rho_{0}$. Then, we reset the critic with the newly sampled parameters while keeping the actor’s policy $\pi_{t}$ unchanged and continue the update.

The restarting mechanism guarantees $\widetilde{W}_{2}(\rho_{t},\rho_{0})\leq\lambda\bar{D}$ by restricting the distribution $\rho_{t}$ of the parameters to be close to $\rho_{0}$. Moreover, by letting $\lambda\geq 3$, we ensure that $\rho_{\pi_{t}}$ is realizable by $\rho_{t}$ since $\widetilde{W}_{2}(\rho_{\pi_{t}},\rho_{0})\leq\bar{D}\leq\lambda\bar{D}$, which means that the neural network is capable of capturing the representation of the action value function $Q^{\pi_{t}}$. We remark that by letting $\widetilde{W}_{0}=O(1)$, we allow $\rho_{t}$ to deviate from $\rho_{0}$ up to $W_{2}(\rho_{t},\rho_{0})\leq O(\alpha^{-1})$ in the restarting mechanism. In contrast, the NTK regime (Cai et al., 2019b) which corresponds to letting $\alpha=\sqrt{M}$ in (3.1) only allows $\rho_{t}$ to deviate from $\rho_{0}$ by the chi-squared divergence $\chi^{2}(\rho_{t}\,\|\,\rho_{0})\leq O(M^{-1})=o(1)$. That is, the NTK regime fails to induce a feature representation significantly different from the initial one. Before moving on, we summarize the construction of the weighting distribution $\widetilde{\Phi}^{\pi_{t}}$ in Theorem 4.1 and 4.5 as follows,

$\displaystyle\widetilde{\Phi}^{\pi_{t}}={\widetilde{\mathcal{E}}}_{\widetilde{\phi}^{\pi_{t}}}^{\pi_{t}},\qquad\widetilde{\phi}^{\pi_{t}}=\frac{1}{2}\widetilde{\phi}_{0}+\frac{1}{2}\phi_{0}\otimes\pi_{t},\qquad\phi_{0}=\int_{\mathcal{A}}\widetilde{\phi}_{0}(\cdot,{\mathrm{d}}a),$

(4.8)

where $\widetilde{\phi}_{0}$ is the base distribution. Now we have the following theorem that characterizes the global optimality and convergence of the two-timescale AC with restarting mechanism.