Mastering Reinforcement Learning: In-Depth Guide to Model-Based Planning using Dynamic Programming

Introduction

In recent years, reinforcement learning (RL) has emerged as one of the most exciting and promising fields of artificial intelligence. At a high level, RL is concerned with building intelligent agents that can learn via interaction with an environment, similar to how humans and animals learn. By taking actions and receiving feedback in the form of rewards or punishments, an RL agent attempts to figure out the optimal way to behave in order to maximize its cumulative reward over time.

Some of the most impressive AI achievements of the past decade, from AlphaGo defeating world champion Go players to OpenAI Five beating professional Dota 2 teams, have utilized deep reinforcement learning under the hood. Tech giants like Google, Facebook, Microsoft, Amazon and Uber have invested heavily in RL research and applications. As RL systems become more sophisticated and scalable, they have the potential to transform industries like robotics, autonomous driving, industrial automation, healthcare, education, and more.

While the most advanced RL systems today leverage deep learning to tackle complex problems with high-dimensional state and action spaces, it‘s important to understand the foundations that modern RL techniques build upon. And there‘s no better place to start than with the classical dynamic programming (DP) algorithms used for model-based planning in environments that can be modeled as Markov Decision Processes.

In this article, we‘ll take a deep dive into the nuts and bolts of DP methods for solving MDPs, with a special focus on policy iteration and improvement. We‘ll explore the math and intuition behind these foundational algorithms, and walk through a concrete code example to solidify your understanding. Whether you‘re an RL researcher, practitioner or enthusiast, this guide will equip you with a solid grasp of the core concepts and techniques that power much of the field. Let‘s get started!

Markov Decision Processes

The first key concept to understand is that of a Markov Decision Process (MDP). An MDP provides a general mathematical framework for modeling sequential decision making under uncertainty, and is defined by:

• A set of possible states S the environment can be in
• A set of possible actions A the agent can take
• A state transition function T(s, a, s‘) mapping a current state s and action a to a probability distribution over next states s‘
• A reward function R(s, a) giving the immediate reward for taking action a in state s
• A discount factor γ between 0 and 1 that weights future rewards

At each timestep, the agent observes the current state, takes an action, receives a reward, and the environment transitions to a new state. This process repeats in a cycle until reaching some terminal state, constituting one episode of interaction. The goal is to find a policy π, or mapping from states to actions, that maximizes the expected cumulative discounted reward:

MDP diagram

A key property of MDPs is that they satisfy the Markov property – the next state and reward depend only on the current state and action, and not on the previous history. This allows us to decompose the overall optimization problem into subproblems for each state. We can define a value function Vπ(s) that represents the expected return starting from state s and following policy π:

Vπ(s) = E[Rt + γRt+1 + γ^2 Rt+2 + … | St=s] = E[Σk=0 to ∞ γ^k Rt+k | St=s]

We can also define an action-value function Qπ(s, a) as the expected return starting from s, taking action a, and then following policy π:

Qπ(s, a) = E[Rt + γ Vπ(St+1) | St=s, At=a] = R(s,a) + γ Σs‘ T(s, a, s‘) * Vπ(s‘)

The value function and Q-function are recursively related by the Bellman expectation equations:

Vπ(s) = Σa π(s,a) Qπ(s,a)
Qπ(s,a) = R(s,a) + γ Σs‘ T(s, a, s‘) * Vπ(s‘)

These equations allow us to express the value of a state in terms of the values of its successor states, and form the basis of the dynamic programming algorithms we‘ll explore next for finding the optimal value function and policy.

Policy Evaluation

Suppose we‘re given an arbitrary policy π that maps states to action probabilities. How can we compute the value function Vπ that predicts the expected return from each state when following this policy? The answer is to turn the Bellman expectation equation into an iterative update:

Vπ(s) ← Σa π(s,a) (R(s,a) + γ Σs‘ T(s, a, s‘) * Vπ(s‘))

By repeatedly applying this update, we can improve our estimate of Vπ until it converges to the true value function for π. This algorithm is known as iterative policy evaluation, and computes the value of each state under the current policy.

Policy evaluation backup diagram

For example, consider the classic FrozenLake environment, where an agent navigates a 4×4 grid of slippery ice and aims to reach the goal state without falling into holes. We can visualize the value of each state after running policy evaluation on a uniform random policy:

Policy evaluation on FrozenLake

The value increases as the agent gets closer to the goal, and decreases near holes. Even this basic policy already produces some useful information about which states are most promising. But we can do even better by improving the policy itself.

Policy Improvement

Given a policy π and its value function Vπ, we can improve the policy by being greedy with respect to the current value function. The idea is to consider selecting a different action a‘ in each state s that maximizes the expected return:

π‘(s) = argmax_a‘ Qπ(s, a‘)
= argmax_a‘ (R(s,a‘) + γ Σs‘ T(s, a‘, s‘) Vπ(s‘))

If this new greedy policy π‘ is an improvement over π, we can continue the process and 1) evaluate π‘ to get Vπ‘, and 2) improve π‘ to get an even better π‘‘. We keep alternating between policy evaluation and policy improvement until the policy converges and no longer changes. At this point, we have arrived at the optimal policy π* that maximizes long-term reward from any starting state.

Policy improvement diagram

This is the key idea behind policy iteration, a dynamic programming algorithm that provably converges to the optimal policy and value function for any finite MDP. In code, we can implement policy iteration as follows:

def policy_iteration(env, gamma=0.99, theta=1e-8):
    policy = np.ones([env.nS, env.nA]) / env.nA
    while True:
        V = policy_evaluation(policy, env, gamma, theta)
        policy_stable = True
        for s in range(env.nS):
            old_action = np.argmax(policy[s])
            Q_vals = one_step_lookahead(env, V, s, gamma)
            best_action = np.argmax(Q_vals)
            policy[s] = np.eye(env.nA)[best_action]
            if old_action != best_action:
                policy_stable = False
        if policy_stable:
            return policy, V

def policy_evaluation(policy, env, gamma, theta):
    V = np.zeros(env.nS)
    while True:
        delta = 0
        for s in range(env.nS):
            v_old = V[s]
            new_v = 0
            for a, a_prob in enumerate(policy[s]):
                for prob, next_state, reward, done in env.P[s][a]:
                    new_v += a_prob * prob * (reward + gamma * V[next_state])
            V[s] = new_v
            delta = max(delta, abs(v_old - V[s]))
        if delta < theta:
            break
    return V

def one_step_lookahead(env, V, state, gamma):
    Q_vals = np.zeros(env.nA)
    for a in range(env.nA):
        for prob, next_state, reward, done in env.P[state][a]:
            Q_vals[a] += prob * (reward + gamma * V[next_state])
    return Q_vals

Here we initialize an arbitrary policy and alternate between policy evaluation and policy improvement steps until convergence. Running this on the FrozenLake environment quickly yields the optimal policy after just a few iterations:

Optimal policy for FrozenLake

The agent has learned to navigate directly to the goal while avoiding holes along the shortest path. Policy iteration is effective and guaranteed to converge, but requires iterating the policy evaluation step to completion, which can be expensive. An alternative is to use value iteration, which combines policy evaluation and improvement into a single step.

Value iteration backup diagram

In value iteration, we start with an arbitrary value function and iteratively apply the Bellman optimality equation as an update:

V(s) ← max_a (R(s,a) + γ Σs‘ T(s, a, s‘) V(s‘))

After convergence, we can extract the optimal policy by being greedy with respect to the optimal value function:

π(s) = argmax_a (R(s,a) + γ Σs‘ T(s, a, s‘) V(s‘))

Value iteration often converges faster than policy iteration in practice. Here is a concise implementation:

def value_iteration(env, gamma=0.99, theta=1e-8):
    V = np.zeros(env.nS)
    while True:
        delta = 0
        for s in range(env.nS):
            v_old = V[s]
            Q_vals = one_step_lookahead(env, V, s, gamma)
            V[s] = max(Q_vals)
            delta = max(delta, abs(v_old - V[s]))
        if delta < theta:
            break
    policy = np.zeros([env.nS, env.nA])
    for s in range(env.nS):
        Q_vals = one_step_lookahead(env, V, s, gamma)
        best_action = np.argmax(Q_vals)
        policy[s, best_action] = 1.0
    return policy, V

On FrozenLake, value iteration produces the same optimal policy as policy iteration, but in fewer iterations. However, neither algorithm is efficient for large state spaces, as they both require sweeping over the entire state space at each iteration. Various asynchronous DP methods have been proposed to alleviate this issue, as well as function approximation techniques to extend DP to high-dimensional environments.

Recent Advances and Applications

Although DP laid the foundations for reinforcement learning, today most practical RL systems use sample-based methods like Q-learning, policy gradients, and actor-critic algorithms that learn from experience without requiring a model of the environment. However, DP still plays an important role in several ways:

• DP serves as the basis for understanding and analyzing RL algorithms. Many sample-based methods can be viewed as stochastic approximations to DP that converge to the same optimal values and policies.

• DP is used as a subroutine for planning in model-based RL, where a learned model of the environment dynamics is used to simulate Experience and perform lookahead search. Notable examples include Dyna architectures, Monte Carlo Tree Search, MuZero, and AlphaZero.

• DP is used to solve subproblems in hierarchical RL, where a high-level policy selects temporally extended actions that are themselves optimized using DP over a low-level model.

• DP can be used offline to compress or bootstrap a value function that is then used as a prior for online RL. This can accelerate learning on complex tasks.

• In situations where a perfect environment model is available (e.g. board games, robotic simulators, operational research), DP alone is sufficient to solve the MDP and is often more efficient than RL.

Some active areas of research that leverage DP include transfer learning, lifelong learning, multi-task RL, and meta RL. As models, algorithms and hardware improve, the scope of problems that can be tackled by DP and RL will continue to expand.

Conclusion

In this article, we took a deep dive into the classical DP algorithms of policy iteration and value iteration for solving MDPs. We walked through the key steps of policy evaluation and policy improvement, building up from the Bellman equations to full implementations of each algorithm.

While DP is not always feasible for large-scale problems, understanding these foundational methods is crucial for anyone serious about RL. They provide a clear, principled framework for sequential decision making under uncertainty, and embody the key ideas of long-term value, bootstrapping, and the policy-value relationship that are at the core of more advanced RL techniques.

I hope this guide helped clarify DP and its role within the broader landscape of RL. Despite the field‘s rapid progress, these classic algorithms remain a vital part of the researcher and practitioner‘s toolkit. Equipped with a firm grasp of DP, you‘re now ready to explore the more flexible sample-based algorithms that power the most impressive AI systems of the modern era. Happy learning!