How can I implement policy evaluation when reward is tied to an action outcome?

Question

I'm following Stanford reinforcement learning videos on youtube. One of the assignments asks to write code for policy evaluation for Gym's FrozenLake-v0 environment.

In the course (and books I have seen), they define policy evaluation as

$$V^\pi_k(s)=r(s,\pi(s))+\gamma\sum_{s'}p(s'|s,\pi(s))V^\pi_{k-1}(s')$$

My confusion is that in the frozen lake example, the reward is tied to the result of the action. So, for each pair state-action, I have a list that contains a possible next-state, the probability to get to that next-state and the reward. For example, being in the target state and performing any action brings a reward of $0$, but being in any state that brings me to the target state gives me a reward of $1$.

Does this mean that, for this example, I need to rewrite $V^\pi_k(s)$ as something like this:

$$V^\pi_k(s)= \sum_{s'} p(s'|s,\pi(s)) [r(s,\pi(s), s')+ \gamma V^\pi_{k-1}(s')]$$

Answer 1

The renowned book Reinforcement Learning: An Introduction (2nd edition), by Sutton and Barto, provides a different update rule than your first update rule for policy evaluation. Their update rule is more similar to your second update rule. See section 4.1. They also provide the pseudocode for policy evaluation on page 75 of the book. You can also find the pseudocode here.

Moreover, note that the update rule doesn't need to change only because the rewards are tied to the outcome of an action. This information is encoded in the functions $p$ (the transition function) and $r$ (the reward function) of the Markov decision process, which is incorporated in the update rule. If you want to understand the update rule, you should read the relevant pages (especially, chapter 4) of the cited book.