opentau.datasets.online_buffer

An online buffer for the online training loop in train.py

Note to maintainers: This duplicates some logic from LeRobotDataset and EpisodeAwareSampler. We should consider converging to one approach. Here we have opted to use numpy.memmap to back the data buffer. It’s much faster than using HuggingFace Datasets as there’s no conversion to an intermediate non-python object. Also it supports in-place slicing and mutation which is very handy for a dynamic buffer.

Functions

compute_sampler_weights(offline_dataset[, ...])

Compute the sampling weights for the online training dataloader in train.py.

Classes

OnlineBuffer(write_dir, data_spec, ...[, ...])

FIFO data buffer for the online training loop in train.py.

class opentau.datasets.online_buffer.OnlineBuffer(write_dir: str | Path, data_spec: dict[str, Any] | None, buffer_capacity: int | None, fps: float | None = None, delta_timestamps: dict[str, list[float]] | dict[str, ndarray] | None = None)

Bases: Dataset

FIFO data buffer for the online training loop in train.py.

Follows the protocol of LeRobotDataset as much as is required to have it be used by the online training loop in the same way that a LeRobotDataset would be used.

The underlying data structure will have data inserted in a circular fashion. Always insert after the last index, and when you reach the end, wrap around to the start.

The data is stored in a numpy memmap.

EPISODE_INDEX_KEY = 'episode_index'

FRAME_INDEX_KEY = 'frame_index'

INDEX_KEY = 'index'

IS_PAD_POSTFIX = '_is_pad'

NEXT_INDEX_KEY = '_next_index'

OCCUPANCY_MASK_KEY = '_occupancy_mask'

TIMESTAMP_KEY = 'timestamp'

__init__(write_dir: str | Path, data_spec: dict[str, Any] | None, buffer_capacity: int | None, fps: float | None = None, delta_timestamps: dict[str, list[float]] | dict[str, ndarray] | None = None)

The online buffer can be provided from scratch or you can load an existing online buffer by passing a write_dir associated with an existing buffer.

Parameters:

• write_dir – Where to keep the numpy memmap files. One memmap file will be stored for each data key. Note that if the files already exist, they are opened in read-write mode (used for training resumption.)

• data_spec – A mapping from data key to data specification, like {data_key: {“shape”: tuple[int], “dtype”: np.dtype}}. This should include all the data that you wish to record into the buffer, but note that “index”, “frame_index” and “episode_index” are already accounted for by this class, so you don’t need to include them.

• buffer_capacity – How many frames should be stored in the buffer as a maximum. Be aware of your system’s available disk space when choosing this.

• fps – Same as the fps concept in LeRobot dataset. Here it needs to be provided for the delta_timestamps logic. You can pass None if you are not using delta_timestamps.

• delta_timestamps – Same as the delta_timestamps concept in LeRobotDataset. This is internally converted to dict[str, np.ndarray] for optimization purposes.

add_data(data: dict[str, ndarray]) → None

Add new data to the buffer, potentially overwriting old data in a circular fashion.

The new data should contain all frames (in order) of any number of episodes. Indices should start from 0. This method shifts incoming data indices and episode indices to continue from the last frame in the buffer.

Parameters:

data – Dictionary mapping data keys to numpy arrays. All arrays must have the same length. Must include all keys from data_keys.

Raises:

ValueError – If data is missing required keys or arrays have different lengths.

Note

This modifies the input data in place by shifting indices. See rollout and eval_policy functions in eval.py for usage examples.

property data_keys: list[str]

property delta_timestamps: dict[str, ndarray] | None

property fps: float | None

get_data_by_key(key: str) → Tensor

Get all occupied data for a given key as a torch tensor.

Parameters:

key – Data key to retrieve.

Returns:

Tensor containing all non-padded data for the specified key.

property num_episodes: int

property num_frames: int

set_delta_timestamps(value: dict[str, list[float]] | None) → None

Set delta_timestamps converting the values to numpy arrays.

The conversion is for an optimization in the __getitem__. The loop is much slower if the arrays need to be converted into numpy arrays on each access.

Parameters:

value – Dictionary mapping feature names to lists of delta timestamps, or None to disable delta timestamps.

opentau.datasets.online_buffer.compute_sampler_weights(offline_dataset: LeRobotDataset, offline_drop_n_last_frames: int = 0, online_dataset: OnlineBuffer | None = None, online_sampling_ratio: float | None = None, online_drop_n_last_frames: int = 0) → Tensor

Compute the sampling weights for the online training dataloader in train.py.

Parameters:

• offline_dataset – The LeRobotDataset used for offline pre-training.

• online_drop_n_last_frames – Number of frames to drop from the end of each offline dataset episode.

• online_dataset – The OnlineBuffer used in online training.

• online_sampling_ratio – The proportion of data that should be sampled from the online dataset. If an online dataset is provided, this value must also be provided.

• online_drop_n_first_frames – See offline_drop_n_last_frames. This is the same, but for the online dataset.

Returns:

Tensor of weights for [offline_dataset; online_dataset], normalized to 1.

Notes to maintainers:

• This duplicates some logic from EpisodeAwareSampler. We should consider converging to one approach.

• When used with torch.utils.data.WeightedRandomSampler, it could completely replace EpisodeAwareSampler as the online dataset related arguments are optional. The only missing feature is the ability to turn shuffling off.

• Options drop_first_n_frames and episode_indices_to_use can be added easily. They were not included here to avoid adding complexity.