Skip to content

Add outcome-aware dataset fragments - #547

Open
yuecideng wants to merge 2 commits into
mainfrom
codex/segment-outcome-resume
Open

Add outcome-aware dataset fragments#547
yuecideng wants to merge 2 commits into
mainfrom
codex/segment-outcome-resume

Conversation

@yuecideng

@yuecideng yuecideng commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

Description

This PR adds stable per-segment outcome and provenance annotations to Task Program demonstration execution and lets collectors persist eligible program segments as independent LeRobot episodes.

The change preserves successful work from partially failed Task Programs without conflating segment quality with episode termination. The storage design follows the current LeRobot Dataset v3 episode/task model, while using EmbodiChain-owned names such as segment_accepted instead of overloading policy/evaluation fields such as next.success.

Key changes:

  • add DemoExecutionCfg with continuous and segment_fragments persistence modes;
  • classify each segment row by its first authoritative failure phase and record dense attempt, acceptance, and continuity annotations;
  • split accepted segments into independent episodes in synchronous and asynchronous LeRobot recorders, with failed-fragment persistence as an explicit opt-in;
  • make fragment commits recorder-locally idempotent by stable fragment_id, preserving earlier commits and preventing duplicate writes after partial post-commit failure;
  • prevent online segment/boundary sampling from selecting rejected frames or crossing continuity boundaries;
  • propagate retry attempt IDs through online generation and the CLI collection path;
  • preserve the single-reset commit_env_ids behavior from main for partial vector batches;
  • update Task Program design/context, public API docs, and focused regression coverage.

Checkpoint capture, restore, and resume are intentionally deferred until an authoritative restore port exists. This PR records continuity_id for forward compatibility, with current executions remaining in continuity region zero.

Design: Task Program segment outcome and resume plan

LeRobot references: Datasets v3 overview, current dataset API

Dependencies: None. The existing LeRobot version constraint is unchanged.

Issue: None.

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • Enhancement (non-breaking change which improves an existing functionality)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (existing functionality will not work without user modification)
  • Documentation update

Screenshots

Not applicable; this change has no visual UI.

Validation

  • black --check --diff --color . — 824 files unchanged
  • python docs/scripts/check_api_docs.py — 1675/1675 exports documented
  • python -m pytest -q tests/gym/envs tests/data_pipeline tests/lab/scripts/test_run_env.py tests/gym/utils/test_gym_utils.py -m "not requires_sim and not gpu" --disable-warnings --maxfail=20 — 754 passed, 1 skipped, 23 deselected
  • python -m pytest -q tests/test_agent_context_map.py --disable-warnings — 7 passed
  • python -m pytest -q tests/docs --confcutdir=tests/docs --disable-warnings — 20 passed
  • git diff main...HEAD --check — passed

The full simulator/GPU suite was not run locally; those markers were excluded from the proportional affected-area test run.

Checklist

  • I have run the project formatter check.
  • I have made corresponding changes to the documentation.
  • Public API changes are reflected in the API docs (python docs/scripts/check_api_docs.py).
  • I have added tests that prove the feature works.
  • Dependencies have been updated, if applicable (not applicable; no dependency changes).

Add stable per-segment outcomes and dense provenance annotations, persist accepted segments as independent LeRobot episodes, and constrain online sampling by acceptance and continuity. Checkpoint capture and resume remain intentionally deferred.
@yuecideng yuecideng added enhancement New feature or request dataset gym robot learning env and its related features data Related to data_pipeline module labels Aug 24, 2026
Resolve the Task Program migration while preserving segment outcome, fragment persistence, and retry semantics.
@yuecideng
yuecideng marked this pull request as ready for review August 31, 2026 10:42
@greptile-apps

greptile-apps Bot commented Aug 31, 2026

Copy link
Copy Markdown

Greptile Summary

This PR adds outcome-aware demonstration metadata and optional persistence of accepted Task Program segments as independent LeRobot episodes.

  • Records dense segment acceptance, attempt, and continuity annotations.
  • Adds synchronous and asynchronous fragment slicing with recorder-local idempotency.
  • Filters online samples by acceptance and causal continuity.
  • Propagates retry attempt identifiers through online and CLI collection paths.
  • Updates design documentation, public API references, and regression tests.

Confidence Score: 5/5

The PR appears safe to merge; no concrete blocking or independently actionable non-blocking issue was identified.

The changed execution, annotation, sampling, vector-row persistence, and recorder idempotency paths remain internally coordinated, and the investigated failure scenarios were contradicted or lacked a reachable erroneous behavior.

Important Files Changed

Filename Overview
embodichain/lab/gym/envs/demo.py Adds execution configuration, outcome classification, attempt and continuity metadata, and fragment-count result semantics.
embodichain/lab/gym/envs/embodied_env.py Writes dense outcome annotations and makes fragment-bearing rows eligible for dataset persistence.
embodichain/lab/gym/envs/managers/datasets.py Splits eligible segment spans into independent payloads and adds recorder-local idempotent fragment commits.
embodichain/lab/gym/envs/managers/async_datasets.py Expands cloned rows into fragment payloads while retaining single-worker ordered persistence and deferred error reporting.
embodichain/data_pipeline/engine/data.py Propagates attempt IDs and prevents sampling rejected segments or windows crossing continuity boundaries.
embodichain/lab/scripts/run_env.py Integrates fragment persistence into collection retries and explicit commit-or-abort handling.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Execute demonstration] --> B[Record per-frame segment metadata]
  B --> C{Persistence mode}
  C -->|continuous| D[Commit eligible complete episode]
  C -->|segment_fragments| E[Slice eligible segment spans]
  E --> F[Generate stable fragment IDs]
  F --> G[Deduplicate recorder-local commits]
  G --> H[Commit independent LeRobot episodes]
  B --> I[Publish completed online rollout]
  I --> J[Filter windows by acceptance and continuity]
Loading

Reviews (1): Last reviewed commit: "Merge main into outcome-aware dataset fr..." | Re-trigger Greptile

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

data Related to data_pipeline module dataset enhancement New feature or request gym robot learning env and its related features

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant