protomotions.agents.utils.data module

Data utilities for experience management and batching.

This module provides utilities for managing experience buffers and creating minibatch datasets for training. Handles efficient storage and retrieval of rollout data collected during environment interaction.

Key Classes:

• ExperienceBuffer: Buffer for storing rollout experience

• DictDataset: Dataset for creating minibatches from experience

Key Functions:

• swap_and_flatten01: Reshape tensors for batching

• get_dict: Extract dictionary view of experience buffer

protomotions.agents.utils.data.swap_and_flatten01(arr)

Swap and flatten first two dimensions of a tensor.

Converts (num_steps, num_envs, …) to (num_steps * num_envs, …). Commonly used to batch experience from parallel environments.

Parameters:

arr (MockTensor) – Tensor with at least 2 dimensions.

Returns:

Tensor with first two dimensions flattened.

class protomotions.agents.utils.data.ExperienceBuffer(num_envs, num_steps, device)

Bases: <Mock object at 0x7fd69425e1d0>[]

Buffer for storing rollout experience from parallel environments.

Collects observations, actions, rewards, and other data during environment rollouts. Provides efficient storage and batching for on-policy algorithms. Uses PyTorch buffers for automatic device management.

Parameters:

• num_envs (int) – Number of parallel environments.

• num_steps (int) – Number of steps per rollout.

store_dict

Dictionary tracking which keys have been populated.

Example

>>> buffer = ExperienceBuffer(num_envs=1024, num_steps=16)
>>> buffer.register_key("obs", shape=(128,))
>>> buffer.update_data("obs", step=0, data=observations)
>>> data_dict = buffer.get_dict()

__init__(num_envs, num_steps, device)

register_key(key, shape=(), dtype=<Mock object>)

update_data(key, index, data)

total_sum()

batch_update_data(key, data)

make_dict()

class protomotions.agents.utils.data.DictDataset(batch_size, tensor_dict, shuffle=False)

Bases: <Mock object at 0x7fd6943d14d0>[]

PyTorch Dataset for dictionary of tensors with minibatching.

Creates minibatches from a dictionary of tensors. Supports shuffling and automatic batching for training. Used to create minibatch iterators from collected experience buffers.

Tensors may have different lengths along dim 0. The longest tensor defines the number of minibatches; shorter tensors are indexed with modulo wrapping. This allows e.g. expert discriminator observations to be computed for fewer samples than the full rollout without padding or duplication at creation time.

Parameters:

• batch_size (int) – Size of each minibatch.

• tensor_dict (Dict[str, MockTensor]) – Dictionary of tensors to batch. All lengths must be <= the maximum and the maximum must be divisible by batch_size.

• shuffle – Whether to shuffle indices before batching.

Example

>>> data = {"obs": obs_tensor, "actions": action_tensor}
>>> dataset = DictDataset(batch_size=256, tensor_dict=data, shuffle=True)
>>> for batch in dataset:
>>>     train_on_batch(batch)

__init__(batch_size, tensor_dict, shuffle=False)

shuffle()

num_batches()