AI Concepts in Python

Reinforcement Learning: Value Iteration and Q-Learning in Python

In reinforcement learning an agent acts in an environment and receives rewards. It learns how good each state or action is, so it can choose the actions that bring the most reward over time.

The concept in plain words

In reinforcement learning an agent acts in an environment and receives rewards. It learns how good each state or action is, so it can choose the actions that bring the most reward over time.

Teach an agent to reach a goal: compute state values with value iteration, then learn them by trial and error with Q-learning.

Key ideas

The discount factor γ makes rewards that come sooner worth more.
Value iteration needs a model of the environment; Q-learning learns from experience.
Q(s, a) is the expected total reward for taking action a in state s.
Exploration (here 20% random moves) is needed so the agent discovers the goal.

When should you use it?

Understand what happens inside neural networks and AI libraries.
Learn the idea before you implement it in the matching challenges.
Prepare for machine learning courses, interviews and projects.

See it in action

Press Play on the left to spread values from the goal through the grid. On the right, watch the agent learn the same route by trial and error.

Example code

Run code →

main.py

import random

# A corridor of 5 cells: start in cell 0, reward +10 for reaching cell 4.
N, GOAL, GAMMA = 5, 4, 0.9

# Value iteration: a cell is worth its reward plus the discounted best neighbour.
values = [0.0] * N
for _ in range(50):
    values = [0.0 if s == GOAL else max((10 if n == GOAL else 0) + GAMMA * values[n]
              for n in (max(0, s - 1), min(N - 1, s + 1))) for s in range(N)]
print("state values:", [round(v, 2) for v in values])

# Q-learning learns the same thing from trial and error.
random.seed(0)
Q = [[0.0, 0.0] for _ in range(N)]   # actions: 0 = left, 1 = right
for episode in range(200):
    s = 0
    while s != GOAL:
        a = random.randrange(2) if random.random() < 0.2 else Q[s].index(max(Q[s]))
        n = max(0, s - 1) if a == 0 else min(N - 1, s + 1)
        reward = 10 if n == GOAL else 0
        target = reward + (0 if n == GOAL else GAMMA * max(Q[n]))
        Q[s][a] += 0.5 * (target - Q[s][a])   # move the estimate toward the target
        s = n
print("best action per cell:", ["right" if q[1] >= q[0] else "left" for q in Q[:GOAL]])
print("Q(0, right):", round(Q[0][1], 2))

Expected output

state values: [7.29, 8.1, 9.0, 10.0, 0.0]
best action per cell: ['right', 'right', 'right', 'right']
Q(0, right): 7.29

How it works

Value iteration repeats the Bellman update, where a cell's value is the best reward plus the discounted value of the next cell; this gives 10, 9, 8.1 and 7.29 moving away from the goal. Q-learning does not know the rules. After every step it moves its estimate toward reward + γ·max Q(next), and it learns to always go right.

Change the values and run the program in the CodeUtility online Python compiler without installing Python locally.

Practice exercises

Modify the runnable example with the exercises below to build understanding beyond copying the result.

Change the inputs or weights and predict the output before running.
Print the intermediate values to follow each step of the calculation.
Solve the linked challenges below to implement the idea yourself.

Run in Python IDE →

Practice with a challenge

Use what this lesson covers to solve a real problem. Each challenge runs in your browser with instant tests and an AI review.

All challenges