Revisit Policy Optimization in Matrix Form

TL;DR

This paper revisits policy optimization in tabular reinforcement learning by disentangling policy and environment dynamics in matrix form, simplifying policy updates and extending to model-based RL.

Contribution

It introduces a matrix formulation that separates policy and environment effects, enabling more straightforward policy optimization and potential extensions to model-based reinforcement learning.

Findings

Reformulation of policy evaluation in matrix form.

Unified framework for policy gradient and TRPO.

Potential extension to model-based RL.

Abstract

In tabular case, when the reward and environment dynamics are known, policy evaluation can be written as $\bm{V}_{\bm{\pi}} = (I - \gamma P_{\bm{\pi}})^{-1} \bm{r}_{\bm{\pi}}$, where $P_{\bm{\pi}}$ is the state transition matrix given policy ${\bm{\pi}}$ and $\bm{r}_{\bm{\pi}}$ is the reward signal given ${\bm{\pi}}$. What annoys us is that $P_{\bm{\pi}}$ and $\bm{r}_{\bm{\pi}}$ are both mixed with ${\bm{\pi}}$, which means every time when we update ${\bm{\pi}}$, they will change together. In this paper, we leverage the notation from \cite{wang2007dual} to disentangle ${\bm{\pi}}$ and environment dynamics which makes optimization over policy more straightforward. We show that policy gradient theorem \cite{sutton2018reinforcement} and TRPO \cite{schulman2015trust} can be put into a more general framework and such notation has good potential to be extended to model-based reinforcement learning.