100% Faster Reinforcement Learning Environments with Cygym | by Jacob Gursky | Towards Data Science

Anyone who at least dabbles in reinforcement learning will likely tell you they have used OpenAI’s Gym package (linked below), and for good reason! It is an easy-to-use, extensible, and well-supported package that reduces much of the overhead in setting up a RL project.
However, when conducting a large-scale project with minimal resources, its essential that every drop of performance be squeezed out of code that is run many times. While working on an RL project, I took a look through the source code of the classic control environments that are often used for benchmarking RL agents to see if we could get some cheap speed gains.
I noticed that all of the gym code is written in pure Python, so the production-code part of my brain immediately went: “What about cythonizing these environments for a speed boost?”
Cython is a superset of Python that enables compilation of interpreted Python code into extremely fast C extensions. What’s nice about this is that we don’t actually have to write any C code, only some minor modifications to our Python script! Depending on the use-case, sometimes you can see 1000x+ speed increases over Python code. The .step() method of gym environments seemed like a good candidate for this, as it is called once for every frame of a game, which can occur millions of times in training an RL agent.
I recently released a public repository containing a Python package called cygym, which contains cythonized versions of the classic control environments from the original gym package. All credit for the original environment code goes to the authors of those scripts, but I did introduce some minor optimizations, as well as leveraging some Cython magic to get up to a 100% speed boost over the original implementations. Hopefully, this number will only go up as we expand our code base and get more contributors.
Installation and setup of cygym are quite easy, simply clone the public package repository and install using pip!
git clone https://github.com/gursky1/cygymcd cygympip install -e . 
Pip should handle all of the setup needed for the package, but in the event that you have issues with installation that way, you can manually build the package yourself and install in an alternative way using setuptools:
python setup.py build_ext --inplace bdist_wheelpip install dist/(wheel file here).whl
If you are still running into issues, please feel free to open an issue or contact me directly!
Cygym is nearly a complete drop-in for OpenAI gym, all you have to do is import the environment classes and get started!
Cygym should also be compatible with any and all OpenAI gym wrappers. One feature we are working on right now is getting the cythonized classes compatible with the OpenAI gym registry, but for Cython reasons this is not super straightforward.
One of the first speed tests we can perform is comparing how long each of the two most common gym methods, .reset() and .step() , take to run in both gym and cygym for our classic control problems.
.reset() is called when we initialize an environment, while .step() is called to apply an action and get the next frame. Most of the run time in interacting with an environment is spent in .step(), so we will place more weight on that comparison:
Unfortunately, it seems that cythonizing only brings a moderate speed boost to initializing environments using .reset(), and sometimes even slows it down! But fear not, this method is only called once per episode, so these results are secondary to our next comparison:
Wow! We can see up to a six times speed boost in iterating our environments using our cythonized variants! Note that most of our environment interaction time is here, meaning our overall model training time should be less dependent on the environment!
Testing two common methods with random actions is all well and good, but how does cygym do in actually training agents? Let’s take a look...
Stable baselines is an excellent fork of the original OpenAI baselines package that contains implementations of numerous RL algorithms such as Deep-Q, TRPO, SAC, PPO, and many more. I won’t dive too deep into this package, but we will be utilizing the vectorized environments available in Stable-Baseslines with our gym and cygym environments to train actual RL agents, and see how long it takes to train our models for 100k frames:
Now we are going to run a similar comparison as the previous sections, looking at all of our classic control environments to see what kind of speed up we can expect using cygym:
Again, wow! It looks like cygym was always faster than gym, with up to a 50% reduction in train time! This can be a massive difference when trying to prototype new model architectures or doing hyperparameter searches.
In the future we plan to look at further compatibility with OpenAI gym, such as using gym.make() with our cythonized environments, and possibly fast rendering (which currently isn’t supported).
We also plan to take a look at cythonizing more complex tasks such as those in the box2d and algorithms modules of the original gym package.
We have shown that simple cythonization with cygym can lead to huge decreases in RL agent train time, meaning less time waiting and more time analyzing results! We are actively looking for contributors and maintainers for cygym, so if you have experience in Cython package development and optimization, please get in touch!
Cygym: Fast gym-compatible classic control RL environments
github.com
OpenAI Gym: A versatile package for reinforcement learning environments
github.com
Cython: A compiled superset of the Python language
cython.org
Stable-Baselines: A fork of OpenAI baselines