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

Implementing the Improved Observation Function, Step 1

Now let's implement our improved state representation by creating a method that transforms the raw environment state into our more informative representation:

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

The first part of our implementation determines the relative direction to the goal. We calculate the direction vector (dr, dc) and then convert it to one of eight compass directions (or a ninth option when the agent is already at the goal). This makes the state representation invariant to the specific coordinates — the agent only needs to know which direction to move relative to its current position, a far more intuitive way to understand spatial relationships than raw coordinates.

Learning with the Improved Representation

Now let's see how our agent learns using this enhanced state representation. The training process itself remains unchanged — we run episodes and collect statistics on rewards, success rate, and mine hits. The magic happens behind the scenes where our improved state representation makes learning more efficient. With the enhanced representation, our agent doesn't have to learn the complex relationship between raw coordinates and optimal actions from scratch — much of this information is already encoded in the representation itself.

Let's train and test the agent:

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)

Running the above code results in the following output:

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

These impressive results — a perfect success rate of 100% and complete avoidance of mines — highlight the power of good state representation. By encoding relevant information in an efficient way, we've enabled our agent to learn a robust navigation policy in a complex environment with significantly fewer training episodes than would be required with raw coordinates.

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