Operator-Theoretic Foundations & Policy Gradient Methods for General MDPs

Modern engineering demands infinite-dimensional control

• Active flow control: state is the fluid flow field around a wing — an infinite-dimensional function space

• Soft robotics: configuration of a Cosserat rod is a map from the rod to \(\mathbb{R}^3\) — a PDE state

• Mean field control: coordinator’s state includes a measure over local agents

The gap in RL theory

• Most RL algorithms designed for finite or (with some stretch of imagination) Euclidean MDPs

• Ad hoc regularizers (KL divergence) without structural justification

• No exploitation of problem structure (monotonicity, Lipschitz, LQG)

\(\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

• \(\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*}\]

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\in\Pi_{\textup{\texttt{ss}}}\): \[\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*}\]

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) \]

Implications

• Necessary condition for optimality: \(\mathfrak{D}_{\pi'}(v_{\pi^*}) = 0\) for all \(\pi' \in \Pi_{\textup{\texttt{ss}}}\)
• First time this result has been established for general MDPs
• Unifies results from Kakade (2002), Schulman (2015), Ju (2022) in one framework

IPM-based policy metric

For every \(s\in\mathcal S\), define \(\mathfrak{F}_{\mathcal A}(s) = \{h:\mathcal K(s)\to\mathbb{R}: h(a) = q(s,a), q\in\mathfrak{F}_{\mathcal Q}\}\) and

\[\textup{\texttt{IPM}}_{\mathfrak{F}_{\mathcal A}(s)} (\pi(s),\pi'(s)) = \sup_{h\in\mathfrak{F}_{\mathcal A}(s)} \left|\int h\, d\pi(s) - \int h\, d\pi'(s)\right|\]

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)}\]

• Assume: some upper bound on \(\beta(q_{\pi_k},\pi') \leq \beta_k\) for all stable \(\pi'\)

• 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\)

\[\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*}\]