opentau.datasets.sampler

Episode-aware sampler for PyTorch DataLoader.

This module provides a sampler that respects episode boundaries in robot learning datasets. It allows filtering specific episodes, dropping frames from episode boundaries, and optional shuffling while maintaining episode structure awareness.

The sampler is designed for use with PyTorch DataLoader to ensure proper sampling behavior when working with sequential episode data, where episode boundaries are important for maintaining temporal coherence or avoiding invalid transitions.

Key Features:

  • Episode filtering: Select specific episodes to include in sampling.

  • Boundary frame dropping: Optionally drop frames from the start or end of episodes (useful for avoiding invalid transitions or edge cases).

  • Optional shuffling: Shuffle indices while maintaining episode awareness.

  • PyTorch compatible: Implements the standard Sampler interface for use with DataLoader.

Classes:
EpisodeAwareSampler: PyTorch-style sampler that respects episode

boundaries, supports episode filtering, frame dropping, and shuffling.

Example

Create a sampler for specific episodes:
>>> episode_data_index = {"from": [0, 100, 200], "to": [99, 199, 299]}
>>> sampler = EpisodeAwareSampler(
...     episode_data_index,
...     episode_indices_to_use=[0, 2],  # Use episodes 0 and 2
...     drop_n_first_frames=5,
...     drop_n_last_frames=5,
...     shuffle=True
... )
>>> dataloader = DataLoader(dataset, sampler=sampler)

Classes

EpisodeAwareSampler(episode_data_index[, ...])

class opentau.datasets.sampler.EpisodeAwareSampler(episode_data_index: dict, episode_indices_to_use: list | None = None, drop_n_first_frames: int = 0, drop_n_last_frames: int = 0, shuffle: bool = False)[source]

Bases: object

__init__(episode_data_index: dict, episode_indices_to_use: list | None = None, drop_n_first_frames: int = 0, drop_n_last_frames: int = 0, shuffle: bool = False)[source]

Sampler that optionally incorporates episode boundary information.

Parameters:

  • episode_data_index – Dictionary with keys ‘from’ and ‘to’ containing the start and end indices of each episode.

  • episode_indices_to_use – List of episode indices to use. If None, all episodes are used. Assumes that episodes are indexed from 0 to N-1.

  • drop_n_first_frames – Number of frames to drop from the start of each episode.

  • drop_n_last_frames – Number of frames to drop from the end of each episode.

  • shuffle – Whether to shuffle the indices.