imitation.data.huggingface_utils#

Helpers to convert between Trajectories and HuggingFace’s datasets library.

Functions

trajectories_to_dataset(trajectories[, info])

Convert a sequence of trajectories to a HuggingFace dataset.

trajectories_to_dict(trajectories)

Convert a sequence of trajectories to a dict.

Classes

TrajectoryDatasetSequence(dataset)

A wrapper to present an HF dataset as a sequence of trajectories.
class imitation.data.huggingface_utils.TrajectoryDatasetSequence(dataset)[source]#

Bases: Sequence[Trajectory]

A wrapper to present an HF dataset as a sequence of trajectories.

Converts the dataset to a sequence of trajectories on the fly.

__init__(dataset)[source]#

Construct a TrajectoryDatasetSequence.

property dataset#

Return the underlying HF dataset.

imitation.data.huggingface_utils.trajectories_to_dataset(trajectories, info=None)[source]#

Convert a sequence of trajectories to a HuggingFace dataset.

Return type

Dataset

imitation.data.huggingface_utils.trajectories_to_dict(trajectories)[source]#

Convert a sequence of trajectories to a dict.

The dict has the following fields:

  • obs: The observations. Shape: (num_trajectories, num_timesteps, obs_dim).

  • acts: The actions. Shape: (num_trajectories, num_timesteps, act_dim).

  • infos: The infos. Shape: (num_trajectories, num_timesteps) as jsonpickled str.

  • terminal: The terminal flags. Shape: (num_trajectories, num_timesteps, ).

  • rews: The rewards. Shape: (num_trajectories, num_timesteps) if applicable.

This dict can be used to construct a HuggingFace dataset.

Parameters

trajectories (Sequence[Trajectory]) – The trajectories to save.

Raises

ValueError – If not all trajectories have the same type, i.e. some are Trajectory and others are TrajectoryWithRew.

Return type

Dict[str, Sequence[Any]]

Returns

A dict representing the trajectories.