imitation.testing.expert_trajectories#

Test utilities to conveniently generate expert trajectories.

Functions

generate_expert_trajectories(env_id, ...)

Generate expert trajectories for the given environment.

lazy_generate_expert_trajectories(...)

Generate or load expert trajectories from cache.

make_expert_transition_loader(cache_dir, ...)

Creates different kinds of PyTorch data loaders for expert transitions.
imitation.testing.expert_trajectories.generate_expert_trajectories(env_id, num_trajectories, rng)[source]#

Generate expert trajectories for the given environment.

Note: will just pull a pretrained policy from the Hugging Face model hub.

Parameters

  • env_id (str) – The environment to generate trajectories for.

  • num_trajectories (int) – The number of trajectories to generate.

  • rng (Generator) – The random number generator to use.

Return type

Sequence[TrajectoryWithRew]

Returns

A list of trajectories with rewards.

imitation.testing.expert_trajectories.lazy_generate_expert_trajectories(cache_path, env_id, num_trajectories, rng)[source]#

Generate or load expert trajectories from cache.

Parameters

  • cache_path (PathLike) – A path to the folder to be used as cache for the expert trajectories.

  • env_id (str) – The environment to generate trajectories for.

  • num_trajectories (int) – The number of trajectories to generate.

  • rng (Generator) – The random number generator to use.

Return type

Sequence[TrajectoryWithRew]

Returns

A list of trajectories with rewards.

imitation.testing.expert_trajectories.make_expert_transition_loader(cache_dir, batch_size, expert_data_type, env_name, rng, num_trajectories=1, shuffle=True)[source]#

Creates different kinds of PyTorch data loaders for expert transitions.

Parameters

  • cache_dir (Path) – The directory to use for caching the expert trajectories.

  • batch_size (int) – The batch size to use for the data loader.

  • expert_data_type (str) – The type of expert data to use. Can be one of “trajectories”, “data_loader”, “ducktyped_data_loader”, “transitions”.

  • env_name (str) – The environment to generate trajectories for.

  • rng (Generator) – The random number generator to use.

  • num_trajectories (int) – The number of trajectories to generate.

  • shuffle (bool) – Whether to shuffle the dataset when creating a data loader.

Raises

ValueError – If expert_data_type is not one of the supported types.

Returns

A pytorch data loader for expert transitions.