embodichain.data_pipeline.datasets

Contents

embodichain.data_pipeline.datasets#

Online datasets and offline inspection with lazy simulation dependencies.

Overview#

Datasets and samplers for live simulation and bounded resident numeric data. OnlineDataset consumes trajectory data streamed through the engine, while the chunk samplers (UniformChunkSampler, ChunkSizeSampler, GMMChunkSampler) carve that stream into training chunks.

Classes

OnlineDataset

Infinite IterableDataset backed by a live OnlineDataEngine shared buffer.

UniformChunkSampler

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

ChunkSizeSampler

Abstract base class for chunk-size samplers.

GMMChunkSampler

Gaussian Mixture Model chunk-size sampler.

SequenceSpecCfg

Specify causal history and current/future supervision.

BatchSource

A source of owned training batches with explicit data semantics.

OnlineDataSource

Adapt a started ODS engine without taking ownership of its lifecycle.

TensorDataSource

Sample causal sequences from resident, explicitly described CPU data.

LeRobotFieldMappingCfg

Declare SDK columns and optional EmbodiChain provenance explicitly.

LeRobotDataSource

Load numeric episodes once and serve owned resident context batches.

OfflineDataset

alias of SequenceDataset

SequenceDataset

Iterate owned causal sequences from a live or resident data source.

DataBudgetExhaustedError

A complete requested batch exceeds eligible source delivery allowances.

DataUnavailableError

No eligible unlocked window exists under the current source/filter state.

Schema utilities

validate_data_schema(schema)

Validate declarations common to live and resident causal sources.

data_schema_fingerprint(schema)

Hash a validated finite JSON schema without changing feature order.

Online datasets and offline inspection with lazy simulation dependencies.

Classes:

BatchSource

A source of owned training batches with explicit data semantics.

ChunkSizeSampler

Abstract base class for chunk-size samplers.

GMMChunkSampler

Gaussian Mixture Model chunk-size sampler.

LeRobotDataSource

Load numeric episodes once and serve owned resident context batches.

LeRobotFieldMappingCfg

Declare SDK columns and optional EmbodiChain provenance explicitly.

OfflineDataset

alias of SequenceDataset

OnlineDataSource

Adapt a started ODS engine without taking ownership of its lifecycle.

OnlineDataset

Infinite IterableDataset backed by a live OnlineDataEngine shared buffer.

SequenceDataset

Iterate owned causal sequences from a live or resident data source.

SequenceSpecCfg

Specify causal history and current/future supervision.

TensorDataSource

Sample causal sequences from resident, explicitly described CPU data.

UniformChunkSampler

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

Exceptions:

DataBudgetExhaustedError

A complete requested batch exceeds eligible source delivery allowances.

DataUnavailableError

No eligible unlocked window exists under the current source/filter state.

Functions:

data_schema_fingerprint(schema)

Hash a validated finite JSON schema without changing feature order.

validate_data_schema(schema)

Validate declarations common to live and resident causal sources.
class embodichain.data_pipeline.datasets.BatchSource[source]#

Bases: Protocol

A source of owned training batches with explicit data semantics.

schema declares ordered observations/actions, their timing and units. source_id identifies the source dataset/run independently of cache slots. Production/refill and loading lifecycles belong to concrete sources and the caller, not to this sampling protocol. Sampling never performs hidden waits.

Methods:

__init__(*args, **kwargs)

sample_context(batch_size, spec, *[, generator])

Return copied obs/actions/valid/loss_mask and source metadata.

Attributes:

schema

source_id

__init__(*args, **kwargs)#
sample_context(batch_size, spec, *, generator=None)[source]#

Return copied obs/actions/valid/loss_mask and source metadata.

Parameters:

  • batch_size (int) – Number of episode-uniform contexts, with replacement.

  • spec (SequenceSpecCfg) – Window, boundary and optional generation requirements.

  • generator (Generator | None) – Sampling RNG; None uses the current Torch RNG.

Return type:

TensorDict

Returns:

TensorDict with batch axes [B, H + F - 1]. Returned storage remains valid after subsequent sampling or source replacement.

Raises:

  • DataBudgetExhaustedError – Eligible delivery credits cannot provide the entire batch. No partial debit is performed.

  • DataUnavailableError – No eligible anchor exists in the current data.

  • RuntimeError – A producer or source read failed.

  • ValueError – Schema or requested capabilities are incompatible.

schema: Mapping[str, Any]#
source_id: str#
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].

exception embodichain.data_pipeline.datasets.DataBudgetExhaustedError[source]#

Bases: RuntimeError

A complete requested batch exceeds eligible source delivery allowances.

exception embodichain.data_pipeline.datasets.DataUnavailableError[source]#

Bases: RuntimeError

No eligible unlocked window exists under the current source/filter state.

This is distinct from quota depletion and a terminal producer/read failure. The caller decides whether new publication can help and whether to wait with a deadline; a fixed offline source may never satisfy the request. Legacy consumers catching RuntimeError remain compatible.

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.LeRobotDataSource[source]#

Bases: TensorDataSource

Load numeric episodes once and serve owned resident context batches.

Parameters:

  • root (Path | str) – Existing local LeRobot dataset directory.

  • schema (Mapping[str, Any]) – Explicit ordered observation/action schema. Actions require a nonempty semantics or representation and positive step_dt.

  • memory_limit_bytes (int) – Positive hard limit for owned resident tensors. SDK caches and bounded read/validation staging are separate costs.

  • source_id (str) – Persisted dataset/origin identity; used for generated frame IDs when the dataset does not already record provenance.

  • mapping (LeRobotFieldMappingCfg | None) – Explicit numeric SDK columns and optional source metadata.

  • repo_id (str) – Dataset identifier required by the installed SDK.

  • max_episode_steps (int | None) – Optional fixed padded capacity. Inferred from source frame indices otherwise; specify it for exact ODS RNG pairing.

  • episode_frames (Sequence[Sequence[int]] | None) – Optional original valid frame indices, including gaps and empty padded rows. SDK episode order must match nonempty rows.

  • read_batch_size (int) – Maximum numeric SDK rows per read.

  • language (Mapping[str, Mapping[int | str, str]] | None) – Optional source-owned task/subtask ID-to-text dictionaries.

Uses the official SDK lazily and never imports simulation modules. Missing backing-column access selects scalar SDK fallback; integrity failures do not. External datasets need explicit semantic mapping, not a showcase sidecar. Timestamps must match the declared fixed-rate action grid.

Methods:

__init__(root, *, schema, ...[, mapping, ...])

__init__(root, *, schema, memory_limit_bytes, source_id, mapping=None, repo_id='local/dataset', max_episode_steps=None, episode_frames=None, read_batch_size=1024, language=None)[source]#
class embodichain.data_pipeline.datasets.LeRobotFieldMappingCfg[source]#

Bases: object

Declare SDK columns and optional EmbodiChain provenance explicitly.

Parameters:

  • observation_key (str) – Numeric vector column mapped to obs.

  • action_key (str) – Numeric vector column mapped to actions.

  • provenance_key (str | None) – Optional int64 vector [origin, episode, frame, generation].

  • metadata_keys (dict[str, str]) – Optional canonical scalar fields mapped to SDK columns; supports continuity/segment IDs, segment acceptance and language IDs.

Column mapping is a declaration of semantics by the caller, not automatic inference from a matching width. Image/video sequence decoding is not part of this resident numeric backend.

Methods:

__init__([observation_key, action_key, ...])

copy(**kwargs)

Return a new object replacing specified fields with new values.

replace(**kwargs)

Return a new object replacing specified fields with new values.

to_dict()

Convert an object into dictionary recursively.

validate([prefix])

Check the validity of configclass object.

Attributes:

action_key

metadata_keys

observation_key

provenance_key

__init__(observation_key=<factory>, action_key=<factory>, provenance_key=<factory>, metadata_keys=<factory>)#
action_key: str#
copy(**kwargs)#

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2
Parameters:

  • obj (object) – The object to replace.

  • **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

metadata_keys: dict[str, str]#
observation_key: str#
provenance_key: str | None#
replace(**kwargs)#

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2
Parameters:

  • obj (object) – The object to replace.

  • **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

to_dict()#

Convert an object into dictionary recursively.

Note

Ignores all names starting with “__” (i.e. built-in methods).

Parameters:

obj (object) – An instance of a class to convert.

Raises:

ValueError – When input argument is not an object.

Return type:

dict[str, Any]

Returns:

Converted dictionary mapping.

validate(prefix='')#

Check the validity of configclass object.

This function checks if the object is a valid configclass object. A valid configclass object contains no MISSING entries.

Parameters:

  • obj (object) – The object to check.

  • prefix (str) – The prefix to add to the missing fields. Defaults to ‘’.

Return type:

list[str]

Returns:

A list of missing fields.

Raises:

TypeError – When the object is not a valid configuration object.

embodichain.data_pipeline.datasets.OfflineDataset#

alias of SequenceDataset

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

Bases: object

Adapt a started ODS engine without taking ownership of its lifecycle.

Parameters:

  • engine (OnlineDataEngine) – Caller-owned engine. Enable provenance for portable frame IDs.

  • schema (Mapping[str, Any]) – Explicit schema of the returned observations and actions.

  • observation_transform (Callable[[Any], Tensor | TensorDictBase] | None) – Optional row-local feature function. It runs on the owned copy, after the producer lock has been released.

Refill, waits, start and stop remain caller operations. Sampling propagates empty-window, quota and worker errors immediately. The adapter preserves additional continuity/language/segment fields when present.

Methods:

__init__(engine, *, schema[, ...])

resolve_language(batch)

Resolve copied task/subtask IDs through the engine's text registry.

sample_context(batch_size, spec, *[, generator])

Sample an owned canonical batch under the engine's normal checks.
__init__(engine, *, schema, observation_transform=None)[source]#
resolve_language(batch)[source]#

Resolve copied task/subtask IDs through the engine’s text registry.

Parameters:

batch (TensorDict) – Copied batch containing task/subtask index tensors.

Return type:

dict[str, object]

Returns:

Overall and segment texts with the batch’s nested shape.

sample_context(batch_size, spec, *, generator=None)[source]#

Sample an owned canonical batch under the engine’s normal checks.

Parameters:

  • batch_size (int) – Number of contexts to deliver atomically.

  • spec (SequenceSpecCfg) – History/future lengths, segment mode and generation filter.

  • generator (Generator | None) – Optional sampling RNG independent of model randomness.

Return type:

TensorDict

Returns:

Canonical TensorDict with obs/actions/valid/loss_mask and any available frame_ids/generations and segment/language fields. Optional acceptance Unix nanoseconds and publication host-monotonic nanoseconds remain provenance, not inputs to feature inference.

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.

Language-conditioned training uses the numeric task_index and subtask_index fields. Call resolve_language() on the copied item or collated batch to get the overall task and current segment text, then tokenize them in the consumer or a transform. The text dictionaries append new entries without changing existing indices, so previously sampled items retain their language when the engine refills trajectory slots.

New context consumers may use SequenceDataset(OnlineDataSource(...), sequence, ...) directly. This wrapper preserves legacy chunk APIs and delegates its opt-in context iteration to that same Dataset implementation.

Parameters:

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

  • chunk_size (int | ChunkSizeSampler | None) – 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.

  • sampling_mode (Literal['episode', 'segment', 'boundary']) – "episode" allows chunks to cross segment boundaries, "segment" keeps chunks within one segment, and "boundary" samples chunks that cross a boundary.

  • sequence (SequenceSpecCfg | None) – Opt-in causal context specification. Omit chunk_size and configure segment mode through this specification when supplied.

  • schema (Mapping[str, Any] | None) – Explicit observation/action semantics required in context mode.

  • observation_transform (Optional[Callable[[Any], Tensor]]) – Optional row-local observation conversion for context mode; canonical padding is restored after conversion.

  • generator (Generator | None) – Optional context-sampling RNG; explicit generators require owner-process iteration to avoid duplicated worker streams.

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.

resolve_language(batch)

Decode the overall and segment instructions for an item or batch.
__init__(engine, chunk_size=None, batch_size=None, transform=None, sampling_mode='episode', *, sequence=None, schema=None, observation_transform=None, generator=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.

resolve_language(batch)[source]#

Decode the overall and segment instructions for an item or batch.

Parameters:

batch (TensorDict) – Copied numeric sample with task/subtask index fields.

Return type:

dict[str, object]

Returns:

task and subtask strings nested to match batch dimensions. Lookup remains valid after trajectory refills and engine shutdown.

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

Bases: IterableDataset

Iterate owned causal sequences from a live or resident data source.

Parameters:

  • source (BatchSource) – BatchSource, such as OnlineDataSource, LeRobotDataSource or TensorDataSource. The caller owns its loading/production lifecycle.

  • sequence (SequenceSpecCfg) – Causal window and boundary requirements.

  • batch_size (int | None) – None yields one context for DataLoader collation; a positive integer yields complete batches (use DataLoader batch_size=None).

  • transform (Callable[[TensorDict], TensorDict] | None) – Optional batch/item transform preserving axes and masks.

  • generator (Generator | None) – Optional CPU sampling RNG. Explicit generators require num_workers=0 to prevent copied workers repeating the same stream.

  • max_batches (int | None) – Optional positive draws per iterator. Stored cardinality does not imply iteration length: random sampling otherwise repeats. Finite iteration requires num_workers=0.

There is no inferred epoch length or automatic refill. max_batches is a bound on each iterator, not a source quota or a checkpointed update count. Sampling charges delivery before any transform or optimizer update; a failed consumer does not refund it. Prefetched batches also count.

Quota exhaustion and source errors propagate immediately. Budgeted sources reject unsupported worker copies; this wrapper never clones allowances, retries permanent failures or guesses an epoch length.

Methods:

__init__(source, sequence, *[, batch_size, ...])

collate_fn(batch)

Stack contexts for item-mode DataLoader use.

passthrough_collate_fn(batch)

Pass through an already batched item.

resolve_language(batch)

Resolve source-local language IDs when the source supplies dictionaries.
__init__(source, sequence, *, batch_size=None, transform=None, generator=None, max_batches=None)[source]#
static collate_fn(batch)[source]#

Stack contexts for item-mode DataLoader use.

Parameters:

batch (list[TensorDict]) – Individual owned contexts.

Return type:

TensorDict

Returns:

TensorDict with a new leading batch axis.

static passthrough_collate_fn(batch)[source]#

Pass through an already batched item.

Parameters:

batch (TensorDict) – Source-owned copied minibatch.

Return type:

TensorDict

Returns:

The supplied TensorDict unchanged.

resolve_language(batch)[source]#

Resolve source-local language IDs when the source supplies dictionaries.

Parameters:

batch (TensorDict) – Copied numeric batch containing task/subtask indices.

Return type:

dict[str, object]

Returns:

Overall and segment text with the batch’s nested shape.

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

Bases: object

Specify causal history and current/future supervision.

Parameters:

  • history_length (int) – Number of observation frames, including the current one.

  • prediction_horizon (int) – Number of target frames, including the current one.

  • sampling_mode (Literal['episode', 'segment']) – Allow segment crossings within an episode, or restrict context to one accepted segment. Continuity changes always pad.

  • generations (tuple[int, ...] | None) – Optional source-generation filter. Missing generation metadata raises rather than silently ignoring this filter.

The current index is history_length - 1 in a sequence of length history_length + prediction_horizon - 1. Numeric padding is zero, source/language padding is -1, and only real current/future frames enter loss_mask. Boundary-crossing chunk sampling remains an engine chunk API.

Methods:

__init__([history_length, ...])

copy(**kwargs)

Return a new object replacing specified fields with new values.

replace(**kwargs)

Return a new object replacing specified fields with new values.

to_dict()

Convert an object into dictionary recursively.

validate([prefix])

Check the validity of configclass object.

Attributes:

generations

history_length

prediction_horizon

sampling_mode

__init__(history_length=<factory>, prediction_horizon=<factory>, sampling_mode=<factory>, generations=<factory>)#
copy(**kwargs)#

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2
Parameters:

  • obj (object) – The object to replace.

  • **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

generations: tuple[int, ...] | None#
history_length: int#
prediction_horizon: int#
replace(**kwargs)#

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2
Parameters:

  • obj (object) – The object to replace.

  • **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

sampling_mode: Literal['episode', 'segment']#
to_dict()#

Convert an object into dictionary recursively.

Note

Ignores all names starting with “__” (i.e. built-in methods).

Parameters:

obj (object) – An instance of a class to convert.

Raises:

ValueError – When input argument is not an object.

Return type:

dict[str, Any]

Returns:

Converted dictionary mapping.

validate(prefix='')#

Check the validity of configclass object.

This function checks if the object is a valid configclass object. A valid configclass object contains no MISSING entries.

Parameters:

  • obj (object) – The object to check.

  • prefix (str) – The prefix to add to the missing fields. Defaults to ‘’.

Return type:

list[str]

Returns:

A list of missing fields.

Raises:

TypeError – When the object is not a valid configuration object.

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

Bases: object

Sample causal sequences from resident, explicitly described CPU data.

Parameters:

  • data (TensorDictBase) – TensorDict with [episode, time] batch axes, obs, floating actions and boolean valid. Optional continuity, segment, language and provenance fields are preserved.

  • schema (Mapping[str, Any]) – Ordered observation/action semantics. No feature conversion is inferred from vector dimensions.

  • source_id (str) – Stable source dataset/run identity. Persist this for datasets without recorded frame_ids; it seeds their numeric source ID.

  • copy_data (bool) – Clone the supplied data by default. False transfers storage responsibility to the caller, which must not mutate it afterwards.

  • metadata (Mapping[str, Any] | None) – Optional detached storage/provenance diagnostics.

Each returned batch owns its storage. Optional source quotas are owner-only and cannot be reset. data is exposed for inspection and fitting frozen training statistics; treat its tensors as immutable after construction.

Methods:

__init__(data, *, schema, source_id[, ...])

configure_sampling_budget(budget, hard_limit)

Assign a permanent per-source chunk allowance before first sampling.

get_sampling_budget_snapshot()

Return detached delivery counts, using one context as one chunk.

resolve_language(batch)

Resolve IDs using source-owned dictionaries in metadata['language'].

sample_context(batch_size, spec, *[, generator])

Draw complete owned contexts, charging quotas only after copying.
__init__(data, *, schema, source_id, copy_data=True, metadata=None)[source]#
configure_sampling_budget(budget, hard_limit)[source]#

Assign a permanent per-source chunk allowance before first sampling.

The stateful sampler belongs to its creating process and thread. Copying it into DataLoader workers cannot duplicate allowances. Valid rows must each identify one distinct immutable (run, episode) source. Empty padded rows receive no allowance. There is no reset or refill operation.

Parameters:

  • budget (int) – Maximum delivered contexts for each retained source episode.

  • hard_limit (int) – Explicit protection ceiling; budget must not exceed it.

Return type:

None

get_sampling_budget_snapshot()[source]#

Return detached delivery counts, using one context as one chunk.

Return type:

dict[str, Any]

Returns:

Per-episode allowances/deliveries, total remaining credits and continuity-qualified eligible remaining credits. A sequence’s generation or segment filter can restrict eligibility further.

resolve_language(batch)[source]#

Resolve IDs using source-owned dictionaries in metadata[‘language’].

Parameters:

batch (TensorDict) – Copied task/subtask indices; select real frames before resolving since padding and unknown IDs are rejected.

Return type:

dict[str, object]

Returns:

Nested task/subtask texts with the corresponding index shapes.

sample_context(batch_size, spec, *, generator=None)[source]#

Draw complete owned contexts, charging quotas only after copying.

Parameters:

  • batch_size (int) – Positive number of contexts with replacement.

  • spec (SequenceSpecCfg) – Window lengths, boundary behavior and generation filter.

  • generator (Generator | None) – Optional CPU sampling RNG.

Return type:

TensorDict

Returns:

Canonical TensorDict of shape [B, H + F - 1] with validity and supervision masks. Gaps and continuity/segment boundaries pad.

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]#
embodichain.data_pipeline.datasets.data_schema_fingerprint(schema)[source]#

Hash a validated finite JSON schema without changing feature order.

Parameters:

schema (Mapping[str, Any]) – Persistable declaration, including any feature-transform, normalization, coordinate-frame and action-contract versions.

Return type:

str

Returns:

SHA-256 of canonical UTF-8 JSON. Mapping key order is ignored; ordered lists and caller-declared versions remain significant. This is a schema identity, not a hash of trajectory content or callable implementation.

Raises:

ValueError – The schema is invalid or not finite JSON data.

embodichain.data_pipeline.datasets.validate_data_schema(schema)[source]#

Validate declarations common to live and resident causal sources.

Parameters:

schema (Mapping[str, Any]) – Observation/action declarations. Actions require explicit semantics or representation and finite positive step_dt. Optional names are ordered and unique; dtype names use Torch’s spelling without its prefix. Optional observation timing must be pre_action. Root and feature version values identify schemas and feature transformations when supplied.

Return type:

None

This validates declarations, not their truth or compatibility with a live robot. Source adapters also validate stored features and action contracts. Units, frames, rotations and normalization remain caller declarations; no conversion or missing semantic information is inferred from vector widths.

Raises:

ValueError – A required semantic declaration is absent or malformed.

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.

Language-conditioned training uses the numeric task_index and subtask_index fields. Call resolve_language() on the copied item or collated batch to get the overall task and current segment text, then tokenize them in the consumer or a transform. The text dictionaries append new entries without changing existing indices, so previously sampled items retain their language when the engine refills trajectory slots.

New context consumers may use SequenceDataset(OnlineDataSource(...), sequence, ...) directly. This wrapper preserves legacy chunk APIs and delegates its opt-in context iteration to that same Dataset implementation.

Parameters:

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

  • chunk_size (int | ChunkSizeSampler | None) – 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.

  • sampling_mode (Literal['episode', 'segment', 'boundary']) – "episode" allows chunks to cross segment boundaries, "segment" keeps chunks within one segment, and "boundary" samples chunks that cross a boundary.

  • sequence (SequenceSpecCfg | None) – Opt-in causal context specification. Omit chunk_size and configure segment mode through this specification when supplied.

  • schema (Mapping[str, Any] | None) – Explicit observation/action semantics required in context mode.

  • observation_transform (Optional[Callable[[Any], Tensor]]) – Optional row-local observation conversion for context mode; canonical padding is restored after conversion.

  • generator (Generator | None) – Optional context-sampling RNG; explicit generators require owner-process iteration to avoid duplicated worker streams.

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.

resolve_language(batch)

Decode the overall and segment instructions for an item or batch.
__init__(engine, chunk_size=None, batch_size=None, transform=None, sampling_mode='episode', *, sequence=None, schema=None, observation_transform=None, generator=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.

resolve_language(batch)[source]#

Decode the overall and segment instructions for an item or batch.

Parameters:

batch (TensorDict) – Copied numeric sample with task/subtask index fields.

Return type:

dict[str, object]

Returns:

task and subtask strings nested to match batch dimensions. Lookup remains valid after trajectory refills and engine shutdown.

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]#

Reusable causal training batches#

BatchSource implementations return independently owned TensorDicts with obs, actions, valid, loss_mask and available source metadata. SequenceSpecCfg defines history and prediction windows, accepted-segment boundaries and optional generation filtering. Sampling is episode-uniform; invalid or discontinuous spans pad rather than become training transitions.

OnlineDataSource preserves the engine’s lock, worker checks and shared delivery quotas. Its caller continues to own start, refill, waits and stop. TensorDataSource serves explicit resident CPU episodes; LeRobotDataSource loads numeric columns through the optional official SDK with an explicit resident-memory limit and schema/column mapping. SDK caches and bounded staging are separate from that resident tensor limit. Large out-of-core visual loading is outside this backend’s contract.

SequenceDataset supplies common item or prebatched PyTorch iteration for live and resident sources. OfflineDataset is its compatibility alias. Stored cardinality does not imply an epoch length; set max_batches for finite iteration. OnlineDataset retains its existing chunk interface and accepts an opt-in sequence plus explicit schema for the common context contract. It delegates context iteration to SequenceDataset. Quota exhaustion, DataUnavailableError and permanent source failures propagate without hidden retry. Production, bounded waiting and refill remain caller-owned.

All source constructors use validate_data_schema: action semantics (or representation) and a positive control period are required. Optional observation timing must be pre-action. Feature names, dtypes and declared versions are validated consistently. data_schema_fingerprint hashes finite JSON schema metadata, including ordered names and supplied transformation/normalization versions; it does not identify trajectory content or infer missing semantics.

Storage-independent contracts for causal training sequences.

Classes:

BatchSource

A source of owned training batches with explicit data semantics.

SequenceSpecCfg

Specify causal history and current/future supervision.

Exceptions:

DataBudgetExhaustedError

A complete requested batch exceeds eligible source delivery allowances.

DataUnavailableError

No eligible unlocked window exists under the current source/filter state.
class embodichain.data_pipeline.datasets.sequence.BatchSource[source]#

A source of owned training batches with explicit data semantics.

schema declares ordered observations/actions, their timing and units. source_id identifies the source dataset/run independently of cache slots. Production/refill and loading lifecycles belong to concrete sources and the caller, not to this sampling protocol. Sampling never performs hidden waits.

Methods:

__init__(*args, **kwargs)

sample_context(batch_size, spec, *[, generator])

Return copied obs/actions/valid/loss_mask and source metadata.
__init__(*args, **kwargs)#
sample_context(batch_size, spec, *, generator=None)[source]#

Return copied obs/actions/valid/loss_mask and source metadata.

Parameters:

  • batch_size (int) – Number of episode-uniform contexts, with replacement.

  • spec (SequenceSpecCfg) – Window, boundary and optional generation requirements.

  • generator (Generator | None) – Sampling RNG; None uses the current Torch RNG.

Return type:

TensorDict

Returns:

TensorDict with batch axes [B, H + F - 1]. Returned storage remains valid after subsequent sampling or source replacement.

Raises:

  • DataBudgetExhaustedError – Eligible delivery credits cannot provide the entire batch. No partial debit is performed.

  • DataUnavailableError – No eligible anchor exists in the current data.

  • RuntimeError – A producer or source read failed.

  • ValueError – Schema or requested capabilities are incompatible.

exception embodichain.data_pipeline.datasets.sequence.DataBudgetExhaustedError[source]#

A complete requested batch exceeds eligible source delivery allowances.

exception embodichain.data_pipeline.datasets.sequence.DataUnavailableError[source]#

No eligible unlocked window exists under the current source/filter state.

This is distinct from quota depletion and a terminal producer/read failure. The caller decides whether new publication can help and whether to wait with a deadline; a fixed offline source may never satisfy the request. Legacy consumers catching RuntimeError remain compatible.

class embodichain.data_pipeline.datasets.sequence.SequenceSpecCfg[source]#

Specify causal history and current/future supervision.

Parameters:

  • history_length (int) – Number of observation frames, including the current one.

  • prediction_horizon (int) – Number of target frames, including the current one.

  • sampling_mode (Literal['episode', 'segment']) – Allow segment crossings within an episode, or restrict context to one accepted segment. Continuity changes always pad.

  • generations (tuple[int, ...] | None) – Optional source-generation filter. Missing generation metadata raises rather than silently ignoring this filter.

The current index is history_length - 1 in a sequence of length history_length + prediction_horizon - 1. Numeric padding is zero, source/language padding is -1, and only real current/future frames enter loss_mask. Boundary-crossing chunk sampling remains an engine chunk API.

Methods:

__init__([history_length, ...])

copy(**kwargs)

Return a new object replacing specified fields with new values.

replace(**kwargs)

Return a new object replacing specified fields with new values.

to_dict()

Convert an object into dictionary recursively.

validate([prefix])

Check the validity of configclass object.
__init__(history_length=<factory>, prediction_horizon=<factory>, sampling_mode=<factory>, generations=<factory>)#
copy(**kwargs)#

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2
Parameters:

  • obj (object) – The object to replace.

  • **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

replace(**kwargs)#

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2
Parameters:

  • obj (object) – The object to replace.

  • **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

to_dict()#

Convert an object into dictionary recursively.

Note

Ignores all names starting with “__” (i.e. built-in methods).

Parameters:

obj (object) – An instance of a class to convert.

Raises:

ValueError – When input argument is not an object.

Return type:

dict[str, Any]

Returns:

Converted dictionary mapping.

validate(prefix='')#

Check the validity of configclass object.

This function checks if the object is a valid configclass object. A valid configclass object contains no MISSING entries.

Parameters:

  • obj (object) – The object to check.

  • prefix (str) – The prefix to add to the missing fields. Defaults to ‘’.

Return type:

list[str]

Returns:

A list of missing fields.

Raises:

TypeError – When the object is not a valid configuration object.

Live ODS and resident CPU implementations of the training-batch contract.

Classes:

OnlineDataSource

Adapt a started ODS engine without taking ownership of its lifecycle.

TensorDataSource

Sample causal sequences from resident, explicitly described CPU data.
class embodichain.data_pipeline.datasets.sources.OnlineDataSource[source]#

Adapt a started ODS engine without taking ownership of its lifecycle.

Parameters:

  • engine (OnlineDataEngine) – Caller-owned engine. Enable provenance for portable frame IDs.

  • schema (Mapping[str, Any]) – Explicit schema of the returned observations and actions.

  • observation_transform (Callable[[Any], Tensor | TensorDictBase] | None) – Optional row-local feature function. It runs on the owned copy, after the producer lock has been released.

Refill, waits, start and stop remain caller operations. Sampling propagates empty-window, quota and worker errors immediately. The adapter preserves additional continuity/language/segment fields when present.

Methods:

__init__(engine, *, schema[, ...])

resolve_language(batch)

Resolve copied task/subtask IDs through the engine's text registry.

sample_context(batch_size, spec, *[, generator])

Sample an owned canonical batch under the engine's normal checks.
__init__(engine, *, schema, observation_transform=None)[source]#
resolve_language(batch)[source]#

Resolve copied task/subtask IDs through the engine’s text registry.

Parameters:

batch (TensorDict) – Copied batch containing task/subtask index tensors.

Return type:

dict[str, object]

Returns:

Overall and segment texts with the batch’s nested shape.

sample_context(batch_size, spec, *, generator=None)[source]#

Sample an owned canonical batch under the engine’s normal checks.

Parameters:

  • batch_size (int) – Number of contexts to deliver atomically.

  • spec (SequenceSpecCfg) – History/future lengths, segment mode and generation filter.

  • generator (Generator | None) – Optional sampling RNG independent of model randomness.

Return type:

TensorDict

Returns:

Canonical TensorDict with obs/actions/valid/loss_mask and any available frame_ids/generations and segment/language fields. Optional acceptance Unix nanoseconds and publication host-monotonic nanoseconds remain provenance, not inputs to feature inference.

class embodichain.data_pipeline.datasets.sources.TensorDataSource[source]#

Sample causal sequences from resident, explicitly described CPU data.

Parameters:

  • data (TensorDictBase) – TensorDict with [episode, time] batch axes, obs, floating actions and boolean valid. Optional continuity, segment, language and provenance fields are preserved.

  • schema (Mapping[str, Any]) – Ordered observation/action semantics. No feature conversion is inferred from vector dimensions.

  • source_id (str) – Stable source dataset/run identity. Persist this for datasets without recorded frame_ids; it seeds their numeric source ID.

  • copy_data (bool) – Clone the supplied data by default. False transfers storage responsibility to the caller, which must not mutate it afterwards.

  • metadata (Mapping[str, Any] | None) – Optional detached storage/provenance diagnostics.

Each returned batch owns its storage. Optional source quotas are owner-only and cannot be reset. data is exposed for inspection and fitting frozen training statistics; treat its tensors as immutable after construction.

Methods:

__init__(data, *, schema, source_id[, ...])

configure_sampling_budget(budget, hard_limit)

Assign a permanent per-source chunk allowance before first sampling.

get_sampling_budget_snapshot()

Return detached delivery counts, using one context as one chunk.

resolve_language(batch)

Resolve IDs using source-owned dictionaries in metadata['language'].

sample_context(batch_size, spec, *[, generator])

Draw complete owned contexts, charging quotas only after copying.
__init__(data, *, schema, source_id, copy_data=True, metadata=None)[source]#
configure_sampling_budget(budget, hard_limit)[source]#

Assign a permanent per-source chunk allowance before first sampling.

The stateful sampler belongs to its creating process and thread. Copying it into DataLoader workers cannot duplicate allowances. Valid rows must each identify one distinct immutable (run, episode) source. Empty padded rows receive no allowance. There is no reset or refill operation.

Parameters:

  • budget (int) – Maximum delivered contexts for each retained source episode.

  • hard_limit (int) – Explicit protection ceiling; budget must not exceed it.

Return type:

None

get_sampling_budget_snapshot()[source]#

Return detached delivery counts, using one context as one chunk.

Return type:

dict[str, Any]

Returns:

Per-episode allowances/deliveries, total remaining credits and continuity-qualified eligible remaining credits. A sequence’s generation or segment filter can restrict eligibility further.

resolve_language(batch)[source]#

Resolve IDs using source-owned dictionaries in metadata[‘language’].

Parameters:

batch (TensorDict) – Copied task/subtask indices; select real frames before resolving since padding and unknown IDs are rejected.

Return type:

dict[str, object]

Returns:

Nested task/subtask texts with the corresponding index shapes.

sample_context(batch_size, spec, *, generator=None)[source]#

Draw complete owned contexts, charging quotas only after copying.

Parameters:

  • batch_size (int) – Positive number of contexts with replacement.

  • spec (SequenceSpecCfg) – Window lengths, boundary behavior and generation filter.

  • generator (Generator | None) – Optional CPU sampling RNG.

Return type:

TensorDict

Returns:

Canonical TensorDict of shape [B, H + F - 1] with validity and supervision masks. Gaps and continuity/segment boundaries pad.

Explicit, bounded numeric loading through the optional official LeRobot SDK.

Classes:

LeRobotDataSource

Load numeric episodes once and serve owned resident context batches.

LeRobotFieldMappingCfg

Declare SDK columns and optional EmbodiChain provenance explicitly.
class embodichain.data_pipeline.datasets.lerobot.LeRobotDataSource[source]#

Load numeric episodes once and serve owned resident context batches.

Parameters:

  • root (Path | str) – Existing local LeRobot dataset directory.

  • schema (Mapping[str, Any]) – Explicit ordered observation/action schema. Actions require a nonempty semantics or representation and positive step_dt.

  • memory_limit_bytes (int) – Positive hard limit for owned resident tensors. SDK caches and bounded read/validation staging are separate costs.

  • source_id (str) – Persisted dataset/origin identity; used for generated frame IDs when the dataset does not already record provenance.

  • mapping (LeRobotFieldMappingCfg | None) – Explicit numeric SDK columns and optional source metadata.

  • repo_id (str) – Dataset identifier required by the installed SDK.

  • max_episode_steps (int | None) – Optional fixed padded capacity. Inferred from source frame indices otherwise; specify it for exact ODS RNG pairing.

  • episode_frames (Sequence[Sequence[int]] | None) – Optional original valid frame indices, including gaps and empty padded rows. SDK episode order must match nonempty rows.

  • read_batch_size (int) – Maximum numeric SDK rows per read.

  • language (Mapping[str, Mapping[int | str, str]] | None) – Optional source-owned task/subtask ID-to-text dictionaries.

Uses the official SDK lazily and never imports simulation modules. Missing backing-column access selects scalar SDK fallback; integrity failures do not. External datasets need explicit semantic mapping, not a showcase sidecar. Timestamps must match the declared fixed-rate action grid.

Methods:

__init__(root, *, schema, ...[, mapping, ...])

__init__(root, *, schema, memory_limit_bytes, source_id, mapping=None, repo_id='local/dataset', max_episode_steps=None, episode_frames=None, read_batch_size=1024, language=None)[source]#
class embodichain.data_pipeline.datasets.lerobot.LeRobotFieldMappingCfg[source]#

Declare SDK columns and optional EmbodiChain provenance explicitly.

Parameters:

  • observation_key (str) – Numeric vector column mapped to obs.

  • action_key (str) – Numeric vector column mapped to actions.

  • provenance_key (str | None) – Optional int64 vector [origin, episode, frame, generation].

  • metadata_keys (dict[str, str]) – Optional canonical scalar fields mapped to SDK columns; supports continuity/segment IDs, segment acceptance and language IDs.

Column mapping is a declaration of semantics by the caller, not automatic inference from a matching width. Image/video sequence decoding is not part of this resident numeric backend.

Methods:

__init__([observation_key, action_key, ...])

copy(**kwargs)

Return a new object replacing specified fields with new values.

replace(**kwargs)

Return a new object replacing specified fields with new values.

to_dict()

Convert an object into dictionary recursively.

validate([prefix])

Check the validity of configclass object.
__init__(observation_key=<factory>, action_key=<factory>, provenance_key=<factory>, metadata_keys=<factory>)#
copy(**kwargs)#

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2
Parameters:

  • obj (object) – The object to replace.

  • **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

replace(**kwargs)#

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2
Parameters:

  • obj (object) – The object to replace.

  • **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

to_dict()#

Convert an object into dictionary recursively.

Note

Ignores all names starting with “__” (i.e. built-in methods).

Parameters:

obj (object) – An instance of a class to convert.

Raises:

ValueError – When input argument is not an object.

Return type:

dict[str, Any]

Returns:

Converted dictionary mapping.

validate(prefix='')#

Check the validity of configclass object.

This function checks if the object is a valid configclass object. A valid configclass object contains no MISSING entries.

Parameters:

  • obj (object) – The object to check.

  • prefix (str) – The prefix to add to the missing fields. Defaults to ‘’.

Return type:

list[str]

Returns:

A list of missing fields.

Raises:

TypeError – When the object is not a valid configuration object.

Compatibility import for the storage-independent sequence Dataset.

Classes:

OfflineDataset

Compatibility alias; new consumers should use SequenceDataset.
embodichain.data_pipeline.datasets.offline_data.OfflineDataset#

Compatibility alias; new consumers should use SequenceDataset.

PyTorch iteration shared by resident and live causal data sources.

Classes:

SequenceDataset

Iterate owned causal sequences from a live or resident data source.
class embodichain.data_pipeline.datasets.sequence_data.SequenceDataset[source]#

Iterate owned causal sequences from a live or resident data source.

Parameters:

  • source (BatchSource) – BatchSource, such as OnlineDataSource, LeRobotDataSource or TensorDataSource. The caller owns its loading/production lifecycle.

  • sequence (SequenceSpecCfg) – Causal window and boundary requirements.

  • batch_size (int | None) – None yields one context for DataLoader collation; a positive integer yields complete batches (use DataLoader batch_size=None).

  • transform (Callable[[TensorDict], TensorDict] | None) – Optional batch/item transform preserving axes and masks.

  • generator (Generator | None) – Optional CPU sampling RNG. Explicit generators require num_workers=0 to prevent copied workers repeating the same stream.

  • max_batches (int | None) – Optional positive draws per iterator. Stored cardinality does not imply iteration length: random sampling otherwise repeats. Finite iteration requires num_workers=0.

There is no inferred epoch length or automatic refill. max_batches is a bound on each iterator, not a source quota or a checkpointed update count. Sampling charges delivery before any transform or optimizer update; a failed consumer does not refund it. Prefetched batches also count.

Quota exhaustion and source errors propagate immediately. Budgeted sources reject unsupported worker copies; this wrapper never clones allowances, retries permanent failures or guesses an epoch length.

Methods:

__init__(source, sequence, *[, batch_size, ...])

collate_fn(batch)

Stack contexts for item-mode DataLoader use.

passthrough_collate_fn(batch)

Pass through an already batched item.

resolve_language(batch)

Resolve source-local language IDs when the source supplies dictionaries.
__init__(source, sequence, *, batch_size=None, transform=None, generator=None, max_batches=None)[source]#
static collate_fn(batch)[source]#

Stack contexts for item-mode DataLoader use.

Parameters:

batch (list[TensorDict]) – Individual owned contexts.

Return type:

TensorDict

Returns:

TensorDict with a new leading batch axis.

static passthrough_collate_fn(batch)[source]#

Pass through an already batched item.

Parameters:

batch (TensorDict) – Source-owned copied minibatch.

Return type:

TensorDict

Returns:

The supplied TensorDict unchanged.

resolve_language(batch)[source]#

Resolve source-local language IDs when the source supplies dictionaries.

Parameters:

batch (TensorDict) – Copied numeric batch containing task/subtask indices.

Return type:

dict[str, object]

Returns:

Overall and segment text with the batch’s nested shape.

Shared semantic validation and identity for numeric training-data schemas.

Functions:

data_schema_fingerprint(schema)

Hash a validated finite JSON schema without changing feature order.

validate_data_schema(schema)

Validate declarations common to live and resident causal sources.
embodichain.data_pipeline.datasets.schema.data_schema_fingerprint(schema)[source]#

Hash a validated finite JSON schema without changing feature order.

Parameters:

schema (Mapping[str, Any]) – Persistable declaration, including any feature-transform, normalization, coordinate-frame and action-contract versions.

Return type:

str

Returns:

SHA-256 of canonical UTF-8 JSON. Mapping key order is ignored; ordered lists and caller-declared versions remain significant. This is a schema identity, not a hash of trajectory content or callable implementation.

Raises:

ValueError – The schema is invalid or not finite JSON data.

embodichain.data_pipeline.datasets.schema.validate_data_schema(schema)[source]#

Validate declarations common to live and resident causal sources.

Parameters:

schema (Mapping[str, Any]) – Observation/action declarations. Actions require explicit semantics or representation and finite positive step_dt. Optional names are ordered and unique; dtype names use Torch’s spelling without its prefix. Optional observation timing must be pre_action. Root and feature version values identify schemas and feature transformations when supplied.

Return type:

None

This validates declarations, not their truth or compatibility with a live robot. Source adapters also validate stored features and action contracts. Units, frames, rotations and normalization remain caller declarations; no conversion or missing semantic information is inferred from vector widths.

Raises:

ValueError – A required semantic declaration is absent or malformed.

Offline dataset inspection#

Validate saved episode/frame metadata, feature values, segment boundaries, and optional media evidence without constructing a simulator. Quality-aware split manifests group source episodes and their derived fragments so one lineage cannot cross train, validation, and test partitions.

See Validate, split, and recover expert datasets for CLI and Python examples, diagnostic reports, quality evidence and split-manifest consumption.

Inspect persisted LeRobot data and create provenance-safe split manifests.

The inspector reads files directly and never imports LeRobot or simulation modules. Numeric frame data is scanned in bounded Arrow batches; decoding videos is an explicit, optional expense. No operation modifies source data.

Classes:

DatasetValidationIssue

One actionable diagnostic with a file, episode, or frame location.

DatasetValidationReport

Result of a file-only dataset scan, including capped issue details.

Functions:

create_split_manifest(root, output, *[, ...])

Write an atomic split manifest without modifying dataset metadata.

main([argv])

Run the standalone dataset validation/split command.

validate_dataset(root, *[, check_media, ...])

Validate on-disk LeRobot numeric, semantic, and persistence invariants.
class embodichain.data_pipeline.datasets.inspection.DatasetValidationIssue[source]#

One actionable diagnostic with a file, episode, or frame location.

Parameters:

  • severity (Literal['error', 'warning']) – error for corrupt data, warning for limited evidence.

  • code (str) – Stable machine-readable diagnostic category.

  • location (str) – Dataset-relative file and, when known, row or episode.

  • message (str) – Description of the violated invariant.

Methods:

__init__(severity, code, location, message)

__init__(severity, code, location, message)#
class embodichain.data_pipeline.datasets.inspection.DatasetValidationReport[source]#

Result of a file-only dataset scan, including capped issue details.

Parameters:

  • root (str) – Absolute dataset directory.

  • episodes (int) – Number of declared unique episodes.

  • frames (int) – Number of frame rows scanned.

  • errors (int) – Total error count, including diagnostics beyond the detail cap.

  • warnings (int) – Total warning count, including diagnostics beyond the cap.

  • issues (list[DatasetValidationIssue]) – Up to 1,000 detailed diagnostics.

Methods:

__init__(root[, episodes, frames, errors, ...])

raise_for_errors()

Raise a concise error when data failed validation.

to_dict()

Return JSON-compatible diagnostics and aggregate counts.

Attributes:

ok

Return whether the scan found no errors.
__init__(root, episodes=0, frames=0, errors=0, warnings=0, issues=<factory>)#
property ok: bool#

Return whether the scan found no errors.

raise_for_errors()[source]#

Raise a concise error when data failed validation.

Raises:

ValueError – If the scan found one or more errors.

Return type:

None

to_dict()[source]#

Return JSON-compatible diagnostics and aggregate counts.

Return type:

dict[str, Any]

Returns:

Report mapping suitable for CLI output or automated gates.

embodichain.data_pipeline.datasets.inspection.create_split_manifest(root, output, *, fractions=(0.8, 0.1, 0.1), seed=0, group_by='lineage', scene_field='scene_id', success_only=False, failure_only=False, max_recoveries=None, require_physical_valid=False, allow_legacy_independent=False)[source]#

Write an atomic split manifest without modifying dataset metadata.

Shared source, parent, and episode identities form transitive groups before quality filtering. Scene grouping adds scene identities to these groups; it cannot separate two derivatives of the same episode. Unknown provenance conservatively groups the recording run, or the whole dataset if no run is known. Fractions are targets: indivisible groups may leave a split empty.

Parameters:

  • root (str | Path) – Dataset directory with episode metadata and optional sidecars.

  • output (str | Path) – Separate JSON manifest destination; existing source files cannot be overwritten.

  • fractions (Sequence[float]) – Train, validation, and test fractions summing to one.

  • seed (int) – Seed for deterministic group assignment.

  • group_by (Literal['lineage', 'scene']) – Group by lineage, optionally also by explicit scene identity.

  • scene_field (str) – Sidecar field containing a scene identifier.

  • success_only (bool) – Select only explicitly successful episodes.

  • failure_only (bool) – Select only explicitly failed episodes.

  • max_recoveries (int | None) – Maximum explicit recovery count; unknown counts are excluded when this filter is requested. Canonical Task Program traces also support counting row-local replanned events plus distinct workflow recovery cycles; retries are counted once.

  • require_physical_valid (bool) – Require an explicit accepted physical validation or an ordered-stable-regions objective with successful valid measured evidence. Declaring an objective without measurements is insufficient.

  • allow_legacy_independent (bool) – Explicitly assume unidentified legacy episodes are independent; recorded as an assumption in the manifest.

Return type:

dict[str, Any]

Returns:

Manifest containing split indices, group identities, exclusions, counts, quality selection, and provenance assumptions.

Raises:

ValueError – If arguments or episode metadata are invalid.

embodichain.data_pipeline.datasets.inspection.main(argv=None)[source]#

Run the standalone dataset validation/split command.

Parameters:

argv (Sequence[str] | None) – Arguments excluding the executable name; defaults to process args.

Return type:

int

Returns:

Zero on success, one for validation errors, two for invalid requests.

embodichain.data_pipeline.datasets.inspection.validate_dataset(root, *, check_media=False, media_frame_counts=False, batch_size=4096)[source]#

Validate on-disk LeRobot numeric, semantic, and persistence invariants.

Parameters:

  • root (str | Path) – Dataset root containing meta/info.json and frame Parquet files.

  • check_media (bool) – Check referenced RGB/image/depth files exist.

  • media_frame_counts (bool) – Also decode every referenced video once to check its frame count and RGB cadence. Implies check_media; requires PyAV.

  • batch_size (int) – Maximum Arrow frame rows loaded per batch.

Return type:

DatasetValidationReport

Returns:

Diagnostic report. Missing optional EmbodiChain provenance is a warning; malformed metadata, pending commits, and broken links are errors.

Raises:

ValueError – If batch_size is not a positive integer.

Contents