embodichain.data_pipeline.datasets

embodichain.data_pipeline.datasets#

Overview#

Datasets and samplers for online streaming training from live simulation.

Functions

OnlineDataset

Infinite IterableDataset backed by a live OnlineDataEngine shared buffer.

Classes:

ChunkSizeSampler

Abstract base class for chunk-size samplers.

GMMChunkSampler

Gaussian Mixture Model chunk-size sampler.

OnlineDataset

Infinite IterableDataset backed by a live OnlineDataEngine shared buffer.

UniformChunkSampler

Discrete-uniform chunk-size sampler over [low, high].
class embodichain.data_pipeline.datasets.ChunkSizeSampler[source]#

Bases: ABC

Abstract base class for chunk-size samplers.

Subclasses implement __call__() to return an integer chunk size on demand. A sampler is called once per OnlineDataset.__iter__() step, so consecutive samples / batches may have different time dimensions.

When used in batch mode the same chunk size is drawn once and applied to every trajectory in the batch so that the resulting TensorDict has a consistent shape [batch_size, chunk_size].

class embodichain.data_pipeline.datasets.GMMChunkSampler[source]#

Bases: ChunkSizeSampler

Gaussian Mixture Model chunk-size sampler.

Selects a mixture component according to weights, samples a value from the corresponding Normal(mean, std) distribution, rounds to the nearest integer, and optionally clamps the result to [low, high].

Parameters:

  • means (List[float]) – Mean of each Gaussian component (number of elements = K).

  • stds (List[float]) – Standard deviation of each component (must be > 0, same length as means).

  • weights (Optional[List[float]]) – Unnormalised mixture weights (same length as means). Defaults to a uniform distribution over all components.

  • low (Optional[int]) – Optional lower bound for clamping the sampled value (inclusive, must be ≥ 1 if provided).

  • high (Optional[int]) – Optional upper bound for clamping the sampled value (inclusive, must be ≥ low if both are provided).

Raises:

ValueError – If means, stds, or weights have mismatched lengths, if any std ≤ 0, or if the bounds are inconsistent.

Example — two-component mixture favouring short and long chunks:

sampler = GMMChunkSampler(
    means=[16.0, 64.0],
    stds=[4.0, 8.0],
    weights=[0.6, 0.4],
    low=8,
    high=96,
)
chunk_size = sampler()  # e.g. 18

Methods:

__init__(means, stds[, weights, low, high])

__init__(means, stds, weights=None, low=None, high=None)[source]#
class embodichain.data_pipeline.datasets.OnlineDataset[source]#

Bases: IterableDataset

Infinite IterableDataset backed by a live OnlineDataEngine shared buffer.

Two sampling modes are supported depending on the batch_size argument:

Item mode (batch_size=None, default)

__iter__ yields one TensorDict of shape [chunk_size] per step. Use with a standard DataLoader(dataset, batch_size=B) so the DataLoader handles collation and worker sharding.

Batch mode (batch_size=N)

__iter__ yields one pre-batched TensorDict of shape [N, chunk_size] per step by calling engine.sample_batch(N, chunk_size) directly. Use with DataLoader(dataset, batch_size=None) to skip DataLoader collation and leverage the engine’s bulk-sampling efficiency.

Dynamic chunk sizes

Pass a ChunkSizeSampler as chunk_size to draw a fresh chunk length on every iteration step. In batch mode the size is sampled once per step and applied uniformly to all trajectories in the batch, ensuring a consistent [batch_size, chunk_size] shape. Two built-in samplers are provided:

Note

__len__ is intentionally absent — IterableDataset does not require it and the stream is infinite.

Note

Multi-worker DataLoader: each worker gets its own iterator; since sampling is independent random draws from shared memory, this is safe.

Parameters:

  • engine (OnlineDataEngine) – A started OnlineDataEngine whose shared buffer is used for sampling.

  • chunk_size (int | ChunkSizeSampler) – Fixed number of consecutive timesteps per chunk (int), or a ChunkSizeSampler that returns a fresh size on every iteration step.

  • batch_size (Optional[int]) – If None, yield single chunks of shape [chunk_size] (item mode). If an int, yield pre-batched TensorDicts of shape [batch_size, chunk_size] (batch mode).

  • transform (Optional[Callable[[TensorDict], TensorDict]]) – Optional (TensorDict) -> TensorDict applied to each yielded item/batch before returning.

Example — fixed chunk size, item mode:

dataset = OnlineDataset(engine, chunk_size=64)
loader  = DataLoader(dataset, batch_size=32, num_workers=4,
                     collate_fn=OnlineDataset.collate_fn)
for batch in loader:
    # batch has shape [32, 64, ...]
    train_step(batch)

Example — fixed chunk size, batch mode:

dataset = OnlineDataset(engine, chunk_size=64, batch_size=32)
loader  = DataLoader(dataset, batch_size=None,
                     collate_fn=OnlineDataset.passthrough_collate_fn)
for batch in loader:
    # batch has shape [32, 64, ...]
    train_step(batch)

Example — dynamic chunk size with uniform sampler:

sampler = UniformChunkSampler(low=16, high=64)
dataset = OnlineDataset(engine, chunk_size=sampler)
loader  = DataLoader(dataset, batch_size=32)
for batch in loader:
    # chunk dimension varies each batch
    train_step(batch)

Example — dynamic chunk size with GMM sampler:

sampler = GMMChunkSampler(
    means=[16.0, 64.0], stds=[4.0, 8.0], weights=[0.6, 0.4],
    low=8, high=96,
)
dataset = OnlineDataset(engine, chunk_size=sampler, batch_size=32)
loader  = DataLoader(dataset, batch_size=None)
for batch in loader:
    train_step(batch)

Methods:

__init__(engine, chunk_size[, batch_size, ...])

collate_fn(batch)

Collate a list of TensorDicts into a single batched TensorDict.

passthrough_collate_fn(batch)

Collate function for batch-mode DataLoaders.
__init__(engine, chunk_size, batch_size=None, transform=None)[source]#
static collate_fn(batch)[source]#

Collate a list of TensorDicts into a single batched TensorDict.

Pass this as collate_fn to DataLoader when using item mode (batch_size not None on the DataLoader side) to avoid the default collation failure with TensorDict objects.

Parameters:

batch (List[TensorDict]) – List of TensorDicts, each of shape [chunk_size, ...].

Return type:

TensorDict

Returns:

Stacked TensorDict of shape [len(batch), chunk_size, ...].

static passthrough_collate_fn(batch)[source]#

Collate function for batch-mode DataLoaders.

When the dataset is in batch mode it already yields pre-batched TensorDicts. With batch_size=None, PyTorch’s DataLoader skips auto-batching and passes each item directly to collate_fn as-is (not wrapped in a list). This function returns the TensorDict unchanged.

Pass this as collate_fn to DataLoader when using batch mode (batch_size=None on the DataLoader side) to avoid the default collation failure with TensorDict objects.

Parameters:

batch (TensorDict) – A pre-batched TensorDict of shape [batch_size, chunk_size, ...] passed directly by the DataLoader.

Return type:

TensorDict

Returns:

The pre-batched TensorDict unchanged.

class embodichain.data_pipeline.datasets.UniformChunkSampler[source]#

Bases: ChunkSizeSampler

Discrete-uniform chunk-size sampler over [low, high].

Draws an integer uniformly at random from the closed interval [low, high] on every call.

Parameters:

  • low (int) – Minimum chunk size (inclusive, must be ≥ 1).

  • high (int) – Maximum chunk size (inclusive, must be ≥ low).

Raises:

ValueError – If low < 1 or high < low.

Example:

sampler = UniformChunkSampler(low=16, high=64)
chunk_size = sampler()  # e.g. 37

Methods:

__init__(low, high)

__init__(low, high)[source]#

Classes:

OnlineDataset

Infinite IterableDataset backed by a live OnlineDataEngine shared buffer.
class embodichain.data_pipeline.datasets.online_data.OnlineDataset[source]#

Bases: IterableDataset

Infinite IterableDataset backed by a live OnlineDataEngine shared buffer.

Two sampling modes are supported depending on the batch_size argument:

Item mode (batch_size=None, default)

__iter__ yields one TensorDict of shape [chunk_size] per step. Use with a standard DataLoader(dataset, batch_size=B) so the DataLoader handles collation and worker sharding.

Batch mode (batch_size=N)

__iter__ yields one pre-batched TensorDict of shape [N, chunk_size] per step by calling engine.sample_batch(N, chunk_size) directly. Use with DataLoader(dataset, batch_size=None) to skip DataLoader collation and leverage the engine’s bulk-sampling efficiency.

Dynamic chunk sizes

Pass a ChunkSizeSampler as chunk_size to draw a fresh chunk length on every iteration step. In batch mode the size is sampled once per step and applied uniformly to all trajectories in the batch, ensuring a consistent [batch_size, chunk_size] shape. Two built-in samplers are provided:

  • UniformChunkSampler — uniform discrete distribution over [low, high].

  • GMMChunkSampler — Gaussian Mixture Model, useful for multi-modal chunk-length curricula.

Note

__len__ is intentionally absent — IterableDataset does not require it and the stream is infinite.

Note

Multi-worker DataLoader: each worker gets its own iterator; since sampling is independent random draws from shared memory, this is safe.

Parameters:

  • engine (OnlineDataEngine) – A started OnlineDataEngine whose shared buffer is used for sampling.

  • chunk_size (int | ChunkSizeSampler) – Fixed number of consecutive timesteps per chunk (int), or a ChunkSizeSampler that returns a fresh size on every iteration step.

  • batch_size (Optional[int]) – If None, yield single chunks of shape [chunk_size] (item mode). If an int, yield pre-batched TensorDicts of shape [batch_size, chunk_size] (batch mode).

  • transform (Optional[Callable[[TensorDict], TensorDict]]) – Optional (TensorDict) -> TensorDict applied to each yielded item/batch before returning.

Example — fixed chunk size, item mode:

dataset = OnlineDataset(engine, chunk_size=64)
loader  = DataLoader(dataset, batch_size=32, num_workers=4,
                     collate_fn=OnlineDataset.collate_fn)
for batch in loader:
    # batch has shape [32, 64, ...]
    train_step(batch)

Example — fixed chunk size, batch mode:

dataset = OnlineDataset(engine, chunk_size=64, batch_size=32)
loader  = DataLoader(dataset, batch_size=None,
                     collate_fn=OnlineDataset.passthrough_collate_fn)
for batch in loader:
    # batch has shape [32, 64, ...]
    train_step(batch)

Example — dynamic chunk size with uniform sampler:

sampler = UniformChunkSampler(low=16, high=64)
dataset = OnlineDataset(engine, chunk_size=sampler)
loader  = DataLoader(dataset, batch_size=32)
for batch in loader:
    # chunk dimension varies each batch
    train_step(batch)

Example — dynamic chunk size with GMM sampler:

sampler = GMMChunkSampler(
    means=[16.0, 64.0], stds=[4.0, 8.0], weights=[0.6, 0.4],
    low=8, high=96,
)
dataset = OnlineDataset(engine, chunk_size=sampler, batch_size=32)
loader  = DataLoader(dataset, batch_size=None)
for batch in loader:
    train_step(batch)

Methods:

__init__(engine, chunk_size[, batch_size, ...])

collate_fn(batch)

Collate a list of TensorDicts into a single batched TensorDict.

passthrough_collate_fn(batch)

Collate function for batch-mode DataLoaders.
__init__(engine, chunk_size, batch_size=None, transform=None)[source]#
static collate_fn(batch)[source]#

Collate a list of TensorDicts into a single batched TensorDict.

Pass this as collate_fn to DataLoader when using item mode (batch_size not None on the DataLoader side) to avoid the default collation failure with TensorDict objects.

Parameters:

batch (List[TensorDict]) – List of TensorDicts, each of shape [chunk_size, ...].

Return type:

TensorDict

Returns:

Stacked TensorDict of shape [len(batch), chunk_size, ...].

static passthrough_collate_fn(batch)[source]#

Collate function for batch-mode DataLoaders.

When the dataset is in batch mode it already yields pre-batched TensorDicts. With batch_size=None, PyTorch’s DataLoader skips auto-batching and passes each item directly to collate_fn as-is (not wrapped in a list). This function returns the TensorDict unchanged.

Pass this as collate_fn to DataLoader when using batch mode (batch_size=None on the DataLoader side) to avoid the default collation failure with TensorDict objects.

Parameters:

batch (TensorDict) – A pre-batched TensorDict of shape [batch_size, chunk_size, ...] passed directly by the DataLoader.

Return type:

TensorDict

Returns:

The pre-batched TensorDict unchanged.

Classes:

ChunkSizeSampler

Abstract base class for chunk-size samplers.

GMMChunkSampler

Gaussian Mixture Model chunk-size sampler.

UniformChunkSampler

Discrete-uniform chunk-size sampler over [low, high].
class embodichain.data_pipeline.datasets.sampler.ChunkSizeSampler[source]#

Bases: ABC

Abstract base class for chunk-size samplers.

Subclasses implement __call__() to return an integer chunk size on demand. A sampler is called once per OnlineDataset.__iter__() step, so consecutive samples / batches may have different time dimensions.

When used in batch mode the same chunk size is drawn once and applied to every trajectory in the batch so that the resulting TensorDict has a consistent shape [batch_size, chunk_size].

class embodichain.data_pipeline.datasets.sampler.GMMChunkSampler[source]#

Bases: ChunkSizeSampler

Gaussian Mixture Model chunk-size sampler.

Selects a mixture component according to weights, samples a value from the corresponding Normal(mean, std) distribution, rounds to the nearest integer, and optionally clamps the result to [low, high].

Parameters:

  • means (List[float]) – Mean of each Gaussian component (number of elements = K).

  • stds (List[float]) – Standard deviation of each component (must be > 0, same length as means).

  • weights (Optional[List[float]]) – Unnormalised mixture weights (same length as means). Defaults to a uniform distribution over all components.

  • low (Optional[int]) – Optional lower bound for clamping the sampled value (inclusive, must be ≥ 1 if provided).

  • high (Optional[int]) – Optional upper bound for clamping the sampled value (inclusive, must be ≥ low if both are provided).

Raises:

ValueError – If means, stds, or weights have mismatched lengths, if any std ≤ 0, or if the bounds are inconsistent.

Example — two-component mixture favouring short and long chunks:

sampler = GMMChunkSampler(
    means=[16.0, 64.0],
    stds=[4.0, 8.0],
    weights=[0.6, 0.4],
    low=8,
    high=96,
)
chunk_size = sampler()  # e.g. 18

Methods:

__init__(means, stds[, weights, low, high])

__init__(means, stds, weights=None, low=None, high=None)[source]#
class embodichain.data_pipeline.datasets.sampler.UniformChunkSampler[source]#

Bases: ChunkSizeSampler

Discrete-uniform chunk-size sampler over [low, high].

Draws an integer uniformly at random from the closed interval [low, high] on every call.

Parameters:

  • low (int) – Minimum chunk size (inclusive, must be ≥ 1).

  • high (int) – Maximum chunk size (inclusive, must be ≥ low).

Raises:

ValueError – If low < 1 or high < low.

Example:

sampler = UniformChunkSampler(low=16, high=64)
chunk_size = sampler()  # e.g. 37

Methods:

__init__(low, high)

__init__(low, high)[source]#