Mario’s Gym Routine
By: Robert Clark
Tags: Reinforcement Learning, Mario, Python, Nvidia, Editors Pick

Can you teach an old dog new tricks?

A bit over a year ago, I saw a video of a Python program which was able to train a machine to play and win a game of Atari’s Pong through a process known as Reinforcement Learning (RL). Since viewing that video and my mind being subsequently blown by the potential of this application of AI, I have gone on a journey to learn how to leverage RL to beat various retro video games, and even documented my early stages in a previous post.

Throughout the past year, I’ve written various projects from scratch which aim to beat levels in the original Super Mario Bros. on the NES, but wanted to find a better way to switch between models and parameters without writing several hundred additional lines in each scenario. After experimenting with many RL libraries available in Python, I found rllib from the Ray Project to be the most effective and flexible, while also being able to scale up and out on machines — a requirement for my desires to take on larger and more complex environments. This article is intended to provide an example of how to build your own application which can be used to train an agent to beat Super Mario Bros. levels.

Baby Steps: Setting up an environment

Before I jump into the code, a development environment needs to be created to install dependencies and get your machine ready to learn. For installing dependencies in an isolated location for my project, I use either a virtualenv or conda environment to ensure I won’t be installing anything that could effect other apps on my computer. While I won’t go over setup for either in this post, you can read the official docs for either virtualenv or conda.

Inside an active environment of your choice, several dependencies need to be installed for the program to function. These can be installed with Python’s PIP:

pip install ray[rllib] torch gym-super-mario-bros

A more concrete list of requirements can be found on the repository I created for this project, but note that while the repository and dependencies may change with time, the dependencies will always stay compatible with the latest upstream code.

While this code might work on other Python versions (such as Python 3.6), I have only tested it on 3.7 and newer versions, which is highly recommended.

View from above

Whenever I look at examples for a new application, framework, method, or library, I always prefer to look at the entire chunk of sample code before diving into specifics. This allows me to see where the application begins, what needs to be initialized, any arguments that are parsed, which libraries are imported, and the overall structure of the code. For me personally, it’s much more valuable to have a high-level understanding of the code before diving into specific lines without any context, and that’s exactly what I’m going to do here.

The following example is a completely self-contained Python program which can train Mario to beat a specified level of the game. For those that like to get their hands dirty and just run with this code and don’t need explanations, enjoy! I will dive into more detail on each block of code for those sticking around. But, without further ado, here’s a sample application to give Mario a new brain:

Focusing in

With the full application out of the way, let’s look at each section in greater detail to understand what it does.

Import modules

We first need to import several libraries to make our lives easier. This is one of the great aspects of Python — there’s a high chance that something you need has already been created and is available in a library. I won’t mention all of the imports, but some of the key ones are as follows:

gym_super_mario_bros : An OpenAI Gym-compatible environment which allows Python programs to interact with Super Mario Bros. seamlessly.

: An OpenAI Gym-compatible environment which allows Python programs to interact with Super Mario Bros. seamlessly. ray : A distributed computing framework which makes it easy to run processes across multiple workers and machines.

: A distributed computing framework which makes it easy to run processes across multiple workers and machines. nes_py : A framework to interface between NES game environments, such as Super Mario Bros. and Python applications.

: A framework to interface between NES game environments, such as Super Mario Bros. and Python applications. rllib : A reinforcement learning library built on top of Ray which includes several robust models and algorithms to easily build and scale RL applications.

Custom wrappers

One of the great features of OpenAI’s gym is how easy it is to create custom wrappers around an environment to alter various aspects of the game or virtual world, such as modifying image shapes and sizes, stacking multiple frames together, changing reset conditions, and many, many more.

In some cases, as is the case with playing games from the Atari environment with OpenAI’s Gym, it is possible to get away with using the various atari_wrappers which are included with most RL libraries, without the need to add any additional wrappers. Most of these wrappers are imported above, but I had to update one of them, the EpisodicLifeEnv wrapper, as it expects the environment to have an Atari Learning Environment (ALE) object, which our Super Mario Bros. environment does not contain. I simply copied the code for this class as-is from the ray.rllib.env.atari_wrappers module, but modified lines 18 and 37 above to point to the unwrapped property instead of the ale . By including this code locally in my application and changing those two lines, I still get the full functionality of the wrapper (which I will explain further below) but it won’t complain that my environment doesn’t have an ale component.

I also modified a snippet of code from Uvipen’s awesome Super-mario-bros-A3C-pytorch repository which updates the reward each learner receives after every step. By default, gym_super_mario_bros determines the reward at each step by calculating Mario’s velocity (positive points while moving right, negative points while moving left, zero while standing still), plus a penalty for every frame that passes to encourage movement, and a penalty if Mario dies for any reason. While this is a fairly robust reward system, the levels could be played out more “normally” (as in, the way most humans would play them) by rewarding Mario for increasing his in-game score by defeating enemies, grabbing coins, and collecting power-ups. This information is saved in the info dictionary, and the previous score can be compared with the current game score to find a difference in the latest step and add that to the reward.

In addition to improving the reward when Mario increases his in-game score, a sizable reward is added if he collects the flag (or defeats Bowser) at the end of the level to encourage him to successfully beat the stage. He also receives a relatively large penalty if he doesn’t make it to the end of the level before he dies.

This new reward is then scaled down to be in a smaller range, which is similar to a reward clipping technique applied in the original DQN paper by Minh et. al.

Parsing arguments

Up next we parse arguments passed by the CLI during runtime. The user can modify a few components of the application by specifying various flags. These flags have the following impacts:

--checkpoint : The program automatically creates a new checkpoint file after every 50 training iterations and saves it to a spot that is indicated during runtime, typically at ~/ray_results/<Training instance>/checkpoint_<n>/ . By providing the full path to this file, a run can use that checkpoint’s weights to start a new training pass which ideally has a bit of progress already, dramatically reducing the amount of time required to train an agent, and also making it possible to use transfer learning to take Mario’s knowledge of one level and apply it to another.

: The program automatically creates a new checkpoint file after every 50 training iterations and saves it to a spot that is indicated during runtime, typically at . By providing the full path to this file, a run can use that checkpoint’s weights to start a new training pass which ideally has a bit of progress already, dramatically reducing the amount of time required to train an agent, and also making it possible to use transfer learning to take Mario’s knowledge of one level and apply it to another. --dimension : Each frame is cropped down to a square NxN image which makes it faster and easier to process details in an image by a neural network. Minh et al. used an image size of 84x84 which is commonly used for most retro video games, though sometimes other dimensions are used, such as 42x42.

: Each frame is cropped down to a square NxN image which makes it faster and easier to process details in an image by a neural network. Minh et al. used an image size of 84x84 which is commonly used for most retro video games, though sometimes other dimensions are used, such as 42x42. --environment : This is the OpenAI gym environment to train against. The gym_super_mario_bros environments have the format SuperMarioBros-<world>-<level>-<variant> , where <world> is the game’s world number (1–8), <level> is the level number within that world (1–4), and <variant> modifies how the game looks, from an untouched state ( v0 ), to colored blocks representing all objects, enemies, and players in game ( v3 ). It is recommended to stick with v0 for the game variant as most robust RL agents are able to handle the natural game state.

: This is the OpenAI gym environment to train against. The environments have the format , where is the game’s world number (1–8), is the level number within that world (1–4), and modifies how the game looks, from an untouched state ( ), to colored blocks representing all objects, enemies, and players in game ( ). It is recommended to stick with for the game variant as most robust RL agents are able to handle the natural game state. --framestack : Multiple frames can be stacked together to indicate motion of various objects. This helps the agent understand how to react to a certain situation as it is able to infer where objects will be in future states. This is similar to how humans interpret environments and states. If you were to take a random screenshot of a game and share it with your friend without any context, they might struggle to determine which direction each object in the screen was moving in. If, however, you were to take four screenshots from four connected frames and share those screenshots in order with your friend, they would likely have an easier time deciphering the direction of motion for each object in the image as they could compare any differences between the sequence. While the default value of 4 is recommended, different levels can be experimented with to see the effect.

: Multiple frames can be stacked together to indicate motion of various objects. This helps the agent understand how to react to a certain situation as it is able to infer where objects will be in future states. This is similar to how humans interpret environments and states. If you were to take a random screenshot of a game and share it with your friend without any context, they might struggle to determine which direction each object in the screen was moving in. If, however, you were to take four screenshots from four connected frames and share those screenshots in order with your friend, they would likely have an easier time deciphering the direction of motion for each object in the image as they could compare any differences between the sequence. While the default value of 4 is recommended, different levels can be experimented with to see the effect. --gpus : If users have compatible NVIDIA GPUs installed on their system, they can use those resources to dramatically improve the overall throughput and decrease the time required to train an agent. At the moment, the implementation outlined in this article supports a maximum of 1 GPU, but this will hopefully be changed in future versions of RLLib.

: If users have compatible NVIDIA GPUs installed on their system, they can use those resources to dramatically improve the overall throughput and decrease the time required to train an agent. At the moment, the implementation outlined in this article supports a maximum of 1 GPU, but this will hopefully be changed in future versions of RLLib. --iterations : This flag tells RLLib how many training passes it should go through prior to terminating. Note that this is independent from the number of steps that have been taken during execution.

: This flag tells RLLib how many training passes it should go through prior to terminating. Note that this is independent from the number of steps that have been taken during execution. --workers : Workers are the system resources that actually deploy the agents which will independently learn and update global weights with any progress (depending on the model). This number can never be more than the available number of CPU cores on a system, and is recommended to be one minus the total number of cores in a system, with the last one being a controller which deploys Ray and listens for all updates and broadcasts changes to the cluster.

Creating an environment

Before feeding it into a neural network, an environment needs to be built and wrapped with our modifications to make it easier for the agent to learn. I will go down these modifications line-by-line.

gym_super_mario_bros.make : This is a simple wrapper around the gym.make function and allows us to build the specified SuperMarioBros environment and returns the respective env.

: This is a simple wrapper around the gym.make function and allows us to build the specified SuperMarioBros environment and returns the respective env. CustomReward : This modifies the reward returned after each step following the changes listed earlier in this article.

: This modifies the reward returned after each step following the changes listed earlier in this article. JoypadSpace : When the network picks an action to take, it typically selects a zero-based number from a list which represents an action that the agent can take. For example, the first item in a list of actions might be for the agent to stay in place, the second item to go right, the third to go left, and so on. Adding the SIMPLE_MOVEMENT constant from the gym_super_mario_bros actions to the wrapper maps Mario’s actions to a list of pre-defined movements. There are three action lists available — RIGHT_ONLY , SIMPLE_MOVEMENT , and COMPLEX_MOVEMENT which increase in the number of possible actions that Mario can make, ranging from only being able to run right and jump to unlocking his entire toolset of abilities and motion. SIMPLE_MOVEMENT is a nice balance between offering Mario a range of movement while not bloating the action space with a large number of possible actions, making it harder to learn.

: When the network picks an action to take, it typically selects a zero-based number from a list which represents an action that the agent can take. For example, the first item in a list of actions might be for the agent to stay in place, the second item to go right, the third to go left, and so on. Adding the constant from the actions to the wrapper maps Mario’s actions to a list of pre-defined movements. There are three action lists available — , , and which increase in the number of possible actions that Mario can make, ranging from only being able to run right and jump to unlocking his entire toolset of abilities and motion. is a nice balance between offering Mario a range of movement while not bloating the action space with a large number of possible actions, making it harder to learn. MonitorEnv : The MonitorEnv keeps track of episode statistics after each run which can be viewed with the overall results. While this wrapper doesn’t change the environment, it does provide many useful metrics to evaluate overall performance.

: The keeps track of episode statistics after each run which can be viewed with the overall results. While this wrapper doesn’t change the environment, it does provide many useful metrics to evaluate overall performance. NoopResetEnv : The initial state of an environment is sampled by taking a random number of no-op, or non-actions in a row after reset. The no-op is generally assumed to be the first action in the action space, which in our case correlates to Mario standing still.

: The initial state of an environment is sampled by taking a random number of no-op, or non-actions in a row after reset. The no-op is generally assumed to be the first action in the action space, which in our case correlates to Mario standing still. EpisodicLifeEnv : This wrapper uses the modified version of the class that was shown above with the custom wrappers to add support to SuperMarioBros environments. The wrapper modifies the environment in a way that any time the player loses a life, it is considered the end of an episode, but a reset only happens when the game is over. This helps with estimation of future policies to aid in learning.

: This wrapper uses the modified version of the class that was shown above with the custom wrappers to add support to SuperMarioBros environments. The wrapper modifies the environment in a way that any time the player loses a life, it is considered the end of an episode, but a reset only happens when the game is over. This helps with estimation of future policies to aid in learning. WarpFrame : This wrapper first converts all images to a single-channel grayscale image, then resizes them down to an NxN square image as specified by the --dimension flag. This makes it quicker and easier for the network to find key features in each observation.

: This wrapper first converts all images to a single-channel grayscale image, then resizes them down to an NxN square image as specified by the flag. This makes it quicker and easier for the network to find key features in each observation. FrameStack : This stacks the specified number of frames from the --framestack flag together in a sequence to show motion of objects. The motivation for this wrapper was outlined in the argument parser section above, but the default setting is to stack four frames together to illuminate subtle differences between frames to allow the network to decipher movement.

Printing progress

This function is not strictly necessary, but I find the standard output which can be printed by RLLib to be fairly verbose, as I prefer to run a training loop in the background and only periodically look at certain values, such as the max, min, and mean rewards. The tabulate library makes it very easy to print clean tables of information which is what I use here, similar to what RLLib does with their output. If desired, this function can be omitted, and the result object below can be pretty_print ‘ed instead.

Main function

The main function in this case is where a lot of the magic happens and there’s a lot going on (which slightly pains me given how I normally write Python applications with a slim main function, but I wanted to be relatively concise in this case).

Breaking it down, the function starts with a env_creator_lambda definition which allows us to build a custom environment. This is necessary as RLLib takes a gym environment’s name by default and builds that environment at the beginning of a training run. As we added several custom wrappers to our environment, they need to be included in a new function which is handled specially by RLLib to let our Python application build the environment instead of RLLib doing it automatically.

I borrowed the IMPALA config tuned specifically for Atari Pong from the official Ray repository as a baseline for the Super Mario Bros. config. I chose IMPALA for my agent as I have a fairly beefy workstation which allows me to leverage a few dozen workers and two GPUs to accelerate the training process. IMPALA tends to be on the faster side amongst some of the common RL-models when given enough resources, hence my usage here. One of the reasons I moved to a RL-library, however, was to be able to seamlessly transition between various models without writing a bunch of new code to support new algorithms. Moving forward, I plan on adding new models in this application.

Eagle-eyed readers will also notice I use PyTorch as my framework of choice. While I haven’t done any comparisons for this specific application between PyTorch and TensorFlow (both are supported by RLLib), I tend to pick PyTorch when given the option as I prefer the framework’s more Pythonic and intuitive style, generally better performance, and seamless support for GPUs. Assuming both packages are installed, switching between the two frameworks is as easy as changing the value of the framework setting.

After initializing the config settings, several Ray-specific functions are called to startup a cluster on the current node, register an environment using our custom environment creation function with wrappers above using register_env , and building an IMPALA trainer based on the config that we created.

If the user specified a checkpoint file during runtime, the trainer will read the requested checkpoint file, and update the trainer’s weights based on that checkpoint prior to starting the run.

Lastly, the training sequence is created by looping through the requested number of iterations and calling the train() method against the trainer to begin that phase of the pipeline. This will automatically update the model weights as progress is made and new policies are learned. The results from each iteration are then printed using the custom function created earlier to display the episode rewards over time. On every 50th iteration, a new checkpoint file is created and saved alongside model parameters, typically in the ~/ray_results/ directory.

Running the application

At this point, you should now be able to replicate the training environment on your own machine by copying the first code snippet above, installing dependencies, and running the code with any necessary parameter changes.

If the code was saved as train.py , the training can be kicked off by simply running the following:

$ python3 train.py

Note that if you have a different number of available CPU cores than the default four or have NVIDIA GPUs, you can specify the result with the --workers and --gpus flags, respectively, similar to the following:

$ python3 train.py --workers 60 --gpus 1

Photo by Victor Freitas on Unsplash

Italian Plumbers Pumping Iron

Congratulations! If you followed this guide, you should now be able to train Mario to beat various levels of the original game that started it all. It may take several hours to complete a level depending on resources, and some are harder than others, but this offers a base foundation to achieve greater heights in the exciting realm of reinforcement learning.

Have any suggestions for improvement? Feel free to drop a new issue or pull request on my repository or comment below. And as always, enjoy the new adventures with everyone’s favorite Italian plumber. Wahoo!