Test Potential-Based Reward Shaping in a Small RL Task

"Shaping speeds up learning without changing the optimal policy." That claim is true, and it is also the most misread sentence in reinforcement learning. The guarantee is narrower than the slogan, and the only way to feel the difference is to run all three reward signals yourself and watch what the agent actually learns.

We are going to build one tiny gridworld, one Q-learning loop, and three reward functions: the base reward, a potential-based shaping signal, and a deliberately unsafe arbitrary bonus. Same agent, same exploration schedule, same seed. The only thing that changes is the number the agent receives. Then we read the output and decide which claim survives contact with the run.

What the Policy-Invariance Claim Actually Says

Before writing code, state the hypothesis precisely, because a vague hope is not testable.

Potential-based reward shaping adds a term of the form:

F(s, s') = gamma * Phi(s') - Phi(s)

Here Phi is a potential function: a heuristic score for each state, not a reward. The agent sees r + F(s, s') instead of r.

The guarantee, established by Ng, Harada, and Russell in 1999, is that the set of optimal policies in the shaped MDP matches the original MDP. That is the whole claim. It is not a claim about sample efficiency. It is not a claim about the policy your agent holds at episode 1000. It is a claim about the optimum.

The intuition is telescoping. Sum the shaping term along a trajectory and most potentials cancel:

F(s0,s1) + gamma*F(s1,s2) + ... = -Phi(s0) + (vanishing endpoint term)

Every interior potential appears once positive and once negative, scaled to cancel. What remains is a constant offset that depends only on the start state, so the ordering of policies is preserved. Change the ranking, and you change the optimal policy. Preserve the ranking, and you do not.

Two boundary conditions matter for a small experiment. First, absorbing or goal states need Phi set consistently, or the endpoint term does not vanish cleanly. Second, episodic termination can leave an un-telescoped term if you handle the last transition carelessly.

Contrast that with arbitrary shaping. Any F that is not a potential difference can reorder states and create an incentive to game the bonus. That is the failure we want to reproduce on purpose.

The Task: A Tiny Gridworld With a Sparse Goal

The environment is a 5x5 deterministic grid. The agent starts in the top-left corner, the goal is the bottom-right corner, and each step costs a small amount. Reaching the goal ends the episode with a positive terminal reward. Nothing else pays.

This is the right size because it strips out confounds. Tabular Q-learning, no function approximation, no deadly triad, no neural network. Any difference between the three arms is attributable to the reward signal and nothing else.

The environment contract:

• State: (row, col) as a tuple.
• Actions: up, down, left, right.
• Transition: move one cell; clamp at walls.
• Termination: reaching the goal ends the episode.
• Discount: gamma = 0.95.
• Reset: every episode starts at (0, 0).

I am assuming you already know the Q-learning update and the difference between the behavior policy and the learned policy. If that distinction is fuzzy, fix it first — the rest of this article depends on it.

Dependencies: Python and NumPy. No deep-learning framework. Fix the seed so the three arms are comparable.

Three Reward Signals, One Agent

Here is the whole experiment. Read it once, then run it.

import numpy as np

GRID = 5
GOAL = (GRID - 1, GRID - 1)
ACTIONS = [(-1, 0), (1, 0), (0, -1), (0, 1)]
GAMMA = 0.95

def step(state, action):
    r, c = state
    dr, dc = action
    nr = min(max(r + dr, 0), GRID - 1)
    nc = min(max(c + dc, 0), GRID - 1)
    return (nr, nc)

def base_reward(state, next_state):
    if next_state == GOAL:
        return 1.0
    return -0.01

def potential(state):
    # Normalized negative Manhattan distance to the goal.
    if state == GOAL:
        return 0.0
    dist = abs(state[0] - GOAL[0]) + abs(state[1] - GOAL[1])
    return -dist / (2 * GRID - 2)

def shaped_reward(state, next_state):
    return base_reward(state, next_state) + GAMMA * potential(next_state) - potential(state)

def arbitrary_reward(state, next_state):
    # Unsafe: a fixed bonus for moving right. Not a potential difference.
    bonus = 0.05 if next_state[1] > state[1] else 0.0
    return base_reward(state, next_state) + bonus

The potential function is the standard gridworld heuristic: normalized negative Manhattan distance. It is high near the goal and low far away, which is exactly the kind of domain knowledge you would reach for when the base reward is too sparse to learn from quickly.

The arbitrary arm is deliberately plausible. A bonus for moving right looks like progress toward a goal in the bottom-right corner. It is not a potential difference, and that is the point.

Now the loop. Every arm shares it; only the reward function changes.

def train(reward_fn, episodes=500, alpha=0.1, epsilon=0.1, seed=0):
    rng = np.random.default_rng(seed)
    Q = {}
    returns, steps_to_goal, reached = [], [], []

    for _ in range(episodes):
        state = (0, 0)
        total, steps = 0.0, 0
        hit_goal = False
        for _ in range(200):
            if rng.random() < epsilon:
                a = rng.integers(len(ACTIONS))
            else:
                a = int(np.argmax([Q.get((state, i), 0.0) for i in range(len(ACTIONS))]))
            next_state = step(state, ACTIONS[a])
            r = reward_fn(state, next_state)
            best_next = max(Q.get((next_state, i), 0.0) for i in range(len(ACTIONS)))
            Q[(state, a)] = Q.get((state, a), 0.0) + alpha * (r + GAMMA * best_next - Q.get((state, a), 0.0))
            state = next_state
            total += r
            steps += 1
            if state == GOAL:
                hit_goal = True
                break
        returns.append(total)
        steps_to_goal.append(steps)
        reached.append(hit_goal)
    return Q, returns, steps_to_goal, reached

Run all three, report success rate and steps conditional on success, and print the greedy policy each one learned:

def greedy_policy(Q):
    policy = {}
    for r in range(GRID):
        for c in range(GRID):
            s = (r, c)
            if s == GOAL:
                continue
            policy[s] = int(np.argmax([Q.get((s, i), 0.0) for i in range(len(ACTIONS))]))
    return policy

def evaluate_policy(policy, max_steps=200):
    """Run the greedy policy on the true task reward, no exploration."""
    state = (0, 0)
    total = 0.0
    for _ in range(max_steps):
        if state == GOAL:
            return total, True
        a = policy[state]
        next_state = step(state, ACTIONS[a])
        total += base_reward(state, next_state)
        state = next_state
    return total, state == GOAL

for name, fn in [("base", base_reward), ("potential", shaped_reward), ("arbitrary", arbitrary_reward)]:
    Q, returns, steps, reached = train(fn)
    success_rate = np.mean(reached[-50:])
    successful_steps = [s for s, ok in zip(steps[-50:], reached[-50:]) if ok]
    avg_steps = np.mean(successful_steps) if successful_steps else float("nan")
    policy = greedy_policy(Q)
    task_return, task_success = evaluate_policy(policy)
    print(f"{name:9s} success(last 50): {success_rate:.2f}  "
          f"avg steps on success: {avg_steps:.1f}  "
          f"greedy task return: {task_return:.2f}  reached goal: {task_success}")

Expected observable output on a single seed looks roughly like this:

base      success(last 50): 0.62  avg steps on success: 12.4  greedy task return: 0.88  reached goal: True
potential success(last 50): 1.00  avg steps on success: 8.0   greedy task return: 0.92  reached goal: True
arbitrary success(last 50): 1.00  avg steps on success: 14.0  greedy task return: 0.86  reached goal: True

The exact numbers will move with the seed. The pattern is what matters: the potential arm reaches the goal more reliably and in fewer steps, while the arbitrary arm also reaches the goal but takes a longer path. That longer path is the distortion showing up in the metric that actually counts.

Reading the Output: Speed, Final Policy, and the Trap

Three signals matter, and they are not the same signal.

Success rate over the last 50 episodes. This is the honest measure of whether the agent learned the task. Average episode length alone is misleading, because an episode that hits the 200-step cap and an episode that reaches the goal in 200 steps look identical in a raw step count. Track reached separately.

Steps conditional on success. Among episodes that actually reached the goal, how many steps did they take? This is where the arbitrary arm's detour becomes visible. It reaches the goal, but it wanders.

Greedy task return. Roll out the learned greedy policy with no exploration and score it on the base reward. This is the diagnostic that separates "learned faster" from "learned something different." If the potential arm's greedy return matches the base arm's, the policy is preserved. If the arbitrary arm's greedy return is lower, the bonus changed the policy.

The expected pattern: base reward is slow but correct. Potential-based shaping is faster and converges to the same greedy policy. Arbitrary shaping reaches the goal but takes a longer path, because the bonus pays for rightward motion regardless of whether it helps.

Common mistake: Reading a faster learning curve as evidence of a better policy. The curve measures the shaped objective. The shaped objective is not the task objective. An agent can look excellent on shaped return while being wrong on the real task.

Why the arbitrary arm fails is worth stating plainly. The bonus is not a potential difference, so it changes the relative value of states. A detour that collects more rightward moves can become optimal under the shaped reward even though it is strictly worse under the true task. That is reward hacking in miniature: the agent optimizes the signal you gave it, not the outcome you meant.

Where the Guarantee Stops

The experiment shows a pattern. The theory explains the boundary. Keep them separate.

The guarantee is asymptotic and assumes sufficient state-action coverage. A finite run with epsilon-greedy exploration may not reach the optimal policy in either arm. So "the potential arm matched the base policy" is evidence, not proof — one seed and 500 episodes is a demonstration, not a theorem.

Potential-based shaping can also slow learning. A bad potential is still policy-invariant, which means it cannot hurt the optimum, but it can absolutely hurt your wall-clock time. Invariance is a statement about where you end up, not how fast you get there.

The guarantee is about the optimal policy set, not about the learned Q-values. If you extend the potential to depend on actions — look-ahead advice, in the Wiewiora sense — the potential must be removed from the final Q-values to recover the original ordering. State-only potentials do not carry that extra step.

Episodic edge cases bite quietly. Un-telescoped end-of-episode terms and absorbing-state potentials can break invariance if you handle them carelessly. Set Phi consistently at the goal and check that the terminal transition does not leave a dangling term.

When to use potential-based shaping: sparse rewards, a trustworthy heuristic for state value, and a need for faster learning. When not to: when you cannot express your intuition as a state potential, or when you are not prepared to verify the final policy against the true task.

One Follow-Up Experiment Worth Running

Change one variable. Keep everything else fixed.

Swap the potential for a deliberately bad one — a potential that rewards moving away from the goal. The final policy should still be correct, because the guarantee holds for any bounded potential. Learning should be slower or noisier. That single run teaches the difference between invariance and usefulness better than any paragraph.

Alternatively, take the arbitrary bonus and try to rewrite it as a potential difference. If you can, the policy distortion should disappear. If you cannot, you have found the exact reason the arbitrary arm failed.

Record the seed, the number of runs, and the metric before you change anything. A single run is not evidence. Run five seeds, average the curves, and only then trust the comparison.

The Decision Rule This Experiment Earns

Potential-based shaping is a tool for changing how fast the agent learns, not what it learns. That holds only when the shaping term is a genuine potential difference and you verify the final policy against the true task. The moment your bonus is not a potential difference, you are back to ordinary reward engineering, with all the reward-hacking risk that implies.

Your next step: re-run the three arms across multiple seeds and evaluate the greedy policy on a held-out set of start states. If a shaped agent looks fast but behaves wrong, the diagnostic is the final policy, not the curve.