This repo is still pre-release. The documentation is still incomplete.

megastep

megastep helps you build 1-million FPS reinforcement learning environments on a single GPU.

Examples

Explorer: 180k FPS, 100 lines of Python.

Deathmatch: 1.2m FPS, 150 lines of Python.

Quoted FPS are for a single RTX 2080 Ti and random actions; visualizations are with a 256-neuron LSTM.

Features

  • Run thousands of environments in parallel, entirely on the GPU.

  • Write your own environments using PyTorch alone, no CUDA necessary.

  • 1D observations. The world is more interesting horizontally than vertically.

  • One or many agents, and one or many cameras per agent.

  • A database of 5000 home layouts to explore, based on Cubicasa5k.

  • A minimal, modular library. Not a framework.

  • (In progress) Extensive documentation, tutorials and explanations.

Install

If you’re working on an Ubuntu machine with CUDA already set up and Python 3.6 or higher, it should be as simple as

pip install megastep[cubicasa,rebar]

To check everything works, run

from megastep.demo import *
env = explorer.Explorer(1)
agent = Agent(env).cuda()
demo(env=env, agent=agent, length=64)

The first time you import megastep on your machine will be slow as it compiles the C++ side of things.

Getting Started

Having installed megastep, you have a few choices.

If you’d like something comprehensive, read Writing a Minimal Environment.

If you’d like a more self-propelled introduction, read Playing With Megastep.

Or, you can simply browse the demos.

Support

I am an independent researcher, so just how much support I can give is limited. But you’re absolutely welcome to open a GitHub issue or see if I’m around on the RL Discord. At the very least I’ll try to point you in the right direction.

You can also find me on Twitter or you can email me.

Of the options, a GitHub issue is likely the best choice as if you have a question, there are likely other people out there with the same exact question! And if they don’t, maybe someone else will have an answer for you.