Add segment-aware demonstration episodes

docs/source/features/online_data.md

@@ -22,30 +22,94 @@ These components live under `embodichain/data_pipeline/` and are designed to wor
 
 Key ideas:
 
-- **Shared buffer**: multiple producers (simulation workers) and multiple consumers (training workers) can read/write concurrently.
+- **Shared buffer**: the simulation producer and training consumers can read/write concurrently.
 - **GPU-friendly**: buffer is designed for efficient sampling and minimal copying.
 - **Chunked sampling**: training samples fixed-length or dynamically sized chunks.
+- **Transactional rows**: the current write window is locked until a complete,
+  successful episode replaces it. Failed generation attempts never become
+  sampleable.
+- **Variable lengths**: every frame has a `valid` flag. Sampling constructs
+  windows only from real frames and never reads zero padding or a stale tail.
+- **Episode-uniform sampling**: an eligible episode row is selected uniformly,
+  then a valid start offset is selected within that row. Longer episodes do
+  not receive extra probability merely because they contain more windows.
+- **Explicit lifecycle**: the engine progresses through `CREATED`, `STARTING`,
+  `READY`, and a terminal `FAILED` or `STOPPED` state. Sampling is allowed only
+  in `READY`, and a stopped or failed instance cannot be restarted.
+- **Fail-fast errors**: initial-fill timeouts, rollout failures, hard worker
+  exits, and shutdown/recorder failures are raised to the owner. The same
+  stable error snapshot is visible to every forked or spawned data consumer.
+
+Each row stores one complete task episode. A row may contain multiple semantic
+segments. In addition to observations, actions, and rewards, the shared buffer
+contains:
+
+```text
+valid, episode_step, segment_id, segment_step,
+segment_start, segment_end, terminated, truncated
+```
+
+Tasks use the same `create_demo_segments()` protocol as `run-env`; legacy
+`create_demo_action_list()` tasks are treated as one segment. The worker checks
+termination after every action and retries a failed episode up to
+`max_generation_attempts` before reporting an error.
+
+`start()` blocks until the initial fill is complete and is bounded by
+`initialization_timeout`. If the worker exits or generation fails, `start()`
+raises the worker error instead of waiting indefinitely. A later worker failure
+is reported on the next `sample_batch()` call rather than silently serving a
+permanently stale buffer.
 
 ### Minimal setup
 
 ```python
-from embodichain.data_pipeline.engine.data import OnlineDataEngine, OnlineDataEngineCfg
+from embodichain.data_pipeline.engine import OnlineDataEngine, OnlineDataEngineCfg
 
 cfg = OnlineDataEngineCfg(
-    buffer_size=2,           # number of trajectories kept in the ring buffer
-    state_dim=6,             # example state dimension
-    gym_config=your_gym_cfg, # parsed gym config for the task (JSON or YAML)
+    buffer_size=2,              # trajectories kept in the shared buffer
+    state_dim=6,                # example state dimension
+    gym_config=your_gym_cfg,    # parsed gym config for the task
+    initialization_timeout=300, # maximum seconds for the initial fill
 )
 engine = OnlineDataEngine(cfg)
 engine.start()
 ```
 
 ### Shutdown
 
+Prefer a context manager so cleanup runs on both success and failure:
+
 ```python
-engine.stop()
+with OnlineDataEngine(cfg) as engine:
+    batch = engine.sample_batch(batch_size=32, chunk_size=64)
+    train_step(batch)
+```
+
+For an explicitly managed lifecycle, call `stop()` in the process that created
+the engine:
+
+```python
+engine = OnlineDataEngine(cfg)
+engine.start()
+try:
+    train(engine)
+finally:
+    engine.stop()
 ```
 
+`stop()` is idempotent after a successful cleanup. It waits for the producer
+and raises any failure reported while the worker closes its environment or
+flushes committed data. If graceful shutdown times out and the worker requires
+`terminate()` or `kill()`, durability cannot be confirmed: the engine remains
+`FAILED` and `stop()` raises instead of reporting a false `STOPPED` state. When
+a `with` body and cleanup both fail, the body exception remains primary and the
+cleanup failure is attached as a note.
+
+Start the engine before constructing multiprocessing `DataLoader` workers.
+Forked and spawned copies may call `sample_batch()`, but only the original
+owner process may call `start()` or `stop()`; consumer destructors never signal
+the shared producer.
+
 ---
 
 ## OnlineDataset
@@ -103,6 +167,25 @@ dataset = OnlineDataset(engine, chunk_size=sampler)
 
 In batch mode, the sampler is called once per step so all trajectories in the batch share the same chunk length.
 
+### Segment-aware sampling
+
+Set `sampling_mode` according to the training objective:
+
+```python
+# Chunks may span adjacent subtasks (default).
+episode_dataset = OnlineDataset(engine, chunk_size=64, sampling_mode="episode")
+
+# Every chunk stays inside one pick/place segment.
+segment_dataset = OnlineDataset(engine, chunk_size=32, sampling_mode="segment")
+
+# Every chunk contains an internal transition between two segments.
+boundary_dataset = OnlineDataset(engine, chunk_size=32, sampling_mode="boundary")
+```
+
+All three modes still require every sampled frame to be valid. `boundary`
+requires a chunk size of at least two and raises a clear error when no internal
+boundary can satisfy the requested length.
+
 ---
 
 ## ChunkSizeSampler

docs/source/guides/run_env.md

@@ -128,9 +128,11 @@ details.
 ## Run and record
 
 Without `--preview` or `--replay`, `run-env` enters offline rollout mode. For
-each episode, it asks the task for a demonstration action list, applies the
-actions through `env.step()`, discards invalid generations, and resets the
-environment. `--max_episodes` overrides the value in the gym config:
+each episode, it asks the task for its demonstration segments, applies every
+action through `env.step()`, and commits the complete episode with one reset.
+An empty plan, truncation, or failed final task validation discards the whole
+attempt. Attempts are bounded by `demo_max_attempts` in the gym config
+(default: 3). `--max_episodes` overrides the value in the gym config:
 
 ```bash
 embodichain run-env \
@@ -145,6 +147,69 @@ Headless execution is normally preferred for throughput. Use
 `--filter_dataset_saving` for a rollout smoke test that should not create a
 structured dataset.
 
+### Multi-segment episodes
+
+An episode is the complete task; a segment is one semantic subtask inside it.
+For example, moving three objects is one episode containing three pick/place
+segments, even if each segment has its own motion trajectory. The task owns
+the number, order, and targets of those segments. The runner only manages the
+episode lifecycle, termination checks, retry, and commit/discard boundary.
+
+Existing tasks implementing `create_demo_action_list()` remain compatible and
+are recorded as a single `legacy` segment. A multi-object task can instead
+implement a lazy segment planner:
+
+```python
+from embodichain.lab.gym.envs import DemoSegment
+
+
+def create_demo_segments(self):
+    for object_uid in self.object_order:
+        # This runs after the preceding segment, so planning sees the latest
+        # scene state.
+        actions = self.plan_pick_and_place(object_uid)
+        yield DemoSegment(
+            actions=actions,
+            name="pick_and_place",
+            target_uid=object_uid,
+            instruction=f"Place {object_uid} in its target bin",
+            # Segment validation is separate from Gym episode termination.
+            validator=lambda uid=object_uid: self.is_object_placed(uid),
+        )
+```
+
+Gym `terminated` and `truncated` always describe the whole episode, never an
+individual segment. A segment normally ends when its action iterable is
+exhausted; its optional zero-argument `validator` then returns one boolean per
+parallel environment. A failed validator aborts the batch. Episode-level
+success termination stops the remaining lazy plan without requesting another
+segment.
+
+The executor checks terminal signals after every action and temporarily
+disables Gym auto-reset. Dataset recording is transactional: only the explicit
+reset after successful final validation saves the episode. Exceptions,
+interrupts, failed attempts, and closing an environment with a live rollout
+abort pending structured data, videos, and trajectories.
+
+Segment actions pass through the same action-dimension normalization used by
+legacy `create_demo_action_list()` tasks. A time-limit truncation is always an
+invalid expert rollout, including when it occurs on the planner's final
+action, so a task's `max_episode_steps` must be greater than the longest valid
+expert plan.
+
+In a vectorized environment, segments and actions remain on one shared planner
+clock, but completion is tracked independently. When one environment reports
+success, its terminal result and recording cursor become sticky; subsequent
+shared actions use a safe hold/no-op command for that row while unfinished rows
+continue. Consequently, rollout and trajectory lengths may differ by row.
+The commit remains batch-atomic: every row must eventually succeed, while any
+failure or truncation aborts the whole attempt.
+
+Failed attempts use `reset(options={"save_data": False})`. This discards
+structured datasets, replay trajectories, and camera-video buffers; recorder
+options such as `save_failed_episodes` do not override the runner's expert
+transaction policy.
+
 ### Choose the recording output you need
 
 EmbodiChain uses "recording" for three related but distinct outputs:
@@ -159,7 +224,8 @@ EmbodiChain uses "recording" for three related but distinct outputs:
   - Intended consumer
 * - Structured dataset
   - A dataset manager in the gym config.
-  - Observations, actions, task metadata, and optionally sensor videos.
+  - Observations, actions, episode/segment annotations, task metadata, and
+    optionally sensor videos.
   - Imitation-learning or data-processing pipelines.
 * - Debug or demo video
   - `record_camera_data` or `record_camera_data_async` as an interval event
@@ -196,8 +262,9 @@ embodichain run-env \
 The recorder stores the robot root pose and complete joint position, the raw
 action before ActionManager preprocessing, and the pose or joint state of scene
 rigid objects and articulations. Each environment in a vectorized rollout is
-tracked independently. A trajectory is auto-saved at episode reset and again
-on close if an unfinished buffer still contains steps.
+tracked independently. A trajectory is saved only at an explicit successful
+episode reset. `close()` is a durability barrier for already committed writes,
+not an implicit commit, so an unfinished trajectory is discarded.
 
 Files are named like `traj_env0_000000.pt`. When
 `--trajectory_save_dir` is omitted, they are written below:
@@ -207,9 +274,11 @@ ${EMBODICHAIN_DATA_ROOT:-~/.cache/embodichain_data}/trajectories/<run_id>/
 ```
 
 Each file contains `states`, `actions`, and `meta`. Metadata includes the actual
-per-environment lengths, timestep, robot identity and DOF, active joint IDs,
-recorded object IDs, and original environment IDs. The final directory is also
-printed when the run finishes.
+per-environment lengths, segment ranges and targets, timestep, robot identity
+and DOF, active joint IDs, recorded object IDs, and original environment IDs.
+LeRobot exports also contain per-frame `annotation.segment_*` fields and a
+`meta/embodichain_episodes.jsonl` sidecar with the complete segment records.
+The final directory is also printed when the run finishes.
 
 ## Replay a trajectory
 

docs/source/overview/gym/dataset_functors.md

@@ -221,6 +221,14 @@ The sync stall grows **linearly** with ``num_envs`` (each reset saves N envs ser
 ```{attention}
 - ``AsyncLeRobotRecorder`` clones each finished episode (including camera frames) to CPU before enqueuing. For very high resolutions or many envs, monitor RSS — the worker normally keeps up, but a slow disk can let the queue grow.
 - A single background worker touches the ``LeRobotDataset`` (which is not thread-safe), so episode order is preserved FIFO. Always let ``env.close()`` / ``dataset_manager.finalize()`` run so the worker drains before the dataset is finalized.
+- Finalization is a durability barrier for episodes already submitted through
+  ``mode="save"``. It does not turn a live rollout into an episode; pending
+  frames are discarded. Background and storage failures are raised to the
+  caller, and repeated finalization is safe.
+- Depth video and EmbodiChain metadata sidecars are part of the same durability
+  result: encoder-close or metadata-write failures are surfaced even when the
+  corresponding LeRobot episode has already committed, and its episode index
+  is never reused by a later queued write.
 - ``env.close()`` calls ``sim.destroy()``, which exits the process without returning to Python. In scripts that build multiple envs, run each in its own subprocess and write results before closing.
 ```
 
@@ -260,15 +268,15 @@ dataset = {
 1. **Initialization**: The Dataset Manager initializes the functor with the configured parameters
 2. **Data Collection**: During episode rollout, the functor receives observations and actions
 3. **Save Trigger**: When an episode completes, call the functor with `mode="save"`
-4. **Finalization**: After all episodes, call `finalize()` to save any remaining data
+4. **Finalization**: After all episodes, call `finalize()` to drain committed writes and finalize storage metadata
 
 ```python
 # Inside environment loop
 if episode_done:
     dataset_manager.apply(mode="save", env_ids=completed_env_ids)
 
-# After training completes
-dataset_manager.apply(mode="finalize")
+# After collection completes. This raises if a committed write failed.
+dataset_manager.finalize()
 ```
 
 ### Parallel Collection (async recorder)
@@ -293,7 +301,9 @@ dataset = {
 }
 ```
 
-The async recorder drains its background worker during ``finalize()``, so make sure ``env.close()`` (or ``dataset_manager.finalize()`) runs at the end of collection.
+The async recorder drains its background worker during ``finalize()``, so make
+sure ``env.close()`` (or ``dataset_manager.finalize()``) runs at the end of
+collection. No new episode can be submitted after finalization begins.
 
 ### Compressed depth recording
 
@@ -334,6 +344,6 @@ dataset = {
 The Dataset Manager supports the following modes:
 
 - ``save``: Save completed episodes for specified environment IDs
-- ``finalize``: Finalize the dataset and save any remaining data
+- ``finalize``: Drain explicitly committed writes and finalize dataset resources
 
 See {class}`~managers.dataset_manager.DatasetManager` for more details.

embodichain/data_pipeline/datasets/online_data.py

@@ -16,7 +16,7 @@
 
 from __future__ import annotations
 
-from typing import Callable, Iterator, List, Optional
+from typing import Callable, Iterator, List, Literal, Optional
 
 from tensordict import TensorDict
 from torch.utils.data import IterableDataset
@@ -77,6 +77,9 @@ class OnlineDataset(IterableDataset):
             ``[batch_size, chunk_size]`` (batch mode).
         transform: Optional ``(TensorDict) -> TensorDict`` applied to each
             yielded item/batch before returning.
+        sampling_mode: ``"episode"`` allows chunks to cross segment boundaries,
+            ``"segment"`` keeps chunks within one segment, and ``"boundary"``
+            samples chunks that cross a boundary.
 
     Example — fixed chunk size, item mode::
 
@@ -123,6 +126,7 @@ def __init__(
         chunk_size: int | ChunkSizeSampler,
         batch_size: Optional[int] = None,
         transform: Optional[Callable[[TensorDict], TensorDict]] = None,
+        sampling_mode: Literal["episode", "segment", "boundary"] = "episode",
     ) -> None:
         if isinstance(chunk_size, int):
             if chunk_size < 1:
@@ -131,10 +135,16 @@ def __init__(
             raise TypeError(
                 f"chunk_size must be an int or a ChunkSizeSampler, got {type(chunk_size).__name__}."
             )
+        if sampling_mode not in {"episode", "segment", "boundary"}:
+            raise ValueError(
+                "sampling_mode must be 'episode', 'segment', or 'boundary', "
+                f"got {sampling_mode!r}."
+            )
         self._engine = engine
         self._chunk_size = chunk_size
         self._batch_size = batch_size
         self._transform = transform
+        self._sampling_mode = sampling_mode
 
     # ------------------------------------------------------------------
     # Internal helpers
@@ -177,7 +187,11 @@ def __iter__(self) -> Iterator[TensorDict]:
 
             while True:
                 # Item mode: draw one trajectory and remove the outer batch dim.
-                raw = self._engine.sample_batch(batch_size=1, chunk_size=chunk_size)
+                raw = self._engine.sample_batch(
+                    batch_size=1,
+                    chunk_size=chunk_size,
+                    sampling_mode=self._sampling_mode,
+                )
                 sample: TensorDict = raw[0]
 
                 if self._transform is not None:
@@ -190,7 +204,9 @@ def __iter__(self) -> Iterator[TensorDict]:
 
             # Batch mode: draw a full pre-batched TensorDict.
             sample = self._engine.sample_batch(
-                batch_size=self._batch_size, chunk_size=chunk_size
+                batch_size=self._batch_size,
+                chunk_size=chunk_size,
+                sampling_mode=self._sampling_mode,
             )
 
             if self._transform is not None: