Can I have different rewards for a single action based on which state it transitions to?

Question

I am working on an MDP where there are four states and ten actions. I am supposed to derive the optimal policy to reach the desired state. At any state, a particular action can take you to any of the other states. For ex. If we begin with state S1 -> performing action A1 on S1 can take you to S2 or S3 or S4 or just stay in the same state S1. Similarly for other actions.

My question is - is it mandatory to have only a single reward value for a single action A? Or is it possible to give a reward of 10 if action a1 on state s1 takes you to s2, give a reward of 50 if action a1 on state s1 takes you to s3, give a reward of 100 if action a1 on state s1 takes you to s4 which is the terminal state or give zero reward if that action results in the state being unchanged.

Can I do this??

Because in my case every state is better than its previous state. i.e S2 is better than S1, S3 is better than S2 and so on. So if an action on S1 is directly taking u to S4 which is the final state i would like to award it the maximal reward.

Answer 1

The reward function can be a function of the current state, current action, and next state: $R(s_t, a_t, s_{t+1})$. It's valid to use the Bellman operator in this setting because it's still a contraction and will yield the optimal value function.

NOTE: I'm assuming that you will be solving the MDP with the Bellman equation.