# Implementing SARSA for a 2-stage Markov Decision Process

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.

• You say "In the paper I am reading" - could you please link that resource, as it may explain why the writers are using what appears to be an over-complicated approach to solve a trivial-looking toy problem. – Neil Slater Jun 28 at 8:45
• @NeilSlater I added the link to the paper. I should mention that the research is roughly about modeling human decision-making, especially that of those who have some impairment. As such, the focus is hardly about finding an optimal solver. But my question is very simple - I want to start very simple by implementing a basic SARSA, which I did; however, the simulations I got look very different and wrong even, so I am guessing that my implementation is wrong. – dezdichado Jun 28 at 16:50
• Thanks for adding the link. I didn't read the paper yet, but would expect the researchers are implementing inverse reinforcement learning given the stated goal. That could be dificult to follow if you want to study classic RL used for optimal control. – Neil Slater Jun 28 at 17:30

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.

• Thank you - this makes things clear. However, I have one more question. The task is just people are playing this two stage MDP for a total of say $T = 200$ trials or episodes. In this view, humans are obviously learning things from episode to episode, albeit to a small extent; therefore, would it better to implement it in the second way you propose, as one continuing Markov Task with one episode and $T$ time steps? – dezdichado Jun 28 at 19:18
• @dezdichado: The learning problem is still solved by the strict episodic view, assuming you keep the Q table between episodes. I assumed this - were you thinking to reset it?. Repeating the episodes allows the learner to explore different approaches and figure out which is best - which is why it is being repeated, because the point is to study learning through trial and error. There are cases where you might want to think about aggregating reward over multiple episodes, but even then it is simpler and easier to analyse if you stick with the strictly episodic view for Q values and update rules. – Neil Slater Jun 28 at 19:22
• Sorry to bother you again but I am now transitioning from matlab to python and in doing so, I am rewriting all my codes in a more scalable, object-oriented fashion. I was wondering if could point me to a code that writes the environment in customizable manner. I did do google searches and all the implementations I could find use the openai gym interface, which as a whole is a rather advanced for a beginner programmer like me. – dezdichado Oct 19 at 20:02
• @dezdichado The basic gym interface for environments is a reasonable design for a standardised framework in my opinion. I have not worked with any other similar libraries in any depth. Sorry I cannot help further. – Neil Slater Oct 19 at 21:02