Add segment-aware demonstration episodes - #460
Open
yuecideng wants to merge 8 commits into
Open
Conversation
yuecideng
added
enhancement
# Conflicts: # embodichain/lab/scripts/run_env.py # tests/data_pipeline/test_online_data.py # tests/lab/scripts/test_run_env.py
Greptile SummaryThe PR introduces segment-aware expert demonstrations and propagates episode, segment, boundary, validity, and terminal metadata through execution, online sampling, replay, and LeRobot recording.
Confidence Score: 5/5The PR appears safe to merge. No blocking failure remains within the eligible follow-up-review scope.
|
| Filename | Overview |
|---|---|
| embodichain/lab/gym/envs/demo.py | Introduces the shared segment-aware demonstration contract and tracks per-environment completion and terminal outcomes. |
| embodichain/data_pipeline/engine/data.py | Adds transactional online rollout generation, validity-aware sampling, bounded initialization and generation, and explicit worker lifecycle handling. |
| embodichain/data_pipeline/datasets/online_data.py | Exposes episode, segment, and boundary sampling modes through the iterable dataset. |
| embodichain/lab/gym/envs/managers/datasets.py | Persists frame-level segment annotations and episode metadata through the synchronous LeRobot recorder. |
| embodichain/lab/gym/envs/managers/async_datasets.py | Extends asynchronous recording so cloned episode jobs retain segment metadata. |
| embodichain/lab/scripts/run_env.py | Migrates offline generation to the shared executor and makes reset the explicit episode commit or discard boundary. |
Sequence Diagram
sequenceDiagram
participant Task
participant Executor
participant Env
participant Recorder
participant ODS
Task->>Executor: create_demo_segments()
loop Each segment and action
Executor->>Env: step(action or hold)
Env-->>Executor: terminated/truncated per environment
Executor->>Recorder: record frame and segment annotations
end
alt Complete and all environments successful
Executor->>Env: "reset(save_data=true)"
Env->>Recorder: commit complete episode
Recorder-->>ODS: valid annotated trajectory
else Invalid, truncated, or failed
Executor->>Env: "reset(save_data=false)"
Env->>Recorder: discard pending episode
end
Reviews (6): Last reviewed commit: "wip" | Re-trigger Greptile
Contributor
There was a problem hiding this comment.
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 onreset(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_numis not being enforced here:num_framesusesmax(rgb.shape[0], max_env_num), which will never cap the number of environments rendered whenrgb.shape[0] > max_env_num. This can accidentally record/merge far more env frames than intended (and increases CPU/GPU copy + video size). Usemin(...)so the recorder respectsmax_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.
Contributor
There was a problem hiding this comment.
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
forkserveris not available on all supported platforms (e.g., Windows), so this can make the test suite fail depending on the runner OS. Preferspawnhere or select the start method conditionally (e.g., fall back tospawnwhenforkserveris unavailable viamp.get_all_start_methods()).
engine._mp_ctx = mp.get_context("forkserver")
embodichain/lab/gym/envs/managers/record.py:139
env_idsis accepted but ignored in the synchronousrecord_camera_dataimplementation, while callers now pass per-resetenv_ids. This is confusing to API consumers and makes it unclear whether partial resets are supported; either (mandatory) implementenv_idssemantics (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 thatenv_idsmust beNone/all envs).
def save_and_clear(self, env_ids: Union[torch.Tensor, None] = None) -> None:
embodichain/lab/gym/envs/managers/record.py:157
env_idsis accepted but ignored in the synchronousrecord_camera_dataimplementation, while callers now pass per-resetenv_ids. This is confusing to API consumers and makes it unclear whether partial resets are supported; either (mandatory) implementenv_idssemantics (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 thatenv_idsmust beNone/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()
Contributor
There was a problem hiding this comment.
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
forkserverstart method, which is unavailable on some platforms (notably Windows). Use a supported start method fallback (e.g., preferforkserverwhen available, elsespawn) 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 potentiallysegment_windows) viaunfold(...).all(...)across the entire buffer on everysample_batchcall while holding the shared lock. For largebuffer_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
Comment on lines
+215
to
+223
| def _normalize_env_ids(self, env_ids: Union[torch.Tensor, None]) -> list[int]: | ||
| """Return recorder-local environment IDs for a transaction boundary.""" | ||
| if env_ids is None: | ||
| return list(range(self._num_envs)) | ||
| if isinstance(env_ids, torch.Tensor): | ||
| values = env_ids.reshape(-1).cpu().tolist() | ||
| else: | ||
| values = list(env_ids) | ||
| return [int(env_id) for env_id in values if int(env_id) < self._num_envs] |
Comment on lines
+630
to
+634
| success_source = ( | ||
| success_fn() | ||
| if success_fn is not None | ||
| else last_info.get("success", True) | ||
| ) |
Comment on lines
+994
to
1001
| if annotations is not None: | ||
| for annotation_key, feature_key in DEMO_FRAME_FEATURES.items(): | ||
| if annotation_key not in annotations: | ||
| continue | ||
| value = torch.as_tensor(annotations[annotation_key]).item() | ||
| frame[feature_key] = torch.tensor([int(value)], dtype=torch.int64) | ||
|
|
||
| return frame |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 27 out of 27 changed files in this pull request and generated no new comments.
Suppressed comments (1)
embodichain/lab/gym/envs/managers/record.py:224
- _normalize_env_ids() filters env_ids only by
< self._num_envs. Negative env IDs (e.g. -1) will pass this check and later index from the end of_frames_list/_committed_env_episodes, corrupting the wrong recorder row.
Filter to 0 <= env_id < self._num_envs before returning.
def _normalize_env_ids(self, env_ids: Union[torch.Tensor, None]) -> list[int]:
"""Return recorder-local environment IDs for a transaction boundary."""
if env_ids is None:
return list(range(self._num_envs))
if isinstance(env_ids, torch.Tensor):
values = env_ids.reshape(-1).cpu().tolist()
else:
values = list(env_ids)
return [int(env_id) for env_id in values if int(env_id) < self._num_envs]
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 33 out of 33 changed files in this pull request and generated no new comments.
Suppressed comments (2)
embodichain/lab/gym/envs/managers/datasets.py:347
- If a caller saves episodes without providing rollout-buffer annotations (i.e.,
annotations=None), frames currently have novalidfield even though the LeRobot schema includes segment/terminal annotations. Adding a defaultvalid=Truekeeps the per-frame annotation contract consistent and avoids consumers having to special-case the absence ofannotation.valid.
frame_annotations = {
"episode_step": frame_index,
"segment_id": 0,
"segment_step": frame_index,
"segment_start": frame_index == 0,
embodichain/lab/gym/envs/managers/datasets.py:76
- DEMO_ANNOTATION_KEYS includes
valid, and LeRobotRecorder collectsvalidfrom the rollout buffer, but DEMO_FRAME_FEATURES does not map it into a LeRobotannotation.*feature. This makes thevalidannotation silently dropped (and the collected tensor effectively unused). Consider addingvalid -> annotation.validso the schema matches the executor/buffer contract (or stop collectingvalidif intentionally omitted).
This issue also appears on line 343 of the same file.
DEMO_FRAME_FEATURES = {
"episode_step": "annotation.episode_step",
"segment_id": "annotation.segment_id",
"segment_step": "annotation.segment_step",
"segment_start": "annotation.segment_start",
12 tasks
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 40 out of 40 changed files in this pull request and generated no new comments.
Suppressed comments (1)
embodichain/lab/gym/envs/demo.py:370
- publish_active_mask() builds
_demo_active_maskusingdevice = getattr(previous, "device", None)and then passes that device intotorch.tensor(...). When there is no existing mask and the env does not implement_set_demo_active_mask, this can passdevice=None(which is not accepted on some Torch builds) and also ignorestarget.devicewhen present. Prefer selectingtarget.deviceas a fallback and omit thedevice=argument when unknown.
previous = getattr(target, "_demo_active_mask", None)
device = getattr(previous, "device", None)
setattr(
target,
"_demo_active_mask",
torch.tensor(active, dtype=torch.bool, device=device),
)
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 40 out of 40 changed files in this pull request and generated no new comments.
Suppressed comments (5)
embodichain/lab/scripts/run_env.py:733
BaseException.add_note()is only available on Python 3.11+, but the package declaresrequires-python >=3.10. This will raiseAttributeErrorwhen cleanup fails while unwinding another exception, potentially hiding the original error. Guard the call and fall back to logging when notes are unavailable.
body_error.add_note(
"Environment cleanup also failed: "
f"{type(cleanup_error).__name__}: {cleanup_error}"
)
embodichain/lab/scripts/run_env.py:657
BaseException.add_note()is only available on Python 3.11+, but the project supports Python >=3.10. Calling it here will raiseAttributeErrorduring cleanup and can mask the real close failure. Prefer exception chaining (or guard withhasattr).
This issue also appears on line 730 of the same file.
if abort_error is not None:
close_error.add_note(
"Pending episode abort also failed: "
f"{type(abort_error).__name__}: {abort_error}"
)
embodichain/lab/scripts/preview_lerobot_data.py:151
_read_episode_sidecar()reads the entiremeta/embodichain_episodes.jsonlinto memory viaread_text().splitlines(). For large datasets this can be very slow and memory-heavy; iterating the file line-by-line avoids loading the full sidecar.
for line in sidecar_path.read_text(encoding="utf-8").splitlines():
if not line.strip():
continue
record = json.loads(line)
record_index = record.get("lerobot_episode_index", record.get("episode_index"))
embodichain/lab/gym/envs/managers/datasets.py:452
BaseException.add_note()is only available on Python 3.11+, but the project supports Python >=3.10. If depth abort also fails on Python 3.10, this will raiseAttributeErrorand hide the original save failure. Guardadd_note()and fall back to logging.
error.add_note(
"Depth sidecar abort also failed: "
f"{type(abort_error).__name__}: {abort_error}"
)
embodichain/main.py:78
- CLI command names in
COMMANDSappear to consistently use kebab-case (e.g.,run-env,preview-asset,preview-scene), but the new command is registered aspreview_lerobot_datawith underscores. This is likely to be surprising/inconsistent for users and makes autocomplete harder. Consider renaming topreview-lerobot-dataand updating docs/tests accordingly.
Command(
name="preview_lerobot_data",
target="embodichain.lab.scripts.preview_lerobot_data:cli",
help="Print and validate a recorded LeRobot dataset episode.",
),
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
DemoSegment, structured episode/segment results, and one shared executor used byrun-envand the ODS simulation worker;meta/embodichain_episodes.jsonl, including cloned metadata in the async recorder;episode,segment, andboundarychunk policies;Existing tasks implementing
create_demo_action_list()remain compatible as one legacy segment. Direct callers that passnum_traj > 1togenerate_functionmust migrate that logic intocreate_demo_segments().No new dependencies are required.
Fixes: N/A
Type of change
generate_function(num_traj > 1)callersScreenshots
N/A — runtime/data-model change with no visual UI.
Validation
black==26.3.1on all changed Python filespytest -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 skippedmake -C docs html— succeededThe 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
black .equivalent on all changed Python files.