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

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)

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])

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)

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

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)

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

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)

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])

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)

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

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)

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

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