Reinforcement Learning for General-State General-Action MDPs
Policy Gradient Theorem and Algorithms
Associate Professor, ECE, The Ohio State University
Founder, Ensemble Control Inc.
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
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”
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
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\)
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 yields \(v_0\leq v_1\leq \cdots\) and 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}}}\): \[\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 \end{align*}\]
Proof follows from perturbation theory for linear operators: \[\begin{align*} (\boldsymbol{A}+\epsilon \boldsymbol{B})^{-1} - \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} \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
This Presentation
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
Future Work
Inverse design using RL: superconductor design, new materials discovery
Investigate the convergence properties of the OTPG and MM-RKHS algorithm
Derive sample complexity guarantees
Bring clarity on simplex methods, policy iteration, value iteration, and policy gradient methods for solving MDPs
Abhishek Gupta | gupta.706@osu.edu | The Ohio State University