Add segment-aware demonstration episodes

docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst

@@ -20,6 +20,7 @@ through :func:`~embodichain.lab.gym.utils.registration.make`.
 
    .. autosummary::
 
+      demo
       managers
       wrapper
 
@@ -48,6 +49,27 @@ Environment Classes
     :members:
     :exclude-members: __init__, class_type
 
+Demonstration Episodes
+----------------------
+
+The segment-aware demonstration API represents a complete task as one episode
+containing one or more semantic subtasks. Segment action iterables may be lazy,
+and the common executor records per-environment lengths, terminal status, and
+segment spans.
+
+.. autoclass:: DemoSegment
+    :members:
+
+.. autoclass:: DemoSegmentResult
+    :members:
+
+.. autoclass:: DemoEpisodeResult
+    :members:
+
+.. autofunction:: execute_demo_episode
+
+.. autofunction:: resolve_demo_segments
+
 Wrappers
 --------
 

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/cli.md

@@ -197,7 +197,7 @@ embodichain run-env --gym_config config.yaml \
 | ``--preview`` | ``False`` | Enter interactive preview mode |
 | ``--filter_visual_rand`` | ``False`` | Filter out visual randomization |
 | ``--filter_dataset_saving`` | ``False`` | Filter out dataset saving |
-| ``--max_episodes`` | *(from config)* | Override the maximum number of rollout episodes |
+| ``--max_episodes`` | *(from config)* | Override the exact number of persisted per-environment episodes; vector batches are trimmed to this count |
 | ``--record_trajectory`` | ``False`` | Record per-object kinematic trajectories during generation (for replay). Episodes auto-save to ``--trajectory_save_dir`` (or ``~/.cache/embodichain_data/trajectories/<run_id>/``) |
 | ``--trajectory_save_dir`` | ``None`` | Directory for auto-saved trajectories (default: ``~/.cache/embodichain_data/trajectories/<run_id>/``) |
 | ``--replay`` | ``False`` | Replay a recorded trajectory (``--replay_trajectory`` required; mutually exclusive with ``--preview``) |
@@ -347,6 +347,41 @@ reported separately as ``window_record_capture`` and
 
 ---
 
+(cli-preview-lerobot-data)=
+## Preview LeRobot Data
+
+Print and validate one recorded LeRobot episode without launching the
+simulator:
+
+```bash
+embodichain preview_lerobot_data \
+    outputs/lerobot/multi_segments \
+    --latest \
+    --episode 0 \
+    --expect-segments 3
+```
+
+The positional path must be an exact dataset root containing
+`meta/info.json`, unless `--latest` is used to select the newest direct child.
+`--expect-segments` is an optional exact-count assertion; it does not select,
+split, or modify segments.
+
+| Argument | Default | Description |
+|---|---|---|
+| ``dataset_root`` | *(required)* | Dataset root, or parent directory with ``--latest`` |
+| ``--episode`` | ``0`` | Episode index to inspect |
+| ``--expect-segments`` | *(unchecked)* | Fail unless the episode has exactly this many segments |
+| ``--latest`` | ``False`` | Select the newest direct child dataset |
+
+The command prints dataset format, robot, FPS, state/action shapes and ranges,
+task text, segment frame ranges, subtask descriptions, and sidecar success. It
+returns status 0 when all checks pass, 1 for a validation mismatch, and 2 when
+the path, episode, or dataset cannot be loaded. For the complete validation
+contract and a comparison with LeRobot's official Rerun visualization, see
+{ref}`Inspect Recorded LeRobot Data <tutorial_data_generation_preview>`.
+
+---
+
 ## Train RL
 
 Launch reinforcement learning training from a JSON or YAML config file.

docs/source/guides/configuration.md

@@ -129,6 +129,12 @@ For RL training and data generation, EmbodiChain uses file-based configs (`.json
 
 Configs are loaded with `embodichain.utils.utility.load_config`, which selects the parser from the file extension. Both formats produce the same in-memory dictionary and are passed to `config_to_cfg()` for environment setup.
 
+For offline expert generation, `max_episodes` counts persisted
+per-environment episodes rather than vector batches. Thus `num_envs: 4` and
+`max_episodes: 10` produce two full four-row commits plus a final two-row
+commit. Failed rows count only when the relevant `DatasetFunctorCfg` sets
+`save_failed_episodes: true`.
+
 Example paths in the repository:
 
 | Use case | JSON example | YAML example |
@@ -257,6 +263,27 @@ forwarding unless the service is behind an authenticated gateway. See
 [Browser visualization with Viser](../overview/sim/viser_visualization.md) for
 the full schema, supported scene content, and deformable-object behavior.
 
+### Robot Preset Configs
+
+Use `class_type` to select a `RobotCfg` subclass from
+`embodichain.lab.sim.robots`. Subclass-specific fields remain in the robot
+configuration and are passed to its `from_dict()` method. For example, this
+selects the canonical UR preset and then specifies the UR5 variant:
+
+```json
+{
+    "robot": {
+        "class_type": "URRobot",
+        "robot_type": "ur5",
+        "uid": "Manipulator"
+    }
+}
+```
+
+For backward compatibility, existing configs may continue to use
+`"robot_type": "CobotMagic"` as the preset-class selector when the selected
+class has no separate variant field.
+
 ### RL Training Config (`train_config.json` / `train_config.yaml`)
 
 ```json