Creating a new dataset
Creation of new Datasets use the Controller
class, which requires two different function to work:
* an enjoy
function, which uses the agent to interact with the environment and record an episode); and
* a collate
function, which puts all files created by the enjoy
function into a dataset file.
Default functions
The imitation_datasets
module provides a set of default functions, so you don't need to implement an enjoy
and a collate
function in every project.
The resulting dataset will be a NpzFile
with the following data:
"""
Data:
obs (list[list[float]): gym environment observation. Size [steps, observations space].
actions (list[float]): agent action. Size [steps, action] (1 if single action, n if multiple actions).
rewards (list[int]): reward from the action with the observations (e.g., r(obs, action)). Size [steps, ].
episode_returns (list[float]): accumulated reward for each episode. Size [number of peisodes, ].
episode_starts (list[bool]): whether the episode started at the current observation. Size [steps, ].
"""
A small functional example of how to use the given functions:
# command: $ python <script> --game cartpole --threads 4 --episodes 1000 --mode all
from imitation_datasets.functions import baseline_enjoy, baseline_collate
from imitation_datasets.controller import Controller
from imitation_datasets.args import get_args
args = get_args()
controller = Controller(baseline_enjoy, baseline_collate, args.episodes, args.threads)
controller.start(args)
Converting to a HuggingFace dataset
If you want to upload the dataset to HuggingFace, the imitation_datasets
module also provides a function for converting the NpzFile
into jsonl
format.
from imitation_datasets.dataset.huggingface import baseline_to_huggingface
baseline_to_huggingface("./path/to/NpZFile", "./new/path")