├── .gitignore
├── LICENSE
├── README.md
├── gym_box.ipynb
├── sprites
    ├── bot_blue.png
    ├── floor.png
    └── package.png
├── v0_warehouse_robot.py
├── v0_warehouse_robot_env.py
└── v0_warehouse_robot_train.py


/.gitignore:
--------------------------------------------------------------------------------
1 | __pycache__
2 | models
3 | logs
4 | *solution*.pkl
5 | *solution*.png


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2025 johnnycode8
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | <a name="readme-top"></a>
 2 | 
 3 | <h2 align="center">Gymnasium Custom Reinforcement Learning Environments</h2>
 4 | 
 5 | Tutorials on how to create custom Gymnasium-compatible Reinforcement Learning environments using the [Gymnasium Library](https://gymnasium.farama.org/), formerly OpenAI’s Gym library. Each tutorial has a companion video explanation and code walkthrough from my YouTube channel [@johnnycode](https://www.youtube.com/@johnnycode). If the code and video helped you, please consider:  
 6 | <a href='https://www.buymeacoffee.com/johnnycode'><img src="https://cdn.buymeacoffee.com/buttons/default-blue.png" alt="Buy Me A Coffee" height="41" width="174"></a>
 7 | 
 8 | ## Custom Gym Environment part 1 - Warehouse Robot v0
 9 | This is a very basic tutorial showing end-to-end how to create a custom Gymnasium-compatible Reinforcement Learning environment. The tutorial is divided into three parts:
10 | 1. Model your problem.
11 | 2. Convert your problem into a Gymnasium-compatible environment.
12 | 3. Train your custom environment in two ways; using Q-Learning and using the Stable Baselines3 library. 
13 | 
14 | ##### Code Reference:
15 | * v0_warehouse_robot*.py
16 | 
17 | ##### YouTube Tutorial:
18 | <a href='https://youtu.be/AoGRjPt-vms&list=PL58zEckBH8fDt-F9LbpVASTor_jZzsxRg'><img src='https://img.youtube.com/vi/AoGRjPt-vms/0.jpg' width='400' alt='Build Custom Gymnasium Env'/></a>
19 | 
20 | 
21 | ## Custom Gym Environment part 2 - Visualization with Pygame
22 | In part 1, we created a very simple custom Reinforcement Learning environment that is compatible with Farama Gymnasium (formerly OpenAI Gym). In this tutorial, we'll do a minor upgrade and visualize our environment using Pygame.
23 | 
24 | ##### Code Reference:
25 | * v0_warehouse_robot*.py
26 | 
27 | ##### YouTube Tutorial:
28 | <a href='https://youtu.be/9t64PFO7hr0&list=PL58zEckBH8fDt-F9LbpVASTor_jZzsxRg'><img src='https://img.youtube.com/vi/9t64PFO7hr0/0.jpg' width='400' alt='Build Custom Gymnasium Env with Pygame'/></a>
29 | 
30 | 
31 | 
32 | # Additional Resources
33 | 
34 | ## How gymnasium.spaces.Box Works
35 | The Box space type is used in many Gymnasium environments and you'll likely need it for your custom environment. The Box action space can be used to validate agent actions or generate random actions. The Box observation space can be used to validate the environment's state. This video explains and demos how to create boxes of different sizes/shapes, lower (low) and upper (high) boundaries, and as int/float data types.
36 | 
37 | ##### YouTube Tutorial:
38 | <a href='https://youtu.be/hI7UDemXVsk&list=PL58zEckBH8fDt-F9LbpVASTor_jZzsxRg'><img src='https://img.youtube.com/vi/hI7UDemXVsk/0.jpg' width='400' alt='How gymnasium.spaces.Box Works'/></a>
39 | 
40 | ## Reinforcement Learning Tutorials
41 | For more Reinforcement Learning and Deep Reinforcement Learning tutorials, check out my
42 | [Gym Solutions](https://github.com/johnnycode8/gym_solutions) repository.
43 | 
44 | <p align="right">(<a href="#readme-top">back to top</a>)</p>
45 | 


--------------------------------------------------------------------------------
/gym_box.ipynb:
--------------------------------------------------------------------------------
 1 | {
 2 |  "cells": [
 3 |   {
 4 |    "cell_type": "code",
 5 |    "execution_count": null,
 6 |    "metadata": {},
 7 |    "outputs": [],
 8 |    "source": [
 9 |     "import gymnasium as gym\n",
10 |     "from gymnasium import spaces\n",
11 |     "import numpy as np"
12 |    ]
13 |   },
14 |   {
15 |    "cell_type": "code",
16 |    "execution_count": 318,
17 |    "metadata": {},
18 |    "outputs": [
19 |     {
20 |      "data": {
21 |       "text/plain": [
22 |        "array([[ 1.54535348, 12.69607217, 24.35233266],\n",
23 |        "       [ 2.30901039, 18.2606693 , -1.68016234]])"
24 |       ]
25 |      },
26 |      "execution_count": 318,
27 |      "metadata": {},
28 |      "output_type": "execute_result"
29 |     }
30 |    ],
31 |    "source": [
32 |     "myspace = spaces.Box( \n",
33 |     "    low  = -np.inf,\n",
34 |     "    high = np.array([[2, 15, 25], \n",
35 |     "                     [3, 19, np.inf]]),\n",
36 |     "    shape=(2,3),\n",
37 |     "    dtype=np.float64\n",
38 |     ")\n",
39 |     "\n",
40 |     "myspace.sample()"
41 |    ]
42 |   }
43 |  ],
44 |  "metadata": {
45 |   "kernelspec": {
46 |    "display_name": "gymenv",
47 |    "language": "python",
48 |    "name": "python3"
49 |   },
50 |   "language_info": {
51 |    "codemirror_mode": {
52 |     "name": "ipython",
53 |     "version": 3
54 |    },
55 |    "file_extension": ".py",
56 |    "mimetype": "text/x-python",
57 |    "name": "python",
58 |    "nbconvert_exporter": "python",
59 |    "pygments_lexer": "ipython3",
60 |    "version": "3.11.3"
61 |   }
62 |  },
63 |  "nbformat": 4,
64 |  "nbformat_minor": 2
65 | }
66 | 


--------------------------------------------------------------------------------
/sprites/bot_blue.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/johnnycode8/gym_custom_env/84262cede528d5dd36847c465a40ba58b0668861/sprites/bot_blue.png


--------------------------------------------------------------------------------
/sprites/floor.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/johnnycode8/gym_custom_env/84262cede528d5dd36847c465a40ba58b0668861/sprites/floor.png


--------------------------------------------------------------------------------
/sprites/package.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/johnnycode8/gym_custom_env/84262cede528d5dd36847c465a40ba58b0668861/sprites/package.png


--------------------------------------------------------------------------------
/v0_warehouse_robot.py:
--------------------------------------------------------------------------------
  1 | '''
  2 | This module models the problem to be solved. In this very simple example, the problem is to optimze a Robot that works in a Warehouse.
  3 | The Warehouse is divided into a rectangular grid. A Target is randomly placed on the grid and the Robot's goal is to reach the Target.
  4 | '''
  5 | import random
  6 | from enum import Enum
  7 | import pygame
  8 | import sys
  9 | from os import path
 10 | 
 11 | # Actions the Robot is capable of performing i.e. go in a certain direction
 12 | class RobotAction(Enum):
 13 |     LEFT=0
 14 |     DOWN=1
 15 |     RIGHT=2
 16 |     UP=3
 17 | 
 18 | # The Warehouse is divided into a grid. Use these 'tiles' to represent the objects on the grid.
 19 | class GridTile(Enum):
 20 |     _FLOOR=0
 21 |     ROBOT=1
 22 |     TARGET=2
 23 | 
 24 |     # Return the first letter of tile name, for printing to the console.
 25 |     def __str__(self):
 26 |         return self.name[:1]
 27 | 
 28 | class WarehouseRobot:
 29 | 
 30 |     # Initialize the grid size. Pass in an integer seed to make randomness (Targets) repeatable.
 31 |     def __init__(self, grid_rows=4, grid_cols=5, fps=1):
 32 |         self.grid_rows = grid_rows
 33 |         self.grid_cols = grid_cols
 34 |         self.reset()
 35 | 
 36 |         self.fps = fps
 37 |         self.last_action=''
 38 |         self._init_pygame()
 39 | 
 40 |     def _init_pygame(self):
 41 |         pygame.init() # initialize pygame
 42 |         pygame.display.init() # Initialize the display module
 43 | 
 44 |         # Game clock
 45 |         self.clock = pygame.time.Clock()
 46 | 
 47 |         # Default font
 48 |         self.action_font = pygame.font.SysFont("Calibre",30)
 49 |         self.action_info_height = self.action_font.get_height()
 50 | 
 51 |         # For rendering
 52 |         self.cell_height = 64
 53 |         self.cell_width = 64
 54 |         self.cell_size = (self.cell_width, self.cell_height)        
 55 | 
 56 |         # Define game window size (width, height)
 57 |         self.window_size = (self.cell_width * self.grid_cols, self.cell_height * self.grid_rows + self.action_info_height)
 58 | 
 59 |         # Initialize game window
 60 |         self.window_surface = pygame.display.set_mode(self.window_size) 
 61 | 
 62 |         # Load & resize sprites
 63 |         file_name = path.join(path.dirname(__file__), "sprites/bot_blue.png")
 64 |         img = pygame.image.load(file_name)
 65 |         self.robot_img = pygame.transform.scale(img, self.cell_size)
 66 | 
 67 |         file_name = path.join(path.dirname(__file__), "sprites/floor.png")
 68 |         img = pygame.image.load(file_name)
 69 |         self.floor_img = pygame.transform.scale(img, self.cell_size)
 70 | 
 71 |         file_name = path.join(path.dirname(__file__), "sprites/package.png")
 72 |         img = pygame.image.load(file_name)
 73 |         self.goal_img = pygame.transform.scale(img, self.cell_size) 
 74 | 
 75 | 
 76 |     def reset(self, seed=None):
 77 |         # Initialize Robot's starting position
 78 |         self.robot_pos = [0,0]
 79 | 
 80 |         # Random Target position
 81 |         random.seed(seed)
 82 |         self.target_pos = [
 83 |             random.randint(1, self.grid_rows-1),
 84 |             random.randint(1, self.grid_cols-1)
 85 |         ]
 86 | 
 87 |     def perform_action(self, robot_action:RobotAction) -> bool:
 88 |         self.last_action = robot_action
 89 | 
 90 |         # Move Robot to the next cell
 91 |         if robot_action == RobotAction.LEFT:
 92 |             if self.robot_pos[1]>0:
 93 |                 self.robot_pos[1]-=1
 94 |         elif robot_action == RobotAction.RIGHT:
 95 |             if self.robot_pos[1]<self.grid_cols-1:
 96 |                 self.robot_pos[1]+=1
 97 |         elif robot_action == RobotAction.UP:
 98 |             if self.robot_pos[0]>0:
 99 |                 self.robot_pos[0]-=1
100 |         elif robot_action == RobotAction.DOWN:
101 |             if self.robot_pos[0]<self.grid_rows-1:
102 |                 self.robot_pos[0]+=1
103 | 
104 |         # Return true if Robot reaches Target
105 |         return self.robot_pos == self.target_pos
106 | 
107 |     def render(self):
108 |         # Print current state on console
109 |         for r in range(self.grid_rows):
110 |             for c in range(self.grid_cols):
111 | 
112 |                 if([r,c] == self.robot_pos):
113 |                     print(GridTile.ROBOT, end=' ')
114 |                 elif([r,c] == self.target_pos):
115 |                     print(GridTile.TARGET, end=' ')
116 |                 else:
117 |                     print(GridTile._FLOOR, end=' ')
118 | 
119 |             print() # new line
120 |         print() # new line
121 | 
122 |         self._process_events()
123 | 
124 |         # clear to white background, otherwise text with varying length will leave behind prior rendered portions
125 |         self.window_surface.fill((255,255,255))
126 | 
127 |         # Print current state on console
128 |         for r in range(self.grid_rows):
129 |             for c in range(self.grid_cols):
130 |                 
131 |                 # Draw floor
132 |                 pos = (c * self.cell_width, r * self.cell_height)
133 |                 self.window_surface.blit(self.floor_img, pos)
134 | 
135 |                 if([r,c] == self.target_pos):
136 |                     # Draw target
137 |                     self.window_surface.blit(self.goal_img, pos)
138 | 
139 |                 if([r,c] == self.robot_pos):
140 |                     # Draw robot
141 |                     self.window_surface.blit(self.robot_img, pos)
142 |                 
143 |         text_img = self.action_font.render(f'Action: {self.last_action}', True, (0,0,0), (255,255,255))
144 |         text_pos = (0, self.window_size[1] - self.action_info_height)
145 |         self.window_surface.blit(text_img, text_pos)       
146 | 
147 |         pygame.display.update()
148 |                 
149 |         # Limit frames per second
150 |         self.clock.tick(self.fps)  
151 | 
152 |     def _process_events(self):
153 |         # Process user events, key presses
154 |         for event in pygame.event.get():
155 |             # User clicked on X at the top right corner of window
156 |             if event.type == pygame.QUIT:
157 |                 pygame.quit()
158 |                 sys.exit()
159 | 
160 |             if(event.type == pygame.KEYDOWN):
161 |                 # User hit escape
162 |                 if(event.key == pygame.K_ESCAPE):
163 |                     pygame.quit()
164 |                     sys.exit()
165 |                 
166 | 
167 | 
168 | # For unit testing
169 | if __name__=="__main__":
170 |     warehouseRobot = WarehouseRobot()
171 |     warehouseRobot.render()
172 | 
173 |     while(True):
174 |         rand_action = random.choice(list(RobotAction))
175 |         print(rand_action)
176 | 
177 |         warehouseRobot.perform_action(rand_action)
178 |         warehouseRobot.render()


--------------------------------------------------------------------------------
/v0_warehouse_robot_env.py:
--------------------------------------------------------------------------------
  1 | '''
  2 | Custom Gym environment
  3 | https://gymnasium.farama.org/tutorials/gymnasium_basics/environment_creation/
  4 | '''
  5 | import gymnasium as gym
  6 | from gymnasium import spaces
  7 | from gymnasium.envs.registration import register
  8 | from gymnasium.utils.env_checker import check_env
  9 | 
 10 | import v0_warehouse_robot as wr
 11 | import numpy as np
 12 | 
 13 | # Register this module as a gym environment. Once registered, the id is usable in gym.make().
 14 | register(
 15 |     id='warehouse-robot-v0',                                # call it whatever you want
 16 |     entry_point='v0_warehouse_robot_env:WarehouseRobotEnv', # module_name:class_name
 17 | )
 18 | 
 19 | # Implement our own gym env, must inherit from gym.Env
 20 | # https://gymnasium.farama.org/api/env/
 21 | class WarehouseRobotEnv(gym.Env):
 22 |     # metadata is a required attribute
 23 |     # render_modes in our environment is either None or 'human'.
 24 |     # render_fps is not used in our env, but we are require to declare a non-zero value.
 25 |     metadata = {"render_modes": ["human"], 'render_fps': 4}
 26 | 
 27 |     def __init__(self, grid_rows=4, grid_cols=5, render_mode=None):
 28 | 
 29 |         self.grid_rows=grid_rows
 30 |         self.grid_cols=grid_cols
 31 |         self.render_mode = render_mode
 32 | 
 33 |         # Initialize the WarehouseRobot problem
 34 |         self.warehouse_robot = wr.WarehouseRobot(grid_rows=grid_rows, grid_cols=grid_cols, fps=self.metadata['render_fps'])
 35 | 
 36 |         # Gym requires defining the action space. The action space is robot's set of possible actions.
 37 |         # Training code can call action_space.sample() to randomly select an action. 
 38 |         self.action_space = spaces.Discrete(len(wr.RobotAction))
 39 | 
 40 |         # Gym requires defining the observation space. The observation space consists of the robot's and target's set of possible positions.
 41 |         # The observation space is used to validate the observation returned by reset() and step().
 42 |         # Use a 1D vector: [robot_row_pos, robot_col_pos, target_row_pos, target_col_pos]
 43 |         self.observation_space = spaces.Box(
 44 |             low=0,
 45 |             high=np.array([self.grid_rows-1, self.grid_cols-1, self.grid_rows-1, self.grid_cols-1]),
 46 |             shape=(4,),
 47 |             dtype=np.int32
 48 |         )
 49 | 
 50 |     # Gym required function (and parameters) to reset the environment
 51 |     def reset(self, seed=None, options=None):
 52 |         super().reset(seed=seed) # gym requires this call to control randomness and reproduce scenarios.
 53 | 
 54 |         # Reset the WarehouseRobot. Optionally, pass in seed control randomness and reproduce scenarios.
 55 |         self.warehouse_robot.reset(seed=seed)
 56 | 
 57 |         # Construct the observation state:
 58 |         # [robot_row_pos, robot_col_pos, target_row_pos, target_col_pos]
 59 |         obs = np.concatenate((self.warehouse_robot.robot_pos, self.warehouse_robot.target_pos))
 60 |         
 61 |         # Additional info to return. For debugging or whatever.
 62 |         info = {}
 63 | 
 64 |         # Render environment
 65 |         if(self.render_mode=='human'):
 66 |             self.render()
 67 | 
 68 |         # Return observation and info
 69 |         return obs, info
 70 | 
 71 |     # Gym required function (and parameters) to perform an action
 72 |     def step(self, action):
 73 |         # Perform action
 74 |         target_reached = self.warehouse_robot.perform_action(wr.RobotAction(action))
 75 | 
 76 |         # Determine reward and termination
 77 |         reward=0
 78 |         terminated=False
 79 |         if target_reached:
 80 |             reward=1
 81 |             terminated=True
 82 | 
 83 |         # Construct the observation state: 
 84 |         # [robot_row_pos, robot_col_pos, target_row_pos, target_col_pos]
 85 |         obs = np.concatenate((self.warehouse_robot.robot_pos, self.warehouse_robot.target_pos))
 86 | 
 87 |         # Additional info to return. For debugging or whatever.
 88 |         info = {}
 89 | 
 90 |         # Render environment
 91 |         if(self.render_mode=='human'):
 92 |             print(wr.RobotAction(action))
 93 |             self.render()
 94 | 
 95 |         # Return observation, reward, terminated, truncated (not used), info
 96 |         return obs, reward, terminated, False, info
 97 | 
 98 |     # Gym required function to render environment
 99 |     def render(self):
100 |         self.warehouse_robot.render()
101 | 
102 | # For unit testing
103 | if __name__=="__main__":
104 |     env = gym.make('warehouse-robot-v0', render_mode='human')
105 | 
106 |     # Use this to check our custom environment
107 |     # print("Check environment begin")
108 |     # check_env(env.unwrapped)
109 |     # print("Check environment end")
110 | 
111 |     # Reset environment
112 |     obs = env.reset()[0]
113 | 
114 |     # Take some random actions
115 |     while(True):
116 |         rand_action = env.action_space.sample()
117 |         obs, reward, terminated, _, _ = env.step(rand_action)
118 | 
119 |         if(terminated):
120 |             obs = env.reset()[0]
121 | 


--------------------------------------------------------------------------------
/v0_warehouse_robot_train.py:
--------------------------------------------------------------------------------
  1 | '''
  2 | Example of using Q-Learning or StableBaseline3 to train our custom environment.
  3 | '''
  4 | import gymnasium as gym
  5 | import numpy as np
  6 | import matplotlib.pyplot as plt
  7 | import random
  8 | import pickle
  9 | from stable_baselines3 import A2C
 10 | import os
 11 | import v0_warehouse_robot_env # Even though we don't use this class here, we should include it here so that it registers the WarehouseRobot environment.
 12 | 
 13 | # Train or test using Q-Learning
 14 | def run_q(episodes, is_training=True, render=False):
 15 | 
 16 |     env = gym.make('warehouse-robot-v0', render_mode='human' if render else None)
 17 | 
 18 |     if(is_training):
 19 |         # If training, initialize the Q Table, a 5D vector: [robot_row_pos, robot_row_col, target_row_pos, target_col_pos, actions]
 20 |         q = np.zeros((env.unwrapped.grid_rows, env.unwrapped.grid_cols, env.unwrapped.grid_rows, env.unwrapped.grid_cols, env.action_space.n))
 21 |     else:
 22 |         # If testing, load Q Table from file.
 23 |         f = open('v0_warehouse_solution.pkl', 'rb')
 24 |         q = pickle.load(f)
 25 |         f.close()
 26 | 
 27 |     # Hyperparameters
 28 |     learning_rate_a = 0.9   # alpha or learning rate
 29 |     discount_factor_g = 0.9 # gamma or discount rate. Near 0: more weight/reward placed on immediate state. Near 1: more on future state.
 30 |     epsilon = 1             # 1 = 100% random actions
 31 | 
 32 |     # Array to keep track of the number of steps per episode for the robot to find the target.
 33 |     # We know that the robot will inevitably find the target, so the reward is always obtained,
 34 |     # so we want to know if the robot is reaching the target efficiently.
 35 |     steps_per_episode = np.zeros(episodes)
 36 | 
 37 |     step_count=0
 38 |     for i in range(episodes):
 39 |         if(render):
 40 |             print(f'Episode {i}')
 41 | 
 42 |         # Reset environment at teh beginning of episode
 43 |         state = env.reset()[0]
 44 |         terminated = False
 45 | 
 46 |         # Robot keeps going until it finds the target
 47 |         while(not terminated):
 48 | 
 49 |             # Select action based on epsilon-greedy
 50 |             if is_training and random.random() < epsilon:
 51 |                 # select random action
 52 |                 action = env.action_space.sample()
 53 |             else:                
 54 |                 # Convert state of [1,2,3,4] to (1,2,3,4), use this to index into the 4th dimension of the 5D array.
 55 |                 q_state_idx = tuple(state) 
 56 | 
 57 |                 # select best action
 58 |                 action = np.argmax(q[q_state_idx])
 59 |             
 60 |             # Perform action
 61 |             new_state,reward,terminated,_,_ = env.step(action)
 62 | 
 63 |             # Convert state of [1,2,3,4] and action of [1] into (1,2,3,4,1), use this to index into the 5th dimension of the 5D array.
 64 |             q_state_action_idx = tuple(state) + (action,)
 65 | 
 66 |             # Convert new_state of [1,2,3,4] into (1,2,3,4), use this to index into the 4th dimension of the 5D array.
 67 |             q_new_state_idx = tuple(new_state)
 68 | 
 69 |             if is_training:
 70 |                 # Update Q-Table
 71 |                 q[q_state_action_idx] = q[q_state_action_idx] + learning_rate_a * (
 72 |                         reward + discount_factor_g * np.max(q[q_new_state_idx]) - q[q_state_action_idx]
 73 |                 )
 74 | 
 75 |             # Update current state
 76 |             state = new_state
 77 | 
 78 |             # Record steps
 79 |             step_count+=1
 80 |             if terminated:
 81 |                 steps_per_episode[i] = step_count
 82 |                 step_count = 0
 83 | 
 84 |         # Decrease epsilon
 85 |         epsilon = max(epsilon - 1/episodes, 0)
 86 | 
 87 |     env.close()
 88 | 
 89 |     # Graph steps
 90 |     sum_steps = np.zeros(episodes)
 91 |     for t in range(episodes):
 92 |         sum_steps[t] = np.mean(steps_per_episode[max(0, t-100):(t+1)]) # Average steps per 100 episodes
 93 |     plt.plot(sum_steps)
 94 |     plt.savefig('v0_warehouse_solution.png')
 95 | 
 96 |     if is_training:
 97 |         # Save Q Table
 98 |         f = open("v0_warehouse_solution.pkl","wb")
 99 |         pickle.dump(q, f)
100 |         f.close()
101 | 
102 | # Train using StableBaseline3. Lots of hardcoding for simplicity i.e. use of the A2C (Advantage Actor Critic) algorithm.
103 | def train_sb3():
104 |     # Where to store trained model and logs
105 |     model_dir = "models"
106 |     log_dir = "logs"
107 |     os.makedirs(model_dir, exist_ok=True)
108 |     os.makedirs(log_dir, exist_ok=True)
109 | 
110 |     env = gym.make('warehouse-robot-v0')
111 | 
112 |     # Use Advantage Actor Critic (A2C) algorithm.
113 |     # Use MlpPolicy for observation space 1D vector.
114 |     model = A2C('MlpPolicy', env, verbose=1, device='cuda', tensorboard_log=log_dir)
115 |    
116 |     # This loop will keep training until you stop it with Ctr-C.
117 |     # Start another cmd prompt and launch Tensorboard: tensorboard --logdir logs
118 |     # Once Tensorboard is loaded, it will print a URL. Follow the URL to see the status of the training.
119 |     # Stop the training when you're satisfied with the status.
120 |     TIMESTEPS = 1000
121 |     iters = 0
122 |     while True:
123 |         iters += 1
124 | 
125 |         model.learn(total_timesteps=TIMESTEPS, reset_num_timesteps=False) # train
126 |         model.save(f"{model_dir}/a2c_{TIMESTEPS*iters}") # Save a trained model every TIMESTEPS
127 | 
128 | # Test using StableBaseline3. Lots of hardcoding for simplicity.
129 | def test_sb3(render=True):
130 | 
131 |     env = gym.make('warehouse-robot-v0', render_mode='human' if render else None)
132 | 
133 |     # Load model
134 |     model = A2C.load('models/a2c_2000', env=env)
135 | 
136 |     # Run a test
137 |     obs = env.reset()[0]
138 |     terminated = False
139 |     while True:
140 |         action, _ = model.predict(observation=obs, deterministic=True) # Turn on deterministic, so predict always returns the same behavior
141 |         obs, _, terminated, _, _ = env.step(action)
142 | 
143 |         if terminated:
144 |             break
145 | 
146 | if __name__ == '__main__':
147 | 
148 |     # Train/test using Q-Learning
149 |     # run_q(1000, is_training=True, render=False)
150 |     # run_q(1, is_training=False, render=True)
151 | 
152 |     # Train/test using StableBaseline3
153 |     # train_sb3()
154 |     test_sb3()
155 | 


--------------------------------------------------------------------------------