warp_drive.training.utils package

Submodules

warp_drive.training.utils.data_loader module

warp_drive.training.utils.data_loader.all_equal(iterable)

Check all elements of an iterable (e.g., list) are identical

warp_drive.training.utils.data_loader.create_and_push_data_placeholders(env_wrapper=None, policy_tag_to_agent_id_map=None, create_separate_placeholders_for_each_policy=False, obs_dim_corresponding_to_num_agents='first', training_batch_size_per_env=1, push_data_batch_placeholders=True)

Create observations, sampled_actions, rewards and done flags placeholders and push to the device; this is required for generating environment roll-outs as well as training. env_wrapper: the wrapped environment object. policy_tag_to_agent_id_map:

a dictionary mapping policy tag to agent ids.

create_separate_placeholders_for_each_policy:

flag indicating whether there exist separate observations, actions and rewards placeholders, for each policy, as designed in the step function. The placeholders will be used in the step() function and during training. When there’s only a single policy, this flag will be False. It can also be True when there are multiple policies, yet all the agents have the same obs and action space shapes, so we can share the same placeholder. Defaults to False.

obs_dim_corresponding_to_num_agents:

indicative of which dimension in the observation corresponds to the number of agents, as designed in the step function. It may be “first” or “last”. In other words, observations may be shaped (num_agents, *feature_dim) or (*feature_dim, num_agents). This is required in order for WarpDrive to process the observations correctly. This is only relevant when a single obs key corresponds to multiple agents. Defaults to “first”.

training_batch_size_per_env: the training batch size for each env. push_data_batch_placeholders: an optional flag to push placeholders

for the batches of actions, rewards and the done flags. Defaults to True.

warp_drive.training.utils.data_loader.get_obs(obs, agent_ids, obs_dim_corresponding_to_num_agents='first', key=None)

Module contents