Skip to content

Large-Scale and Comprehensive Data Hub for Reinforcement Learning

License

Notifications You must be signed in to change notification settings

RLE-Foundation/rllte-hub

Repository files navigation



RLLTE Hub: Large-Scale and Comprehensive Data Hub for Reinforcement Learning

Contents

Overview

RLLTE Hub is a repository of multifarious trained models and datasets of reinforcement learning (RL). The following table illustrates its architecture:

Module Function
rllte.hub.*** 📊 .load_curves: Load learning curves of an RL algorithm on a task.
💯 .load_scores: Load test scores of an RL algorithm on a task.
🗃️ .load_models: Load a trained RL agent on a task.
🎮 .load_apis: Load a training API.

A complete support list for RL algorithms and environments can be found in https://docs.rllte.dev/hub.

Installation

Developers can invoke the hub module in rllte directly. Open a terminal and install rllte with pip:

pip install rllte-core

We Provide

Trained RL Models

The following example illustrates how to download an PPO agent trained the Atari benchmark:

from rllte.hub import Atari

agent = Atari().load_models(agent='ppo',
                            env_id='BeamRider-v5',
                            seed=0,
                            device='cuda')
print(agent)

Use the trained agent to play the game:

from rllte.env import make_envpool_atari_env
from rllte.common.utils import get_episode_statistics
import numpy as np

envs = make_envpool_atari_env(env_id="BeamRider-v5",
                              num_envs=1,
                              seed=0,
                              device="cuda",
                              asynchronous=False)

obs, infos = envs.reset(seed=0)
episode_rewards, episode_steps = list(), list()
while len(episode_rewards) < 10:
    # The agent outputs logits of the action distribution
    actions = th.softmax(agent(obs), dim=1).argmax(dim=1)
    obs, rewards, terminateds, truncateds, infos = envs.step(actions)

    eps_r, eps_l = get_episode_statistics(infos)
    episode_rewards.extend(eps_r)
    episode_steps.extend(eps_l)    

print(f"mean episode reward: {np.mean(episode_rewards)}")
print(f"mean episode length: {np.mean(episode_steps)}")

# Output:
# mean episode reward: 3249.8
# mean episode length: 3401.1

RL Training Logs

Download training logs of various RL algorithms on well-recognized benchmarks for academic research.

Training Curves

The following example illustrates how to download training curves of the SAC agent on the DeepMind Control Suite benchmark:

from rllte.hub import DMControl

curves = DMControl().load_curves(agent='sac', env_id="cheetah_run")

This will return a Python Dict of NumPy array like:

curves
├── train: np.ndarray(shape=(N_SEEDS, N_POINTS))
└── eval:  np.ndarray(shape=(N_SEEDS, N_POINTS))

Visualize the training curves:

Test Scores

Similarly, download the final test scores via

scores = DMControl().load_scores(agent='sac', env_id="cheetah_run")

This will return a data array with shape (N_SEEDS, N_POINTS).

RL Training Applications

Developers can also train RL agents on well-recognized benchmarks rapidly using simple interfaces. Suppose we want to train an PPO agent on Procgen benchmark, it suffices to write a train.py like:

from rllte.hub import Procgen

app = Procgen().load_apis(agent="PPO", env_id="coinrun", seed=1, device="cuda")
app.train(num_train_steps=2.5e+7)

All the curves, scores, and models were trained via .load_apis(), and all the hyper-parameters can be found in the reference of the support list.

Cite the Project

If you use this project in your research, please cite this project like this:

@article{yuan2023rllte,
  title={RLLTE: Long-Term Evolution Project of Reinforcement Learning}, 
  author={Mingqi Yuan and Zequn Zhang and Yang Xu and Shihao Luo and Bo Li and Xin Jin and Wenjun Zeng},
  year={2023},
  journal={arXiv preprint arXiv:2309.16382}
}