Reward Shaping for Faster Learning in Reinforcement Learning

Introduction

Welcome back to our second lesson in " Navigating RL Challenges: Strategies and Future Directions "! In the previous lesson, we enhanced our grid world environment with random goals, making our Reinforcement Learning problem more dynamic and realistic. You learned how to provide the agent with both its position and the goal position in every observation, enabling it to develop flexible, goal-conditioned policies. Today, we're taking the next crucial step by exploring reward shaping — a powerful technique that can dramatically improve learning speed and efficiency. While our previous environment provided only binary rewards (1 for reaching the goal, 0 otherwise), you'll see how providing more informative feedback helps your agents learn faster and more effectively, especially in sparse reward environments.

The Sparse Reward Problem

Imagine teaching a child to play basketball, but only telling them whether they scored a basket or not — no tips about proper form, distance, or aiming. This is essentially what happens with sparse rewards in Reinforcement Learning. In our current grid world implementation, the agent receives a reward only upon reaching the goal. This creates several critical challenges: Exploration inefficiency : The agent must stumble upon the goal by random exploration before it can start learning. Delayed learning signals : Feedback comes only at the end of successful episodes, making credit assignment (understanding the contribution of each action towards reaching the goal) difficult. Slow convergence : With minimal guidance, the agent requires many episodes to develop effective policies. Consider a robot learning to navigate a warehouse. If it receives feedback only upon reaching its destination, it might wander aimlessly for hours before getting any useful learning signal. This mirrors what happens in our grid world — as the environment gets larger, learning becomes exponentially more difficult with sparse rewards. The problem becomes particularly severe in environments with large state spaces, long episodes, or costly exploration. Reward shaping offers a solution by providing intermediate feedback that guides the agent toward desirable behaviors.

What is Reward Shaping?

Reward shaping augments the original sparse rewards with additional signals that guide the learning process. It's like transforming a binary "success/failure" game into a continuous "getting warmer or colder" feedback system that provides helpful hints throughout the agent's journey. The key principles of effective reward shaping include: Providing intermediate feedback that indicates progress towards the goal; Maintaining the same optimal solution as the original problem; Balancing the strength of shaping rewards against the main goal reward. In our grid world, we'll use a distance-based approach — giving small positive rewards when the agent moves closer to the goal and small penalties when it moves farther away. This creates a more informative learning signal while still keeping the main reward (reaching the goal) as the primary objective.

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Let's modify our grid world environment to support reward shaping. First, we'll update the constructor to include a shaping option:

Python

def __init__(self, size=5, use_reward_shaping=False):
    """
    GridWorldEnv with optional reward shaping.
    
    Args:
        size (int): Size of the grid (size x size).
        use_reward_shaping (bool): If True, use distance-based shaping rewards.
    """
    self.size = size
    self.use_reward_shaping = use_reward_shaping
    self.state = None
    self.goal_state = None
    self.prev_distance = None  # Will track previous distance to goal
    self.action_space = 4
    self.max_steps = size * 3
    self.steps_count = 0
    self.rng = np.random.default_rng(42)

def __init__(self, size=5, use_reward_shaping=False):
    """
    GridWorldEnv with optional reward shaping.
    
    Args:
        size (int): Size of the grid (size x size).
        use_reward_shaping (bool): If True, use distance-based shaping rewards.
    """
    self.size = size
    self.use_reward_shaping = use_reward_shaping
    self.state = None
    self.goal_state = None
    self.prev_distance = None  # Will track previous distance to goal
    self.action_space = 4
    self.max_steps = size * 3
    self.steps_count = 0
    self.rng = np.random.default_rng(42)

This constructor introduces two important additions:

The use_reward_shaping flag that lets you enable or disable shaping
The prev_distance attribute that will help you track changes in distance to the goal

Now, to calculate the distance between positions, we'll implement a Manhattan distance method:

d_{\text{Manhattan}}(\text{pos}_a, \text{pos}_b) = |x_a - x_b| + |y_a - y_b|

Which in code translates to:

Python

def _manhattan_distance(self, pos_a, pos_b):
    """
    Calculate Manhattan distance between pos_a and pos_b.
    """
    return abs(pos_a[0] - pos_b[0]) + abs(pos_a[1] - pos_b[1])

The Manhattan distance (also called "taxicab" distance) is perfect for our grid world because it represents the minimum number of steps needed to move from one position to another when diagonal moves aren't allowed.

Next, we need to modify our reset method to initialize the distance tracking:

Python

def reset(self):
    """
    Reset the environment to state (0, 0), sample a random goal,
    and initialize distance if using shaping.

    Returns:
        tuple: ((row, col), (goal_row, goal_col))
    """
    self.state = (0, 0)
    self.goal_state = self._sample_random_goal()
    self.prev_distance = self._manhattan_distance(self.state, self.goal_state)
    self.steps_count = 0
    return (self.state, self.goal_state)

def reset(self):
    """
    Reset the environment to state (0, 0), sample a random goal,
    and initialize distance if using shaping.

    Returns:
        tuple: ((row, col), (goal_row, goal_col))
    """
    self.state = (0, 0)
    self.goal_state = self._sample_random_goal()
    self.prev_distance = self._manhattan_distance(self.state, self.goal_state)
    self.steps_count = 0
    return (self.state, self.goal_state)

The key addition here is calculating and storing the initial distance to the goal. This serves as your baseline when computing shaped rewards — you'll reward the agent for decreasing this distance and penalize it for increasing the distance.

Implementing the Shaped Reward Function

Now for the heart of reward shaping — updating the step method to provide intermediate rewards based on progress:

Python

def step(self, action):
    """
    Take an action in the environment.

    Args:
        action (int): 0=up, 1=down, 2=left, 3=right

    Returns:
        tuple: (next_state, reward, done, info)
    """
    
    # Execute action and update self.state
    # ...

    # Calculate reward shaping component
    shaped_reward = 0.0
    if self.use_reward_shaping:
        current_distance = self._manhattan_distance(self.state, self.goal_state)
        distance_diff = self.prev_distance - current_distance
        step_penalty = -0.01  # Small penalty for each step
        distance_reward = 0.05 * distance_diff  # Reward for getting closer (or penalty for moving away)
        shaped_reward = step_penalty + distance_reward
        self.prev_distance = current_distance  # Update for next step

    # Determine final reward and done status
    if self.state == self.goal_state:
        reward = 1.0  # Main reward for reaching goal
        done = True
        info = {"success": True}
    else:
        reward = shaped_reward  # Only shaped rewards when not at goal
        done = (self.steps_count >= self.max_steps)
        info = {"timeout": done}

    return (self.state, self.goal_state), reward, done, info

def step(self, action):
    """
    Take an action in the environment.

    Args:
        action (int): 0=up, 1=down, 2=left, 3=right

    Returns:
        tuple: (next_state, reward, done, info)
    """
    
    # Execute action and update self.state
    # ...

    # Calculate reward shaping component
    shaped_reward = 0.0
    if self.use_reward_shaping:
        current_distance = self._manhattan_distance(self.state, self.goal_state)
        distance_diff = self.prev_distance - current_distance
        step_penalty = -0.01  # Small penalty for each step
        distance_reward = 0.05 * distance_diff  # Reward for getting closer (or penalty for moving away)
        shaped_reward = step_penalty + distance_reward
        self.prev_distance = current_distance  # Update for next step

    # Determine final reward and done status
    if self.state == self.goal_state:
        reward = 1.0  # Main reward for reaching goal
        done = True
        info = {"success": True}
    else:
        reward = shaped_reward  # Only shaped rewards when not at goal
        done = (self.steps_count >= self.max_steps)
        info = {"timeout": done}

    return (self.state, self.goal_state), reward, done, info

This implementation includes several thoughtfully designed elements:

The distance_diff calculation (previous distance minus current distance) is positive when you get closer to the goal and negative when you move away — creating an intuitive progress signal.
The small -0.01 step penalty encourages efficiency by discouraging unnecessarily long paths.
The 0.05 scaling factor balances the shaping signal strength — strong enough to guide learning but not so strong that it overwhelms the main goal reward.
We preserve the original reward structure by still providing the full +1.0 when reaching the goal, maintaining focus on the primary objective.

Comparing Learning Performance

To see the impact of reward shaping, you can run an experiment comparing agents with and without shaped rewards:

Python

def main():
    """
    Compare learning with standard vs shaped rewards.
    """
    env_size = 5
    actions = [0, 1, 2, 3]
    n_episodes = 1000

    # Train agent with standard rewards
    print("Training with standard rewards...")
    env_standard = GridWorldEnv(size=env_size, use_reward_shaping=False)
    agent_standard = QLearningAgent(actions=actions)
    standard_results = train_agent(env_standard, agent_standard, n_episodes)

    # Train agent with shaped rewards
    print("\nTraining with shaped rewards...")
    env_shaped = GridWorldEnv(size=env_size, use_reward_shaping=True)
    agent_shaped = QLearningAgent(actions=actions)
    shaped_results = train_agent(env_shaped, agent_shaped, n_episodes)

    # Compare final performance
    print("\n=== Performance Comparison ===")
    print("Standard Agent:")
    test_agent(env_standard, agent_standard, n_episodes=10)
    print("Shaped Agent:")
    test_agent(env_shaped, agent_shaped, n_episodes=10)

def main():
    """
    Compare learning with standard vs shaped rewards.
    """
    env_size = 5
    actions = [0, 1, 2, 3]
    n_episodes = 1000

    # Train agent with standard rewards
    print("Training with standard rewards...")
    env_standard = GridWorldEnv(size=env_size, use_reward_shaping=False)
    agent_standard = QLearningAgent(actions=actions)
    standard_results = train_agent(env_standard, agent_standard, n_episodes)

    # Train agent with shaped rewards
    print("\nTraining with shaped rewards...")
    env_shaped = GridWorldEnv(size=env_size, use_reward_shaping=True)
    agent_shaped = QLearningAgent(actions=actions)
    shaped_results = train_agent(env_shaped, agent_shaped, n_episodes)

    # Compare final performance
    print("\n=== Performance Comparison ===")
    print("Standard Agent:")
    test_agent(env_standard, agent_standard, n_episodes=10)
    print("Shaped Agent:")
    test_agent(env_shaped, agent_shaped, n_episodes=10)

Running this experiment produces the following results:

text

Standard Agent:
Testing for 10 episodes...
Testing complete.
Average Reward: 0.20
Success Rate: 0.20
Shaped Agent:
Testing for 10 episodes...
Testing complete.
Average Reward: 1.06
Success Rate: 0.90

Standard Agent:
Testing for 10 episodes...
Testing complete.
Average Reward: 0.20
Success Rate: 0.20
Shaped Agent:
Testing for 10 episodes...
Testing complete.
Average Reward: 1.06
Success Rate: 0.90

These results dramatically illustrate the power of reward shaping. The agent trained with standard rewards managed to reach the goal in only 20% of test episodes, while the agent trained with shaped rewards achieved an impressive 90% success rate. This stark difference in performance comes from the additional learning signals provided by shaped rewards.

This performance gap would become even more pronounced in larger environments, where random exploration with sparse rewards becomes increasingly inefficient. Reward shaping essentially provides a "compass" that guides exploration in promising directions, significantly reducing the number of episodes needed to develop effective policies.

While reward shaping is powerful, it requires careful design to avoid common pitfalls:

1. Beware of reward hacking: If your shaped rewards don't align perfectly with your true objective, agents may find ways to maximize the shaped rewards without achieving the intended goal. For example, if you reward getting closer to the goal without a step penalty, an agent might learn to oscillate back and forth near the goal rather than reaching it.

2. Keep it simple: Start with the simplest shaping function that captures the task objective. Complex shaping functions introduce more hyperparameters and can be brittle.

3. Test empirically: Always compare performance with and without shaping to ensure your shaped rewards are actually helping.

4. Consider potential-based shaping: This mathematically guarantees that the optimal policy remains unchanged, preventing distortion of the original task.

5. Use domain knowledge wisely: Effective shaping requires understanding what constitutes progress in your specific environment. In navigation, distance is natural, but other domains may need different progress metrics.

In this lesson, you've explored reward shaping, a powerful technique for accelerating learning by providing more informative feedback to your agents. You've seen how shaped rewards help overcome the sparse reward problem by giving intermediate signals based on progress toward goals. By implementing distance-based reward shaping in our grid world, you've gained practical experience with one of the most important tools in a Reinforcement Learning practitioner's toolkit.

As you move forward, you'll discover that reward shaping is just one of many techniques for improving learning efficiency in Reinforcement Learning. The principles you've learned here — providing meaningful intermediate feedback, balancing exploration and exploitation, and carefully designing reward functions — will serve as a foundation for more advanced methods you'll encounter in upcoming lessons.

Introduction

The Sparse Reward Problem

What is Reward Shaping?

Join the 1M+ learners on CodeSignal

Setting Up the Environment for Reward Shaping

The Manhattan Distance Function

Tracking Progress Toward Goals

Implementing the Shaped Reward Function

Comparing Learning Performance

Reward Shaping Best Practices

Conclusion and Next Steps