AI Research ResidencyCurriculum ↗

Reinforcement Learning

Policy Iteration

An interactive guide to how AI learns the best strategy

World
2×2 grid
Step cost
−1
Discount
0.9
Converges in
2 rounds

How does an AI figure out the best way to do something? One of the oldest and most elegant answers is policy iteration: try a strategy, honestly evaluate how well it performs, make it a little better, and repeat. Same loop you'd use to improve a recipe. Cook it, taste it, adjust, cook again. Below, you'll step through a concrete example and watch it converge on the optimal answer.

The grid world used throughout, and the worked example it is built on, come from Professor George A. Lentzas's machine learning and artificial intelligence course at Columbia Business School.

01

The Big Idea

Decision-Making as a Game

Imagine you're navigating a city you've never visited. At every intersection you have to choose: go left, right, straight, or turn around. Some choices bring you closer to your hotel, others take you in circles. Every wrong turn wastes time (a small penalty), and arriving at the hotel is the reward.

In AI, this setup has a name: a Markov Decision Process. The intersections are states, your turns are actions, the wasted time is a negative reward, and your plan for which way to turn at each intersection is your policy. The goal: find the policy that gets you to the hotel with the least total cost.

There's one more ingredient: a discount factor. It captures the common-sense idea that a reward right now is worth more than the same reward far in the future. A discount of 0.9 means that a reward two steps away is only worth 81% of a reward right now (0.9 × 0.9 = 0.81).

The Evaluate-Improve Loop

Policy iteration solves this problem in two alternating phases, repeated until nothing changes:

Phase 01 · Evaluate

Evaluate the current strategy

“If I follow this plan forever, how much total cost will I rack up from each starting point?” This gives every room a score: a single number capturing how good or bad it is to be there.

Phase 02 · Improve

Look for better options

At each room, try every possible direction and ask: “Would going this way lead somewhere better?” If yes, switch to that direction.

When the improvement step finds nothing to change, the strategy is the best it can possibly be.

EvaluateScore each state
ImprovePick better actions
Done?Nothing changed
02

The Example: A Tiny Grid World

To see policy iteration in action, we'll use the simplest possible world: a 2×2 grid. Think of it as four rooms. The bottom-right room (G) is the goal, like finding your hotel. The other three rooms (A, B, C) are just places you might be.

The rules are simple: you can move Up, Down, Left, or Right. If you try to walk through a wall, you stay put (but still pay the penalty). Every move costs −1 point, except stepping into the goal which is free. The discount factor is 0.9. One thing to watch for: if a state's policy points it into a wall or a loop so it never reaches the goal, the penalties pile up forever and its value drops to about −10 (the worst possible score in this grid).

Astart
B
C
Ggoal
4 actions (Up, Down, Left, Right) · Walls = bounce back · Each move costs −1 · Goal is free · Discount = 0.9
03

Interactive Walkthrough

Step through the algorithm below. The grid shows each cell's score (how bad it is to be there) and an arrow for the current action. Use the buttons to navigate, or hit Play to auto-advance.

Aunscored
Bunscored
Cunscored
G★0
not scored yet
stuck: never reaches the goal
reaches the goal at a cost
at the goal, or free to reach
SETUP

The Grid

Here’s a simple 2×2 grid with four rooms: A (top-left), B (top-right), C (bottom-left), and G (bottom-right). G is the goal, and it has a score of 0 because you’re already there. Our job: figure out which direction each room should point so you reach the goal as fast as possible.

01 / 18

Press Play to watch the algorithm run, or step through one at a time

04

Key Takeaways

  1. 01

    It converges fast

    Even starting from a terrible strategy (every action bumps into a wall), the algorithm found the optimal answer in just two rounds. This isn't a fluke. Policy iteration typically converges in very few iterations, even on much larger problems.

  2. 02

    Improvement is guaranteed

    Each round of improvement is mathematically guaranteed to produce a strategy at least as good as the one before. Since there are only so many possible strategies, the algorithm must eventually land on the best one. You can't get worse, only better or the same.

  3. 03

    Policy iteration vs. value iteration

    There's a sibling algorithm called value iterationthat skips the full evaluation step and instead makes small updates to the scores in every pass. It's simpler to implement but usually needs more passes to finish. Policy iteration needs fewer rounds but each round is more work. In practice, the choice depends on the problem.

05

Policy Iteration in the Real World

Our grid had 4 rooms. Real problems have millions or billions of possible situations. But the algorithm is the same: score your current strategy, look for improvements, repeat. Here are three systems that use this exact loop at massive scale.

Language Models

How ChatGPT & Claude Learn to Write Better

Language models start by learning patterns from text on the internet, but they need refinement to be genuinely helpful. The method is called Reinforcement Learning from Human Feedback (RLHF), and it's the evaluate-improve loop in disguise.

The model generates responses to thousands of prompts. Human raters read pairs of responses and pick the better one. That's evaluation. Then the model is updated to produce more responses like the preferred ones. That's improvement. One common algorithm for this update is PPO(Proximal Policy Optimization), which makes small, careful changes to avoid breaking what already works, like adjusting a recipe one ingredient at a time. The cycle repeats until the model's responses consistently score well.

Generate responsesPolicy in action
Humans rate themEvaluate
Update modelImprove
RepeatUntil stable
Grid WorldLanguage Model
Rooms (A, B, C)Possible conversations
Which direction to goWhich word to write next
−1 per extra moveHuman preference score
Direction for each roomHow the model generates text
Score for each roomPredicted response quality

Games

How AlphaGo Mastered the World's Hardest Board Game

In 2016, DeepMind's AlphaGo defeated the world champion at Go, a game with more possible board positions than atoms in the observable universe. Brute-force search was impossible. Instead, AlphaGo used policy iteration.

It trained two neural networks: a policy network (which move to play) and a value network (how good is this board position, just like our V-values). It played millions of games against itself. After each batch of games, it evaluated positions with the value network, improved the policy network based on which moves led to wins, and repeated.

The same loop that found the optimal path through a 2×2 grid found superhuman Go strategy. The only difference is scale.

Evaluate positionValue network
Select best movePolicy network
Play gameSelf-play
Update networksImprove
Grid WorldGo
Rooms (A, B, C)Board positions (10170 possibilities)
Which direction to goWhere to place a stone (361 intersections)
−1 per extra moveWin (+1) or loss (−1)
Direction for each roomNeural network that picks moves
Score for each roomNeural network that judges positions

Transportation

How Self-Driving Cars Learn to Navigate

A self-driving car faces a continuous stream of decisions: brake, accelerate, change lanes, yield to a pedestrian. Companies like Waymo train driving policies in simulation before deploying them to real roads.

The car's sensors read the road (the state). The policy decides what to do (the action). A simulator scores the outcome. Was the ride safe? Efficient? Comfortable? That is the reward. The policy is updated based on millions of simulated drives, improving its decision-making for every scenario it might encounter. This is policy iteration at industrial scale.

Read sensorsState
Decide actionPolicy
Simulate driveEvaluate
Score & updateImprove
Grid WorldSelf-Driving Car
Rooms (A, B, C)Sensor readings (cameras, lidar, radar)
Which direction to goSteering, braking, acceleration
−1 per extra moveSafety, efficiency, comfort scores
Direction for each roomDriving decisions for each road situation
Score for each roomPredicted safety of each road situation

The deepest lesson isn't the math. It's the structure. Evaluate honestly, improve incrementally, repeat until stable. This pattern shows up wherever complex decisions need to be optimized, from training AI systems to how companies run A/B tests, athletes review game film, and scientists refine hypotheses.

06

Try It Yourself

Now try a bigger challenge. Below is a 4×4 grid with 15 states and a goal in the bottom-right corner. Click a cell to select it, then click again (or use the arrow keys) to cycle its direction.

Then run the algorithm on whatever policy you wrote. Evaluate only scores it without changing anything. Improve once runs a single greedy sweep, switching each state to its best-looking action, and the changed cells are outlined. Run policy iteration alternates the two until nothing changes, which is the entire algorithm. However bad your starting policy, it lands on the optimal one in a handful of rounds.

1↑
2↑
3↑
4↑
5↑
6↑
7↑
8↑
9↑
10↑
11↑
12↑
13↑
14↑
15↑
G★0
Click a cell to cycle its direction

Direction for state 1

Click a selected cell to cycle it, or use the arrow keys.

Where does this direction go?

1↑1 (wall, stays put)

Current policy

↑↑↑↑↑↑↑↑↑↑↑↑↑↑↑★