Introducing Markov Decision Processes, Setting up Gymnasium Environments and Solving them via Dynamic Programming Methods

In a previous post we started our series about Reinforcement Learning (RL) following Sutton’s great book [1]. In that post we introduced RL in general, and discussed Multi-armed Bandits as a nonassociative toy problem.

Here, we will build on this – but go significantly beyond. In particular, we will introduce our first associative problem, which might feel much more like "real" RL to many readers – and introduce a simple but general solution technique. Furthermore, we will introduce Gymnasium [2], a powerful library providing a multitude of environments (e.g. Atari or MuJoCo games) and allowing us to quickly experiment with solving them.

The previously mentioned associative setting is the "standard" in RL: as opposed to the previously introduced nonassociative setting where there is only a single state, and we only have to decide on what action to take, here we have multiple states – and for every state we might decide for a different best action.

In particular, in this post we will discuss:

Markov decision processes (MDPs) to formalize RL problems
the Gymnasium library
and the first solution method given by Sutton, namely Dynamic Programming (DP)

The full code can be found on github.

With that said, let’s dive in!

MDPs

In this section we will introduce Markov decision processes, a crucial concept in the field of RL. MDPs are a framework for modelling sequential decision making. Important terms are agent and environment, as well as states, actions, and rewards. In the views of an MDP, we want to model and control an agent, who moves within an environment. At each timestep, the agent is in a certain state, and can chose to execute one of several available actions. Based on the state and action, the agent will end up in a new state (which can include randomness) – and observe a reward. Sutton visualizes this relations as follows:

Mathematically, an MDP is a 4-tuple (S, A, P, R):

S and A denote the available states
P is the transition function:

R the reward function:

One distinguishes between episodic and continuous tasks: episodic tasks are guaranteed to terminate, while continuous ones go on indefinitely. However, it is possible to turn all episodic tasks into continuous ones, thus unifying notation and allowing us to treat both task types identically. For that, we simply add a novel "sink" state, from which all actions end up in this state again yielding 0 reward.

Goal of the agent is to maximize its accumulated reward – thus defining a good reward function is crucial when solving any RL problem. And – this is by no means simple: one can decide whether to only pass final rewards and 0 otherwise (e.g. 1 for winning a game of chess), or intermediate ones (e.g. -0.1 for each timestep an agent is stuck in a maze, to encourage fast escapes). Further, one needs to be careful – it is common that agents learn to "cheat", finding a reward-maximizing but unexpected behavior.

To formalize the notation of "accumulated reward", we introduce the notion return – which is the sum of all collected rewards during an episode:

But, we need another concept, namely that of discounting. Each future reward is discounted via a power of γ, with 0 < γ < 1:

We discount to express a preference for the present (a reward r now is better than a reward r many steps in the future ) – and: to be able to handle continuous tasks (otherwise all rewards would go to infinity).

Next, let’s define policies and value functions: a policy defines an action distribution per state, formally:

Thus, a policy describes preferences over which actions to take in a state. If π is a one-hot vector, the policy is deterministic.

The value function describes the expected return when starting in a state s and following policy π from then on:

Similar, the action-value function measures the expected return when starting in state s, taking action a, and then following π from then on:

Introduction to Gymnasium

Gymnasium [2], formerly known as Gym, is an amazing simulator which allows everyone to quickly develop and test RL algorithms on a multitude of problems. Originally implemented by OpenAI (the masterminds behind, e.g. ChatGPT), it is now known as Gymnasium and maintained by the Farama foundation. Around it, there are several top-notch tutorials, e.g. about RL algorithms in general, or how to use it.

With Gymnasium, you can quickly spin up a multitude of environments, such as the classical cart pole and mountain car problems, or MuJoCo and Atari tasks.

For our introduction, we will use "GridWorld" – a discrete maze environment with some obstacles, in which one has to find the exit. To be precise, we find our agent placed on an icy surface split into fields. On one such field, there is a present we have to reach. Along the way there are several open lakes – falling into which will end the episode:

We instantiate the environment as such:

import gymnasium as gym

env = gym.make(
    "FrozenLake-v1",
    desc=None,
    map_name="4x4",
    is_slippery=False,
    render_mode="human",
)

We will use this instance / specific environment to show-case and introduce Gymnasium. First note how the environment is selected – passing a different name will give us a different environment. render_mode "human" yields a live rendering of the game we can follow. We will use this to demonstrate our found solution – however, when in the process of finding a good policy, the rendering naturally is disabled.

Let’s discover the available observation and action spaces:

ipdb> env.observation_space Discrete(16)

ipdb> env.action_space Discrete(4)

As we can see, this environment describes an observation (similar to "state") as a discrete number in the range 0..15 – one for each of the 4×4 grid cells. Further, there are four actions – up, down, left, right.

Via env.step(action) we execute said action. This returns a 5 tuple, which here in particular has the following meaning:

observation: the new state
reward: the observed reward (1 if goal is reached, 0 otherwise)
terminated: true when we reached the goal or fell into a lake
truncated: true when termination was caused by a condition outside the MDP, such as a time limit (not relevant to us here)
info: additional debugging output (not relevant to us here)

Let’s write a simple program picking actions at random and playing the episode until we either reach the goal or fall into a lake:

import gymnasium as gym
import numpy as np

env = gym.make(
    "FrozenLake-v1",
    desc=None,
    map_name="4x4",
    is_slippery=False,
    render_mode="human",
)

observation, _ = env.reset()

while True:
    action = np.random.randint(env.action_space.n)
    observation, reward, terminated, truncated, info = env.step(action)
    print(
        f"Step {t} === action: {action}, observation: {observation}, reward: {reward}, terminated: {terminated}, truncated: {truncated}, info: {info}"
    )
    if terminated or truncated:
        break

env.close()

Have fun with your first Gymnasium environment, and enjoy the agent running around the icy world (and falling into a lake very quickly – the chances of reaching the goal at random are slim, ~0.02%).

So with that, let’s move to the next section: we now know how to use Gymnasium, and have also seen that – even for such a small problem – we cannot rely on chance or brute-force to find a good solution. In the next section we will start with the first batch of solution methods, namely DP.

Dynamic Programming

Dynamic Programming is a general solution principle in the world of computer science – denoting the idea of breaking down a larger problems recursively into subproblems.

In the context of RL, it essentially means turning the Bellman equation into an update rule and computing action-value functions using this. Due to this we require a perfect model of the environment. Further this process is rather expensive computationally. Still, DP methods are essential for understanding and provide important theoretical foundations. In fact, all methods to come can be understood as approximations of DP methods – only not requiring a perfect world model and with less compute.

Policy Evaluation (Prediction)

We begin this section by defining how to compute the action-value function. From the previous section we know:

Here, we will compute this value – and do so iteratively. We apply the Bellmann equation as an update rule, and compute for each state:

v_k = v_pi is a fixed point of this formula (the Bellmann equation assures us of this), thus in the limit we indeed converge to the correct state-value estimate.

Let’s have a look at Sutton’s pseudocode:

We will see how this looks like in Python in the next section, when we use it as a subroutine for finding better policies.

Policy Improvement

Knowing the value of a certain policy, of course now it is interesting to know how to improve this policy – a process called policy improvement. Or, in other words: from the previously introduced prediction problem we now move on to the control problem.

For this, we change one action a != π(s). Now, the value of behaving in this manner is:

If this is better than the value of π, we are better off changing our policy in such a way. And we can indeed answer such question via the policy improvement theorem, which states that if, for all states s:

The improved policy is indeed "better", meaning:

Policy Iteration

Knowing how one can do one step of policy improvement, it is natural to chain these improvement steps with policy evaluation, yielding the following scheme of steps:

Each of the resulting policies are monotonically improving, and – since finite MDPs only have finite states – this process must converge to the optimal policy eventually!

Thus, we have our first algorithm for solving an RL problem – called policy iteration.

Let’s concretize this some more by showing Sutton’s pseudocode:

And here’s the corresponding Python code:

def policy_iteration(env: ParametrizedEnv) -> np.ndarray:
    """Uses 'Policy Iteration' to solve the RL problem
    specified by the passed Gymnasium env.

    Args:
        env: env containing the problem

    Returns:
        found policy
    """
    # Make mypy happy ...
    assert isinstance(env.env.observation_space, Discrete)
    observation_space: Discrete = env.env.observation_space
    assert isinstance(env.env.action_space, Discrete)
    action_space: Discrete = env.env.action_space

    pi = np.zeros(observation_space.n).astype(np.int32)

    def _policy_evaluation() -> np.ndarray:
        """Run's policy evaluation - i.e. evaluates the current
        policy pi, and updates the value estimate V.
        """
        V = np.zeros(observation_space.n)
        while True:
            delta = 0
            for s in range(observation_space.n):
                v = V[s]
                V[s] = sum(
                    [
                        p * (r + env.gamma * V[s_next])
                        for p, s_next, r, _ in env.env.P[s][pi[s]]  # type: ignore
                    ]
                )
                delta = max(delta, abs(v - V[s]))
            if delta < env.eps:
                break
        return V

    while True:
        V = _policy_evaluation()

        policy_stable = True
        for s in range(observation_space.n):
            old_a = pi[s]
            pi[s] = np.argmax(
                [
                    p * (r + env.gamma * V[s_next])
                    for a in range(action_space.n)
                    for p, s_next, r, _ in env.env.P[s][a]  # type: ignore
                ]
            )
            if old_a != pi[s]:
                policy_stable = False

        if policy_stable:
            return pi

Value Iteration

One drawback of policy iteration is that each step contains a complete policy evaluation step, a tedious and time-consuming process itself. One might wonder – can we do better – can we maybe get away without full policy evaluation? Sutton run some experiments in which the number of steps in policy evaluation is varied – and as one can see, going beyond three update steps hardly makes any difference.

A special case of policy iteration is when policy evaluation is run only for a single step each time. This is called value iteration:

This can be viewed as turning the Bellman equation (see above) into an update rule. And, in fact, this algorithms works, and we do not lose any of the previously established convergence guarantees.

Let’s have a look at pseudocode:

And Python code:

def value_iteration(env: ParametrizedEnv) -> np.ndarray:
    assert isinstance(env.env.observation_space, Discrete)
    observation_space: Discrete = env.env.observation_space
    assert isinstance(env.env.action_space, Discrete)
    action_space: Discrete = env.env.action_space

    V = np.zeros(observation_space.n)

    while True:
        delta = 0
        for s in range(observation_space.n):
            v = V[s]
            V[s] = max(
                [
                    p * (r + env.gamma * V[s_next])
                    for a in range(action_space.n)
                    for p, s_next, r, _ in env.env.P[s][a]  # type: ignore
                ]
            )
            delta = max(delta, abs(v - V[s]))
        if delta < env.eps:
            break

    return np.asarray(
        [
            np.argmax(
                [
                    p * (r + env.gamma * V[s_next])
                    for a in range(action_space.n)
                    for p, s_next, r, _ in env.env.P[s][a]  # type: ignore
                ]
            )
            for s in range(observation_space.n)
        ]
    )

Results

Both methods above we can run e.g. via this main function:

import argparse

import gymnasium as gym

from dp import policy_iteration, value_iteration
from env import ParametrizedEnv

GAMMA = 0.97
EPS = 0.001
NUM_STEPS = 100

def solve_grid_world(method: str) -> None:
    """Solve the grid world problem using the chosen solving method.

    Args:
        method: solving method
    """
    gym_env = gym.make(
        "FrozenLake-v1",
        desc=None,
        map_name="4x4",
        is_slippery=False,
        render_mode="human",
    )
    env = ParametrizedEnv(gym_env, GAMMA, EPS)

    # Find policy
    if method == "policy_iteration":
        pi = policy_iteration(env)
    elif method == "value_iteration":
        pi = value_iteration(env)
    else:
        raise ValueError(f"Unknown solution method {method}")

    # Test policy and visualize found solution
    observation, _ = env.env.reset()
    for _ in range(NUM_STEPS):
        action = pi[observation]
        observation, _, terminated, truncated, _ = env.env.step(action)
        if terminated or truncated:
            break
    env.env.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process a string input.")
    parser.add_argument("--method", type=str, required=True, help="A string input")
    args = parser.parse_args()

    solve_grid_world(args.method)

Note that you can also find the complete code on github.

In such a simple setting, we cannot observe much of a difference between these two methods (e.g. w.r.t. convergence speed) – both always find the correct and shortest solution.

Conclusion

Let’s recap what we have seen in this post: this is part two in a series covering the book "Reinforcement Learning" by Sutton and Barto. Whereas we considered nonassociative tasks in part one, here we moved to associative tasks – which feel a lot more like "real" RL.

We started by introducing Markov decision processes (MDPs) to model sequential decision problems in general. We introduced general concepts like agents and returns, and introduced the value and action-value function.

Next, we introduced Gymnasium [2], a great library for deploying RL algorithms and playing around with different solution methods. As a special example we used throughout the post we selected the environment "GridWorld", in one has to escape from an icy maze.

Lastly, we introduced our first set of solution methods for solving general RL problems, namely dynamic programming (DP). DP methods are essential to understand the basics of RL: they require a perfect world model and iteratively construct the optimal solution.

In the next post we will cover Monte Carlo methods – which do not need perfect knowledge of the environment, but are able to learn from experience alone.

Other Posts in this Series