Custom Environments in OpenAI’s Gym | Towards Data Science

OpenAI’s Gym is (citing their website): “... a toolkit for developing and comparing reinforcement learning algorithms”. It includes simulated environments, ranging from very simple games to complex physics-based engines, that you can use to train reinforcement learning algorithms. OpenAI’s other package, Baselines, comes with a number of algorithms, so training a reinforcement learning agent is really straightforward with these two libraries, it only takes a couple of lines in Python.
Of course, there is not much fun in only using the built-in elements, you either use the algorithms as benchmarks when you test your own algorithm, or you build a custom environment to solve your own problem with any of the available algorithms. Right now, we are interested in the latter: we are going to set up a custom environment in Gym with Python.
You might assume you can just follow guidelines in the Gym Documentation, but that is not entirely correct. I had to hunt down and compile the information from multiple sources (documentation, GitHub, Stack Overflow, etc), so I figured I should write a clean and simple summary.
Apart from the OpenAI Gym library, we are also going to use a package called Stable Baselines — a project that started as a fork of the OpenAI Baseline library’s reinforcement learning algorithms, with the intention to make it more documented and more user-friendly.
One more source that I used is deeplizard’s step-by-step guide on implementing a simple Q-learning algorithm — we will use this code to test whether our environment can be trained on.
All the code can be found on my GitHub, most importantly in gym_basic_env_test.ipynb. The rest of the repo is a Gym custom environment that you can register, but, as we will see later, you don’t necessarily need to do this step.
After working through the guide, you’ll be able to:
Set up a custom environment that is consistent with Gym.Develop and register different versions of your environment.Decide what exactly needs to go in the environment class — in my opinion, this is the trickiest part of the whole project!Train a simple Q-learning algorithm on your new environment.Utilise Stable Baselines in the environment, use the package to verify that your custom environment is Gym-consistent, and train a more sophisticated algorithm.
Set up a custom environment that is consistent with Gym.
Develop and register different versions of your environment.
Decide what exactly needs to go in the environment class — in my opinion, this is the trickiest part of the whole project!
Train a simple Q-learning algorithm on your new environment.
Utilise Stable Baselines in the environment, use the package to verify that your custom environment is Gym-consistent, and train a more sophisticated algorithm.
(Just to manage expectations, I’d like to mention that the environment we set up is going to be the simplest “game” I could come up with, with one singular choice — our aim is just to understand the structure.)
I won’t go through the reinforcement learning basics, you will need to be familiar with at least the following terms: agent, environment, state, action. The following concepts will also indirectly come up when we do the Q-learning algorithm, but they are less important: Q-values, epsilon-greedy strategy, learning rate, exploration vs. exploitation.
If you are not familiar with these terms, there are tons of good materials out there. The Stable Baselines Documentation has a list of useful links, and I personally used the code of the Q-learning section on deeplizard, they also have nice summaries of the basic stuff.
We need to install Gym and Stable Baselines before we start working.
Quite straightforward, you can follow the documentation. The most convenient method is to use pip:
pip install gym
That’s it.
This one is trickier. For starters, you’ll need some prerequisites, have a look at the documentation. Once that’s done, you can pip install the package. If you are on a Mac, this is all you need to do:
brew install cmake openmpipip install stable-baselines[mpi]
Now, the issue comes when you are starting to use deep learning algorithms that are built on top of tensorflow — as of 2020 December, the Stable Baselines package only supports tensorflow up until version 1.15.0. There is a warning message on the Installation guide page, and if you are interested, you can read more about the cause of the problem on GitHub.
Bottom line is, you need an older version of tensorflow. I did not use tensorflow for anything else, so I just installed an older version like so:
pip install tensorflow==1.15.0
Now that we have the libraries, we can start working on the custom environment.
If you try doing a custom environment from the official guide, chances are, you will start with building this file structure:
gym-basic/  README.md  setup.py  gym_basic/    __init__.py    envs/      __init__.py      basic_env.py      basic_env_2.py
The thing is, it’s not... You don’t actually need to worry about this whole file structure thing, the only thing that really matters is basic_env.py. When I started working on this project, I assumed that when you later build your environment from a Gym command:
env = gym.make(“gym_basic:basic-v0”)
something magical happens in the background, but it seems to me you get the same result if you simply initiate an object from your environment class:
env = BasicEnv()
(You need to import it first of course.) My understanding is that the main advantage of registering your environment in your Gym package is that you can keep all your environments in one place, and you can also run Baseline RL algorithms (original OpenAI and Stable Baselines likewise) in one line:
python -m baselines.run --alg=<name of the algorithm> --env=<environment_id> [additional arguments]
So if you want to register your Gym environment, follow this section, otherwise, skip ahead to the next section, The Environment Class.
You can use the documentation for this part, or my GitHub repository is basically also a Gym custom environment (if you ignore the two Jupyter Notebooks). You need to set up a setup.py, two __init__.py's, and most importantly, your environment file(s), which we will do in the next section. Pay attention to the difference between the two __init__.py’s, they handle the different versions of your environment.
For your convenience, I’m copying the content of the three files here. In this example, we are setting up an environment called basic, and do two versions of it. Naturally, just change the string to whatever you want to call your environment.
from setuptools import setup setup(name=’gym_basic’, version=’0.0.1', install_requires=[‘gym’] )
from gym.envs.registration import register register(id='basic-v0',entry_point='gym_basic.envs:BasicEnv',) register(id='basic-v2',entry_point='gym_basic.envs:BasicEnv2',)
(We will have two environments, BasicEnv, with id = 'basic-v0', and BasicEnv2, with id = basic-v2. This way, you can keep different versions of the environment in the same registry.)
from gym_basic.envs.basic_env import BasicEnvfrom gym_basic.envs.basic_env_2 import BasicEnv2
Hang on, there was another file in there, the README markdown, what do I put in there? Well, that is just your git repo readme file, it has nothing to do with your Gym environment.
As long as you have the corresponding environment files in the envs folder (basic_env.py and basic_env_2.py in our case), you can register your environment. Just navigate up in your terminal to the folder that contains gym-basic, and use pip:
pip install -e gym-basic
If you update the environment .py files later, it should update your environment automatically.
All right, we registered the Gym environment. We can finally concentrate on the important part: the environment class.
Here, I think the Gym documentation is quite misleading. There are parts in the foo_env.py template they provide that you don’t need to include, and there are missing parts that you need. Stable Baselines has a better template, but it was still not entirely clear to me.
Once again, my aim is to create the absolute minimum environment here.
The environment needs to be a class inherited from gym.Env.
In __init__, you need to create two variables with fixed names and types. You need a self.action_space, and a self.observation_space. These two need to be of Gym’s special class, space, which is not outright tricky, but not entirely straightforward either. Have a look at the documentation, there are basically two types: a one-dimensional called Discrete, and an n-dimensional called Box. In this example, we will keep things simple and only use Discrete. (One could argue that you can always re-code your spaces to one-dimensional id codes.)
The reset function returns a value that is within self.observation_space. This is the function that will re-start the environment, for example, at the beginning of a game. In my environments, I call the variable that keeps track of the current state... well, state, but there is no restriction there, you can call it whatever. If we use a Discrete observation space, this state value needs to be an integer. With a Box, it can be a numpy.array.
The step function has one input parameter, needs an action value, usually called action, that is within self.action_space. Similarly to state in the previous point, action can be an integer or a numpy.array. This is the response function, the agent provides the action they want to take, and the environment returns the state it brought them to. The return value is a 4-tuple, in the following order (the naming does not matter, but the variable type does):- state, same type of variable as the return of the reset function, of self.observation_space;- reward, a number that informs the agent about the immediate consequences of its action;- done, a boolean, value that is TRUE if the environment reached an endpoint, and should be reset, or FALSE otherwise;- info, a dictionary that can be used in bug fixing, and as far as I know, the baselines algorithms don’t use it.You need to provide these variables, but they don’t all have to have “proper“ values. For example, the info can be just an empty dictionary.
It’s important to note that you can have as many helping functions within your class as you want, and it doesn’t matter what you call them, as long as they don’t use any of the reserved names.
Now, I’m not saying that having these in an environment is a bad idea, but I think it’s important to distinguish between “absolute requirements” and “things that you will probably need later”. I’m a big fan of trying not to mindlessly copy code and clutter our projects with unnecessary stuff.
So next is a list of things in the Gym documentation that you as a beginner probably won’t need right away when you start playing with custom environments:
metadata = {‘render.modes’: [‘human’]}: This line simply defines possible types for your render function (see next point).
You don’t actually need a render function. I mean, naturally, it comes handy when you try to do bug fixing, but in terms of training a reinforcement algorithm, as far as I noticed, the baseline algorithms don’t assume it to be there, and they are not using it. If you think about it, the purpose of render is simply to convert your state variable to something more readable.
Similarly, you don’t actually need a close function either. The GitHub page of the core environment has some comments on this function, this would be called for additional cleanup when the environment is closed or picked up by the garbage collector.
Same with the seed function. This will be useful when you want to have reproducibility in an environment that uses random number generators.
Please note that the class from which you inherit (gym.env) will have these functions implemented, with a single pass line in each of them. So strictly speaking, you will always have these functions, but don’t need them to do anything.
All right, and now, let’s see a custom environment! As I mentioned earlier, this is going to be very simple. We start at state = 0, and the agent can choose between 5 actions, numbered 0 to 4. If they pick action = 2, they get a reward of 1, otherwise, the reward is -1. (If you would like to see something that’s a tiny bit more complicated, have a look at basic_env_2.py in my GitHub repo.)
The way you can code it as a Gym environment:
import gymclass BasicEnv(gym.Env):def __init__(self):        self.action_space = gym.spaces.Discrete(5)        self.observation_space = gym.spaces.Discrete(2)def step(self, action):        state = 1            if action == 2:            reward = 1        else:            reward = -1                    done = True        info = {}        return state, reward, done, infodef reset(self):        state = 0        return state
That’s it, you have a Gym environment!
Note that the state is always kept within the scope of the function that uses it, and it’s implied that if the environment takes a step (via the step function), it moves from state 0 to 1, and then the “game” immediately ends. So there is no need to keep track of the state within the object in this simple game, but generally, you would want to have a self.state variable.
As mentioned above, you can create an object with either
env = gym.make(“gym_basic:basic-v0”)
if you registered the environment (see File Structure section above), or simply with
env = BasicEnv()
if you have not registered it.
In the remainder of the post, I will assume that env is initiated and is a BasicEnv Gym environment.
One cool thing about Stable Baselines is that it comes with a function that verifies if your environment is Gym-compatible.
Running this checker is as simple as:
from stable_baselines.common.env_checker import check_envcheck_env(env)
If you followed the tutorial, the function will not return anything. Which is a bit weird, but it means that everything is OK.
You can test what happens if you change any of the elements in your environment. For example, if you don’t have a reset function, you will get a NotImplementedError. If you don’t have a proper self.action_space and self.observation_space, for example if you defined them as a regular list instead of the special space class, you will get this error: AssertionError: The action space must inherit from gym.spaces.
Notice how the checker function did not complain about not having any of the stuff we listed as unnecessary earlier!
I am following the Q-learning algorithm implementation from deeplizard. In this section, we are repeating the tutorial, but we replace the environment with our own.
Just like with the built-in environment, the following section works properly on the custom environment. The Gym space class has an n attribute that you can use to gather the dimensions:
action_space_size = env.action_space.nstate_space_size = env.observation_space.nq_table = np.zeros((state_space_size, action_space_size))print(q_table)
returns:
[[0. 0. 0. 0. 0.] [0. 0. 0. 0. 0.]]
These are the environment’s state-action Q-values. Please note that the first row, corresponding with state 0, is the only one we care about, the second row of state 1 is the end state, no action will be taken there. I’m not sure if this is usual practice or not, it seemed logical to me to keep the end state separate.
I’m not going to copy the whole code here, please refer to my GitHub repo.
After the model is trained, the updated Q-values look like this:
[[-0.94185026 -0.89058101  1.         -0.92023356 -0.91137062] [ 0.          0.          0.          0.          0.        ]]
Which indicates that the model indeed learned the environment, and identified action = 2 as the optimal action to take.
If you think that this was a slight overkill, you are right, but I think it is interesting to see how the environment would work in practice. Plus, wait until what we do in the next section!
In this section, we are going to dial up the overkill a bit, and train a Proximal Policy Optimizing Algorithm on our humble environment. (Full disclosure: I have no idea how this algorithm works, or what its advantages are, I simply picked it because it is in the tutorial.)
Long story short, taking the code in the Getting Started Guide on Stable Baselines, and replacing the env with our own (once again, by running either env = gym.make(“gym_basic:basic-v0”), or env = BasicEnv(), this code works perfectly:
from stable_baselines.common.policies import MlpPolicyfrom stable_baselines.common.vec_env import DummyVecEnvfrom stable_baselines import PPO2model = PPO2(MlpPolicy, env, verbose=1)model.learn(total_timesteps=10000)obs = env.reset()for i in range(10):    action, _states = model.predict(obs)    print(action)    obs, rewards, dones, info = env.step(action)    env.render()
(Our render function does not do anything, so I added in a line to print out our action.)
Quite reassuringly, the model was able to learn how to play our game correctly, and picks action = 2 in every case.