Operator-Theoretic Foundations & Policy Gradient Methods for General MDPs
Abhishek Gupta
(joint work with Prof. Aditya Mahajan (McGill University))
Associate Professor & Graduate Studies Chair, ECE
Co-Director, IITB-OSU Frontier Center
The Ohio State University
Founder, Ensemble Control Inc.
Modern Engineering Demands Control of Infinite-dimensional Systems
Active flow control: State is the fluid flow field around a body — an infinite-dimensional function space.
Goal is to reduce the aerodynamic drag and reduce cabin noise using actuators that manipulate the boundary layer.
DFT based New Material Design
New material discovery: State is the electron density field in the lattice — an infinite-dimensional function space.
Goal is to search for continuous chemical composition and crystal structure spaces for target properties (critical temperatures (\(T_c\)) & electronic properties) using density functional theory.
Solving Inverse Problems through Reinforcement Learning
- Metamaterial & Photonic Design: Finding continuous permittivity fields \(\epsilon(\mathbf{x})\) to maximize target wave transmission (Maxwell PDEs).
- Battery Fast Charging: Designing charging current profiles \(I(t)\) to minimize solid electrolyte interphase (SEI) growth & thermal degradation (Electrochemical PDEs).
Previous work on Existence Results
Previous Work on Policy Gradient and RL Algorithms
Finite MDPs: Watkins (1989), Borkar and Konda (1997), 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)
Operator-Theoretic Viewpoint + IPM
→ New Existence Result
Operator-Theoretic Viewpoint
Key innovation: decompose the transition into two operators
\(\boldsymbol{P}:\mathcal F_{\mathcal S}\to \mathcal F_{\mathcal K}\) — integrates over next state \[[\boldsymbol{P} v](s,a) = \int v(s') P(ds'|s,a)\]
\(\boldsymbol{\pi}:\mathcal F_{\mathcal K}\to \mathcal F_{\mathcal S}\) — integrates over action \[[\boldsymbol{\pi }q](s) = \int q(s,a)\pi(da|s)\]
Composite: \(\boldsymbol{P}_\pi = \boldsymbol{\pi}\boldsymbol{P}:\mathcal F_{\mathcal S}\to\mathcal F_{\mathcal S}\)
Bellman: \(\boldsymbol{T}_\pi v = c_\pi + \gamma \boldsymbol{P}_\pi v\)
Integral Probability Metrics (IPM)
\(\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:
\[\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: \[\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*}\]
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 \boldsymbol{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{\boldsymbol{\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\)
Upshot: \(v_k\) is increasing and pointwise bounded by \(v_{\pi}\)
Convergence of the Value Iteration Algorithm
\(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).
Key difference from prior work: No contraction mapping — uses monotone convergence in structured spaces. Subsumes LQG, finite MDPs, Lipschitz MDPs within one framework.
Comparison with Prior Existence Results
Perturbation Theory for Linear Operators
→ Policy Gradient Theorem
A Perturbation Result of Inverse Operators
The fundamental lemma (Kato, 1966; Atkinson, 2005)
Let \(\boldsymbol{A}, \boldsymbol{B} \in \mathcal B(\mathcal X, \mathcal X)\) with \(\boldsymbol{A}^{-1}\) bounded. Then \(\exists\;\bar\epsilon>0\) such that \(\forall\;\epsilon\in[-\bar\epsilon,\bar\epsilon]\):
First resolvent identity: \[\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}
\end{align*}\]
Second-order expansion: \[\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}
\end{align*}\]
Application to MDPs: Set \(\boldsymbol{A} = \textup{\texttt{id}}_{\mathcal F_{\mathcal S}}- \gamma\boldsymbol{P}_\pi\), \(\epsilon\boldsymbol{B} = -\gamma\boldsymbol{P}_\Delta\) where \(\boldsymbol{P}_\Delta = (\boldsymbol{\pi}' - \boldsymbol{\pi})\boldsymbol{P}\). Then \(\boldsymbol{A}^{-1} = \boldsymbol{\sigma}_\pi\) and \((\boldsymbol{A} + \epsilon\boldsymbol{B})^{-1} = \boldsymbol{\sigma}_{\pi'}\).
Policy Difference Lemma
Advantage and Policy Advantage Function
Advantage function: \(A_\pi = q_\pi - v_\pi\) where \(q_\pi = c + \gamma\boldsymbol{P} v_\pi\)
Policy Advantage Function (PAF) at \(\pi,\pi'\in\Pi_{\textup{\texttt{ss}}}\) from Kakade and Langford (2002):
\[\begin{align*}
PAF_{\pi}(\pi') = \boldsymbol{\sigma}_\pi \boldsymbol{\pi}' A_\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]
\end{align*}\]
Policy Difference Lemma (via perturbation theory): For \(\pi,\pi'\in\Pi_{\textup{\texttt{ss}}}\) with \(\boldsymbol{\Delta }=\boldsymbol{\pi}'- \boldsymbol{\pi}\): \[\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{\Delta}\boldsymbol{P} \boldsymbol{\sigma}_{\pi} \boldsymbol{\Delta }q_\pi + o(\|\boldsymbol{\Delta}\|^2).
\end{align*}\]
Proof sketch: Apply the resolvent expansion to \(\boldsymbol{\sigma}_{\pi'} = (\textup{\texttt{id}}_{\mathcal F_{\mathcal S}}- \gamma\boldsymbol{P}_\pi - \gamma\boldsymbol{P}_\Delta)^{-1}\) and collect first- and second-order terms. The compact operator-theoretic proof replaces page-long algebraic arguments in prior work.
Gâteaux Derivative of the Value Function
Directional derivative in policy space
For \(\pi, \pi' \in \Pi_{\textup{\texttt{ss}}}\), the Gâteaux derivative of \(v_\pi\) in the direction of \(\pi'\): \[
\mathfrak{D}_{\pi'}(v_\pi) = \lim_{\epsilon\downarrow 0} \frac{1}{\epsilon}(v_{\pi+\epsilon(\pi'-\pi)} - v_\pi)
\]
![]()
Gâteaux Derivative of the Value Function
Directional derivative in policy space
For \(\pi, \pi' \in \Pi_{\textup{\texttt{ss}}}\), the Gâteaux derivative of \(v_\pi\) in the direction of \(\pi'\): \[
\mathfrak{D}_{\pi'}(v_\pi) = \lim_{\epsilon\downarrow 0} \frac{1}{\epsilon}(v_{\pi+\epsilon(\pi'-\pi)} - v_\pi)
\]
Theorem: \(\mathfrak{D}_{\pi'}(v_\pi) = PAF_\pi(\pi')\)
The Policy Advantage Function is precisely the Gâteaux derivative of the value function w.r.t. the policy operator.
Implications
- Necessary condition for optimality: \(\mathfrak{D}_{\pi'}(v_{\pi^*}) = 0\) for all \(\pi' \in \Pi_{\textup{\texttt{ss}}}\)
- Unifies results from Kakade (2002), Schulman (2015), Ju (2022) in one framework
Policy Gradient + Majorization Minimization
→ Class of Policy Gradient Methods
Majorization Bound
- 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{\Delta }q_\pi\)
An upper bound on the second-order term:
\[[\boldsymbol{\Delta }\boldsymbol{P} \boldsymbol{\sigma}_{\pi'} \boldsymbol{\Delta }q_\pi](s)
\leq \beta \textup{\texttt{w}}_{\mathcal S}(s) \textup{\texttt{IPM}}_{\Pi,\textup{\texttt{w}}_{\mathcal S}}(\pi,\pi')^2\]
where \(\beta = \kappa_{\boldsymbol{P}}\|\boldsymbol{\sigma}_{\pi'}\|\varrho_{\mathfrak{F}_{\mathcal Q}}(q)\) and the weighted policy metric is: \[\textup{\texttt{IPM}}_{\Pi,\textup{\texttt{w}}_{\mathcal S}}(\pi,\pi') = \sup_{s\in\mathcal S} \frac{\textup{\texttt{IPM}}_{\mathfrak{F}_{\mathcal A}(s)} (\pi(\cdot|s), \pi'(\cdot|s))}{\textup{\texttt{w}}_{\mathcal S}(s)}\]
Majorization Theorem: \[\begin{align*}
J_{\pi'}(\rho) - J_{\pi}(\rho) &\leq \left\langle \boldsymbol{\pi}'A_\pi + \beta \textup{\texttt{w}}_{\mathcal S}\textup{\texttt{IPM}}_{\Pi,\textup{\texttt{w}}_{\mathcal S}}(\pi,\pi')^2,\boldsymbol{\sigma}_\pi^*\rho \right\rangle = M(\pi';\pi)
\end{align*}\]
IPPO Algorithm
Minimize the majorization bound \(\pi_{k+1} = \arg\min_{\pi'} M(\pi';\pi_k)\)
IPPO Algorithm (IPM based Policy Optimization): Under T&C \[\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 \text{ for all } \underbrace{s\in supp(\boldsymbol{\sigma}_{\pi_k}^*\rho)}_{\textbf{not entire state space}}
\end{align*}\]
Generalizes TRPO to general MDPs: IPM square is replaced with KL divergence in TRPO \[\begin{align*}
J_{\pi'}(\rho) - J_{\pi}(\rho) &\leq \left\langle \boldsymbol{\pi}'A_\pi + \beta \:\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*}\]
In PPO, due to estimation error in advantage function, clipping of policy update is used.
What if we pick IPM to be maximum mean discrepancy? Would it simplify the algorithm?
IPPO-MMD 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{IPPO Algorithm}} + \underbrace{\frac{1}{\eta_k}D_{KL}(p\|\pi_k(s))}_{\text{Mirror descent}}
\end{align*}\]
IPPO-MMD 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
Summary and Future Work
→ Potential Collaborative Activities
Summary of Contributions
Linear Operator Formulation
- Kernel & policy as linear operators over weighted function spaces \(\mathcal{F}_{\mathcal{S}}\) and \(\mathcal{F}_{\mathcal{K}}\)
- Appropriate assumptions \(\implies\) can use linear operator theory
- Monotone convergence, does not invoke contraction mapping
- Uses spectral theory of linear operators
- Unified: subsumes LQG, Finite, & Lipschitz
A New Policy Gradient Method: IPPO Algorithm
- Perturbation theory of linear operators yields policy gradient
- Majorization bound obtained using integral probability metrics
- Yields IPPO algorithm for general MDPs
IPPO-MMD Algorithm for finite MDPs
- Use MMD + mirror descent \(\implies\) Low-complexity policy optimization IPPO-MMD
- Superlinear convergence in simulations
- Anecdotally appears to be better than PPO in multiple environments
Future Work
Solve inverse design problems using operator learning and RL
Derive sample complexity guarantees for IPPO and IPPO-MMD by exploiting dimension independence in concentration of measures for IPM
Establish global optimality, convergence properties and gradient dominance for general MDPs
Seeking collaboration
- Solve complex problems using RL in materials discovery, drug design, robotics, etc.
- Unify different strands of work in statistics, parameter identification, reinforcement learning, and dynamic optimization into one coherent theory that can be exploited for AI for Scientific Discovery
