State Space

We define a finite set of states:

• State 0 is a terminal trap (absorbing, negative reward)
• State 4 is a terminal treasure (absorbing, positive reward)
• States 1, 2, 3 are non-terminal corridor states

Action Space

At every non-terminal state, the agent may choose:

These actions attempt (we use ‘attempt’ here because of the stochasticity within our MDP) to move the agent along the corridor.

Transition Model

This world is stochastic, meaning that the agents intended actions do not always happen. In this case, the corridor is slippery, and as a result he may slip and move in the other direction.

For any non-terminal state:

This introduces uncertainty and forces the agent to reason in expectation (i.e. most of the time it will go where intended, but sometimes, it may slip and move in the other direction).

Reward Function

The agent shall be rewarded with ‘+1’ when it reaches the treasure. If it becomes immobilised during its expedition, then it receives a reward of ‘-1’. Furthermore, there is a ‘living cost’, which is a very small negative value. The reason for the reward being negative is so the agent understands there is a cost for wandering, and thus is motivated to choose shorter paths.

Discount Factor

Future rewards are discounted by:

Full Python Implementation (Value Iteration)

Now, I shall implement the algorithm in Python. For anyone curious and wanting to implement the code yourself, I recommend going through the code section by section to understand what is being done, and then writing the code line by line, and lastly (which I think is the most important part), implement the code base yourself from scratch.

It builds the transition tensor (which is the probability distribution for the transition process), the reward tensor (which specifies the immediate reward received when the agent takes action in state , and transitions to the next state ); performs value iteration (as per the algorithm above), and then extracts the greedy optimal policy, with respect to for all .

import numpy as np

# --------
# MDP Setup
# --------

S = 5  # we have 5 states
A = 2  # 2 actions: left and right

LEFT, RIGHT = 0, 1
terminal_states = {0, 4}

gamma = 0.95
p_intended = 0.8
p_slip = 0.2
living_cost = -0.04


def is_terminal(s: int) -> bool:
    """Check whether we are in a terminal state."""
    return s in terminal_states


def next_state(s: int, action: int) -> int:
    """
    Terminal state is absorbing.
    Otherwise next state is s-1 (LEFT) or s+1 (RIGHT).
    """
    if is_terminal(s):
        return s
    return s - 1 if action == LEFT else s + 1


# --------
# Build out Transition Distribution and Reward Table
# --------

Tau = np.zeros((S, A, S))
R = np.zeros((S, A, S))


def create_TauR():
    for s in range(S):
        for a in range(A):

            if is_terminal(s):
                Tau[s, a, s] = 1.0
                R[s, a, s] = 0.0
                continue

            INTENDED = LEFT if a == LEFT else RIGHT
            SLIP = RIGHT if INTENDED == LEFT else LEFT

            sp_intended = next_state(s, INTENDED)
            sp_slip = next_state(s, SLIP)

            Tau[s, a, sp_intended] += p_intended
            Tau[s, a, sp_slip] += p_slip

            for sp in [sp_intended, sp_slip]:
                if sp == 4:
                    R[s, a, sp] = 1.0
                elif sp == 0:
                    R[s, a, sp] = -1.0
                else:
                    R[s, a, sp] = living_cost


# --------
# One-Step Bellman Backup: Q(s,a)
# --------

def compute_Q(V: np.ndarray) -> np.ndarray:
    """
    Q[s,a] = sum_{s'} Tau(s'|s,a) * ( R(s,a,s') + gamma * V(s') )
    """
    Q = np.zeros((S, A))
    for s in range(S):
        for a in range(A):
            Q[s, a] = np.sum(Tau[s, a, :] * (R[s, a, :] + gamma * V))
    return Q


# --------
# Value Iteration
# --------

def value_iteration(tol: float = 1e-8, max_iters: int = 10_000):
    V = np.zeros(S)

    for _ in range(max_iters):
        Q = compute_Q(V)
        V_new = V.copy()

        for s in range(S):
            if is_terminal(s):
                V_new[s] = 0.0
            else:
                V_new[s] = np.max(Q[s])

        if np.max(np.abs(V_new - V)) < tol:
            V = V_new
            break

        V = V_new

    # Extract greedy policy
    Q_final = compute_Q(V)
    pi = np.zeros(S, dtype=int)
    for s in range(S):
        if is_terminal(s):
            pi[s] = 0
        else:
            pi[s] = np.argmax(Q_final[s])

    return V, pi


# --------
# Display the Policy with the Values
# --------

def action_name(a: int) -> str:
    return "L" if a == LEFT else "R"


create_TauR()
V, pi = value_iteration()

print("VALUE ITERATION SOLUTION:")
print("----------")
for s in range(S):
    if is_terminal(s):
        print(f"s={s}: V={V[s]: .6f} terminal")
    else:
        print(f"s={s}: V={V[s]: .6f} pi = {action_name(pi[s])}")
print("----------")
print("The Policy is:")
print(pi)

OUTPUT:

VALUE ITERATION SOLUTION:
----------
s=0: V= 0.000000 terminal
s=1: V= 0.321372  pi = R
s=2: V= 0.728121  pi = R
s=3: V= 0.930343  pi = R
s=4: V= 0.000000 terminal
----------
The Policy is:
[0 1 1 1 0]

Explanation of the Output

Now, as a recap, our goal at a high level was simple: guide the agent to the treasure without stepping on the trap and becoming immobilised. Recall that, by definition, we set all terminal states to have value - similar to a Dirichlet-style boundary condition.

So, from this, our task was to recover an optimal policy, which is just a mapping from states to actions that maximises the agent’s expected, discounted future rewards. Also, an important side note: this policy need not be unique. There can be multiple paths that lead to the same optimal outcome - as we find to be true in life.

Looking at the output, we first see that the terminal states , and , the trap and the treasure, respectively, both have values equal to zero, exactly as defined. The lowest non-terminal value appears at . Intuitively, this makes sense since this is the most dangerous place to be for the agent. From here, the agent is close to the trap, and the 20% slip probability introduces a high risk of becoming immobilised. As a result, the optimal action in this state is to move RIGHT, away from danger.

At , our Dr Jones is further from the trap and closer to the treasure. The risk has dropped, and so the value of the state increases. Taking another step RIGHT to yields the highest value among the non-terminal states: the agent is now as far from the trap as possible and one step away from the goal. From here, a final move to the right delivers the agent into the desired terminal state - the treasure!!

In other words, the optimal policy is rather simple: from any non-terminal state, go right until you reach the treasure.

Now what makes this result quite satisfying is that it didn’t come from the designer hard-coding the agents of the action (which would be glaringly obvious when looking at this problem); or trial-and-error’ing’ the solution in a more complicated scenario…instead the behaviour of the agent resulted from the mathematics - effectively, from repeatedly applying the Bellman optimality operator to our value function. I want to delve into the mathematics a bit more now:

At each state, we are asking the following question: “if I was standing in this state, and I could choose the best possible action, what is the maximum expected return that I could obtain from here onward?” This question is encoded into the Bellman optimality equation, which is defined as:

In the code, I broke this up into two separate computations for simplicity:

First, I computed the action-value function, defined as:

This can be interpreted as the agent saying: ‘for each action, I’ll average over all possible next states and obtain a value weighted by the immediate reward, and the discounted value of future rewards.’

Then, for each state, I applied the Bellman Optimality operator, which yields:

i.e. we only keep the best possible action that promises the highest expected return.

Each iteration shrinks the gap between our current estimate, and the true optimal value function . The reason for this is that the Bellman Operator is a contraction mapping, for (which just means that values are being pulled closer and closer together). From this, we can apply the Banach fixed-point theorem (which is why one has to show that the operator one uses, is actually a contraction mapping), to show repeated application of the Bellman operator converges to a unique fixed point, which is .

Now once we check that the values were within a very small tolerance of each other, we can just extract the policy out from there. We simply fix a state, take the most current , and then figure out the action that achieved the maximum value such that:

Conclusion

This was the first toy example I implemented which allowed me to gain a better understanding of the theory in practice.

Moreover, it is important to realise that the dynamic programming algorithm is not feasible for large MDPs, given the effective computational complexity of order (as during the Bellman backup step, for each state and action, we need to sum over all possible next states).

I look forward to continuing my learning and pushing this framework beyond a single corridor, into settings where the state space is much larger, the dynamics are unknwon and the structure has to be learned rather than assumed.

Temporal Difference Learning - c’est parti!

References

• Albrecht, S. V., Christianos, F., & Schäfer, L. Multi-Agent Reinforcement Learning: Foundations and Modern Approaches. MIT Press, 2024.
• Sutton, R. S., & Barto, A. G. Reinforcement Learning: An Introduction