Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Fast Two-Time-Scale Stochastic Gradient Method with Applications in Reinforcement Learning

Sihan Zeng, Thinh T. Doan

TL;DR

This work addresses slow convergence in two-time-scale stochastic gradient methods for RL-style problems where the gradient of the upper objective $h(\theta)$ depends on solving a lower-level monotone operator. It introduces an averaging-based fast two-time-scale algorithm that denoises the stochastic operator estimates $F$ and $G$, decoupling the updates of $\theta$ and $\omega$ and enabling larger steps. Theoretical results establish finite-time rates under strong convexity, the Polyak-Lojasiewicz condition, and general nonconvexity, achieving ${\cal O}(1/k)$ or ${\cal O}(1/\sqrt{k})$ rates and significantly improving over prior ${\cal O}(k^{-2/3})$ and ${\cal O}(k^{-2/5})$ baselines. Applications to TDC policy evaluation, LQR policy optimization, and entropy-regularized MDPs demonstrate practical RL benefits, with simulations confirming faster convergence than standard two-time-scale SA.

Abstract

Two-time-scale optimization is a framework introduced in Zeng et al. (2024) that abstracts a range of policy evaluation and policy optimization problems in reinforcement learning (RL). Akin to bi-level optimization under a particular type of stochastic oracle, the two-time-scale optimization framework has an upper level objective whose gradient evaluation depends on the solution of a lower level problem, which is to find the root of a strongly monotone operator. In this work, we propose a new method for solving two-time-scale optimization that achieves significantly faster convergence than the prior arts. The key idea of our approach is to leverage an averaging step to improve the estimates of the operators in both lower and upper levels before using them to update the decision variables. These additional averaging steps eliminate the direct coupling between the main variables, enabling the accelerated performance of our algorithm. We characterize the finite-time convergence rates of the proposed algorithm under various conditions of the underlying objective function, including strong convexity, Polyak-Lojasiewicz condition, and general non-convexity. These rates significantly improve over the best-known complexity of the standard two-time-scale stochastic approximation algorithm. When applied to RL, we show how the proposed algorithm specializes to novel online sample-based methods that surpass or match the performance of the existing state of the art. Finally, we support our theoretical results with numerical simulations in RL.

Fast Two-Time-Scale Stochastic Gradient Method with Applications in Reinforcement Learning

TL;DR

This work addresses slow convergence in two-time-scale stochastic gradient methods for RL-style problems where the gradient of the upper objective depends on solving a lower-level monotone operator. It introduces an averaging-based fast two-time-scale algorithm that denoises the stochastic operator estimates and , decoupling the updates of and and enabling larger steps. Theoretical results establish finite-time rates under strong convexity, the Polyak-Lojasiewicz condition, and general nonconvexity, achieving or rates and significantly improving over prior and baselines. Applications to TDC policy evaluation, LQR policy optimization, and entropy-regularized MDPs demonstrate practical RL benefits, with simulations confirming faster convergence than standard two-time-scale SA.

Abstract

Two-time-scale optimization is a framework introduced in Zeng et al. (2024) that abstracts a range of policy evaluation and policy optimization problems in reinforcement learning (RL). Akin to bi-level optimization under a particular type of stochastic oracle, the two-time-scale optimization framework has an upper level objective whose gradient evaluation depends on the solution of a lower level problem, which is to find the root of a strongly monotone operator. In this work, we propose a new method for solving two-time-scale optimization that achieves significantly faster convergence than the prior arts. The key idea of our approach is to leverage an averaging step to improve the estimates of the operators in both lower and upper levels before using them to update the decision variables. These additional averaging steps eliminate the direct coupling between the main variables, enabling the accelerated performance of our algorithm. We characterize the finite-time convergence rates of the proposed algorithm under various conditions of the underlying objective function, including strong convexity, Polyak-Lojasiewicz condition, and general non-convexity. These rates significantly improve over the best-known complexity of the standard two-time-scale stochastic approximation algorithm. When applied to RL, we show how the proposed algorithm specializes to novel online sample-based methods that surpass or match the performance of the existing state of the art. Finally, we support our theoretical results with numerical simulations in RL.
Paper Structure (35 sections, 18 theorems, 161 equations, 2 figures, 1 table, 2 algorithms)

This paper contains 35 sections, 18 theorems, 161 equations, 2 figures, 1 table, 2 algorithms.

Table of Contents

  1. Introduction
  2. Main Contributions
  3. Related Work
  4. Faster Two-Time-Scale Stochastic Gradient Method
  5. Finite-Time Complexity
  6. Strong Convexity
  7. Non-Convexity with Polyak-Lojasiewicz Condition
  8. General Nonconvexity
  9. Motivating Applications
  10. Gradient-Based Policy Evaluation in Discounted-Reward MDP
  11. Policy Optimization for Linear-Quadratic Regulators
  12. Policy Optimization for Entropy-Regularized MDPs
  13. Numerical Simulations
  14. Analysis for Strongly Convex Functions
  15. Proof of Theorem \ref{['thm:stronglyconvex']}
  16. ...and 20 more sections

Key Result

Theorem 1

We recall the definition of residual variables in eq:residuals and denote $z_k = \|\theta_k-\theta^{\star}\|^2$. Under Assumptions assump:Lipschitz-assump:stronglyconvex and the step sizes in eq:step_sizes:stronglyconvex, the iterates of Algorithm alg:main satisfy for all $k$

Figures (2)

  • Figure 1: Fast TDC Algorithm for Random Policy Evaluation
  • Figure 2: Performance of Fast Actor-Critic Algorithm for LQR

Theorems & Definitions (20)

  • Remark 1
  • Theorem 1
  • Remark 2
  • Theorem 2
  • Theorem 3
  • Lemma 1
  • Lemma 2
  • Lemma 3
  • Lemma 4
  • Lemma 5
  • ...and 10 more