Transition Recorder Wrapper for Gym Environments

A simple and ubiquitous transition recorder wrapper for Gym Environments to facilitate offline Reinforcement Learning (RL) dataset manufacturing. Transitions (observations, actions, rewards, dones & infos) and episodes (set of transitions) are buffered with LZ4 compression and each episode is saved as a JSON line (.jsonl). Output files may be further processed to be compatible with offline RL libraries like Ray RLlib (see Usage).

Requirements

• Python >= 3.6
• Pypi packages: pip install -r requirements.txt

Usage

Simply import & wrap your Gym environment:

import gym
from gym_recorder import TransitionRecorderWrapper  # import the wrapper

env = gym.make("CartPole-v1")
env = TransitionRecorderWrapper(env)  # wrap your environment
env.reset()

# Use your environment as you would
while True:
    env.render()
    action = env.action_space.sample()
    obs, reward, done, info = env.step(action)
    if done:
        env.reset()

You may also use the save_folder option to customize where the transitions are saved, min_transitions_per_file to customize the output file size, and disable compression with the compress option.

Conversion

The generated .jsonl files may be further processed to be used by offline RL libraries like Ray RLlib. To convert a .jsonl dataset generated by the wrapper to a RLlib offline dataset you may run the following command:

python -m gym_recorder.converters.ray -i data/raw -o data/ray

More options can be found with python -m gym_recorder.converters.ray --help.