Implementing SARSA for a 2-stage Markov Decision Process

Question

I am a bit confused as to how exactly I should be implementing SARSA (or Q-learning too) on what is a simple 2-stage Markov Decision Task. The structure of the task is as follows:

Basically, there are three states $\{S_1,S_2,S_3\}$ with $S_1$ is in the first stage for which the two possible actions are the two yellow airplanes. $S_2$ and $S_3$ are the possible states for the second stage and the feasible actions are the blue and red background pictures, respectively. There is only a reward at the end of the second stage choice. If I call the two first stage actions $\{a_{11},a_{12}\}$ and the four possible second stage actions $\{a_{21},a_{22},a_{23},a_{24}\}$, from left to right, then a sample trial/episode will look like: $$S_1, a_{11}, S_2, a_{22},R \quad \text{ or }\quad S_1, a_{11}, S_3, a_{24}, R.$$

In the paper I am reading, where the figure is from, they used a complicated version of TD$(\lambda)$ in which they maintained two action-value functions $Q_1$ and $Q_2$ for each stages. On the other hand, I am trying to implement a simple SARSA update for each episode $t$: $$Q_{t+1}(s,a)= Q_t(s,a) + \alpha\left(r + \gamma\cdot Q_t(s',a') - Q_t(s,a)\right).$$

In the first-stage, there is no reward so an actual realization will look like: $$Q_{t+1}(S_1, a_{11}) = Q_t(S_1,a_{11})+\alpha\left( \gamma\cdot Q_t(S_3,a_{23}) - Q_t(S_1,a_{11})\right).$$

I guess my confusion is then how should it look like for the second stage of an episode? That is, if we continue the above realization of the task above, $S_1, a_{11}, S_3, a_{23}, R$, would should fill in the $?$: $$Q_{t+1}(S_3,a_{23}) = Q_t(S_3,a_{23}) + \alpha\left(R +\gamma\cdot Q_t(\cdot,\cdot)-Q_t(s_3,a_{23}) \right)$$

One on hand, it seems to me that since this is the end of an episode, we assign $0$ to the $Q_t(\cdot,\cdot).$ On the other hand, the nature of this task is that it repeats the same episode over time for a total of $T$, a large number, times we need $Q_t(\cdot,\cdot) = Q_t(S_1,\cdot),$ with the additional action-selection in the first stage there.

I will greatly appreciate if someone can tell me what is the right way to go here.

The link to paper

Answer 1

In this game you can view end of an episode two ways:

• There is an implied, terminal, fourth state $s_4$ representing the end of the game.

• You could view the process as a continuous repeating one, where no matter what the choice is made in $s_2$ or $s_3$, the following state is $s_1$.

The first, terminating, view is a simpler and entirely natural view since nothing that the agent does in one episode can influence the next. It will result in a Q table that predicts future rewards within a single episode for the current agent (as opposed to discounted view over multiple episodes).

You are over-complicating things for yourself by ignoring that a zero reward is still a reward (of $0$). There is no need to remove $R$ from your initial update rule. In many environments there are rewards collected before the end of an episode.

In addition, to complete the standard episodic view, you can note that $Q(s_4, \cdot) = 0$ always by definition, hence so does $\text{max}_{a'}[Q(s_4, a'] = 0$. It is common here though to have a branch based on detecting a terminal state, and use a different update rule:

$$Q_{t+1}(S_3,a_{23}) = Q_t(S_3,a_{23}) + \alpha\left(R - Q_t(s_3,a_{23}) \right)$$

In brief, most implementations of TD algorithms do this:

• Always assume a reward on each time step, which can be set to $0$

• Special case for end of episode with a simplified update rule, to avoid needing to store, look up or calculate the $0$ value associated with terminal states

When implementing the environment, it is common to have a step function that always returns reward, next state and whether or not it is terminal e.g.

reward, next_state, done = env.step(action)

Details may vary around this. If you are working with an environment that does not have such a function (many will not have an inherent reward), then it is common to implement a similar function as a convenient wrapper to the environment so that the agent code does not have to include calculations of what the reward should be or whether the state is terminal.