Math of Q-Learning — Python. Understand where the Bellman equation... | by Omar Aflak | Towards Data Science

Q-Learning is a type of Reinforcement Learning which is a type of Machine Learning. Reinforcement learning has been used lately (typically) to teach an AI to play a game (Google DeepMind Atari, etc). Our goal is to understand a simple version of Reinforcement learning called Q-Learning, and write a program that will learn how to play a simple “game”. Let’s dive in!
A Markov chain is a mathematical model that experiences transition of states with probabilistic rules.
Here we have two states E and A, and the probabilities of going from one state to another (e.g. 70% chance of going to state A, starting from state E).
A Markov Decision Process (MDP) is an extension of the Markov chain and it is used to model more complex environments. In this extension, we add the possibility to make a choice at every state which is called an action. We also add a reward which is a feedback from the environment for going from one state to another taking one action.
In the image above, we are in the initial state don’t understand, where we have two possible actions, study and don’t study. For the study action, we may end up in different states according to a probabilistic rule. This is what we call a stochastic environment (random), in the sense that for one same action taken in the same state, we might have different results (understand and don’t understand).
In reinforcement learning, this is how we model a game or environment, and our goal will be to maximize the reward we get from that environment.
The reward is the feedback from the environment that tells us how good we are doing. It can be the number of coins you grab in a game for example. Our goal is to maximize the total reward. Therefore, we need to write it down.
This is the total reward we can get starting at some point t in time.
For example, if we use the MDP presented above. We’re initially in the state don’t understand, we take the study action which takes us randomly to don’t understand. Therefore we experienced the reward r(t+1)=-1. Now we can decide to take another action which will give r(t+2) and so on. The total reward is the sum of all the immediate rewards we get for taking actions in the environment.
Defining the reward this way leads to two major problems :
This sum can potentially go to infinity, which doesn’t make sense since we want to maximize it.
We are accounting as much for future rewards as we do for immediate rewards.
One way to fix up these problems is to use a decreasing factor for future rewards.
Setting γ=1 takes us back to the first expression where every reward is equally important. Setting γ=0 results in only looking for the immediate reward (always acting for the optimal next step). Setting γ between 0 and 1 is a compromise to look more for immediate reward but still account for future rewards.
We can rewrite that expression in a recursive manner, that will come handy later on.
A policy is a function that tells what action to take when in a certain state. This function is usually denoted π(s,a) and yields the probability of taking action a in state s. We want to find the policy that maximizes the reward function.
If we get back to the previous MDP for example, the policy can tell you the probability of taking action study when you’re in the state don’t understand.
Moreover, because this is a probability distribution, the sum over all the possible actions must be equal to 1.
We are going to start playing around with some equations, and for that we need to introduce new notations.
This is the expected immediate reward for going from state s to state s’ through action a.
This is the transition probability of going from state s to state s’ through action a.
In this example :
The expected immediate reward for going from state don’t understand to state don’t understand through action don’t study is equal to 0.
The probability of going from state don’t understand to state understand through action study is equal to 80%.
Two so-called “value functions” exist. The state value function, and the action value function. These functions are a way to measure the “value”, or how good some state is, or how good some action is, respectively.
The value of a state is the expected total reward we can get starting from that state. It depends on the policy which tells how to take decisions.
The value of an action taken in some state is the expected total reward we can get, starting from that state and taking that action. It also depends on the policy.
Now that we are settled with notations we can finally start playing around with the math! Looking at the following diagram during the calculation can help you understand.
We will start by expanding the state value function. The expected operator is linear.
Next, we can expand the action value function.
This form of the Q-Value is very generic. It handles stochastic environments, but we could write it in a deterministic one. Meaning, whenever you take an action you always end up in the same next state and receive the same reward. In that case, we simply do not need to make a weighted sum with probabilities, and the equation becomes:
Where s’ is the state you end up in for taking action a in state s. Written, more formally, this is:
You probably already came across greedy policy reading on the internet. A greedy policy is a policy where you always choose the optimal next step.
In a greedy policy context, we can write a relation between the state value and the action value functions.
Therefore, plugging this into the previous equation, we get the Q-Value of a (state, action) pair in a deterministic environment, following a greedy policy.
Or simply,
And this is the Bellman equation in the Q-Learning context ! It tells that the value of an action a in some state s is the immediate reward you get for taking that action, to which you add the maximum expected reward you can get in the next state.
It actually makes sens when you think about it.
Here, if you only look at the immediate reward, you surely choose to go left. Unfortunately, the game ends after and you cannot get more points.
If you add the maximum expected reward of the next state, then you will most probably go to the right since the maximum expected reward of S1 is equal to zero and the maximum expected reward of S2 is probably higher than 10–5=5.
You can also tweak γ to specify how important are the next rewards.
Here is a simple environment which consists of a 5-by-5 grid. A treasure (T) is placed at the bottom right corner of the grid. The agent (O) starts at the top left corner of the grid.
O.......................T
The agent needs to get to the treasure using the 4 available actions : left, right, up, down.
If the agent takes an action that leads him directly to T, he gets a reward of 1, otherwise a reward of 0.
The code is well commented and it is simply what we just discussed. Now the interesting part, the Q-Learning algorithm !
I almost commented every single line of this code, so hopefully, it will be easy to understand!
Put both of the above files in the same directory, and run :
python3 medium_qlearning_rl.py
Around the epoch number 40, the agent should have learned to get to the treasure using one of the shortest paths (8 steps).
We have seen how to derive statistical formulas to find the Bellman equation and used it to teach an AI how to play a simple game. Notice that in this game, the number of possible states is finite (the number of different cells you might end up in), which is why building a Q-Table (a table of values that approaches the real value of the Q function for discrete values) is still manageable. What about a graphical game, such as Flappy Bird, Mario Bros, or Call Of Duty ? Every frame displayed by the game can be considered as a different state. In that case it’s impossible to build a Q-Table, and what we do instead is use a neural network who’s goal will be to learn the Q function. That neural network will typically take as input the current state of the game, and output the best possible action to take in that state. This is known as Deep Q Learning and is exactly how AIs such as Deep Blue or Alpha Go managed to beat world champions at Chess or Go.