From e465d9fd91fd62ad32a14632015e69c9009c8cf1 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Mon, 24 Aug 2026 18:04:55 +0800 Subject: [PATCH] feat(gym): add outcome-aware dataset fragments 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. --- .../design/segment_outcome_and_resume_plan.md | 114 +++++--- .../embodichain/embodichain.lab.gym.envs.rst | 7 + docs/source/api_reference/public_api.rst | 3 + embodichain/data_pipeline/engine/data.py | 28 +- embodichain/lab/gym/envs/demo.py | 270 +++++++++++++++++- embodichain/lab/gym/envs/embodied_env.py | 102 ++++++- .../lab/gym/envs/managers/async_datasets.py | 15 +- embodichain/lab/gym/envs/managers/datasets.py | 223 +++++++++++++-- embodichain/lab/gym/utils/gym_utils.py | 30 ++ embodichain/lab/scripts/run_env.py | 80 +++++- tests/data_pipeline/test_online_data.py | 61 ++++ tests/gym/envs/expert_program/test_bridge.py | 5 + .../managers/test_async_dataset_functors.py | 54 ++++ .../envs/managers/test_dataset_functors.py | 164 +++++++++++ .../gym/envs/managers/test_dataset_manager.py | 26 ++ tests/gym/envs/test_demo.py | 110 +++++++ tests/gym/utils/test_gym_utils.py | 4 + tests/lab/scripts/test_run_env.py | 64 ++++- 18 files changed, 1290 insertions(+), 70 deletions(-) diff --git a/docs/design/segment_outcome_and_resume_plan.md b/docs/design/segment_outcome_and_resume_plan.md index fc02908c8..c92da8b25 100644 --- a/docs/design/segment_outcome_and_resume_plan.md +++ b/docs/design/segment_outcome_and_resume_plan.md @@ -1,8 +1,9 @@ # Program-Segment Outcome and Checkpoint Resume -- Status: proposed +- Status: outcome propagation and natural segment fragments implemented; + checkpoint resume deferred - Scope: Expert Program demonstration generation in simulation -- Date: 2026-08-23 +- Date: 2026-08-24 - Related design: `docs/design/declarative_expert_program_plan.md` ## 1. Decision @@ -19,9 +20,9 @@ The current layers already have the right semantic boundary: application validators; - `DemoSegmentResult` records the resulting per-environment outcome. -The missing behavior is downstream of that decision: +The original gaps were downstream of that decision: -1. trajectory buffers and online sampling do not expose the segment outcome as +1. trajectory buffers and online sampling did not expose the segment outcome as a dense frame annotation; 2. failed rows are permanently removed from the current program run; 3. there is no typed, auditable way to restore the entry state of a later @@ -29,7 +30,7 @@ The missing behavior is downstream of that decision: 4. a discontinuous restore cannot currently be represented without making the resulting data look like one continuous episode. -The proposal therefore has two parts: +The design therefore has two parts: 1. persist the existing program-segment acceptance result as trajectory/data quality metadata; @@ -39,6 +40,12 @@ The proposal therefore has two parts: The default continuous-episode path remains fail-closed and unchanged. +Implementation scope for the current change stops after outcome propagation, +causal-continuity-aware sampling, and persistence of naturally executed +segments as independent fragments. Checkpoint values, simulator state restore, +and resume-after-failure execution remain future work. Consequently every +currently recorded frame has ``continuity_id == 0``. + ## 2. Terminology Three existing meanings of "segment" must remain separate. @@ -81,7 +88,7 @@ The public outcome should preserve both `accepted` and the first authoritative failure phase. A boolean alone is insufficient for diagnostics and recovery policy. -Recommended stable outcome kinds are: +Implemented stable outcome kinds are: ```text succeeded @@ -91,9 +98,11 @@ validation_failed cancelled truncated not_attempted -restore_failed ``` +``restore_failed`` remains reserved for the deferred checkpoint implementation; +it is not exposed by the current result type because no restore can occur. + `DemoSegmentResult.successes` remains the compatibility boolean view. Its source must stay the bridge's accepted mask rather than a new call to `is_task_success()`. @@ -108,21 +117,32 @@ ordinary Expert Program execution. ## 4. Data contract +The implementation remains compatible with EmbodiChain's current LeRobot +``>=0.4.4,<0.5`` dependency. It does not require a dataset-format upgrade: +dense qualification is stored as additive numeric ``annotation.*`` features, +the fragment's segment instruction becomes its LeRobot task, and richer +program provenance stays in EmbodiChain's JSONL sidecar. This also keeps the +program-segment acceptance meaning separate from LeRobot reward/task success. + ### 4.1 Dense frame annotations -Extend expert rollout annotations with: +Expert rollout annotations include: ```text -segment_success: bool +segment_accepted: bool segment_attempt_id: int64 continuity_id: int64 ``` -`segment_success` is filled retroactively for the segment's complete frame span +`segment_accepted` is filled retroactively for the segment's complete frame span when `_end_demo_segment_recording()` receives the terminal result. It does not replace `valid`; `valid` continues to mean that a buffer slot contains a real transition. +The name intentionally avoids collision with LeRobot's task/reward success +fields: this value means that the owning Expert Program segment passed runtime, +post-policy, and validator qualification. + `segment_attempt_id` distinguishes retries or repeated attempts of the same compiled segment occurrence. `continuity_id` increments whenever state is restored without an environment action. @@ -142,7 +162,7 @@ The online sampler's normal segment mode selects only windows for which: ```text valid == true -segment_success == true +segment_accepted == true segment_id is constant continuity_id is constant ``` @@ -283,29 +303,40 @@ written for the state jump. ## 7. Execution modes -Add a collector-owned configuration, not an Expert Program field: +The implemented collector-owned configuration, rather than an Expert Program +field, is: ```python @configclass class DemoExecutionCfg: mode: Literal["continuous", "segment_fragments"] = "continuous" - on_segment_failure: Literal["stop", "resume_next"] = "stop" save_failed_fragments: bool = False ``` -The checkpoint port is passed as a live dependency rather than serialized in -this config. +Both modes stop at the current fail-closed execution boundary. A future resume +extension should add an explicit failure policy only together with the live +checkpoint port; the current config deliberately cannot request unsupported +restore behavior. + +The mode is selected through the collection API: + +```python +generate_function( + env, + execution_cfg=DemoExecutionCfg(mode="segment_fragments"), +) +``` -`continuous + stop` is the current default. `resume_next` is accepted only in -`segment_fragments` mode and only when a checkpoint port is installed. It is -rejected for real devices unless that integration supplies an authoritative -restore implementation. +One source program row may produce multiple LeRobot episodes in fragment mode. +Accepted fragments are saved by default; ``save_failed_fragments=True`` is an +explicit diagnostic-data opt-in. The existing CLI ``max_episodes`` accounting +therefore remains on continuous mode until a fragment-count quota is defined. The existing `DemoSegment.failure_policy` continues to mean batch behavior (`batch_abort` versus `row_independent`). It must not be overloaded with cross-segment recovery semantics. -### State machine +### Future resume state machine ```text READY(k) @@ -325,30 +356,31 @@ shared segment barrier. Rows that cannot be restored remain inactive. ## 8. Public result semantics -Do not redefine `DemoEpisodeResult.completed`. Add explicit collection-level -views instead: +Do not redefine `DemoEpisodeResult.completed`. The current implementation adds +``successful_fragment_count_by_env`` and retains the existing continuous +success fields. Checkpoint-dependent views remain deferred: ```text -program_exhausted_by_env -continuous_success_by_env recovered_by_env -successful_fragment_count_by_env +program_exhausted_by_env ``` -Each segment/fragment result additionally records: +Each current segment/fragment result additionally records: ```text attempt_id continuity_id outcome_kind -resumed_from_checkpoint -checkpoint_id ``` -This makes the following cases distinguishable: +``resumed_from_checkpoint`` and ``checkpoint_id`` are deferred with resume. + +Current fields distinguish the first three natural-execution cases below; +deferred checkpoint fields will distinguish the final two: - a fully successful continuous expert trajectory; - a failed complete episode retained for diagnostics; +- a successful independent natural segment fragment; - a successful independent suffix fragment restored from a checkpoint; - a failed restore that emitted no controller commands. @@ -364,7 +396,7 @@ The first implementation should remain outside atomic-action planning. | `embodichain/lab/gym/envs/expert_program/bridge.py` | Allow one selected compiled segment to run from an explicit initial eligibility mask and `TaskState`; keep validator as the acceptance commit point | | `embodichain/lab/gym/envs/expert_program/environment.py` | Assemble a fresh runtime from an explicit initial `TaskState` and selected segment index | | simulation Expert Program integration | Implement the checkpoint port and post-restore synchronization/entry validation | -| dataset recorders | Persist fragment provenance and `segment_success` without treating a restored suffix as a continuous episode | +| dataset recorders | Persist fragment provenance and `segment_accepted` without treating a fragment as a continuous episode | | online data engine | Filter unsuccessful segments and reject windows crossing `continuity_id` | `TrajectorySegment`, `ActionPlan`, `ExecutionSession`, atomic recovery policy, @@ -375,15 +407,17 @@ and semantic effect monitors require no API change for the initial feature. ### Phase A: outcome propagation 1. Expose the bridge's accepted mask as the sole segment outcome. -2. Add dense `segment_success`, `segment_attempt_id`, and `continuity_id` +2. Add dense `segment_accepted`, `segment_attempt_id`, and `continuity_id` annotations. 3. Update recorders and online sampling to consume them. 4. Keep execution fail-closed. -This phase is independently useful and low risk. +This phase is implemented. ### Phase B: checkpoint primitives +Deferred; not part of the current implementation. + 1. Add immutable checkpoint/result values and an explicit port. 2. Extend simulation state capture/restore for selected rows and controller targets. @@ -393,20 +427,28 @@ This phase is independently useful and low risk. ### Phase C: fragment collector -1. Execute one selected compiled segment through a fresh runtime. -2. Commit accepted segments independently. +1. Commit naturally executed accepted segments independently. **Implemented.** +2. Execute one selected compiled segment through a fresh runtime. **Deferred.** 3. On failure, optionally restore entry of the next segment and continue under a new `continuity_id`. + **Deferred.** 4. Preserve failed attempt metadata without promoting the collection session - to continuous episode success. + to continuous episode success. **Implemented.** ## 11. Validation surface -Focused unit tests must cover: +Implemented focused unit tests cover: - runtime, post-policy, and validator masks combining into one accepted mask; - retroactive per-frame segment outcome annotation; - online sampling excluding failed segments and cross-restore windows; +- natural accepted segments being saved as separate synchronous and + asynchronous LeRobot episodes; +- failed fragments requiring explicit opt-in; +- accepted prefix fragments not promoting whole-episode success. + +Deferred checkpoint work must additionally cover: + - missing, incompatible, and stale checkpoints failing before commands; - row-selective restore leaving healthy peers unchanged; - controller targets matching restored qpos on the first subsequent step; diff --git a/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst b/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst index b80a2b364..aa9da5f23 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst @@ -79,6 +79,13 @@ segment spans. .. currentmodule:: embodichain.lab.gym.envs.demo +.. autoclass:: DemoExecutionCfg + :members: + +.. autodata:: DemoOutputMode + +.. autodata:: DemoSegmentOutcomeKind + .. autoclass:: DemoSegment :members: diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index b5890e243..86bcfd0a9 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -85,8 +85,11 @@ embodichain.lab.gym.envs.demo DEMO_ANNOTATION_KEYS DEMO_SCHEMA_VERSION + DemoExecutionCfg DemoEpisodeResult + DemoOutputMode DemoSegment + DemoSegmentOutcomeKind DemoSegmentResult execute_demo_episode resolve_demo_segments diff --git a/embodichain/data_pipeline/engine/data.py b/embodichain/data_pipeline/engine/data.py index 069140a5e..fb73a94d6 100644 --- a/embodichain/data_pipeline/engine/data.py +++ b/embodichain/data_pipeline/engine/data.py @@ -354,6 +354,7 @@ def _run_sim_worker( result = execute_demo_episode( env, episode_index=rollout_idx, + attempt_id=attempt - 1, should_stop=close_signal.is_set, progress=lambda actions, description: tqdm( actions, @@ -916,9 +917,10 @@ def sample_batch( Only fully valid windows are candidates, so padding or stale tail frames are never returned. ``episode`` mode allows a window to cross - segment boundaries, ``segment`` keeps every window inside one segment, - and ``boundary`` deliberately samples windows crossing an internal - segment boundary. + segment boundaries within one causal-continuity region, ``segment`` + keeps every window inside one accepted segment, and ``boundary`` + deliberately samples windows crossing a boundary between accepted + segments. No mode crosses a discontinuous state-restore boundary. After sampling the internal :attr:`_sample_count` is incremented by *batch_size*; if the count exceeds @@ -990,6 +992,17 @@ def sample_batch( if segment_ids is None: segment_ids = torch.zeros_like(valid, dtype=torch.int64) + continuity_ids = self.shared_buffer.get("continuity_id", None) + if continuity_ids is None: + # Schema-v2 and earlier buffers contain no out-of-band state + # restore, so the complete row belongs to continuity region 0. + continuity_ids = torch.zeros_like(valid, dtype=torch.int64) + continuity_windows = continuity_ids.unfold(1, chunk_size, 1) + same_continuity = (continuity_windows == continuity_windows[..., :1]).all( + dim=-1 + ) & (continuity_windows[..., 0] >= 0) + valid_windows &= same_continuity + if sampling_mode == "segment": segment_windows = segment_ids.unfold(1, chunk_size, 1) same_segment = (segment_windows == segment_windows[..., :1]).all( @@ -1003,6 +1016,15 @@ def sample_batch( ).any(dim=-1) valid_windows &= crosses_boundary + if sampling_mode in {"segment", "boundary"}: + segment_accepted = self.shared_buffer.get("segment_accepted", None) + if segment_accepted is None: + # Older successful-only online buffers predate explicit + # segment qualification and remain fully eligible. + segment_accepted = torch.ones_like(valid, dtype=torch.bool) + accepted_windows = segment_accepted.bool().unfold(1, chunk_size, 1) + valid_windows &= accepted_windows.all(dim=-1) + candidate_rows = ( valid_windows.any(dim=1).nonzero(as_tuple=False).squeeze(-1) ) diff --git a/embodichain/lab/gym/envs/demo.py b/embodichain/lab/gym/envs/demo.py index 6a5e7b432..010f724da 100644 --- a/embodichain/lab/gym/envs/demo.py +++ b/embodichain/lab/gym/envs/demo.py @@ -25,19 +25,24 @@ import torch +from embodichain.utils import configclass + from ._json import json_safe_copy as _json_safe_copy __all__ = [ "DEMO_ANNOTATION_KEYS", "DEMO_SCHEMA_VERSION", + "DemoExecutionCfg", "DemoEpisodeResult", + "DemoOutputMode", "DemoSegment", + "DemoSegmentOutcomeKind", "DemoSegmentResult", "execute_demo_episode", "resolve_demo_segments", ] -DEMO_SCHEMA_VERSION = 2 +DEMO_SCHEMA_VERSION = 3 """Current version of the segment-aware demonstration metadata schema.""" DEMO_ANNOTATION_KEYS = ( @@ -47,11 +52,121 @@ "segment_step", "segment_start", "segment_end", + "segment_accepted", + "segment_attempt_id", + "continuity_id", "terminated", "truncated", ) """Per-frame annotation keys stored in expert rollout buffers.""" +DemoOutputMode = Literal["continuous", "segment_fragments"] +"""Supported persistence layouts for one demonstration execution.""" + +DemoSegmentOutcomeKind = Literal[ + "succeeded", + "runtime_failed", + "post_policy_failed", + "validation_failed", + "cancelled", + "truncated", + "not_attempted", +] +"""Stable, first-failure-phase outcome for one program segment row.""" + + +@configclass +class DemoExecutionCfg: + """Collector-owned settings for demonstration persistence. + + ``segment_fragments`` persists each eligible program segment as an + independent LeRobot episode. It does not resume execution after a failed + segment; checkpoint capture and resume are intentionally outside this + configuration until an authoritative restore port exists. + + Args: + mode: Continuous episode or independent segment-fragment persistence. + save_failed_fragments: Whether failed segments with recorded frames are + retained in fragment mode. Failed fragments remain explicitly + annotated and are excluded by successful-segment sampling. + """ + + mode: DemoOutputMode = "continuous" + save_failed_fragments: bool = False + + def __post_init__(self) -> None: + if self.mode not in {"continuous", "segment_fragments"}: + raise ValueError( + "mode must be 'continuous' or 'segment_fragments', " + f"got {self.mode!r}." + ) + if not isinstance(self.save_failed_fragments, bool): + raise TypeError("save_failed_fragments must be a bool.") + if self.mode == "continuous" and self.save_failed_fragments: + raise ValueError( + "save_failed_fragments is only valid in segment_fragments mode." + ) + + +def _validation_mask_value( + validation: Mapping[str, Any], key: str, env_id: int +) -> bool | None: + """Return one optional row value from bridge validation metadata.""" + values = validation.get(key) + if not isinstance(values, (list, tuple)) or env_id >= len(values): + return None + value = values[env_id] + return bool(value) if isinstance(value, bool) else None + + +def _segment_outcome_kind( + *, + participant: bool, + success: bool, + failure_reason: str | None, + metadata: Mapping[str, Any], + env_id: int, +) -> DemoSegmentOutcomeKind: + """Derive the first authoritative failure phase without re-validating.""" + if not participant: + return "not_attempted" + if success: + return "succeeded" + if failure_reason == "truncated": + return "truncated" + if failure_reason in {"interrupted", "batch_aborted", "empty_segment"}: + return "cancelled" + + validation = metadata.get("validation") + if isinstance(validation, Mapping): + if _validation_mask_value(validation, "runtime_success_mask", env_id) is False: + return "runtime_failed" + if ( + _validation_mask_value(validation, "post_policy_success_mask", env_id) + is False + ): + return "post_policy_failed" + validators = validation.get("validators") + if isinstance(validators, (list, tuple)): + for validator in validators: + if not isinstance(validator, Mapping): + continue + result_mask = validator.get("result_mask") + if ( + isinstance(result_mask, (list, tuple)) + and env_id < len(result_mask) + and result_mask[env_id] is False + ): + return "validation_failed" + if _validation_mask_value(validation, "accepted_mask", env_id) is False: + return "validation_failed" + + if failure_reason == "segment_validation_failed": + return "validation_failed" + if failure_reason is None: + return "not_attempted" + return "runtime_failed" + @dataclass(frozen=True) class DemoSegment: @@ -128,6 +243,10 @@ class DemoSegmentResult: end_steps: Per-environment exclusive ends. successes: Per-environment segment status. failure_reasons: Per-environment failure reasons. + attempt_id: Collection attempt that produced this segment. + continuity_id: Causal-continuity region containing this segment. + outcome_kind: Aggregate first-failure-phase outcome. + outcome_kinds: Per-environment first-failure-phase outcomes. """ segment_id: int @@ -144,16 +263,73 @@ class DemoSegmentResult: end_steps: tuple[int, ...] = () successes: tuple[bool, ...] = () failure_reasons: tuple[str | None, ...] = () + attempt_id: int = 0 + continuity_id: int = 0 + outcome_kind: DemoSegmentOutcomeKind | None = None + outcome_kinds: tuple[DemoSegmentOutcomeKind, ...] = () def __post_init__(self) -> None: if not isinstance(self.metadata, Mapping): raise TypeError("metadata must be a mapping.") + if self.attempt_id < 0: + raise ValueError("attempt_id must be non-negative.") + if self.continuity_id < 0: + raise ValueError("continuity_id must be non-negative.") owned_metadata = _json_safe_copy( self.metadata, field_name="segment result metadata", ) object.__setattr__(self, "metadata", MappingProxyType(owned_metadata)) + if self.outcome_kinds: + expected = len(self.successes) or len(self.active) + if expected and len(self.outcome_kinds) != expected: + raise ValueError( + "outcome_kinds must contain one value per environment row." + ) + row_outcomes = self.outcome_kinds + elif self.successes: + active = self.active or (True,) * len(self.successes) + row_outcomes = tuple( + _segment_outcome_kind( + participant=(active[env_id] if env_id < len(active) else True), + success=self.successes[env_id], + failure_reason=( + self.failure_reasons[env_id] + if env_id < len(self.failure_reasons) + else self.failure_reason + ), + metadata=owned_metadata, + env_id=env_id, + ) + for env_id in range(len(self.successes)) + ) + object.__setattr__(self, "outcome_kinds", row_outcomes) + else: + row_outcomes = () + + if self.outcome_kind is None: + if self.success: + aggregate_outcome: DemoSegmentOutcomeKind = "succeeded" + elif row_outcomes: + aggregate_outcome = next( + ( + outcome + for outcome in row_outcomes + if outcome not in {"succeeded", "not_attempted"} + ), + "not_attempted", + ) + else: + aggregate_outcome = _segment_outcome_kind( + participant=True, + success=False, + failure_reason=self.failure_reason, + metadata=owned_metadata, + env_id=0, + ) + object.__setattr__(self, "outcome_kind", aggregate_outcome) + def to_metadata(self, env_id: int | None = None) -> dict[str, Any]: """Return a JSON-compatible aggregate or per-environment representation. @@ -169,6 +345,8 @@ def to_metadata(self, env_id: int | None = None) -> dict[str, Any]: "name": self.name, "target_uid": self.target_uid, "instruction": self.instruction, + "attempt_id": self.attempt_id, + "continuity_id": self.continuity_id, "metadata": _json_safe_copy( self.metadata, field_name="segment result metadata", @@ -178,9 +356,26 @@ def to_metadata(self, env_id: int | None = None) -> dict[str, Any]: metadata.update( { "start_step": self.start_steps[env_id], - "end_step": self.end_steps[env_id], - "success": self.successes[env_id], - "failure_reason": self.failure_reasons[env_id], + "end_step": ( + self.end_steps[env_id] + if env_id < len(self.end_steps) + else self.end_step + ), + "success": ( + self.successes[env_id] + if env_id < len(self.successes) + else self.success + ), + "failure_reason": ( + self.failure_reasons[env_id] + if env_id < len(self.failure_reasons) + else self.failure_reason + ), + "outcome_kind": ( + self.outcome_kinds[env_id] + if env_id < len(self.outcome_kinds) + else self.outcome_kind + ), } ) return metadata @@ -191,6 +386,7 @@ def to_metadata(self, env_id: int | None = None) -> dict[str, Any]: "end_step": self.end_step, "success": self.success, "failure_reason": self.failure_reason, + "outcome_kind": self.outcome_kind, } ) if self.active: @@ -201,6 +397,7 @@ def to_metadata(self, env_id: int | None = None) -> dict[str, Any]: "end_steps": list(self.end_steps), "successes": list(self.successes), "failure_reasons": list(self.failure_reasons), + "outcome_kinds": list(self.outcome_kinds), } ) return metadata @@ -222,6 +419,8 @@ class DemoEpisodeResult: lengths: Independent per-environment recorded lengths. completed_by_env: Independent valid-completion flags. terminal_reasons: Independent terminal reasons. + execution_mode: Persistence layout selected for this execution. + attempt_id: Zero-based collection attempt identifier. """ episode_index: int @@ -235,6 +434,8 @@ class DemoEpisodeResult: lengths: tuple[int, ...] = () completed_by_env: tuple[bool, ...] = () terminal_reasons: tuple[str, ...] = () + execution_mode: DemoOutputMode = "continuous" + attempt_id: int = 0 @property def all_success(self) -> bool: @@ -246,11 +447,47 @@ def any_success(self) -> bool: """Whether at least one parallel environment completed successfully.""" return any(self.success) + @property + def successful_fragment_count_by_env(self) -> tuple[int, ...]: + """Count accepted, non-empty program segments for each environment.""" + if not self.success: + return () + counts = [0] * len(self.success) + for segment in self.segments: + if not segment.successes: + if segment.success and segment.end_step > segment.start_step: + counts[0] += 1 + continue + for env_id, accepted in enumerate(segment.successes): + start = ( + segment.start_steps[env_id] + if env_id < len(segment.start_steps) + else segment.start_step + ) + end = ( + segment.end_steps[env_id] + if env_id < len(segment.end_steps) + else segment.end_step + ) + if ( + accepted + and ( + not segment.active + or env_id >= len(segment.active) + or segment.active[env_id] + ) + and end > start + ): + counts[env_id] += 1 + return tuple(counts) + def to_metadata(self) -> dict[str, Any]: """Return a JSON-compatible representation.""" metadata = { "schema_version": DEMO_SCHEMA_VERSION, "episode_index": self.episode_index, + "execution_mode": self.execution_mode, + "attempt_id": self.attempt_id, "length": self.length, "completed": self.completed, "success": list(self.success), @@ -258,6 +495,9 @@ def to_metadata(self) -> dict[str, Any]: "truncated": list(self.truncated), "terminal_reason": self.terminal_reason, "segments": [segment.to_metadata() for segment in self.segments], + "successful_fragment_count_by_env": list( + self.successful_fragment_count_by_env + ), } if self.lengths: metadata.update( @@ -405,6 +645,8 @@ def execute_demo_episode( env: Any, *, episode_index: int = 0, + execution_cfg: DemoExecutionCfg | None = None, + attempt_id: int = 0, should_stop: StopPredicate | None = None, progress: ProgressWrapper | None = None, **plan_kwargs: Any, @@ -419,6 +661,9 @@ def execute_demo_episode( Args: env: Gym environment or wrapper. episode_index: Logical episode identifier used in metadata and logs. + execution_cfg: Collector-owned output settings. Defaults to continuous + episode persistence. + attempt_id: Zero-based identifier for this collection attempt. should_stop: Optional callback checked before every action. progress: Optional wrapper such as ``tqdm`` for action iterables. **plan_kwargs: Arguments forwarded to the task's planning method. @@ -427,6 +672,13 @@ def execute_demo_episode( A :class:`DemoEpisodeResult` describing segment spans and terminal state. """ + if execution_cfg is None: + execution_cfg = DemoExecutionCfg() + elif not isinstance(execution_cfg, DemoExecutionCfg): + raise TypeError("execution_cfg must be a DemoExecutionCfg or None.") + if attempt_id < 0: + raise ValueError("attempt_id must be non-negative.") + target = _env_target(env) num_envs = int(getattr(target, "num_envs", 1)) begin_episode = _get_env_callable(env, "_begin_demo_episode_recording") @@ -454,7 +706,11 @@ def publish_active_mask() -> None: ) if begin_episode is not None: - begin_episode(episode_index=episode_index) + begin_episode( + episode_index=episode_index, + execution_cfg=execution_cfg, + attempt_id=attempt_id, + ) publish_active_mask() previous_no_auto_reset = bool(getattr(target, "_demo_no_auto_reset", False)) @@ -805,6 +1061,8 @@ def publish_active_mask() -> None: end_steps=end_steps, successes=tuple(segment_successes), failure_reasons=tuple(segment_failure_reasons), + attempt_id=attempt_id, + continuity_id=0, ) segment_results.append(segment_result) if end_segment is not None: @@ -864,6 +1122,8 @@ def publish_active_mask() -> None: lengths=tuple(lengths), completed_by_env=tuple(completed_by_env), terminal_reasons=tuple(terminal_reasons), + execution_mode=execution_cfg.mode, + attempt_id=attempt_id, ) if end_episode is not None: end_episode(result=result) diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index 3b4fa198b..23c1a0ea6 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -59,6 +59,7 @@ from embodichain.lab.gym.envs import BaseEnv, EnvCfg from embodichain.lab.gym.envs.demo import ( DEMO_SCHEMA_VERSION, + DemoExecutionCfg, DemoEpisodeResult, DemoSegment, DemoSegmentResult, @@ -402,6 +403,10 @@ def __init__(self, cfg: EmbodiedEnvCfg, **kwargs): # common demo executor updates this context while the regular rollout # writer turns it into per-frame annotations. self._demo_episode_index = 0 + self._demo_execution_cfg = DemoExecutionCfg() + self._demo_attempt_id = 0 + self._demo_continuity_id = 0 + self._demo_program_run_id = "0:0" self._demo_active_segment_id = 0 self._demo_active_segment_ids = torch.zeros( self.num_envs, dtype=torch.long, device=self.device @@ -858,8 +863,18 @@ def _initialize_episode( env_ids_to_save = env_ids_to_process else: successful_envs = self.episode_success_status | self._task_success + fragment_envs = torch.tensor( + [ + EmbodiedEnv._demo_fragment_row_is_persistable( + self, int(env_id) + ) + for env_id in env_ids_to_process.cpu().tolist() + ], + dtype=torch.bool, + device=status_device, + ) env_ids_to_save = env_ids_to_process[ - successful_envs[env_ids_to_process] + successful_envs[env_ids_to_process] | fragment_envs ] if env_ids_to_save.numel() > 0: @@ -968,20 +983,58 @@ def _clear_expert_rollout_rows(self, env_ids: torch.Tensor) -> None: "valid", "segment_start", "segment_end", + "segment_accepted", "terminated", "truncated", ): if key in self.rollout_buffer.keys(): self.rollout_buffer[key][buffer_ids] = False - for key in ("episode_step", "segment_id", "segment_step"): + for key in ( + "episode_step", + "segment_id", + "segment_step", + "segment_attempt_id", + "continuity_id", + ): if key in self.rollout_buffer.keys(): self.rollout_buffer[key][buffer_ids] = -1 + def _demo_fragment_row_is_persistable(self, env_id: int) -> bool: + """Return whether one fragment-mode row owns an eligible frame span.""" + episode_metadata = getattr(self, "_demo_episode_metadata", None) + if episode_metadata is None or env_id >= len(episode_metadata): + return False + metadata = episode_metadata[env_id] + if metadata.get("output_mode") != "segment_fragments": + return False + include_failed = bool(metadata.get("save_failed_fragments", False)) + for segment in metadata.get("segments", []): + if int(segment.get("end_step", 0)) <= int(segment.get("start_step", 0)): + continue + if bool(segment.get("success", False)) or include_failed: + return True + return False + def _new_demo_episode_metadata(self, env_id: int) -> dict[str, Any]: """Create an empty metadata record for one environment row.""" + execution_cfg = getattr(self, "_demo_execution_cfg", None) + if execution_cfg is None: + execution_cfg = DemoExecutionCfg() + attempt_id = int(getattr(self, "_demo_attempt_id", 0)) return { "schema_version": DEMO_SCHEMA_VERSION, "episode_index": int(getattr(self, "_demo_episode_index", 0)), + "output_mode": execution_cfg.mode, + "save_failed_fragments": execution_cfg.save_failed_fragments, + "attempt_id": attempt_id, + "continuity_id": int(getattr(self, "_demo_continuity_id", 0)), + "program_run_id": str( + getattr( + self, + "_demo_program_run_id", + f"{int(getattr(self, '_demo_episode_index', 0))}:{attempt_id}", + ) + ), "env_id": env_id, "length": 0, "completed": False, @@ -992,9 +1045,20 @@ def _new_demo_episode_metadata(self, env_id: int) -> dict[str, Any]: "segments": [], } - def _begin_demo_episode_recording(self, episode_index: int = 0) -> None: + def _begin_demo_episode_recording( + self, + episode_index: int = 0, + execution_cfg: DemoExecutionCfg | None = None, + attempt_id: int = 0, + ) -> None: """Start annotation metadata for a new demonstration episode.""" + if execution_cfg is None: + execution_cfg = DemoExecutionCfg() self._demo_episode_index = episode_index + self._demo_execution_cfg = execution_cfg + self._demo_attempt_id = attempt_id + self._demo_continuity_id = 0 + self._demo_program_run_id = f"{episode_index}:{attempt_id}" self._demo_active_segment_id = 0 self._demo_active_segment_ids.zero_() self._demo_active_mask.fill_(True) @@ -1057,6 +1121,17 @@ def _end_demo_segment_recording(self, result: DemoSegmentResult) -> None: and rollout_end > rollout_start ): self.rollout_buffer["segment_end"][env_id, rollout_end - 1] = True + if ( + self.rollout_buffer is not None + and "segment_accepted" in self.rollout_buffer.keys() + and rollout_end > rollout_start + ): + accepted = ( + result.successes[env_id] if result.successes else result.success + ) + self.rollout_buffer["segment_accepted"][ + env_id, rollout_start:rollout_end + ] = accepted metadata = result.to_metadata(env_id if result.start_steps else None) metadata["start_step"] = start @@ -1179,6 +1254,13 @@ def get_demo_episode_metadata(self, env_id: int) -> dict[str, Any]: "start_step": 0, "end_step": length, "success": success, + "attempt_id": int(metadata.get("attempt_id", 0)), + "continuity_id": int(metadata.get("continuity_id", 0)), + "outcome_kind": ( + "succeeded" + if success + else ("truncated" if truncated else "validation_failed") + ), "target_uid": None, "instruction": instruction, "failure_reason": None if success else terminal_reason, @@ -1285,6 +1367,20 @@ def _write_episode_rollout_step( self.rollout_buffer["segment_start"][buffer_env_ids, buffer_step_ids] = ( segment_steps.to(buffer_device) == 0 ) + if "segment_accepted" in buffer_keys: + # Acceptance is unknown until the segment validator commits. The + # end-segment hook fills this complete span retroactively. + self.rollout_buffer["segment_accepted"][ + buffer_env_ids, buffer_step_ids + ] = False + if "segment_attempt_id" in buffer_keys: + self.rollout_buffer["segment_attempt_id"][ + buffer_env_ids, buffer_step_ids + ] = int(getattr(self, "_demo_attempt_id", 0)) + if "continuity_id" in buffer_keys: + self.rollout_buffer["continuity_id"][buffer_env_ids, buffer_step_ids] = int( + getattr(self, "_demo_continuity_id", 0) + ) if terminateds is None: terminateds = torch.zeros( diff --git a/embodichain/lab/gym/envs/managers/async_datasets.py b/embodichain/lab/gym/envs/managers/async_datasets.py index 45e1fb320..7093e305d 100644 --- a/embodichain/lab/gym/envs/managers/async_datasets.py +++ b/embodichain/lab/gym/envs/managers/async_datasets.py @@ -237,8 +237,8 @@ def __call__( if metadata_getter is not None else None ) - self._save_queue.put( - ( + payloads = list( + self._episode_payloads( env_id, obs_clone, action_clone, @@ -246,6 +246,17 @@ def __call__( episode_metadata, ) ) + if ( + episode_metadata is not None + and episode_metadata.get("output_mode") == "segment_fragments" + and not payloads + ): + raise RuntimeError( + f"Committed fragment collection for env {env_id} had no " + "eligible segment spans." + ) + for payload in payloads: + self._save_queue.put(payload) def finalize(self) -> Optional[str]: """Drain committed writes, finalize storage, and surface all failures. diff --git a/embodichain/lab/gym/envs/managers/datasets.py b/embodichain/lab/gym/envs/managers/datasets.py index dd655d048..b8cc51f1f 100644 --- a/embodichain/lab/gym/envs/managers/datasets.py +++ b/embodichain/lab/gym/envs/managers/datasets.py @@ -18,6 +18,7 @@ from __future__ import annotations +import copy import json import math import threading @@ -72,6 +73,9 @@ "segment_step": "annotation.segment_step", "segment_start": "annotation.segment_start", "segment_end": "annotation.segment_end", + "segment_accepted": "annotation.segment_accepted", + "segment_attempt_id": "annotation.segment_attempt_id", + "continuity_id": "annotation.continuity_id", "terminated": "annotation.terminated", "truncated": "annotation.truncated", } @@ -281,17 +285,167 @@ def _save_episodes( episode_metadata = ( metadata_getter(env_id) if metadata_getter is not None else None ) - saved = self._save_single_episode( - env_id, - obs_list, - action_list, - annotations=annotations, - episode_metadata=episode_metadata, + payloads = list( + self._episode_payloads( + env_id, + obs_list, + action_list, + annotations, + episode_metadata, + ) ) - if not saved: + if ( + episode_metadata is not None + and episode_metadata.get("output_mode") == "segment_fragments" + and not payloads + ): raise RuntimeError( - f"Committed episode for env {env_id} was not persisted." + f"Committed fragment collection for env {env_id} had no " + "eligible segment spans." ) + for payload in payloads: + saved = self._save_single_episode(*payload) + if not saved: + raise RuntimeError( + f"Committed episode for env {env_id} was not persisted." + ) + + def _episode_payloads( + self, + env_id: int, + obs_list: Any, + action_list: Any, + annotations: Mapping[str, Any], + episode_metadata: Mapping[str, Any] | None, + ) -> Iterable[tuple[int, Any, Any, Mapping[str, Any], Mapping[str, Any] | None]]: + """Yield one continuous payload or independent natural-segment slices.""" + if ( + episode_metadata is None + or episode_metadata.get("output_mode") != "segment_fragments" + ): + yield env_id, obs_list, action_list, annotations, episode_metadata + return + + episode_length = min(len(obs_list), len(action_list)) + include_failed = bool(episode_metadata.get("save_failed_fragments", False)) + for segment in episode_metadata.get("segments", []): + if not isinstance(segment, Mapping): + raise TypeError("Segment sidecar metadata must be a mapping.") + accepted = bool(segment.get("success", False)) + if not accepted and not include_failed: + continue + start = int(segment.get("start_step", 0)) + end = int(segment.get("end_step", 0)) + if start < 0 or end > episode_length or end <= start: + raise ValueError( + "Fragment span must be a non-empty subset of the buffered " + f"episode; got [{start}, {end}) for length {episode_length}." + ) + + length = end - start + fragment_annotations = { + key: values[start:end].clone() for key, values in annotations.items() + } + reference = next(iter(fragment_annotations.values()), None) + device = getattr(reference, "device", None) + fragment_annotations["episode_step"] = torch.arange( + length, dtype=torch.int64, device=device + ) + fragment_annotations["segment_step"] = torch.arange( + length, dtype=torch.int64, device=device + ) + fragment_annotations["segment_start"] = torch.zeros( + length, dtype=torch.bool, device=device + ) + fragment_annotations["segment_start"][0] = True + fragment_annotations["segment_end"] = torch.zeros( + length, dtype=torch.bool, device=device + ) + fragment_annotations["segment_end"][-1] = True + fragment_annotations["segment_accepted"] = torch.full( + (length,), accepted, dtype=torch.bool, device=device + ) + fragment_annotations["segment_attempt_id"] = torch.full( + (length,), + int(segment.get("attempt_id", episode_metadata.get("attempt_id", 0))), + dtype=torch.int64, + device=device, + ) + fragment_annotations["continuity_id"] = torch.full( + (length,), + int( + segment.get( + "continuity_id", episode_metadata.get("continuity_id", 0) + ) + ), + dtype=torch.int64, + device=device, + ) + + fragment_segment = copy.deepcopy(dict(segment)) + fragment_segment.update({"start_step": 0, "end_step": length}) + fragment_metadata = copy.deepcopy(dict(episode_metadata)) + program_run_id = str(fragment_metadata.get("program_run_id", "unknown")) + segment_provenance = segment.get("metadata", {}) + if not isinstance(segment_provenance, Mapping): + segment_provenance = {} + segment_attempt_id = int( + segment.get("attempt_id", episode_metadata.get("attempt_id", 0)) + ) + continuity_id = int( + segment.get("continuity_id", episode_metadata.get("continuity_id", 0)) + ) + fragment_terminated = bool( + torch.as_tensor( + fragment_annotations.get( + "terminated", torch.zeros(length, dtype=torch.bool) + )[-1] + ).item() + ) + fragment_truncated = bool( + torch.as_tensor( + fragment_annotations.get( + "truncated", torch.zeros(length, dtype=torch.bool) + )[-1] + ).item() + ) + fragment_metadata.update( + { + "fragment": True, + "fragment_origin": "natural_segment", + "fragment_id": ( + f"{program_run_id}:{env_id}:" + f"{int(segment.get('segment_id', 0))}:" + f"{segment_attempt_id}:{continuity_id}" + ), + "source_program_id": segment_provenance.get("expert_program_id"), + "program_segment_id": segment_provenance.get("program_segment_id"), + "source_episode_index": fragment_metadata.get("episode_index"), + "source_env_id": env_id, + "source_start_step": start, + "source_end_step": end, + "length": length, + "completed": accepted, + "success": accepted, + "terminated": fragment_terminated, + "truncated": fragment_truncated, + "terminal_reason": ( + "segment_succeeded" + if accepted + else segment.get("failure_reason") + or segment.get("outcome_kind") + or "segment_failed" + ), + "segments": [fragment_segment], + } + ) + yield ( + env_id, + obs_list[start:end].clone(), + action_list[start:end].clone(), + fragment_annotations, + fragment_metadata, + ) def _episode_length(self, env_id: int) -> int: """Return the valid buffered length for one environment.""" @@ -332,6 +486,12 @@ def _save_single_episode( if self.instruction else "unknown_task" ) + if episode_metadata is not None and episode_metadata.get("fragment"): + segments = episode_metadata.get("segments", []) + if segments and isinstance(segments[0], Mapping): + task = self._normalize_subtask_description( + segments[0].get("instruction") or task + ) if len(obs_list) == 0: logger.log_warning(f"No episode data to save for env {env_id}") @@ -362,6 +522,8 @@ def _save_single_episode( depth_prefix = f"{LeRobotKey.OBS_PREFIX.value}depth." episode_index = self.curr_episode dataset_committed = False + episode_attempt_id = int((episode_metadata or {}).get("attempt_id", 0)) + episode_continuity_id = int((episode_metadata or {}).get("continuity_id", 0)) try: frame_subtasks = [ self._subtask_for_frame(task, episode_metadata, frame_index) @@ -379,12 +541,28 @@ def _save_single_episode( desc=f"Converting env {env_id} episode to LeRobot format", ) ): + frame_segment = self._segment_for_frame(episode_metadata, frame_index) frame_annotations = { "episode_step": frame_index, "segment_id": 0, "segment_step": frame_index, "segment_start": frame_index == 0, "segment_end": frame_index == len(obs_list) - 1, + "segment_accepted": ( + bool(frame_segment.get("success", True)) + if frame_segment is not None + else True + ), + "segment_attempt_id": ( + int(frame_segment.get("attempt_id", episode_attempt_id)) + if frame_segment is not None + else episode_attempt_id + ), + "continuity_id": ( + int(frame_segment.get("continuity_id", episode_continuity_id)) + if frame_segment is not None + else episode_continuity_id + ), "terminated": False, "truncated": False, } @@ -533,23 +711,36 @@ def _normalize_scalar_episode_buffer(self) -> None: episode_buffer[feature_key] = normalized_values @staticmethod - def _subtask_for_frame( - default_subtask: str, + def _segment_for_frame( episode_metadata: Mapping[str, Any] | None, frame_index: int, - ) -> str: - """Resolve a segment-specific instruction for one LeRobot frame.""" + ) -> Mapping[str, Any] | None: + """Return the sidecar segment owning one frame, when available.""" if episode_metadata is None: - return LeRobotRecorder._normalize_subtask_description(default_subtask) + return None for segment in episode_metadata.get("segments", []): + if not isinstance(segment, Mapping): + continue if ( int(segment.get("start_step", 0)) <= frame_index < int(segment.get("end_step", 0)) ): - return LeRobotRecorder._normalize_subtask_description( - segment.get("instruction") or default_subtask - ) + return segment + return None + + @staticmethod + def _subtask_for_frame( + default_subtask: str, + episode_metadata: Mapping[str, Any] | None, + frame_index: int, + ) -> str: + """Resolve a segment-specific instruction for one LeRobot frame.""" + segment = LeRobotRecorder._segment_for_frame(episode_metadata, frame_index) + if segment is not None: + return LeRobotRecorder._normalize_subtask_description( + segment.get("instruction") or default_subtask + ) return LeRobotRecorder._normalize_subtask_description(default_subtask) @staticmethod diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index e524b6765..995ba17ec 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -1196,6 +1196,21 @@ def _init_buffer_from_space( "segment_end": torch.zeros( (num_envs, max_episode_steps), dtype=torch.bool, device=device ), + "segment_accepted": torch.zeros( + (num_envs, max_episode_steps), dtype=torch.bool, device=device + ), + "segment_attempt_id": torch.full( + (num_envs, max_episode_steps), + -1, + dtype=torch.int64, + device=device, + ), + "continuity_id": torch.full( + (num_envs, max_episode_steps), + -1, + dtype=torch.int64, + device=device, + ), "terminated": torch.zeros( (num_envs, max_episode_steps), dtype=torch.bool, device=device ), @@ -1422,6 +1437,21 @@ def init_rollout_buffer_from_config( "segment_end": torch.zeros( (batch_size, max_episode_steps), dtype=torch.bool, device=device ), + "segment_accepted": torch.zeros( + (batch_size, max_episode_steps), dtype=torch.bool, device=device + ), + "segment_attempt_id": torch.full( + (batch_size, max_episode_steps), + -1, + dtype=torch.int64, + device=device, + ), + "continuity_id": torch.full( + (batch_size, max_episode_steps), + -1, + dtype=torch.int64, + device=device, + ), "terminated": torch.zeros( (batch_size, max_episode_steps), dtype=torch.bool, device=device ), diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index bef780469..770fea16b 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/embodichain/lab/scripts/run_env.py @@ -31,7 +31,11 @@ import torch import tqdm -from embodichain.lab.gym.envs.demo import DemoEpisodeResult, execute_demo_episode +from embodichain.lab.gym.envs.demo import ( + DemoEpisodeResult, + DemoExecutionCfg, + execute_demo_episode, +) from embodichain.lab.gym.envs.expert_program.loader import ( load_expert_program as _load_expert_program, ) @@ -167,6 +171,47 @@ def _selected_rows_have_frames( return result.length > 0 +def _persistable_fragment_env_ids( + env: Any, + result: DemoEpisodeResult, + save_env_ids: Sequence[int], + *, + include_failed: bool, +) -> tuple[int, ...]: + """Select rows containing at least one eligible non-empty segment span.""" + persistable: list[int] = [] + metadata_getter = getattr(_env_target(env), "get_demo_episode_metadata", None) + for env_id in save_env_ids: + if callable(metadata_getter): + metadata = metadata_getter(env_id) + if isinstance(metadata, dict): + eligible = any( + int(segment.get("end_step", 0)) > int(segment.get("start_step", 0)) + and (bool(segment.get("success", False)) or include_failed) + for segment in metadata.get("segments", []) + if isinstance(segment, dict) + ) + if eligible: + persistable.append(env_id) + continue + for segment in result.segments: + if segment.active and not segment.active[env_id]: + continue + start = ( + segment.start_steps[env_id] + if segment.start_steps + else segment.start_step + ) + end = segment.end_steps[env_id] if segment.end_steps else segment.end_step + accepted = ( + segment.successes[env_id] if segment.successes else segment.success + ) + if end > start and (accepted or include_failed): + persistable.append(env_id) + break + return tuple(persistable) + + def generate_and_execute_action_list( env: gymnasium.Env, idx: int, @@ -215,6 +260,7 @@ def generate_function( save_video: bool = False, debug_mode: bool = False, save_env_ids: Sequence[int] | torch.Tensor | None = None, + execution_cfg: DemoExecutionCfg | None = None, **kwargs: Any, ) -> bool: """Generate, execute, and transactionally save one task episode batch. @@ -234,11 +280,14 @@ def generate_function( debug_mode (bool, optional): Enable debug mode for visualization and logging. save_env_ids: Environment rows to persist from this vector batch. Other rows are explicitly discarded after the selected rows commit. + execution_cfg: Continuous or independent segment-fragment persistence + settings. Checkpoint resume is not performed in either mode. **kwargs: Additional keyword arguments for data generation. Returns: - True if one episode per selected environment row was committed. With - ``save_failed_episodes`` enabled, committed episodes may be unsuccessful. + True if continuous episodes, or at least one eligible fragment row, + were committed. With ``save_failed_episodes`` enabled, committed + continuous episodes may be unsuccessful. """ if num_traj not in (None, 1): raise ValueError( @@ -252,6 +301,10 @@ def generate_function( raise ValueError(f"max_attempts must be at least 1, got {max_attempts}.") normalized_save_env_ids = _normalize_save_env_ids(env, save_env_ids) save_failed_episodes = _save_failed_episodes_enabled(env) + if execution_cfg is None: + execution_cfg = DemoExecutionCfg() + elif not isinstance(execution_cfg, DemoExecutionCfg): + raise TypeError("execution_cfg must be a DemoExecutionCfg or None.") if reset_before: _abort_pending_episode(env) @@ -262,16 +315,35 @@ def generate_function( result: DemoEpisodeResult = execute_demo_episode( env, episode_index=time_id, + execution_cfg=execution_cfg, + attempt_id=attempt - 1, progress=_progress_wrapper, **kwargs, ) successful = result.completed and result.all_success + fragment_env_ids = _persistable_fragment_env_ids( + env, + result, + normalized_save_env_ids, + include_failed=execution_cfg.save_failed_fragments, + ) persistable_failure = ( not successful and save_failed_episodes and _selected_rows_have_frames(result, normalized_save_env_ids) ) - if successful or persistable_failure: + if execution_cfg.mode == "segment_fragments" and fragment_env_ids: + _commit_pending_episode(env, fragment_env_ids) + commit_succeeded = True + if not successful: + log_warning( + f"Program run {time_id} stopped ({result.terminal_reason}); " + f"saved eligible segments from env rows {fragment_env_ids}." + ) + return True + if execution_cfg.mode == "continuous" and ( + successful or persistable_failure + ): # reset() is the commit boundary: dataset functors consume the # whole episode once, then buffers and scene state are reset. _commit_pending_episode(env, normalized_save_env_ids) diff --git a/tests/data_pipeline/test_online_data.py b/tests/data_pipeline/test_online_data.py index dd1641f8e..55947aa25 100644 --- a/tests/data_pipeline/test_online_data.py +++ b/tests/data_pipeline/test_online_data.py @@ -124,6 +124,12 @@ def _make_fake_engine( "rewards": torch.randn(buffer_size, max_episode_steps, 1), "valid": torch.ones(buffer_size, max_episode_steps, dtype=torch.bool), "segment_id": torch.zeros(buffer_size, max_episode_steps, dtype=torch.long), + "segment_accepted": torch.ones( + buffer_size, max_episode_steps, dtype=torch.bool + ), + "continuity_id": torch.zeros( + buffer_size, max_episode_steps, dtype=torch.long + ), }, batch_size=[buffer_size, max_episode_steps], ) @@ -731,6 +737,61 @@ def test_boundary_sampling_crosses_segment_boundary(self) -> None: .all() ) + def test_segment_sampling_excludes_failed_segment_frames(self) -> None: + """A valid transition remains ineligible when its segment was rejected.""" + midpoint = MAX_EPISODE_STEPS // 2 + self.engine.shared_buffer["segment_id"][:, midpoint:] = 1 + self.engine.shared_buffer["segment_accepted"][:, midpoint:] = False + + result = self.engine.sample_batch( + batch_size=64, + chunk_size=8, + sampling_mode="segment", + ) + + assert result["segment_accepted"].all() + assert (result["segment_id"] == 0).all() + + def test_boundary_sampling_rejects_failed_side_of_boundary(self) -> None: + """Boundary chunks cannot promote a rejected segment into training data.""" + midpoint = MAX_EPISODE_STEPS // 2 + self.engine.shared_buffer["segment_id"][:, midpoint:] = 1 + self.engine.shared_buffer["segment_accepted"][:, midpoint:] = False + + with pytest.raises(RuntimeError, match="No unlocked valid chunk"): + self.engine.sample_batch( + batch_size=1, + chunk_size=8, + sampling_mode="boundary", + ) + + def test_sampling_never_crosses_continuity_boundary(self) -> None: + """An out-of-band state jump is not exposed as a learnable transition.""" + midpoint = MAX_EPISODE_STEPS // 2 + self.engine.shared_buffer["continuity_id"][:, midpoint:] = 1 + + result = self.engine.sample_batch( + batch_size=64, + chunk_size=10, + sampling_mode="episode", + ) + + assert (result["continuity_id"] == result["continuity_id"][:, :1]).all() + + def test_boundary_sampling_rejects_cross_continuity_segment_boundary( + self, + ) -> None: + midpoint = MAX_EPISODE_STEPS // 2 + self.engine.shared_buffer["segment_id"][:, midpoint:] = 1 + self.engine.shared_buffer["continuity_id"][:, midpoint:] = 1 + + with pytest.raises(RuntimeError, match="No unlocked valid chunk"): + self.engine.sample_batch( + batch_size=1, + chunk_size=8, + sampling_mode="boundary", + ) + def test_no_valid_window_raises(self) -> None: """Sampling fails clearly when all real episodes are too short.""" self.engine.shared_buffer["valid"][:, 3:] = False diff --git a/tests/gym/envs/expert_program/test_bridge.py b/tests/gym/envs/expert_program/test_bridge.py index 378ca4066..7cb258244 100644 --- a/tests/gym/envs/expert_program/test_bridge.py +++ b/tests/gym/envs/expert_program/test_bridge.py @@ -1295,6 +1295,7 @@ def test_zero_command_terminal_runtime_failure_preserves_trace_and_validates_onc assert len(result.segments) == 1 segment_result = result.segments[0] assert segment_result.failure_reason == "segment_validation_failed" + assert segment_result.outcome_kinds == ("runtime_failed", "runtime_failed") runtime_trace = segment_result.metadata["runtime"] assert runtime_trace["status"] == "failed" assert ( @@ -1408,6 +1409,10 @@ def test_post_policy_timeout_is_row_local_and_preserved_in_segment_result() -> N None, "segment_validation_failed", ) + assert result.segments[0].outcome_kinds == ( + "succeeded", + "post_policy_failed", + ) assert result.segments[0].metadata["post_policies"][0]["result_mask"] == [ True, False, diff --git a/tests/gym/envs/managers/test_async_dataset_functors.py b/tests/gym/envs/managers/test_async_dataset_functors.py index ba75c2b19..1423df05d 100644 --- a/tests/gym/envs/managers/test_async_dataset_functors.py +++ b/tests/gym/envs/managers/test_async_dataset_functors.py @@ -114,6 +114,9 @@ def __init__(self, num_envs: int = 2, num_joints: int = 6, steps: int = 5): "segment_step": episode_step.clone(), "segment_start": episode_step == 0, "segment_end": torch.zeros(num_envs, steps, dtype=torch.bool), + "segment_accepted": torch.ones(num_envs, steps, dtype=torch.bool), + "segment_attempt_id": torch.zeros(num_envs, steps, dtype=torch.long), + "continuity_id": torch.zeros(num_envs, steps, dtype=torch.long), "terminated": torch.zeros(num_envs, steps, dtype=torch.bool), "truncated": torch.zeros(num_envs, steps, dtype=torch.bool), }, @@ -200,6 +203,57 @@ def test_call_enqueues_without_blocking(self): assert mock_ds.save_episode_calls == 2 assert len(mock_ds.add_frame_calls) == 8 + def test_call_enqueues_each_accepted_segment_as_independent_fragment(self): + env = _MockEnv(num_envs=1, steps=4) + env.rollout_buffer["segment_id"][0] = torch.tensor([0, 0, 1, 1]) + env.rollout_buffer["segment_step"][0] = torch.tensor([0, 1, 0, 1]) + env.rollout_buffer["segment_start"][0] = torch.tensor( + [True, False, True, False] + ) + env.rollout_buffer["segment_end"][0] = torch.tensor([False, True, False, True]) + env.rollout_buffer["segment_accepted"][0] = torch.tensor( + [True, True, False, False] + ) + env.episode_metadata = { + "output_mode": "segment_fragments", + "save_failed_fragments": False, + "episode_index": 4, + "attempt_id": 0, + "program_run_id": "4:0", + "segments": [ + { + "segment_id": 0, + "start_step": 0, + "end_step": 2, + "success": True, + "instruction": "pick task", + "metadata": {}, + }, + { + "segment_id": 1, + "start_step": 2, + "end_step": 4, + "success": False, + "instruction": "place task", + "metadata": {}, + }, + ], + } + mock_ds = _MockDataset() + recorder = _make_recorder(env, mock_ds) + + recorder(env, env_ids=torch.tensor([0])) + env.current_rollout_step = 0 + recorder.finalize() + + assert mock_ds.save_episode_calls == 1 + assert len(mock_ds.add_frame_calls) == 2 + assert {frame["task"] for frame in mock_ds.add_frame_calls} == {"pick task"} + assert all( + frame["annotation.segment_accepted"].tolist() == [1] + for frame in mock_ds.add_frame_calls + ) + def test_worker_operates_on_clone_not_live_buffer(self): """Mutating the rollout buffer after __call__ must not corrupt the save. diff --git a/tests/gym/envs/managers/test_dataset_functors.py b/tests/gym/envs/managers/test_dataset_functors.py index a8ef7fe0b..41d052730 100644 --- a/tests/gym/envs/managers/test_dataset_functors.py +++ b/tests/gym/envs/managers/test_dataset_functors.py @@ -339,6 +339,13 @@ def test_build_features_creates_correct_structure(self, mock_lerobot_dataset): "shape": (1,), "names": ["segment_id"], } + assert features["annotation.segment_accepted"] == { + "dtype": "int64", + "shape": (1,), + "names": ["segment_accepted"], + } + assert "annotation.segment_attempt_id" in features + assert "annotation.continuity_id" in features @patch("embodichain.lab.gym.envs.managers.datasets.LeRobotDataset") def test_build_features_with_sensor(self, mock_lerobot_dataset): @@ -849,6 +856,120 @@ def test_episode_metadata_sidecar_appends_json_lines(tmp_path) -> None: assert [record["segments"][0]["name"] for record in records] == ["pick", "place"] +@pytest.mark.skipif(not LEROBOT_AVAILABLE, reason="LeRobot not installed") +def test_segment_fragment_payloads_are_independent_and_keep_provenance() -> None: + """Only accepted natural segments are sliced by default.""" + recorder = LeRobotRecorder.__new__(LeRobotRecorder) + obs = TensorDict( + {"state": torch.arange(5, dtype=torch.float32).unsqueeze(-1)}, + batch_size=[5], + ) + actions = torch.arange(5, dtype=torch.float32).unsqueeze(-1) + annotations = { + "valid": torch.ones(5, dtype=torch.bool), + "episode_step": torch.arange(5), + "segment_id": torch.tensor([0, 0, 1, 1, 1]), + "segment_step": torch.tensor([0, 1, 0, 1, 2]), + "segment_start": torch.tensor([True, False, True, False, False]), + "segment_end": torch.tensor([False, True, False, False, True]), + "segment_accepted": torch.tensor([True, True, False, False, False]), + "segment_attempt_id": torch.full((5,), 2), + "continuity_id": torch.zeros(5, dtype=torch.long), + "terminated": torch.zeros(5, dtype=torch.bool), + "truncated": torch.zeros(5, dtype=torch.bool), + } + metadata = { + "output_mode": "segment_fragments", + "save_failed_fragments": False, + "episode_index": 9, + "attempt_id": 2, + "program_run_id": "9:2", + "terminated": False, + "truncated": True, + "segments": [ + { + "segment_id": 0, + "start_step": 0, + "end_step": 2, + "success": True, + "instruction": "Pick the cube", + "attempt_id": 2, + "continuity_id": 0, + "outcome_kind": "succeeded", + "metadata": { + "expert_program_id": "repeat_pick_place", + "program_segment_id": "pick_0", + }, + }, + { + "segment_id": 1, + "start_step": 2, + "end_step": 5, + "success": False, + "failure_reason": "segment_validation_failed", + "outcome_kind": "validation_failed", + "instruction": "Place the cube", + "attempt_id": 2, + "continuity_id": 0, + "metadata": {}, + }, + ], + } + + payloads = list(recorder._episode_payloads(0, obs, actions, annotations, metadata)) + + assert len(payloads) == 1 + _, fragment_obs, fragment_actions, fragment_annotations, fragment_metadata = ( + payloads[0] + ) + assert fragment_obs.batch_size == torch.Size([2]) + assert fragment_actions.shape == (2, 1) + assert fragment_annotations["episode_step"].tolist() == [0, 1] + assert fragment_annotations["segment_start"].tolist() == [True, False] + assert fragment_annotations["segment_end"].tolist() == [False, True] + assert fragment_annotations["segment_accepted"].all() + assert fragment_metadata["fragment_origin"] == "natural_segment" + assert fragment_metadata["source_program_id"] == "repeat_pick_place" + assert fragment_metadata["program_segment_id"] == "pick_0" + assert fragment_metadata["source_start_step"] == 0 + assert fragment_metadata["source_end_step"] == 2 + assert fragment_metadata["terminated"] is False + assert fragment_metadata["truncated"] is False + assert fragment_metadata["segments"][0]["start_step"] == 0 + assert fragment_metadata["segments"][0]["end_step"] == 2 + + +@pytest.mark.skipif(not LEROBOT_AVAILABLE, reason="LeRobot not installed") +def test_failed_segment_fragment_requires_explicit_opt_in() -> None: + recorder = LeRobotRecorder.__new__(LeRobotRecorder) + obs = TensorDict({"state": torch.zeros(2, 1)}, batch_size=[2]) + actions = torch.zeros(2, 1) + annotations = { + "segment_id": torch.zeros(2, dtype=torch.long), + "segment_accepted": torch.zeros(2, dtype=torch.bool), + } + metadata = { + "output_mode": "segment_fragments", + "save_failed_fragments": True, + "segments": [ + { + "segment_id": 0, + "start_step": 0, + "end_step": 2, + "success": False, + "outcome_kind": "runtime_failed", + "metadata": {}, + } + ], + } + + payloads = list(recorder._episode_payloads(0, obs, actions, annotations, metadata)) + + assert len(payloads) == 1 + assert not payloads[0][3]["segment_accepted"].any() + assert payloads[0][4]["success"] is False + + @pytest.mark.skipif(not LEROBOT_AVAILABLE, reason="LeRobot not installed") def test_subtask_registry_writes_stable_deduplicated_indices(tmp_path) -> None: """Repeated descriptions retain one stable row in subtasks.parquet.""" @@ -985,6 +1106,49 @@ def test_post_commit_metadata_failure_does_not_reuse_episode_index() -> None: assert recorder.curr_episode == 1 +@pytest.mark.skipif(not LEROBOT_AVAILABLE, reason="LeRobot not installed") +def test_missing_dense_annotations_fall_back_to_segment_sidecar_outcome() -> None: + """Legacy buffers do not silently label a retained failed segment accepted.""" + recorder = LeRobotRecorder.__new__(LeRobotRecorder) + recorder._env = MockEnvForDataset(has_sensors=False) + recorder.instruction = None + recorder.extra = {} + recorder.total_time = 0.0 + recorder.curr_episode = 0 + recorder.dataset_full_path = Path("/tmp/test_dataset") + recorder.dataset = MagicMock() + recorder.dataset.meta.info = {"fps": 30} + recorder._depth_manager = None + recorder._register_subtasks = MagicMock(return_value={"unknown_task": 0}) + recorder._convert_frame_to_lerobot = MagicMock(return_value={}) + recorder._write_episode_metadata = MagicMock() + + assert recorder._save_single_episode( + 0, + [object()], + [object()], + episode_metadata={ + "attempt_id": 4, + "continuity_id": 2, + "segments": [ + { + "start_step": 0, + "end_step": 1, + "success": False, + "failure_reason": "segment_validation_failed", + } + ], + }, + ) + + frame_annotations = recorder._convert_frame_to_lerobot.call_args.kwargs[ + "annotations" + ] + assert frame_annotations["segment_accepted"] is False + assert frame_annotations["segment_attempt_id"] == 4 + assert frame_annotations["continuity_id"] == 2 + + @pytest.mark.skipif(not LEROBOT_AVAILABLE, reason="LeRobot not installed") def test_save_episodes_skips_empty_rollout() -> None: """An initial reset with no recorded frames is not a failed commit.""" diff --git a/tests/gym/envs/managers/test_dataset_manager.py b/tests/gym/envs/managers/test_dataset_manager.py index 0f9ba0cb8..f1bc61d6f 100644 --- a/tests/gym/envs/managers/test_dataset_manager.py +++ b/tests/gym/envs/managers/test_dataset_manager.py @@ -210,6 +210,32 @@ def test_initialize_episode_saves_failed_reset_envs_when_enabled() -> None: assert torch.equal(manager.saved_env_ids, torch.tensor([1, 2])) +def test_initialize_episode_saves_row_with_accepted_segment_fragment() -> None: + """Fragment eligibility is independent of whole-episode success.""" + env, manager = make_env_for_episode_selection( + save_failed_episodes=False, + successful_env_ids=[], + ) + env._demo_episode_metadata = [ + {"output_mode": "continuous", "segments": []}, + { + "output_mode": "segment_fragments", + "save_failed_fragments": False, + "segments": [{"start_step": 0, "end_step": 2, "success": True}], + }, + {"output_mode": "continuous", "segments": []}, + ] + env._new_demo_episode_metadata = lambda env_id: { + "output_mode": "continuous", + "env_id": env_id, + "segments": [], + } + + EmbodiedEnv._initialize_episode(env, env_ids=[1, 2]) + + assert torch.equal(manager.saved_env_ids, torch.tensor([1])) + + def test_discard_reset_does_not_auto_save_trajectory() -> None: """save_data=False clears trajectory state without writing a file.""" env, _ = make_env_for_episode_selection( diff --git a/tests/gym/envs/test_demo.py b/tests/gym/envs/test_demo.py index 219113bdd..b7cd2db00 100644 --- a/tests/gym/envs/test_demo.py +++ b/tests/gym/envs/test_demo.py @@ -27,6 +27,7 @@ from tensordict import TensorDict from embodichain.lab.gym.envs.demo import ( + DemoExecutionCfg, DemoSegment, DemoSegmentResult, execute_demo_episode, @@ -69,6 +70,71 @@ def test_demo_segment_result_rejects_non_json_metadata() -> None: ) +@pytest.mark.parametrize( + ("validation", "expected"), + [ + ( + { + "runtime_success_mask": [False], + "post_policy_success_mask": None, + "validators": [], + "accepted_mask": [False], + }, + "runtime_failed", + ), + ( + { + "runtime_success_mask": [True], + "post_policy_success_mask": [False], + "validators": [], + "accepted_mask": [False], + }, + "post_policy_failed", + ), + ( + { + "runtime_success_mask": [True], + "post_policy_success_mask": [True], + "validators": [{"result_mask": [False]}], + "accepted_mask": [False], + }, + "validation_failed", + ), + ], +) +def test_demo_segment_result_preserves_first_authoritative_failure_phase( + validation: dict[str, object], expected: str +) -> None: + result = DemoSegmentResult( + segment_id=0, + name="place", + start_step=0, + end_step=2, + success=False, + failure_reason="segment_validation_failed", + metadata={"validation": validation}, + active=(True,), + start_steps=(0,), + end_steps=(2,), + successes=(False,), + failure_reasons=("segment_validation_failed",), + ) + + assert result.outcome_kind == expected + assert result.outcome_kinds == (expected,) + assert result.to_metadata(0)["outcome_kind"] == expected + + +def test_demo_execution_cfg_rejects_unknown_mode() -> None: + with pytest.raises(ValueError, match="mode must be"): + DemoExecutionCfg(mode="resume") # type: ignore[arg-type] + + +def test_demo_execution_cfg_rejects_failed_fragment_policy_in_continuous_mode() -> None: + with pytest.raises(ValueError, match="only valid in segment_fragments"): + DemoExecutionCfg(save_failed_fragments=True) + + def _controller_action_env() -> EmbodiedEnv: env = object.__new__(EmbodiedEnv) env._num_envs = 2 @@ -792,6 +858,9 @@ def _make_rollout_buffer(num_envs: int, steps: int) -> TensorDict: "segment_step": torch.full((num_envs, steps), -1, dtype=torch.long), "segment_start": torch.zeros(num_envs, steps, dtype=torch.bool), "segment_end": torch.zeros(num_envs, steps, dtype=torch.bool), + "segment_accepted": torch.zeros(num_envs, steps, dtype=torch.bool), + "segment_attempt_id": torch.full((num_envs, steps), -1, dtype=torch.long), + "continuity_id": torch.full((num_envs, steps), -1, dtype=torch.long), "terminated": torch.zeros(num_envs, steps, dtype=torch.bool), "truncated": torch.zeros(num_envs, steps, dtype=torch.bool), }, @@ -812,6 +881,8 @@ def __init__(self) -> None: self.rollout_buffer = _make_rollout_buffer(2, 5) self.rollout_steps = torch.tensor([0, 2], dtype=torch.long) self.current_rollout_step = 2 + self._demo_attempt_id = 3 + self._demo_continuity_id = 0 self._demo_active_segment_start_steps = torch.tensor([0, 2]) self.rollout_buffer["obs"]["state"][0, 0] = 10.0 self.rollout_buffer["obs"]["state"][1, 2] = 20.0 @@ -837,6 +908,8 @@ def test_expert_rollout_writer_uses_independent_per_env_lengths() -> None: assert env.rollout_buffer["segment_id"][0, 0].item() == 4 assert env.rollout_buffer["segment_step"][1, 2].item() == 0 assert env.rollout_buffer["segment_end"][1, 2] + assert env.rollout_buffer["segment_attempt_id"][0, 0].item() == 3 + assert env.rollout_buffer["continuity_id"][1, 2].item() == 0 assert torch.equal( env.rollout_buffer["obs"]["state"][0, 0], torch.tensor([10.0, 10.0]) ) @@ -852,6 +925,43 @@ def test_expert_rollout_writer_uses_independent_per_env_lengths() -> None: assert env.current_rollout_step == 3 +def test_end_segment_retroactively_annotates_accepted_frame_spans() -> None: + """The bridge result, not a new evaluator, qualifies every segment frame.""" + env = _RolloutWriterStub() + env._demo_segment_participants = torch.tensor([True, True]) + env._demo_active_segment_start_steps = torch.tensor([0, 0]) + env._demo_active_rollout_start_steps = torch.tensor([0, 0]) + env._demo_steps = torch.tensor([2, 2]) + env.rollout_steps = torch.tensor([2, 2]) + env._demo_episode_metadata = [{"segments": []}, {"segments": []}] + env.rollout_buffer["valid"][:, :2] = True + + result = DemoSegmentResult( + segment_id=0, + name="pick", + start_step=0, + end_step=2, + success=False, + active=(True, True), + start_steps=(0, 0), + end_steps=(2, 2), + successes=(True, False), + failure_reasons=(None, "segment_validation_failed"), + attempt_id=3, + continuity_id=0, + ) + + EmbodiedEnv._end_demo_segment_recording(env, result) + + assert env.rollout_buffer["segment_accepted"][0, :2].all() + assert not env.rollout_buffer["segment_accepted"][1, :2].any() + assert env.rollout_buffer["segment_end"][:, 1].all() + assert env._demo_episode_metadata[0]["segments"][0]["outcome_kind"] == ("succeeded") + assert env._demo_episode_metadata[1]["segments"][0]["outcome_kind"] == ( + "validation_failed" + ) + + def test_expert_rollout_writer_freezes_inactive_demo_row() -> None: """Sticky terminal rows do not receive frames from later shared actions.""" env = _RolloutWriterStub() diff --git a/tests/gym/utils/test_gym_utils.py b/tests/gym/utils/test_gym_utils.py index 6376859a8..ef925b0be 100644 --- a/tests/gym/utils/test_gym_utils.py +++ b/tests/gym/utils/test_gym_utils.py @@ -77,6 +77,10 @@ def test_basic_rollout_buffer(self): # Check actions and rewards assert buffer["actions"].shape == (4, 100, 7) assert buffer["rewards"].shape == (4, 100) + assert buffer["segment_accepted"].shape == (4, 100) + assert not buffer["segment_accepted"].any() + assert (buffer["segment_attempt_id"] == -1).all() + assert (buffer["continuity_id"] == -1).all() def test_extra_observation_with_shape_tuple(self): """Test that extra observations with shape tuple are added correctly.""" diff --git a/tests/lab/scripts/test_run_env.py b/tests/lab/scripts/test_run_env.py index 53ceda2d3..02714a435 100644 --- a/tests/lab/scripts/test_run_env.py +++ b/tests/lab/scripts/test_run_env.py @@ -23,7 +23,11 @@ import pytest import torch -from embodichain.lab.gym.envs.demo import DemoEpisodeResult +from embodichain.lab.gym.envs.demo import ( + DemoEpisodeResult, + DemoExecutionCfg, + DemoSegmentResult, +) from embodichain.lab.gym.utils.gym_utils import merge_args_with_gym_config from embodichain.lab.scripts import run_env from embodichain.lab.scripts.run_env import ( @@ -380,6 +384,64 @@ def test_generate_function_commits_failed_episode_when_configured(monkeypatch) - assert env.reset_options == [None] +def test_generate_function_commits_accepted_prefix_as_segment_fragment( + monkeypatch, +) -> None: + """Fragment mode keeps useful accepted data without claiming episode success.""" + env = _ResetTrackingEnv() + result = DemoEpisodeResult( + episode_index=0, + length=4, + completed=False, + success=(False,), + terminated=(False,), + truncated=(False,), + terminal_reason="segment_validation_failed", + segments=( + DemoSegmentResult( + segment_id=0, + name="pick", + start_step=0, + end_step=2, + success=True, + active=(True,), + start_steps=(0,), + end_steps=(2,), + successes=(True,), + failure_reasons=(None,), + ), + DemoSegmentResult( + segment_id=1, + name="place", + start_step=2, + end_step=4, + success=False, + failure_reason="segment_validation_failed", + active=(True,), + start_steps=(2,), + end_steps=(4,), + successes=(False,), + failure_reasons=("segment_validation_failed",), + ), + ), + lengths=(4,), + ) + monkeypatch.setattr( + "embodichain.lab.scripts.run_env.execute_demo_episode", + lambda *args, **kwargs: result, + ) + + generated = generate_function( + env, + execution_cfg=DemoExecutionCfg(mode="segment_fragments"), + max_attempts=1, + reset_before=False, + ) + + assert generated + assert env.reset_options == [None] + + def test_generate_function_retries_empty_failure_even_when_saving_failures( monkeypatch, ) -> None: