A library to build and train reinforcement learning agents in OpenAI Gym environments.

Build Status Documentation Status

Read full documentation here.

Getting Started

An agent has to implement the act() method which takes the current state as input and returns an action:

from train import Agent

class RandomAgent(Agent):

    def act(self, state):
        return self.env.action_space.sample()

Create an environment using OpenAI Gym:

import gym

env = gym.make('CartPole-v0')

Initialize your agent using the environment:

agent = RandomAgent(env=env)

Now you can start training your agent (in this example, the agent acts randomly always and doesn’t learn anything):

scores = agent.train(episodes=100)

You can also visualize how the training progresses but it will slow down the process:

scores = agent.train(episodes=100, render=True)

Once you are done with the training, you can test it:

scores = agent.test(episodes=10)

Alternatively, visualize how it performs:

scores = agent.test(episodes=10, render=True)

To learn more about how to build an agent that learns see agents documentation.

See examples directory to see implementations of some algorithms (DQN, A3C, PPO etc.) created using TensorFlow, PyTorch and NN libraries.



  • Python >= 3.6

Install from PyPI (recommended):

pip install train

Alternatively, install from source:

git clone
cd train
pip install -e .

To run examples and tests, install from source.

Other libraries such as Gym, TensorFlow, PyTorch and NN should be installed separately.


To run examples, install TensorFlow, PyTorch and install other dependencies:

pip install -e .[examples]

and run an example in examples directory:

cd examples


To run tests, install dependencies:

pip install -e .[tests]

and run:

pytest tests