Building a Grid World Environment for Reinforcement Learning | CodeSignal Learn

Building a Grid World Environment for Reinforcement Learning | CodeSignal Learn

Introduction

Welcome to the second lesson of " Environment Engineering: The Foundation of RL Systems "! In our previous lesson, we explored the fundamental concepts of Reinforcement Learning, including states, actions, rewards, and transitions. Today, we'll take an exciting step forward by starting the actual implementation of our very own Grid World environment from scratch. As you may recall, the environment is where our RL agent lives and interacts. It defines the rules of the world, manages the agent's state, and provides feedback through rewards. Building a well-structured environment is crucial for developing effective reinforcement learning systems, and that's exactly what we'll focus on today. By the end of this lesson, you'll have laid the foundations of our Grid World environment that we will employ throughout this path. Think of this as building your agent's first training ground — an essential skill for any RL practitioner!

Understanding RL Environments Structure

Before diving into code, let's understand how reinforcement learning environments are typically structured. RL environments generally follow a common interface that allows agents to interact with them in a standardized way, which includes these core functionalities: Initialization : Setting up the environment parameters and initial conditions. Reset : Returning the environment to an initial state to start a new episode. Step : Executing an action and returning the new state, reward, and episode status (and additional info). Render : Providing a human-readable visualization of the current state. This structure is inspired by the OpenAI Gym framework (now Gymnasium), which has become the standard interface for RL environments. By following this pattern, we create environments that can work with a wide variety of RL algorithms. In real-world applications, this standardized structure is tremendously valuable: for example, robotics researchers can swap different control algorithms on the same robot interface, and game AI developers can test various learning strategies without rewriting environment code. For our Grid World, we'll implement these components incrementally. Today, we'll focus on the first two: initialization and reset. In the next lesson, we'll tackle the step and render methods.

Creating the GridWorldEnv Class

Sign up

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 start by creating our GridWorldEnv class and implementing the initialization method: class GridWorldEnv: def __init__(self, size=5): """ Basic environment setup for a grid of shape (size x size). We'll store the current agent position (state) and a goal state. """ self.size = size self.state = None self.goal_state = (size - 1, size - 1) self.action_space = {0, 1, 2, 3} # possible actions self.max_steps = size * 3 # reasonable limit for episode length In this initialization method, we're setting up several important attributes: size: Determines the dimensions of our grid (size × size). state: Represents the current position of our agent (initially set to None). goal_state: Defines which cell the agent is trying to reach; it is set to (size - 1, size - 1), which corresponds to the bottom-right corner of our grid. action_space: Specifies the set of actions the agent can take (in our case, actions 0, 1, 2, and 3,representing up, down, left and right respectively). max_steps: Sets a limit on how many steps can be taken in an episode before it terminates. This is important to prevent what's known in RL as the "wandering agent problem" — without this limit, an untrained agent might wander indefinitely without finding the goal. Think of these parameters as the rules of the game you're creating: just as chess has a board size and movement rules, your environment needs clear boundaries and action definitions.

Implementing the Reset Method

Let's see how we might implement the reset method for our GridWorldEnv : def reset(self): """ Reset the environment to an initial state. By default, agent starts at (0, 0). Returns the initial state. """ self.state = (0, 0) self.steps_count = 0 # We'll just return (row, col) for now. return self.state Our reset method performs three important functions: It places the agent at the starting position, which we've defined as the top-left corner (0, 0) It resets the step counter to zero, tracking a fresh episode It returns the initial state so the agent knows where it's starting from In more complex environments, the reset method might also randomize certain elements to provide variety in training. For example, in a robotic simulation, you might randomly vary the starting position to ensure your agent learns a robust policy.

Testing Our Environment

Now that we've implemented the basic structure of our environment, let's write a simple test function to verify that everything is working correctly. This will help us understand how to use our environment in practice. def main(): env = GridWorldEnv(size=5) initial_state = env.reset() print("Environment initialized. Initial state:", initial_state) print("Action space:", env.action_space) print("Maximum episode steps:", env.max_steps) if __name__ == "__main__": main() When you run this code, here's what happens: We create a new instance of our GridWorldEnv with a grid size of 5\times5 We call the reset method to initialize the environment and get the starting state We print out some information about our environment to verify it's set up correctly

Environment Parameters and Flexibility

One important aspect of good environment design is flexibility. Notice how we've parameterized our Grid World with a size parameter in the initialization. This allows us to easily create grids of different dimensions without changing our code. For example, you could create a smaller 3×3 grid: small_env = GridWorldEnv(size=3) initial_state = small_env.reset() print(f"Small environment created. Grid size: 3x3") print(f"Initial state: {initial_state}") # (0, 0) print(f"Goal state: {small_env.goal_state}") # (2, 2) Or a larger 10×10 grid: large_env = GridWorldEnv(size=10) initial_state = large_env.reset() print(f"Large environment created. Grid size: 10x10") print(f"Initial state: {initial_state}") # (0, 0) print(f"Goal state: {large_env.goal_state}") # (9, 9) This flexibility is valuable when experimenting with reinforcement learning algorithms. Different environment sizes present varying levels of difficulty, allowing you to test your agent's capabilities under different conditions. In industry applications, this type of parameterization is essential. For example, a drone navigation system might need to adapt to different physical spaces, or a game AI might need to handle varying map sizes.

Python

class GridWorldEnv:
    def __init__(self, size=5):
        """
        Basic environment setup for a grid of shape (size x size).
        We'll store the current agent position (state) and a goal state.
        """
        self.size = size
        self.state = None
        self.goal_state = (size - 1, size - 1)
        self.action_space = {0, 1, 2, 3}  # possible actions
        self.max_steps = size * 3  # reasonable limit for episode length

Python

def reset(self):
    """
    Reset the environment to an initial state.
    By default, agent starts at (0, 0).
    Returns the initial state.
    """
    self.state = (0, 0)
    self.steps_count = 0
    # We'll just return (row, col) for now.
    return self.state

Python

def main():
    env = GridWorldEnv(size=5)
    initial_state = env.reset()
    print("Environment initialized. Initial state:", initial_state)
    print("Action space:", env.action_space)
    print("Maximum episode steps:", env.max_steps)

if __name__ == "__main__":
    main()

Python

small_env = GridWorldEnv(size=3)
initial_state = small_env.reset()
print(f"Small environment created. Grid size: 3x3")
print(f"Initial state: {initial_state}")      # (0, 0)
print(f"Goal state: {small_env.goal_state}")  # (2, 2)

Python

large_env = GridWorldEnv(size=10)
initial_state = large_env.reset()
print(f"Large environment created. Grid size: 10x10")
print(f"Initial state: {initial_state}")      # (0, 0)
print(f"Goal state: {large_env.goal_state}")  # (9, 9)