Designing Effective State Representations in Reinforcement Learning

Introduction

Welcome to Lesson 4 of our course " Navigating RL Challenges: Strategies and Future Directions "! So far, we've enhanced our grid world environment with random goals, implemented reward shaping to accelerate learning, and added mines as obstacles that our agent must learn to avoid. In this lesson, we'll tackle one of the most fundamental challenges in Reinforcement Learning: designing effective state representations . The way we represent the environment state to our agent can dramatically impact learning efficiency, generalization ability, and overall performance. As we've seen in previous lessons, our agent has been receiving raw state information, such as the agent position and the goal position. While this approach works, it's not the most efficient representation for learning: Real-world RL applications often require careful feature engineering to help agents learn more effectively from the same experiences. By the end of this lesson, we'll have implemented an improved state representation that extracts meaningful features from our Grid World environment, allowing our agent to learn more efficiently and generalize better across different scenarios.

The Importance of State Representation in RL

As we mentioned before, the observation (or state representation) is the agent's "actual view" of the environment — it's all the information the agent can use to make decisions. Poor state representations can lead to several problems: Slow learning : When states contain irrelevant information or aren't structured optimally, agents require more samples to learn effective policies. Limited generalization : Raw state coordinates don't generalize or scale well. For example, being at position (3,4) with a goal at (5,5) is conceptually the same as being at (1,1) with a goal at (3,2) — both require moving southeast. Large state spaces : Using raw coordinates creates a state space that scales with the environment size, making learning more difficult. Difficulty capturing relevant relationships : Important information like "am I moving toward the goal?" isn't explicit in raw coordinates. A well-designed state representation should include only relevant information for decision-making, encode meaningful relationships and patterns, support generalization across similar situations, and be compact enough for efficient learning. In our grid world navigation task, instead of using raw coordinates, we can represent the state in terms of relationships between the agent, goal, and obstacles — information that directly informs the optimal policy.

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 _get_improved_observation(self):
    """
    Create an improved observation representation:
    - Direction to goal (as a compass direction: N, NE, E, SE, S, SW, W, NW)
    - Manhattan distance to goal (normalized by grid size)
    - Mine danger in each direction (binary)
    """
    # Calculate relative direction to goal
    agent_r, agent_c = self.state
    goal_r, goal_c = self.goal_state

    # Direction vector to goal
    dr = goal_r - agent_r
    dc = goal_c - agent_c

    # Convert to compass direction (0-7)
    # 0:N, 1:NE, 2:E, 3:SE, 4:S, 5:SW, 6:W, 7:NW
    compass_direction = -1
    if dr < 0 and dc == 0:
        compass_direction = 0  # North
    elif dr < 0 and dc > 0:
        compass_direction = 1  # Northeast
    elif dr == 0 and dc > 0:
        compass_direction = 2  # East
    # ... additional compass directions

def _get_improved_observation(self):
    """
    Create an improved observation representation:
    - Direction to goal (as a compass direction: N, NE, E, SE, S, SW, W, NW)
    - Manhattan distance to goal (normalized by grid size)
    - Mine danger in each direction (binary)
    """
    # Calculate relative direction to goal
    agent_r, agent_c = self.state
    goal_r, goal_c = self.goal_state

    # Direction vector to goal
    dr = goal_r - agent_r
    dc = goal_c - agent_c

    # Convert to compass direction (0-7)
    # 0:N, 1:NE, 2:E, 3:SE, 4:S, 5:SW, 6:W, 7:NW
    compass_direction = -1
    if dr < 0 and dc == 0:
        compass_direction = 0  # North
    elif dr < 0 and dc > 0:
        compass_direction = 1  # Northeast
    elif dr == 0 and dc > 0:
        compass_direction = 2  # East
    # ... additional compass directions

Python

    # Calculate normalized Manhattan distance
    distance = self._manhattan_distance(self.state, self.goal_state) / (2 * self.size)

    # Check for mines in adjacent cells
    mines_near = self._check_nearby_mines()

    # Determine if we're at an edge of the grid
    at_edge = (
        agent_r == 0 or agent_r == self.size - 1 or
        agent_c == 0 or agent_c == self.size - 1
    )

    # Combine into a compact state representation
    return (compass_direction, distance, mines_near, at_edge)

    # Calculate normalized Manhattan distance
    distance = self._manhattan_distance(self.state, self.goal_state) / (2 * self.size)

    # Check for mines in adjacent cells
    mines_near = self._check_nearby_mines()

    # Determine if we're at an edge of the grid
    at_edge = (
        agent_r == 0 or agent_r == self.size - 1 or
        agent_c == 0 or agent_c == self.size - 1
    )

    # Combine into a compact state representation
    return (compass_direction, distance, mines_near, at_edge)

Python

def reset(self):
    """
    Reset environment:
    - Agent at (0, 0)
    - Random goal
    - Random mines (if any)
    - Initialize distance for shaping
    """
    self.state = (0, 0)
    self.goal_state = self._sample_random_goal()
    self.mines = self._sample_mines() if self.n_mines > 0 else set()
    self.prev_distance = self._manhattan_distance(self.state, self.goal_state)
    self.steps_count = 0

    # Return the improved observation
    return self._get_improved_observation()

def reset(self):
    """
    Reset environment:
    - Agent at (0, 0)
    - Random goal
    - Random mines (if any)
    - Initialize distance for shaping
    """
    self.state = (0, 0)
    self.goal_state = self._sample_random_goal()
    self.mines = self._sample_mines() if self.n_mines > 0 else set()
    self.prev_distance = self._manhattan_distance(self.state, self.goal_state)
    self.steps_count = 0

    # Return the improved observation
    return self._get_improved_observation()

Python

def step(self, action):
    """
    Apply the action and return improved observation, reward, done, info.
    Mines cause an immediate -1 reward and end the episode.
    """
    self.steps_count += 1
    r, c = self.state

    # Apply action to update state
    if action == 0:  # up
        r = max(r - 1, 0)
    elif action == 1:  # down
        r = min(r + 1, self.size - 1)
    # ... handling other actions
    
    self.state = (r, c)

    # Check for mine
    if self.state in self.mines:
        return self._get_improved_observation(), -1.0, True, {"hit_mine": True}

    # Calculate shaped reward
    # ... reward shaping logic
    
    return self._get_improved_observation(), reward, done, info

def step(self, action):
    """
    Apply the action and return improved observation, reward, done, info.
    Mines cause an immediate -1 reward and end the episode.
    """
    self.steps_count += 1
    r, c = self.state

    # Apply action to update state
    if action == 0:  # up
        r = max(r - 1, 0)
    elif action == 1:  # down
        r = min(r + 1, self.size - 1)
    # ... handling other actions
    
    self.state = (r, c)

    # Check for mine
    if self.state in self.mines:
        return self._get_improved_observation(), -1.0, True, {"hit_mine": True}

    # Calculate shaped reward
    # ... reward shaping logic
    
    return self._get_improved_observation(), reward, done, info

Python

def main():
    env_size = 5
    actions = [0, 1, 2, 3]
    n_episodes = 500
    n_mines = 3

    # Train the agent
    env = GridWorldEnv(size=env_size, use_reward_shaping=True, n_mines=n_mines)
    agent = QLearningAgent(actions=actions)
    train_agent(env, agent, n_episodes=n_episodes)

    # Test the agent
    test_env = GridWorldEnv(size=env_size, use_reward_shaping=True, n_mines=n_mines)
    test_agent(test_env, agent, n_episodes=10, render=False)

def main():
    env_size = 5
    actions = [0, 1, 2, 3]
    n_episodes = 500
    n_mines = 3

    # Train the agent
    env = GridWorldEnv(size=env_size, use_reward_shaping=True, n_mines=n_mines)
    agent = QLearningAgent(actions=actions)
    train_agent(env, agent, n_episodes=n_episodes)

    # Test the agent
    test_env = GridWorldEnv(size=env_size, use_reward_shaping=True, n_mines=n_mines)
    test_agent(test_env, agent, n_episodes=10, render=False)

text

Average Reward: 1.14
Success Rate: 1.00
Mine Hits Rate: 0.00

Average Reward: 1.14
Success Rate: 1.00
Mine Hits Rate: 0.00

Python

def _get_improved_observation(self):
    """
    Create an improved observation representation:
    - Direction to goal (as a compass direction: N, NE, E, SE, S, SW, W, NW)
    - Manhattan distance to goal (normalized by grid size)
    - Mine danger in each direction (binary)
    """
    # Calculate relative direction to goal
    agent_r, agent_c = self.state
    goal_r, goal_c = self.goal_state

    # Direction vector to goal
    dr = goal_r - agent_r
    dc = goal_c - agent_c

    # Convert to compass direction (0-7)
    # 0:N, 1:NE, 2:E, 3:SE, 4:S, 5:SW, 6:W, 7:NW
    compass_direction = -1
    if dr < 0 and dc == 0:
        compass_direction = 0  # North
    elif dr < 0 and dc > 0:
        compass_direction = 1  # Northeast
    elif dr == 0 and dc > 0:
        compass_direction = 2  # East
    # ... additional compass directions

Python

    # Calculate normalized Manhattan distance
    distance = self._manhattan_distance(self.state, self.goal_state) / (2 * self.size)

    # Check for mines in adjacent cells
    mines_near = self._check_nearby_mines()

    # Determine if we're at an edge of the grid
    at_edge = (
        agent_r == 0 or agent_r == self.size - 1 or
        agent_c == 0 or agent_c == self.size - 1
    )

    # Combine into a compact state representation
    return (compass_direction, distance, mines_near, at_edge)

Python

def reset(self):
    """
    Reset environment:
    - Agent at (0, 0)
    - Random goal
    - Random mines (if any)
    - Initialize distance for shaping
    """
    self.state = (0, 0)
    self.goal_state = self._sample_random_goal()
    self.mines = self._sample_mines() if self.n_mines > 0 else set()
    self.prev_distance = self._manhattan_distance(self.state, self.goal_state)
    self.steps_count = 0

    # Return the improved observation
    return self._get_improved_observation()

Python

def step(self, action):
    """
    Apply the action and return improved observation, reward, done, info.
    Mines cause an immediate -1 reward and end the episode.
    """
    self.steps_count += 1
    r, c = self.state

    # Apply action to update state
    if action == 0:  # up
        r = max(r - 1, 0)
    elif action == 1:  # down
        r = min(r + 1, self.size - 1)
    # ... handling other actions
    
    self.state = (r, c)

    # Check for mine
    if self.state in self.mines:
        return self._get_improved_observation(), -1.0, True, {"hit_mine": True}

    # Calculate shaped reward
    # ... reward shaping logic
    
    return self._get_improved_observation(), reward, done, info

Python

def main():
    env_size = 5
    actions = [0, 1, 2, 3]
    n_episodes = 500
    n_mines = 3

    # Train the agent
    env = GridWorldEnv(size=env_size, use_reward_shaping=True, n_mines=n_mines)
    agent = QLearningAgent(actions=actions)
    train_agent(env, agent, n_episodes=n_episodes)

    # Test the agent
    test_env = GridWorldEnv(size=env_size, use_reward_shaping=True, n_mines=n_mines)
    test_agent(test_env, agent, n_episodes=10, render=False)

text

Average Reward: 1.14
Success Rate: 1.00
Mine Hits Rate: 0.00

Introduction

The Importance of State Representation in RL

Join the 1M+ learners on CodeSignal

Designing an Improved State Representation

Implementing the Improved Observation Function, Step 1

Implementing the Improved Observation Function, Step 2

Integrating the New Representation: Reset Method

Integrating the New Representation: Step Method

Learning with the Improved Representation

Conclusion and Next Steps