How do I convert an MDP with the reward function in the form $R(s,a,s')$ to and an MDP with a reward function in the form $R(s,a)$?

Question

The AIMA book has an exercise about showing that an MDP with rewards of the form $r(s, a, s')$ can be converted to an MDP with rewards $r(s, a)$, and to an MDP with rewards $r(s)$ with equivalent optimal policies.

In the case of converting to $r(s)$ I see the need to include a post-state, as the author's solution suggests. However, my immediate approach to transform from $r(s,a,s')$ to $r(s,a)$ was to simply take the expectation of $r(s,a,s')$ with respect to s' (*). That is:

$$ r(s,a) = \sum_{s'} r(s,a,s') \cdot p(s'|s,a) $$

The authors however suggest a pre-state transformation, similar to the post-state one. I believe that the expectation-based method is much more elegant and shows a different kind of reasoning that complements the introduction of artificial states. However, another resource I found also talks about pre-states.

Is there any flaw in my reasoning that prevents taking the expectation of the reward and allow a much simpler transformation? I would be inclined to say no since the accepted answer here seems to support this. This answer mentions Sutton and Barto's book, by the way, which also seems to be fine with taking the expectation of $r(s, a, s')$.

This is the kind of existential question that bugs me from time to time and I wanted to get some confirmation.

(*) Of course, that doesn't work in the $r(s, a)$ to $r(s)$ case, as we do not have a probability distribution over the actions (that would be a policy, in fact, and that's what we are after).

Answer 1

I think I may be in position to answer my own question. The Bellman equation (for the optimal policy) for a MDP with $r(s,a,s')$ rewards would look like this:

$$V(s) = \max_a \left\{ \sum_{s'} p(s'|s,a)(r(s,a,s') + \gamma V(s')) \right\} $$ $$V(s) = \max_a \left\{ \sum_{s'} p(s'|s,a) \cdot r(s,a,s') + \gamma \sum_{s'} p(s'|a,s) \cdot V(s') \right\} $$

Now, $ \sum_{s'} p(s'|s,a) \cdot r(s,a,s') $ is precisely $ \mathbb{E}\left[ r(s,a,s') | s,a \right] = r(s,a) $.

So all in all, the resulting Bellman equation looks like this:

$$V(s) = \max_a \left\{r(s,a) + \gamma \sum_{s'} p(s'|s,a) \cdot V(s') \right\} $$

It is clear, then, that a process with $ r(s,a,s') $ rewards can be transformed to a $ r(s,a) $ process without introducing artificial states and maintaining the optimal policies.

As a side note unrelated to the question itself, that leads me to believe that $ r(s,a,s') $ functions may be convenient in some scenarios, but they do not add "expressive power", in the sense that they do not allow to model problems more compactly (as it happens when we extend $ r(s) $ to $ r(s,a) $).