embodichain.data_pipeline.engine

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

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

class embodichain.data_pipeline.engine.OnlineDataEngine

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.

Initialisation barrier

The is_init property returns False until the subprocess completes the very first full buffer fill, after which it becomes permanently True. Training code should wait on this flag before calling sample_batch() to avoid drawing all-zero data.

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.

is_init

True once the buffer has been populated at least once.

Methods:

__init__(cfg)
sample_batch(batch_size, chunk_size)	Sample a batch of trajectory chunks from the shared rollout buffer.
start()
stop()	Terminate the simulation subprocess and release resources.

Attributes:

buffer_size
is_init	Whether the shared buffer has been fully populated at least once.
shared_buffer

__init__(cfg)

buffer_size: int

property is_init: bool

Whether the shared buffer has been fully populated at least once.

Returns True after the simulation subprocess completes its first full buffer fill, False while that initial fill is still in progress. Callers that must not sample stale (all-zero) data can poll or block on this property before entering their training loop:

while not engine.is_init:
    time.sleep(0.5)

Returns:

True once the buffer contains valid trajectory data.

sample_batch(batch_size, chunk_size)

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

Randomly draws batch_size environment trajectories from the portion of the buffer that has been written at least once, skipping any rows currently being overwritten by the simulation subprocess. For each selected trajectory a contiguous window of chunk_size timesteps is chosen at a uniformly random offset.

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.

Return type:

TensorDict

Returns:

TensorDict with batch size [batch_size, chunk_size].

Raises:

ValueError – If chunk_size exceeds max_episode_steps.

shared_buffer: TensorDict

start()

Return type:

None

stop()

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

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

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).
max_episode_steps	Maximum number of timesteps per episode.
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>)

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.

max_episode_steps: int

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

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.

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

class embodichain.data_pipeline.engine.data.OnlineDataEngine

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.

Initialisation barrier

The is_init property returns False until the subprocess completes the very first full buffer fill, after which it becomes permanently True. Training code should wait on this flag before calling sample_batch() to avoid drawing all-zero data.

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.

is_init

True once the buffer has been populated at least once.

Methods:

__init__(cfg)
sample_batch(batch_size, chunk_size)	Sample a batch of trajectory chunks from the shared rollout buffer.
start()
stop()	Terminate the simulation subprocess and release resources.

Attributes:

buffer_size
is_init	Whether the shared buffer has been fully populated at least once.
shared_buffer

__init__(cfg)

buffer_size: int

property is_init: bool

Whether the shared buffer has been fully populated at least once.

Returns True after the simulation subprocess completes its first full buffer fill, False while that initial fill is still in progress. Callers that must not sample stale (all-zero) data can poll or block on this property before entering their training loop:

while not engine.is_init:
    time.sleep(0.5)

Returns:

True once the buffer contains valid trajectory data.

sample_batch(batch_size, chunk_size)

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

Randomly draws batch_size environment trajectories from the portion of the buffer that has been written at least once, skipping any rows currently being overwritten by the simulation subprocess. For each selected trajectory a contiguous window of chunk_size timesteps is chosen at a uniformly random offset.

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.

Return type:

TensorDict

Returns:

TensorDict with batch size [batch_size, chunk_size].

Raises:

ValueError – If chunk_size exceeds max_episode_steps.

shared_buffer: TensorDict

start()

Return type:

None

stop()

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

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

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).
max_episode_steps	Maximum number of timesteps per episode.
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>)

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.

max_episode_steps: int

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

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.