Add segment-aware demonstration episodes - #460
Merged
Merged
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 their lifecycle and annotations through environment execution, online sampling, recording, replay, and documentation.
Confidence Score: 5/5The PR appears safe to merge because no blocking failure remains within the eligible follow-up-review scope. No blocking failure remains.
|
| Filename | Overview |
|---|---|
| embodichain/lab/gym/envs/demo.py | Adds the shared segment-aware demonstration executor, structured results, action normalization, terminal handling, and recording callbacks. |
| embodichain/data_pipeline/engine/data.py | Reworks online rollout production, shared-buffer sampling, lifecycle management, retry bounds, and cross-process error propagation. |
| embodichain/lab/gym/envs/embodied_env.py | Adds per-environment rollout validity and segment metadata while integrating explicit demonstration recording boundaries. |
| embodichain/lab/gym/envs/managers/datasets.py | Extends dataset persistence with episode and segment annotations, sidecar metadata, and selective row commits. |
| embodichain/lab/scripts/run_env.py | Migrates offline generation to the shared executor with bounded retries and exact per-environment episode accounting. |
| embodichain/data_pipeline/datasets/online_data.py | Exposes episode-, segment-, and boundary-aware online chunk sampling through the iterable dataset. |
| embodichain/lab/gym/envs/wrapper/replay.py | Preserves segment-aware trajectory metadata through replay. |
| embodichain/lab/scripts/preview_lerobot_data.py | Adds validation and inspection tooling for recorded LeRobot episodes and segment metadata. |
Sequence Diagram
sequenceDiagram
participant Task
participant Executor
participant Env
participant Recorder
participant Sampler
Task->>Executor: create_demo_segments()
loop Each lazy segment
Executor->>Env: step(action)
Env->>Recorder: record valid frame and segment metadata
Env-->>Executor: observation, reward, terminal signals
Executor->>Executor: update per-environment completion
end
Executor-->>Env: structured episode result
alt selected episode rows are committed
Env->>Recorder: reset and commit selected rows
Recorder-->>Sampler: publish valid annotated episode rows
else attempt is discarded
Env->>Recorder: "reset(save_data=false)"
end
Sampler->>Sampler: select episode, segment, or boundary windows
Reviews (12): 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 (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.",
),
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 (2)
embodichain/lab/scripts/preview_lerobot_data.py:363
build_episode_preview()stacks the fullobservation.stateandactionmatrices just to compute shape and min/max. This duplicates the already-materializedsamplesand can spike memory for long episodes; computing shape/range incrementally avoids allocating a(frames, dim)array.
state = _feature_matrix(samples, "observation.state")
action = _feature_matrix(samples, "action")
return EpisodePreview(
dataset_root=dataset_root,
episode_index=episode_index,
codebase_version=str(info.get("codebase_version", "unknown")),
robot_type=str(info.get("robot_type", "unknown")),
fps=fps,
total_episodes=int(info.get("total_episodes", 0)),
total_frames=int(info.get("total_frames", 0)),
episode_frames=len(samples),
task=task,
state_shape=tuple(state.shape),
action_shape=tuple(action.shape),
state_range=(float(np.min(state)), float(np.max(state))),
action_range=(float(np.min(action)), float(np.max(action))),
embodichain/lab/scripts/preview_lerobot_data.py:154
_read_episode_sidecar()reads the entire JSONL sidecar into memory viaread_text().splitlines(). For large datasets this can be unnecessarily memory-heavy; iterating the file line-by-line avoids the full-file load while keeping the same behavior.
This issue also appears on line 348 of the same file.
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"))
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 (2)
embodichain/lab/gym/envs/managers/record.py:224
- record_camera_data_async._normalize_env_ids() currently allows negative env IDs (e.g., -1) because it only checks "env_id < self._num_envs". Negative indices will then alias the last env's buffers and can corrupt which episode is committed/discarded.
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]
embodichain/lab/scripts/run_env.py:163
- In generate_function(), the abort reset in the finally block can mask the original exception (e.g., KeyboardInterrupt/SystemExit/execute_demo_episode failure) if env.reset(options={"save_data": False}) raises during abort. That makes debugging and correct error reporting harder, because the abort error replaces the real failure that triggered cleanup.
# ``finally`` also covers KeyboardInterrupt, SystemExit, and
# GeneratorExit. A failed commit is aborted as well, so close()
# can never implicitly persist the pending partial episode.
if not commit_succeeded:
_abort_pending_episode(env)
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 (2)
pyproject.toml:52
- Core code now imports pandas (LeRobotRecorder writes meta/subtasks.parquet), but pandas is not declared in project dependencies. If lerobot does not install pandas, EmbodiChain installs will break at runtime when dataset recording is configured. Declare pandas as a dependency (or remove the hard pandas requirement).
"h5py",
"tensordict",
"viser==1.0.21",
"lerobot>=0.4.4,<0.5"
]
embodichain/lab/gym/envs/managers/datasets.py:181
- LeRobotRecorder always calls _initialize_dataset(), even when LEROBOT_AVAILABLE is False (e.g., lerobot or pandas import failed in the module-level try/except). That leads to a runtime crash (LeRobotDataset is unavailable) instead of a clear ImportError. Add a hard guard before initializing the dataset.
# Initialize dataset
self._initialize_dataset()
…Force/EmbodiChain into feat/segmented-demo-episodes
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 52 out of 52 changed files in this pull request and generated no new comments.
Suppressed comments (3)
embodichain/lab/scripts/run_env.py:150
- generate_function() accepts debug_mode but never uses it; the argument is also not forwarded into execute_demo_episode(), so debug_mode cannot influence task planning even if a task expects it via create_demo_segments(**kwargs). Forward debug_mode as a planning kwarg (or remove the parameter entirely) to avoid a silent no-op CLI flag.
result: DemoEpisodeResult = execute_demo_episode(
env,
episode_index=time_id,
progress=_progress_wrapper,
**kwargs,
embodichain_tasks/embodichain_tasks/tableware/stack_blocks_two.py:133
- The return type annotation
tuple[DemoSegment]describes a 1-element tuple type, not a variable-length tuple. Since create_demo_segments returns a tuple literal (and may evolve to multiple segments), the annotation should betuple[DemoSegment, ...](orIterable[DemoSegment]) to match Python typing semantics.
embodichain/lab/scripts/preview_lerobot_data.py:154 - _read_episode_sidecar() can crash when a JSONL record is missing both lerobot_episode_index and episode_index (int(None) raises TypeError), and it reads the whole JSONL into memory via read_text(). Iterating the file line-by-line and skipping malformed records makes the preview more robust for large datasets.
continue
record = json.loads(line)
record_index = record.get("lerobot_episode_index", record.get("episode_index"))
if int(record_index) == episode_index:
return record
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 55 out of 55 changed files in this pull request and generated no new comments.
Suppressed comments (1)
embodichain/lab/scripts/preview_lerobot_data.py:187
- build_episode_preview() records missing required features as an error but still proceeds to index into those features (e.g., sample["frame_index"]). When a dataset is missing any required field this will raise KeyError and turn a validation mismatch into a hard load failure (exit code 2) instead of returning a structured EpisodePreview with errors.
features = set(info.get("features", {}))
missing_features = sorted(REQUIRED_FEATURES - features)
if missing_features:
errors.append(f"Missing required features: {', '.join(missing_features)}")
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 58 out of 58 changed files in this pull request and generated no new comments.
Suppressed comments (4)
embodichain/lab/gym/utils/gym_utils.py:1570
- After dropping
"lengths"from the required meta keys (see above),load_trajectory()should also defaultlengthsto a uniform[num_steps] * num_envswhen the key is absent; otherwise replay/validation will still fail withKeyErroron legacy files.
lengths = meta["lengths"]
embodichain/lab/gym/utils/gym_utils.py:1538
load_trajectory()currently treatsmeta["lengths"]as a required key, which makes trajectories recorded by older versions (uniform-length, no per-env lengths) unreadable. If backward compatibility is intended, only requirenum_steps/num_envsand compute a defaultlengthswhen missing.
This issue also appears on line 1570 of the same file.
for key in ("num_steps", "num_envs", "lengths"):
if key not in meta:
raise ValueError(f"Trajectory meta is missing key: {key!r}")
embodichain/lab/gym/utils/gym_utils.py:1551
load_trajectory()allowsnum_steps == 0, butReplayWrapper.reset()unconditionally indexesstates[:, 0], which will crash for an empty trajectory. Reject zero-step trajectories here with a clearValueErrorso callers get an actionable error instead of an index failure later.
if num_steps < 0 or num_envs <= 0:
raise ValueError(
f"Invalid trajectory dimensions: num_envs={num_envs}, "
f"num_steps={num_steps}."
)
embodichain/main.py:78
- CLI command names in
COMMANDSare consistently kebab-case (e.g.preview-asset,run-env,workspace-cache), but the new command ispreview_lerobot_data(snake_case). For consistency and discoverability, consider renaming the command topreview-lerobot-dataand updating docs/tests accordingly (the module name can remainpreview_lerobot_data.py).
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.