DataCollectorV0#

minari.DataCollectorV0#

class minari.DataCollectorV0(env: Env, step_data_callback: Type[StepDataCallback] = StepDataCallback, episode_metadata_callback: Type[EpisodeMetadataCallback] = EpisodeMetadataCallback, record_infos: bool = False, max_buffer_steps: int | None = None, max_buffer_episodes: int | None = None)[source]#

Gymnasium environment wrapper that collects step data.

This wrapper is meant to work as a temporary buffer of the environment data before creating a Minari dataset. The creation of the buffers that will be convert to a Minari dataset is agnostic to the user:

import minari
import gymnasium as gym

env = minari.DataCollectorV0(gym.make('EnvID'))

env.reset()

for _ in range(num_steps):
    action = env.action_space.sample()
    obs, rew, terminated, truncated, info = env.step()

    if terminated or truncated:
        env.reset()

dataset = minari.create_dataset_from_collector_env(dataset_id="env_name-dataset_name-v(version)", collector_env=env, **kwargs)

Some of the characteristics of this wrapper:

  • The step data is stored per episode in dictionaries. This dictionaries are then stored in-memory in a global list buffer. The episode dictionaries contain items with list buffers as values for the main episode step datasets observations, actions, terminations, and truncations, the infos key can be a list or another nested dictionary with extra datasets. Separate data keys can be added by passing a custom StepDataCallback to the wrapper. When creating the HDF5 file the list values in the episode dictionary will be stored as datasets and the nested dictionaries will generate a new HDF5 group.

  • A new episode dictionary buffer is created if the env.step(action) call returns truncated or terminated, or if the environment calls env.reset(). If calling reset and the previous episode was not truncated or terminated, this will automatically be truncated.

  • To perform caching the user can set the max_buffer_steps or max_buffer_episodes before saving the in-memory buffers to a temporary HDF5 file in disk. If non of max_buffer_steps or max_buffer_episodes are set, the data will move from in-memory to a permanent location only when the Minari dataset is created. To move all the stored data to a permanent location use DataCollectorV0.save_to_disK(path_to_permanent_location).

Initialize the data colletor attributes and create the temporary directory for caching.

Parameters:

  • env (gym.Env) – Gymnasium environment

  • step_data_callback (type[StepDataCallback], optional) – Callback class to edit/update step databefore storing to buffer. Defaults to StepDataCallback.

  • episode_metadata_callback (type[EpisodeMetadataCallback], optional) – Callback class to add custom metadata to episode group in HDF5 file. Defaults to EpisodeMetadataCallback.

  • record_infos (bool, optional) – If True record the info return key of each step. Defaults to False.

  • max_buffer_steps (Optional[int], optional) – number of steps saved in-memory buffers before dumping to HDF5 file in disk. Defaults to None.

  • max_buffer_episodes (Optional[int], optional) – number of episodes saved in-memory buffers before dumping to HDF5 file in disk. Defaults to None.

Raises:

ValueError – max_buffer_steps and max_buffer_episodes can’t be passed at the same time

Methods#

minari.DataCollectorV0.step(self, action: ActType) → tuple[ObsType, SupportsFloat, bool, bool, dict[str, Any]]#

Gymnasium step method.

minari.DataCollectorV0.reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None) → tuple[ObsType, dict[str, Any]]#

Gymnasium environment reset.

minari.DataCollectorV0.close(self)#

Close the Gymnasium environment.

Clear buffer and close temporary directory.

minari.DataCollectorV0.save_to_disk(self, path: str, dataset_metadata: Dict = {})#

Save all in-memory buffer data and move temporary HDF5 file to a permanent location in disk.

Parameters:

  • path (str) – path to store permanent HDF5, i.e: ‘/home/foo/datasets/data.hdf5’

  • dataset_metadata (Dict, optional) – additional metadata to add to HDF5 dataset file as attributes. Defaults to {}.

minari.DataCollectorV0.clear_buffer_to_tmp_file(self, truncate_last_episode: bool = False)#

Save the global buffer in-memory to a temporary HDF5 file in disk.

Parameters:

truncate_last_episode (bool, optional) – If True the last episode from the buffer will be truncated before saving to disk. Defaults to False.