embodichain.data_pipeline.engine

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>)

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

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>)
`OnlineDataEngineState`	Lifecycle states for `OnlineDataEngine`.

Exceptions:

OnlineDataWorkerError

Fallback error for a worker exception that cannot be reconstructed.

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)
`sample_batch`(batch_size, chunk_size[, ...])	Sample a batch of trajectory chunks from the shared rollout buffer.
`start`()	Start the worker and block until its first fill completes.
`stop`()	Terminate the simulation subprocess and release resources.

Attributes:

`is_init`	Whether the engine is ready to serve initialized data.
`state`	Return the engine's current lifecycle state.

__init__(cfg)[source]#

property is_init: bool#: Whether the engine is ready to serve initialized data.

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

Raises:

ValueError – If an argument is invalid.
RuntimeError – If no unlocked valid window satisfies the policy.

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 briefly for the subprocess to exit gracefully (it checks the signal between rollout steps). If the subprocess is still alive after the grace period it is force-terminated.

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

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>)

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.
`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.
`gym_config`	Gym environment configuration dictionary (already loaded, not a file path).
`initialization_timeout`	Maximum seconds to wait for the worker's initial buffer fill.
`max_episode_steps`	Maximum number of timesteps per episode.
`max_generation_attempts`	Maximum planning/execution attempts for each buffer write transaction.
`refill_threshold`	Total number of samples (refill_threshold * buffer_size) drawn from the shared buffer before a refill is triggered.
`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>)#

action_config: dict#: Action configuration dictionary. The contents depend on the specific environment and robot being used.

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.

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.

initialization_timeout: float#: Maximum seconds to wait for the worker’s initial buffer fill.

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_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.

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:

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

Bases: RuntimeError

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>)
`OnlineDataEngineState`	Lifecycle states for `OnlineDataEngine`.

Exceptions:

OnlineDataWorkerError

Fallback error for a worker exception that cannot be reconstructed.

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)
`sample_batch`(batch_size, chunk_size[, ...])	Sample a batch of trajectory chunks from the shared rollout buffer.
`start`()	Start the worker and block until its first fill completes.
`stop`()	Terminate the simulation subprocess and release resources.

Attributes:

`is_init`	Whether the engine is ready to serve initialized data.
`state`	Return the engine's current lifecycle state.

__init__(cfg)[source]#

property is_init: bool#: Whether the engine is ready to serve initialized data.

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

Raises:

ValueError – If an argument is invalid.
RuntimeError – If no unlocked valid window satisfies the policy.

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 briefly for the subprocess to exit gracefully (it checks the signal between rollout steps). If the subprocess is still alive after the grace period it is force-terminated.

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

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>)

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.
`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.
`gym_config`	Gym environment configuration dictionary (already loaded, not a file path).
`initialization_timeout`	Maximum seconds to wait for the worker's initial buffer fill.
`max_episode_steps`	Maximum number of timesteps per episode.
`max_generation_attempts`	Maximum planning/execution attempts for each buffer write transaction.
`refill_threshold`	Total number of samples (refill_threshold * buffer_size) drawn from the shared buffer before a refill is triggered.
`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>)#

action_config: dict#: Action configuration dictionary. The contents depend on the specific environment and robot being used.

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.

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.

initialization_timeout: float#: Maximum seconds to wait for the worker’s initial buffer fill.

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_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.

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.

embodichain.data_pipeline.engine

Contents

embodichain.data_pipeline.engine#

Overview#