Does maximizing the average reward also maximizes the expected return in the initial state?

Question

Suppose I have an episodic Markov Decision Process where all episodes start in the same state, $s_0$. I also have a parameterized policy $\pi_\theta$, and I'm trying to find a $\theta$ such that the performance of the policy is maximized in this environment.

Let's examine two different ways of defining performance for the policy. The first one is simply the value (expected accumulated reward) of the policy in the initial state: $$ J_1(\theta)=V_{\pi_\theta}(s_0) $$

The second one is the average reward for the same policy. $$ J_2(\theta)=\mathbb E_{S\sim d_{\pi_\theta}, A\sim \pi_\theta(\cdot|S)}\big[R(S,A)\big] $$ where $d_{\pi_\theta}$ is the on-policy distribution for policy $\pi_\theta$.

If I remember Sutton's book correctly, then $\nabla_\theta J_1(\theta)\propto\nabla_\theta J_2(\theta)$. Is this really the case? In simpler words, does maximizing the average reward imply that we're also maximizing the expected return at the starting state? If so, how can one prove this?

Note

This question might seem similar to Average expected reward vs expected reward for start-state, but the similarity is superficial. I am aware that $V_{\pi_\theta}(s_0)\neq\mathbb E_{S\sim d_{\pi_\theta}}[V_{\pi_\theta}(S)]$, which seems to be the source of confusion for the asker. Instead of asking if such quantities are equal, I'm asking if maximizing one leads to maximization of the other.

Answer 1

Your question is about the relationship between different metrics for the policy gradient methods. Let $v_\pi(s_0)$ be the state value of a starting state. Let $\bar{r}_\pi$ be the average reward (sorry that I am used to my own notations). Consider the discounted case.

The gradient of $v_\pi(s_0)$ is $$\nabla_{\theta} v_\pi(s_0)=\mathbb{E}\big[\nabla_{\theta}\ln\pi(A|S,\theta)q_{\pi}(S,A)\big]$$ where $S\sim \rho_\pi$ and $A\sim \pi(s,\theta)$. Here, $\rho_\pi$ is a special long-run distribution of the states (not stationary distribution).

The gradient of $\bar{r}_\pi$ is $$ \nabla_{\theta} \bar{r}_{\pi}\approx \mathbb{E}\big[\nabla_{\theta}\ln\pi(A|S,\theta)q_{\pi}(S,A)\big] $$ where $S\sim d_\pi$ and $A\sim \pi(s,\theta)$. Here, $d_\pi$ is the stationary distribution.

It is clear that they are not the same because $S$ obeys different distributions. However, they are very similar. Moreover, if we consider the stochastic gradient (that is the value after removing the expectation), then their stochastic gradient are the same.

In summary, maximizing the two metrics is not equivalent. However, people usually do not distinguish them.

Details of the above two equations can be found in Theorem~9.2 and Theorem~9.3 in the book Mathematical foundation of reinforcement learning.