Add segment-aware demonstration episodes

[yuecideng]

Description

Adds a segment-aware expert demonstration contract so one Gym episode can contain multiple semantic subtrajectories, such as several pick-and-place operations, while preserving legacy single-action-list tasks.

The change:

• introduces DemoSegment, structured episode/segment results, and one shared executor used by run-env and the ODS simulation worker;
• makes reset the explicit commit/discard boundary, checks termination after every action, and bounds failed generation retries;
• records independent per-env lengths plus validity, segment, and terminal annotations in expert rollout buffers;
• persists LeRobot frame annotations, dynamic segment instructions, and meta/embodichain_episodes.jsonl, including cloned metadata in the async recorder;
• makes ODS sampling validity-aware and adds episode, segment, and boundary chunk policies;
• documents the task-owned segment API and ODS schema.

Existing tasks implementing create_demo_action_list() remain compatible as one legacy segment. Direct callers that pass num_traj > 1 to generate_function must migrate that logic into create_demo_segments().

No new dependencies are required.

Fixes: N/A

Type of change

• New feature (non-breaking for existing task implementations)
• Breaking behavior for direct generate_function(num_traj > 1) callers
• Documentation update

Screenshots

N/A — runtime/data-model change with no visual UI.

Validation

• black==26.3.1 on all changed Python files
• pytest -q tests/gym/envs tests/gym/utils/test_gym_utils.py tests/data_pipeline/test_online_data.py tests/lab/scripts/test_run_env.py — 291 passed, 3 skipped
• Focused demo/ODS/LeRobot/async/replay regression suite — 87 passed
• Final affected-path regression suite — 66 passed, 1 skipped
• make -C docs html — succeeded

The repository-wide strict Sphinx build remains blocked by 671 existing warnings; the normal HTML build succeeds and the changed Markdown pages produced no path-specific warnings.

Checklist

• I have run the black . equivalent on all changed Python files.
• I have made corresponding changes to the documentation.
• I have added tests that prove the feature works.
• Dependencies have been updated, if applicable (no dependency changes required).

[greptile-apps]

Greptile Summary

The PR introduces segment-aware expert demonstrations and carries their lifecycle and annotations through execution, online sampling, and dataset persistence.

• Adds a shared segment-aware demonstration executor with bounded retries and transactional reset behavior.
• Tracks independent vector-environment lengths, validity, segment boundaries, and terminal state.
• Adds episode-, segment-, and boundary-aware online-data sampling.
• Persists segment annotations and metadata through synchronous and asynchronous LeRobot recorders.
• Strengthens worker and recorder lifecycle error propagation and documents the new contracts.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
embodichain/lab/gym/envs/demo.py	Defines the segment contract and central executor, including per-environment completion, validation, retry, and terminal-state handling.
embodichain/lab/gym/envs/embodied_env.py	Adds independent expert-rollout cursors, segment annotations, inactive-row action masking, and transactional recording hooks.
embodichain/data_pipeline/engine/data.py	Adds validity-aware segment sampling, bounded rollout generation, and explicit producer lifecycle and error propagation.
embodichain/data_pipeline/datasets/online_data.py	Exposes episode, segment, and boundary sampling policies through OnlineDataset.
embodichain/lab/gym/envs/managers/datasets.py	Persists per-frame demonstration annotations and episode metadata while making finalization a durability barrier.
embodichain/lab/gym/envs/managers/async_datasets.py	Clones segment annotations and metadata before asynchronous FIFO persistence and reports background failures during finalization.
embodichain/lab/scripts/run_env.py	Migrates offline demonstration rollout to the shared segment-aware executor and explicit commit-or-discard resets.

Sequence Diagram

sequenceDiagram
    participant Task
    participant Executor
    participant Env
    participant Buffer
    participant Recorder
    Task->>Executor: create_demo_segments()
    loop Each segment and action
        Executor->>Env: step(masked action)
        Env->>Buffer: write valid frame and annotations
        Env-->>Executor: terminated, truncated, info
    end
    alt Complete successful episode
        Executor->>Env: "reset(save_data=true)"
        Env->>Recorder: commit episode and metadata
    else Failed or incomplete attempt
        Executor->>Env: "reset(save_data=false)"
        Env->>Buffer: discard pending episode
    end

Loading

_{Reviews (2): Last reviewed commit: "test(data): tolerate spawned consumer st..." | Re-trigger Greptile}

[Copilot]

Pull request overview

Adds a segment-aware expert demonstration contract so a single Gym episode can contain multiple semantic sub-trajectories (“segments”), while keeping legacy create_demo_action_list() tasks compatible. This integrates a shared episode executor across run-env, dataset recorders (LeRobot + sidecars), and the OnlineDataEngine shared-buffer worker, with reset becoming the explicit commit/discard boundary.

Changes:

• Introduces segment-aware demo types (DemoSegment, DemoEpisodeResult, etc.) and a common executor that supports lazy segment planning, vector-env staggered completion, and per-frame annotations.
• Makes offline generation and online shared-buffer filling transactional (commit on reset(), discard on reset(save_data=False)), adds bounded retry, and propagates durability/flush failures.
• Extends recording/sampling schema to include valid + segment/terminal annotations, updates LeRobot exports + metadata sidecar, and documents new APIs/sampling modes.

Reviewed changes

Copilot reviewed 27 out of 27 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
tests/lab/scripts/test_run_env.py	Expands CLI runner tests for transactional generation, replay ownership, and cleanup semantics.
tests/gym/envs/test_replay.py	Updates replay close behavior expectation (no implicit autosave on close).
tests/gym/envs/test_demo.py	New comprehensive tests for segment-aware demo execution and annotation behavior.
tests/gym/envs/tasks/test_stay_still_save.py	Adjusts registered time limit to avoid truncation on a 100-step expert plan.
tests/gym/envs/managers/test_dataset_manager.py	Adds finalize idempotency/error aggregation and commit/discard camera/trajectory tests.
tests/gym/envs/managers/test_dataset_functors.py	Adds LeRobot finalize idempotency, frame annotation, and JSONL sidecar tests.
tests/gym/envs/managers/test_async_dataset_functors.py	Extends async recorder tests for cloned annotations/metadata and finalize failure aggregation.
tests/data_pipeline/test_online_data.py	Adds engine lifecycle/state/error-channel tests and segment/boundary sampling tests.
tests/data_pipeline/depth_video/test_writer.py	Ensures depth encoder close failures surface as durability errors.
embodichain/lab/scripts/run_env.py	Switches to segment-aware executor; makes reset the commit boundary; improves replay/cleanup semantics.
embodichain/lab/gym/utils/gym_utils.py	Extends rollout buffer schema with validity/segment/terminal annotation fields.
embodichain/lab/gym/envs/managers/record.py	Makes camera recorders transactional + idempotent finalize/close; async recorder supports explicit per-env commit queues.
embodichain/lab/gym/envs/managers/datasets.py	LeRobotRecorder now persists per-frame annotations + episode JSONL sidecar and has idempotent durability-aware finalize.
embodichain/lab/gym/envs/managers/dataset_manager.py	Makes dataset functor finalization idempotent and aggregates failures across functors.
embodichain/lab/gym/envs/managers/async_datasets.py	Async LeRobot recorder now clones annotations + metadata and aggregates background failures at finalize barrier.
embodichain/lab/gym/envs/embodied_env.py	Adds per-env rollout cursors, segment metadata hooks, transactional discard, and safe masking for staggered vector demos.
embodichain/lab/gym/envs/demo.py	New segment-aware demonstration protocol + executor implementation.
embodichain/lab/gym/envs/base_env.py	Disables auto-reset during demo execution (similar to replay).
embodichain/lab/gym/envs/init.py	Exposes demo protocol symbols from envs package.
embodichain/data_pipeline/engine/data.py	Adds explicit lifecycle states, worker error broadcast channel, transactional writes, and segment-aware sampling modes.
embodichain/data_pipeline/engine/init.py	Re-exports new engine state/error types.
embodichain/data_pipeline/depth_video/writer.py	Treats depth sidecar close/meta failures as raised durability errors.
embodichain/data_pipeline/datasets/online_data.py	Forwards segment/boundary sampling modes from OnlineDataset into OnlineDataEngine.
embodichain_tasks/embodichain_tasks/special/stay_still_save.py	Updates env time limit rationale for segment-aware executor commit behavior.
docs/source/overview/gym/dataset_functors.md	Documents finalize as a durability barrier (no implicit commit) and error surfacing semantics.
docs/source/guides/run_env.md	Documents segment API, retry/commit behavior, and annotation outputs.
docs/source/features/online_data.md	Documents transactional rows, lifecycle states, validity masks, and segment/boundary sampling.

Suppressed comments (1)

embodichain/lab/gym/envs/managers/record.py:198

• max_env_num is not being enforced here: num_frames uses max(rgb.shape[0], max_env_num), which will never cap the number of environments rendered when rgb.shape[0] > max_env_num. This can accidentally record/merge far more env frames than intended (and increases CPU/GPU copy + video size). Use min(...) so the recorder respects max_env_num.

        rgb = data["color"]

        num_frames = max(rgb.shape[0], max_env_num)
        rgb = rgb[:num_frames]
        rgb = self._draw_frames_into_one_image(rgb)[..., :3].cpu().numpy()

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 27 out of 27 changed files in this pull request and generated no new comments.

Suppressed comments (5)

tests/data_pipeline/test_online_data.py:134

• forkserver is not available on all supported platforms (e.g., Windows), so this can make the test suite fail depending on the runner OS. Prefer spawn here or select the start method conditionally (e.g., fall back to spawn when forkserver is unavailable via mp.get_all_start_methods()).

    engine._mp_ctx = mp.get_context("forkserver")

embodichain/lab/gym/envs/managers/record.py:139

• env_ids is accepted but ignored in the synchronous record_camera_data implementation, while callers now pass per-reset env_ids. This is confusing to API consumers and makes it unclear whether partial resets are supported; either (mandatory) implement env_ids semantics (e.g., per-env frame buffers like the async recorder) or (alternative) remove the parameter and have the caller only invoke this recorder on full resets (or document/enforce that env_ids must be None/all envs).

    def save_and_clear(self, env_ids: Union[torch.Tensor, None] = None) -> None:

embodichain/lab/gym/envs/managers/record.py:157

• env_ids is accepted but ignored in the synchronous record_camera_data implementation, while callers now pass per-reset env_ids. This is confusing to API consumers and makes it unclear whether partial resets are supported; either (mandatory) implement env_ids semantics (e.g., per-env frame buffers like the async recorder) or (alternative) remove the parameter and have the caller only invoke this recorder on full resets (or document/enforce that env_ids must be None/all envs).

    def discard_and_clear(self, env_ids: Union[torch.Tensor, None] = None) -> None:
        """Discard recorded frames without creating an episode video."""
        self._frames = []

embodichain/data_pipeline/engine/data.py:991

• sample_batch() holds the producer window lock while doing relatively expensive tensor ops (unfold, window filtering, and candidate selection) and then cloning the result. This can block the producer from advancing the write window (it acquires the same lock) and may reduce throughput under load. Consider shrinking the critical section by reading (lock_start, lock_end) under the lock, releasing it during window computation, and only re-acquiring to clone (with a re-check that (lock_start, lock_end) is unchanged), or by moving window precomputation to a cheaper structure.

        with self._lock_index.get_lock():
            lock_start: int = self._lock_index[0]
            lock_end: int = self._lock_index[1]

            if "valid" in self.shared_buffer.keys():
                valid = self.shared_buffer["valid"].bool()
            else:
                # Schema-v1 buffers are one fully valid segment per row.
                valid = torch.ones(
                    self.buffer_size,
                    max_steps,
                    dtype=torch.bool,
                    device=self.shared_buffer.device,
                )

            all_rows = torch.arange(self.buffer_size, device=valid.device)
            is_locked = (all_rows >= lock_start) & (all_rows < lock_end)
            valid_windows = valid.unfold(1, chunk_size, 1).all(dim=-1)
            valid_windows[is_locked] = False

            segment_ids = self.shared_buffer.get("segment_id", None)
            if segment_ids is None:
                segment_ids = torch.zeros_like(valid, dtype=torch.int64)

embodichain/data_pipeline/engine/data.py:1032

• sample_batch() holds the producer window lock while doing relatively expensive tensor ops (unfold, window filtering, and candidate selection) and then cloning the result. This can block the producer from advancing the write window (it acquires the same lock) and may reduce throughput under load. Consider shrinking the critical section by reading (lock_start, lock_end) under the lock, releasing it during window computation, and only re-acquiring to clone (with a re-check that (lock_start, lock_end) is unchanged), or by moving window precomputation to a cheaper structure.

            result = self.shared_buffer[row_indices[:, None], time_indices].clone()

[Copilot]

Pull request overview

Copilot reviewed 27 out of 27 changed files in this pull request and generated 3 comments.

Suppressed comments (2)

tests/data_pipeline/test_online_data.py:134

• The test helper unconditionally requests the forkserver start method, which is unavailable on some platforms (notably Windows). Use a supported start method fallback (e.g., prefer forkserver when available, else spawn) so the test suite remains portable.

    engine._mp_ctx = mp.get_context("forkserver")

embodichain/data_pipeline/engine/data.py:1004

• This computes valid_windows (and potentially segment_windows) via unfold(...).all(...) across the entire buffer on every sample_batch call while holding the shared lock. For large buffer_size / max_episode_steps, this can become a major CPU bottleneck and also delay the producer from advancing the lock window. Consider a more O(buffer_size) approach (e.g., track per-row valid lengths, and for segment/boundary modes precompute per-row boundary indices or segment spans) so sampling remains fast under training load.

            all_rows = torch.arange(self.buffer_size, device=valid.device)
            is_locked = (all_rows >= lock_start) & (all_rows < lock_end)
            valid_windows = valid.unfold(1, chunk_size, 1).all(dim=-1)
            valid_windows[is_locked] = False

            segment_ids = self.shared_buffer.get("segment_id", None)
            if segment_ids is None:
                segment_ids = torch.zeros_like(valid, dtype=torch.int64)

            if sampling_mode == "segment":
                segment_windows = segment_ids.unfold(1, chunk_size, 1)
                same_segment = (segment_windows == segment_windows[..., :1]).all(
                    dim=-1
                ) & (segment_windows[..., 0] >= 0)
                valid_windows &= same_segment
            elif sampling_mode == "boundary":
                segment_windows = segment_ids.unfold(1, chunk_size, 1)
                crosses_boundary = (
                    segment_windows[..., 1:] != segment_windows[..., :-1]
                ).any(dim=-1)
                valid_windows &= crosses_boundary