Reinforcement Learning for General-State General-Action MDPs

Policy Gradient Theorem and Algorithms

Abhishek Gupta
Joint work with Prof. Aditya Mahajan (McGill Univ)

Associate Professor, ECE, The Ohio State University

Founder, Ensemble Control Inc.

Research Highlights

  • Theory

    • Games and with asymmetric information
    • Reinforcement learning
    • Perturbation theory for dynamic program and fast algorithms
    • Random non-linear operator theory
      • Risk sensitive decision theory and dynamic optimization
    • Model-free anomaly detection algorithms

  • Applications

    • Electric vehicle charging optimization
    • Pricing under demand-supply mismatch in electricity and transportation markets
    • EV charger location optimization
    • Energy optimization in vehicles
    • Dynamic watermarking for autonomous systems

General MDPs \((\mathcal S, \mathcal A, \mathcal K\subset\mathcal S\times\mathcal A, \gamma, P, c, \rho)\)

\(\mathcal{S},\mathcal A\) are Polish, \(\mathcal K\subset\mathcal S\times\mathcal A\) is measurable, and \(\pi:s\mapsto\mathcal P(\mathcal K(s))\). Compute \(\pi^*\) such that \(J_{\pi^*}(\rho) \leq J_{\pi}(\rho)\): \[ \begin{align*} J_{\pi}(\rho) = \langle v_{\pi}, \rho \rangle \text{ where } v_{\pi}(s_0) = \mathop{\mathbb{E}}_{a_t\sim\pi(s_t)}\left[\sum_{t=0}^\infty \gamma^t c(s_t,a_t)\Big| s_0\right]. \end{align*} \]

The Bellman operators are defined as

\[ \begin{align*} [\boldsymbol{T} v](s) = \inf_{p \in \mathcal P(\mathcal K(s))} \int \left(c(s,a)+\gamma \int v(s') P(ds'|s,a)\right) p(da)\\ [\boldsymbol{T}_\pi v](s) = \int \left(c(s,a)+\gamma \int v(s') P(ds'|s,a)\right) \pi(da|s). \end{align*} \]

Two widely used algorithms to compute \(\pi^*\) iteratively are

  • Value iteration: \(v_{k+1} = \boldsymbol{T} v_k\) with \(v_0 = 0\)

  • Policy iteration: \(\pi_{k+1} = \mathop{\mathrm{\arg\min}}_{\pi} \left\langle \boldsymbol{T}_{\pi} v_{\pi_k},\rho \right\rangle\) with \(\pi_0\) picked arbitrarily

Inverse Problems through Reinforcement Learning

  • State and action spaces are high/infinite dimensional

  • One operator evaluation \(f(s_t)\) may take several hours to days to weeks

  • Operator learning takes samples of \((s,f(s))\) and “learns the operator”

  • Operator learning is speeding up and revolutionizing reward computation

Key Challenge: Design efficient and flexible RL algorithm that can solve general MDPs with fewer samples and some guarantee on convergence

Prior Existence Results

Previous Work

Current state of affairs

  • Finite MDPs: Watkins (1989), Haarnoja et. al. (2018)

  • Finite MDPs with Function Approximation: Williams (1992), Mnih et. al. (2015), Schulman et. al. (2015), Schulman et. al. (2017)

  • Continuous MDPs: Lillicrap et. al. (2015), Hafner et. al. (2023)


Trouble with the current state of the affairs:

  • Extreme assumptions for existence

  • No exploitation of structure (monotonicity, Lipschitz continuity, invariance, LQG etc.)

  • Using ad hoc regularizer (KL divergence, Bregman divergence)


Our solution:

  • Use integral probability metrics, adapted to natural structure in MDP

  • Existence leading to natural regularlizer leading to algorithm

  • Sample complexity can exploit recent work on dimension independence in the concentration of measures phenomena for IPM (future work)

Background on Integral Probability Metrics

  • \(\mathcal X\) Polish, weight function \(\texttt{w}(x) \ge 1\) for all \(x \in \mathcal{X}\)

  • Weighted norm \(\|f\|_{\texttt{w}} = \sup_{x\in\mathcal X} |f(x)|/\texttt{w}(x)\), \(\mathcal F_{\texttt{w}}:=\{ f \in \mathcal{M}_{\mathcal X} : \|f\|_{\texttt{w}} < \infty \}\) is a Banach space

  • Generator set \(\mathfrak{F}\subset\mathcal F_{\textup{\texttt{w}}}\) separates points. The integral probability metric is defined as \[\begin{align*} \textup{\texttt{IPM}}_{\mathfrak{F}} (\mu,\mu') = \sup_{f\in\mathfrak F} \Bigg| \int f d\mu - \int f d\mu'\Bigg| \end{align*}\]

  • Minkowski functional \(\varrho_{\mathfrak{F}}:\mathcal F_{\texttt{w}}\to[0,\infty]\) is defined as \[\begin{align*} \varrho_{\mathfrak{F}}(f) = \inf\left\{r>0: \frac{f}{r}\in \texttt{absconv}(\mathfrak{F})\right\} \text{ where } \texttt{absconv}(\mathfrak{F}) = \bar{co}(\{\alpha f: f\in \mathfrak{F}, \alpha\in[-1,1]\}) \end{align*}\]

Key Result: For \(f\) with finite Minkowski functional, we have \[\begin{align*} \Bigg| \int f d\mu - \int f d\mu'\Bigg| \leq \varrho_{\mathfrak{F}}(f) \textup{\texttt{IPM}}_{\mathfrak{F}} (\mu,\mu') \end{align*}\]

Spaces and Mappings

  • Weight functions determines \(\mathcal F_{\mathcal S}\) and \(\mathcal F_{\mathcal K}\)

  • Generator set determines \(\mathcal V_b\) and \(\mathcal Q_b\)

  • Problem structure determines \(\mathcal V\) and \(\mathcal Q\)

We pick these so that

  • \(\|\boldsymbol{P}\|<\infty\)

  • \(\boldsymbol{P}\mathcal V\subset\mathcal Q\) and \(\boldsymbol{\pi }\mathcal Q\subset\mathcal V\) for all \(\pi\in\Pi\)

Bellman Operators and the Value Iteration Algorithm

Define the composite operator \(\boldsymbol{P}_{\pi}:\mathcal F_{\mathcal S}\to \mathcal F_{\mathcal S}\) \[\boldsymbol{P}_{\pi} v(s) = \int \pi(da|s) \int v(s') P(ds'|s,a)\] which yields the Bellman operator: \[\boldsymbol{T}_\pi v= c_{\pi}+\gamma P_{\pi} v\] If \(\|\gamma \boldsymbol{P}_\pi\|<1\), then policy evaluation is written as \[v_\pi = \boldsymbol{T}^\infty_\pi 0 = \sum_{t\geq 0} (\gamma\boldsymbol{P}_{\pi})^t c_{\pi} = \underbrace{\sigma_\pi}_{(\textup{\texttt{id}}_{\mathcal F_{\mathcal S}}- \gamma \boldsymbol{P}_\pi)^{-1}} c_\pi\]

Value iteration algorithm: \(v_0 = 0\), \(v_{k+1} = \boldsymbol{T} v_k\)

Lemma: Under certain reasonable assumptions, if \(c\geq 0\), then \(\boldsymbol{T}^{k-1} 0 \leq \boldsymbol{T}^k 0 \leq \boldsymbol{T}^k_\pi 0\leq v_\pi\)

Convergence of the Value Iteration Algorithm

Assumptions

  • \(c\in\mathcal Q, c\geq 0\) is inf-compact

  • \(\Pi_{\textup{\texttt{ss}}}= \{\pi\in\Pi:\textup{\texttt{spec}}(\gamma\boldsymbol{P}_\pi)<1\}\) is nonempty

  • \(\boldsymbol{P}\mathcal V\subset\mathcal Q\) and \(\boldsymbol{\pi }\mathcal Q\subset\mathcal V\) for all \(\pi\in\Pi\)

  • For every \(q\in\mathcal Q\) such that \(q\) is inf-compact, there exists a \(\pi\in\Pi\) such that \(\pi(s) \in \mathop{\mathrm{\arg\min}}_{p\in\mathcal P(\mathcal K(s))} \int q(s,a) p(da)\)

  • The space \(\mathcal V\) is closed under bounded increasing limits

  • The transition kernel \(P\) is continuous with respect to \(\textup{\texttt{IPM}}_{\mathfrak{F}_{\mathcal V}}\): \((s_n,a_n)\to (s,a)\implies\lim_{n\to\infty} \textup{\texttt{IPM}}_{\mathfrak{F}_{\mathcal V}}(P(\cdot|s,a),P(\cdot|s_n,a_n)) = 0\)


Main Theorem: Value iteration converges to a limit if \(v_0 = 0\). The limit is the optimal value function. Greedy policy with respect to the optimal value function is the optimal policy.


Proof mainly follows from Ch 9 in Bertsekas and Shreve (1978) and Feinberg, Kasyanov and Zadoianchuk (2012).

Comparison with Prior Existence Results

This existence result readily leads to a policy gradient type algorithm

Policy Difference Lemma

Policy advantage function \(PAF_\pi:\Pi_{\textup{\texttt{ss}}}\to\mathcal V_{\textup{\texttt{b}}}\) at a policy \(\pi\in\Pi_{\textup{\texttt{ss}}}\) is defined as \[\begin{align*} PAF_{\pi}(\pi') & = \mathop{\mathbb{E}}_{s_t\sim P_\pi(\cdot|s_{t-1})}\left[\sum_{t=0}^\infty \gamma^t[\boldsymbol{\pi}' A_\pi](s_t)\right] = \sum_{t = 0}^\infty \gamma^t \boldsymbol{P}_\pi^t \boldsymbol{\pi}' A_\pi =\boldsymbol{\sigma}_\pi \boldsymbol{\pi}' A_\pi. \end{align*}\]

Policy Gradient Theorem: For spectrally stable \(\pi,\pi'\) with \(\boldsymbol{\Delta }=\boldsymbol{\pi}'- \boldsymbol{\pi}\), we have \[\begin{align*} v_{\pi'} - v_{\pi} & = PAF_{\pi}(\pi') + \gamma \boldsymbol{\sigma}_\pi \boldsymbol{\Delta}\boldsymbol{P} \boldsymbol{\sigma}_{\pi'} \boldsymbol{\pi}' A_\pi\\ & = PAF_{\pi}(\pi') + \gamma \boldsymbol{\sigma}_\pi \boldsymbol{P}_\Delta \boldsymbol{\sigma}_{\pi} \boldsymbol{\Delta }q_\pi + o(\|\boldsymbol{\Delta}\|^2). \end{align*}\]

Proof follows from perturbation theory for linear operators: \[\begin{align*} (\boldsymbol{A}+\epsilon \boldsymbol{B})^{-1} &= \boldsymbol{A}^{-1}-\epsilon (\boldsymbol{A}+\epsilon \boldsymbol{B})^{-1} \boldsymbol{B} \boldsymbol{A}^{-1} = \boldsymbol{A}^{-1}-\epsilon \boldsymbol{A}^{-1}\boldsymbol{B} (\boldsymbol{A}+\epsilon \boldsymbol{B})^{-1} \label{eqn:diffinvop}\\ & = \boldsymbol{A}^{-1}-\epsilon \boldsymbol{A}^{-1}\boldsymbol{B} \boldsymbol{A}^{-1} + \epsilon^2 \boldsymbol{A}^{-1}\boldsymbol{B} (\boldsymbol{A}+\epsilon \boldsymbol{B})^{-1} \boldsymbol{B} \boldsymbol{A}^{-1}\label{eqn:diffinvop2} % & = \op A^{-1} - \epsilon \op A^{-1} \op B \op A^{-1} + o(\epsilon). \end{align*}\]

Majorization Bound and MM Algorithm

Operator Theoretic Policy Gradient Algorithm (OTPG)

  • Assume: some upper bound on Upper Collatz–Wielandt number

  • Policy difference lemma implies \(v_{\pi'} - v_{\pi} = \underbrace{PAF_{\pi}(\pi')}_{\boldsymbol{\sigma}_\pi \boldsymbol{\pi}' A_\pi} + \gamma \boldsymbol{\sigma}_\pi \boldsymbol{\Delta}\boldsymbol{P} \boldsymbol{\sigma}_{\pi'} \boldsymbol{\pi}' A_\pi\)

Majorization Bound: \[\begin{align*} J_{\pi'}(\rho) - J_{\pi}(\rho) &= \left\langle v_{\pi'} - v_{\pi},\rho \right\rangle\leq \left\langle \boldsymbol{\pi}'A_\pi + \beta(\pi,\pi') \underbrace{\textup{\texttt{IPM}}_{\mathfrak{F}_{\mathcal A}(s)} (\pi_k(s),\pi'(s))^2}_{\leq D_{KL}(\pi,\pi')},\boldsymbol{\sigma}_\pi^*\rho \right\rangle =: M(\pi';\pi) \end{align*}\]

OTPG Algorithm: \[\begin{align*} \pi_{k+1}(s)&= \mathop{\mathrm{\arg\min}}_{p\in\mathcal P(\mathcal K(s))} \int A_{\pi_k}(s,a) p(da) + \beta_k \textup{\texttt{IPM}}_{\mathfrak{F}_{\mathcal A}(s)} (\pi_k(s),p)^2.\label{eqn:ipmppo} \end{align*}\]

MM-RKHS Algorithm for Finite MDPs

Define kernel function \(\mathfrak{K}_{\mathcal S}(s,s') = \frac{1}{2}Q_{ss'}\) and \(\mathfrak{K}_{\mathcal K}((s,a),(s',a')) = \frac{1}{2}(Q_{ss'}+R_{aa'})\) and \(\mathfrak{F}_{\mathcal V}\) and \(\mathfrak{F}_{\mathcal Q}\) are unit balls in the respective RKHS.

\[\begin{align*} \textup{\texttt{IPM}}_{\mathfrak{F}_{\mathcal A}(s)} (p_1,p_2)^2 = \frac{1}{2} \big(p_1^T R p_1 + p_2^T R p_2 - 2 p_1^T R p_2\big) \text{ for } p_1,p_2\in \mathcal P(\mathcal A) \end{align*}\]

\[\begin{align*} \pi_{k+1}(s)&= \mathop{\mathrm{\arg\min}}_{p\in\mathcal P(\mathcal A)}\ \underbrace{A_{\pi_k}(s,\cdot)^T p + \beta_k \textup{\texttt{IPM}}_{\mathfrak{F}_{\mathcal A}(s)} (\pi_k(s),p)^2}_{\text{OTPG Algorithm}} + \underbrace{\frac{1}{\eta_k}D_{KL}(p\|\pi_k(s))}_{\text{Mirror descent}} \end{align*}\]

MM-RKHS Algorithm: At each step \(k\), run the following for \(\bar l\) steps with \(p_0=\pi_k(s)\) \[\begin{align*} \vartheta_{k,l}(s,a) &= -\eta_k \Big(A_{\pi_k}(s,a) - \beta_k R_a^T \big(\pi_k(s)-p_l\big)\Big) \tag{clipped}\\ p_{l+1,a} &= \frac{1}{Z_k(p_l)} p_{l,a} \exp(\vartheta_{k,l}(s,a)), \text{ where } Z_k(p_l) = \sum_{a\in\mathcal A} p_{l,a} e^{\vartheta_{k,l}(s,a)} \end{align*}\] and set \(\pi_{k+1}(s) = p_{\bar l}\).

Simulation Results

(1000 states, 200 actions, branching factor 20)

Simulation Results (log axis): Superlinear Convergence

Concluding Thoughts


Finite MDPs:

  • The value function is gradient dominated as a function of policy

  • Consequently, policy gradient methods (and perhaps the MM algorithm) converges to the optimal policy

  • PPO has been wildly successful in many applications across a wide range of engineering fields

  • But PPO requires a lot of samples that are just not available in inverse design problems


General MDPs

  • Rederive MM algorithm for the general MDPs with enough flexibility for the algorithm design

  • Introduced OTPG algorithm for general MDPs and MM-RKHS for finite MDPs

  • Demonstrated superior performance in simulations

Open Problems


Future Work

  • Inverse design problems using operator learning and RL

  • Derive sample complexity guarantees and investigate the convergence properties of the OTPG and MM-RKHS algorithm

  • Identify complexity of simplex methods, policy iteration, value iteration, and policy gradient methods for solving MDPs


Potential Applications

  • Real time control of soft robotics

  • Design of high temperature superconductors

  • Design of solid electrolytes