nemo_rl.data.datasets.response_datasets.nemogym_dataset#

Module Contents#

Classes#

NemoGymDataset

Simple wrapper around the Nemo Gym dataset.

API#

class nemo_rl.data.datasets.response_datasets.nemogym_dataset.NemoGymDataset(data_path: str, repeat: int = 1, **kwargs)#

Bases: nemo_rl.data.datasets.raw_dataset.RawDataset

Simple wrapper around the Nemo Gym dataset.

Parameters:

  • data_path – Path to a JSONL file or a pre-converted Arrow/Parquet dataset.

  • repeat – Number of times to repeat the dataset, default is 1

Initialization