diff --git a/agent_context/topics/env-framework/env-framework.md b/agent_context/topics/env-framework/env-framework.md index da489a69b..f34b1020f 100644 --- a/agent_context/topics/env-framework/env-framework.md +++ b/agent_context/topics/env-framework/env-framework.md @@ -13,6 +13,8 @@ | `embodichain/lab/gym/envs/base_env.py` | `BaseEnv(gym.Env)` + `EnvCfg` — low-level env loop | | `embodichain/lab/gym/envs/types.py` | `ControllerAction` — owned controller-ready action boundary | | `embodichain/lab/gym/envs/embodied_env.py` | `EmbodiedEnv(BaseEnv)` + `EmbodiedEnvCfg` — modular task base class | +| `embodichain/lab/gym/envs/demo.py` | Segment-aware demonstration execution, result, and persistence-mode contracts | +| `embodichain/lab/scripts/run_env.py` | Offline collection retries and explicit dataset commit/abort boundaries | | `embodichain/lab/gym/utils/registration.py` | `@register_env` decorator + `REGISTERED_ENVS` registry + `make()` | | `embodichain/lab/gym/utils/gym_utils.py` | Gym config parsing and config-owned runtime registration | | `embodichain/lab/gym/utils/_component_composition.py` | Reusable physical environment, embodiment, and standalone scene resolution | @@ -74,6 +76,8 @@ gym.Env dataset save, and manager resets. - Overrides `_update_sim_state()` to run event-manager `interval` mode. - Manages rollout buffer (expert or RL mode) via `_hook_after_sim_step()`. +- Owns per-row demonstration metadata and dense segment acceptance, attempt, + and causal-continuity annotations. - `extensions` dict entries are set as attributes on both cfg and env instance. ### Action boundary (`types.py`, `embodied_env.py`) @@ -109,6 +113,10 @@ gym.Env results. - `reset()` lets `BaseEnv` read that final mask before clearing the active bridge, preventing completed state from leaking into the next episode. +- `DemoSegmentResult.successes` is the persisted segment acceptance authority. + `DemoExecutionCfg` keeps continuous episodes as the default and can instead + make each eligible natural segment an independent dataset fragment. This + mode does not restore state or resume after failure. - Environments without a Task Program keep the ordinary `BaseEnv.is_task_success()` behavior and task-specific overrides. @@ -279,9 +287,9 @@ reset(options) ├── is_task_success() → save status before resetting ├── sim.reset_objects_state(env_ids, excluded_uids) ├── _initialize_episode(env_ids) - │ ├── dataset_manager.apply("save") for successful episodes (or an - │ │ explicit dataset-only ``commit_env_ids`` subset during a final - │ │ vector batch) + │ ├── dataset_manager.apply("save") for successful episodes, eligible + │ │ segment-fragment rows, or an explicit dataset-only + │ │ ``commit_env_ids`` subset during a final vector batch │ ├── event_manager.apply("reset", env_ids) │ ├── observation_manager.reset(env_ids) │ └── reward_manager.reset(env_ids) diff --git a/agent_context/topics/manager-functor/manager-functor.md b/agent_context/topics/manager-functor/manager-functor.md index b4d91963d..a6bbd370e 100644 --- a/agent_context/topics/manager-functor/manager-functor.md +++ b/agent_context/topics/manager-functor/manager-functor.md @@ -46,9 +46,20 @@ At init, the manager resolves every `FunctorCfg.func` (string → callable or cl | Recorder | File | Behavior | |----------|------|----------| | `LeRobotRecorder` | `managers/datasets.py` | Synchronous. `__call__` runs convert + `add_frame` + `save_episode` inline, blocking `env.reset()`. Default; base class for the async variant. | -| `AsyncLeRobotRecorder` | `managers/async_datasets.py` | Subclass. `__call__` clones the rollout-buffer slice (obs+actions) to CPU, enqueues it, and returns immediately. A single daemon worker thread drains the queue and runs the same `_save_single_episode` path. `finalize()` drains then calls `dataset.finalize()`. | +| `AsyncLeRobotRecorder` | `managers/async_datasets.py` | Subclass. `__call__` clones the rollout-buffer slice (obs+actions) to CPU, enqueues it, and returns immediately. A single daemon worker thread drains the queue through the same `_persist_episode_payload` path. `finalize()` drains then calls `dataset.finalize()`. | -**Save flow**: `env.step` writes each frame into `rollout_buffer` (`_hook_after_sim_step`). On truncation the caller does `env.reset(options={"save_data": True})` -> `_initialize_episode` -> `DatasetManager.apply("save", env_ids)`. For a final partial vector batch, `run-env` performs one full reset with `save_data=False` and `commit_env_ids`, so only selected **dataset** rows are persisted while whole-world reset events remain safe; camera and trajectory recorders retain their normal discard behavior. `DatasetFunctorCfg.save_failed_episodes=True` saves every env on every ordinary reset (not only successes). `env.close()` -> `dataset_manager.finalize()` flushes any remaining buffer. +**Save flow**: `env.step` writes each frame into `rollout_buffer` (`_hook_after_sim_step`). On truncation the caller does `env.reset(options={"save_data": True})` -> `_initialize_episode` -> `DatasetManager.apply("save", env_ids)`. For a final partial vector batch, `run-env` performs one full reset with `save_data=False` and `commit_env_ids`, so only selected **dataset** rows are persisted while whole-world reset events remain safe; camera and trajectory recorders retain their normal discard behavior. `DatasetFunctorCfg.save_failed_episodes=True` saves every env on every ordinary reset (not only successes). `env.close()` -> `dataset_manager.finalize()` drains explicitly committed async work; it never commits the live rollout implicitly. + +`DemoExecutionCfg(mode="segment_fragments")` changes one buffered Task Program +row into independent natural-segment payloads. Accepted segments are retained +by default; failed segments require `save_failed_fragments=True`. Every +fragment carries dense `segment_accepted`, `segment_attempt_id`, and +`continuity_id` features plus Task Program provenance in the JSONL sidecar. +Fragment commits are append-only: failure of a later fragment does not roll +back earlier ones. The recorder-local deterministic `fragment_id` registry +deduplicates same-run retries. A LeRobot commit followed by sidecar/depth +failure is sticky and rejects retry rather than writing a duplicate. This is +not a cross-process recovery journal. **Two independent speed levers** (both honor `image_writer_threads` / `image_writer_processes` in `params`, wired through to `LeRobotDataset.create()` -> lerobot `AsyncImageWriter`): - Opt A: `LeRobotRecorder` + `image_writer_threads=4` - per-frame PNG writes offloaded to a thread pool. ~2.5x faster, no background thread, bounded memory. @@ -61,6 +72,8 @@ At init, the manager resolves every `FunctorCfg.func` (string → callable or cl **Correctness invariants for the async recorder** (do not break these when editing): - The buffer slice is **cloned in the caller thread** before enqueue - the worker must not hold a view into `rollout_buffer` (it is cleared/reused on reset). - **Single worker** only - `LeRobotDataset` is not thread-safe and FIFO order must be preserved for `episode_index`. +- Duplicate fragment ids pass through `_persist_episode_payload` so sync and + async writers share the same idempotency rule. - `finalize()` must drain the queue before `dataset.finalize()`. - `__call__` accepts `**kwargs` because `DatasetManager.apply` passes `**functor_cfg.params` (includes construction-only params like `image_writer_threads`); `manager_base._resolve_common_functor_cfg` tolerates `**kwargs`. diff --git a/agent_context/topics/task-programs/task-programs.md b/agent_context/topics/task-programs/task-programs.md index 8a1c2b694..8aae251b8 100644 --- a/agent_context/topics/task-programs/task-programs.md +++ b/agent_context/topics/task-programs/task-programs.md @@ -198,6 +198,25 @@ Keep these boundaries separate: reward, reset, and persistence. Final success is published only after every segment lifecycle completes normally. +## Demonstration outcome and persistence + +The bridge's accepted mask remains the sole segment-quality authority. The +common demo executor records it in `DemoSegmentResult.successes`, classifies the +first authoritative failure phase in `outcome_kinds`, and writes +`segment_accepted`, `segment_attempt_id`, and `continuity_id` for every real +rollout frame. Task Program segment metadata uses `task_program_id`; LeRobot +fragment sidecars copy it to the provider-neutral `source_program_id` field. + +`DemoExecutionCfg` is collector-owned, not part of the Task Program language. +Its default `continuous` mode preserves episode semantics. Its +`segment_fragments` mode persists eligible naturally executed segments as +independent LeRobot episodes, with failed fragments requiring explicit opt-in. +Current execution remains fail-closed and every frame has `continuity_id == 0`; +checkpoint capture, state restore, and suffix resume are deferred. Fragment +commits are individually durable and recorder-locally idempotent by stable +`fragment_id`, so a later write failure does not duplicate an earlier fragment +on retry. + ## Parallel and registered calls Parallel execution requires disjoint resource claims and runtime targets, @@ -245,6 +264,8 @@ language. | Live simulation binding | `integrations/simulation/` | | Gym action/segment lifecycle | `gym/envs/task_program/bridge.py` | | Episode program selection/final success | `gym/envs/embodied_env.py` | +| Outcome annotations and persistence mode | `gym/envs/demo.py`, `gym/envs/embodied_env.py` | +| LeRobot fragment slicing/idempotency | `gym/envs/managers/datasets.py`, `async_datasets.py` | ## Focused validation @@ -252,6 +273,7 @@ language. pytest -q tests/lab/task_program pytest -q tests/gym/envs/task_program pytest -q tests/gym/envs/test_embodied_env_task_program.py +pytest -q tests/gym/envs/test_demo.py tests/gym/envs/managers pytest -q tests/agents/mllm/test_task_program.py pytest -q tests/sim/atomic_actions python docs/scripts/check_api_docs.py diff --git a/docs/design/task_program_segment_outcome_and_resume_plan.md b/docs/design/task_program_segment_outcome_and_resume_plan.md new file mode 100644 index 000000000..6a7edd242 --- /dev/null +++ b/docs/design/task_program_segment_outcome_and_resume_plan.md @@ -0,0 +1,507 @@ +# Task Program Segment Outcome and Checkpoint Resume + +- Status: outcome propagation and natural segment fragments implemented; + checkpoint resume deferred +- Scope: Task Program demonstration generation in simulation +- Date: 2026-08-31 +- Related design: `docs/design/task_program_integration_plan.md` + +## 1. Decision + +EmbodiChain should support segment-aware data qualification and optional +checkpoint-based suffix collection, but it should not add success semantics to +atomic-action `TrajectorySegment` values. + +The current layers already have the right semantic boundary: + +- an Atomic Skill verifies one Semantic Call through the semantic executor and + its effect-assurance policy; +- a Task Program segment combines runtime success, post-policy success, and + application validators; +- `DemoSegmentResult` records the resulting per-environment outcome. + +The original gaps were downstream of that decision: + +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 + segment; +4. a discontinuous restore cannot currently be represented without making the + resulting data look like one continuous episode. + +The design therefore has two parts: + +1. persist the existing program-segment acceptance result as trajectory/data + quality metadata; +2. add an opt-in simulation-only fragment collector that may restore a + qualified entry checkpoint for the next segment after the current segment + fails. + +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. + +| Term | Owner | Meaning | Success boundary | +|---|---|---|---| +| Atomic trajectory segment | `ActionPlan.segments` / `TrajectorySegment` | Named frame range such as approach, close, lift, or retreat inside one action | None; the enclosing action owns planning, recovery, and terminal effect verification | +| Program segment | `SegmentCfg` / `CompiledTaskProgramSegment` | Logical transaction containing Semantic Calls, post-policies, and validators | Runtime success AND post-policy success AND all validators | +| Dataset fragment | Demo/trajectory recorder | One continuous, independently usable sequence of state-action transitions | The owning program segment was accepted and no restore occurs inside the fragment | + +This design uses **entry checkpoint of segment K** to mean a physical and +symbolic state from which segment K may begin. For `K > 0`, that checkpoint is +qualified only after segment `K - 1` has succeeded. It is not the terminal +success state of segment K. + +If an application literally wants to start after segment K's terminal state, +that is a different operation (`resume_after(K)`) and skips segment K. It is +outside the first implementation. + +## 3. Existing canonical segment outcome + +No second segment-success evaluator should be introduced. For each +participating environment row, the bridge already has the required inputs: + +```text +runtime_ok + = SkillResult.success_mask + +post_ok + = AND of every post-policy result + +validator_ok + = AND of every compiled segment validator result + +accepted + = participant AND runtime_ok AND post_ok AND validator_ok +``` + +The public outcome should preserve both `accepted` and the first authoritative +failure phase. A boolean alone is insufficient for diagnostics and recovery +policy. + +Implemented stable outcome kinds are: + +```text +succeeded +runtime_failed +post_policy_failed +validation_failed +cancelled +truncated +not_attempted +``` + +``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()`. + +### Open-loop calls + +An atomic call marked `open_loop=True` proves only command completion. A +program segment containing such a call must have an explicit application +validator before its accepted state may qualify a later checkpoint. This rule +is enforced only when checkpoint capture/resume is enabled; it does not change +ordinary Task 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 + +Demonstration rollout annotations include: + +```text +segment_accepted: bool +segment_attempt_id: int64 +continuity_id: int64 +``` + +`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 Task 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. + +String-valued provenance stays in sidecar metadata: + +```text +fragment_id +source_program_id +program_segment_id +``` + +The following fields remain deferred with checkpoint resume: + +```text +checkpoint_id +checkpoint_source +resume_reason +``` + +### 4.2 Sampling + +The online sampler's normal segment mode selects only windows for which: + +```text +valid == true +segment_accepted == true +segment_id is constant +continuity_id is constant +``` + +Boundary sampling must also require a constant `continuity_id`. A state restore +is never a learnable transition and must not appear inside a sampled window. + +### 4.3 Persistence + +Two output modes remain intentionally different: + +- **continuous episode**: current behavior; commit only a causally continuous + episode according to existing episode policy; +- **segment fragments**: every accepted segment is committed as an independent + dataset episode/fragment with source-program provenance. Checkpoint + provenance is added only when resume exists. + +A failed segment followed by a restored successful segment must never be saved +as one successful continuous episode. Episode-level `completed` and `success` +must remain false for a row that crossed a restore boundary, even if the +program cursor later reaches the end. + +Failed fragments are retained only when `save_failed_fragments=True` explicitly +requests them. Continuous failed episodes retain the separate existing +`save_failed_episodes` policy. + +Fragment commits are append-only and independent rather than one source-row +transaction. Every fragment receives a deterministic ``fragment_id`` derived +from ``program_run_id``, environment row, local segment index, attempt, and +continuity region. If a later fragment fails, earlier fragments remain durable. +Within one recorder/dataset, retrying the same source collection skips ids that +already completed. If LeRobot committed an episode but depth or metadata +finalization then failed, that id becomes a sticky partial commit and retry is +rejected instead of silently writing a duplicate. + +This idempotency registry is recorder-local; it is not a cross-process recovery +journal. Callers starting a new recorder receive a new dataset and must still +provide unique logical episode indices. Synchronous writers surface the failing +fragment immediately. Asynchronous writers continue draining independent +fragments and report every failure from ``finalize()``. + +## 5. Checkpoint contract + +### 5.1 Checkpoint value + +Introduce an immutable environment-bound value conceptually equivalent to: + +```python +@dataclass(frozen=True, slots=True) +class SegmentEntryCheckpoint: + schema_version: int + checkpoint_id: str + program_id: str + program_fingerprint: str + segment_id: str + segment_index: int + scene_registry_id: str + robot_profile_id: str + environment_fingerprint: str + physical_state: TensorDictBase + task_state: TaskState + metadata: Mapping[str, JSONValue] +``` + +The exact serialized physical-state type may reuse trajectory-state storage, +but the checkpoint abstraction is stricter than a trajectory frame. + +It must contain or reconstruct: + +- robot root pose, complete qpos/qvel, and controller targets; +- registered articulation root pose, qpos/qvel; +- registered rigid-object pose and linear/angular velocity; +- any task-specific attachment or constraint state required at the boundary; +- the verified semantic `TaskState` used for later call grounding; +- compatibility and provenance fingerprints. + +Camera frames, rewards, runner cursors, pending effects, and command buffers are +not checkpoint state. + +### 5.2 Source of checkpoints + +The entry checkpoint for the next segment cannot be derived from a row whose +current segment failed. It must come from one of: + +1. a previously qualified successful rollout; +2. another compatible successful vector-environment row; +3. an environment-owned deterministic state materializer. + +The first implementation should use a pre-qualified checkpoint store and fail +closed when no compatible entry exists. Cross-row copying is allowed only when +environment and randomization fingerprints match exactly. + +Segment 0 uses the ordinary environment reset state. Entry checkpoint K is +captured immediately before segment K starts, after the prior segment's +post-policy and validator have accepted the source row. + +### 5.3 Runtime port + +Stateful storage and simulator mutation remain outside the declarative Task +Program schema. The fragment executor receives an explicit environment port: + +```python +@runtime_checkable +class SegmentCheckpointPort(Protocol): + def resolve_entry( + self, + *, + program: CompiledProgram, + segment_index: int, + env_mask: torch.Tensor, + ) -> SegmentEntryCheckpoint: + ... + + def restore_entry( + self, + checkpoint: SegmentEntryCheckpoint, + *, + env_mask: torch.Tensor, + ) -> SegmentRestoreResult: + ... +``` + +`SegmentRestoreResult` owns disjoint `restored_mask` and `failed_mask`, the +full-batch merged `TaskState`, and JSON-safe provenance. A partial restore must +not mutate healthy rows. + +## 6. Restore barrier + +A restore is legal only after the previous segment has reached a terminal +boundary and its action iterator, cancellation handshake, post-policies, and +validators are complete. + +For failed rows, the restore sequence is: + +1. confirm that no runtime command, acknowledgement, or effect request is + pending; +2. safe-hold every target armed by the failed segment; +3. resolve and validate the next segment's compatible entry checkpoint; +4. restore physical state for selected rows, including controller targets; +5. synchronize simulation-side caches and obtain a fresh measured scene; +6. install the checkpoint's verified `TaskState` for restored rows; +7. validate the restored entry state through the checkpoint port; +8. reseed pending rollout observations and trajectory pre-action state; +9. increment `continuity_id` and begin a new fragment; +10. construct a fresh Task Program runtime/bridge for the selected segment. + +A fresh runtime is preferred to mutating the completed bridge. It resets +observation-provider baselines, scene revisions, evidence collectors, command +buffers, and clocks while reusing the immutable compiled program and profile +contracts. `SemanticCallExecutor` already accepts an initial `TaskState`; the +Task Program integration assembly needs to expose that existing construction +parameter at the selected-segment bridge boundary. + +No restore step is represented as `env.step()`, and no synthetic action is +written for the state jump. + +## 7. Execution modes + +The implemented collector-owned configuration, rather than a Task Program +field, is: + +```python +@configclass +class DemoExecutionCfg: + mode: Literal["continuous", "segment_fragments"] = "continuous" + save_failed_fragments: bool = False +``` + +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"), +) +``` + +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. + +### Future resume state machine + +```text +READY(k) + -> RUNNING(k) + -> POST_AND_VALIDATE(k) + -> accepted -----------------------> READY(k + 1) + -> failed + continuous mode -------> STOPPED + -> failed + fragment resume + -> resolve entry(k + 1) + -> restore + entry validation + -> restored --------------> READY(k + 1), new continuity_id + -> restore failed --------> STOPPED +``` + +Successful and restored rows may join the same next-segment batch at the +shared segment barrier. Rows that cannot be restored remain inactive. + +## 8. Public result semantics + +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 +recovered_by_env +program_exhausted_by_env +``` + +Each current segment/fragment result additionally records: + +```text +attempt_id +continuity_id +outcome_kind +``` + +``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 demonstration 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. + +## 9. Recommended change sites + +The first implementation should remain outside atomic-action planning. + +| Area | Recommended change | +|---|---| +| `embodichain/lab/gym/envs/demo.py` | Add execution config, outcome/provenance fields, and a fragment-oriented executor path; keep current continuous defaults | +| `embodichain/lab/gym/envs/embodied_env.py` | Retroactively write segment outcome annotations and provide a safe reseed hook after out-of-band state restore | +| `embodichain/lab/gym/utils/trajectory_state.py` | Add row-selective checkpoint capture/restore and controller-target restoration; do not silently claim unsupported constraints | +| `embodichain/lab/gym/envs/task_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/task_program/integrations/environment.py` | Assemble a fresh semantic executor and bridge from an explicit initial `TaskState` and selected segment index | +| `embodichain/lab/task_program/integrations/simulation/` | Implement the checkpoint port and post-restore synchronization/entry validation | +| 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, +and semantic effect monitors require no API change for the initial feature. + +## 10. Delivery phases + +### Phase A: outcome propagation + +1. Expose the bridge's accepted mask as the sole segment outcome. +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 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. +3. Add compatibility fingerprints and restored-entry validation. +4. Support only restore-safe boundaries with no unrepresentable live + constraint state, unless a task-specific materializer handles it. + +### Phase C: fragment collector + +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. **Implemented.** + +## 11. Validation surface + +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; +- a later fragment failure preserving and deduplicating earlier commits; +- post-commit sidecar failure refusing a duplicate fragment write; +- 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; +- fresh runtime receiving the checkpoint `TaskState`; +- a restored Place/Handover entry preserving required held-object state; +- restore entry validation failure remaining terminal; +- restored suffix data being saved as a separate fragment, never a successful + continuous episode; +- safe stop and no pending buffered action at every restore boundary. + +An end-to-end simulation qualification should deliberately fail one middle +segment, restore a pre-qualified next entry, complete the suffix, and verify: + +1. the failed segment is not sampled as successful data; +2. no sampled chunk crosses the state jump; +3. the suffix has the correct segment and checkpoint provenance; +4. the overall continuous episode remains unsuccessful; +5. the restored suffix fragment is independently accepted. + +## 12. First supported task + +`TaskProgramRepeatedPickPlace-v1` is the safest initial qualification target. +Its segment boundary occurs after Place, settling, and object-near-target +validation, where the gripper should no longer own a held-object relation. + +Open Drawer should not be the first checkpoint-resume target: the passive +articulation, handle contact, and open-loop Slide semantics require stronger +entry-state reconstruction and validation. Hardware execution remains out of +scope until a device integration can provide an authoritative restore port. 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 423456640..79d1b3614 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 cdf2f0611..ea1aa59e5 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -276,8 +276,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 a96e88bbe..ff06f3632 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 d23fe8cd9..b8de4176e 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, @@ -435,6 +436,10 @@ def __init__( # 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 @@ -917,8 +922,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: @@ -1027,20 +1042,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, @@ -1051,9 +1104,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) @@ -1116,6 +1180,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 @@ -1238,6 +1313,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, @@ -1344,6 +1426,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..9b1b917fb 100644 --- a/embodichain/lab/gym/envs/managers/async_datasets.py +++ b/embodichain/lab/gym/envs/managers/async_datasets.py @@ -76,7 +76,7 @@ class AsyncLeRobotRecorder(LeRobotRecorder): 5. Returns immediately - the sim is free to reset and keep stepping. A single daemon worker thread drains the queue and runs the standard - :meth:`LeRobotRecorder._save_single_episode` on each cloned payload. + :meth:`LeRobotRecorder._persist_episode_payload` on each cloned payload. Correctness ----------- @@ -146,8 +146,13 @@ def _worker_loop(self) -> None: # Sentinel: finalize() is draining. Exit the worker. break env_id, obs_clone, action_clone, annotations, episode_metadata = item + fragment_id = None + is_fragment = bool( + episode_metadata is not None and episode_metadata.get("fragment", False) + ) try: - saved = self._save_single_episode( + fragment_id = self._fragment_id_from_metadata(episode_metadata) + saved = self._persist_episode_payload( env_id, obs_clone, action_clone, @@ -157,14 +162,19 @@ def _worker_loop(self) -> None: if not saved: self._record_background_error(env_id, "episode save returned False") except BaseException as error: # noqa: BLE001 - worker must not die - self._record_background_error(env_id, str(error)) + label = ( + f"fragment {fragment_id!r}" + if fragment_id is not None + else ("fragment payload" if is_fragment else "episode") + ) + self._record_background_error(env_id, f"{label}: {error}") logger.log_warning( f"[AsyncLeRobotRecorder] Background worker failed on " - f"env {env_id}: {error}" + f"env {env_id} {label}: {error}" ) def _record_background_error(self, env_id: int, message: str) -> None: - """Remember one failed committed episode for the final durability check.""" + """Remember one failed committed payload for the durability barrier.""" with self._background_error_lock: self._background_errors.append((env_id, message)) @@ -237,8 +247,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 +256,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..545cda2dc 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", } @@ -172,6 +176,9 @@ def __init__(self, cfg: DatasetFunctorCfg, env: EmbodiedEnv): self.curr_episode: int = 0 self._metadata_lock = threading.Lock() self._subtask_to_index: dict[str, int] = {} + self._fragment_commit_lock = threading.RLock() + self._committed_fragment_ids: dict[str, int] = {} + self._partial_fragment_commits: dict[str, tuple[int, str]] = {} self._finalize_lock = threading.Lock() self._finalized = False self._finalize_result: Optional[str] = None @@ -258,7 +265,7 @@ def _save_episodes( """Save completed episodes for specified environments. This reads each env's slice from the rollout buffer and delegates to - :meth:`_save_single_episode`. The slice read happens in the caller + :meth:`_persist_episode_payload`. The slice read happens in the caller thread so that subclasses (e.g. :class:`AsyncLeRobotRecorder`) can clone the slice and defer the actual conversion/disk-write to a background worker without racing the buffer reuse on reset. @@ -281,17 +288,277 @@ def _save_episodes( episode_metadata = ( metadata_getter(env_id) if metadata_getter is not None else None ) - saved = self._save_single_episode( + payloads = list( + self._episode_payloads( + env_id, + obs_list, + action_list, + annotations, + 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." + ) + resolved_fragment_ids: list[str] = [] + for payload in payloads: + fragment_id = self._fragment_id_from_metadata(payload[-1]) + try: + saved = self._persist_episode_payload(*payload) + except Exception as error: + if fragment_id is None: + raise + prior = ( + f" Earlier fragments {resolved_fragment_ids!r} remain " + "committed and will be deduplicated on retry." + if resolved_fragment_ids + else "" + ) + raise RuntimeError( + f"Failed to persist fragment {fragment_id!r} for env " + f"{env_id}.{prior}" + ) from error + if not saved: + label = ( + f"fragment {fragment_id!r}" + if fragment_id is not None + else f"episode for env {env_id}" + ) + raise RuntimeError(f"Committed {label} was not persisted.") + if fragment_id is not None: + resolved_fragment_ids.append(fragment_id) + + 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 = {} + source_program_id = segment_provenance.get("task_program_id") + if source_program_id is None: + # Schema-v2 Expert Program sidecars remain readable. + source_program_id = segment_provenance.get("expert_program_id") + 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": source_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, + ) + + @staticmethod + def _fragment_id_from_metadata( + episode_metadata: Mapping[str, Any] | None, + ) -> str | None: + """Return the stable idempotency key for one fragment payload.""" + if episode_metadata is None or not episode_metadata.get("fragment", False): + return None + fragment_id = episode_metadata.get("fragment_id") + if ( + not isinstance(fragment_id, str) + or not fragment_id + or fragment_id != fragment_id.strip() + ): + raise ValueError( + "Fragment metadata must contain a non-empty fragment_id without " + "outer whitespace." + ) + return fragment_id + + def _ensure_fragment_commit_tracking(self) -> None: + """Initialize fragment commit state for lightweight test instances.""" + if not hasattr(self, "_fragment_commit_lock"): + self._fragment_commit_lock = threading.RLock() + if not hasattr(self, "_committed_fragment_ids"): + self._committed_fragment_ids = {} + if not hasattr(self, "_partial_fragment_commits"): + self._partial_fragment_commits = {} + + def _persist_episode_payload( + self, + env_id: int, + obs_list: Any, + action_list: Any, + annotations: Mapping[str, Any] | None = None, + episode_metadata: Mapping[str, Any] | None = None, + ) -> bool: + """Persist one payload, deduplicating completed fragment commits. + + Fragment ids are scoped to the current recorder/dataset. A successful + fragment is an independent commit and is skipped if the same source + collection is retried. A post-commit failure is sticky: the LeRobot + episode already exists, so retrying raises instead of creating a + duplicate with incomplete sidecar durability. + """ + fragment_id = self._fragment_id_from_metadata(episode_metadata) + if fragment_id is None: + return self._save_single_episode( env_id, obs_list, action_list, annotations=annotations, episode_metadata=episode_metadata, ) - if not saved: + + self._ensure_fragment_commit_tracking() + with self._fragment_commit_lock: + partial_commit = self._partial_fragment_commits.get(fragment_id) + if partial_commit is not None: + episode_index, message = partial_commit raise RuntimeError( - f"Committed episode for env {env_id} was not persisted." + f"Fragment {fragment_id!r} already reached LeRobot episode " + f"{episode_index}, but post-commit finalization failed: " + f"{message}. Refusing to write a duplicate." ) + committed_episode = self._committed_fragment_ids.get(fragment_id) + if committed_episode is not None: + logger.log_info( + f"[LeRobotRecorder] Skipping duplicate fragment " + f"{fragment_id!r}; already saved as episode " + f"{committed_episode}." + ) + return True + + episode_index = int(getattr(self, "curr_episode", 0)) + saved = self._save_single_episode( + env_id, + obs_list, + action_list, + annotations=annotations, + episode_metadata=episode_metadata, + ) + if saved: + self._committed_fragment_ids[fragment_id] = episode_index + return saved def _episode_length(self, env_id: int) -> int: """Return the valid buffered length for one environment.""" @@ -332,6 +599,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 +635,9 @@ def _save_single_episode( depth_prefix = f"{LeRobotKey.OBS_PREFIX.value}depth." episode_index = self.curr_episode dataset_committed = False + fragment_id = self._fragment_id_from_metadata(episode_metadata) + 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 +655,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, } @@ -464,6 +756,13 @@ def _save_single_episode( return True except Exception as error: + if dataset_committed and fragment_id is not None: + self._ensure_fragment_commit_tracking() + with self._fragment_commit_lock: + self._partial_fragment_commits[fragment_id] = ( + episode_index, + f"{type(error).__name__}: {error}", + ) if not dataset_committed: self.total_time = previous_total_time if self._depth_manager is not None and not dataset_committed: @@ -533,23 +832,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 5439e2a2d..d5d20d286 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -1331,6 +1331,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 ), @@ -1557,6 +1572,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 5fe517b4a..ac6cda2cf 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.wrapper import ReplayWrapper from embodichain.lab.gym.utils.gym_utils import ( add_env_launcher_args_to_parser, @@ -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,15 +260,18 @@ 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. + """Generate, execute, and commit one demonstration collection batch. A task owns its segment count through ``create_demo_segments``. The legacy ``num_traj`` parameter is accepted only as ``None`` or ``1`` so callers do not accidentally repeat a one-grasp planner inside the same episode. When a dataset functor enables ``save_failed_episodes``, a failed result with at least one frame in every selected row is committed instead of retried. + Continuous mode has one reset commit boundary; fragment mode delegates + independent idempotent fragment commits to the dataset recorder. Args: env: The environment instance. @@ -234,11 +282,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 +303,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 +317,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 44cf41f5a..eca05afed 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], ) @@ -734,6 +740,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/managers/test_async_dataset_functors.py b/tests/gym/envs/managers/test_async_dataset_functors.py index ba75c2b19..60e9408d2 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,87 @@ 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_duplicate_fragment_enqueue_is_idempotent(self): + """The worker writes one stable fragment id at most once.""" + env = _MockEnv(num_envs=1, steps=2) + 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": {}, + } + ], + } + recorder = _make_recorder(env, _MockDataset()) + recorder._save_single_episode = Mock(return_value=True) + + recorder(env, env_ids=torch.tensor([0])) + recorder(env, env_ids=torch.tensor([0])) + env.current_rollout_step = 0 + recorder.finalize() + + recorder._save_single_episode.assert_called_once() + 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..690dab74e 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,175 @@ 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": { + "task_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_fragment_retry_skips_an_earlier_independent_commit() -> None: + """A later fragment failure cannot duplicate an earlier committed slice.""" + env = Mock() + env.rollout_steps = torch.tensor([2]) + env.rollout_buffer = TensorDict( + { + "obs": torch.zeros(1, 2, 1), + "actions": torch.zeros(1, 2, 1), + }, + batch_size=[1, 2], + ) + env.get_demo_episode_metadata.return_value = {"output_mode": "segment_fragments"} + first_payload = ( + 0, + [object()], + [object()], + {}, + {"fragment": True, "fragment_id": "run:0:first"}, + ) + second_payload = ( + 0, + [object()], + [object()], + {}, + {"fragment": True, "fragment_id": "run:0:second"}, + ) + recorder = LeRobotRecorder.__new__(LeRobotRecorder) + recorder._env = env + recorder.curr_episode = 0 + recorder._episode_payloads = Mock(return_value=[first_payload, second_payload]) + recorder._save_single_episode = Mock( + side_effect=[True, OSError("second fragment disk failure")] + ) + + with pytest.raises(RuntimeError, match="run:0:second") as error_info: + recorder._save_episodes(torch.tensor([0])) + + assert "run:0:first" in str(error_info.value) + assert recorder._committed_fragment_ids == {"run:0:first": 0} + + recorder._save_single_episode.reset_mock(side_effect=True) + recorder._save_episodes(torch.tensor([0])) + + recorder._save_single_episode.assert_called_once() + assert recorder._save_single_episode.call_args.kwargs["episode_metadata"] == { + "fragment": True, + "fragment_id": "run:0:second", + } + assert recorder._committed_fragment_ids == { + "run:0:first": 0, + "run:0:second": 0, + } + + @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 +1161,92 @@ 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_post_commit_fragment_failure_is_sticky_and_not_duplicated() -> None: + """A committed fragment with a failed sidecar cannot be written twice.""" + 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(side_effect=OSError("disk full")) + fragment_metadata = { + "fragment": True, + "fragment_id": "run:0:pick", + "segments": [], + } + + with pytest.raises(OSError, match="disk full"): + recorder._persist_episode_payload( + 0, + [object()], + [object()], + episode_metadata=fragment_metadata, + ) + + recorder._write_episode_metadata.side_effect = None + with pytest.raises(RuntimeError, match="Refusing to write a duplicate"): + recorder._persist_episode_payload( + 0, + [object()], + [object()], + episode_metadata=fragment_metadata, + ) + + recorder.dataset.save_episode.assert_called_once_with() + 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 8a3b421c6..a4fb1b64b 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_initialize_episode_commits_only_explicit_vector_rows() -> None: """An explicit commit subset persists only requested dataset rows.""" env, manager = make_env_for_episode_selection( diff --git a/tests/gym/envs/task_program/test_bridge.py b/tests/gym/envs/task_program/test_bridge.py index e0e780a38..b56dbb7ec 100644 --- a/tests/gym/envs/task_program/test_bridge.py +++ b/tests/gym/envs/task_program/test_bridge.py @@ -1428,6 +1428,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 ( @@ -1541,6 +1542,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/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 d080d98eb..458530887 100644 --- a/tests/gym/utils/test_gym_utils.py +++ b/tests/gym/utils/test_gym_utils.py @@ -99,6 +99,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 d0b351ad6..f1d85d2ab 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.task_program.language.loader import ( load_task_program as _load_task_program, ) @@ -381,6 +385,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: