embodichain.data_pipeline.engine

Contents

embodichain.data_pipeline.engine#

Process-safe shared buffer (OnlineDataEngine) that decouples simulation producers from training consumers.

Overview#

Online data streaming engine: a process-safe shared buffer for trajectory data.

Functions

OnlineDataEngine

Engine for managing Online Data Streaming (ODS) and environment rollouts.

OnlineDataEngineCfg

OnlineDataEngineCfg(buffer_size: 'int' = <factory>, max_episode_steps: 'int' = <factory>, state_dim: 'int' = <factory>, buffer_device: 'str' = <factory>, gym_config: 'dict' = <factory>, action_config: 'dict' = <factory>, refill_threshold: 'int' = <factory>, max_generation_attempts: 'int' = <factory>, initialization_timeout: 'float' = <factory>, language_buffer_bytes: 'int' = <factory>, action_dim: 'int | None' = <factory>, refill_policy: "Literal['sample_threshold', 'manual']" = <factory>, include_provenance: 'bool' = <factory>, episode_sample_budget: 'int | None' = <factory>, episode_sample_hard_limit: 'int | None' = <factory>, lifecycle_event_capacity: 'int' = <factory>, action_semantics: "Literal['expert_controller_command', 'policy_request']" = <factory>, action_contracts: 'tuple[str, ...] | None' = <factory>, shutdown_timeout: 'float' = <factory>)

Process-safe shared buffer (OnlineDataEngine) that decouples simulation producers from training consumers.

Exceptions:

OnlineDataBudgetExhaustedError

alias of DataBudgetExhaustedError

OnlineDataWorkerError

Fallback error for a worker exception that cannot be reconstructed.

Classes:

OnlineDataEngine

Engine for managing Online Data Streaming (ODS) and environment rollouts.

OnlineDataEngineCfg

OnlineDataEngineCfg(buffer_size: 'int' = <factory>, max_episode_steps: 'int' = <factory>, state_dim: 'int' = <factory>, buffer_device: 'str' = <factory>, gym_config: 'dict' = <factory>, action_config: 'dict' = <factory>, refill_threshold: 'int' = <factory>, max_generation_attempts: 'int' = <factory>, initialization_timeout: 'float' = <factory>, language_buffer_bytes: 'int' = <factory>, action_dim: 'int | None' = <factory>, refill_policy: "Literal['sample_threshold', 'manual']" = <factory>, include_provenance: 'bool' = <factory>, episode_sample_budget: 'int | None' = <factory>, episode_sample_hard_limit: 'int | None' = <factory>, lifecycle_event_capacity: 'int' = <factory>, action_semantics: "Literal['expert_controller_command', 'policy_request']" = <factory>, action_contracts: 'tuple[str, ...] | None' = <factory>, shutdown_timeout: 'float' = <factory>)

OnlineDataEngineState

Lifecycle states for OnlineDataEngine.

embodichain.data_pipeline.engine.OnlineDataBudgetExhaustedError#

alias of DataBudgetExhaustedError

class embodichain.data_pipeline.engine.OnlineDataEngine[source]#

Bases: object

Engine for managing Online Data Streaming (ODS) and environment rollouts.

Creates a shared rollout buffer in CPU shared memory, spawns a dedicated simulation subprocess that fills the buffer with demonstration trajectories, and exposes a sample_batch() method for the training process to draw batches of trajectory chunks.

Subprocess lifecycle

The simulation subprocess is started in start() and immediately receives a fill signal so the buffer is populated before the first call to sample_batch(). The subprocess loops indefinitely: it waits for fill_signal, runs buffer_size // num_envs rollouts to overwrite every buffer slot, then goes back to waiting.

Concurrency and lock protection

_lock_index [write_start, write_end) is updated by the subprocess after each rollout so that sample_batch() can skip the slot currently being written to, preventing partial reads.

Refill criterion

sample_batch() accumulates the total number of individual trajectory samples drawn into _sample_count. When this counter exceeds refill_threshold the fill signal is raised and the counter resets to zero. This amortises the cost of GPU-accelerated simulation across many training iterations.

Lifecycle state

Every instance starts in OnlineDataEngineState.CREATED, passes through STARTING while the first fill is running, and only serves data in READY. Worker failures transition to FAILED and explicit cleanup transitions to terminal STOPPED; failed or stopped instances cannot be restarted.

Parameters:

cfg (OnlineDataEngineCfg) – Engine configuration.

shared_buffer#

Shared-memory TensorDict of shape [buffer_size, max_episode_steps, ...].

buffer_size#

Total number of trajectory slots in the shared buffer.

device#

Device of the shared buffer.

state#

Current OnlineDataEngineState.

is_init#

True only while the engine is ready to sample.

Methods:

__init__(cfg)

check_health()

Require a ready engine and surface existing worker failures.

drain_lifecycle_events()

Return and remove pending publication/eviction records in order.

get_flow_status()

Return detached counters for READY-only host flow control.

get_sampling_budget_snapshot()

Return a detached quota audit under the producer/sampling copy lock.

get_stats()

Return a detached snapshot of production, consumption and valid rows.

request_refill()

Request a fill, coalescing an outstanding manual request.

resolve_language(batch)

Decode overall tasks and current subtasks for a copied sample batch.

sample_batch(batch_size, chunk_size[, ...])

Sample a batch of trajectory chunks from the shared rollout buffer.

sample_batch_with_context(batch_size, ...[, ...])

Sample causal observations and future supervision around valid anchors.

start()

Start the worker and block until its first fill completes.

stop()

Terminate the simulation subprocess and release resources.

wait_for_generation(generation[, timeout])

Wait for a requested generation while propagating worker failures.

Attributes:

is_init

Whether the engine is ready to serve initialized data.

state

Return the engine's current lifecycle state.

__init__(cfg)[source]#
check_health()[source]#

Require a ready engine and surface existing worker failures.

This check neither samples data nor requests a refill. It checks worker liveness through the existing owner/consumer error channel.

Raises:
  • RuntimeError – If the engine is not READY or its worker has exited.

  • BaseException – The original reported worker or startup exception.

Return type:

None

drain_lifecycle_events()[source]#

Return and remove pending publication/eviction records in order.

Owner-only, including after stop or failure, so diagnostics are not lost. Disabled logging returns an empty list. Each record contains kind, sequence, run_id, episode_id, generation, slot, monotonic_ns, valid_frame_indices, eligible_anchor_indices, valid_frames and eligible_anchors. Eligible anchors refer to episode-mode context sampling (valid frames with nonnegative continuity). This is a host-local monotonic clock, not a wall-clock timestamp.

Drain regularly while producing, including during owner-side waits. wait_for_generation does not consume these records and may time out if an enabled log is full. Records are detached from later row reuse.

Return type:

list[dict[str, object]]

Returns:

Detached lifecycle records in publication order, or an empty list when logging is disabled or no records are pending.

Raises:

RuntimeError – If called outside the engine’s owner process.

get_flow_status()[source]#

Return detached counters for READY-only host flow control.

Producer counters and quota totals use separate lock snapshots. The quota sum includes only published, unlocked rows with a valid anchor and nonnegative continuity when continuity metadata is present. It is a coarse bound; sequence lengths and filters may reduce eligibility. This method creates no per-row audit records or payload copies.

Return type:

dict[str, int | None]

Returns:

Requested/completed generations and successful episodes, plus unlocked_remaining_chunks and lifetime total_delivered_chunks. Quota counters are None when protection is disabled.

Raises:
  • RuntimeError – If the engine is not READY or its worker has exited.

  • BaseException – The original reported worker or startup exception.

get_sampling_budget_snapshot()[source]#

Return a detached quota audit under the producer/sampling copy lock.

rows describes currently retained source identities and includes run_id, episode_id, generation, slot, assigned, delivered, remaining and locked. Lifetime total_delivered_chunks includes retired sources. The unlocked sum covers sources with at least one valid continuity-qualified anchor; a requested chunk length or sampling mode may further restrict eligibility. Known worker failures are propagated; clean STOPPED engines remain auditable. Disabled quotas return empty rows and unavailable counters.

Return type:

dict[str, object]

get_stats()[source]#

Return a detached snapshot of production, consumption and valid rows.

Return type:

dict[str, int | float]

Returns:

Engine identity, requested/completed generations, episode/frame counts, cumulative production seconds and trajectory sample count.

property is_init: bool#

Whether the engine is ready to serve initialized data.

request_refill()[source]#

Request a fill, coalescing an outstanding manual request.

Return type:

int

Returns:

Generation that will contain the requested fresh data.

Raises:

RuntimeError – If called outside READY or from a reader process.

resolve_language(batch)[source]#

Decode overall tasks and current subtasks for a copied sample batch.

The lookup remains valid after a ring-buffer refill or engine shutdown. Returned strings are separate from the shared numeric TensorDict; callers can tokenize them in a training transform without changing sampling locks or collation behavior.

Parameters:

batch (TensorDict) – Sampled TensorDict containing task_index/subtask_index.

Return type:

dict[str, object]

Returns:

task and subtask strings nested to match the batch dimensions (for example [batch_size][chunk_size]).

Raises:

ValueError – If language indices are absent, unknown, or padding.

sample_batch(batch_size, chunk_size, sampling_mode='episode')[source]#

Sample a batch of trajectory chunks from the shared rollout buffer.

Only fully valid windows are candidates, so padding or stale tail frames are never returned. episode mode allows a window to cross segment boundaries within one causal-continuity region, segment keeps every window inside one accepted segment, and boundary deliberately samples windows crossing a boundary between accepted segments. No mode crosses a discontinuous state-restore boundary.

After sampling the internal _sample_count is incremented by batch_size; if the count exceeds refill_threshold a buffer refill is triggered automatically.

Parameters:
  • batch_size (int) – Number of trajectory chunks to include in the batch.

  • chunk_size (int) – Number of consecutive timesteps in each chunk.

  • sampling_mode (Literal['episode', 'segment', 'boundary']) – Segment-boundary policy for candidate windows.

Return type:

TensorDict

Returns:

TensorDict with batch size [batch_size, chunk_size]. Integer task_index and subtask_index identify immutable language snapshots; resolve_language() decodes their descriptions.

Raises:
  • ValueError – If an argument is invalid.

  • DataUnavailableError – If no unlocked valid window satisfies the policy. Subclasses RuntimeError for existing callers.

  • RuntimeError – If a producer fails or the engine is not ready.

  • OnlineDataBudgetExhaustedError – If eligible source quotas cannot supply the complete batch. No partial batch is charged.

sample_batch_with_context(batch_size, history_length, prediction_horizon, *, generator=None, sampling_mode='episode', generations=None)[source]#

Sample causal observations and future supervision around valid anchors.

The current frame is at history_length - 1. Every valid frame, including episode starts and ends, can be an anchor. Context never crosses invalid frames or continuity boundaries. Out-of-range numeric payloads are zero; language/provenance IDs are -1. valid marks real frames and loss_mask selects only valid current/future targets. Selection, provenance and payload copies share the producer lock.

Parameters:
  • batch_size (int) – Number of contexts, with replacement.

  • history_length (int) – Number of causal observation frames including now.

  • prediction_horizon (int) – Number of target frames including now.

  • generator (Generator | None) – Optional RNG isolated from augmentation/model randomness.

  • sampling_mode (Literal['episode', 'segment']) – Episode context or one accepted segment.

  • generations (tuple[int, ...] | None) – Optional positive generation IDs; requires provenance.

Return type:

TensorDict

Returns:

TensorDict of shape [batch_size, history_length + prediction_horizon - 1]. Consumers must mask padded supervision and explicitly handle missing initial observation history.

Raises:
  • ValueError – If context lengths are not positive integers.

  • DataUnavailableError – If no unlocked valid anchor exists.

  • RuntimeError – If the worker fails or the engine is not ready.

  • OnlineDataBudgetExhaustedError – If eligible source quotas cannot supply the complete batch. Each context costs one chunk.

start()[source]#

Start the worker and block until its first fill completes.

Raises:
  • RuntimeError – If the engine was already started or stopped.

  • TimeoutError – If the first fill exceeds initialization_timeout.

  • BaseException – The original exception raised by the worker.

Return type:

None

property state: OnlineDataEngineState#

Return the engine’s current lifecycle state.

stop()[source]#

Terminate the simulation subprocess and release resources.

Sets the close signal and waits up to cfg.shutdown_timeout for the subprocess to exit gracefully (it checks the signal between rollout steps). If it remains alive after that wait, it is force-terminated.

Safe to call multiple times — subsequent calls are no-ops if the subprocess has already been terminated.

Return type:

None

wait_for_generation(generation, timeout=300.0)[source]#

Wait for a requested generation while propagating worker failures.

Parameters:
  • generation (int) – Positive generation returned by request_refill.

  • timeout (float) – Finite positive timeout in seconds.

Raises:
  • ValueError – If the generation was not requested or timeout is invalid.

  • TimeoutError – If publication does not finish before the deadline.

Return type:

None

class embodichain.data_pipeline.engine.OnlineDataEngineCfg[source]#

Bases: object

OnlineDataEngineCfg(buffer_size: ‘int’ = <factory>, max_episode_steps: ‘int’ = <factory>, state_dim: ‘int’ = <factory>, buffer_device: ‘str’ = <factory>, gym_config: ‘dict’ = <factory>, action_config: ‘dict’ = <factory>, refill_threshold: ‘int’ = <factory>, max_generation_attempts: ‘int’ = <factory>, initialization_timeout: ‘float’ = <factory>, language_buffer_bytes: ‘int’ = <factory>, action_dim: ‘int | None’ = <factory>, refill_policy: “Literal[‘sample_threshold’, ‘manual’]” = <factory>, include_provenance: ‘bool’ = <factory>, episode_sample_budget: ‘int | None’ = <factory>, episode_sample_hard_limit: ‘int | None’ = <factory>, lifecycle_event_capacity: ‘int’ = <factory>, action_semantics: “Literal[‘expert_controller_command’, ‘policy_request’]” = <factory>, action_contracts: ‘tuple[str, …] | None’ = <factory>, shutdown_timeout: ‘float’ = <factory>)

Methods:

__init__([buffer_size, max_episode_steps, ...])

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_config

Action configuration dictionary.

action_contracts

Expected ActionManager contract IDs in flat policy order.

action_dim

Stored action width; None preserves the historical state_dim width.

action_semantics

Stored action meaning; the default preserves expert-controller commands.

buffer_device

Device on which the shared buffer is allocated.

buffer_size

Number of episodes (environment trajectories) that can be stored in the shared buffer at once.

episode_sample_budget

Maximum delivered chunks assigned to each fresh source episode.

episode_sample_hard_limit

Positive signed-int64 ceiling for episode_sample_budget.

gym_config

Gym environment configuration dictionary (already loaded, not a file path).

include_provenance

Include numeric run/episode/generation IDs, episode frame indices, and successful-publication Unix and host-monotonic timestamps in nanoseconds.

initialization_timeout

Maximum seconds to wait for the worker's initial buffer fill.

language_buffer_bytes

Append-only task/subtask dictionary capacity in UTF-8 bytes.

lifecycle_event_capacity

Optional bounded publication/eviction log; zero disables logging.

max_episode_steps

Maximum number of timesteps per episode.

max_generation_attempts

Maximum planning/execution attempts for each buffer write transaction.

refill_policy

Automatic sample-count refill or owner-requested refill.

refill_threshold

Total number of samples (refill_threshold * buffer_size) drawn from the shared buffer before a refill is triggered.

shutdown_timeout

Maximum seconds to wait for graceful worker close before forcing shutdown.

state_dim

Dimensionality of the state space.

__init__(buffer_size=<factory>, max_episode_steps=<factory>, state_dim=<factory>, buffer_device=<factory>, gym_config=<factory>, action_config=<factory>, refill_threshold=<factory>, max_generation_attempts=<factory>, initialization_timeout=<factory>, language_buffer_bytes=<factory>, action_dim=<factory>, refill_policy=<factory>, include_provenance=<factory>, episode_sample_budget=<factory>, episode_sample_hard_limit=<factory>, lifecycle_event_capacity=<factory>, action_semantics=<factory>, action_contracts=<factory>, shutdown_timeout=<factory>)#
action_config: dict#

Action configuration dictionary. The contents depend on the specific environment and robot being used.

action_contracts: tuple[str, ...] | None#

Expected ActionManager contract IDs in flat policy order.

Required with an explicit action_dim for policy_request recording. The worker rejects a different live contract sequence, width, or joint order.

action_dim: int | None#

Stored action width; None preserves the historical state_dim width.

action_semantics: Literal['expert_controller_command', 'policy_request']#

Stored action meaning; the default preserves expert-controller commands.

buffer_device: str#

Device on which the shared buffer is allocated.

buffer_size: int#

Number of episodes (environment trajectories) that can be stored in the shared buffer at once. Must be ≥ num_envs and ideally a multiple of num_envs.

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.

episode_sample_budget: int | None#

Maximum delivered chunks assigned to each fresh source episode.

Each returned batch row costs one, regardless of context length or repeated selection of the same episode. Requires episode_sample_hard_limit and a CPU shared buffer. None disables per-source quotas without changing sampling. Exhaustion does not itself request a refill; choose manual replenishment or an automatic threshold that can be reached within the assigned quotas.

episode_sample_hard_limit: int | None#

Positive signed-int64 ceiling for episode_sample_budget.

Both quota fields must be supplied together and budget must not exceed this ceiling. The assigned budget is fixed at engine construction; no source may borrow another source’s unused quota or receive more on generation changes.

gym_config: dict#

Gym environment configuration dictionary (already loaded, not a file path). The contents depend on the specific environment being used. Default is None.

include_provenance: bool#

Include numeric run/episode/generation IDs, episode frame indices, and successful-publication Unix and host-monotonic timestamps in nanoseconds. Invalid frames use -1 for every provenance field.

initialization_timeout: float#

Maximum seconds to wait for the worker’s initial buffer fill.

language_buffer_bytes: int#

Append-only task/subtask dictionary capacity in UTF-8 bytes. Descriptions remain available for copied samples after trajectory slots are reused. Exhaustion fails the worker explicitly; increase this for many unique tasks.

lifecycle_event_capacity: int#

Optional bounded publication/eviction log; zero disables logging.

When enabled, capacity must fit twice buffer_size and all initial-fill records (including overlapping final-batch evictions). The owner must drain events while refilling; a full log backpressures the producer outside the trajectory-copy lock. Stop interrupts this wait.

max_episode_steps: int#

Maximum number of timesteps per episode. Must be ≥ chunk_size used by OnlineDataset.

max_generation_attempts: int#

Maximum planning/execution attempts for each buffer write transaction.

refill_policy: Literal['sample_threshold', 'manual']#

Automatic sample-count refill or owner-requested refill.

refill_threshold: int#

Total number of samples (refill_threshold * buffer_size) drawn from the shared buffer before a refill is triggered. Accumulates across all calls to OnlineDataEngine.sample_batch(). When this threshold is exceeded the engine signals the simulation subprocess to regenerate the entire buffer, amortising the cost of environment simulation over many training steps.

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.

shutdown_timeout: float#

Maximum seconds to wait for graceful worker close before forcing shutdown.

Must be finite and positive. Forced termination still fails the engine’s recorder-durability contract, regardless of this timeout.

state_dim: int#

Dimensionality of the state space.

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

Bases: str, Enum

Lifecycle states for OnlineDataEngine.

Attributes:

Methods:

__new__(value)

CREATED = 'CREATED'#
FAILED = 'FAILED'#
READY = 'READY'#
STARTING = 'STARTING'#
STOPPED = 'STOPPED'#
__new__(value)#
exception embodichain.data_pipeline.engine.OnlineDataWorkerError[source]#

Bases: RuntimeError

Fallback error for a worker exception that cannot be reconstructed.

Exceptions:

OnlineDataBudgetExhaustedError

alias of DataBudgetExhaustedError

OnlineDataWorkerError

Fallback error for a worker exception that cannot be reconstructed.

Classes:

OnlineDataEngine

Engine for managing Online Data Streaming (ODS) and environment rollouts.

OnlineDataEngineCfg

OnlineDataEngineCfg(buffer_size: 'int' = <factory>, max_episode_steps: 'int' = <factory>, state_dim: 'int' = <factory>, buffer_device: 'str' = <factory>, gym_config: 'dict' = <factory>, action_config: 'dict' = <factory>, refill_threshold: 'int' = <factory>, max_generation_attempts: 'int' = <factory>, initialization_timeout: 'float' = <factory>, language_buffer_bytes: 'int' = <factory>, action_dim: 'int | None' = <factory>, refill_policy: "Literal['sample_threshold', 'manual']" = <factory>, include_provenance: 'bool' = <factory>, episode_sample_budget: 'int | None' = <factory>, episode_sample_hard_limit: 'int | None' = <factory>, lifecycle_event_capacity: 'int' = <factory>, action_semantics: "Literal['expert_controller_command', 'policy_request']" = <factory>, action_contracts: 'tuple[str, ...] | None' = <factory>, shutdown_timeout: 'float' = <factory>)

OnlineDataEngineState

Lifecycle states for OnlineDataEngine.

embodichain.data_pipeline.engine.data.OnlineDataBudgetExhaustedError

alias of DataBudgetExhaustedError

class embodichain.data_pipeline.engine.data.OnlineDataEngine[source]

Bases: object

Engine for managing Online Data Streaming (ODS) and environment rollouts.

Creates a shared rollout buffer in CPU shared memory, spawns a dedicated simulation subprocess that fills the buffer with demonstration trajectories, and exposes a sample_batch() method for the training process to draw batches of trajectory chunks.

Subprocess lifecycle

The simulation subprocess is started in start() and immediately receives a fill signal so the buffer is populated before the first call to sample_batch(). The subprocess loops indefinitely: it waits for fill_signal, runs buffer_size // num_envs rollouts to overwrite every buffer slot, then goes back to waiting.

Concurrency and lock protection

_lock_index [write_start, write_end) is updated by the subprocess after each rollout so that sample_batch() can skip the slot currently being written to, preventing partial reads.

Refill criterion

sample_batch() accumulates the total number of individual trajectory samples drawn into _sample_count. When this counter exceeds refill_threshold the fill signal is raised and the counter resets to zero. This amortises the cost of GPU-accelerated simulation across many training iterations.

Lifecycle state

Every instance starts in OnlineDataEngineState.CREATED, passes through STARTING while the first fill is running, and only serves data in READY. Worker failures transition to FAILED and explicit cleanup transitions to terminal STOPPED; failed or stopped instances cannot be restarted.

Parameters:

cfg (OnlineDataEngineCfg) – Engine configuration.

shared_buffer

Shared-memory TensorDict of shape [buffer_size, max_episode_steps, ...].

buffer_size

Total number of trajectory slots in the shared buffer.

device

Device of the shared buffer.

state

Current OnlineDataEngineState.

is_init

True only while the engine is ready to sample.

Methods:

__init__(cfg)

check_health()

Require a ready engine and surface existing worker failures.

drain_lifecycle_events()

Return and remove pending publication/eviction records in order.

get_flow_status()

Return detached counters for READY-only host flow control.

get_sampling_budget_snapshot()

Return a detached quota audit under the producer/sampling copy lock.

get_stats()

Return a detached snapshot of production, consumption and valid rows.

request_refill()

Request a fill, coalescing an outstanding manual request.

resolve_language(batch)

Decode overall tasks and current subtasks for a copied sample batch.

sample_batch(batch_size, chunk_size[, ...])

Sample a batch of trajectory chunks from the shared rollout buffer.

sample_batch_with_context(batch_size, ...[, ...])

Sample causal observations and future supervision around valid anchors.

start()

Start the worker and block until its first fill completes.

stop()

Terminate the simulation subprocess and release resources.

wait_for_generation(generation[, timeout])

Wait for a requested generation while propagating worker failures.

Attributes:

is_init

Whether the engine is ready to serve initialized data.

state

Return the engine's current lifecycle state.

__init__(cfg)[source]
check_health()[source]

Require a ready engine and surface existing worker failures.

This check neither samples data nor requests a refill. It checks worker liveness through the existing owner/consumer error channel.

Raises:
  • RuntimeError – If the engine is not READY or its worker has exited.

  • BaseException – The original reported worker or startup exception.

Return type:

None

drain_lifecycle_events()[source]

Return and remove pending publication/eviction records in order.

Owner-only, including after stop or failure, so diagnostics are not lost. Disabled logging returns an empty list. Each record contains kind, sequence, run_id, episode_id, generation, slot, monotonic_ns, valid_frame_indices, eligible_anchor_indices, valid_frames and eligible_anchors. Eligible anchors refer to episode-mode context sampling (valid frames with nonnegative continuity). This is a host-local monotonic clock, not a wall-clock timestamp.

Drain regularly while producing, including during owner-side waits. wait_for_generation does not consume these records and may time out if an enabled log is full. Records are detached from later row reuse.

Return type:

list[dict[str, object]]

Returns:

Detached lifecycle records in publication order, or an empty list when logging is disabled or no records are pending.

Raises:

RuntimeError – If called outside the engine’s owner process.

get_flow_status()[source]

Return detached counters for READY-only host flow control.

Producer counters and quota totals use separate lock snapshots. The quota sum includes only published, unlocked rows with a valid anchor and nonnegative continuity when continuity metadata is present. It is a coarse bound; sequence lengths and filters may reduce eligibility. This method creates no per-row audit records or payload copies.

Return type:

dict[str, int | None]

Returns:

Requested/completed generations and successful episodes, plus unlocked_remaining_chunks and lifetime total_delivered_chunks. Quota counters are None when protection is disabled.

Raises:
  • RuntimeError – If the engine is not READY or its worker has exited.

  • BaseException – The original reported worker or startup exception.

get_sampling_budget_snapshot()[source]

Return a detached quota audit under the producer/sampling copy lock.

rows describes currently retained source identities and includes run_id, episode_id, generation, slot, assigned, delivered, remaining and locked. Lifetime total_delivered_chunks includes retired sources. The unlocked sum covers sources with at least one valid continuity-qualified anchor; a requested chunk length or sampling mode may further restrict eligibility. Known worker failures are propagated; clean STOPPED engines remain auditable. Disabled quotas return empty rows and unavailable counters.

Return type:

dict[str, object]

get_stats()[source]

Return a detached snapshot of production, consumption and valid rows.

Return type:

dict[str, int | float]

Returns:

Engine identity, requested/completed generations, episode/frame counts, cumulative production seconds and trajectory sample count.

property is_init: bool

Whether the engine is ready to serve initialized data.

request_refill()[source]

Request a fill, coalescing an outstanding manual request.

Return type:

int

Returns:

Generation that will contain the requested fresh data.

Raises:

RuntimeError – If called outside READY or from a reader process.

resolve_language(batch)[source]

Decode overall tasks and current subtasks for a copied sample batch.

The lookup remains valid after a ring-buffer refill or engine shutdown. Returned strings are separate from the shared numeric TensorDict; callers can tokenize them in a training transform without changing sampling locks or collation behavior.

Parameters:

batch (TensorDict) – Sampled TensorDict containing task_index/subtask_index.

Return type:

dict[str, object]

Returns:

task and subtask strings nested to match the batch dimensions (for example [batch_size][chunk_size]).

Raises:

ValueError – If language indices are absent, unknown, or padding.

sample_batch(batch_size, chunk_size, sampling_mode='episode')[source]

Sample a batch of trajectory chunks from the shared rollout buffer.

Only fully valid windows are candidates, so padding or stale tail frames are never returned. episode mode allows a window to cross segment boundaries within one causal-continuity region, segment keeps every window inside one accepted segment, and boundary deliberately samples windows crossing a boundary between accepted segments. No mode crosses a discontinuous state-restore boundary.

After sampling the internal _sample_count is incremented by batch_size; if the count exceeds refill_threshold a buffer refill is triggered automatically.

Parameters:
  • batch_size (int) – Number of trajectory chunks to include in the batch.

  • chunk_size (int) – Number of consecutive timesteps in each chunk.

  • sampling_mode (Literal['episode', 'segment', 'boundary']) – Segment-boundary policy for candidate windows.

Return type:

TensorDict

Returns:

TensorDict with batch size [batch_size, chunk_size]. Integer task_index and subtask_index identify immutable language snapshots; resolve_language() decodes their descriptions.

Raises:
  • ValueError – If an argument is invalid.

  • DataUnavailableError – If no unlocked valid window satisfies the policy. Subclasses RuntimeError for existing callers.

  • RuntimeError – If a producer fails or the engine is not ready.

  • OnlineDataBudgetExhaustedError – If eligible source quotas cannot supply the complete batch. No partial batch is charged.

sample_batch_with_context(batch_size, history_length, prediction_horizon, *, generator=None, sampling_mode='episode', generations=None)[source]

Sample causal observations and future supervision around valid anchors.

The current frame is at history_length - 1. Every valid frame, including episode starts and ends, can be an anchor. Context never crosses invalid frames or continuity boundaries. Out-of-range numeric payloads are zero; language/provenance IDs are -1. valid marks real frames and loss_mask selects only valid current/future targets. Selection, provenance and payload copies share the producer lock.

Parameters:
  • batch_size (int) – Number of contexts, with replacement.

  • history_length (int) – Number of causal observation frames including now.

  • prediction_horizon (int) – Number of target frames including now.

  • generator (Generator | None) – Optional RNG isolated from augmentation/model randomness.

  • sampling_mode (Literal['episode', 'segment']) – Episode context or one accepted segment.

  • generations (tuple[int, ...] | None) – Optional positive generation IDs; requires provenance.

Return type:

TensorDict

Returns:

TensorDict of shape [batch_size, history_length + prediction_horizon - 1]. Consumers must mask padded supervision and explicitly handle missing initial observation history.

Raises:
  • ValueError – If context lengths are not positive integers.

  • DataUnavailableError – If no unlocked valid anchor exists.

  • RuntimeError – If the worker fails or the engine is not ready.

  • OnlineDataBudgetExhaustedError – If eligible source quotas cannot supply the complete batch. Each context costs one chunk.

start()[source]

Start the worker and block until its first fill completes.

Raises:
  • RuntimeError – If the engine was already started or stopped.

  • TimeoutError – If the first fill exceeds initialization_timeout.

  • BaseException – The original exception raised by the worker.

Return type:

None

property state: OnlineDataEngineState

Return the engine’s current lifecycle state.

stop()[source]

Terminate the simulation subprocess and release resources.

Sets the close signal and waits up to cfg.shutdown_timeout for the subprocess to exit gracefully (it checks the signal between rollout steps). If it remains alive after that wait, it is force-terminated.

Safe to call multiple times — subsequent calls are no-ops if the subprocess has already been terminated.

Return type:

None

wait_for_generation(generation, timeout=300.0)[source]

Wait for a requested generation while propagating worker failures.

Parameters:
  • generation (int) – Positive generation returned by request_refill.

  • timeout (float) – Finite positive timeout in seconds.

Raises:
  • ValueError – If the generation was not requested or timeout is invalid.

  • TimeoutError – If publication does not finish before the deadline.

Return type:

None

class embodichain.data_pipeline.engine.data.OnlineDataEngineCfg[source]

Bases: object

OnlineDataEngineCfg(buffer_size: ‘int’ = <factory>, max_episode_steps: ‘int’ = <factory>, state_dim: ‘int’ = <factory>, buffer_device: ‘str’ = <factory>, gym_config: ‘dict’ = <factory>, action_config: ‘dict’ = <factory>, refill_threshold: ‘int’ = <factory>, max_generation_attempts: ‘int’ = <factory>, initialization_timeout: ‘float’ = <factory>, language_buffer_bytes: ‘int’ = <factory>, action_dim: ‘int | None’ = <factory>, refill_policy: “Literal[‘sample_threshold’, ‘manual’]” = <factory>, include_provenance: ‘bool’ = <factory>, episode_sample_budget: ‘int | None’ = <factory>, episode_sample_hard_limit: ‘int | None’ = <factory>, lifecycle_event_capacity: ‘int’ = <factory>, action_semantics: “Literal[‘expert_controller_command’, ‘policy_request’]” = <factory>, action_contracts: ‘tuple[str, …] | None’ = <factory>, shutdown_timeout: ‘float’ = <factory>)

Methods:

__init__([buffer_size, max_episode_steps, ...])

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_config

Action configuration dictionary.

action_contracts

Expected ActionManager contract IDs in flat policy order.

action_dim

Stored action width; None preserves the historical state_dim width.

action_semantics

Stored action meaning; the default preserves expert-controller commands.

buffer_device

Device on which the shared buffer is allocated.

buffer_size

Number of episodes (environment trajectories) that can be stored in the shared buffer at once.

episode_sample_budget

Maximum delivered chunks assigned to each fresh source episode.

episode_sample_hard_limit

Positive signed-int64 ceiling for episode_sample_budget.

gym_config

Gym environment configuration dictionary (already loaded, not a file path).

include_provenance

Include numeric run/episode/generation IDs, episode frame indices, and successful-publication Unix and host-monotonic timestamps in nanoseconds.

initialization_timeout

Maximum seconds to wait for the worker's initial buffer fill.

language_buffer_bytes

Append-only task/subtask dictionary capacity in UTF-8 bytes.

lifecycle_event_capacity

Optional bounded publication/eviction log; zero disables logging.

max_episode_steps

Maximum number of timesteps per episode.

max_generation_attempts

Maximum planning/execution attempts for each buffer write transaction.

refill_policy

Automatic sample-count refill or owner-requested refill.

refill_threshold

Total number of samples (refill_threshold * buffer_size) drawn from the shared buffer before a refill is triggered.

shutdown_timeout

Maximum seconds to wait for graceful worker close before forcing shutdown.

state_dim

Dimensionality of the state space.

__init__(buffer_size=<factory>, max_episode_steps=<factory>, state_dim=<factory>, buffer_device=<factory>, gym_config=<factory>, action_config=<factory>, refill_threshold=<factory>, max_generation_attempts=<factory>, initialization_timeout=<factory>, language_buffer_bytes=<factory>, action_dim=<factory>, refill_policy=<factory>, include_provenance=<factory>, episode_sample_budget=<factory>, episode_sample_hard_limit=<factory>, lifecycle_event_capacity=<factory>, action_semantics=<factory>, action_contracts=<factory>, shutdown_timeout=<factory>)
action_config: dict

Action configuration dictionary. The contents depend on the specific environment and robot being used.

action_contracts: tuple[str, ...] | None

Expected ActionManager contract IDs in flat policy order.

Required with an explicit action_dim for policy_request recording. The worker rejects a different live contract sequence, width, or joint order.

action_dim: int | None

Stored action width; None preserves the historical state_dim width.

action_semantics: Literal['expert_controller_command', 'policy_request']

Stored action meaning; the default preserves expert-controller commands.

buffer_device: str

Device on which the shared buffer is allocated.

buffer_size: int

Number of episodes (environment trajectories) that can be stored in the shared buffer at once. Must be ≥ num_envs and ideally a multiple of num_envs.

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.

episode_sample_budget: int | None

Maximum delivered chunks assigned to each fresh source episode.

Each returned batch row costs one, regardless of context length or repeated selection of the same episode. Requires episode_sample_hard_limit and a CPU shared buffer. None disables per-source quotas without changing sampling. Exhaustion does not itself request a refill; choose manual replenishment or an automatic threshold that can be reached within the assigned quotas.

episode_sample_hard_limit: int | None

Positive signed-int64 ceiling for episode_sample_budget.

Both quota fields must be supplied together and budget must not exceed this ceiling. The assigned budget is fixed at engine construction; no source may borrow another source’s unused quota or receive more on generation changes.

gym_config: dict

Gym environment configuration dictionary (already loaded, not a file path). The contents depend on the specific environment being used. Default is None.

include_provenance: bool

Include numeric run/episode/generation IDs, episode frame indices, and successful-publication Unix and host-monotonic timestamps in nanoseconds. Invalid frames use -1 for every provenance field.

initialization_timeout: float

Maximum seconds to wait for the worker’s initial buffer fill.

language_buffer_bytes: int

Append-only task/subtask dictionary capacity in UTF-8 bytes. Descriptions remain available for copied samples after trajectory slots are reused. Exhaustion fails the worker explicitly; increase this for many unique tasks.

lifecycle_event_capacity: int

Optional bounded publication/eviction log; zero disables logging.

When enabled, capacity must fit twice buffer_size and all initial-fill records (including overlapping final-batch evictions). The owner must drain events while refilling; a full log backpressures the producer outside the trajectory-copy lock. Stop interrupts this wait.

max_episode_steps: int

Maximum number of timesteps per episode. Must be ≥ chunk_size used by OnlineDataset.

max_generation_attempts: int

Maximum planning/execution attempts for each buffer write transaction.

refill_policy: Literal['sample_threshold', 'manual']

Automatic sample-count refill or owner-requested refill.

refill_threshold: int

Total number of samples (refill_threshold * buffer_size) drawn from the shared buffer before a refill is triggered. Accumulates across all calls to OnlineDataEngine.sample_batch(). When this threshold is exceeded the engine signals the simulation subprocess to regenerate the entire buffer, amortising the cost of environment simulation over many training steps.

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.

shutdown_timeout: float

Maximum seconds to wait for graceful worker close before forcing shutdown.

Must be finite and positive. Forced termination still fails the engine’s recorder-durability contract, regardless of this timeout.

state_dim: int

Dimensionality of the state space.

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.engine.data.OnlineDataEngineState[source]

Bases: str, Enum

Lifecycle states for OnlineDataEngine.

Attributes:

CREATED

FAILED

READY

STARTING

STOPPED

Methods:

__new__(value)

CREATED = 'CREATED'
FAILED = 'FAILED'
READY = 'READY'
STARTING = 'STARTING'
STOPPED = 'STOPPED'
__new__(value)
exception embodichain.data_pipeline.engine.data.OnlineDataWorkerError[source]

Bases: RuntimeError

Fallback error for a worker exception that cannot be reconstructed.

Recorded actions and provenance#

OnlineDataEngineCfg.action_dim optionally separates stored action width from state_dim; omitting it preserves the historical shared width. action_semantics="expert_controller_command" is the default and follows the environment’s ExpertActionSpec. action_semantics="policy_request" requires an explicit width and ordered action_contracts matching the live ActionManager. The worker validates resolved semantics, dimensions, term slices and controlled-joint order before generating transitions. Explicit policy recording uses raw requests independently of the expert schema and does not require a DatasetManager.

refill_policy="manual" lets the owner request complete generations with request_refill() and observe publication with wait_for_generation(). The default remains sample-threshold refill. get_stats() returns detached counters and excludes the active write window from valid-frame counts. check_health() requires READY and surfaces worker failures without sampling or refill. get_flow_status() supplies compact generation, accepted-episode and optional quota counters for frequent scheduling checks. Its producer and quota counters use separate lock snapshots; unlocked quota is a coarse bound that does not apply every sequence/filter requirement. Use get_sampling_budget_snapshot() for the retained-source audit. shutdown_timeout is the finite, positive graceful-close wait in seconds (default 5.0). Increasing it allows longer worker teardown; it does not change forced terminate/kill waits or their terminal failure semantics.

include_provenance=True adds ods_run_id, ods_episode_id, ods_generation, ods_frame_index and ods_accepted_at_ns to each frame. The frame index is its zero-based stored episode position. The timestamp is successful-publication Unix nanoseconds, stamped under the publication lock and copied with its payload. Invalid frames and causal-history padding use -1 for all provenance fields. These engine-local identities do not restore shared-buffer contents or replace offline corpus content identities.

Shared language snapshots#

Generated batches include numeric task_index and subtask_index columns. OnlineDataEngine.resolve_language(batch) returns nested task/subtask strings for tokenization, including in DataLoader transforms. IDs are engine-local and append-only, so a copied sample still resolves after its trajectory slot is refilled. They are independent of an offline recorder’s vocabulary indices. History padding retains -1 in both language columns. The resolver rejects invalid IDs, so select real valid or supervised loss_mask frames before decoding padded histories.

The registry uses a bounded shared UTF-8 buffer instead of mutable strings in TensorDict storage. language_buffer_bytes limits its lifetime capacity; overflow fails the producer before publishing that rollout. Consumers use the same immutable indices and never depend on the current contents of reused rows.

The package-level alias is available for convenience; the canonical class and its methods are indexed under the language module.

class embodichain.data_pipeline.engine.SharedLanguageRegistry[source]

Share immutable text indices without a multiprocessing manager server.

Task and subtask indices have separate dictionaries. New descriptions append UTF-8 JSON records to a bounded shared byte buffer; previous indices never change, even when their trajectory slots are overwritten. Each process keeps its own incremental lookup cache and uses the shared length lock to publish or read complete records.

Parameters:
  • capacity_bytes (int) – Maximum encoded registry size for the engine’s lifetime.

  • context (BaseContext | None) – Multiprocessing context used by the producer and consumers.

Raises:

ValueError – If capacity is not a positive integer.

Methods:

__init__([capacity_bytes, context])

__new__(**kwargs)

__init__(capacity_bytes=1048576, context=None)[source]
__new__(**kwargs)

Append-only shared language dictionaries for online trajectory consumers.

Classes:

SharedLanguageRegistry

Share immutable text indices without a multiprocessing manager server.

class embodichain.data_pipeline.engine.language.SharedLanguageRegistry[source]#

Share immutable text indices without a multiprocessing manager server.

Task and subtask indices have separate dictionaries. New descriptions append UTF-8 JSON records to a bounded shared byte buffer; previous indices never change, even when their trajectory slots are overwritten. Each process keeps its own incremental lookup cache and uses the shared length lock to publish or read complete records.

Parameters:
  • capacity_bytes (int) – Maximum encoded registry size for the engine’s lifetime.

  • context (BaseContext | None) – Multiprocessing context used by the producer and consumers.

Raises:

ValueError – If capacity is not a positive integer.

Methods:

__init__([capacity_bytes, context])

intern(description, *[, kind])

Return an existing text index or append one immutable dictionary entry.

resolve(indices, *[, kind])

Decode a tensor of indices while preserving its nested batch shape.

__init__(capacity_bytes=1048576, context=None)[source]#
intern(description, *, kind='task')[source]#

Return an existing text index or append one immutable dictionary entry.

Parameters:
  • description (str) – Nonempty language snapshot, preserved exactly.

  • kind (Literal['task', 'subtask']) – The task or subtask dictionary to update.

Return type:

int

Returns:

Stable zero-based index in the selected dictionary.

Raises:
  • ValueError – If the description or dictionary kind is invalid.

  • OverflowError – If a new description exceeds the shared capacity. Existing descriptions remain resolvable after overflow.

resolve(indices, *, kind='task')[source]#

Decode a tensor of indices while preserving its nested batch shape.

Parameters:
  • indices (Tensor) – Integer task/subtask indices from a copied trajectory batch.

  • kind (Literal['task', 'subtask']) – Dictionary that produced these indices.

Return type:

Any

Returns:

A string for a scalar index, otherwise nested Python lists of strings with the same shape as indices.

Raises:
  • TypeError – If indices are not an integer tensor.

  • ValueError – If any index is unknown, including invalid padding -1.