Accelerated Multi-Time-Scale Stochastic Approximation: Optimal Complexity and Applications in Reinforcement Learning and Multi-Agent Games

The objective of a range of problems in reinforcement learning and games can be expressed as a coupled system of $N$ equations, with each equation defined through a nonlinear operator that can be queried with noise. This work proposes a single-loop stochastic approximation (SA) algorithm which solves such a system using a (near) optimal number of queries. Given a compact statistical sample space ${\cal X}$ and operators $\{F_{i}:\mathbb{R}^{d_{1}}\times\cdots\times\mathbb{R}^{d_{N}}\times{\cal X}\rightarrow\mathbb{R}^{d_{i}}\}_{i=1,\cdots,N}$, our aim is to find a solution tuple $\boldsymbol{\theta}^{\star}=(\theta_{1}^{\star},\cdots,\theta_{N}^{\star})\in\mathbb{R}^{d_{1}}\times\cdots\times\mathbb{R}^{d_{N}}$ such that

where $\mu_{\boldsymbol{\theta}}$ denotes a distribution over ${\cal X}$ (possibly) as a function of $\boldsymbol{\theta}$. We do not assume having direct access to $\mu_{\boldsymbol{\theta}}$. Rather, our sampling oracle is that we may query an operator $P_{\boldsymbol{\theta}}:{\cal X}\rightarrow\Delta_{{\cal X}}$, which induces $\mu_{\boldsymbol{\theta}}$ as the stationary distribution. More precisely, for any $\boldsymbol{\theta}$ and $X$, we can draw a sample $X^{\prime}\sim P_{\boldsymbol{\theta}}(X)$. We consider this relaxed sampling oracle as it realistically captures how samples are generated in reinforcement learning (RL). We use $\Omega=\mathbb{R}^{d_{1}}\times\cdots\times\mathbb{R}^{d_{N}}$ to denote the space of the decision variable $\boldsymbol{\theta}$ and denote $\widebar{F}_{i}(\theta_{1},\cdots,\theta_{N})=\mathbb{E}_{X\sim\mu_{\boldsymbol{\theta}}}[F_{i}(\theta_{1},\cdots,\theta_{N},X)]$ for any $\boldsymbol{\theta}\in\Omega$.

It may appear that the system (5) is symmetric and each equation is identically coupled with any other equation. This is not true in the context of RL and games, where the problems usually have a specific hierarchical structure. The structure dictates the order in which the system can be most naturally solved and is formalized later through Assumption 3.

Exploiting the hierarchical structure, the standard multi-time-scale stochastic approximation algorithm for solving (5), thereafter referred to as MSA, maintains parameters $\theta_{1}^{[t]},\cdots,\theta_{N}^{[t]}$ as an estimate of $(\theta_{1}^{\star},\cdots,\theta_{N}^{\star})$ and simultaneously updates them, each using a unique step size. Given (possibly biased) stochastic samples $\{X^{[t]}\}$, MSA iteratively carries out

With properly selected step sizes, this simple method is guaranteed to converge to the solution of (5) under strong monotonicity and Lipschitz continuity assumptions on $\widebar{F}_{i}$ to be introduced later (Hong et al., 2023; Zeng et al., 2024b; Shen and Chen, 2022; Han et al., 2024). However, its time and sample complexity degrades as $N$ increases and is sub-optimal as known so far. When $N=1$, MSA reduces to the classic stochastic approximation algorithm (Robbins and Monro, 1951) and has a complexity of ${\cal O}(1/t)$ when $X^{[t]}$ are i.i.d. samples from a fixed distribution (and $\widetilde{{\cal O}}(1/t)$ if the samples $X^{[t]}$ are consecutively generated from a Markov chain, where $\widetilde{{\cal O}}$ hides any logarithm factor). More precisely, the last iterate of (6) converges to an $\widetilde{{\cal O}}(1/t)$ neighborhood around $\boldsymbol{\theta}^{\star}$ in $t$ iterations, matching the worst-case lower bound (Chen et al., 2022). When $N=2$, the best-known complexity, as established in Hong et al. (2023); Zeng et al. (2024b), becomes $\widetilde{{\cal O}}(1/t^{2/3})$, inferior to that under $N=1$ by a factor of $t^{1/3}$. The authors in Hong et al. (2023); Zeng et al. (2024b) only analyze the case of two equations, but we show in Section 3.2 that their techniques can be extended to establish a complexity of $\widetilde{{\cal O}}(1/t^{1/2})$ for $N=3$ and further worse bounds for $N=4,5,\cdots$. This apparent gap in the performance of MSA for different $N$ makes us question whether the problem becomes fundamentally harder as $N$ increases or whether any algorithmic/analytical improvement can be made. In this work, we answer the question in the positive direction by proposing an accelerated SA algorithm which achieves the optimal complexity $\widetilde{{\cal O}}(1/t)$ across all $N$ while maintaining the favorable properties of (6) such as simplicity and single loop structure. The same algorithm has been introduced earlier in our prior works (Doan, 2024; Zeng and Doan, 2024) but we make important contributions over them as discussed in details in Section 1.2. We now summarize our main contributions.

The MSA method (6) updates each $\theta_{i}$ in the direction of the corresponding operator $F_{i}$ evaluated at the latest iterates. Our proposed method, formally presented in Algorithm 1, modifies the update by introducing the auxiliary variables $f_{i}^{[t]}$ to estimate $\widebar{F}_{i}^{[t]}(\boldsymbol{\theta}^{[t]})$ by forming a weighted average of the operator samples controlled by a parameter $\lambda^{[t]}\in(0,1)$. These auxiliary variables will then be used to update the decision variables. If we set $\lambda^{[t]}=1$ for all $t$, Algorithm 1 reduces to MSA in (6). However, a large $\lambda^{[t]}$ leads to a high-variance estimate, whereas if $\lambda^{[t]}$ is too small $f_{i}^{[t]}$ cannot track the moving target $\widebar{F}_{i}^{[t]}(\boldsymbol{\theta}^{[t]})$. Decaying $\lambda^{[t]}$ appropriately with respect to $\alpha_{i}^{[t]}$, $f_{i}^{t}$ helps to control the variance of our estimates, which is crucial to derive the optimal convergence rate.

Note that in our algorithm the samples are drawn from the transition kernel $P_{\boldsymbol{\theta}}$ parameterized by $\boldsymbol{\theta}$ (a.k.a (9)). Thus, as $\boldsymbol{\theta}^{[t]}$ changes over time, the samples are temporally dependent and Markovian with time-varying stationary distributions. To ensure the stable behavior of the Markovian samples we make the following assumptions. Our first assumption is on the fast mixing of the Markov chain generated under $P_{\boldsymbol{\theta}}$. Given two distributions $u_{1},u_{2}\in\Delta_{{\cal X}}$, their total variation (TV) distance and the mixing time of the Markov chain are defined as

The geometric ergodicity assumption states that the distribution of a sample from the Markov chain approaches the stationary distribution geometrically fast, and is again standard in the literature on RL and SA under Markovian noise (Srikant and Ying, 2019; Wu et al., 2020; Zeng et al., 2024b).

Denoting $\tau(a)\triangleq\sup_{\boldsymbol{\theta}\in\mathbb{R}^{d}}\tau_{\boldsymbol{\theta}}(a)$, this assumption implies that there exists a positive constant $C$ depending only on $\rho$ and $m$ such that

For convenience, we denote by $\tau_{t}\triangleq\tau(\lambda^{[t]})$ the mixing time corresponding to $\lambda^{[t]}$ in Algorithm 1.

Assumption 2 represents a regularity condition on $P_{\boldsymbol{\theta}}$, which can be interpreted as a type of the triangle inequality for the TV distance defined over the transition kernels w.r.t different parameters. This assumption can be shown to hold in standard MDPs (See Wu et al. (2020, Lemma A1)).

We denote by ${\cal F}_{t}=\{X^{[0]},X^{[1]},\cdots,X^{[t]}\}$ the filtration containing all randomness generated by Algorithm 1 up to time $t$.

In this section, we establish the finite-sample complexity of A-MSA to the solution of (5) when the operators $F_{i}$ have a nested structure. For convenience, we refer to the first and last equations in (5) as the highest- and lowest-level problems, respectively, i.e., the problem levels increase from $N$ to 1. Levels here indicate the order in which the problem can be solved.

With some abuse of notation, we denote $\boldsymbol{\theta}_{1:N-1}=\{\theta_{1},\cdots,\theta_{N-1}\}$. Given $\boldsymbol{\theta}_{1:N-1}$, we use $y_{N}(\boldsymbol{\theta}_{1:N-1})$ to denote the solution of the lowest level problem

Similarly, we denote by $\{y_{N-1}(\boldsymbol{\theta}_{1:N-2}),y_{N}(\boldsymbol{\theta}_{1:N-2})\}$ the solution of the last two levels corresponding to $\widebar{F}_{N-1},\widebar{F}_{N}$ given $\boldsymbol{\theta}_{1:N-2}$

Conceptually, $y_{N-1}(\boldsymbol{\theta}_{1:N-2})$ and $y_{N}(\boldsymbol{\theta}_{1:N-2})$ represent the optimal solutions of the last and second last equation with respect to the last and second last decision variables when $\theta_{1},\cdots,\theta_{N-2}$ are given.

Generalizing this line of discussion, we introduce $y_{j}(\boldsymbol{\theta}_{1:i-1})$ to represent the optimal solution of the $j^{\text{th}}$ equation with respect to the $j^{\text{th}}$ variable when $\theta_{1},\cdots,\theta_{i-1}$ are given, for any $i\in[2,N]$ and $j\in[i,N]$. For $i=1$ (i.e. no $\theta$ is given), we additionally define $y_{j}(\emptyset)=\theta_{j}^{\star}\in\mathbb{R}^{d_{k}}$, which is the $j^{\text{th}}$ component of the optimal solution to (5). We introduce the aggregate notation

and note that ${\bf y}_{i:N}$ is the solution to the system

We now introduce the main assumption that drives our analysis, referred to as the nested strong monotonicity of the operators $\widebar{F}_{i}$ corresponding to the increasing levels defined above.

Assumption 3 represents a nested variant of strong monotonicity for the operators $\widebar{F}_{i}$, which does not require each $\widebar{F}_{i}$ being strongly monotone w.r.t to all of its variables. It states that given $\boldsymbol{\theta}_{i-1}$, $\widebar{F}_{i}$ is strongly monotone w.r.t $\theta_{i}$ when all lower level decision variables are at the corresponding optimal solutions. When $N=2$, the assumption reduces to the same condition made in the existing literature (Doan, 2022; Shen and Chen, 2022; Han et al., 2024) and can be verified to hold in RL applications including the gradient-based TD learning (Xu et al., 2019) and distributed TD learning learning¹¹1To model the distributed TD learning in the two-time-scale SA framework, the outer loop equation enforces the consensus among agents whereas in the inner loop each agent learns its local value function.. As an implication of the assumption, $y_{j}(\boldsymbol{\theta}_{1:i-1})$ is always unique for any $i,j$. Without any loss of generality, we assume $\delta\leq 1$, which makes it convenient for us to simplify terms.

Eqs. (15) and (16) are standard Lipschitz continuity assumptions on the operator $F_{i}$ and learning target (Doan, 2022; Zeng et al., 2024b; Shen and Chen, 2022). Eq. 17 assumes that the learning targets are bounded and guarantees algorithm stability in the presence time-varying Markovian noise. This assumption is made in Kaledin et al. (2020); Zeng et al. (2024b) when $N=2$. Without any loss of generality, we assume $L\geq\max_{i\in\{1,\cdots,N\},X\in{\cal X}}\|F_{i}(0,X)\|$, which is finite due to the compactness of ${\cal X}$. Assumption 4 implies that the operator $F_{i}$ can be bounded affinely by the decision variable

We use the same constant $L$ in Assumptions 2 and 4 to simplify the notation.

Finally, we will study the convergence of Algorithm 1 under the following choice of step sizes

where $\lambda^{[t]}\geq\alpha_{i}^{[t]}$ as the operator estimates are implemented at a “faster” time scale than the updates of decision variables. In particular, these step sizes satisfies the following conditions for all $t\geq 0$

where $D$ is an absolute constant depending on $L$, $B$, $\rho$, and $m$. We note that there always exist constants $c_{\lambda},c_{i},h$ such that the step sizes in (21) satisfy (22).

For our analysis, we consider the following Lyapunov function

Under the choice of step sizes above, the main result of this paper is presented in the following theorem, where we characterize the convergence complexity of Algorithm 1

Theorem 1 shows that the last iterate of the A-MSA algorithm converges to the unique solution of (5) in the mean-squared sense at an optimal rate $\widetilde{{\cal O}}(1/t)$. As the algorithm draws exactly one sample in each iteration, this rate translates to a sample complexity of $\widetilde{{\cal O}}(\epsilon^{-1})$ for finding an $\epsilon$-optimal solution. The same complexity has been shown to be achievable by the MSA algorithm under an additional smoothness assumption on $F_{i}$ (Shen and Chen, 2022). We match the rate without requiring this restrictive assumption. In the absence of the smoothness condition, our rate is order-optimal (up to an logarithm factor) and significantly improves over the best-known rate of MSA, which is $\widetilde{{\cal O}}(1/t^{2/3})$ as established in Hong et al. (2023); Zeng et al. (2024b); Han et al. (2024) when $N=2$. In Section 3.2, we further show that the convergence rate of MSA is $\widetilde{{\cal O}}(1/t^{1/2})$ when $N=3$ and deteriorates as $N$ increases. This observation highlights the advantage of our proposed algorithm where its convergence complexity is optimal for arbitrary $N$ .