From 9871355447cd4c3d38b28d6546f0bd5550f4e885 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Wed, 5 Aug 2026 21:15:35 +0800 Subject: [PATCH 1/3] Add segment-aware demonstration episodes --- docs/source/features/online_data.md | 40 +- docs/source/guides/run_env.md | 53 ++- .../data_pipeline/datasets/online_data.py | 22 +- embodichain/data_pipeline/engine/data.py | 153 +++++-- embodichain/lab/gym/envs/__init__.py | 1 + embodichain/lab/gym/envs/base_env.py | 5 +- embodichain/lab/gym/envs/demo.py | 381 ++++++++++++++++++ embodichain/lab/gym/envs/embodied_env.py | 335 +++++++++++++-- .../lab/gym/envs/managers/async_datasets.py | 40 +- embodichain/lab/gym/envs/managers/datasets.py | 193 ++++++++- embodichain/lab/gym/utils/gym_utils.py | 66 +++ embodichain/lab/scripts/run_env.py | 132 +++--- tests/data_pipeline/test_online_data.py | 67 ++- .../managers/test_async_dataset_functors.py | 31 +- .../envs/managers/test_dataset_functors.py | 22 +- tests/gym/envs/test_demo.py | 239 +++++++++++ tests/lab/scripts/test_run_env.py | 57 ++- 17 files changed, 1684 insertions(+), 153 deletions(-) create mode 100644 embodichain/lab/gym/envs/demo.py create mode 100644 tests/gym/envs/test_demo.py diff --git a/docs/source/features/online_data.md b/docs/source/features/online_data.md index de54ff569..c13dc8c43 100644 --- a/docs/source/features/online_data.md +++ b/docs/source/features/online_data.md @@ -22,9 +22,28 @@ These components live under `embodichain/data_pipeline/` and are designed to wor Key ideas: -- **Shared buffer**: multiple producers (simulation workers) and multiple consumers (training workers) can read/write concurrently. +- **Shared buffer**: the simulation producer and training consumers can read/write concurrently. - **GPU-friendly**: buffer is designed for efficient sampling and minimal copying. - **Chunked sampling**: training samples fixed-length or dynamically sized chunks. +- **Transactional rows**: the current write window is locked until a complete, + successful episode replaces it. Failed generation attempts never become + sampleable. +- **Variable lengths**: every frame has a `valid` flag. Sampling constructs + windows only from real frames and never reads zero padding or a stale tail. + +Each row stores one complete task episode. A row may contain multiple semantic +segments. In addition to observations, actions, and rewards, the shared buffer +contains: + +```text +valid, episode_step, segment_id, segment_step, +segment_start, segment_end, terminated, truncated +``` + +Tasks use the same `create_demo_segments()` protocol as `run-env`; legacy +`create_demo_action_list()` tasks are treated as one segment. The worker checks +termination after every action and retries a failed episode up to +`max_generation_attempts` before reporting an error. ### Minimal setup @@ -103,6 +122,25 @@ dataset = OnlineDataset(engine, chunk_size=sampler) In batch mode, the sampler is called once per step so all trajectories in the batch share the same chunk length. +### Segment-aware sampling + +Set `sampling_mode` according to the training objective: + +```python +# Chunks may span adjacent subtasks (default). +episode_dataset = OnlineDataset(engine, chunk_size=64, sampling_mode="episode") + +# Every chunk stays inside one pick/place segment. +segment_dataset = OnlineDataset(engine, chunk_size=32, sampling_mode="segment") + +# Every chunk contains an internal transition between two segments. +boundary_dataset = OnlineDataset(engine, chunk_size=32, sampling_mode="boundary") +``` + +All three modes still require every sampled frame to be valid. `boundary` +requires a chunk size of at least two and raises a clear error when no internal +boundary can satisfy the requested length. + --- ## ChunkSizeSampler diff --git a/docs/source/guides/run_env.md b/docs/source/guides/run_env.md index 5d8d8665d..b86676872 100644 --- a/docs/source/guides/run_env.md +++ b/docs/source/guides/run_env.md @@ -128,9 +128,11 @@ details. ## Run and record Without `--preview` or `--replay`, `run-env` enters offline rollout mode. For -each episode, it asks the task for a demonstration action list, applies the -actions through `env.step()`, discards invalid generations, and resets the -environment. `--max_episodes` overrides the value in the gym config: +each episode, it asks the task for its demonstration segments, applies every +action through `env.step()`, and commits the complete episode with one reset. +An empty plan, truncation, or failed final task validation discards the whole +attempt. Attempts are bounded by `demo_max_attempts` in the gym config +(default: 3). `--max_episodes` overrides the value in the gym config: ```bash embodichain run-env \ @@ -145,6 +147,40 @@ Headless execution is normally preferred for throughput. Use `--filter_dataset_saving` for a rollout smoke test that should not create a structured dataset. +### Multi-segment episodes + +An episode is the complete task; a segment is one semantic subtask inside it. +For example, moving three objects is one episode containing three pick/place +segments, even if each segment has its own motion trajectory. The task owns +the number, order, and targets of those segments. The runner only manages the +episode lifecycle, termination checks, retry, and commit/discard boundary. + +Existing tasks implementing `create_demo_action_list()` remain compatible and +are recorded as a single `legacy` segment. A multi-object task can instead +implement a lazy segment planner: + +```python +from embodichain.lab.gym.envs import DemoSegment + + +def create_demo_segments(self): + for object_uid in self.object_order: + # This runs after the preceding segment, so planning sees the latest + # scene state. + actions = self.plan_pick_and_place(object_uid) + yield DemoSegment( + actions=actions, + name="pick_and_place", + target_uid=object_uid, + instruction=f"Place {object_uid} in its target bin", + ) +``` + +The executor checks `terminated` and `truncated` after every action and stops +immediately. It temporarily disables Gym auto-reset, so dataset recording is +transactional: only the explicit reset after successful final validation +saves the episode. + ### Choose the recording output you need EmbodiChain uses "recording" for three related but distinct outputs: @@ -159,7 +195,8 @@ EmbodiChain uses "recording" for three related but distinct outputs: - Intended consumer * - Structured dataset - A dataset manager in the gym config. - - Observations, actions, task metadata, and optionally sensor videos. + - Observations, actions, episode/segment annotations, task metadata, and + optionally sensor videos. - Imitation-learning or data-processing pipelines. * - Debug or demo video - `record_camera_data` or `record_camera_data_async` as an interval event @@ -207,9 +244,11 @@ ${EMBODICHAIN_DATA_ROOT:-~/.cache/embodichain_data}/trajectories// ``` Each file contains `states`, `actions`, and `meta`. Metadata includes the actual -per-environment lengths, timestep, robot identity and DOF, active joint IDs, -recorded object IDs, and original environment IDs. The final directory is also -printed when the run finishes. +per-environment lengths, segment ranges and targets, timestep, robot identity +and DOF, active joint IDs, recorded object IDs, and original environment IDs. +LeRobot exports also contain per-frame `annotation.segment_*` fields and a +`meta/embodichain_episodes.jsonl` sidecar with the complete segment records. +The final directory is also printed when the run finishes. ## Replay a trajectory diff --git a/embodichain/data_pipeline/datasets/online_data.py b/embodichain/data_pipeline/datasets/online_data.py index 647760f5b..36ae36f57 100644 --- a/embodichain/data_pipeline/datasets/online_data.py +++ b/embodichain/data_pipeline/datasets/online_data.py @@ -16,7 +16,7 @@ from __future__ import annotations -from typing import Callable, Iterator, List, Optional +from typing import Callable, Iterator, List, Literal, Optional from tensordict import TensorDict from torch.utils.data import IterableDataset @@ -77,6 +77,9 @@ class OnlineDataset(IterableDataset): ``[batch_size, chunk_size]`` (batch mode). transform: Optional ``(TensorDict) -> TensorDict`` applied to each yielded item/batch before returning. + sampling_mode: ``"episode"`` allows chunks to cross segment boundaries, + ``"segment"`` keeps chunks within one segment, and ``"boundary"`` + samples chunks that cross a boundary. Example — fixed chunk size, item mode:: @@ -123,6 +126,7 @@ def __init__( chunk_size: int | ChunkSizeSampler, batch_size: Optional[int] = None, transform: Optional[Callable[[TensorDict], TensorDict]] = None, + sampling_mode: Literal["episode", "segment", "boundary"] = "episode", ) -> None: if isinstance(chunk_size, int): if chunk_size < 1: @@ -131,10 +135,16 @@ def __init__( raise TypeError( f"chunk_size must be an int or a ChunkSizeSampler, got {type(chunk_size).__name__}." ) + if sampling_mode not in {"episode", "segment", "boundary"}: + raise ValueError( + "sampling_mode must be 'episode', 'segment', or 'boundary', " + f"got {sampling_mode!r}." + ) self._engine = engine self._chunk_size = chunk_size self._batch_size = batch_size self._transform = transform + self._sampling_mode = sampling_mode # ------------------------------------------------------------------ # Internal helpers @@ -177,7 +187,11 @@ def __iter__(self) -> Iterator[TensorDict]: while True: # Item mode: draw one trajectory and remove the outer batch dim. - raw = self._engine.sample_batch(batch_size=1, chunk_size=chunk_size) + raw = self._engine.sample_batch( + batch_size=1, + chunk_size=chunk_size, + sampling_mode=self._sampling_mode, + ) sample: TensorDict = raw[0] if self._transform is not None: @@ -190,7 +204,9 @@ def __iter__(self) -> Iterator[TensorDict]: # Batch mode: draw a full pre-batched TensorDict. sample = self._engine.sample_batch( - batch_size=self._batch_size, chunk_size=chunk_size + batch_size=self._batch_size, + chunk_size=chunk_size, + sampling_mode=self._sampling_mode, ) if self._transform is not None: diff --git a/embodichain/data_pipeline/engine/data.py b/embodichain/data_pipeline/engine/data.py index 1b330b914..5b37a86c6 100644 --- a/embodichain/data_pipeline/engine/data.py +++ b/embodichain/data_pipeline/engine/data.py @@ -20,6 +20,7 @@ import torch import multiprocessing as mp +from typing import Literal from multiprocessing.sharedctypes import Synchronized, SynchronizedArray from multiprocessing.synchronize import Event as MpEvent from tensordict import TensorDict @@ -61,6 +62,9 @@ class OnlineDataEngineCfg: amortising the cost of environment simulation over many training steps. """ + max_generation_attempts: int = 3 + """Maximum planning/execution attempts for each buffer write transaction.""" + # --------------------------------------------------------------------------- # Subprocess entry point (module-level so it can be pickled by multiprocessing) @@ -100,6 +104,7 @@ def _sim_worker_fn( config_to_cfg, get_manager_modules, ) + from embodichain.lab.gym.envs.demo import execute_demo_episode from embodichain.lab.sim import SimulationManagerCfg from embodichain.utils.logger import log_info, log_warning, log_error @@ -165,27 +170,37 @@ def _sim_worker_fn( return tmp_buffer = shared_buffer[lock_index[0] : lock_index[1], :] - env.get_wrapper_attr("set_rollout_buffer")(tmp_buffer) - - _, _ = env.reset() - action_list = env.get_wrapper_attr("create_demo_action_list")() - - if action_list is None or len(action_list) == 0: + result = None + for attempt in range(1, cfg.max_generation_attempts + 1): + # set_rollout_buffer invalidates the locked rows before + # reuse, so stale tail frames can never become sampleable + # after a shorter replacement episode. + env.get_wrapper_attr("set_rollout_buffer")(tmp_buffer) + env.reset(options={"save_data": False}) + result = execute_demo_episode( + env, + episode_index=rollout_idx, + should_stop=close_signal.is_set, + progress=lambda actions, description: tqdm( + actions, + desc=description, + unit="step", + leave=False, + ), + ) + if result.completed and result.all_success: + break log_warning( - f"[Simulation Process] Rollout {rollout_idx + 1}/{num_rollouts_per_fill}: " - "action list is empty, skipping episode." + f"[Simulation Process] Rollout {rollout_idx + 1}/{num_rollouts_per_fill} " + f"attempt {attempt}/{cfg.max_generation_attempts} failed: " + f"{result.terminal_reason}." + ) + + if result is None or not (result.completed and result.all_success): + raise RuntimeError( + f"Failed to generate rollout {rollout_idx + 1} after " + f"{cfg.max_generation_attempts} attempts." ) - continue - - for action in tqdm( - action_list, - desc=f"[Sim] rollout {rollout_idx + 1}/{num_rollouts_per_fill}", - unit="step", - leave=False, - ): - if close_signal.is_set(): - return - env.step(action) rollout_idx += 1 @@ -287,6 +302,12 @@ class OnlineDataEngine: def __init__(self, cfg: OnlineDataEngineCfg) -> None: self.cfg = cfg + if cfg.max_generation_attempts < 1: + raise ValueError( + "max_generation_attempts must be at least 1, " + f"got {cfg.max_generation_attempts}." + ) + # Allocate the shared buffer (shape: [buffer_size, max_episode_steps, ...]). self.shared_buffer: TensorDict = self._create_buffer() self.buffer_size: int = self.shared_buffer.batch_size[0] @@ -412,14 +433,19 @@ def is_init(self) -> bool: # Sampling # ----------------------------------------------------------------------- - def sample_batch(self, batch_size: int, chunk_size: int) -> TensorDict: + def sample_batch( + self, + batch_size: int, + chunk_size: int, + sampling_mode: Literal["episode", "segment", "boundary"] = "episode", + ) -> TensorDict: """Sample a batch of trajectory chunks from the shared rollout buffer. - Randomly draws *batch_size* environment trajectories from the portion - of the buffer that has been written at least once, skipping any rows - currently being overwritten by the simulation subprocess. For each - selected trajectory a contiguous window of *chunk_size* timesteps is - chosen at a uniformly random offset. + 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. After sampling the internal :attr:`_sample_count` is incremented by *batch_size*; if the count exceeds @@ -429,47 +455,86 @@ def sample_batch(self, batch_size: int, chunk_size: int) -> TensorDict: Args: batch_size: Number of trajectory chunks to include in the batch. chunk_size: Number of consecutive timesteps in each chunk. + sampling_mode: Segment-boundary policy for candidate windows. Returns: TensorDict with batch size ``[batch_size, chunk_size]``. Raises: - ValueError: If ``chunk_size`` exceeds ``max_episode_steps``. + ValueError: If an argument is invalid. + RuntimeError: If no unlocked valid window satisfies the policy. """ max_steps: int = self.shared_buffer.batch_size[1] + if batch_size < 1: + raise ValueError(f"batch_size must be at least 1, got {batch_size}.") if chunk_size > max_steps: log_error( f"chunk_size ({chunk_size}) exceeds max_episode_steps ({max_steps}).", error_type=ValueError, ) + if chunk_size < 1: + raise ValueError(f"chunk_size must be at least 1, got {chunk_size}.") + if sampling_mode not in {"episode", "segment", "boundary"}: + raise ValueError( + "sampling_mode must be 'episode', 'segment', or 'boundary', " + f"got {sampling_mode!r}." + ) + if sampling_mode == "boundary" and chunk_size < 2: + raise ValueError("boundary sampling requires chunk_size >= 2.") # Build the set of rows that are safe to sample from: all valid rows # minus the slice currently being written by the subprocess. lock_start: int = self._lock_index[0] lock_end: int = self._lock_index[1] - all_valid = torch.arange(self.buffer_size) - is_locked = (all_valid >= lock_start) & (all_valid < lock_end) - available = all_valid[~is_locked] - - if len(available) == 0: - # Edge case: the entire valid region is locked. Sampling a batch - # is not possible in this state and will result in a hard failure. - log_error( - "[OnlineDataEngine] All valid buffer rows are currently locked. " - "Cannot sample a batch at this time; sampling fails because no " - "unlocked rows are available.", - error_type=RuntimeError, + if "valid" in self.shared_buffer.keys(): + valid = self.shared_buffer["valid"].bool() + else: + # Schema-v1 buffers are one fully valid segment per row. + valid = torch.ones( + self.buffer_size, + max_steps, + dtype=torch.bool, + device=self.shared_buffer.device, ) - # Sample row indices and chunk start offsets. - row_sample_idx = torch.randint(0, len(available), (batch_size,)) - row_indices = available[row_sample_idx] + all_rows = torch.arange(self.buffer_size, device=valid.device) + is_locked = (all_rows >= lock_start) & (all_rows < lock_end) + valid_windows = valid.unfold(1, chunk_size, 1).all(dim=-1) + valid_windows[is_locked] = False + + segment_ids = self.shared_buffer.get("segment_id", None) + if segment_ids is None: + segment_ids = torch.zeros_like(valid, dtype=torch.int64) - max_start = max_steps - chunk_size - start_indices = torch.randint(0, max_start + 1, (batch_size,)) + if sampling_mode == "segment": + segment_windows = segment_ids.unfold(1, chunk_size, 1) + same_segment = (segment_windows == segment_windows[..., :1]).all(dim=-1) & ( + segment_windows[..., 0] >= 0 + ) + valid_windows &= same_segment + elif sampling_mode == "boundary": + segment_windows = segment_ids.unfold(1, chunk_size, 1) + crosses_boundary = ( + segment_windows[..., 1:] != segment_windows[..., :-1] + ).any(dim=-1) + valid_windows &= crosses_boundary + + candidates = valid_windows.nonzero(as_tuple=False) + if candidates.numel() == 0: + raise RuntimeError( + "[OnlineDataEngine] No unlocked valid chunk satisfies " + f"sampling_mode={sampling_mode!r} and chunk_size={chunk_size}." + ) + + sampled_candidate_ids = torch.randint( + 0, candidates.shape[0], (batch_size,), device=candidates.device + ) + selected = candidates[sampled_candidate_ids] + row_indices = selected[:, 0] + start_indices = selected[:, 1] - time_offsets = torch.arange(chunk_size) + time_offsets = torch.arange(chunk_size, device=start_indices.device) time_indices = start_indices[:, None] + time_offsets[None, :] result = self.shared_buffer[row_indices[:, None], time_indices] diff --git a/embodichain/lab/gym/envs/__init__.py b/embodichain/lab/gym/envs/__init__.py index 90f8ef396..14c7e98bf 100644 --- a/embodichain/lab/gym/envs/__init__.py +++ b/embodichain/lab/gym/envs/__init__.py @@ -19,6 +19,7 @@ from __future__ import annotations from .base_env import * +from .demo import * from .embodied_env import * from .wrapper import * diff --git a/embodichain/lab/gym/envs/base_env.py b/embodichain/lab/gym/envs/base_env.py index d8348a818..c84270407 100644 --- a/embodichain/lab/gym/envs/base_env.py +++ b/embodichain/lab/gym/envs/base_env.py @@ -716,7 +716,10 @@ def step( **kwargs, ) - if not getattr(self, "_replay_no_auto_reset", False): + if not ( + getattr(self, "_replay_no_auto_reset", False) + or getattr(self, "_demo_no_auto_reset", False) + ): reset_env_ids = dones.nonzero(as_tuple=False).squeeze(-1) if len(reset_env_ids) > 0: with self._profiler.section("auto_reset"): diff --git a/embodichain/lab/gym/envs/demo.py b/embodichain/lab/gym/envs/demo.py new file mode 100644 index 000000000..7424b4416 --- /dev/null +++ b/embodichain/lab/gym/envs/demo.py @@ -0,0 +1,381 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Segment-aware expert demonstration protocol and executor.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass, field +from typing import Any + +import torch + +__all__ = [ + "DEMO_ANNOTATION_KEYS", + "DEMO_SCHEMA_VERSION", + "DemoEpisodeResult", + "DemoSegment", + "DemoSegmentResult", + "execute_demo_episode", + "resolve_demo_segments", +] + +DEMO_SCHEMA_VERSION = 2 +"""Current version of the segment-aware demonstration metadata schema.""" + +DEMO_ANNOTATION_KEYS = ( + "valid", + "episode_step", + "segment_id", + "segment_step", + "segment_start", + "segment_end", + "terminated", + "truncated", +) +"""Per-frame annotation keys stored in expert rollout buffers.""" + + +@dataclass(frozen=True) +class DemoSegment: + """One semantic subtask inside a demonstration episode. + + The action iterable may be lazy. This lets a task yield one segment, wait + for it to execute, inspect the updated scene, and only then plan the next + segment. + + Args: + actions: Actions for this segment. + name: Stable human-readable segment name. + target_uid: Optional scene entity manipulated by this segment. + instruction: Optional language instruction specific to this segment. + metadata: Additional JSON-compatible task metadata. + """ + + actions: Iterable[Any] + name: str = "segment" + target_uid: str | None = None + instruction: str | None = None + metadata: Mapping[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class DemoSegmentResult: + """Execution result and half-open frame range for one segment.""" + + segment_id: int + name: str + start_step: int + end_step: int + success: bool + target_uid: str | None = None + instruction: str | None = None + failure_reason: str | None = None + metadata: Mapping[str, Any] = field(default_factory=dict) + + def to_metadata(self) -> dict[str, Any]: + """Return a JSON-compatible representation.""" + return { + "segment_id": self.segment_id, + "name": self.name, + "start_step": self.start_step, + "end_step": self.end_step, + "success": self.success, + "target_uid": self.target_uid, + "instruction": self.instruction, + "failure_reason": self.failure_reason, + "metadata": dict(self.metadata), + } + + +@dataclass(frozen=True) +class DemoEpisodeResult: + """Result of executing all planned segments for one batched episode.""" + + episode_index: int + length: int + completed: bool + success: tuple[bool, ...] + terminated: tuple[bool, ...] + truncated: tuple[bool, ...] + terminal_reason: str + segments: tuple[DemoSegmentResult, ...] = () + + @property + def all_success(self) -> bool: + """Whether every parallel environment completed successfully.""" + return bool(self.success) and all(self.success) + + @property + def any_success(self) -> bool: + """Whether at least one parallel environment completed successfully.""" + return any(self.success) + + def to_metadata(self) -> dict[str, Any]: + """Return a JSON-compatible representation.""" + return { + "schema_version": DEMO_SCHEMA_VERSION, + "episode_index": self.episode_index, + "length": self.length, + "completed": self.completed, + "success": list(self.success), + "terminated": list(self.terminated), + "truncated": list(self.truncated), + "terminal_reason": self.terminal_reason, + "segments": [segment.to_metadata() for segment in self.segments], + } + + +ProgressWrapper = Callable[[Iterable[Any], str], Iterable[Any]] +StopPredicate = Callable[[], bool] + + +def _env_target(env: Any) -> Any: + """Return the unwrapped environment when available.""" + return getattr(env, "unwrapped", env) + + +def _get_env_callable(env: Any, name: str) -> Callable[..., Any] | None: + """Resolve an environment method through Gym wrappers when necessary.""" + getter = getattr(env, "get_wrapper_attr", None) + if getter is not None: + try: + value = getter(name) + except AttributeError: + value = None + if callable(value): + return value + + value = getattr(_env_target(env), name, None) + return value if callable(value) else None + + +def _as_bool_tuple(value: Any, num_envs: int) -> tuple[bool, ...]: + """Normalize a scalar, sequence, or tensor to one flag per environment.""" + if value is None: + return (False,) * num_envs + tensor = torch.as_tensor(value, dtype=torch.bool).reshape(-1).cpu() + if tensor.numel() == 1 and num_envs > 1: + tensor = tensor.repeat(num_envs) + if tensor.numel() != num_envs: + raise ValueError( + f"Expected {num_envs} environment flags, got {tensor.numel()}." + ) + return tuple(bool(item) for item in tensor.tolist()) + + +def resolve_demo_segments(env: Any, **kwargs: Any) -> Iterable[DemoSegment]: + """Resolve a task's segment plan with legacy single-action-list fallback. + + Tasks implementing ``create_demo_segments`` own the number, order, and + targets of segments. Older tasks that only implement + ``create_demo_action_list`` are represented as one ``legacy`` segment. + + Args: + env: Gym environment or wrapper. + **kwargs: Planning arguments forwarded to the task method. + + Returns: + A possibly lazy iterable of :class:`DemoSegment` objects. + + Raises: + AttributeError: If the environment exposes neither planning API. + TypeError: If a segment planner yields a value of the wrong type. + """ + creator = _get_env_callable(env, "create_demo_segments") + if creator is not None: + segments = creator(**kwargs) + else: + legacy_creator = _get_env_callable(env, "create_demo_action_list") + if legacy_creator is None: + raise AttributeError( + "Environment must implement create_demo_segments() or " + "create_demo_action_list()." + ) + actions = legacy_creator(**kwargs) + segments = None if actions is None else (DemoSegment(actions, name="legacy"),) + + if segments is None: + return () + if isinstance(segments, DemoSegment): + segments = (segments,) + + def _validate() -> Iterable[DemoSegment]: + for segment in segments: + if not isinstance(segment, DemoSegment): + raise TypeError( + "create_demo_segments() must yield DemoSegment objects, " + f"got {type(segment).__name__}." + ) + yield segment + + return _validate() + + +def execute_demo_episode( + env: Any, + *, + episode_index: int = 0, + should_stop: StopPredicate | None = None, + progress: ProgressWrapper | None = None, + **plan_kwargs: Any, +) -> DemoEpisodeResult: + """Plan and execute every segment in one environment episode. + + Auto-reset is suspended for the duration of execution. The caller owns the + transaction boundary and must explicitly call ``env.reset()`` to commit a + successful episode or ``env.reset(options={"save_data": False})`` to + discard an invalid attempt. + + Args: + env: Gym environment or wrapper. + episode_index: Logical episode identifier used in metadata and logs. + 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. + + Returns: + A :class:`DemoEpisodeResult` describing segment spans and terminal + state. + """ + target = _env_target(env) + num_envs = int(getattr(target, "num_envs", 1)) + begin_episode = _get_env_callable(env, "_begin_demo_episode_recording") + begin_segment = _get_env_callable(env, "_begin_demo_segment_recording") + end_segment = _get_env_callable(env, "_end_demo_segment_recording") + end_episode = _get_env_callable(env, "_end_demo_episode_recording") + + if begin_episode is not None: + begin_episode(episode_index=episode_index) + + previous_no_auto_reset = bool(getattr(target, "_demo_no_auto_reset", False)) + setattr(target, "_demo_no_auto_reset", True) + + total_steps = 0 + segment_results: list[DemoSegmentResult] = [] + terminated = (False,) * num_envs + truncated = (False,) * num_envs + last_info: Mapping[str, Any] = {} + terminal_reason = "completed" + completed = True + + try: + segment_count = 0 + for segment_id, segment in enumerate(resolve_demo_segments(env, **plan_kwargs)): + segment_count += 1 + start_step = total_steps + if begin_segment is not None: + begin_segment(segment_id=segment_id, segment=segment) + + action_count = 0 + segment_reason: str | None = None + actions: Iterable[Any] = segment.actions + if actions is None: + actions = () + if progress is not None: + actions = progress( + actions, f"Executing segment {segment_id}: {segment.name}" + ) + + for action in actions: + if should_stop is not None and should_stop(): + completed = False + terminal_reason = "interrupted" + segment_reason = terminal_reason + break + + _, _, terminated_value, truncated_value, info = env.step(action) + action_count += 1 + total_steps += 1 + last_info = info + terminated = _as_bool_tuple(terminated_value, num_envs) + truncated = _as_bool_tuple(truncated_value, num_envs) + + if any(terminated) or any(truncated): + completed = False + terminal_reason = "truncated" if any(truncated) else "terminated" + segment_reason = terminal_reason + break + + if action_count == 0 and segment_reason is None: + completed = False + terminal_reason = "empty_segment" + segment_reason = terminal_reason + + segment_success = segment_reason is None + if segment_reason == "terminated": + segment_success = all( + _as_bool_tuple(last_info.get("success"), num_envs) + ) + + segment_result = DemoSegmentResult( + segment_id=segment_id, + name=segment.name, + start_step=start_step, + end_step=total_steps, + success=segment_success, + target_uid=segment.target_uid, + instruction=segment.instruction, + failure_reason=None if segment_success else segment_reason, + metadata=segment.metadata, + ) + segment_results.append(segment_result) + if end_segment is not None: + end_segment(result=segment_result) + + if segment_reason is not None: + break + + if segment_count == 0: + completed = False + terminal_reason = "empty_plan" + + success_source = last_info.get("success") + if success_source is None: + success_fn = _get_env_callable(env, "is_task_success") + success_source = success_fn() if success_fn is not None else completed + success = _as_bool_tuple(success_source, num_envs) + + if any(truncated): + success = tuple(False for _ in success) + terminal_reason = "truncated" + elif any(terminated): + if all(success): + completed = True + terminal_reason = "success" + else: + terminal_reason = "failure" + elif completed and all(success): + terminal_reason = "success" + elif completed: + terminal_reason = "task_incomplete" + + result = DemoEpisodeResult( + episode_index=episode_index, + length=total_steps, + completed=completed, + success=success, + terminated=terminated, + truncated=truncated, + terminal_reason=terminal_reason, + segments=tuple(segment_results), + ) + if end_episode is not None: + end_episode(result=result) + return result + finally: + setattr(target, "_demo_no_auto_reset", previous_no_auto_reset) diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index 3c9173c45..c0d399c95 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + from math import log from functools import wraps from datetime import datetime @@ -23,7 +25,7 @@ import gymnasium as gym from dataclasses import MISSING -from typing import Dict, Union, Sequence, Tuple, Any, List, Optional +from typing import Dict, Union, Sequence, Tuple, Any, Iterable, List, Optional from tensordict import TensorDict from embodichain.lab.sim.cfg import ( @@ -43,6 +45,12 @@ from embodichain.lab.sim.sensors import BaseSensor, SensorCfg from embodichain.lab.sim.types import EnvObs, EnvAction from embodichain.lab.gym.envs import BaseEnv, EnvCfg +from embodichain.lab.gym.envs.demo import ( + DEMO_SCHEMA_VERSION, + DemoEpisodeResult, + DemoSegment, + DemoSegmentResult, +) from embodichain.lab.gym.envs.managers import ( EventManager, ObservationManager, @@ -339,8 +347,28 @@ def __init__(self, cfg: EmbodiedEnvCfg, **kwargs): self.num_envs, dtype=torch.long, device=self.device ) + self.rollout_steps = torch.zeros( + self.num_envs, dtype=torch.long, device=self.device + ) + self._demo_steps = torch.zeros( + self.num_envs, dtype=torch.long, device=self.device + ) self.current_rollout_step = 0 + # Segment recording is intentionally separate from task planning. The + # common demo executor updates this context while the regular rollout + # writer turns it into per-frame annotations. + self._demo_episode_index = 0 + self._demo_active_segment_id = 0 + self._demo_active_segment_ids = torch.zeros( + self.num_envs, dtype=torch.long, device=self.device + ) + self._demo_active_segment_start_steps = self._demo_steps.clone() + self._demo_active_rollout_start_steps = self.rollout_steps.clone() + self._demo_episode_metadata: list[dict[str, Any]] = [ + self._new_demo_episode_metadata(env_id) for env_id in range(self.num_envs) + ] + self.episode_success_status: torch.Tensor = torch.zeros( self.num_envs, dtype=torch.bool, device=self.device ) @@ -358,6 +386,12 @@ def set_rollout_buffer(self, rollout_buffer: TensorDict) -> None: transition-only fields is reserved as padding. Expert buffers keep the legacy `[num_envs, time]` batch layout. """ + if rollout_buffer.shape[0] != self.num_envs: + raise ValueError( + "Rollout buffer rows must match env.num_envs: " + f"got {rollout_buffer.shape[0]} rows for {self.num_envs} envs." + ) + self.rollout_buffer = rollout_buffer self._rollout_buffer_mode = self._infer_rollout_buffer_mode(rollout_buffer) if self._rollout_buffer_mode == "rl": @@ -376,7 +410,12 @@ def set_rollout_buffer(self, rollout_buffer: TensorDict) -> None: f"Invalid rollout buffer shape: {rollout_buffer.shape}. The expected shape is (num_envs, max_episode_steps) for each key." ) self._max_rollout_steps = self.rollout_buffer.shape[1] + self.rollout_steps.zero_() self.current_rollout_step = 0 + if self._rollout_buffer_mode != "rl": + self._clear_expert_rollout_rows( + torch.arange(self.num_envs, device=self.device) + ) def _init_sim_state(self, **kwargs): """Initialize the simulation state at the beginning of scene creation.""" @@ -492,8 +531,8 @@ def _hook_after_sim_step( # TODO: We may make the data collection customizable for rollout buffer. if self.rollout_buffer is not None: with self._profiler.section("rollout_write"): - if self.current_rollout_step < self._max_rollout_steps: - if self._rollout_buffer_mode == "rl": + if self._rollout_buffer_mode == "rl": + if self.current_rollout_step < self._max_rollout_steps: self._write_rl_rollout_step( obs=obs, rewards=rewards, @@ -502,17 +541,23 @@ def _hook_after_sim_step( truncateds=kwargs.get("truncateds"), ) else: - self._write_episode_rollout_step( - obs=obs, - action=action, - rewards=rewards, + logger.log_warning( + f"Current rollout step {self.current_rollout_step} exceeds max rollout steps {self._max_rollout_steps}. " + "Data will not be recorded in the rollout buffer." ) + self.current_rollout_step += 1 else: - logger.log_warning( - f"Current rollout step {self.current_rollout_step} exceeds max rollout steps {self._max_rollout_steps}. \ - Data will not be recorded in the rollout buffer." + self._write_episode_rollout_step( + obs=obs, + action=action, + rewards=rewards, + terminateds=kwargs.get("terminateds"), + truncateds=kwargs.get("truncateds"), ) - self.current_rollout_step += 1 + + demo_steps = getattr(self, "_demo_steps", None) + if demo_steps is not None: + demo_steps += 1 with self._profiler.section("trajectory_write"): self._write_trajectory_step() @@ -608,10 +653,6 @@ def _initialize_episode( if isinstance(functor_cfg.func, record_camera_data): functor_cfg.func.save_and_clear() - # Clear episode buffers and reset success status for environments being reset - if self.rollout_buffer is not None and self._rollout_buffer_mode != "rl": - self.current_rollout_step = 0 - # Auto-save + reset the per-env trajectory buffer for environments being # reset. Use getattr so this no-ops on envs/subclasses that don't allocate # a _traj_buffer (e.g. unit-test stubs of _initialize_episode). @@ -627,6 +668,27 @@ def _initialize_episode( if _traj_steps is not None: _traj_steps[env_ids_to_process] = 0 + # Clear episode buffers only after every recorder has consumed them. + if self.rollout_buffer is not None and self._rollout_buffer_mode != "rl": + self._clear_expert_rollout_rows(env_ids_to_process) + rollout_steps = getattr(self, "rollout_steps", None) + if rollout_steps is not None: + rollout_ids = env_ids_to_process.to(rollout_steps.device) + rollout_steps[rollout_ids] = 0 + self.current_rollout_step = int(rollout_steps.max().item()) + + episode_metadata = getattr(self, "_demo_episode_metadata", None) + if episode_metadata is not None: + for env_id in env_ids_to_process.cpu().tolist(): + episode_metadata[env_id] = self._new_demo_episode_metadata(env_id) + active_segment_ids = getattr(self, "_demo_active_segment_ids", None) + if active_segment_ids is not None: + demo_ids = env_ids_to_process.to(active_segment_ids.device) + active_segment_ids[demo_ids] = 0 + self._demo_active_segment_start_steps[demo_ids] = 0 + self._demo_active_rollout_start_steps[demo_ids] = 0 + self._demo_steps[demo_ids] = 0 + self.episode_success_status[env_ids_to_process] = False # apply events such as randomization for environments that need a reset @@ -653,6 +715,139 @@ def _initialize_episode( with self._profiler.section("dataset_reset"): self.dataset_manager.reset(env_ids=env_ids) + def _clear_expert_rollout_rows(self, env_ids: torch.Tensor) -> None: + """Invalidate selected expert-buffer rows without clearing large frames.""" + if self.rollout_buffer is None or len(env_ids) == 0: + return + buffer_ids = env_ids.to(self.rollout_buffer.device, dtype=torch.long) + if "valid" not in self.rollout_buffer.keys(): + # Preserve schema-v1 behavior for external buffers that have no + # validity mask and therefore cannot hide a stale tail. + for key in self.rollout_buffer.keys(include_nested=True, leaves_only=True): + self.rollout_buffer[key][buffer_ids] = 0 + return + + for key in ( + "valid", + "segment_start", + "segment_end", + "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"): + if key in self.rollout_buffer.keys(): + self.rollout_buffer[key][buffer_ids] = -1 + + def _new_demo_episode_metadata(self, env_id: int) -> dict[str, Any]: + """Create an empty metadata record for one environment row.""" + return { + "schema_version": DEMO_SCHEMA_VERSION, + "episode_index": int(getattr(self, "_demo_episode_index", 0)), + "env_id": env_id, + "length": 0, + "completed": False, + "success": False, + "terminated": False, + "truncated": False, + "terminal_reason": "unknown", + "segments": [], + } + + def _begin_demo_episode_recording(self, episode_index: int = 0) -> None: + """Start annotation metadata for a new demonstration episode.""" + self._demo_episode_index = episode_index + self._demo_active_segment_id = 0 + self._demo_active_segment_ids.zero_() + self._demo_active_segment_start_steps = self._demo_steps.clone() + self._demo_active_rollout_start_steps = self.rollout_steps.clone() + self._demo_episode_metadata = [ + self._new_demo_episode_metadata(env_id) for env_id in range(self.num_envs) + ] + + def _begin_demo_segment_recording( + self, segment_id: int, segment: DemoSegment + ) -> None: + """Set the segment context used by subsequent rollout writes.""" + self._demo_active_segment_id = segment_id + self._demo_active_segment_ids.fill_(segment_id) + self._demo_active_segment_start_steps = self._demo_steps.clone() + self._demo_active_rollout_start_steps = self.rollout_steps.clone() + + def _end_demo_segment_recording(self, result: DemoSegmentResult) -> None: + """Close the active segment span and mark its final valid frame.""" + for env_id in range(self.num_envs): + start = int(self._demo_active_segment_start_steps[env_id].item()) + end = int(self._demo_steps[env_id].item()) + rollout_start = int(self._demo_active_rollout_start_steps[env_id].item()) + rollout_end = int(self.rollout_steps[env_id].item()) + if ( + self.rollout_buffer is not None + and "segment_end" in self.rollout_buffer.keys() + and rollout_end > rollout_start + ): + self.rollout_buffer["segment_end"][env_id, rollout_end - 1] = True + + metadata = result.to_metadata() + metadata["start_step"] = start + metadata["end_step"] = end + self._demo_episode_metadata[env_id]["segments"].append(metadata) + + def _end_demo_episode_recording(self, result: DemoEpisodeResult) -> None: + """Finalize per-environment metadata after demonstration execution.""" + for env_id in range(self.num_envs): + metadata = self._demo_episode_metadata[env_id] + metadata.update( + { + "episode_index": result.episode_index, + "length": int(self._demo_steps[env_id].item()), + "completed": result.completed, + "success": result.success[env_id], + "terminated": result.terminated[env_id], + "truncated": result.truncated[env_id], + "terminal_reason": result.terminal_reason, + } + ) + + def get_demo_episode_metadata(self, env_id: int) -> dict[str, Any]: + """Return segment-aware metadata for one buffered episode. + + Legacy collection paths that do not use the common executor are + represented as one segment spanning every valid frame. + + Args: + env_id: Parallel environment row. + + Returns: + A JSON-compatible metadata dictionary. + """ + metadata = dict(self._demo_episode_metadata[env_id]) + metadata["segments"] = [ + dict(segment) + for segment in self._demo_episode_metadata[env_id].get("segments", []) + ] + length = int(self._demo_steps[env_id].item()) + metadata["length"] = length + if length > 0 and not metadata["segments"]: + metadata["segments"] = [ + { + "segment_id": 0, + "name": "legacy", + "start_step": 0, + "end_step": length, + "success": bool( + self.episode_success_status[env_id] + or self._task_success[env_id] + ), + "target_uid": None, + "instruction": None, + "failure_reason": None, + "metadata": {}, + } + ] + return metadata + def _infer_rollout_buffer_mode(self, rollout_buffer: TensorDict) -> str: """Infer whether the rollout buffer is expert recording or RL training data.""" if { @@ -672,11 +867,25 @@ def _write_episode_rollout_step( obs: EnvObs, action: EnvAction, rewards: torch.Tensor, + terminateds: torch.Tensor | None = None, + truncateds: torch.Tensor | None = None, ) -> None: - """Write one step into the legacy episode recording rollout buffer.""" + """Write one expert step at each environment's independent cursor.""" buffer_device = self.rollout_buffer.device - self.rollout_buffer["obs"][:, self.current_rollout_step, ...].copy_( - obs.to(buffer_device), non_blocking=True + active = self.rollout_steps < self._max_rollout_steps + if not bool(active.any()): + logger.log_warning( + "All expert rollout rows reached max_episode_steps; new frames are dropped." + ) + return + + env_ids = active.nonzero(as_tuple=False).squeeze(-1) + step_ids = self.rollout_steps[env_ids] + buffer_env_ids = env_ids.to(buffer_device) + buffer_step_ids = step_ids.to(buffer_device) + + self.rollout_buffer["obs"][buffer_env_ids, buffer_step_ids] = obs[env_ids].to( + buffer_device ) if isinstance(action, TensorDict): action_to_store = ( @@ -693,12 +902,69 @@ def _write_episode_rollout_step( ) action_to_store = None if action_to_store is not None: - self.rollout_buffer["actions"][:, self.current_rollout_step, ...].copy_( - action_to_store.to(buffer_device), non_blocking=True + self.rollout_buffer["actions"][buffer_env_ids, buffer_step_ids] = ( + action_to_store[env_ids].to(buffer_device) ) - self.rollout_buffer["rewards"][:, self.current_rollout_step].copy_( - rewards.to(buffer_device), non_blocking=True - ) + self.rollout_buffer["rewards"][buffer_env_ids, buffer_step_ids] = rewards[ + env_ids + ].to(buffer_device) + + segment_start_steps = getattr( + self, + "_demo_active_rollout_start_steps", + self._demo_active_segment_start_steps, + )[env_ids] + segment_steps = step_ids - segment_start_steps + buffer_keys = set(self.rollout_buffer.keys()) + if "valid" in buffer_keys: + self.rollout_buffer["valid"][buffer_env_ids, buffer_step_ids] = True + if "episode_step" in buffer_keys: + self.rollout_buffer["episode_step"][ + buffer_env_ids, buffer_step_ids + ] = buffer_step_ids + if "segment_id" in buffer_keys: + active_segment_ids = getattr(self, "_demo_active_segment_ids", None) + segment_ids = ( + active_segment_ids[env_ids].to(buffer_device) + if active_segment_ids is not None + else self._demo_active_segment_id + ) + self.rollout_buffer["segment_id"][ + buffer_env_ids, buffer_step_ids + ] = segment_ids + if "segment_step" in buffer_keys: + self.rollout_buffer["segment_step"][buffer_env_ids, buffer_step_ids] = ( + segment_steps.to(buffer_device) + ) + if "segment_start" in buffer_keys: + self.rollout_buffer["segment_start"][buffer_env_ids, buffer_step_ids] = ( + segment_steps.to(buffer_device) == 0 + ) + + if terminateds is None: + terminateds = torch.zeros( + self.num_envs, dtype=torch.bool, device=self.device + ) + if truncateds is None: + truncateds = torch.zeros( + self.num_envs, dtype=torch.bool, device=self.device + ) + terminal = terminateds | truncateds + if "terminated" in buffer_keys: + self.rollout_buffer["terminated"][buffer_env_ids, buffer_step_ids] = ( + terminateds[env_ids].to(buffer_device) + ) + if "truncated" in buffer_keys: + self.rollout_buffer["truncated"][buffer_env_ids, buffer_step_ids] = ( + truncateds[env_ids].to(buffer_device) + ) + if "segment_end" in buffer_keys: + self.rollout_buffer["segment_end"][buffer_env_ids, buffer_step_ids] = ( + terminal[env_ids].to(buffer_device) + ) + + self.rollout_steps[env_ids] += 1 + self.current_rollout_step = int(self.rollout_steps.max().item()) def _write_trajectory_step(self) -> None: """Write one step of per-env ``states`` + pre-process ``action`` into ``_traj_buffer``.""" @@ -1177,6 +1443,26 @@ def create_demo_action_list(self, *args, **kwargs) -> Sequence[EnvAction] | None "The method 'create_demo_action_list' must be implemented in subclasses." ) + def create_demo_segments(self, *args, **kwargs) -> Iterable[DemoSegment] | None: + """Create the semantic segments that make up one task episode. + + The default adapter preserves existing tasks by wrapping their single + ``create_demo_action_list`` result in one segment. Multi-object tasks + should override this method and may return a lazy generator so each + segment can be planned from the scene state left by the previous one. + + Args: + *args: Positional arguments forwarded to the legacy planner. + **kwargs: Keyword arguments forwarded to the legacy planner. + + Returns: + Segment sequence, or ``None`` when planning fails. + """ + actions = self.create_demo_action_list(*args, **kwargs) + if actions is None: + return None + return (DemoSegment(actions=actions, name="task"),) + def save_trajectory(self, path: str, env_ids: Sequence[int] | None = None) -> str: """Save recorded trajectory (states + actions) to a ``.pt`` file. @@ -1214,6 +1500,9 @@ def save_trajectory(self, path: str, env_ids: Sequence[int] | None = None) -> st }, "rigid_object_uids": list(self.sim._rigid_objects.keys()), "env_ids": [int(e) for e in env_ids], + "demo_episodes": [ + self.get_demo_episode_metadata(int(env_id)) for env_id in env_ids + ], } torch.save({"states": states, "actions": actions, "meta": meta}, path) return path diff --git a/embodichain/lab/gym/envs/managers/async_datasets.py b/embodichain/lab/gym/envs/managers/async_datasets.py index 5da4a12f8..521b4a6fd 100644 --- a/embodichain/lab/gym/envs/managers/async_datasets.py +++ b/embodichain/lab/gym/envs/managers/async_datasets.py @@ -26,6 +26,7 @@ from __future__ import annotations +import copy import queue import threading from typing import TYPE_CHECKING, Dict, Optional, Union @@ -33,6 +34,7 @@ import torch from embodichain.utils import logger +from embodichain.lab.gym.envs.demo import DEMO_ANNOTATION_KEYS from .datasets import LeRobotRecorder if TYPE_CHECKING: @@ -67,8 +69,9 @@ class AsyncLeRobotRecorder(LeRobotRecorder): 1. Reads each env's rollout-buffer slice (``obs``/``actions``). 2. **Clones** the slice to CPU (detached from the live buffer). - 3. Pushes ``(env_id, obs_clone, action_clone)`` onto a queue. - 4. Returns immediately - the sim is free to reset and keep stepping. + 3. Clones frame annotations and episode/segment metadata with the payload. + 4. Pushes the detached payload onto a queue. + 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. @@ -133,9 +136,15 @@ def _worker_loop(self) -> None: if item is None: # Sentinel: finalize() is draining. Exit the worker. break - env_id, obs_clone, action_clone = item + env_id, obs_clone, action_clone, annotations, episode_metadata = item try: - self._save_single_episode(env_id, obs_clone, action_clone) + self._save_single_episode( + env_id, + obs_clone, + action_clone, + annotations=annotations, + episode_metadata=episode_metadata, + ) except Exception as e: # noqa: BLE001 - worker must not die logger.log_error( f"[AsyncLeRobotRecorder] Background worker failed on " @@ -180,8 +189,8 @@ def __call__( if len(env_ids) == 0: return - step = env.current_rollout_step for env_id in env_ids.cpu().tolist(): + step = self._episode_length(env_id) obs_view = env.rollout_buffer["obs"][env_id, :step] action_view = env.rollout_buffer["actions"][env_id, :step] # Clone in the caller thread: the rollout buffer is cleared and @@ -189,7 +198,26 @@ def __call__( # a view into it. obs_clone = obs_view.clone().cpu() action_clone = action_view.clone().cpu() - self._save_queue.put((env_id, obs_clone, action_clone)) + annotations = { + key: env.rollout_buffer[key][env_id, :step].clone().cpu() + for key in DEMO_ANNOTATION_KEYS + if key in env.rollout_buffer.keys() + } + metadata_getter = getattr(env, "get_demo_episode_metadata", None) + episode_metadata = ( + copy.deepcopy(metadata_getter(env_id)) + if metadata_getter is not None + else None + ) + self._save_queue.put( + ( + env_id, + obs_clone, + action_clone, + annotations, + episode_metadata, + ) + ) def finalize(self) -> Optional[str]: """Drain the background worker, then finalize the dataset. diff --git a/embodichain/lab/gym/envs/managers/datasets.py b/embodichain/lab/gym/envs/managers/datasets.py index 4ae14dacb..b458e3e62 100644 --- a/embodichain/lab/gym/envs/managers/datasets.py +++ b/embodichain/lab/gym/envs/managers/datasets.py @@ -18,6 +18,10 @@ from __future__ import annotations +import json +import threading + +from collections.abc import Mapping from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, Optional, Union @@ -37,6 +41,7 @@ detect_depth_encoder, ) from embodichain.lab.sim.sensors import Camera, ContactSensor +from embodichain.lab.gym.envs.demo import DEMO_ANNOTATION_KEYS, DEMO_SCHEMA_VERSION from .manager_base import Functor from .cfg import DatasetFunctorCfg @@ -58,6 +63,16 @@ } CAMERA_AUXILIARY_FRAMES = CAMERA_DEPTH_FRAMES | CAMERA_MASK_FRAMES +DEMO_FRAME_FEATURES = { + "episode_step": "annotation.episode_step", + "segment_id": "annotation.segment_id", + "segment_step": "annotation.segment_step", + "segment_start": "annotation.segment_start", + "segment_end": "annotation.segment_end", + "terminated": "annotation.terminated", + "truncated": "annotation.truncated", +} + if TYPE_CHECKING: from embodichain.lab.gym.envs import EmbodiedEnv @@ -75,6 +90,7 @@ class LeRobotRecorder(Functor): """Functor for recording episodes in LeRobot format. This functor handles: + - Recording observation-action pairs during episodes - Converting data to LeRobot format - Saving episodes when they complete @@ -146,6 +162,7 @@ def __init__(self, cfg: DatasetFunctorCfg, env: EmbodiedEnv): # Tracking self.total_time: float = 0.0 self.curr_episode: int = 0 + self._metadata_lock = threading.Lock() # Initialize dataset self._initialize_dataset() @@ -208,17 +225,43 @@ def _save_episodes( clone the slice and defer the actual conversion/disk-write to a background worker without racing the buffer reuse on reset. """ - step = self._env.current_rollout_step for env_id in env_ids.cpu().tolist(): + step = self._episode_length(env_id) obs_list = self._env.rollout_buffer["obs"][env_id, :step] action_list = self._env.rollout_buffer["actions"][env_id, :step] - self._save_single_episode(env_id, obs_list, action_list) + annotations = { + key: self._env.rollout_buffer[key][env_id, :step] + for key in DEMO_ANNOTATION_KEYS + if key in self._env.rollout_buffer.keys() + } + metadata_getter = getattr(self._env, "get_demo_episode_metadata", None) + episode_metadata = ( + metadata_getter(env_id) if metadata_getter is not None else None + ) + self._save_single_episode( + env_id, + obs_list, + action_list, + annotations=annotations, + episode_metadata=episode_metadata, + ) + + def _episode_length(self, env_id: int) -> int: + """Return the valid buffered length for one environment.""" + rollout_steps = getattr(self._env, "rollout_steps", None) + if rollout_steps is not None: + return int(rollout_steps[env_id].item()) + if "valid" in self._env.rollout_buffer.keys(): + return int(self._env.rollout_buffer["valid"][env_id].sum().item()) + return int(self._env.current_rollout_step) def _save_single_episode( self, env_id: int, obs_list: Any, action_list: Any, + annotations: Mapping[str, Any] | None = None, + episode_metadata: Mapping[str, Any] | None = None, ) -> bool: """Convert and persist one episode already sliced from the buffer. @@ -231,6 +274,8 @@ def _save_single_episode( env_id: Environment id (used for logging only). obs_list: Per-frame observations for the episode. action_list: Per-frame actions for the episode. + annotations: Optional per-frame segment and terminal annotations. + episode_metadata: Optional episode/segment sidecar metadata. Returns: True if the episode was saved successfully, False otherwise. @@ -248,6 +293,13 @@ def _save_single_episode( # Align obs and action (obs may be one longer than action) if len(obs_list) > len(action_list): obs_list = obs_list[:-1] + episode_length = min(len(obs_list), len(action_list)) + obs_list = obs_list[:episode_length] + action_list = action_list[:episode_length] + if annotations is not None: + annotations = { + key: values[:episode_length] for key, values in annotations.items() + } # Update metadata extra_info = self.extra.copy() if self.extra else {} @@ -264,12 +316,37 @@ def _save_single_episode( self._depth_manager.start_episode( self.curr_episode, list(self._depth_sensor_specs.keys()) ) - for obs, action in tqdm.tqdm( - zip(obs_list, action_list), - total=len(obs_list), - desc=f"Converting env {env_id} episode to LeRobot format", + for frame_index, (obs, action) in enumerate( + tqdm.tqdm( + zip(obs_list, action_list), + total=len(obs_list), + desc=f"Converting env {env_id} episode to LeRobot format", + ) ): - frame = self._convert_frame_to_lerobot(obs, action, task) + 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, + "terminated": False, + "truncated": False, + } + if annotations is not None: + frame_annotations.update( + { + key: values[frame_index] + for key, values in annotations.items() + if key in DEMO_FRAME_FEATURES + } + ) + frame_task = self._task_for_frame(task, episode_metadata, frame_index) + frame = self._convert_frame_to_lerobot( + obs, + action, + frame_task, + annotations=frame_annotations, + ) # Offload depth to the sidecar writer and drop it from the frame # so LeRobot's RGB-only image/video path never sees it. With # ``keep_numeric_fallback`` the numeric feature is retained too. @@ -286,6 +363,39 @@ def _save_single_episode( if self._depth_manager is not None: self._depth_manager.end_episode(self.curr_episode) + sidecar_metadata = dict(episode_metadata or {}) + if not sidecar_metadata.get("segments"): + sidecar_metadata["segments"] = [ + { + "segment_id": 0, + "name": "legacy", + "start_step": 0, + "end_step": len(obs_list), + "success": True, + "target_uid": None, + "instruction": task, + "failure_reason": None, + "metadata": {}, + } + ] + sidecar_metadata.update(episode_extra_info) + sidecar_metadata.update( + { + "schema_version": DEMO_SCHEMA_VERSION, + "lerobot_episode_index": self.curr_episode, + "env_id": env_id, + "length": len(obs_list), + "instruction": task, + } + ) + try: + self._write_episode_metadata(sidecar_metadata) + except OSError as error: + logger.log_warning( + "Saved the LeRobot episode but could not append EmbodiChain " + f"segment metadata: {error}" + ) + logger.log_info( f"[LeRobotRecorder] Saved dataset to: {self.dataset_path}\n" f" Episode {self.curr_episode} (env {env_id}): {len(obs_list)} frames" @@ -299,11 +409,57 @@ def _save_single_episode( self._depth_manager.abort_episode() return False + @staticmethod + def _task_for_frame( + default_task: str, + episode_metadata: Mapping[str, Any] | None, + frame_index: int, + ) -> str: + """Resolve a segment-specific instruction for one LeRobot frame.""" + if episode_metadata is None: + return default_task + for segment in episode_metadata.get("segments", []): + if ( + int(segment.get("start_step", 0)) + <= frame_index + < int(segment.get("end_step", 0)) + ): + return str(segment.get("instruction") or default_task) + return default_task + + @staticmethod + def _json_default(value: Any) -> Any: + """Convert common tensor/array values for metadata serialization.""" + if isinstance(value, torch.Tensor): + return value.detach().cpu().tolist() + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + return str(value) + + def _write_episode_metadata(self, metadata: Mapping[str, Any]) -> None: + """Append one episode record to EmbodiChain's LeRobot sidecar.""" + if self.dataset_full_path is None: + return + metadata_dir = self.dataset_full_path / "meta" + metadata_dir.mkdir(parents=True, exist_ok=True) + metadata_path = metadata_dir / "embodichain_episodes.jsonl" + with self._metadata_lock, metadata_path.open("a", encoding="utf-8") as stream: + json.dump(dict(metadata), stream, default=self._json_default) + stream.write("\n") + def finalize(self) -> Optional[str]: """Finalize the dataset.""" # Save any remaining episodes - if self._env.current_rollout_step > 0: + rollout_steps = getattr(self._env, "rollout_steps", None) + if rollout_steps is not None: + active_env_ids = (rollout_steps > 0).nonzero(as_tuple=False).squeeze(-1) + elif self._env.current_rollout_step > 0: active_env_ids = torch.arange(self._env.num_envs, device=self._env.device) + else: + active_env_ids = torch.empty(0, dtype=torch.long, device=self._env.device) + if len(active_env_ids) > 0: self._save_episodes(active_env_ids) try: @@ -489,6 +645,13 @@ def _build_features(self) -> Dict: "names": joint_names, } + for feature_key in DEMO_FRAME_FEATURES.values(): + features[feature_key] = { + "dtype": "int64", + "shape": (1,), + "names": [feature_key.rsplit(".", 1)[-1]], + } + # Setup sensor observation features based env.observation.sensor if self._env.has_sensors: sensor_obs_space: dict = self._env.single_observation_space["sensor"] @@ -682,7 +845,11 @@ def _modify_feature_names(self, features: dict[str, Any]) -> None: ) def _convert_frame_to_lerobot( - self, obs: TensorDict, action: TensorDict | torch.Tensor, task: str + self, + obs: TensorDict, + action: TensorDict | torch.Tensor, + task: str, + annotations: Mapping[str, Any] | None = None, ) -> Dict: """Convert a single frame to LeRobot format. @@ -690,6 +857,7 @@ def _convert_frame_to_lerobot( obs: Single environment observation (already extracted from batch) action: Single environment action (already extracted from batch) task: Task name + annotations: Optional segment and terminal fields for this frame. Returns: Frame dict in LeRobot format with numpy arrays @@ -766,6 +934,13 @@ def _convert_frame_to_lerobot( frame[LeRobotKey.ACTION.value] = action_data + if annotations is not None: + for annotation_key, feature_key in DEMO_FRAME_FEATURES.items(): + if annotation_key not in annotations: + continue + value = torch.as_tensor(annotations[annotation_key]).item() + frame[feature_key] = torch.tensor([int(value)], dtype=torch.int64) + return frame def _add_nested_obs_to_frame( diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 4f0dfce39..f9bbaa37b 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -1122,6 +1122,39 @@ def _init_buffer_from_space( "rewards": torch.zeros( (num_envs, max_episode_steps), dtype=torch.float32, device=device ), + "valid": torch.zeros( + (num_envs, max_episode_steps), dtype=torch.bool, device=device + ), + "episode_step": torch.full( + (num_envs, max_episode_steps), + -1, + dtype=torch.int64, + device=device, + ), + "segment_id": torch.full( + (num_envs, max_episode_steps), + -1, + dtype=torch.int64, + device=device, + ), + "segment_step": torch.full( + (num_envs, max_episode_steps), + -1, + dtype=torch.int64, + device=device, + ), + "segment_start": torch.zeros( + (num_envs, max_episode_steps), dtype=torch.bool, device=device + ), + "segment_end": torch.zeros( + (num_envs, max_episode_steps), dtype=torch.bool, device=device + ), + "terminated": torch.zeros( + (num_envs, max_episode_steps), dtype=torch.bool, device=device + ), + "truncated": torch.zeros( + (num_envs, max_episode_steps), dtype=torch.bool, device=device + ), }, batch_size=[num_envs, max_episode_steps], device=device, @@ -1315,6 +1348,39 @@ def init_rollout_buffer_from_config( "rewards": torch.zeros( (batch_size, max_episode_steps), dtype=torch.float32, device=device ), + "valid": torch.zeros( + (batch_size, max_episode_steps), dtype=torch.bool, device=device + ), + "episode_step": torch.full( + (batch_size, max_episode_steps), + -1, + dtype=torch.int64, + device=device, + ), + "segment_id": torch.full( + (batch_size, max_episode_steps), + -1, + dtype=torch.int64, + device=device, + ), + "segment_step": torch.full( + (batch_size, max_episode_steps), + -1, + dtype=torch.int64, + device=device, + ), + "segment_start": torch.zeros( + (batch_size, max_episode_steps), dtype=torch.bool, device=device + ), + "segment_end": torch.zeros( + (batch_size, max_episode_steps), dtype=torch.bool, device=device + ), + "terminated": torch.zeros( + (batch_size, max_episode_steps), dtype=torch.bool, device=device + ), + "truncated": torch.zeros( + (batch_size, max_episode_steps), dtype=torch.bool, device=device + ), }, batch_size=[batch_size, max_episode_steps], device=device, diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index 4fc5e6600..0fcbe7824 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/embodichain/lab/scripts/run_env.py @@ -22,13 +22,15 @@ import sys import time -from collections.abc import Sequence +from collections.abc import Iterable, Sequence +from typing import Any import gymnasium import numpy as np import torch import tqdm +from embodichain.lab.gym.envs.demo import DemoEpisodeResult, 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, @@ -42,47 +44,53 @@ from embodichain.utils.logger import log_warning, log_info, log_error -def generate_and_execute_action_list(env, idx, debug_mode, **kwargs): +def _progress_wrapper(actions: Iterable[Any], description: str) -> Iterable[Any]: + """Wrap a segment action iterable in the run-env progress bar.""" + return tqdm.tqdm(actions, desc=description, unit="step") - action_list = env.get_wrapper_attr("create_demo_action_list")( - action_sentence=idx, **kwargs - ) - - if action_list is None or len(action_list) == 0: - log_warning("Action is invalid. Skip to next generation.") - return False - for action in tqdm.tqdm( - action_list, desc=f"Executing action list #{idx}", unit="step" - ): - # Step the environment with the current action - # The environment will automatically detect truncation based on action_length - obs, reward, terminated, truncated, info = env.step(action) - - # TODO: We may assume in export demonstration rollout, there is no truncation from the env. - # but truncation is useful to improve the generation efficiency. +def generate_and_execute_action_list( + env: Any, idx: int, debug_mode: bool, **kwargs: Any +) -> bool: + """Execute one legacy planner result through the common episode executor. + This compatibility helper now represents one complete task episode. New + multi-object tasks should implement ``create_demo_segments`` instead of + calling this function repeatedly. + """ + result = execute_demo_episode( + env, + episode_index=idx, + progress=_progress_wrapper, + action_sentence=idx, + **kwargs, + ) + if not result.completed or not result.all_success: + log_warning( + f"Demo episode {idx} is invalid ({result.terminal_reason}); it will not be saved." + ) + return False return True def generate_function( - env, - num_traj, + env: Any, + num_traj: int | None = None, time_id: int = 0, save_path: str = "", save_video: bool = False, debug_mode: bool = False, - **kwargs, -): - """Generate and execute a sequence of actions in the environment. + **kwargs: Any, +) -> bool: + """Generate, execute, and transactionally save one task episode. - This function resets the environment, generates and executes action trajectories, - collects data, and optionally saves videos of the episodes. It supports both online - and offline data generation modes. + 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. Args: env: The environment instance. - num_traj (int): Number of trajectories to generate per episode. + num_traj: Deprecated compatibility value. Must be ``None`` or ``1``. time_id (int, optional): Identifier for the current time step or episode. save_path (str, optional): Path to save generated videos. save_video (bool, optional): Whether to save episode videos. @@ -90,30 +98,42 @@ def generate_function( **kwargs: Additional keyword arguments for data generation. Returns: - bool: True if data generation is successful, False otherwise. + True if a successful episode was committed, otherwise False. """ + if num_traj not in (None, 1): + raise ValueError( + "num_traj no longer controls sub-trajectories. Implement " + "create_demo_segments() in the task to define multiple segments." + ) - valid = True - _, _ = env.reset() - while True: - ret = [] - for trajectory_idx in range(num_traj): - valid = generate_and_execute_action_list( - env, trajectory_idx, debug_mode, **kwargs - ) + max_attempts = int(kwargs.pop("max_attempts", 3)) + reset_before = bool(kwargs.pop("reset_before", True)) + if max_attempts < 1: + raise ValueError(f"max_attempts must be at least 1, got {max_attempts}.") - if not valid: - # Failed execution: reset without saving invalid data - _, _ = env.reset(options={"save_data": False}) - break + if reset_before: + env.reset(options={"save_data": False}) - if valid: - break - else: - log_warning("Reset valid flag to True.") - valid = True + for attempt in range(1, max_attempts + 1): + result: DemoEpisodeResult = execute_demo_episode( + env, + episode_index=time_id, + progress=_progress_wrapper, + **kwargs, + ) + if result.completed and result.all_success: + # reset() is the commit boundary: dataset functors consume the + # whole episode once, then buffers and scene state are reset. + env.reset() + return True - return True + log_warning( + f"Episode {time_id} attempt {attempt}/{max_attempts} failed: " + f"{result.terminal_reason}. Discarding {result.length} frames." + ) + env.reset(options={"save_data": False}) + + return False def replay(env, trajectory_path: str, mode: str = "kinematic") -> None: @@ -371,22 +391,26 @@ def main(args, env, gym_config): return log_info("Start offline data generation.", color="green") - # TODO: Support multiple trajectories per episode generation. - num_traj = 1 try: + # Prepare one clean scene. Every successful generate_function call + # commits via reset(), leaving the next episode ready to plan. + env.reset(options={"save_data": False}) for i in range(gym_config.get("max_episodes", 1)): - generate_function( + generated = generate_function( env, - num_traj, - i, + time_id=i, save_path=getattr(args, "save_path", ""), save_video=getattr(args, "save_video", False), debug_mode=getattr(args, "debug_mode", False), regenerate=getattr(args, "regenerate", False), + max_attempts=gym_config.get("demo_max_attempts", 3), + reset_before=False, ) - - # Final reset (saves the last completed episode). - _, _ = env.reset() + if not generated: + raise RuntimeError( + f"Failed to generate episode {i} after " + f"{gym_config.get('demo_max_attempts', 3)} attempts." + ) # Log the trajectory save location BEFORE env.close() (in the finally # below) tears down the sim and, by default, os._exit()s the process. diff --git a/tests/data_pipeline/test_online_data.py b/tests/data_pipeline/test_online_data.py index 4b931e8a6..e2b9cf491 100644 --- a/tests/data_pipeline/test_online_data.py +++ b/tests/data_pipeline/test_online_data.py @@ -104,6 +104,8 @@ def _make_fake_engine( "obs": torch.randn(buffer_size, max_episode_steps, OBS_DIM), "actions": torch.randn(buffer_size, max_episode_steps, ACTION_DIM), "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), }, batch_size=[buffer_size, max_episode_steps], ) @@ -119,8 +121,7 @@ def _make_fake_engine( engine._init_signal.set() # mark as initialised engine._close_signal = engine._mp_ctx.Event() engine._sample_count = engine._mp_ctx.Value("i", 0) - - engine.start() + engine._sim_process = None return engine @@ -197,6 +198,53 @@ def test_chunk_size_exceeds_max_steps_raises(self) -> None: with pytest.raises(ValueError): self.engine.sample_batch(batch_size=1, chunk_size=MAX_EPISODE_STEPS + 1) + def test_sample_batch_never_reads_invalid_tail(self) -> None: + """Variable-length rows never expose padding or stale tail frames.""" + self.engine.shared_buffer["valid"][:, 7:] = False + self.engine.shared_buffer["obs"][:, 7:] = 999.0 + + for _ in range(20): + result = self.engine.sample_batch(batch_size=4, chunk_size=5) + assert result["valid"].all() + assert not (result["obs"] == 999.0).any() + + def test_segment_sampling_never_crosses_boundary(self) -> None: + """Segment mode keeps every sampled chunk inside one subtask.""" + midpoint = MAX_EPISODE_STEPS // 2 + self.engine.shared_buffer["segment_id"][:, midpoint:] = 1 + + result = self.engine.sample_batch( + batch_size=32, + chunk_size=10, + sampling_mode="segment", + ) + + assert (result["segment_id"] == result["segment_id"][:, :1]).all() + + def test_boundary_sampling_crosses_segment_boundary(self) -> None: + """Boundary mode returns chunks containing both adjacent segments.""" + midpoint = MAX_EPISODE_STEPS // 2 + self.engine.shared_buffer["segment_id"][:, midpoint:] = 1 + + result = self.engine.sample_batch( + batch_size=16, + chunk_size=8, + sampling_mode="boundary", + ) + + assert ( + (result["segment_id"][:, 1:] != result["segment_id"][:, :-1]) + .any(dim=1) + .all() + ) + + 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 + + with pytest.raises(RuntimeError, match="No unlocked valid chunk"): + self.engine.sample_batch(batch_size=1, chunk_size=4) + def test_refill_triggered_after_threshold(self) -> None: """_fill_signal is set once accumulated sample count exceeds the threshold.""" # Use a very small threshold so we can trigger it quickly. @@ -312,6 +360,21 @@ def scale_rewards(td: TensorDict) -> TensorDict: sample["rewards"].abs().max().item() > 1.0 ), "scaled rewards should have large absolute values" + def test_sampling_mode_is_forwarded_to_engine(self) -> None: + """OnlineDataset exposes segment-aware engine sampling.""" + midpoint = MAX_EPISODE_STEPS // 2 + self.engine.shared_buffer["segment_id"][:, midpoint:] = 1 + dataset = OnlineDataset( + self.engine, + chunk_size=12, + batch_size=16, + sampling_mode="segment", + ) + + sample = next(iter(dataset)) + + assert (sample["segment_id"] == sample["segment_id"][:, :1]).all() + def test_dataloader_item_mode(self) -> None: """DataLoader with batch_size=4 produces [4, chunk_size] batches.""" BATCH = 4 diff --git a/tests/gym/envs/managers/test_async_dataset_functors.py b/tests/gym/envs/managers/test_async_dataset_functors.py index 53842b073..a69ada8ac 100644 --- a/tests/gym/envs/managers/test_async_dataset_functors.py +++ b/tests/gym/envs/managers/test_async_dataset_functors.py @@ -102,10 +102,35 @@ def __init__(self, num_envs: int = 2, num_joints: int = 6, steps: int = 5): batch_size=[num_envs, steps], ) actions = torch.zeros(num_envs, steps, num_joints) + episode_step = torch.arange(steps).repeat(num_envs, 1) self.rollout_buffer = TensorDict( - {"obs": obs, "actions": actions}, batch_size=[num_envs, steps] + { + "obs": obs, + "actions": actions, + "valid": torch.ones(num_envs, steps, dtype=torch.bool), + "episode_step": episode_step, + "segment_id": torch.zeros(num_envs, steps, dtype=torch.long), + "segment_step": episode_step.clone(), + "segment_start": episode_step == 0, + "segment_end": episode_step == steps - 1, + "terminated": torch.zeros(num_envs, steps, dtype=torch.bool), + "truncated": torch.zeros(num_envs, steps, dtype=torch.bool), + }, + batch_size=[num_envs, steps], ) self.current_rollout_step = steps + self.episode_metadata = { + "segments": [ + { + "start_step": 0, + "end_step": steps, + "instruction": "original segment task", + } + ] + } + + def get_demo_episode_metadata(self, env_id: int): + return self.episode_metadata def _make_recorder(env: _MockEnv, mock_dataset: _MockDataset) -> AsyncLeRobotRecorder: @@ -190,6 +215,8 @@ def test_worker_operates_on_clone_not_live_buffer(self): # Corrupt the live buffer after enqueue (simulates reset reuse). env.rollout_buffer["obs"][0, :3] = 999.0 env.rollout_buffer["actions"][0, :3] = 999.0 + env.rollout_buffer["segment_id"][0, :3] = 999 + env.episode_metadata["segments"][0]["instruction"] = "corrupted" recorder.finalize() @@ -197,6 +224,8 @@ def test_worker_operates_on_clone_not_live_buffer(self): for frame in mock_ds.add_frame_calls: assert (frame[LeRobotKey.OBS_STATE.value] == 0).all() assert (frame[LeRobotKey.ACTION.value] == 0).all() + assert frame["annotation.segment_id"].tolist() == [0] + assert frame["task"] == "original segment task" def test_finalize_drains_and_finalizes_dataset(self): """finalize() must drain the worker then call dataset.finalize().""" diff --git a/tests/gym/envs/managers/test_dataset_functors.py b/tests/gym/envs/managers/test_dataset_functors.py index d1c4aeac0..b0ad3b374 100644 --- a/tests/gym/envs/managers/test_dataset_functors.py +++ b/tests/gym/envs/managers/test_dataset_functors.py @@ -247,6 +247,11 @@ def test_build_features_creates_correct_structure(self, mock_lerobot_dataset): # Check shapes assert features[LeRobotKey.OBS_STATE.value]["shape"] == (6,) assert features[LeRobotKey.ACTION.value]["shape"] == (6,) + assert features["annotation.segment_id"] == { + "dtype": "int64", + "shape": (1,), + "names": ["segment_id"], + } @patch("embodichain.lab.gym.envs.managers.datasets.LeRobotDataset") def test_build_features_with_sensor(self, mock_lerobot_dataset): @@ -645,12 +650,27 @@ def test_convert_frame_with_tensor_action(self, mock_lerobot_dataset): # Create mock action action = torch.zeros(6) - frame = recorder._convert_frame_to_lerobot(obs, action, "test_task") + frame = recorder._convert_frame_to_lerobot( + obs, + action, + "test_task", + annotations={ + "episode_step": 4, + "segment_id": 2, + "segment_step": 1, + "segment_start": False, + "segment_end": True, + "terminated": False, + "truncated": False, + }, + ) assert "task" in frame assert frame["task"] == "test_task" assert LeRobotKey.OBS_STATE.value in frame assert LeRobotKey.ACTION.value in frame + assert frame["annotation.segment_id"].tolist() == [2] + assert frame["annotation.segment_end"].tolist() == [1] @patch("embodichain.lab.gym.envs.managers.datasets.LeRobotDataset") def test_convert_frame_with_depth_and_mask(self, mock_lerobot_dataset): diff --git a/tests/gym/envs/test_demo.py b/tests/gym/envs/test_demo.py new file mode 100644 index 000000000..a3884080c --- /dev/null +++ b/tests/gym/envs/test_demo.py @@ -0,0 +1,239 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for segment-aware demonstration execution and annotations.""" + +from __future__ import annotations + +from typing import Any + +import torch +from tensordict import TensorDict + +from embodichain.lab.gym.envs.demo import DemoSegment, execute_demo_episode +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnv + + +class _SegmentedEnv: + """Small environment stub that supports lazy two-segment planning.""" + + def __init__(self) -> None: + self.num_envs = 1 + self.state = 0 + self.actions: list[int] = [] + self.no_auto_reset_during_steps: list[bool] = [] + self.segment_results = [] + + def create_demo_segments(self): + yield DemoSegment( + actions=(1, 2), + name="pick_a", + target_uid="object_a", + instruction="place object a", + ) + assert self.state == 2 + yield DemoSegment( + actions=(3,), + name="pick_b", + target_uid="object_b", + instruction="place object b", + ) + + def step(self, action: int): + self.actions.append(action) + self.state = action + self.no_auto_reset_during_steps.append(self._demo_no_auto_reset) + success = torch.tensor([action == 3]) + return ( + None, + torch.zeros(1), + torch.zeros(1, dtype=torch.bool), + torch.zeros(1, dtype=torch.bool), + {"success": success}, + ) + + def is_task_success(self) -> torch.Tensor: + return torch.tensor([self.state == 3]) + + def _end_demo_segment_recording(self, result) -> None: + self.segment_results.append(result) + + +def test_execute_demo_episode_runs_lazy_segments_as_one_episode() -> None: + """A task can plan the second object after the first segment executes.""" + env = _SegmentedEnv() + + result = execute_demo_episode(env, episode_index=7) + + assert env.actions == [1, 2, 3] + assert result.length == 3 + assert result.completed + assert result.all_success + assert result.terminal_reason == "success" + assert [(item.start_step, item.end_step) for item in result.segments] == [ + (0, 2), + (2, 3), + ] + assert [item.target_uid for item in result.segments] == ["object_a", "object_b"] + assert all(env.no_auto_reset_during_steps) + assert not env._demo_no_auto_reset + + +class _TerminatingEnv(_SegmentedEnv): + def create_demo_segments(self): + return (DemoSegment(actions=(1, 2, 3), name="pick"),) + + def step(self, action: int): + self.actions.append(action) + self.state = action + self.no_auto_reset_during_steps.append(self._demo_no_auto_reset) + terminated = torch.tensor([action == 2]) + return ( + None, + torch.zeros(1), + terminated, + torch.zeros(1, dtype=torch.bool), + {"success": terminated.clone()}, + ) + + +def test_execute_demo_episode_stops_immediately_on_success_termination() -> None: + """No action after a terminal transition leaks into the next episode.""" + env = _TerminatingEnv() + + result = execute_demo_episode(env) + + assert env.actions == [1, 2] + assert result.length == 2 + assert result.completed + assert result.all_success + assert result.terminal_reason == "success" + + +class _TruncatingEnv(_TerminatingEnv): + def step(self, action: int): + self.actions.append(action) + return ( + None, + torch.zeros(1), + torch.zeros(1, dtype=torch.bool), + torch.ones(1, dtype=torch.bool), + {"success": torch.ones(1, dtype=torch.bool)}, + ) + + +def test_execute_demo_episode_never_accepts_truncated_rollout() -> None: + """Truncation wins over a conflicting success flag and discards the episode.""" + env = _TruncatingEnv() + + result = execute_demo_episode(env) + + assert env.actions == [1] + assert not result.completed + assert not result.any_success + assert result.terminal_reason == "truncated" + + +class _LegacyEnv(_SegmentedEnv): + create_demo_segments = None + + def create_demo_action_list(self): + return (3,) + + +def test_execute_demo_episode_adapts_legacy_action_list() -> None: + """Existing tasks remain a one-segment episode without code changes.""" + env = _LegacyEnv() + + result = execute_demo_episode(env) + + assert result.all_success + assert len(result.segments) == 1 + assert result.segments[0].name == "legacy" + + +def _make_rollout_buffer(num_envs: int, steps: int) -> TensorDict: + return TensorDict( + { + "obs": {"state": torch.zeros(num_envs, steps, 2)}, + "actions": torch.zeros(num_envs, steps, 2), + "rewards": torch.zeros(num_envs, steps), + "valid": torch.zeros(num_envs, steps, dtype=torch.bool), + "episode_step": torch.full((num_envs, steps), -1, dtype=torch.long), + "segment_id": torch.full((num_envs, steps), -1, dtype=torch.long), + "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), + "terminated": torch.zeros(num_envs, steps, dtype=torch.bool), + "truncated": torch.zeros(num_envs, steps, dtype=torch.bool), + }, + batch_size=[num_envs, steps], + ) + + +class _RolloutWriterStub: + """Attributes required by EmbodiedEnv's pure rollout writer method.""" + + num_envs = 2 + _max_rollout_steps = 5 + _demo_active_segment_id = 4 + + 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_active_segment_start_steps = torch.tensor([0, 2]) + + +def test_expert_rollout_writer_uses_independent_per_env_lengths() -> None: + """Partial resets do not overwrite another environment's active episode.""" + env = _RolloutWriterStub() + obs = TensorDict({"state": torch.tensor([[1.0, 1.0], [2.0, 2.0]])}, batch_size=[2]) + + EmbodiedEnv._write_episode_rollout_step( + env, + obs=obs, + action=torch.tensor([[3.0, 3.0], [4.0, 4.0]]), + rewards=torch.tensor([0.5, 1.0]), + terminateds=torch.tensor([False, True]), + truncateds=torch.tensor([False, False]), + ) + + assert env.rollout_steps.tolist() == [1, 3] + assert env.rollout_buffer["valid"][0, 0] + assert env.rollout_buffer["valid"][1, 2] + 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.current_rollout_step == 3 + + +def test_clear_expert_rows_preserves_unrelated_environment() -> None: + """Clearing a completed row is an actual in-place selective mutation.""" + env = _RolloutWriterStub() + for key in env.rollout_buffer.keys(include_nested=True, leaves_only=True): + value: Any = env.rollout_buffer[key] + if value.dtype == torch.bool: + value[:] = True + else: + value[:] = 5 + + EmbodiedEnv._clear_expert_rollout_rows(env, torch.tensor([0])) + + assert not env.rollout_buffer["valid"][0].any() + assert env.rollout_buffer["valid"][1].all() + assert (env.rollout_buffer["segment_id"][0] == -1).all() + assert (env.rollout_buffer["segment_id"][1] == 5).all() diff --git a/tests/lab/scripts/test_run_env.py b/tests/lab/scripts/test_run_env.py index 7950d5a56..4adca539a 100644 --- a/tests/lab/scripts/test_run_env.py +++ b/tests/lab/scripts/test_run_env.py @@ -16,8 +16,11 @@ from __future__ import annotations +import pytest + +from embodichain.lab.gym.envs.demo import DemoEpisodeResult from embodichain.lab.gym.utils.gym_utils import merge_args_with_gym_config -from embodichain.lab.scripts.run_env import _create_parser +from embodichain.lab.scripts.run_env import _create_parser, generate_function GYM_CONFIG_PATH = "task.yaml" GYM_ID = "Dummy-v0" @@ -64,3 +67,55 @@ def test_run_env_preserves_configured_viser_image_fps() -> None: ) assert merged["visualization"]["sensor_image_fps"] == configured_fps + + +class _ResetTrackingEnv: + def __init__(self) -> None: + self.reset_options = [] + + def reset(self, options=None): + self.reset_options.append(options) + return None, {} + + +def _episode_result(*, success: bool, reason: str) -> DemoEpisodeResult: + return DemoEpisodeResult( + episode_index=0, + length=2, + completed=success, + success=(success,), + terminated=(False,), + truncated=(False,), + terminal_reason=reason, + ) + + +def test_generate_function_discards_retry_then_commits_once(monkeypatch) -> None: + """Failed attempts are discarded and only the complete episode is saved.""" + env = _ResetTrackingEnv() + results = iter( + [ + _episode_result(success=False, reason="empty_plan"), + _episode_result(success=True, reason="success"), + ] + ) + monkeypatch.setattr( + "embodichain.lab.scripts.run_env.execute_demo_episode", + lambda *args, **kwargs: next(results), + ) + + generated = generate_function( + env, + time_id=3, + max_attempts=2, + reset_before=False, + ) + + assert generated + assert env.reset_options == [{"save_data": False}, None] + + +def test_generate_function_rejects_runner_owned_segment_count() -> None: + """The task, not run-env, defines how many segments an episode contains.""" + with pytest.raises(ValueError, match="create_demo_segments"): + generate_function(_ResetTrackingEnv(), num_traj=2) From 57768ad767ec2720fde5b0deca050bc9ae29f034 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Thu, 6 Aug 2026 14:40:02 +0000 Subject: [PATCH 2/3] wip --- docs/source/features/online_data.md | 54 +- docs/source/guides/run_env.md | 36 +- docs/source/overview/gym/dataset_functors.md | 20 +- .../data_pipeline/depth_video/writer.py | 33 +- embodichain/data_pipeline/engine/__init__.py | 9 +- embodichain/data_pipeline/engine/data.py | 796 ++++++++++++++++-- embodichain/lab/gym/envs/demo.py | 428 ++++++++-- embodichain/lab/gym/envs/embodied_env.py | 393 +++++++-- .../lab/gym/envs/managers/async_datasets.py | 161 ++-- .../lab/gym/envs/managers/dataset_manager.py | 68 +- embodichain/lab/gym/envs/managers/datasets.py | 147 ++-- embodichain/lab/gym/envs/managers/record.py | 175 ++-- embodichain/lab/scripts/run_env.py | 202 +++-- .../data_pipeline/depth_video/test_writer.py | 15 + tests/data_pipeline/test_online_data.py | 486 ++++++++++- .../managers/test_async_dataset_functors.py | 54 +- .../envs/managers/test_dataset_functors.py | 55 ++ .../gym/envs/managers/test_dataset_manager.py | 224 ++++- tests/gym/envs/test_demo.py | 516 ++++++++++++ tests/gym/envs/test_replay.py | 4 +- tests/lab/scripts/test_run_env.py | 211 +++++ 21 files changed, 3587 insertions(+), 500 deletions(-) diff --git a/docs/source/features/online_data.md b/docs/source/features/online_data.md index f41d02eb3..607ab5947 100644 --- a/docs/source/features/online_data.md +++ b/docs/source/features/online_data.md @@ -33,6 +33,12 @@ Key ideas: - **Episode-uniform sampling**: an eligible episode row is selected uniformly, then a valid start offset is selected within that row. Longer episodes do not receive extra probability merely because they contain more windows. +- **Explicit lifecycle**: the engine progresses through `CREATED`, `STARTING`, + `READY`, and a terminal `FAILED` or `STOPPED` state. Sampling is allowed only + in `READY`, and a stopped or failed instance cannot be restarted. +- **Fail-fast errors**: initial-fill timeouts, rollout failures, hard worker + exits, and shutdown/recorder failures are raised to the owner. The same + stable error snapshot is visible to every forked or spawned data consumer. Each row stores one complete task episode. A row may contain multiple semantic segments. In addition to observations, actions, and rewards, the shared buffer @@ -48,19 +54,22 @@ Tasks use the same `create_demo_segments()` protocol as `run-env`; legacy termination after every action and retries a failed episode up to `max_generation_attempts` before reporting an error. -If the simulation worker exits during initial fill, `start()` raises instead -of waiting indefinitely. A later worker failure is reported on the next -`sample_batch()` call rather than silently serving a permanently stale buffer. +`start()` blocks until the initial fill is complete and is bounded by +`initialization_timeout`. If the worker exits or generation fails, `start()` +raises the worker error instead of waiting indefinitely. A later worker failure +is reported on the next `sample_batch()` call rather than silently serving a +permanently stale buffer. ### Minimal setup ```python -from embodichain.data_pipeline.engine.data import OnlineDataEngine, OnlineDataEngineCfg +from embodichain.data_pipeline.engine import OnlineDataEngine, OnlineDataEngineCfg cfg = OnlineDataEngineCfg( - buffer_size=2, # number of trajectories kept in the ring buffer - state_dim=6, # example state dimension - gym_config=your_gym_cfg, # parsed gym config for the task (JSON or YAML) + buffer_size=2, # trajectories kept in the shared buffer + state_dim=6, # example state dimension + gym_config=your_gym_cfg, # parsed gym config for the task + initialization_timeout=300, # maximum seconds for the initial fill ) engine = OnlineDataEngine(cfg) engine.start() @@ -68,10 +77,39 @@ engine.start() ### Shutdown +Prefer a context manager so cleanup runs on both success and failure: + ```python -engine.stop() +with OnlineDataEngine(cfg) as engine: + batch = engine.sample_batch(batch_size=32, chunk_size=64) + train_step(batch) ``` +For an explicitly managed lifecycle, call `stop()` in the process that created +the engine: + +```python +engine = OnlineDataEngine(cfg) +engine.start() +try: + train(engine) +finally: + engine.stop() +``` + +`stop()` is idempotent after a successful cleanup. It waits for the producer +and raises any failure reported while the worker closes its environment or +flushes committed data. If graceful shutdown times out and the worker requires +`terminate()` or `kill()`, durability cannot be confirmed: the engine remains +`FAILED` and `stop()` raises instead of reporting a false `STOPPED` state. When +a `with` body and cleanup both fail, the body exception remains primary and the +cleanup failure is attached as a note. + +Start the engine before constructing multiprocessing `DataLoader` workers. +Forked and spawned copies may call `sample_batch()`, but only the original +owner process may call `start()` or `stop()`; consumer destructors never signal +the shared producer. + --- ## OnlineDataset diff --git a/docs/source/guides/run_env.md b/docs/source/guides/run_env.md index 70a0731f1..2c43b8143 100644 --- a/docs/source/guides/run_env.md +++ b/docs/source/guides/run_env.md @@ -173,13 +173,23 @@ def create_demo_segments(self): name="pick_and_place", target_uid=object_uid, instruction=f"Place {object_uid} in its target bin", + # Segment validation is separate from Gym episode termination. + validator=lambda uid=object_uid: self.is_object_placed(uid), ) ``` -The executor checks `terminated` and `truncated` after every action and stops -immediately. It temporarily disables Gym auto-reset, so dataset recording is -transactional: only the explicit reset after successful final validation -saves the episode. +Gym `terminated` and `truncated` always describe the whole episode, never an +individual segment. A segment normally ends when its action iterable is +exhausted; its optional zero-argument `validator` then returns one boolean per +parallel environment. A failed validator aborts the batch. Episode-level +success termination stops the remaining lazy plan without requesting another +segment. + +The executor checks terminal signals after every action and temporarily +disables Gym auto-reset. Dataset recording is transactional: only the explicit +reset after successful final validation saves the episode. Exceptions, +interrupts, failed attempts, and closing an environment with a live rollout +abort pending structured data, videos, and trajectories. Segment actions pass through the same action-dimension normalization used by legacy `create_demo_action_list()` tasks. A time-limit truncation is always an @@ -187,12 +197,13 @@ invalid expert rollout, including when it occurs on the planner's final action, so a task's `max_episode_steps` must be greater than the longest valid expert plan. -In a vectorized environment, the shared executor currently treats the batch -as one transaction: the first environment to terminate or truncate stops the -batch, and every environment must pass final validation before the reset is -committed. Tasks whose parallel environments finish asynchronously should use -`num_envs: 1` for expert generation until masked per-environment execution is -supported. +In a vectorized environment, segments and actions remain on one shared planner +clock, but completion is tracked independently. When one environment reports +success, its terminal result and recording cursor become sticky; subsequent +shared actions use a safe hold/no-op command for that row while unfinished rows +continue. Consequently, rollout and trajectory lengths may differ by row. +The commit remains batch-atomic: every row must eventually succeed, while any +failure or truncation aborts the whole attempt. Failed attempts use `reset(options={"save_data": False})`. This discards structured datasets, replay trajectories, and camera-video buffers; recorder @@ -251,8 +262,9 @@ embodichain run-env \ The recorder stores the robot root pose and complete joint position, the raw action before ActionManager preprocessing, and the pose or joint state of scene rigid objects and articulations. Each environment in a vectorized rollout is -tracked independently. A trajectory is auto-saved at episode reset and again -on close if an unfinished buffer still contains steps. +tracked independently. A trajectory is saved only at an explicit successful +episode reset. `close()` is a durability barrier for already committed writes, +not an implicit commit, so an unfinished trajectory is discarded. Files are named like `traj_env0_000000.pt`. When `--trajectory_save_dir` is omitted, they are written below: diff --git a/docs/source/overview/gym/dataset_functors.md b/docs/source/overview/gym/dataset_functors.md index 6acde58a6..9d24e3f70 100644 --- a/docs/source/overview/gym/dataset_functors.md +++ b/docs/source/overview/gym/dataset_functors.md @@ -221,6 +221,14 @@ The sync stall grows **linearly** with ``num_envs`` (each reset saves N envs ser ```{attention} - ``AsyncLeRobotRecorder`` clones each finished episode (including camera frames) to CPU before enqueuing. For very high resolutions or many envs, monitor RSS — the worker normally keeps up, but a slow disk can let the queue grow. - A single background worker touches the ``LeRobotDataset`` (which is not thread-safe), so episode order is preserved FIFO. Always let ``env.close()`` / ``dataset_manager.finalize()`` run so the worker drains before the dataset is finalized. +- Finalization is a durability barrier for episodes already submitted through + ``mode="save"``. It does not turn a live rollout into an episode; pending + frames are discarded. Background and storage failures are raised to the + caller, and repeated finalization is safe. +- Depth video and EmbodiChain metadata sidecars are part of the same durability + result: encoder-close or metadata-write failures are surfaced even when the + corresponding LeRobot episode has already committed, and its episode index + is never reused by a later queued write. - ``env.close()`` calls ``sim.destroy()``, which exits the process without returning to Python. In scripts that build multiple envs, run each in its own subprocess and write results before closing. ``` @@ -260,15 +268,15 @@ dataset = { 1. **Initialization**: The Dataset Manager initializes the functor with the configured parameters 2. **Data Collection**: During episode rollout, the functor receives observations and actions 3. **Save Trigger**: When an episode completes, call the functor with `mode="save"` -4. **Finalization**: After all episodes, call `finalize()` to save any remaining data +4. **Finalization**: After all episodes, call `finalize()` to drain committed writes and finalize storage metadata ```python # Inside environment loop if episode_done: dataset_manager.apply(mode="save", env_ids=completed_env_ids) -# After training completes -dataset_manager.apply(mode="finalize") +# After collection completes. This raises if a committed write failed. +dataset_manager.finalize() ``` ### Parallel Collection (async recorder) @@ -293,7 +301,9 @@ dataset = { } ``` -The async recorder drains its background worker during ``finalize()``, so make sure ``env.close()`` (or ``dataset_manager.finalize()`) runs at the end of collection. +The async recorder drains its background worker during ``finalize()``, so make +sure ``env.close()`` (or ``dataset_manager.finalize()``) runs at the end of +collection. No new episode can be submitted after finalization begins. ### Compressed depth recording @@ -334,6 +344,6 @@ dataset = { The Dataset Manager supports the following modes: - ``save``: Save completed episodes for specified environment IDs -- ``finalize``: Finalize the dataset and save any remaining data +- ``finalize``: Drain explicitly committed writes and finalize dataset resources See {class}`~managers.dataset_manager.DatasetManager` for more details. diff --git a/embodichain/data_pipeline/depth_video/writer.py b/embodichain/data_pipeline/depth_video/writer.py index fbab673d5..50326e318 100644 --- a/embodichain/data_pipeline/depth_video/writer.py +++ b/embodichain/data_pipeline/depth_video/writer.py @@ -335,23 +335,28 @@ def add_frame(self, sensor_key: str, depth: Any) -> None: def end_episode(self, episode_index: int) -> None: """Close all writers for the current episode and update metadata. - On failure of any writer the partial files are discarded and that - sensor/episode is skipped (with a warning) rather than aborting the - whole recording. + Every writer is attempted. Partial files from failed writers are + discarded, metadata for successful writers is flushed, and any error + is then raised so a committed LeRobot episode cannot silently lose its + depth sidecar. Args: episode_index: Global episode index. + + Raises: + RuntimeError: If a depth video or its metadata cannot be finalized. """ + errors: list[str] = [] for key, writer in self._writers.items(): try: writer.close() frame_count = writer.frame_count - except Exception as e: - logger.log_warning( - f"Failed to finalize depth video for sensor '{key}' " - f"episode {episode_index}: {e}" - ) - writer.abort() + except Exception as error: + errors.append(f"sensor {key!r}: {error}") + try: + writer.abort() + except Exception as abort_error: # noqa: BLE001 - aggregate cleanup + errors.append(f"sensor {key!r} abort: {abort_error}") continue rel_file = writer.final_path.relative_to(self._root).as_posix() self._meta["sensors"][key]["episodes"][str(episode_index)] = { @@ -361,7 +366,15 @@ def end_episode(self, episode_index: int) -> None: } self._writers = {} self._episode_sensors = [] - self._write_meta() + try: + self._write_meta() + except Exception as error: + errors.append(f"metadata: {error}") + if errors: + raise RuntimeError( + f"Failed to finalize depth episode {episode_index}: " + + "; ".join(errors) + ) def abort_episode(self) -> None: """Abort the current episode: discard all partial depth videos. diff --git a/embodichain/data_pipeline/engine/__init__.py b/embodichain/data_pipeline/engine/__init__.py index 0f4781f23..d680cb5d9 100644 --- a/embodichain/data_pipeline/engine/__init__.py +++ b/embodichain/data_pipeline/engine/__init__.py @@ -16,9 +16,16 @@ """Process-safe shared buffer (``OnlineDataEngine``) that decouples simulation producers from training consumers.""" -from .data import OnlineDataEngine, OnlineDataEngineCfg +from .data import ( + OnlineDataEngine, + OnlineDataEngineCfg, + OnlineDataEngineState, + OnlineDataWorkerError, +) __all__ = [ "OnlineDataEngine", "OnlineDataEngineCfg", + "OnlineDataEngineState", + "OnlineDataWorkerError", ] diff --git a/embodichain/data_pipeline/engine/data.py b/embodichain/data_pipeline/engine/data.py index 36384cbf9..26f0f693a 100644 --- a/embodichain/data_pipeline/engine/data.py +++ b/embodichain/data_pipeline/engine/data.py @@ -16,13 +16,22 @@ from __future__ import annotations +import os +import math +import pickle +import sys +import threading import time -import torch import multiprocessing as mp +from dataclasses import dataclass +from enum import Enum +from types import TracebackType from typing import Literal from multiprocessing.sharedctypes import Synchronized, SynchronizedArray from multiprocessing.synchronize import Event as MpEvent + +import torch from tensordict import TensorDict from tqdm import tqdm @@ -30,7 +39,154 @@ from embodichain.utils.logger import log_info, log_error from embodichain.utils import configclass -__all__ = ["OnlineDataEngine", "OnlineDataEngineCfg"] +__all__ = [ + "OnlineDataEngine", + "OnlineDataEngineCfg", + "OnlineDataEngineState", + "OnlineDataWorkerError", +] + +_ERROR_BUFFER_SIZE = 64 * 1024 + + +class OnlineDataEngineState(str, Enum): + """Lifecycle states for :class:`OnlineDataEngine`.""" + + CREATED = "CREATED" + STARTING = "STARTING" + READY = "READY" + FAILED = "FAILED" + STOPPED = "STOPPED" + + +_STATE_TO_CODE = {state: index for index, state in enumerate(OnlineDataEngineState)} +_CODE_TO_STATE = {code: state for state, code in _STATE_TO_CODE.items()} + + +class OnlineDataWorkerError(RuntimeError): + """Fallback error for a worker exception that cannot be reconstructed.""" + + +def _forced_shutdown_error() -> OnlineDataWorkerError: + """Build the error used when graceful worker durability is unknown.""" + return OnlineDataWorkerError( + "OnlineDataEngine worker required terminate()/kill(); graceful close " + "and recorder durability could not be confirmed." + ) + + +@dataclass(frozen=True) +class _WorkerErrorEnvelope: + """Stable, pickle-safe description of a worker exception.""" + + module: str + qualname: str + message: str + representation: str + exception_payload: bytes | None + + +def _make_worker_error_envelope(error: BaseException) -> _WorkerErrorEnvelope: + """Create a stable envelope, retaining the exception when it is picklable.""" + exception_payload = None + try: + candidate = pickle.dumps(error, protocol=pickle.HIGHEST_PROTOCOL) + if len(candidate) < _ERROR_BUFFER_SIZE // 2: + exception_payload = candidate + except BaseException: + pass + + return _WorkerErrorEnvelope( + module=type(error).__module__, + qualname=type(error).__qualname__, + message=str(error)[:4096], + representation=repr(error)[:4096], + exception_payload=exception_payload, + ) + + +def _error_from_envelope(envelope: _WorkerErrorEnvelope) -> BaseException: + """Reconstruct an original exception or return the stable fallback type.""" + if envelope.exception_payload is not None: + try: + error = pickle.loads(envelope.exception_payload) + except BaseException as decode_error: + return OnlineDataWorkerError( + f"Worker raised {envelope.module}.{envelope.qualname}: " + f"{envelope.message} (exception reconstruction failed: " + f"{type(decode_error).__name__}: {decode_error})" + ) + if isinstance(error, BaseException): + return error + + return OnlineDataWorkerError( + f"Worker raised {envelope.module}.{envelope.qualname}: {envelope.message}" + ) + + +def _publish_worker_error( + error_buffer: SynchronizedArray, + error_length: Synchronized, + failed_signal: MpEvent, + state_value: Synchronized, + error: BaseException, +) -> bool: + """Publish a worker exception once to every process sharing the engine.""" + try: + envelope = _make_worker_error_envelope(error) + payload = pickle.dumps(envelope, protocol=pickle.HIGHEST_PROTOCOL) + if len(payload) > len(error_buffer): + envelope = _WorkerErrorEnvelope( + module=envelope.module, + qualname=envelope.qualname, + message=envelope.message[:1024], + representation=envelope.representation[:1024], + exception_payload=None, + ) + payload = pickle.dumps(envelope, protocol=pickle.HIGHEST_PROTOCOL) + if len(payload) > len(error_buffer): + return False + + with error_buffer.get_lock(): + error_buffer[: len(payload)] = payload + error_length.value = len(payload) + failed_signal.set() + with state_value.get_lock(): + state_value.value = _STATE_TO_CODE[OnlineDataEngineState.FAILED] + except BaseException: + return False + return True + + +def _monitor_worker_process( + process: mp.Process, + close_signal: MpEvent, + error_buffer: SynchronizedArray, + error_length: Synchronized, + failed_signal: MpEvent, + state_value: Synchronized, +) -> None: + """Broadcast hard worker exits without retaining the engine instance.""" + try: + process.join() + if close_signal.is_set() or failed_signal.is_set(): + return + error = RuntimeError( + "OnlineDataEngine simulation worker exited unexpectedly " + f"(exit code {process.exitcode})." + ) + except BaseException as caught_error: + if close_signal.is_set(): + return + error = caught_error + + _publish_worker_error( + error_buffer, + error_length, + failed_signal, + state_value, + error, + ) @configclass @@ -67,19 +223,27 @@ class OnlineDataEngineCfg: max_generation_attempts: int = 3 """Maximum planning/execution attempts for each buffer write transaction.""" + initialization_timeout: float = 300.0 + """Maximum seconds to wait for the worker's initial buffer fill.""" + # --------------------------------------------------------------------------- # Subprocess entry point (module-level so it can be pickled by multiprocessing) # --------------------------------------------------------------------------- -def _sim_worker_fn( +def _run_sim_worker( cfg: OnlineDataEngineCfg, shared_buffer: TensorDict, lock_index: SynchronizedArray, fill_signal: MpEvent, init_signal: MpEvent, close_signal: MpEvent, + error_buffer: SynchronizedArray, + error_length: Synchronized, + failed_signal: MpEvent, + state_value: Synchronized, + error_reported: list[bool], ) -> None: """Simulation subprocess entry point. @@ -108,7 +272,7 @@ def _sim_worker_fn( ) from embodichain.lab.gym.envs.demo import execute_demo_episode from embodichain.lab.sim import SimulationManagerCfg - from embodichain.utils.logger import log_info, log_warning, log_error + from embodichain.utils.logger import log_info, log_warning gym_config: dict = cfg.gym_config action_config: dict = cfg.action_config @@ -253,14 +417,65 @@ def _sim_worker_fn( "[Simulation Process] Initial buffer fill complete. Engine is ready.", color="cyan", ) - except KeyboardInterrupt: - log_warning("[Simulation Process] Stopping (KeyboardInterrupt).") - except Exception as e: - log_error(f"[Simulation Process] Unhandled error: {e}") finally: + error = sys.exc_info()[1] + if error is not None and not error_reported[0]: + error_reported[0] = _publish_worker_error( + error_buffer, + error_length, + failed_signal, + state_value, + error, + ) env.close() +def _sim_worker_fn( + cfg: OnlineDataEngineCfg, + shared_buffer: TensorDict, + lock_index: SynchronizedArray, + fill_signal: MpEvent, + init_signal: MpEvent, + close_signal: MpEvent, + error_buffer: SynchronizedArray, + error_length: Synchronized, + failed_signal: MpEvent, + state_value: Synchronized, +) -> None: + """Run the simulation worker and publish a stable exception envelope. + + Picklable exceptions are reconstructed with their original type and + message. Every consumer process reads the same shared snapshot; exceptions + that cannot be reconstructed are represented by + :class:`OnlineDataWorkerError`. + """ + error_reported = [False] + try: + _run_sim_worker( + cfg, + shared_buffer, + lock_index, + fill_signal, + init_signal, + close_signal, + error_buffer, + error_length, + failed_signal, + state_value, + error_reported, + ) + except BaseException as error: + if not error_reported[0]: + _publish_worker_error( + error_buffer, + error_length, + failed_signal, + state_value, + error, + ) + raise + + # --------------------------------------------------------------------------- # OnlineDataEngine # --------------------------------------------------------------------------- @@ -296,12 +511,13 @@ class OnlineDataEngine: and the counter resets to zero. This amortises the cost of GPU-accelerated simulation across many training iterations. - **Initialisation barrier** + **Lifecycle state** - The :attr:`is_init` property returns ``False`` until the subprocess - completes the very first full buffer fill, after which it becomes - permanently ``True``. Training code should wait on this flag before - calling :meth:`sample_batch` to avoid drawing all-zero data. + Every instance starts in :attr:`OnlineDataEngineState.CREATED`, passes + through ``STARTING`` while the first fill is running, and only serves data + in ``READY``. Worker failures transition to ``FAILED`` and explicit cleanup + transitions to terminal ``STOPPED``; failed or stopped instances cannot be + restarted. Args: cfg: Engine configuration. @@ -311,10 +527,16 @@ class OnlineDataEngine: ``[buffer_size, max_episode_steps, ...]``. buffer_size: Total number of trajectory slots in the shared buffer. device: Device of the shared buffer. - is_init: ``True`` once the buffer has been populated at least once. + state: Current :class:`OnlineDataEngineState`. + is_init: ``True`` only while the engine is ready to sample. """ def __init__(self, cfg: OnlineDataEngineCfg) -> None: + self._owner_pid = os.getpid() + self._lifecycle_condition = threading.Condition(threading.RLock()) + self._stop_requested = False + self._cleanup_complete = False + self._forced_shutdown_attempted = False self.cfg = cfg if cfg.max_generation_attempts < 1: @@ -322,6 +544,14 @@ def __init__(self, cfg: OnlineDataEngineCfg) -> None: "max_generation_attempts must be at least 1, " f"got {cfg.max_generation_attempts}." ) + if ( + not math.isfinite(cfg.initialization_timeout) + or cfg.initialization_timeout <= 0 + ): + raise ValueError( + "initialization_timeout must be finite and greater than zero, " + f"got {cfg.initialization_timeout}." + ) # Allocate the shared buffer (shape: [buffer_size, max_episode_steps, ...]). self.shared_buffer: TensorDict = self._create_buffer() @@ -361,46 +591,267 @@ def __init__(self, cfg: OnlineDataEngineCfg) -> None: # Accumulated sample count used by the refill criterion. self._sample_count: Synchronized = self._mp_ctx.Value("i", 0) + # State and worker failures are shared so every DataLoader consumer + # observes the same terminal state and exception snapshot. + self._state_value: Synchronized = self._mp_ctx.Value( + "i", _STATE_TO_CODE[OnlineDataEngineState.CREATED] + ) + self._worker_failed_signal: MpEvent = self._mp_ctx.Event() + self._worker_error_buffer: SynchronizedArray = self._mp_ctx.Array( + "B", _ERROR_BUFFER_SIZE + ) + self._worker_error_length: Synchronized = self._mp_ctx.Value("i", 0) + # Handle to the simulation subprocess, set in start() and used in stop(). self._sim_process: mp.Process | None = None + self._monitor_thread: threading.Thread | None = None + self._channel_error: BaseException | None = None + self._worker_error: BaseException | None = None def start(self) -> None: - self._sim_process: mp.Process = self._mp_ctx.Process( - target=_sim_worker_fn, - args=( - self.cfg, - self.shared_buffer, - self._lock_index, - self._fill_signal, - self._init_signal, - self._close_signal, - ), - daemon=True, - ) - self._sim_process.start() - log_info( - f"[OnlineDataEngine] Simulation subprocess started (PID={self._sim_process.pid}).", - color="green", - ) + """Start the worker and block until its first fill completes. - # Trigger the initial fill so data is ready before the first sample. - self._fill_signal.set() + Raises: + RuntimeError: If the engine was already started or stopped. + TimeoutError: If the first fill exceeds ``initialization_timeout``. + BaseException: The original exception raised by the worker. + """ + self._require_owner_process("start") + with self._lifecycle_condition: + self._require_state(OnlineDataEngineState.CREATED, "start") + self._stop_requested = False + self._set_state(OnlineDataEngineState.STARTING) + + try: + with self._lifecycle_condition: + self._sim_process = self._mp_ctx.Process( + target=_sim_worker_fn, + args=( + self.cfg, + self.shared_buffer, + self._lock_index, + self._fill_signal, + self._init_signal, + self._close_signal, + self._worker_error_buffer, + self._worker_error_length, + self._worker_failed_signal, + self._state_value, + ), + # Some planners create their own process pool. A daemonic + # producer would make those nested workers illegal. + daemon=False, + ) + self._sim_process.start() + process = self._sim_process + self._monitor_thread = threading.Thread( + target=_monitor_worker_process, + args=( + process, + self._close_signal, + self._worker_error_buffer, + self._worker_error_length, + self._worker_failed_signal, + self._state_value, + ), + name="online-data-worker-monitor", + daemon=True, + ) + self._monitor_thread.start() + log_info( + f"[OnlineDataEngine] Simulation subprocess started (PID={self._sim_process.pid}).", + color="green", + ) + + # Trigger the initial fill so data is ready before the first sample. + self._fill_signal.set() + deadline = time.monotonic() + self.cfg.initialization_timeout + + while not self._init_signal.wait(timeout=0.1): + with self._lifecycle_condition: + if self._stop_requested: + raise RuntimeError( + "OnlineDataEngine.start() was cancelled by stop()." + ) + self._ensure_worker_alive() + if time.monotonic() >= deadline: + raise TimeoutError( + "OnlineDataEngine initial buffer fill exceeded " + f"{self.cfg.initialization_timeout} seconds." + ) - while not self.is_init: self._ensure_worker_alive() - time.sleep(0.5) + with self._lifecycle_condition: + if self._stop_requested: + raise RuntimeError( + "OnlineDataEngine.start() was cancelled by stop()." + ) + self._set_state(OnlineDataEngineState.READY) + self._lifecycle_condition.notify_all() + except BaseException as error: + cleanup_error = None + forced_shutdown = False + with self._lifecycle_condition: + stop_requested = self._stop_requested + if not self._worker_failed_signal.is_set(): + self._worker_error = error + try: + forced_shutdown = self._shutdown_worker() + except BaseException as caught_cleanup_error: + cleanup_error = caught_cleanup_error + error.add_note( + f"Worker cleanup also failed: {caught_cleanup_error}" + ) + else: + self._cleanup_complete = True + + forced_shutdown = forced_shutdown or self._forced_shutdown_attempted + + # The worker may only publish an env.close()/recorder failure + # while the join above is in progress. Keep the start error as + # primary, but never lose that late durability error. + channel_error = self._receive_worker_error() + if channel_error is not None and channel_error is not error: + error.add_note( + "Worker also failed during cleanup: " + f"{type(channel_error).__name__}: {channel_error}" + ) + + if forced_shutdown: + durability_error = _forced_shutdown_error() + if channel_error is None: + self._record_worker_error(durability_error) + channel_error = durability_error + error.add_note(str(durability_error)) + + if ( + stop_requested + and channel_error is None + and cleanup_error is None + and not forced_shutdown + ): + self._set_state(OnlineDataEngineState.STOPPED) + else: + self._set_state(OnlineDataEngineState.FAILED) + self._lifecycle_condition.notify_all() + raise def _ensure_worker_alive(self) -> None: - """Raise when a started simulation worker exits unexpectedly.""" - if self._sim_process is None or self._sim_process.is_alive(): + """Fail the engine immediately when its worker reports or exits.""" + worker_error = self._receive_worker_error() + if worker_error is not None: + self._set_state(OnlineDataEngineState.FAILED) + raise worker_error + + if not self._is_owner_process(): + if self._close_signal.is_set(): + raise RuntimeError( + "OnlineDataEngine owner has stopped the simulation worker." + ) + return + + if self._sim_process is None: + error = RuntimeError("OnlineDataEngine simulation worker was not created.") + self._record_worker_error(error) + raise error + if self._sim_process.is_alive(): return self._sim_process.join(timeout=0) + worker_error = self._receive_worker_error() + if worker_error is None: + worker_error = RuntimeError( + "OnlineDataEngine simulation worker exited unexpectedly " + f"(exit code {self._sim_process.exitcode})." + ) + self._record_worker_error(worker_error) + self._set_state(OnlineDataEngineState.FAILED) + raise worker_error + + def _receive_worker_error(self) -> BaseException | None: + """Return the broadcast worker exception when one has been published.""" + if not self._worker_failed_signal.is_set(): + return None + if self._channel_error is not None: + return self._channel_error + + try: + with self._worker_error_buffer.get_lock(): + payload_length = self._worker_error_length.value + payload = bytes(self._worker_error_buffer[:payload_length]) + envelope = pickle.loads(payload) + if not isinstance(envelope, _WorkerErrorEnvelope): + raise TypeError(f"invalid envelope type {type(envelope).__name__}") + error = _error_from_envelope(envelope) + except BaseException as decode_error: + error = OnlineDataWorkerError( + "OnlineDataEngine could not decode the worker error channel: " + f"{type(decode_error).__name__}: {decode_error}" + ) + self._channel_error = error + if self._worker_error is None: + self._worker_error = error + self._set_state(OnlineDataEngineState.FAILED) + return error + + def _record_worker_error(self, error: BaseException) -> None: + """Publish an owner-detected worker failure to every consumer.""" + if self._worker_failed_signal.is_set(): + self._receive_worker_error() + return + + self._channel_error = error + if self._worker_error is None: + self._worker_error = error + published = _publish_worker_error( + self._worker_error_buffer, + self._worker_error_length, + self._worker_failed_signal, + self._state_value, + error, + ) + if not published: + fallback = OnlineDataWorkerError( + "OnlineDataEngine worker error serialization failed for " + f"{type(error).__module__}.{type(error).__qualname__}." + ) + _publish_worker_error( + self._worker_error_buffer, + self._worker_error_length, + self._worker_failed_signal, + self._state_value, + fallback, + ) + self._set_state(OnlineDataEngineState.FAILED) + + def _require_state(self, expected: OnlineDataEngineState, operation: str) -> None: + """Require an exact lifecycle state for a public operation.""" + current_state = self.state + if current_state is expected: + return raise RuntimeError( - "OnlineDataEngine simulation worker exited unexpectedly " - f"(exit code {self._sim_process.exitcode})." + f"OnlineDataEngine.{operation}() requires state {expected.value}; " + f"current state is {current_state.value}." + ) from self._worker_error + + def _is_owner_process(self) -> bool: + """Whether the current process owns the producer lifecycle.""" + return os.getpid() == self._owner_pid + + def _require_owner_process(self, operation: str) -> None: + """Reject lifecycle operations from forked or spawned consumers.""" + if self._is_owner_process(): + return + raise RuntimeError( + f"OnlineDataEngine.{operation}() may only be called by owner process " + f"{self._owner_pid}; current process is {os.getpid()}." ) + def _set_state(self, state: OnlineDataEngineState) -> None: + """Publish a lifecycle transition to every sharing process.""" + with self._state_value.get_lock(): + self._state_value.value = _STATE_TO_CODE[state] + # ----------------------------------------------------------------------- # Buffer initialisation # ----------------------------------------------------------------------- @@ -440,21 +891,16 @@ def _create_buffer(self) -> TensorDict: # ----------------------------------------------------------------------- @property - def is_init(self) -> bool: - """Whether the shared buffer has been fully populated at least once. + def state(self) -> OnlineDataEngineState: + """Return the engine's current lifecycle state.""" + with self._state_value.get_lock(): + state_code = self._state_value.value + return _CODE_TO_STATE[state_code] - Returns ``True`` after the simulation subprocess completes its first - full buffer fill, ``False`` while that initial fill is still in - progress. Callers that must not sample stale (all-zero) data can - poll or block on this property before entering their training loop:: - - while not engine.is_init: - time.sleep(0.5) - - Returns: - ``True`` once the buffer contains valid trajectory data. - """ - return self._init_signal.is_set() + @property + def is_init(self) -> bool: + """Whether the engine is ready to serve initialized data.""" + return self.state is OnlineDataEngineState.READY # ----------------------------------------------------------------------- # Sampling @@ -491,7 +937,12 @@ def sample_batch( ValueError: If an argument is invalid. RuntimeError: If no unlocked valid window satisfies the policy. """ - self._ensure_worker_alive() + with self._lifecycle_condition: + worker_error = self._receive_worker_error() + if worker_error is not None: + raise worker_error + self._require_state(OnlineDataEngineState.READY, "sample_batch") + self._ensure_worker_alive() max_steps: int = self.shared_buffer.batch_size[1] if batch_size < 1: @@ -621,6 +1072,107 @@ def _trigger_refill_if_needed(self, count: int = 1) -> None: # Lifecycle # ----------------------------------------------------------------------- + def _detect_preexisting_worker_exit(self) -> BaseException | None: + """Return and broadcast a worker exit observed before shutdown starts.""" + process = self._sim_process + if process is None: + error = RuntimeError("OnlineDataEngine simulation worker was not created.") + self._record_worker_error(error) + return error + + try: + if process.is_alive(): + return None + process.join(timeout=0) + exit_code = process.exitcode + except BaseException as inspection_error: + error = RuntimeError( + "OnlineDataEngine could not inspect its simulation worker before " + f"shutdown: {type(inspection_error).__name__}: {inspection_error}" + ) + self._record_worker_error(error) + return error + + worker_error = self._receive_worker_error() + if worker_error is not None: + return worker_error + + error = RuntimeError( + "OnlineDataEngine simulation worker exited unexpectedly " + f"before stop() (exit code {exit_code})." + ) + self._record_worker_error(error) + return error + + def _shutdown_worker(self) -> bool: + """Stop and reap the worker, returning whether force was required.""" + self._require_owner_process("stop") + self._close_signal.set() + self._fill_signal.set() + + process = self._sim_process + if process is None: + return False + + forced_shutdown = False + + try: + is_alive = process.is_alive() + except ValueError: + is_alive = False + + try: + process.join(timeout=5.0 if is_alive else 0) + except (AssertionError, ValueError): + pass + + try: + is_alive = process.is_alive() + except ValueError: + is_alive = False + if is_alive: + forced_shutdown = True + self._forced_shutdown_attempted = True + process.terminate() + process.join(timeout=3.0) + + try: + is_alive = process.is_alive() + except ValueError: + is_alive = False + if is_alive and hasattr(process, "kill"): + forced_shutdown = True + self._forced_shutdown_attempted = True + process.kill() + process.join(timeout=1.0) + + try: + is_alive = process.is_alive() + except ValueError: + is_alive = False + if is_alive: + raise RuntimeError( + "OnlineDataEngine simulation worker remained alive after " + "graceful shutdown, terminate(), and kill()." + ) + + monitor_thread = self._monitor_thread + if ( + monitor_thread is not None + and monitor_thread is not threading.current_thread() + ): + monitor_thread.join(timeout=1.0) + if monitor_thread.is_alive(): + raise RuntimeError( + "OnlineDataEngine worker monitor did not stop after the worker exited." + ) + + if hasattr(process, "close"): + process.close() + self._sim_process = None + self._monitor_thread = None + return forced_shutdown + def stop(self) -> None: """Terminate the simulation subprocess and release resources. @@ -631,21 +1183,127 @@ def stop(self) -> None: Safe to call multiple times — subsequent calls are no-ops if the subprocess has already been terminated. """ - if self._sim_process is None or not self._sim_process.is_alive(): - return - - # Ask the subprocess to stop and unblock it if it is waiting on fill_signal. - self._close_signal.set() - self._fill_signal.set() - - # Allow time for a graceful exit (close_signal is checked between steps). - self._sim_process.join(timeout=5.0) - - if self._sim_process.is_alive(): - self._sim_process.terminate() - self._sim_process.join(timeout=3.0) + self._require_owner_process("stop") + with self._lifecycle_condition: + if self.state is OnlineDataEngineState.STOPPED: + return + if self.state is OnlineDataEngineState.FAILED and self._cleanup_complete: + # A successful join makes future publications impossible, but + # still decode anything already visible before honoring the + # idempotent stop contract. + self._receive_worker_error() + return + + if self.state is OnlineDataEngineState.STARTING: + self._stop_requested = True + self._close_signal.set() + self._fill_signal.set() + while self.state is OnlineDataEngineState.STARTING: + self._lifecycle_condition.wait(timeout=0.1) + if self.state is OnlineDataEngineState.STOPPED: + return - log_info("[OnlineDataEngine] Simulation subprocess terminated.", color="green") + state_before_shutdown = self.state + worker_error = self._receive_worker_error() + if ( + worker_error is None + and state_before_shutdown is OnlineDataEngineState.READY + ): + worker_error = self._detect_preexisting_worker_exit() + + forced_shutdown = False + try: + forced_shutdown = self._shutdown_worker() + except BaseException as cleanup_error: + # The worker can fail while handling the close signal (for + # example, while flushing a recorder in ``env.close()``). + # Re-read the shared channel after waiting for it so that a + # durability failure is not hidden behind the cleanup path. + worker_error = worker_error or self._receive_worker_error() + if self._forced_shutdown_attempted: + durability_error = _forced_shutdown_error() + if worker_error is None: + self._record_worker_error(durability_error) + worker_error = durability_error + else: + worker_error.add_note(str(durability_error)) + self._set_state(OnlineDataEngineState.FAILED) + self._lifecycle_condition.notify_all() + if worker_error is not None: + worker_error.add_note( + f"Worker cleanup also failed: {cleanup_error}" + ) + raise worker_error + self._worker_error = cleanup_error + raise + + # ``_shutdown_worker`` joins the producer, so any exception raised + # during its final ``env.close()`` has now been published. + worker_error = worker_error or self._receive_worker_error() + forced_shutdown = forced_shutdown or self._forced_shutdown_attempted + if forced_shutdown: + durability_error = _forced_shutdown_error() + if worker_error is None: + self._record_worker_error(durability_error) + worker_error = durability_error + else: + worker_error.add_note(str(durability_error)) + + self._cleanup_complete = True + if worker_error is not None: + self._set_state(OnlineDataEngineState.FAILED) + self._lifecycle_condition.notify_all() + raise worker_error + + self._set_state(OnlineDataEngineState.STOPPED) + self._lifecycle_condition.notify_all() + log_info("[OnlineDataEngine] Engine stopped.", color="green") + + def __enter__(self) -> "OnlineDataEngine": + """Start the engine and return it for a managed lifecycle block.""" + self.start() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> bool | None: + """Stop the engine while preserving any exception from the block.""" + try: + self.stop() + except BaseException as cleanup_error: + if exc_value is None: + raise + exc_value.add_note( + "OnlineDataEngine cleanup also failed: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + return None + + def __getstate__(self) -> dict: + """Serialize only consumer-safe fields for spawned DataLoader workers.""" + state = self.__dict__.copy() + state["_sim_process"] = None + state["_monitor_thread"] = None + state["_lifecycle_condition"] = None + state["_channel_error"] = None + state["_worker_error"] = None + return state + + def __setstate__(self, state: dict) -> None: + """Restore process-local synchronization after consumer deserialization.""" + self.__dict__.update(state) + self._lifecycle_condition = threading.Condition(threading.RLock()) def __del__(self) -> None: - self.stop() + try: + if getattr(self, "_owner_pid", None) != os.getpid(): + return + if self.state is not OnlineDataEngineState.STOPPED: + self.stop() + except BaseException: + # Destructors run during partially initialized objects and interpreter + # shutdown, where cleanup must never mask the original exception. + pass diff --git a/embodichain/lab/gym/envs/demo.py b/embodichain/lab/gym/envs/demo.py index d635ad756..006b3318b 100644 --- a/embodichain/lab/gym/envs/demo.py +++ b/embodichain/lab/gym/envs/demo.py @@ -64,6 +64,11 @@ class DemoSegment: target_uid: Optional scene entity manipulated by this segment. instruction: Optional language instruction specific to this segment. metadata: Additional JSON-compatible task metadata. + validator: Optional zero-argument callback that validates this segment + after its actions are exhausted. It must return one boolean per + parallel environment (or one scalar broadcast to every environment). + Gym ``terminated`` and ``truncated`` remain episode-level signals; + use this callback for subtask-level validation. """ actions: Iterable[Any] @@ -71,11 +76,32 @@ class DemoSegment: target_uid: str | None = None instruction: str | None = None metadata: Mapping[str, Any] = field(default_factory=dict) + validator: Callable[[], Any] | None = field(default=None, repr=False, compare=False) @dataclass(frozen=True) class DemoSegmentResult: - """Execution result and half-open frame range for one segment.""" + """Execution result and half-open frame range for one segment. + + Scalar span and status fields are batch aggregates kept for compatibility. + The tuple fields preserve each vector-environment row independently. + + Args: + segment_id: Zero-based segment index within the episode. + name: Stable segment name supplied by the task. + start_step: Earliest participating row start, inclusive. + end_step: Latest participating row end, exclusive. + success: Whether every participating row completed the segment. + target_uid: Optional manipulated scene entity. + instruction: Optional language instruction. + failure_reason: First aggregate failure reason, if any. + metadata: Additional JSON-compatible task metadata. + active: Participation mask captured at segment start. + start_steps: Per-environment inclusive starts. + end_steps: Per-environment exclusive ends. + successes: Per-environment segment status. + failure_reasons: Per-environment failure reasons. + """ segment_id: int name: str @@ -86,25 +112,78 @@ class DemoSegmentResult: instruction: str | None = None failure_reason: str | None = None metadata: Mapping[str, Any] = field(default_factory=dict) - - def to_metadata(self) -> dict[str, Any]: - """Return a JSON-compatible representation.""" - return { + active: tuple[bool, ...] = () + start_steps: tuple[int, ...] = () + end_steps: tuple[int, ...] = () + successes: tuple[bool, ...] = () + failure_reasons: tuple[str | None, ...] = () + + def to_metadata(self, env_id: int | None = None) -> dict[str, Any]: + """Return a JSON-compatible aggregate or per-environment representation. + + Args: + env_id: Optional parallel-environment index. When provided, scalar + spans and status are selected from the per-environment fields. + + Returns: + JSON-compatible segment metadata. + """ + metadata = { "segment_id": self.segment_id, "name": self.name, - "start_step": self.start_step, - "end_step": self.end_step, - "success": self.success, "target_uid": self.target_uid, "instruction": self.instruction, - "failure_reason": self.failure_reason, "metadata": dict(self.metadata), } + if env_id is not None and self.start_steps: + 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], + } + ) + return metadata + + metadata.update( + { + "start_step": self.start_step, + "end_step": self.end_step, + "success": self.success, + "failure_reason": self.failure_reason, + } + ) + if self.active: + metadata.update( + { + "active": list(self.active), + "start_steps": list(self.start_steps), + "end_steps": list(self.end_steps), + "successes": list(self.successes), + "failure_reasons": list(self.failure_reasons), + } + ) + return metadata @dataclass(frozen=True) class DemoEpisodeResult: - """Result of executing all planned segments for one batched episode.""" + """Result of executing all planned segments for one batched episode. + + Args: + episode_index: Logical episode identifier. + length: Maximum recorded row length. + completed: Whether every environment completed successfully. + success: Sticky per-environment success flags. + terminated: Sticky per-environment Gym termination flags. + truncated: Sticky per-environment Gym truncation flags. + terminal_reason: Aggregate terminal reason. + segments: Executed segment results. + lengths: Independent per-environment recorded lengths. + completed_by_env: Independent valid-completion flags. + terminal_reasons: Independent terminal reasons. + """ episode_index: int length: int @@ -114,6 +193,9 @@ class DemoEpisodeResult: truncated: tuple[bool, ...] terminal_reason: str segments: tuple[DemoSegmentResult, ...] = () + lengths: tuple[int, ...] = () + completed_by_env: tuple[bool, ...] = () + terminal_reasons: tuple[str, ...] = () @property def all_success(self) -> bool: @@ -127,7 +209,7 @@ def any_success(self) -> bool: def to_metadata(self) -> dict[str, Any]: """Return a JSON-compatible representation.""" - return { + metadata = { "schema_version": DEMO_SCHEMA_VERSION, "episode_index": self.episode_index, "length": self.length, @@ -138,6 +220,15 @@ def to_metadata(self) -> dict[str, Any]: "terminal_reason": self.terminal_reason, "segments": [segment.to_metadata() for segment in self.segments], } + if self.lengths: + metadata.update( + { + "lengths": list(self.lengths), + "completed_by_env": list(self.completed_by_env), + "terminal_reasons": list(self.terminal_reasons), + } + ) + return metadata ProgressWrapper = Callable[[Iterable[Any], str], Iterable[Any]] @@ -259,31 +350,71 @@ def execute_demo_episode( end_segment = _get_env_callable(env, "_end_demo_segment_recording") end_episode = _get_env_callable(env, "_end_demo_episode_recording") normalize_action = _get_env_callable(env, "_normalize_demo_action") + mask_action = _get_env_callable(env, "_mask_demo_action") + set_active_mask = _get_env_callable(env, "_set_demo_active_mask") + success_fn = _get_env_callable(env, "is_task_success") + + active = [True] * num_envs + + def publish_active_mask() -> None: + """Publish executor liveness to recording hooks and action masking.""" + if set_active_mask is not None: + set_active_mask(tuple(active)) + return + previous = getattr(target, "_demo_active_mask", None) + device = getattr(previous, "device", None) + setattr( + target, + "_demo_active_mask", + torch.tensor(active, dtype=torch.bool, device=device), + ) if begin_episode is not None: begin_episode(episode_index=episode_index) + publish_active_mask() previous_no_auto_reset = bool(getattr(target, "_demo_no_auto_reset", False)) setattr(target, "_demo_no_auto_reset", True) - total_steps = 0 + lengths = [0] * num_envs + success = [False] * num_envs + completed_by_env = [False] * num_envs + terminated = [False] * num_envs + truncated = [False] * num_envs + terminal_reasons = ["pending"] * num_envs segment_results: list[DemoSegmentResult] = [] - terminated = (False,) * num_envs - truncated = (False,) * num_envs last_info: Mapping[str, Any] = {} - terminal_reason = "completed" - completed = True + fatal_reason: str | None = None try: segment_count = 0 - for segment_id, segment in enumerate(resolve_demo_segments(env, **plan_kwargs)): + segments = iter(resolve_demo_segments(env, **plan_kwargs)) + while any(active): + if should_stop is not None and should_stop(): + fatal_reason = "interrupted" + for env_id, is_active in enumerate(active): + if is_active: + terminal_reasons[env_id] = fatal_reason + active[env_id] = False + publish_active_mask() + break + try: + segment = next(segments) + except StopIteration: + break + + segment_id = segment_count segment_count += 1 - start_step = total_steps + participants = tuple(active) + start_steps = tuple(lengths) + segment_successes = [False] * num_envs + segment_failure_reasons: list[str | None] = [None] * num_envs + segment_reason: str | None = None if begin_segment is not None: begin_segment(segment_id=segment_id, segment=segment) action_count = 0 - segment_reason: str | None = None + actions_exhausted = True actions: Iterable[Any] = segment.actions if actions is None: actions = () @@ -296,99 +427,246 @@ def execute_demo_episode( for action in actions: if should_stop is not None and should_stop(): - completed = False - terminal_reason = "interrupted" - segment_reason = terminal_reason + actions_exhausted = False + fatal_reason = "interrupted" + segment_reason = fatal_reason + for env_id, is_active in enumerate(active): + if is_active: + terminal_reasons[env_id] = fatal_reason + segment_failure_reasons[env_id] = fatal_reason + active[env_id] = False + publish_active_mask() break if normalize_action is not None: action = normalize_action(action) + if not all(active): + if mask_action is None: + raise RuntimeError( + "A vector demo environment completed asynchronously but " + "does not implement _mask_demo_action(action, active_mask)." + ) + action = mask_action(action, tuple(active)) + + active_before_step = tuple(active) _, _, terminated_value, truncated_value, info = env.step(action) action_count += 1 - total_steps += 1 last_info = info - terminated = _as_bool_tuple(terminated_value, num_envs) - truncated = _as_bool_tuple(truncated_value, num_envs) + for env_id, was_active in enumerate(active_before_step): + if was_active: + lengths[env_id] += 1 + + step_terminated = _as_bool_tuple(terminated_value, num_envs) + step_truncated = _as_bool_tuple(truncated_value, num_envs) + step_success_source = last_info.get("success") + if step_success_source is None and any( + step_terminated[env_id] + for env_id, was_active in enumerate(active_before_step) + if was_active + ): + step_success_source = ( + success_fn() if success_fn is not None else False + ) + step_success = _as_bool_tuple(step_success_source, num_envs) + step_failure = _as_bool_tuple(last_info.get("fail"), num_envs) + + step_failed = False + active_step_truncated = False + for env_id, was_active in enumerate(active_before_step): + if not was_active: + continue + # Preserve both raw Gym flags when an environment reports + # terminated and truncated on the same transition. + terminated[env_id] |= step_terminated[env_id] + truncated[env_id] |= step_truncated[env_id] + if step_truncated[env_id]: + active_step_truncated = True + success[env_id] = False + terminal_reasons[env_id] = "truncated" + segment_failure_reasons[env_id] = "truncated" + active[env_id] = False + step_failed = True + elif step_failure[env_id]: + success[env_id] = False + terminal_reasons[env_id] = "failure" + segment_failure_reasons[env_id] = "failure" + active[env_id] = False + step_failed = True + elif step_terminated[env_id]: + active[env_id] = False + if step_success[env_id]: + success[env_id] = True + completed_by_env[env_id] = True + terminal_reasons[env_id] = "success" + segment_successes[env_id] = True + else: + success[env_id] = False + terminal_reasons[env_id] = "failure" + segment_failure_reasons[env_id] = "failure" + step_failed = True + + if step_failed: + actions_exhausted = False + fatal_reason = "truncated" if active_step_truncated else "failure" + segment_reason = fatal_reason + for env_id, is_active in enumerate(active): + if is_active: + terminal_reasons[env_id] = "batch_aborted" + segment_failure_reasons[env_id] = "batch_aborted" + active[env_id] = False + publish_active_mask() + break - if any(terminated) or any(truncated): - completed = False - terminal_reason = "truncated" if any(truncated) else "terminated" - segment_reason = terminal_reason + publish_active_mask() + if not any(active): + # Every row reached episode-level success. Stop this segment + # and do not request another lazy segment. + actions_exhausted = False + break + if should_stop is not None and should_stop(): + actions_exhausted = False + fatal_reason = "interrupted" + segment_reason = fatal_reason + for env_id, is_active in enumerate(active): + if is_active: + terminal_reasons[env_id] = fatal_reason + segment_failure_reasons[env_id] = fatal_reason + active[env_id] = False + publish_active_mask() break if action_count == 0 and segment_reason is None: - completed = False - terminal_reason = "empty_segment" - segment_reason = terminal_reason - - segment_success = segment_reason is None - if segment_reason == "terminated": - segment_success = all( - _as_bool_tuple(last_info.get("success"), num_envs) + fatal_reason = "empty_segment" + segment_reason = fatal_reason + for env_id, is_participant in enumerate(participants): + if is_participant: + terminal_reasons[env_id] = fatal_reason + segment_failure_reasons[env_id] = fatal_reason + active[env_id] = False + publish_active_mask() + + if actions_exhausted and segment_reason is None: + validation = ( + _as_bool_tuple(segment.validator(), num_envs) + if segment.validator is not None + else (True,) * num_envs ) - + validation_failed = False + for env_id, is_participant in enumerate(participants): + if not is_participant: + continue + if completed_by_env[env_id] and success[env_id]: + segment_successes[env_id] = True + elif active[env_id] and validation[env_id]: + segment_successes[env_id] = True + elif active[env_id]: + validation_failed = True + segment_failure_reasons[env_id] = "segment_validation_failed" + terminal_reasons[env_id] = "segment_validation_failed" + + if validation_failed: + fatal_reason = "segment_validation_failed" + segment_reason = fatal_reason + for env_id, is_active in enumerate(active): + if is_active: + if segment_failure_reasons[env_id] is None: + segment_failure_reasons[env_id] = "batch_aborted" + terminal_reasons[env_id] = "batch_aborted" + segment_successes[env_id] = False + active[env_id] = False + publish_active_mask() + + participant_ids = [ + env_id + for env_id, is_participant in enumerate(participants) + if is_participant + ] + segment_ok = bool(participant_ids) and all( + segment_successes[env_id] for env_id in participant_ids + ) + aggregate_failure = segment_reason or next( + ( + segment_failure_reasons[env_id] + for env_id in participant_ids + if segment_failure_reasons[env_id] is not None + ), + None, + ) + end_steps = tuple(lengths) segment_result = DemoSegmentResult( segment_id=segment_id, name=segment.name, - start_step=start_step, - end_step=total_steps, - success=segment_success, + start_step=min(start_steps[env_id] for env_id in participant_ids), + end_step=max(end_steps[env_id] for env_id in participant_ids), + success=segment_ok, target_uid=segment.target_uid, instruction=segment.instruction, - failure_reason=None if segment_success else segment_reason, + failure_reason=aggregate_failure, metadata=segment.metadata, + active=participants, + start_steps=start_steps, + end_steps=end_steps, + successes=tuple(segment_successes), + failure_reasons=tuple(segment_failure_reasons), ) segment_results.append(segment_result) if end_segment is not None: end_segment(result=segment_result) - if segment_reason is not None: + if segment_reason is not None or not any(active): break - if segment_count == 0: - completed = False - terminal_reason = "empty_plan" - - success_fn = _get_env_callable(env, "is_task_success") - if any(terminated) or any(truncated): - success_source = last_info.get("success") - if success_source is None: - success_source = success_fn() if success_fn is not None else False - else: - # Expert tasks historically validate the final planned state with - # is_task_success(). EmbodiedEnv.get_info() always contains a - # compute_task_state()-based success field, which remains false for - # legacy tasks that only implement the expert-policy hook. + if segment_count == 0 and fatal_reason is None: + fatal_reason = "empty_plan" + for env_id, is_active in enumerate(active): + if is_active: + terminal_reasons[env_id] = fatal_reason + active[env_id] = False + publish_active_mask() + elif fatal_reason is None and any(active): + # Normal plan exhaustion validates only rows that have not already + # reached sticky episode success. Legacy expert tasks use + # is_task_success() for this final validation. success_source = ( success_fn() if success_fn is not None - else last_info.get("success", completed) + else last_info.get("success", True) ) - success = _as_bool_tuple(success_source, num_envs) - - if any(truncated): - success = tuple(False for _ in success) - terminal_reason = "truncated" - elif any(terminated): - if all(success): - completed = True - terminal_reason = "success" - else: - terminal_reason = "failure" + final_success = _as_bool_tuple(success_source, num_envs) + for env_id, is_active in enumerate(active): + if not is_active: + continue + success[env_id] = final_success[env_id] + completed_by_env[env_id] = final_success[env_id] + terminal_reasons[env_id] = ( + "success" if final_success[env_id] else "task_incomplete" + ) + active[env_id] = False + publish_active_mask() + + completed = bool(completed_by_env) and all(completed_by_env) + if fatal_reason is not None: + terminal_reason = fatal_reason elif completed and all(success): terminal_reason = "success" - elif completed: - terminal_reason = "task_incomplete" + else: + terminal_reason = next( + (reason for reason in terminal_reasons if reason != "success"), + "task_incomplete", + ) result = DemoEpisodeResult( episode_index=episode_index, - length=total_steps, + length=max(lengths, default=0), completed=completed, - success=success, - terminated=terminated, - truncated=truncated, + success=tuple(success), + terminated=tuple(terminated), + truncated=tuple(truncated), terminal_reason=terminal_reason, segments=tuple(segment_results), + lengths=tuple(lengths), + completed_by_env=tuple(completed_by_env), + terminal_reasons=tuple(terminal_reasons), ) 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 4e32bacc5..0ce655589 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -20,6 +20,7 @@ from functools import wraps from datetime import datetime import os +import threading import torch import numpy as np import gymnasium as gym @@ -363,6 +364,10 @@ def __init__(self, cfg: EmbodiedEnvCfg, **kwargs): self._demo_active_segment_ids = torch.zeros( self.num_envs, dtype=torch.long, device=self.device ) + self._demo_active_mask = torch.ones( + self.num_envs, dtype=torch.bool, device=self.device + ) + self._demo_segment_participants = self._demo_active_mask.clone() self._demo_active_segment_start_steps = self._demo_steps.clone() self._demo_active_rollout_start_steps = self.rollout_steps.clone() self._demo_episode_metadata: list[dict[str, Any]] = [ @@ -372,6 +377,9 @@ def __init__(self, cfg: EmbodiedEnvCfg, **kwargs): self.episode_success_status: torch.Tensor = torch.zeros( self.num_envs, dtype=torch.bool, device=self.device ) + self._closed = False + self._close_error: BaseException | None = None + self._close_lock = threading.RLock() def set_rollout_buffer(self, rollout_buffer: TensorDict) -> None: """Set the rollout buffer for episode data collection. @@ -561,15 +569,36 @@ def _hook_after_sim_step( demo_steps = getattr(self, "_demo_steps", None) if demo_steps is not None: - demo_steps += 1 + active_mask = getattr(self, "_demo_active_mask", None) + if active_mask is None: + demo_steps += 1 + else: + demo_steps += active_mask.to(demo_steps.device, dtype=demo_steps.dtype) with self._profiler.section("trajectory_write"): self._write_trajectory_step() - # Update success status for all environments where episode is done - if "success" in info: - # info["success"] should be a tensor or array of shape (num_envs,) - self.episode_success_status[dones] = info["success"][dones] + self._update_episode_success_status(info, dones) + + def _update_episode_success_status( + self, info: Dict[str, Any], dones: torch.Tensor + ) -> None: + """Update terminal success without mutating already frozen demo rows.""" + if "success" not in info: + return + update_mask = dones.to( + device=self.episode_success_status.device, dtype=torch.bool + ) + if getattr(self, "_demo_no_auto_reset", False): + active_mask = getattr(self, "_demo_active_mask", None) + if active_mask is not None: + update_mask = update_mask & active_mask.to(update_mask.device) + success = torch.as_tensor( + info["success"], + dtype=torch.bool, + device=self.episode_success_status.device, + ) + self.episode_success_status[update_mask] = success[update_mask] def _extend_obs(self, obs: EnvObs, **kwargs) -> EnvObs: if self.observation_manager: @@ -656,9 +685,13 @@ def _initialize_episode( for functor_cfg in mode_cfgs: if isinstance(functor_cfg.func, record_camera_data): if save_data: - functor_cfg.func.save_and_clear() + functor_cfg.func.save_and_clear( + env_ids=env_ids_to_process + ) else: - functor_cfg.func.discard_and_clear() + functor_cfg.func.discard_and_clear( + env_ids=env_ids_to_process + ) # Auto-save + reset the per-env trajectory buffer for environments being # reset. Use getattr so this no-ops on envs/subclasses that don't allocate @@ -694,6 +727,8 @@ def _initialize_episode( if active_segment_ids is not None: demo_ids = env_ids_to_process.to(active_segment_ids.device) active_segment_ids[demo_ids] = 0 + self._demo_active_mask[demo_ids] = True + self._demo_segment_participants[demo_ids] = False self._demo_active_segment_start_steps[demo_ids] = 0 self._demo_active_rollout_start_steps[demo_ids] = 0 self._demo_steps[demo_ids] = 0 @@ -769,6 +804,8 @@ def _begin_demo_episode_recording(self, episode_index: int = 0) -> None: self._demo_episode_index = episode_index self._demo_active_segment_id = 0 self._demo_active_segment_ids.zero_() + self._demo_active_mask.fill_(True) + self._demo_segment_participants.fill_(False) self._demo_active_segment_start_steps = self._demo_steps.clone() self._demo_active_rollout_start_steps = self.rollout_steps.clone() self._demo_episode_metadata = [ @@ -780,13 +817,43 @@ def _begin_demo_segment_recording( ) -> None: """Set the segment context used by subsequent rollout writes.""" self._demo_active_segment_id = segment_id - self._demo_active_segment_ids.fill_(segment_id) + self._demo_segment_participants = self._demo_active_mask.clone() + self._demo_active_segment_ids[self._demo_segment_participants] = segment_id self._demo_active_segment_start_steps = self._demo_steps.clone() self._demo_active_rollout_start_steps = self.rollout_steps.clone() + def _set_demo_active_mask(self, active_mask: Sequence[bool]) -> None: + """Set rows that still participate in the current demo episode. + + The executor updates this sticky mask after each terminal transition. + Recording hooks use it to freeze completed rows while unfinished rows + continue on the shared vector-environment clock. + + Args: + active_mask: One liveness flag per parallel environment. + + Raises: + ValueError: If the mask length does not match ``num_envs``. + """ + mask = torch.as_tensor( + active_mask, dtype=torch.bool, device=self._demo_active_mask.device + ).reshape(-1) + if mask.numel() != self.num_envs: + raise ValueError( + f"Expected {self.num_envs} demo activity flags, got {mask.numel()}." + ) + self._demo_active_mask.copy_(mask) + def _end_demo_segment_recording(self, result: DemoSegmentResult) -> None: """Close the active segment span and mark its final valid frame.""" for env_id in range(self.num_envs): + is_participant = ( + result.active[env_id] + if result.active + else bool(self._demo_segment_participants[env_id]) + ) + if not is_participant: + continue start = int(self._demo_active_segment_start_steps[env_id].item()) end = int(self._demo_steps[env_id].item()) rollout_start = int(self._demo_active_rollout_start_steps[env_id].item()) @@ -798,7 +865,7 @@ def _end_demo_segment_recording(self, result: DemoSegmentResult) -> None: ): self.rollout_buffer["segment_end"][env_id, rollout_end - 1] = True - metadata = result.to_metadata() + metadata = result.to_metadata(env_id if result.start_steps else None) metadata["start_step"] = start metadata["end_step"] = end self._demo_episode_metadata[env_id]["segments"].append(metadata) @@ -807,15 +874,30 @@ def _end_demo_episode_recording(self, result: DemoEpisodeResult) -> None: """Finalize per-environment metadata after demonstration execution.""" for env_id in range(self.num_envs): metadata = self._demo_episode_metadata[env_id] + length = ( + result.lengths[env_id] + if result.lengths + else int(self._demo_steps[env_id].item()) + ) + completed = ( + result.completed_by_env[env_id] + if result.completed_by_env + else result.completed + ) + terminal_reason = ( + result.terminal_reasons[env_id] + if result.terminal_reasons + else result.terminal_reason + ) metadata.update( { "episode_index": result.episode_index, - "length": int(self._demo_steps[env_id].item()), - "completed": result.completed, + "length": length, + "completed": completed, "success": result.success[env_id], "terminated": result.terminated[env_id], "truncated": result.truncated[env_id], - "terminal_reason": result.terminal_reason, + "terminal_reason": terminal_reason, } ) @@ -920,9 +1002,13 @@ def _write_episode_rollout_step( """Write one expert step at each environment's independent cursor.""" buffer_device = self.rollout_buffer.device active = self.rollout_steps < self._max_rollout_steps + demo_active = getattr(self, "_demo_active_mask", None) + if demo_active is not None: + active &= demo_active.to(active.device) if not bool(active.any()): logger.log_warning( - "All expert rollout rows reached max_episode_steps; new frames are dropped." + "No active expert rollout row can accept another frame; " + "new frames are dropped." ) return @@ -1021,8 +1107,10 @@ def _write_trajectory_step(self) -> None: env_idx = torch.arange(self.num_envs, device=self.device) step = self._traj_steps mask = step < max_steps + demo_active = getattr(self, "_demo_active_mask", None) + if demo_active is not None: + mask &= demo_active.to(mask.device) if not bool(mask.any()): - self._traj_steps = (self._traj_steps + 1).clamp(max=max_steps) return idx = env_idx[mask] st = step[mask] @@ -1046,7 +1134,7 @@ def _write_trajectory_step(self) -> None: ] = obj.get_local_pose()[idx] if self._traj_raw_action is not None: self._traj_buffer["actions"][idx, st] = self._traj_raw_action[idx] - self._traj_steps = (self._traj_steps + 1).clamp(max=max_steps) + self._traj_steps[idx] += 1 def _write_rl_rollout_step( self, @@ -1086,6 +1174,87 @@ def _normalize_demo_action(self, action: EnvAction) -> EnvAction: expected_dim = int(np.prod(self.single_action_space.shape)) return self._normalize_demo_action_tensor(action, expected_dim) + def _mask_demo_action( + self, action: EnvAction, active_mask: Sequence[bool] + ) -> EnvAction: + """Accept an asynchronously completed vector-demo action. + + Raw actions may still require :class:`ActionManager` preprocessing, so + the actual hold/no-op substitution is applied to the processed command + in :meth:`_preprocess_action`. Subclasses with a specialized safe action + may override this hook and return a replacement raw action. + + Args: + action: Raw normalized action for the shared demo step. + active_mask: Rows that still participate in the episode. + + Returns: + The raw action to preprocess. + """ + self._set_demo_active_mask(active_mask) + return action + + def _mask_processed_demo_action(self, action: EnvAction) -> EnvAction: + """Replace inactive processed commands with a safe hold or no-op.""" + active_mask = self._demo_active_mask + if bool(active_mask.all()): + return action + + def replace_rows( + value: torch.Tensor, replacement: torch.Tensor, key: str + ) -> torch.Tensor: + if value.ndim < 2 or value.shape[0] != self.num_envs: + raise ValueError( + f"Cannot mask demo {key} action with shape {tuple(value.shape)}; " + f"expected a leading vector-env dimension of {self.num_envs}." + ) + masked = value.clone() + inactive = ~active_mask.to(value.device) + masked[inactive] = replacement.to(device=value.device, dtype=value.dtype)[ + inactive + ] + return masked + + measured_qpos = self.robot.get_qpos() + active_qpos = measured_qpos[:, self.active_joint_ids] + + def qpos_replacement(value: torch.Tensor) -> torch.Tensor: + """Select the measured qpos layout matching the processed command.""" + if value.shape[1:] == active_qpos.shape[1:]: + return active_qpos + if value.shape[1:] == measured_qpos.shape[1:]: + return measured_qpos + raise ValueError( + "Cannot construct a qpos hold command for processed demo action " + f"shape {tuple(value.shape)}; measured active/full layouts are " + f"{tuple(active_qpos.shape)} and {tuple(measured_qpos.shape)}." + ) + + if isinstance(action, torch.Tensor): + return replace_rows(action, qpos_replacement(action), "qpos") + if not isinstance(action, TensorDict): + raise TypeError( + "Processed demo actions must be torch.Tensor or TensorDict, " + f"got {type(action).__name__}." + ) + + masked_action = action.clone() + supported_key = False + for key in ("qpos", "qvel", "qf"): + if key not in masked_action: + continue + supported_key = True + value = masked_action[key] + replacement = ( + qpos_replacement(value) if key == "qpos" else torch.zeros_like(value) + ) + masked_action[key] = replace_rows(value, replacement, key) + if not supported_key: + raise ValueError( + "Cannot mask a processed demo TensorDict without qpos, qvel, or qf." + ) + return masked_action + def _normalize_demo_action_list( self, action_list: Sequence[EnvAction] | torch.Tensor | None ) -> Sequence[EnvAction] | torch.Tensor | None: @@ -1211,24 +1380,39 @@ def _step_action(self, action: EnvAction) -> EnvAction: Returns: The action return. """ + + def active_joint_command(command: torch.Tensor, key: str) -> torch.Tensor: + """Normalize full-robot commands to the active-joint layout.""" + active_dim = len(self.active_joint_ids) + if command.shape[-1] == active_dim: + return command.to(self.device) + full_dim = int(self.robot.get_qpos().shape[-1]) + if command.shape[-1] == full_dim: + return command[..., self.active_joint_ids].to(self.device) + raise ValueError( + f"Processed {key} action has dim {command.shape[-1]}; expected " + f"active-joint dim {active_dim} or full robot dim {full_dim}." + ) + if isinstance(action, TensorDict): # Support multiple control modes simultaneously + action = action.clone() if "qpos" in action: + action["qpos"] = active_joint_command(action["qpos"], "qpos") self.robot.set_qpos( - qpos=action["qpos"].to(self.device), joint_ids=self.active_joint_ids + qpos=action["qpos"], joint_ids=self.active_joint_ids ) if "qvel" in action: + action["qvel"] = active_joint_command(action["qvel"], "qvel") self.robot.set_qvel( - qvel=action["qvel"].to(self.device), joint_ids=self.active_joint_ids + qvel=action["qvel"], joint_ids=self.active_joint_ids ) if "qf" in action: - self.robot.set_qf( - qf=action["qf"].to(self.device), joint_ids=self.active_joint_ids - ) + action["qf"] = active_joint_command(action["qf"], "qf") + self.robot.set_qf(qf=action["qf"], joint_ids=self.active_joint_ids) elif isinstance(action, torch.Tensor): - self.robot.set_qpos( - qpos=action.to(self.device), joint_ids=self.active_joint_ids - ) + action = active_joint_command(action, "qpos") + self.robot.set_qpos(qpos=action, joint_ids=self.active_joint_ids) else: logger.log_error(f"Unsupported action type: {type(action)}") @@ -1290,8 +1474,12 @@ def _preprocess_action(self, action: EnvAction) -> EnvAction: if self._traj_buffer is not None: self._traj_raw_action = action if self.action_manager is not None: - return self.action_manager.process_action(action, mode="pre") - return super()._preprocess_action(action) + action = self.action_manager.process_action(action, mode="pre") + else: + action = super()._preprocess_action(action) + if getattr(self, "_demo_no_auto_reset", False): + action = self._mask_processed_demo_action(action) + return action def _postprocess_action(self, action): if self.action_manager is not None: @@ -1559,40 +1747,135 @@ def save_trajectory(self, path: str, env_ids: Sequence[int] | None = None) -> st return path def _save_trajectory_for_env(self, env_id: int) -> str | None: - """Auto-save one env's trajectory (best-effort; never crashes the episode).""" + """Persist one explicitly committed environment trajectory. + + I/O errors are intentionally propagated: reset is the commit boundary, + and clearing the in-memory trajectory after a failed write would report + success while losing committed data. + """ if self._traj_buffer is None or not self.cfg.trajectory_auto_save: return None if int(self._traj_steps[env_id].item()) == 0: return None - try: - base = self.cfg.trajectory_save_dir - if base is None: - base = os.path.join( - EMBODICHAIN_DEFAULT_DATA_ROOT, "trajectories", self._traj_run_id - ) - os.makedirs(base, exist_ok=True) - path = os.path.join( - base, f"traj_env{env_id}_{self._traj_save_count:06d}.pt" + base = self.cfg.trajectory_save_dir + if base is None: + base = os.path.join( + EMBODICHAIN_DEFAULT_DATA_ROOT, "trajectories", self._traj_run_id ) - self._traj_save_count += 1 - return self.save_trajectory(path, env_ids=[env_id]) - except OSError as e: - logger.log_warning( - f"Auto-save failed for env {env_id} ({e}); skipping. " - "Use save_trajectory(path) explicitly to surface IO errors." + os.makedirs(base, exist_ok=True) + path = os.path.join(base, f"traj_env{env_id}_{self._traj_save_count:06d}.pt") + saved_path = self.save_trajectory(path, env_ids=[env_id]) + self._traj_save_count += 1 + return saved_path + + def _discard_pending_recordings(self) -> None: + """Abort recorder state that has not crossed an explicit reset commit.""" + errors: list[str] = [] + if self.cfg.events and self.event_manager is not None: + from embodichain.lab.gym.envs.managers.record import record_camera_data + + for mode_cfgs in self.event_manager._mode_functor_cfgs.values(): + for functor_cfg in mode_cfgs: + if isinstance(functor_cfg.func, record_camera_data): + recorder = functor_cfg.func + recorder_name = type(recorder).__name__ + try: + recorder.discard_and_clear() + except Exception as error: + errors.append(f"{recorder_name} discard: {error}") + try: + recorder.finalize() + except Exception as error: + errors.append(f"{recorder_name} finalize: {error}") + + if self._traj_steps is not None: + try: + self._traj_steps.zero_() + except Exception as error: + errors.append(f"trajectory discard: {error}") + if self.rollout_buffer is not None and self._rollout_buffer_mode != "rl": + try: + env_ids = torch.arange(self.num_envs, device=self.device) + self._clear_expert_rollout_rows(env_ids) + self.rollout_steps.zero_() + self.current_rollout_step = 0 + except Exception as error: + errors.append(f"expert rollout discard: {error}") + + if errors: + raise RuntimeError( + f"Failed to abort {len(errors)} pending recorder operation(s): " + + "; ".join(errors) ) - return None - def close(self) -> None: - """Close the environment and release resources.""" - if self._traj_buffer is not None and self.cfg.trajectory_auto_save: - for env_id in range(self.num_envs): - self._save_trajectory_for_env(env_id) - # Finalize dataset if present - if self.dataset_manager: - self.dataset_manager.finalize() - - # Report before sim.destroy(): destroy() exits the process without - # returning to Python, so the report must be flushed first. - self._profiler.report() - self.sim.destroy() + def close(self, *, exit_process: bool | None = None) -> None: + """Abort pending data, finalize committed writes, and release resources. + + Closing is idempotent and is never an implicit episode commit. A demo + episode enters the dataset only through ``reset(save_data=True)``; + partial data left by failure, cancellation, or interpreter shutdown is + discarded before recorder finalization. + + Args: + exit_process: Forwarded to :meth:`SimulationManager.destroy` after + successful cleanup. Error paths always disable process exit so + the durability exception can propagate. + + Raises: + RuntimeError: If one or more recorders fail their durability barrier. + """ + close_lock = getattr(self, "_close_lock", None) + if close_lock is None: + close_lock = threading.RLock() + self._close_lock = close_lock + + with close_lock: + if getattr(self, "_closed", False): + close_error = getattr(self, "_close_error", None) + if close_error is not None: + raise close_error + return + + errors: list[Exception] = [] + try: + self._discard_pending_recordings() + except Exception as error: + errors.append(error) + if self.dataset_manager: + try: + self.dataset_manager.finalize() + except Exception as error: + errors.append(error) + + # Report before sim.destroy(): the default destroy path exits the + # process, so diagnostics and durability checks must finish first. + try: + self._profiler.report() + except Exception as error: + errors.append(error) + + if errors: + # Queue simulator cleanup without hiding persistence failures + # behind SimulationManager's default os._exit(0). + try: + self.sim.destroy(exit_process=False) + except Exception as error: + errors.append(error) + messages = "; ".join(str(error) for error in errors) + close_error = RuntimeError( + f"Failed to close EmbodiedEnv cleanly: {messages}" + ) + self._close_error = close_error + self._closed = True + raise close_error from errors[0] + + try: + if exit_process is None: + self.sim.destroy() + else: + self.sim.destroy(exit_process=exit_process) + except Exception as error: + self._close_error = error + self._closed = True + raise + self._closed = True diff --git a/embodichain/lab/gym/envs/managers/async_datasets.py b/embodichain/lab/gym/envs/managers/async_datasets.py index 3aa82d7ec..ebffe1edd 100644 --- a/embodichain/lab/gym/envs/managers/async_datasets.py +++ b/embodichain/lab/gym/envs/managers/async_datasets.py @@ -120,6 +120,13 @@ def __init__(self, cfg, env: EmbodiedEnv): # only ever touched from one thread (it is not thread-safe) and keeps # episode ordering deterministic. self._save_queue: "queue.Queue[Optional[tuple]]" = queue.Queue() + self._background_error_lock = threading.Lock() + self._background_errors: list[tuple[int, str]] = [] + self._async_finalize_lock = threading.Lock() + self._accepting_commits = True + self._async_finalized = False + self._async_finalize_result: Optional[str] = None + self._async_finalize_error: Optional[str] = None self._worker: threading.Thread = threading.Thread( target=self._worker_loop, name="AsyncLeRobotRecorder-worker", @@ -140,19 +147,27 @@ def _worker_loop(self) -> None: break env_id, obs_clone, action_clone, annotations, episode_metadata = item try: - self._save_single_episode( + saved = self._save_single_episode( env_id, obs_clone, action_clone, annotations=annotations, episode_metadata=episode_metadata, ) - except Exception as e: # noqa: BLE001 - worker must not die - logger.log_error( + 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)) + logger.log_warning( f"[AsyncLeRobotRecorder] Background worker failed on " - f"env {env_id}: {e}" + f"env {env_id}: {error}" ) + def _record_background_error(self, env_id: int, message: str) -> None: + """Remember one failed committed episode for the final durability check.""" + with self._background_error_lock: + self._background_errors.append((env_id, message)) + def __call__( self, env: EmbodiedEnv, @@ -183,53 +198,101 @@ def __call__( passed through by ``DatasetManager.apply``; honored in :meth:`__init__`, ignored here. """ - if env_ids is None: - env_ids = torch.arange(env.num_envs, device=env.device) - elif isinstance(env_ids, (list, range)): - env_ids = torch.tensor(list(env_ids), device=env.device) - - if len(env_ids) == 0: - return - - for env_id in env_ids.cpu().tolist(): - step = self._episode_length(env_id) - obs_view = env.rollout_buffer["obs"][env_id, :step] - action_view = env.rollout_buffer["actions"][env_id, :step] - # Clone in the caller thread: the rollout buffer is cleared and - # reused by the next episode on reset, so the worker must not hold - # a view into it. - obs_clone = obs_view.clone().cpu() - action_clone = action_view.clone().cpu() - annotations = { - key: env.rollout_buffer[key][env_id, :step].clone().cpu() - for key in DEMO_ANNOTATION_KEYS - if key in env.rollout_buffer.keys() - } - metadata_getter = getattr(env, "get_demo_episode_metadata", None) - episode_metadata = ( - copy.deepcopy(metadata_getter(env_id)) - if metadata_getter is not None - else None - ) - self._save_queue.put( - ( - env_id, - obs_clone, - action_clone, - annotations, - episode_metadata, + # Serializing enqueue with finalize prevents an episode from being + # placed behind the worker's shutdown sentinel. + with self._async_finalize_lock: + if not self._accepting_commits: + raise RuntimeError("AsyncLeRobotRecorder is already finalized") + + if env_ids is None: + env_ids = torch.arange(env.num_envs, device=env.device) + elif isinstance(env_ids, (list, range)): + env_ids = torch.tensor(list(env_ids), device=env.device) + + if len(env_ids) == 0: + return + + for env_id in env_ids.cpu().tolist(): + step = self._episode_length(env_id) + obs_view = env.rollout_buffer["obs"][env_id, :step] + action_view = env.rollout_buffer["actions"][env_id, :step] + # Clone in the caller thread: the rollout buffer is cleared and + # reused by the next episode on reset, so the worker must not hold + # a view into it. + obs_clone = obs_view.clone().cpu() + action_clone = action_view.clone().cpu() + annotations = { + key: env.rollout_buffer[key][env_id, :step].clone().cpu() + for key in DEMO_ANNOTATION_KEYS + if key in env.rollout_buffer.keys() + } + metadata_getter = getattr(env, "get_demo_episode_metadata", None) + episode_metadata = ( + copy.deepcopy(metadata_getter(env_id)) + if metadata_getter is not None + else None + ) + self._save_queue.put( + ( + env_id, + obs_clone, + action_clone, + annotations, + episode_metadata, + ) ) - ) def finalize(self) -> Optional[str]: - """Drain the background worker, then finalize the dataset. + """Drain committed writes, finalize storage, and surface all failures. + + The queue is a durability barrier for episodes explicitly committed by + ``reset(save_data=True)``. A live rollout that was never enqueued is not + saved during close. - Signals the worker to stop after finishing all queued episodes, waits - for it to exit, and delegates to the parent to flush any remaining - buffer, stop the async image writer, and finalize dataset metadata. + Returns: + The finalized dataset path. + + Raises: + RuntimeError: If any queued episode or dataset resource failed. """ - # Signal the worker to exit once the queue is drained. - self._save_queue.put(None) - self._worker.join() - # Parent flushes leftover episodes, stops image writer, finalizes. - return super().finalize() + with self._async_finalize_lock: + if self._async_finalized: + if self._async_finalize_error is not None: + raise RuntimeError(self._async_finalize_error) + return self._async_finalize_result + + self._accepting_commits = False + # FIFO ordering guarantees every commit queued before this sentinel + # is processed before the worker exits. + self._save_queue.put(None) + self._worker.join() + + parent_error: Optional[str] = None + try: + self._async_finalize_result = super().finalize() + except Exception as error: # noqa: BLE001 - combine all failures + parent_error = str(error) + + with self._background_error_lock: + background_errors = list(self._background_errors) + + failures: list[str] = [] + if background_errors: + episode_details = "; ".join( + f"env {env_id}: {message}" for env_id, message in background_errors + ) + failures.append( + f"failed to persist {len(background_errors)} committed " + f"episode(s): {episode_details}" + ) + if parent_error is not None: + failures.append(f"storage finalization: {parent_error}") + + self._async_finalized = True + if failures: + self._async_finalize_error = ( + "AsyncLeRobotRecorder finalization failed: " + "; ".join(failures) + ) + raise RuntimeError(self._async_finalize_error) + + return self._async_finalize_result diff --git a/embodichain/lab/gym/envs/managers/dataset_manager.py b/embodichain/lab/gym/envs/managers/dataset_manager.py index fcb98dafc..f0174ec4a 100644 --- a/embodichain/lab/gym/envs/managers/dataset_manager.py +++ b/embodichain/lab/gym/envs/managers/dataset_manager.py @@ -19,6 +19,7 @@ from __future__ import annotations import inspect +import threading from typing import TYPE_CHECKING, Any, Dict, Optional, Union from collections.abc import Sequence @@ -81,6 +82,10 @@ def __init__(self, cfg: object, env: EmbodiedEnv): self._mode_functor_names: dict[str, list[str]] = {} self._mode_functor_cfgs: dict[str, list[DatasetFunctorCfg]] = {} self._mode_class_functor_cfgs: dict[str, list[DatasetFunctorCfg]] = {} + self._finalize_lock = threading.Lock() + self._finalized = False + self._finalize_result: Optional[str] = None + self._finalize_error: Optional[str] = None # Call base class to parse functors super().__init__(cfg, env) @@ -243,32 +248,61 @@ def apply( ) def finalize(self) -> Optional[str]: - """Finalize all dataset functors. + """Finalize every dataset functor exactly once. - Called when the environment is closed. Saves any remaining episodes - and finalizes all datasets. + Finalization is a storage barrier only; individual recorders must not + implicitly commit live episode buffers here. All functors are attempted + even when one fails, and their failures are reported together. Returns: - Path to the first saved dataset, or None if failed. - """ - dataset_paths = [] + Path to the first finalized dataset, or ``None`` if none was returned. - # Call finalize on all functor instances across all modes - for mode_cfgs in self._mode_functor_cfgs.values(): - for functor_cfg in mode_cfgs: - if hasattr(functor_cfg.func, "finalize"): + Raises: + RuntimeError: If one or more functors fail to finalize. + """ + with self._finalize_lock: + if self._finalized: + if self._finalize_error is not None: + raise RuntimeError(self._finalize_error) + return self._finalize_result + + dataset_paths: list[str] = [] + errors: list[str] = [] + + # Call every functor even when an earlier cleanup failed. + for mode, mode_cfgs in self._mode_functor_cfgs.items(): + names = self._mode_functor_names.get(mode, []) + for index, functor_cfg in enumerate(mode_cfgs): + functor = functor_cfg.func + if not hasattr(functor, "finalize"): + continue + functor_name = ( + names[index] if index < len(names) else type(functor).__name__ + ) try: - path = functor_cfg.func.finalize() + path = functor.finalize() if path: dataset_paths.append(path) - except Exception as e: - logger.log_error(f"Failed to finalize functor: {e}") + except Exception as error: # noqa: BLE001 - aggregate cleanup + errors.append(f"{functor_name}: {error}") + + self._finalize_result = dataset_paths[0] if dataset_paths else None + self._finalized = True + + if errors: + self._finalize_error = ( + f"Failed to finalize {len(errors)} dataset functor(s): " + + "; ".join(errors) + ) + raise RuntimeError(self._finalize_error) - if dataset_paths: - logger.log_info(f"Finalized {len(dataset_paths)} datasets") - return dataset_paths[0] + if dataset_paths: + logger.log_info(f"Finalized {len(dataset_paths)} datasets") + return self._finalize_result - return None + def close(self) -> Optional[str]: + """Finalize all dataset functors; repeated calls are safe.""" + return self.finalize() def get_cached_data(self) -> list[Dict[str, Any]]: """Get cached data from all dataset functors (for online training). diff --git a/embodichain/lab/gym/envs/managers/datasets.py b/embodichain/lab/gym/envs/managers/datasets.py index 2f1bc34fd..ef32709f4 100644 --- a/embodichain/lab/gym/envs/managers/datasets.py +++ b/embodichain/lab/gym/envs/managers/datasets.py @@ -165,6 +165,10 @@ def __init__(self, cfg: DatasetFunctorCfg, env: EmbodiedEnv): self.total_time: float = 0.0 self.curr_episode: int = 0 self._metadata_lock = threading.Lock() + self._finalize_lock = threading.Lock() + self._finalized = False + self._finalize_result: Optional[str] = None + self._finalize_error: Optional[str] = None # Initialize dataset self._initialize_dataset() @@ -205,15 +209,19 @@ def __call__( ``DatasetManager.apply`` via ``**functor_cfg.params``. They are read in :meth:`__init__` and ignored here. """ - # If env_ids is None, check all environments for completed episodes - if env_ids is None: - env_ids = torch.arange(env.num_envs, device=env.device) - elif isinstance(env_ids, (list, range)): - env_ids = torch.tensor(list(env_ids), device=env.device) + with self._finalize_lock: + if self._finalized: + raise RuntimeError("LeRobotRecorder is already finalized") - # Save episodes for specified environments - if len(env_ids) > 0: - self._save_episodes(env_ids) + # If env_ids is None, check all environments for completed episodes + if env_ids is None: + env_ids = torch.arange(env.num_envs, device=env.device) + elif isinstance(env_ids, (list, range)): + env_ids = torch.tensor(list(env_ids), device=env.device) + + # Save episodes for specified environments + if len(env_ids) > 0: + self._save_episodes(env_ids) def _save_episodes( self, @@ -240,13 +248,17 @@ def _save_episodes( episode_metadata = ( metadata_getter(env_id) if metadata_getter is not None else None ) - self._save_single_episode( + saved = self._save_single_episode( env_id, obs_list, action_list, annotations=annotations, episode_metadata=episode_metadata, ) + if not saved: + raise RuntimeError( + f"Committed episode for env {env_id} was not persisted." + ) def _episode_length(self, env_id: int) -> int: """Return the valid buffered length for one environment.""" @@ -309,14 +321,17 @@ def _save_single_episode( current_episode_time = len(obs_list) / fps if fps > 0 else 0 episode_extra_info = extra_info.copy() + previous_total_time = self.total_time self.total_time += current_episode_time episode_extra_info["total_time"] = self.total_time depth_prefix = f"{LeRobotKey.OBS_PREFIX.value}depth." + episode_index = self.curr_episode + dataset_committed = False try: if self._depth_manager is not None: self._depth_manager.start_episode( - self.curr_episode, list(self._depth_sensor_specs.keys()) + episode_index, list(self._depth_sensor_specs.keys()) ) for frame_index, (obs, action) in enumerate( tqdm.tqdm( @@ -366,8 +381,13 @@ def _save_single_episode( self.dataset.add_frame(frame) self.dataset.save_episode() + # LeRobot has committed this index. Advance immediately so a later + # depth/metadata failure cannot make the next queued episode reuse + # and overwrite the same sidecar filename. + dataset_committed = True + self.curr_episode += 1 if self._depth_manager is not None: - self._depth_manager.end_episode(self.curr_episode) + self._depth_manager.end_episode(episode_index) sidecar_metadata = dict(episode_metadata or {}) if not sidecar_metadata.get("segments"): @@ -388,32 +408,32 @@ def _save_single_episode( sidecar_metadata.update( { "schema_version": DEMO_SCHEMA_VERSION, - "lerobot_episode_index": self.curr_episode, + "lerobot_episode_index": episode_index, "env_id": env_id, "length": len(obs_list), "instruction": task, } ) - try: - self._write_episode_metadata(sidecar_metadata) - except OSError as error: - logger.log_warning( - "Saved the LeRobot episode but could not append EmbodiChain " - f"segment metadata: {error}" - ) + self._write_episode_metadata(sidecar_metadata) logger.log_info( f"[LeRobotRecorder] Saved dataset to: {self.dataset_path}\n" - f" Episode {self.curr_episode} (env {env_id}): {len(obs_list)} frames" + f" Episode {episode_index} (env {env_id}): {len(obs_list)} frames" ) - self.curr_episode += 1 return True - except Exception as e: - logger.log_error(f"Failed to save episode {env_id}: {e}") - if self._depth_manager is not None: - self._depth_manager.abort_episode() - return False + except Exception as error: + if not dataset_committed: + self.total_time = previous_total_time + if self._depth_manager is not None and not dataset_committed: + try: + self._depth_manager.abort_episode() + except Exception as abort_error: # noqa: BLE001 - preserve primary + error.add_note( + "Depth sidecar abort also failed: " + f"{type(abort_error).__name__}: {abort_error}" + ) + raise @staticmethod def _task_for_frame( @@ -456,40 +476,71 @@ def _write_episode_metadata(self, metadata: Mapping[str, Any]) -> None: stream.write("\n") def finalize(self) -> Optional[str]: - """Finalize the dataset.""" - # Save any remaining episodes - rollout_steps = getattr(self._env, "rollout_steps", None) - if rollout_steps is not None: - active_env_ids = (rollout_steps > 0).nonzero(as_tuple=False).squeeze(-1) - elif self._env.current_rollout_step > 0: - active_env_ids = torch.arange(self._env.num_envs, device=self._env.device) - else: - active_env_ids = torch.empty(0, dtype=torch.long, device=self._env.device) - if len(active_env_ids) > 0: - self._save_episodes(active_env_ids) + """Finalize resources without implicitly committing a partial episode. - try: + Episodes are committed only when :meth:`__call__` is invoked by an + explicit ``reset(save_data=True)``. Closing the environment therefore + leaves any still-live rollout buffer uncommitted. + + Returns: + The finalized dataset path, or ``None`` when no dataset exists. + + Raises: + RuntimeError: If one or more dataset resources cannot be finalized. + """ + with self._finalize_lock: + if self._finalized: + if self._finalize_error is not None: + raise RuntimeError(self._finalize_error) + return self._finalize_result + + errors: list[str] = [] if self.dataset is not None: # Flush + stop the async image writer (if enabled) so every - # queued PNG write lands on disk before metadata is finalized. + # explicitly committed frame lands on disk before metadata is + # finalized. if self.dataset.image_writer is not None: - self.dataset.stop_image_writer() - self.dataset.finalize() - # Finalize the depth sidecar metadata (depth videos are already - # written per episode; this flushes depth_meta.json). - if self._depth_manager is not None: + try: + self.dataset.stop_image_writer() + except Exception as error: # noqa: BLE001 - aggregate cleanup + errors.append(f"image writer: {error}") + try: + self.dataset.finalize() + except Exception as error: # noqa: BLE001 - aggregate cleanup + errors.append(f"LeRobot dataset: {error}") + + # Depth videos are written per committed episode; this only flushes + # their metadata and must still be attempted if LeRobot cleanup fails. + if self._depth_manager is not None: + try: self._depth_manager.finalize() + except Exception as error: # noqa: BLE001 - aggregate cleanup + errors.append(f"depth sidecar: {error}") + + self._finalize_result = ( + self.dataset_path if self.dataset is not None else None + ) + self._finalized = True + + if errors: + self._finalize_error = ( + "LeRobotRecorder failed to finalize " + f"{len(errors)} resource(s): {'; '.join(errors)}" + ) + raise RuntimeError(self._finalize_error) + + if self.dataset is not None: logger.log_info( f"[LeRobotRecorder] Dataset finalized successfully\n" f" Path: {self.dataset_path}\n" f" Total episodes: {self.curr_episode}\n" f" Total time: {self.total_time:.2f}s" ) - return self.dataset_path - except Exception as e: - logger.log_error(f"[LeRobotRecorder] Failed to finalize dataset: {e}") + return self._finalize_result - return None + def close(self) -> Optional[str]: + """Finalize the recorder; repeated calls are safe.""" + return self.finalize() def _parse_depth_video_cfg(self, params: Dict) -> DepthVideoCfg: """Parse the optional ``depth_video`` parameter into a config. diff --git a/embodichain/lab/gym/envs/managers/record.py b/embodichain/lab/gym/envs/managers/record.py index b197b2a7a..ae4ea15c2 100644 --- a/embodichain/lab/gym/envs/managers/record.py +++ b/embodichain/lab/gym/envs/managers/record.py @@ -16,12 +16,15 @@ from __future__ import annotations -import torch import os import random -import numpy as np +import threading +from collections import deque from typing import TYPE_CHECKING, Literal, Union, List +import numpy as np +import torch + from dexsim.utility import images_to_video from embodichain.lab.gym.envs.managers import Functor, FunctorCfg from embodichain.lab.sim.sensors.camera import CameraCfg, Camera @@ -87,6 +90,13 @@ def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv): self._save_path = cfg.params.get("save_path", "./outputs/videos") self._current_episode = 0 self._frames: List[np.ndarray] = [] + self._finalize_lock = threading.Lock() + self._finalized = False + + def _ensure_open(self) -> None: + """Raise when capture or commit is attempted after finalization.""" + if getattr(self, "_finalized", False): + raise RuntimeError(f"{type(self).__name__} is already finalized") def _draw_frames_into_one_image(self, frames: torch.Tensor) -> torch.Tensor: """ @@ -126,7 +136,7 @@ def _draw_frames_into_one_image(self, frames: torch.Tensor) -> torch.Tensor: return result - def save_and_clear(self) -> None: + def save_and_clear(self, env_ids: Union[torch.Tensor, None] = None) -> None: """Save recorded frames as video and clear the buffer. This method is called from :meth:`EmbodiedEnv._initialize_episode` to ensure @@ -134,6 +144,7 @@ def save_and_clear(self) -> None: final episode's frames are lost because the save previously relied on detecting a reset inside :meth:`__call__`. """ + self._ensure_open() if len(self._frames) > 0: video_name = f"episode_{self._current_episode}_{self._name}" images_to_video(self._frames, self._save_path, video_name, fps=20) @@ -141,10 +152,24 @@ def save_and_clear(self) -> None: self._current_episode += 1 self._frames = [] - def discard_and_clear(self) -> None: + def discard_and_clear(self, env_ids: Union[torch.Tensor, None] = None) -> None: """Discard recorded frames without creating an episode video.""" self._frames = [] + def finalize(self) -> None: + """Discard uncommitted frames and close the recorder exactly once.""" + if not hasattr(self, "_finalize_lock"): + self._finalize_lock = threading.Lock() + with self._finalize_lock: + if getattr(self, "_finalized", False): + return + self.discard_and_clear() + self._finalized = True + + def close(self) -> None: + """Finalize the recorder; repeated calls are safe.""" + self.finalize() + def __call__( self, env: EmbodiedEnv, @@ -163,6 +188,7 @@ def __call__( max_env_num: int = 16, save_path: str = "./outputs/videos", ): + self._ensure_open() self.camera.update(fetch_only=True) data = self.camera.get_data() rgb = data["color"] @@ -181,17 +207,98 @@ def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv): self._num_envs = min(4, getattr(env, "num_envs", 1)) self._frames_list = [[] for _ in range(self._num_envs)] self._ep_idx = [0 for _ in range(self._num_envs)] + self._committed_env_episodes = [deque() for _ in range(self._num_envs)] + self._async_camera_finalize_lock = threading.Lock() + self._async_camera_finalized = False + self._async_camera_finalize_error: str | None = None + + def _normalize_env_ids(self, env_ids: Union[torch.Tensor, None]) -> list[int]: + """Return recorder-local environment IDs for a transaction boundary.""" + if env_ids is None: + return list(range(self._num_envs)) + if isinstance(env_ids, torch.Tensor): + values = env_ids.reshape(-1).cpu().tolist() + else: + values = list(env_ids) + return [int(env_id) for env_id in values if int(env_id) < self._num_envs] + + def _flush_committed_episodes(self) -> None: + """Persist every complete FIFO set of explicitly committed env episodes.""" + while all(self._committed_env_episodes): + episode_frames = [queue[0] for queue in self._committed_env_episodes] + min_len = min(len(frames) for frames in episode_frames) + big_frames = [] + for frame_id in range(min_len): + frames = [frames[frame_id] for frames in episode_frames] + frames_tensor = torch.from_numpy(np.stack(frames)).to(torch.uint8) + big_frame = ( + self._draw_frames_into_one_image(frames_tensor)[..., :3] + .cpu() + .numpy() + ) + big_frames.append(big_frame) - def save_and_clear(self) -> None: - """No-op for the async variant; saving is handled inside :meth:`__call__`.""" - pass + video_name = f"ep{self._current_episode}_{self._name}_allenvs" + # Peek above and pop only after persistence succeeds. A failed + # write therefore remains observable and retryable at finalize(). + images_to_video(big_frames, self._save_path, video_name, fps=20) + for queue in self._committed_env_episodes: + queue.popleft() + self._current_episode += 1 - def discard_and_clear(self) -> None: - """Discard in-flight frames and pending invalid environment episodes.""" - self._frames_list = [[] for _ in range(self._num_envs)] - pending = getattr(self, "_pending_env_episodes", None) - if pending is not None: - pending.clear() + def save_and_clear(self, env_ids: Union[torch.Tensor, None] = None) -> None: + """Commit selected rows immediately instead of waiting for a later step.""" + self._ensure_open() + for env_id in self._normalize_env_ids(env_ids): + frames = self._frames_list[env_id] + if frames: + self._committed_env_episodes[env_id].append(frames) + self._frames_list[env_id] = [] + self._ep_idx[env_id] += 1 + self._flush_committed_episodes() + + def discard_and_clear(self, env_ids: Union[torch.Tensor, None] = None) -> None: + """Discard live frames while preserving already committed episodes.""" + super().discard_and_clear() + for env_id in self._normalize_env_ids(env_ids): + self._frames_list[env_id] = [] + + def finalize(self) -> None: + """Flush committed frame sets and reject incomplete committed batches.""" + if not hasattr(self, "_async_camera_finalize_lock"): + self._async_camera_finalize_lock = threading.Lock() + self._async_camera_finalized = False + self._async_camera_finalize_error = None + with self._async_camera_finalize_lock: + if self._async_camera_finalized: + if self._async_camera_finalize_error is not None: + raise RuntimeError(self._async_camera_finalize_error) + return + + errors: list[str] = [] + try: + self._flush_committed_episodes() + except Exception as error: # noqa: BLE001 - finish recorder cleanup + errors.append(f"video persistence: {error}") + + pending_counts = [len(queue) for queue in self._committed_env_episodes] + if any(pending_counts): + errors.append( + "incomplete committed environment batch " + f"(pending episodes per env: {pending_counts})" + ) + + try: + super().finalize() + except Exception as error: # noqa: BLE001 - aggregate cleanup + errors.append(f"recorder cleanup: {error}") + + self._async_camera_finalized = True + if errors: + self._async_camera_finalize_error = ( + "Async camera recorder finalization failed: " + "; ".join(errors) + ) + raise RuntimeError(self._async_camera_finalize_error) def __call__( self, @@ -211,6 +318,7 @@ def __call__( max_env_num: int = 16, save_path: str = "./outputs/videos", ): + self._ensure_open() self.camera.update(fetch_only=True) data = self.camera.get_data() rgb = data["color"] # shape: (num_envs, H, W, 4) @@ -222,47 +330,6 @@ def __call__( for i in range(self._num_envs): self._frames_list[i].append(rgb_np[i][..., :]) - # Check if elapsed_steps==1 (just reset) - elapsed = env.elapsed_steps - if isinstance(elapsed, torch.Tensor): - elapsed_np = elapsed.cpu().numpy() - else: - elapsed_np = elapsed - # Only check reset for the first 4 environments - ready_envs = [ - i - for i in range(self._num_envs) - if elapsed_np[i] == 1 and len(self._frames_list[i]) > 1 - ] - # Used to temporarily store episode frames for each env - if not hasattr(self, "_pending_env_episodes"): - self._pending_env_episodes = {} - for i in ready_envs: - if i not in self._pending_env_episodes: - self._pending_env_episodes[i] = self._frames_list[i][:-1] - self._frames_list[i] = [ - self._frames_list[i][-1] - ] # Only keep the first frame after reset - self._ep_idx[i] += 1 - # If all specified envs have collected frames, concatenate and save - if len(self._pending_env_episodes) == self._num_envs: - min_len = min(len(frames) for frames in self._pending_env_episodes.values()) - big_frames = [] - for j in range(min_len): - frames = [ - self._pending_env_episodes[i][j] for i in range(self._num_envs) - ] - frames_tensor = torch.from_numpy(np.stack(frames)).to(torch.uint8) - big_frame = ( - self._draw_frames_into_one_image(frames_tensor)[..., :3] - .cpu() - .numpy() - ) - big_frames.append(big_frame) - video_name = f"ep{self._ep_idx[0]-1}_{self._name}_allenvs" - images_to_video(big_frames, save_path, video_name, fps=20) - self._pending_env_episodes.clear() - class validation_cameras(Functor): """ diff --git a/embodichain/lab/scripts/run_env.py b/embodichain/lab/scripts/run_env.py index 4209d8889..65453bd62 100644 --- a/embodichain/lab/scripts/run_env.py +++ b/embodichain/lab/scripts/run_env.py @@ -49,6 +49,11 @@ def _progress_wrapper(actions: Iterable[Any], description: str) -> Iterable[Any] return tqdm.tqdm(actions, desc=description, unit="step") +def _abort_pending_episode(env: Any) -> None: + """Discard buffered data before retrying or closing an environment.""" + env.reset(options={"save_data": False}) + + def generate_and_execute_action_list( env: gymnasium.Env, idx: int, @@ -128,9 +133,10 @@ def generate_function( raise ValueError(f"max_attempts must be at least 1, got {max_attempts}.") if reset_before: - env.reset(options={"save_data": False}) + _abort_pending_episode(env) for attempt in range(1, max_attempts + 1): + commit_succeeded = False try: result: DemoEpisodeResult = execute_demo_episode( env, @@ -138,22 +144,23 @@ def generate_function( progress=_progress_wrapper, **kwargs, ) - except Exception: - # Never let close()/finalize() commit a partially written rollout - # when planning, normalization, or stepping raises. - env.reset(options={"save_data": False}) - raise - if result.completed and result.all_success: - # reset() is the commit boundary: dataset functors consume the - # whole episode once, then buffers and scene state are reset. - env.reset() - return True + if result.completed and result.all_success: + # reset() is the commit boundary: dataset functors consume the + # whole episode once, then buffers and scene state are reset. + env.reset() + commit_succeeded = True + return True + finally: + # ``finally`` also covers KeyboardInterrupt, SystemExit, and + # GeneratorExit. A failed commit is aborted as well, so close() + # can never implicitly persist the pending partial episode. + if not commit_succeeded: + _abort_pending_episode(env) log_warning( f"Episode {time_id} attempt {attempt}/{max_attempts} failed: " f"{result.terminal_reason}. Discarding {result.length} frames." ) - env.reset(options={"save_data": False}) return False @@ -161,6 +168,9 @@ def generate_function( def replay(env, trajectory_path: str, mode: str = "kinematic") -> None: """Replay a recorded trajectory. + The caller retains ownership of ``env``. Wrapper-specific replay state is + restored before returning, but the environment is not closed here. + Args: env: The environment built from the same config that recorded the trajectory (wrapped via :class:`ReplayWrapper`). @@ -183,7 +193,12 @@ def replay(env, trajectory_path: str, mode: str = "kinematic") -> None: else: replay_auto(replay_env, mode) finally: - replay_env.close() + # ReplayWrapper.close() also closes its wrapped environment. Restore + # its local state here and leave the single close() to cli(). + try: + replay_env.env.sim.enable_physics(True) + finally: + replay_env.env._replay_no_auto_reset = False def replay_auto(replay_env: ReplayWrapper, mode: str) -> None: @@ -395,7 +410,8 @@ def replay_control(replay_env: ReplayWrapper) -> None: _run_replay_control_loop(replay_env, control_input) -def main(args, env, gym_config): +def main(args: Any, env: Any, gym_config: dict[str, Any]) -> None: + """Run the selected workflow without taking ownership of ``env``.""" if getattr(args, "replay", False): log_info("Replay mode.", color="green") replay( @@ -413,56 +429,45 @@ def main(args, env, gym_config): return log_info("Start offline data generation.", color="green") - try: - # Prepare one clean scene. Every successful generate_function call - # commits via reset(), leaving the next episode ready to plan. - env.reset(options={"save_data": False}) - for i in range(gym_config.get("max_episodes", 1)): - generated = generate_function( - env, - time_id=i, - save_path=getattr(args, "save_path", ""), - save_video=getattr(args, "save_video", False), - debug_mode=getattr(args, "debug_mode", False), - regenerate=getattr(args, "regenerate", False), - max_attempts=gym_config.get("demo_max_attempts", 3), - reset_before=False, + # Prepare one clean scene. Every successful generate_function call commits + # via reset(), leaving the next episode ready to plan. + _abort_pending_episode(env) + for i in range(gym_config.get("max_episodes", 1)): + generated = generate_function( + env, + time_id=i, + save_path=getattr(args, "save_path", ""), + save_video=getattr(args, "save_video", False), + debug_mode=getattr(args, "debug_mode", False), + regenerate=getattr(args, "regenerate", False), + max_attempts=gym_config.get("demo_max_attempts", 3), + reset_before=False, + ) + if not generated: + raise RuntimeError( + f"Failed to generate episode {i} after " + f"{gym_config.get('demo_max_attempts', 3)} attempts." ) - if not generated: - raise RuntimeError( - f"Failed to generate episode {i} after " - f"{gym_config.get('demo_max_attempts', 3)} attempts." - ) - # Log the trajectory save location BEFORE env.close() (in the finally - # below) tears down the sim and, by default, os._exit()s the process. - if getattr(args, "record_trajectory", False): - save_dir = args.trajectory_save_dir - if save_dir is None: - import os + # Log the trajectory save location before cli() tears down the sim and, by + # default, os._exit()s the process. + if getattr(args, "record_trajectory", False): + save_dir = args.trajectory_save_dir + if save_dir is None: + import os - from embodichain.data.constants import EMBODICHAIN_DEFAULT_DATA_ROOT + from embodichain.data.constants import EMBODICHAIN_DEFAULT_DATA_ROOT - save_dir = os.path.join( - EMBODICHAIN_DEFAULT_DATA_ROOT, - "trajectories", - env.unwrapped._traj_run_id, - ) - log_info( - f"Trajectories recorded to: {save_dir} " - "(replay with --replay --replay_trajectory )", - color="green", + save_dir = os.path.join( + EMBODICHAIN_DEFAULT_DATA_ROOT, + "trajectories", + env.unwrapped._traj_run_id, ) - finally: - # Drain the dataset recorder and finalize the LeRobot dataset before - # the process exits. This is REQUIRED for AsyncLeRobotRecorder: its - # background worker only *enqueues* episodes during reset, so without - # close() the worker is killed at exit and no data reaches disk. - # env.close() -> dataset_manager.finalize() drains the worker + flushes - # meta/stats, then sim.destroy() tears down the sim. sim.destroy() exits - # the process without returning, so this MUST be the last thing main() - # does. - env.close() + log_info( + f"Trajectories recorded to: {save_dir} " + "(replay with --replay --replay_trajectory )", + color="green", + ) def preview(env: gymnasium.Env) -> None: @@ -507,7 +512,7 @@ def preview(env: gymnasium.Env) -> None: elif txt == "q": end = True - exit(0) + return def _create_parser() -> argparse.ArgumentParser: @@ -542,6 +547,45 @@ def _create_parser() -> argparse.ArgumentParser: return parser +def _abort_and_close_env(env: Any, *, exit_process: bool | None = None) -> None: + """Abort pending data, then close the environment exactly once. + + Args: + env: Gym environment or wrapper. + exit_process: Optional process-exit policy for ``EmbodiedEnv.close``. + An abort failure always forces ``False`` so the error can propagate. + """ + abort_error: BaseException | None = None + close_error: BaseException | None = None + try: + _abort_pending_episode(env) + except BaseException as error: + abort_error = error + try: + if exit_process is None and abort_error is None: + env.close() + else: + target = getattr(env, "unwrapped", env) + target.close( + exit_process=False if abort_error is not None else exit_process + ) + except BaseException as error: + close_error = error + + # Recorder finalization is a durability barrier. Never turn a failed flush + # into a warning that lets an apparently successful data-generation run + # continue. + if close_error is not None: + if abort_error is not None: + close_error.add_note( + "Pending episode abort also failed: " + f"{type(abort_error).__name__}: {abort_error}" + ) + raise close_error + if abort_error is not None: + raise abort_error + + def cli(argv: Sequence[str] | None = None) -> None: """Command-line interface for environment runner. @@ -586,19 +630,41 @@ def cli(argv: Sequence[str] | None = None) -> None: # to skip the unsafe teardown entirely. When that env var is disabled (e.g. # dev/test), ``flush_cleanup_queue`` drains the queue and runs the deferred # destruction + GC + scene-barrier so we still exit cleanly. + body_error: BaseException | None = None try: main(args, env, gym_config) + except BaseException as error: + body_error = error + raise finally: try: - env.close() - except Exception as e: - log_warning(f"Failed to close environment: {e}") - try: - from embodichain.lab.sim.sim_manager import SimulationManager + # close() may auto-save trajectory state and finalizes asynchronous + # dataset writers. Resolve the pending transaction as an abort + # first, including when main() exits via an interrupt or SystemExit. + _abort_and_close_env( + env, + # Successful CLI runs keep the existing fast-exit default. + # While unwinding, cleanup must return so the original error + # (including Ctrl-C) remains observable to the caller/shell. + exit_process=False if body_error is not None else None, + ) + except BaseException as cleanup_error: + if body_error is None: + raise + if isinstance(body_error, SystemExit) and body_error.code in (None, 0): + # A nominal zero exit must not hide a failed recorder barrier. + raise cleanup_error + body_error.add_note( + "Environment cleanup also failed: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + finally: + try: + from embodichain.lab.sim.sim_manager import SimulationManager - SimulationManager.flush_cleanup_queue() - except Exception as e: - log_warning(f"Failed to flush simulation cleanup queue: {e}") + SimulationManager.flush_cleanup_queue() + except Exception as error: + log_warning(f"Failed to flush simulation cleanup queue: {error}") if __name__ == "__main__": diff --git a/tests/data_pipeline/depth_video/test_writer.py b/tests/data_pipeline/depth_video/test_writer.py index 0f7127605..7e19f7c98 100644 --- a/tests/data_pipeline/depth_video/test_writer.py +++ b/tests/data_pipeline/depth_video/test_writer.py @@ -19,6 +19,7 @@ import json from pathlib import Path +from unittest.mock import MagicMock import av import numpy as np @@ -205,6 +206,20 @@ def test_multiple_episodes_accumulate_in_metadata(self, tmp_path: Path): meta = json.loads((tmp_path / "depth_meta.json").read_text()) assert sorted(meta["sensors"]["camera"]["episodes"]) == ["0", "1", "2"] + def test_end_episode_surfaces_encoder_close_failure(self, tmp_path: Path): + """A missing committed depth video is a visible durability failure.""" + mgr = self._make_manager(tmp_path) + writer = MagicMock() + writer.close.side_effect = OSError("encoder failed") + mgr._writers = {"camera": writer} + mgr._episode_sensors = ["camera"] + + with pytest.raises(RuntimeError, match="encoder failed"): + mgr.end_episode(0) + + writer.abort.assert_called_once_with() + assert mgr._writers == {} + class TestCodec: """Codec detection tests.""" diff --git a/tests/data_pipeline/test_online_data.py b/tests/data_pipeline/test_online_data.py index f832bb83f..70ef32d53 100644 --- a/tests/data_pipeline/test_online_data.py +++ b/tests/data_pipeline/test_online_data.py @@ -16,11 +16,10 @@ """Unit tests for OnlineDataset and OnlineDataEngine. -These tests do **not** start a real simulation subprocess. Instead, -``_make_fake_engine`` builds an ``OnlineDataEngine`` instance, directly injects -a pre-filled ``shared_buffer`` TensorDict with known random data, sets the -``_init_signal``, and sets ``_lock_index`` to ``[-1, -1]`` (no locked rows), -bypassing ``start()`` entirely. +These tests do **not** start a real simulation subprocess. ``_make_fake_engine`` +injects a pre-filled shared buffer for sampling tests, while lifecycle tests +drive ``start()`` with a lightweight process mock and a real multiprocessing +error pipe. This exercises all public logic in ``sample_batch``, ``_trigger_refill_if_needed``, and ``OnlineDataset.__iter__`` without GPU or @@ -30,6 +29,8 @@ from __future__ import annotations import multiprocessing as mp +import os +import threading import unittest from unittest.mock import MagicMock @@ -39,13 +40,21 @@ from tensordict import TensorDict from torch.utils.data import DataLoader +from embodichain.data_pipeline.engine import ( + OnlineDataEngineState, + OnlineDataWorkerError, +) +from embodichain.data_pipeline.engine import data as engine_module from embodichain.data_pipeline.datasets import ( ChunkSizeSampler, GMMChunkSampler, OnlineDataset, UniformChunkSampler, ) -from embodichain.data_pipeline.engine.data import OnlineDataEngine, OnlineDataEngineCfg +from embodichain.data_pipeline.engine.data import ( + OnlineDataEngine, + OnlineDataEngineCfg, +) # --------------------------------------------------------------------------- # Constants @@ -69,6 +78,7 @@ def _make_fake_engine( refill_threshold: int = 1000, lock_start: int = -1, lock_end: int = -1, + state: OnlineDataEngineState = OnlineDataEngineState.READY, ) -> OnlineDataEngine: """Build an OnlineDataEngine with a pre-filled shared buffer, bypassing start(). @@ -82,6 +92,7 @@ def _make_fake_engine( accidental refill triggers in most tests. lock_start: Write-lock range start (``-1`` means no lock). lock_end: Write-lock range end. + state: Initial lifecycle state for the synthetic engine. Returns: A configured OnlineDataEngine whose ``shared_buffer`` contains valid @@ -116,22 +127,78 @@ def _make_fake_engine( engine.device = shared_buffer.device # Interprocess primitives — use the same mp context consistently to avoid - engine._mp_ctx = mp.get_context("spawn") + engine._mp_ctx = mp.get_context("forkserver") engine._lock_index = engine._mp_ctx.Array("i", [lock_start, lock_end]) engine._fill_signal = engine._mp_ctx.Event() engine._init_signal = engine._mp_ctx.Event() - engine._init_signal.set() # mark as initialised + if state is OnlineDataEngineState.READY: + engine._init_signal.set() engine._close_signal = engine._mp_ctx.Event() engine._sample_count = engine._mp_ctx.Value("i", 0) + engine._state_value = engine._mp_ctx.Value("i", engine_module._STATE_TO_CODE[state]) + engine._worker_failed_signal = engine._mp_ctx.Event() + engine._worker_error_buffer = engine._mp_ctx.Array( + "B", engine_module._ERROR_BUFFER_SIZE + ) + engine._worker_error_length = engine._mp_ctx.Value("i", 0) + engine._channel_error = None + engine._worker_error = None + engine._owner_pid = os.getpid() + engine._lifecycle_condition = threading.Condition(threading.RLock()) + engine._stop_requested = False + engine._cleanup_complete = False + engine._forced_shutdown_attempted = False + engine._monitor_thread = None # Sampling tests intentionally bypass the simulation worker. Starting it # here would exercise unrelated environment setup and make teardown wait # for the subprocess timeout when the synthetic config cannot run a task. engine._sim_process = None + if state is OnlineDataEngineState.READY: + engine._sim_process = MagicMock() + engine._sim_process.is_alive.return_value = True + + def finish_join(*args, **kwargs) -> None: + engine._sim_process.is_alive.return_value = False + + engine._sim_process.join.side_effect = finish_join return engine +def _install_fake_starting_worker( + engine: OnlineDataEngine, *, initialize: bool = True +) -> MagicMock: + """Install a process mock that optionally completes the initial fill.""" + process = MagicMock(pid=1234, exitcode=None) + process.is_alive.return_value = True + + def join(timeout=None) -> None: + if timeout is None: + engine._close_signal.wait(timeout=5.0) + if engine._close_signal.is_set(): + process.is_alive.return_value = False + + process.join.side_effect = join + if initialize: + process.start.side_effect = engine._init_signal.set + process_context = MagicMock() + process_context.Process.return_value = process + engine._mp_ctx = process_context + return process + + +def _sample_engine_in_subprocess( + engine: OnlineDataEngine, result_queue: mp.Queue +) -> None: + """Sample once in a real consumer process and return a pickle-safe result.""" + try: + sample = engine.sample_batch(batch_size=1, chunk_size=1) + result_queue.put(("ok", tuple(sample.shape))) + except BaseException as error: + result_queue.put(("error", type(error).__name__, str(error))) + + # =========================================================================== # TestOnlineDataEngine # =========================================================================== @@ -143,29 +210,397 @@ class TestOnlineDataEngine: def setup_method(self) -> None: self.engine = _make_fake_engine() - def test_start_raises_if_worker_exits_before_initialization(self) -> None: - """A failed initial fill is reported instead of waiting forever.""" - process = MagicMock(pid=1234, exitcode=1) - process.is_alive.return_value = False - process_context = MagicMock() - process_context.Process.return_value = process - self.engine._mp_ctx = process_context - self.engine._init_signal.clear() + def _use_created_engine(self) -> OnlineDataEngine: + """Replace the ready sampling fixture with a fresh lifecycle fixture.""" + self.engine.stop() + self.engine = _make_fake_engine(state=OnlineDataEngineState.CREATED) + return self.engine + + def test_worker_forwards_original_exception(self, monkeypatch) -> None: + """The subprocess entry point sends and re-raises the same exception.""" + error = ValueError("planner failed") + monkeypatch.setattr( + engine_module, + "_run_sim_worker", + MagicMock(side_effect=error), + ) + with pytest.raises(ValueError) as raised: + engine_module._sim_worker_fn( + MagicMock(), + MagicMock(), + MagicMock(), + MagicMock(), + MagicMock(), + MagicMock(), + self.engine._worker_error_buffer, + self.engine._worker_error_length, + self.engine._worker_failed_signal, + self.engine._state_value, + ) + + assert raised.value is error + forwarded = self.engine._receive_worker_error() + assert isinstance(forwarded, ValueError) + assert str(forwarded) == "planner failed" + + def test_start_transitions_through_starting_to_ready(self) -> None: + """A successful initial fill publishes the explicit lifecycle states.""" + engine = self._use_created_engine() + process = _install_fake_starting_worker(engine) + observed_states = [] + + def initialize() -> None: + observed_states.append(engine.state) + engine._init_signal.set() + + process.start.side_effect = initialize + + engine.start() + + assert observed_states == [OnlineDataEngineState.STARTING] + assert engine.state is OnlineDataEngineState.READY + assert engine.is_init + + def test_double_start_is_rejected(self) -> None: + """A ready engine cannot create a second producer process.""" + engine = self._use_created_engine() + _install_fake_starting_worker(engine) + engine.start() + + with pytest.raises(RuntimeError, match="requires state CREATED"): + engine.start() + + def test_sample_before_ready_is_rejected(self) -> None: + """Created engines cannot expose their zero-filled shared buffer.""" + engine = self._use_created_engine() + + with pytest.raises(RuntimeError, match="current state is CREATED"): + engine.sample_batch(batch_size=1, chunk_size=1) + + def test_start_reraises_original_worker_exception(self) -> None: + """Initial-fill failures retain their original exception type and text.""" + engine = self._use_created_engine() + process = _install_fake_starting_worker(engine, initialize=False) + process.start.side_effect = lambda: engine._record_worker_error( + ValueError("initial planner failure") + ) + + with pytest.raises(ValueError, match="initial planner failure"): + engine.start() + + assert engine.state is OnlineDataEngineState.FAILED - with pytest.raises(RuntimeError, match="worker exited unexpectedly"): - self.engine.start() + def test_start_times_out_when_live_worker_never_initializes(self) -> None: + """A hung initial fill cannot leave start blocked indefinitely.""" + engine = self._use_created_engine() + engine.cfg.initialization_timeout = 0.001 + _install_fake_starting_worker(engine, initialize=False) - process.join.assert_called_once_with(timeout=0) + with pytest.raises(TimeoutError, match="initial buffer fill exceeded"): + engine.start() - def test_sampling_raises_after_worker_failure(self) -> None: - """A failed refill worker cannot silently serve stale data forever.""" - process = MagicMock(exitcode=1) + assert engine.state is OnlineDataEngineState.FAILED + + def test_start_drains_worker_error_published_during_cleanup(self) -> None: + """A late recorder error annotates rather than disappears behind timeout.""" + engine = self._use_created_engine() + engine.cfg.initialization_timeout = 0.001 + _install_fake_starting_worker(engine, initialize=False) + shutdown_worker = engine._shutdown_worker + + def shutdown_and_publish_error() -> bool: + forced_shutdown = shutdown_worker() + assert engine_module._publish_worker_error( + engine._worker_error_buffer, + engine._worker_error_length, + engine._worker_failed_signal, + engine._state_value, + ValueError("late recorder flush failure"), + ) + return forced_shutdown + + engine._shutdown_worker = shutdown_and_publish_error + + with pytest.raises(TimeoutError) as raised: + engine.start() + + assert any( + "late recorder flush failure" in note for note in raised.value.__notes__ + ) + channel_error = engine._receive_worker_error() + assert isinstance(channel_error, ValueError) + assert str(channel_error) == "late recorder flush failure" + assert engine.state is OnlineDataEngineState.FAILED + assert engine._cleanup_complete + + @pytest.mark.parametrize("timeout", [float("nan"), float("inf")]) + def test_non_finite_initialization_timeout_is_rejected( + self, timeout: float + ) -> None: + """A non-finite timeout cannot turn startup into an infinite wait.""" + cfg = OnlineDataEngineCfg(initialization_timeout=timeout) + + with pytest.raises(ValueError, match="must be finite"): + OnlineDataEngine(cfg) + + def test_sampling_reraises_original_runtime_worker_exception(self) -> None: + """A refill failure is surfaced before any stale batch is returned.""" + self.engine._record_worker_error(ValueError("refill planner failure")) + + with pytest.raises(ValueError, match="refill planner failure"): + self.engine.sample_batch(batch_size=1, chunk_size=1) + + assert self.engine.state is OnlineDataEngineState.FAILED + + def test_sampling_detects_worker_exit_without_error_payload(self) -> None: + """An unexplained worker exit still fails the ready engine immediately.""" + process = MagicMock(exitcode=17) process.is_alive.return_value = False self.engine._sim_process = process - with pytest.raises(RuntimeError, match="worker exited unexpectedly"): + with pytest.raises(RuntimeError, match="exit code 17"): + self.engine.sample_batch(batch_size=1, chunk_size=1) + + assert self.engine.state is OnlineDataEngineState.FAILED + + def test_stopped_engine_rejects_sampling(self) -> None: + """Stopping permanently closes the sampling surface.""" + engine = self._use_created_engine() + engine.stop() + + with pytest.raises(RuntimeError, match="current state is STOPPED"): + engine.sample_batch(batch_size=1, chunk_size=1) + + def test_stopped_engine_rejects_restart(self) -> None: + """A stopped engine cannot reuse closed process primitives.""" + engine = self._use_created_engine() + engine.stop() + + with pytest.raises(RuntimeError, match="current state is STOPPED"): + engine.start() + + def test_failed_engine_rejects_restart(self) -> None: + """A failed engine is terminal until a new instance is constructed.""" + engine = self._use_created_engine() + engine._set_state(OnlineDataEngineState.FAILED) + engine._worker_error = ValueError("worker failed") + + with pytest.raises(RuntimeError, match="current state is FAILED"): + engine.start() + + def test_context_manager_starts_and_stops_engine(self) -> None: + """Managed use always closes the worker when the block exits.""" + engine = self._use_created_engine() + _install_fake_starting_worker(engine) + + with engine as active_engine: + assert active_engine is engine + assert active_engine.state is OnlineDataEngineState.READY + + assert engine.state is OnlineDataEngineState.STOPPED + + def test_monitor_target_does_not_retain_engine(self) -> None: + """The live monitor owns shared primitives, not the engine instance.""" + engine = self._use_created_engine() + _install_fake_starting_worker(engine) + + engine.start() + + assert engine._monitor_thread._target is engine_module._monitor_worker_process + assert all(arg is not engine for arg in engine._monitor_thread._args) + + def test_non_owner_cannot_stop_or_signal_owner_worker(self) -> None: + """A DataLoader copy cannot mutate its owner's producer lifecycle.""" + owner_pid = self.engine._owner_pid + self.engine._owner_pid = owner_pid + 1 + try: + with pytest.raises(RuntimeError, match="only be called by owner process"): + self.engine.stop() + self.engine.__del__() + + assert not self.engine._close_signal.is_set() + assert not self.engine._fill_signal.is_set() + finally: + self.engine._owner_pid = owner_pid + + @pytest.mark.parametrize("start_method", ["fork", "spawn"]) + def test_real_consumer_process_can_sample(self, start_method: str) -> None: + """Forked and spawned DataLoader-style consumers receive safe copies.""" + if start_method not in mp.get_all_start_methods(): + pytest.skip(f"multiprocessing start method {start_method!r} unavailable") + context = mp.get_context(start_method) + result_queue = context.Queue() + process = context.Process( + target=_sample_engine_in_subprocess, + args=(self.engine, result_queue), + ) + + process.start() + result = result_queue.get(timeout=10.0) + process.join(timeout=10.0) + + assert process.exitcode == 0 + assert result == ("ok", (1, 1)) + assert not self.engine._close_signal.is_set() + + def test_worker_error_is_broadcast_to_all_consumers(self) -> None: + """Every consumer reconstructs the same exception snapshot.""" + context = mp.get_context("fork") + assert engine_module._publish_worker_error( + self.engine._worker_error_buffer, + self.engine._worker_error_length, + self.engine._worker_failed_signal, + self.engine._state_value, + ValueError("broadcast planner failure"), + ) + result_queue = context.Queue() + processes = [ + context.Process( + target=_sample_engine_in_subprocess, + args=(self.engine, result_queue), + ) + for _ in range(2) + ] + + for process in processes: + process.start() + results = [result_queue.get(timeout=10.0) for _ in processes] + for process in processes: + process.join(timeout=10.0) + + assert results == [ + ("error", "ValueError", "broadcast planner failure"), + ("error", "ValueError", "broadcast planner failure"), + ] + assert all(process.exitcode == 0 for process in processes) + + def test_unpickleable_worker_error_uses_stable_fallback(self) -> None: + """Unpickleable exception arguments still produce a useful failure.""" + error = ValueError("unpickleable planner failure", lambda: None) + assert engine_module._publish_worker_error( + self.engine._worker_error_buffer, + self.engine._worker_error_length, + self.engine._worker_failed_signal, + self.engine._state_value, + error, + ) + + with pytest.raises(OnlineDataWorkerError, match="unpickleable planner failure"): self.engine.sample_batch(batch_size=1, chunk_size=1) + def test_stop_failure_does_not_publish_stopped(self) -> None: + """A worker that survives kill keeps the engine failed and retryable.""" + process = self.engine._sim_process + process.join.side_effect = None + process.is_alive.return_value = True + + with pytest.raises( + OnlineDataWorkerError, match="durability could not be confirmed" + ): + self.engine.stop() + + assert self.engine.state is OnlineDataEngineState.FAILED + assert not self.engine._cleanup_complete + process.is_alive.return_value = False + with pytest.raises( + OnlineDataWorkerError, match="durability could not be confirmed" + ): + self.engine.stop() + assert self.engine.state is OnlineDataEngineState.FAILED + assert self.engine._cleanup_complete + self.engine.stop() + + def test_stop_detects_worker_that_died_before_shutdown(self) -> None: + """An already-dead producer cannot be reclassified as a clean stop.""" + process = MagicMock(exitcode=23) + process.is_alive.return_value = False + self.engine._sim_process = process + + with pytest.raises(RuntimeError, match="before stop.*exit code 23"): + self.engine.stop() + + assert self.engine.state is OnlineDataEngineState.FAILED + assert self.engine._worker_failed_signal.is_set() + assert self.engine._cleanup_complete + + def test_forced_termination_fails_durability_contract(self) -> None: + """Needing terminate means graceful recorder durability is unconfirmed.""" + process = self.engine._sim_process + process.join.side_effect = None + process.is_alive.return_value = True + + def terminate() -> None: + process.is_alive.return_value = False + + process.terminate.side_effect = terminate + + with pytest.raises( + OnlineDataWorkerError, match="durability could not be confirmed" + ): + self.engine.stop() + + assert self.engine.state is OnlineDataEngineState.FAILED + assert self.engine._cleanup_complete + + def test_stop_surfaces_worker_failure_during_shutdown(self) -> None: + """A recorder failure during worker close cannot be published as stopped.""" + + def shutdown_and_fail() -> None: + self.engine._record_worker_error(ValueError("recorder flush failed")) + + self.engine._shutdown_worker = MagicMock(side_effect=shutdown_and_fail) + + with pytest.raises(ValueError, match="recorder flush failed"): + self.engine.stop() + + assert self.engine.state is OnlineDataEngineState.FAILED + assert self.engine._cleanup_complete + + def test_stop_cancels_start_without_state_reversal(self) -> None: + """Concurrent stop owns cleanup and leaves the terminal state stopped.""" + engine = self._use_created_engine() + _install_fake_starting_worker(engine, initialize=False) + start_errors = [] + + def start_engine() -> None: + try: + engine.start() + except BaseException as error: + start_errors.append(error) + + start_thread = threading.Thread(target=start_engine) + start_thread.start() + assert engine._fill_signal.wait(timeout=2.0) + + engine.stop() + start_thread.join(timeout=2.0) + + assert not start_thread.is_alive() + assert len(start_errors) == 1 + assert "cancelled by stop" in str(start_errors[0]) + assert engine.state is OnlineDataEngineState.STOPPED + + def test_context_exit_preserves_body_exception(self) -> None: + """Cleanup failures annotate rather than replace a with-body error.""" + body_error = KeyError("training failed") + self.engine.stop = MagicMock(side_effect=RuntimeError("cleanup failed")) + + result = self.engine.__exit__(KeyError, body_error, None) + + assert result is None + assert any("cleanup failed" in note for note in body_error.__notes__) + + def test_context_exit_surfaces_background_worker_failure(self) -> None: + """A worker failure after READY cannot be hidden by context cleanup.""" + engine = self._use_created_engine() + _install_fake_starting_worker(engine) + + with pytest.raises(ValueError, match="late planner failure"): + with engine: + engine._record_worker_error(ValueError("late planner failure")) + + assert engine.state is OnlineDataEngineState.FAILED + # ----------------------------------------------------------------------- def test_sample_batch_shape(self) -> None: @@ -309,7 +744,10 @@ def test_refill_not_double_triggered(self) -> None: ), "_fill_signal should remain set without reset" def teardown_method(self) -> None: - self.engine.stop() + try: + self.engine.stop() + except BaseException: + pass # =========================================================================== diff --git a/tests/gym/envs/managers/test_async_dataset_functors.py b/tests/gym/envs/managers/test_async_dataset_functors.py index 20b6d417b..77320edfe 100644 --- a/tests/gym/envs/managers/test_async_dataset_functors.py +++ b/tests/gym/envs/managers/test_async_dataset_functors.py @@ -237,7 +237,59 @@ def test_finalize_drains_and_finalizes_dataset(self): recorder(env, env_ids=torch.tensor([0])) env.current_rollout_step = 0 # mirror post-save reset - recorder.finalize() + first_path = recorder.finalize() + second_path = recorder.close() assert mock_ds.save_episode_calls == 1 assert mock_ds.finalize_calls == 1 + assert first_path == second_path == recorder.dataset_path + + def test_finalize_does_not_enqueue_uncommitted_live_rollout(self): + """A rollout not explicitly passed to the recorder is discarded on close.""" + env = _MockEnv(num_envs=1, steps=2) + mock_ds = _MockDataset() + recorder = _make_recorder(env, mock_ds) + + recorder.finalize() + + assert mock_ds.save_episode_calls == 0 + assert mock_ds.finalize_calls == 1 + + def test_finalize_aggregates_background_failures_after_draining(self): + """All queued commits run and their failures surface at the barrier.""" + env = _MockEnv(num_envs=2, steps=2) + mock_ds = _MockDataset() + recorder = _make_recorder(env, mock_ds) + + def fail_save(env_id, *args, **kwargs): + if env_id == 0: + return False + raise OSError("disk unavailable") + + recorder._save_single_episode = Mock(side_effect=fail_save) + recorder(env, env_ids=torch.tensor([0, 1])) + + with pytest.raises(RuntimeError) as error_info: + recorder.finalize() + + message = str(error_info.value) + assert "2 committed episode(s)" in message + assert "env 0" in message + assert "env 1" in message + assert "disk unavailable" in message + assert recorder._save_single_episode.call_count == 2 + assert mock_ds.finalize_calls == 1 + + with pytest.raises(RuntimeError, match="2 committed episode"): + recorder.close() + assert recorder._save_single_episode.call_count == 2 + assert mock_ds.finalize_calls == 1 + + def test_call_after_finalize_is_rejected(self): + """No commit may be queued behind the shutdown sentinel.""" + env = _MockEnv(num_envs=1, steps=1) + recorder = _make_recorder(env, _MockDataset()) + recorder.finalize() + + with pytest.raises(RuntimeError, match="already finalized"): + recorder(env, env_ids=torch.tensor([0])) diff --git a/tests/gym/envs/managers/test_dataset_functors.py b/tests/gym/envs/managers/test_dataset_functors.py index beaf51540..0355c502a 100644 --- a/tests/gym/envs/managers/test_dataset_functors.py +++ b/tests/gym/envs/managers/test_dataset_functors.py @@ -19,6 +19,7 @@ import json import threading +from pathlib import Path import numpy as np import pytest @@ -214,6 +215,38 @@ def test_initialization_with_videos(self, mock_lerobot_dataset): assert recorder.use_videos == True + @patch("embodichain.lab.gym.envs.managers.datasets.LeRobotDataset") + def test_finalize_is_idempotent_and_does_not_commit_live_rollout( + self, mock_lerobot_dataset, tmp_path + ): + """Closing finalizes storage but never turns a partial rollout into an episode.""" + env = MockEnvForDataset(num_envs=2) + env.current_rollout_step = 3 + mock_dataset_instance = Mock() + mock_dataset_instance.image_writer = None + mock_dataset_instance.meta = Mock() + mock_dataset_instance.meta.info = {"fps": 30} + mock_lerobot_dataset.create.return_value = mock_dataset_instance + cfg = MockFunctorCfg( + params={ + "save_path": str(tmp_path), + "robot_meta": {"robot_type": "test_robot", "control_freq": 30}, + "instruction": {"lang": "test task"}, + "extra": {"task_description": "test"}, + } + ) + recorder = LeRobotRecorder(cfg, env) + recorder._save_episodes = Mock() + + first_path = recorder.finalize() + second_path = recorder.close() + + assert first_path == second_path == recorder.dataset_path + recorder._save_episodes.assert_not_called() + mock_dataset_instance.finalize.assert_called_once_with() + with pytest.raises(RuntimeError, match="already finalized"): + recorder(env, torch.tensor([0])) + @pytest.mark.skipif(not LEROBOT_AVAILABLE, reason="LeRobot not installed") class TestLeRobotRecorderFeatures: @@ -763,6 +796,28 @@ 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_post_commit_metadata_failure_does_not_reuse_episode_index() -> None: + """A sidecar failure after LeRobot commit advances the global index first.""" + recorder = LeRobotRecorder.__new__(LeRobotRecorder) + 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._convert_frame_to_lerobot = MagicMock(return_value={}) + recorder._write_episode_metadata = MagicMock(side_effect=OSError("disk full")) + + with pytest.raises(OSError, match="disk full"): + recorder._save_single_episode(0, [object()], [object()]) + + recorder.dataset.save_episode.assert_called_once_with() + assert recorder.curr_episode == 1 + + class TestDatasetFunctorCfg: """Tests for dataset functor configuration.""" diff --git a/tests/gym/envs/managers/test_dataset_manager.py b/tests/gym/envs/managers/test_dataset_manager.py index e33793383..0f9ba0cb8 100644 --- a/tests/gym/envs/managers/test_dataset_manager.py +++ b/tests/gym/envs/managers/test_dataset_manager.py @@ -16,15 +16,22 @@ from __future__ import annotations +import threading +from collections import deque from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch +import numpy as np +import pytest import torch from embodichain.lab.gym.envs.embodied_env import EmbodiedEnv from embodichain.lab.gym.envs.managers.cfg import DatasetFunctorCfg from embodichain.lab.gym.envs.managers.dataset_manager import DatasetManager -from embodichain.lab.gym.envs.managers.record import record_camera_data +from embodichain.lab.gym.envs.managers.record import ( + record_camera_data, + record_camera_data_async, +) from embodichain.lab.gym.utils.profiler import EnvProfiler @@ -124,6 +131,63 @@ def recorder(env, env_ids, *, use_videos: bool) -> None: assert calls == [(env, env_ids, True)] +def make_manager_for_finalize(*named_functors) -> DatasetManager: + """Create a manager lifecycle fixture without constructing a simulator.""" + manager = DatasetManager.__new__(DatasetManager) + manager._mode_functor_names = {"save": [name for name, _ in named_functors]} + manager._mode_functor_cfgs = { + "save": [SimpleNamespace(func=functor) for _, functor in named_functors] + } + manager._finalize_lock = threading.Lock() + manager._finalized = False + manager._finalize_result = None + manager._finalize_error = None + return manager + + +def test_finalize_is_idempotent_after_success() -> None: + """Repeated manager finalization returns the cached path without re-closing.""" + first = SimpleNamespace(finalize=MagicMock(return_value="/dataset/first")) + second = SimpleNamespace(finalize=MagicMock(return_value="/dataset/second")) + manager = make_manager_for_finalize(("first", first), ("second", second)) + + assert manager.close() == "/dataset/first" + assert manager.finalize() == "/dataset/first" + first.finalize.assert_called_once_with() + second.finalize.assert_called_once_with() + + +def test_finalize_aggregates_errors_and_attempts_every_functor() -> None: + """One broken recorder cannot prevent the remaining recorders from closing.""" + first = SimpleNamespace( + finalize=MagicMock(side_effect=OSError("first disk failure")) + ) + successful = SimpleNamespace(finalize=MagicMock(return_value="/dataset/good")) + third = SimpleNamespace( + finalize=MagicMock(side_effect=RuntimeError("third writer failure")) + ) + manager = make_manager_for_finalize( + ("first", first), ("successful", successful), ("third", third) + ) + + with pytest.raises(RuntimeError) as error_info: + manager.finalize() + + message = str(error_info.value) + assert "2 dataset functor(s)" in message + assert "first: first disk failure" in message + assert "third: third writer failure" in message + first.finalize.assert_called_once_with() + successful.finalize.assert_called_once_with() + third.finalize.assert_called_once_with() + + with pytest.raises(RuntimeError, match="2 dataset functor"): + manager.finalize() + first.finalize.assert_called_once_with() + successful.finalize.assert_called_once_with() + third.finalize.assert_called_once_with() + + def test_initialize_episode_limits_successes_to_reset_envs() -> None: env, manager = make_env_for_episode_selection( save_failed_episodes=False, @@ -163,6 +227,23 @@ def test_discard_reset_does_not_auto_save_trajectory() -> None: assert env._traj_steps.tolist() == [0, 2, 1] +def test_committed_trajectory_write_failure_preserves_buffer_and_raises() -> None: + """A failed trajectory commit is visible and its live cursor is not cleared.""" + env, _ = make_env_for_episode_selection( + save_failed_episodes=False, + successful_env_ids=[], + ) + env._traj_buffer = object() + env._traj_steps = torch.tensor([3, 2, 1]) + env.cfg.trajectory_auto_save = True + env._save_trajectory_for_env = MagicMock(side_effect=OSError("disk full")) + + with pytest.raises(OSError, match="disk full"): + EmbodiedEnv._initialize_episode(env, env_ids=[0], save_data=True) + + assert env._traj_steps.tolist() == [3, 2, 1] + + def test_discard_reset_clears_camera_frames_without_saving() -> None: """save_data=False drops video frames through the recorder abort hook.""" env, _ = make_env_for_episode_selection( @@ -182,3 +263,142 @@ def test_discard_reset_clears_camera_frames_without_saving() -> None: recorder.save_and_clear.assert_not_called() assert recorder._frames == [] + + +def test_camera_finalize_discards_uncommitted_frames_idempotently() -> None: + """Closing a camera recorder cannot resurrect an explicitly discarded episode.""" + recorder = record_camera_data.__new__(record_camera_data) + recorder._name = "test" + recorder._save_path = "/tmp/test-camera-recorder" + recorder._current_episode = 0 + recorder._frames = [object()] + recorder._finalize_lock = threading.Lock() + recorder._finalized = False + + recorder.discard_and_clear() + with patch( + "embodichain.lab.gym.envs.managers.record.images_to_video" + ) as save_video: + recorder.finalize() + recorder.close() + + save_video.assert_not_called() + assert recorder._frames == [] + assert recorder._current_episode == 0 + with pytest.raises(RuntimeError, match="already finalized"): + recorder.save_and_clear() + + +def test_camera_explicit_commit_is_not_repeated_by_finalize() -> None: + """Only save_and_clear commits frames; finalize merely closes the empty buffer.""" + recorder = record_camera_data.__new__(record_camera_data) + frame = object() + recorder._name = "test" + recorder._save_path = "/tmp/test-camera-recorder" + recorder._current_episode = 0 + recorder._frames = [frame] + recorder._finalize_lock = threading.Lock() + recorder._finalized = False + + with patch( + "embodichain.lab.gym.envs.managers.record.images_to_video" + ) as save_video: + recorder.save_and_clear() + recorder.finalize() + recorder.close() + + save_video.assert_called_once_with( + [frame], + "/tmp/test-camera-recorder", + "episode_0_test", + fps=20, + ) + assert recorder._frames == [] + assert recorder._current_episode == 1 + + +def _make_async_camera_recorder(num_envs: int = 2) -> record_camera_data_async: + """Build an async camera recorder without constructing simulation sensors.""" + recorder = record_camera_data_async.__new__(record_camera_data_async) + recorder._name = "test" + recorder._save_path = "/tmp/test-async-camera-recorder" + recorder._current_episode = 0 + recorder._frames = [] + recorder._finalize_lock = threading.Lock() + recorder._finalized = False + recorder._num_envs = num_envs + recorder._frames_list = [[] for _ in range(num_envs)] + recorder._ep_idx = [0 for _ in range(num_envs)] + recorder._committed_env_episodes = [deque() for _ in range(num_envs)] + recorder._async_camera_finalize_lock = threading.Lock() + recorder._async_camera_finalized = False + recorder._async_camera_finalize_error = None + return recorder + + +def test_async_camera_commit_persists_without_a_later_step() -> None: + """The final committed vector episode is durable before close/finalize.""" + recorder = _make_async_camera_recorder() + recorder._frames_list = [ + [np.zeros((2, 2, 4), dtype=np.uint8)], + [np.ones((2, 2, 4), dtype=np.uint8)], + ] + + with patch( + "embodichain.lab.gym.envs.managers.record.images_to_video" + ) as save_video: + recorder.save_and_clear(env_ids=torch.tensor([0, 1])) + recorder.finalize() + + save_video.assert_called_once() + _, save_path, video_name = save_video.call_args.args + assert save_path == "/tmp/test-async-camera-recorder" + assert video_name == "ep0_test_allenvs" + assert save_video.call_args.kwargs == {"fps": 20} + assert recorder._current_episode == 1 + assert all(not queue for queue in recorder._committed_env_episodes) + + +def test_async_camera_finalize_rejects_incomplete_committed_batch() -> None: + """A partial vector commit is reported instead of silently discarded.""" + recorder = _make_async_camera_recorder() + recorder._frames_list[0] = [np.zeros((2, 2, 4), dtype=np.uint8)] + recorder.save_and_clear(env_ids=torch.tensor([0])) + recorder.discard_and_clear(env_ids=torch.tensor([1])) + + with pytest.raises(RuntimeError, match="pending episodes per env: \\[1, 0\\]"): + recorder.finalize() + + with pytest.raises(RuntimeError, match="incomplete committed environment batch"): + recorder.close() + + +def test_pending_discard_attempts_every_camera_and_clears_trajectory() -> None: + """One camera failure cannot skip later durability barriers or live cleanup.""" + first = record_camera_data.__new__(record_camera_data) + first.discard_and_clear = MagicMock(side_effect=OSError("camera one failed")) + first.finalize = MagicMock(side_effect=OSError("camera one flush failed")) + second = record_camera_data.__new__(record_camera_data) + second.discard_and_clear = MagicMock() + second.finalize = MagicMock() + env = SimpleNamespace( + cfg=SimpleNamespace(events=True), + event_manager=SimpleNamespace( + _mode_functor_cfgs={ + "interval": [ + SimpleNamespace(func=first), + SimpleNamespace(func=second), + ] + } + ), + _traj_steps=torch.tensor([3, 2]), + rollout_buffer=None, + ) + + with pytest.raises(RuntimeError, match="camera one failed"): + EmbodiedEnv._discard_pending_recordings(env) + + first.finalize.assert_called_once_with() + second.discard_and_clear.assert_called_once_with() + second.finalize.assert_called_once_with() + assert env._traj_steps.tolist() == [0, 0] diff --git a/tests/gym/envs/test_demo.py b/tests/gym/envs/test_demo.py index ec656f45e..036899596 100644 --- a/tests/gym/envs/test_demo.py +++ b/tests/gym/envs/test_demo.py @@ -18,8 +18,10 @@ from __future__ import annotations +import threading from typing import Any +import pytest import torch from tensordict import TensorDict @@ -123,6 +125,136 @@ def test_execute_demo_episode_stops_immediately_on_success_termination() -> None assert result.terminal_reason == "success" +class _StaggeredVectorEnv: + """Two-row stub whose environments reach success on different steps.""" + + def __init__(self) -> None: + self.num_envs = 2 + self.step_count = 0 + self.actions: list[int] = [] + self.masked_actions: list[tuple[int, tuple[bool, ...]]] = [] + self.requested_second_segment = False + + def create_demo_segments(self): + yield DemoSegment(actions=(1, 2, 3, 4), name="shared") + self.requested_second_segment = True + yield DemoSegment(actions=(5,), name="must_not_run") + + def _mask_demo_action(self, action: int, active_mask: tuple[bool, ...]) -> int: + self.masked_actions.append((action, active_mask)) + return action + + def step(self, action: int): + self.actions.append(action) + self.step_count += 1 + terminated = torch.tensor( + [self.step_count == 1, self.step_count == 3], dtype=torch.bool + ) + return ( + None, + torch.zeros(self.num_envs), + terminated, + torch.zeros(self.num_envs, dtype=torch.bool), + {"success": terminated.clone()}, + ) + + def is_task_success(self) -> torch.Tensor: + return torch.ones(self.num_envs, dtype=torch.bool) + + +def test_execute_demo_episode_supports_staggered_vector_success() -> None: + """A completed row freezes while unfinished rows continue their shared plan.""" + env = _StaggeredVectorEnv() + + result = execute_demo_episode(env) + + assert env.actions == [1, 2, 3] + assert env.masked_actions == [(2, (False, True)), (3, (False, True))] + assert not env.requested_second_segment + assert result.length == 3 + assert result.lengths == (1, 3) + assert result.completed_by_env == (True, True) + assert result.terminal_reasons == ("success", "success") + assert result.terminated == (True, True) + assert result.all_success + assert result.segments[0].end_steps == (1, 3) + assert result.segments[0].successes == (True, True) + + +class _VectorFailureEnv(_StaggeredVectorEnv): + def create_demo_segments(self): + return (DemoSegment(actions=(1, 2, 3), name="shared"),) + + def step(self, action: int): + self.actions.append(action) + self.step_count += 1 + terminated = torch.tensor([self.step_count == 2, False], dtype=torch.bool) + return ( + None, + torch.zeros(self.num_envs), + terminated, + torch.zeros(self.num_envs, dtype=torch.bool), + {"success": torch.zeros(self.num_envs, dtype=torch.bool)}, + ) + + +def test_vector_failure_aborts_peer_and_preserves_per_env_reason() -> None: + """Batch-atomic failure distinguishes the failing row from its peer.""" + env = _VectorFailureEnv() + + result = execute_demo_episode(env) + + assert env.actions == [1, 2] + assert not result.completed + assert result.terminal_reason == "failure" + assert result.terminal_reasons == ("failure", "batch_aborted") + assert result.terminated == (True, False) + assert result.success == (False, False) + assert result.lengths == (2, 2) + + +class _ValidatedSegmentEnv(_SegmentedEnv): + def __init__(self, validation: bool) -> None: + super().__init__() + self.validation = validation + self.requested_second_segment = False + + def create_demo_segments(self): + yield DemoSegment( + actions=(1,), + name="pick", + validator=lambda: torch.tensor([self.validation]), + ) + self.requested_second_segment = True + yield DemoSegment(actions=(3,), name="place") + + +def test_segment_validator_stops_invalid_lazy_plan() -> None: + """Subtask validation is independent from episode-level Gym termination.""" + env = _ValidatedSegmentEnv(validation=False) + + result = execute_demo_episode(env) + + assert env.actions == [1] + assert not env.requested_second_segment + assert not result.completed + assert result.terminal_reason == "segment_validation_failed" + assert result.terminal_reasons == ("segment_validation_failed",) + assert result.segments[0].failure_reason == "segment_validation_failed" + + +def test_segment_validator_allows_next_lazy_segment() -> None: + """A validated subtask advances the lazy planner without abusing terminated.""" + env = _ValidatedSegmentEnv(validation=True) + + result = execute_demo_episode(env) + + assert env.actions == [1, 3] + assert env.requested_second_segment + assert result.all_success + assert [segment.name for segment in result.segments] == ["pick", "place"] + + class _TruncatingEnv(_TerminatingEnv): def step(self, action: int): self.actions.append(action) @@ -147,6 +279,147 @@ def test_execute_demo_episode_never_accepts_truncated_rollout() -> None: assert result.terminal_reason == "truncated" +class _ConflictingTerminalEnv(_TerminatingEnv): + def step(self, action: int): + self.actions.append(action) + return ( + None, + torch.zeros(1), + torch.ones(1, dtype=torch.bool), + torch.ones(1, dtype=torch.bool), + { + "success": torch.ones(1, dtype=torch.bool), + "fail": torch.zeros(1, dtype=torch.bool), + }, + ) + + +def test_terminated_and_truncated_flags_are_both_sticky() -> None: + """A conflicting Gym transition keeps both flags while truncation wins.""" + result = execute_demo_episode(_ConflictingTerminalEnv()) + + assert result.terminated == (True,) + assert result.truncated == (True,) + assert result.success == (False,) + assert result.terminal_reason == "truncated" + + +class _SuccessAndFailureEnv(_TerminatingEnv): + def step(self, action: int): + self.actions.append(action) + return ( + None, + torch.zeros(1), + torch.ones(1, dtype=torch.bool), + torch.zeros(1, dtype=torch.bool), + { + "success": torch.ones(1, dtype=torch.bool), + "fail": torch.ones(1, dtype=torch.bool), + }, + ) + + +def test_explicit_failure_wins_over_conflicting_success() -> None: + """The task failure signal cannot be committed as a successful episode.""" + result = execute_demo_episode(_SuccessAndFailureEnv()) + + assert not result.completed + assert result.success == (False,) + assert result.terminal_reason == "failure" + + +class _InactiveStaleTruncationEnv(_StaggeredVectorEnv): + def create_demo_segments(self): + return (DemoSegment(actions=(1, 2), name="shared"),) + + def step(self, action: int): + self.actions.append(action) + self.step_count += 1 + if self.step_count == 1: + terminated = torch.tensor([True, False]) + truncated = torch.tensor([False, False]) + success = torch.tensor([True, False]) + else: + terminated = torch.tensor([False, True]) + truncated = torch.tensor([True, False]) + success = torch.tensor([True, False]) + return None, torch.zeros(2), terminated, truncated, {"success": success} + + +def test_inactive_stale_truncation_does_not_change_active_failure_reason() -> None: + """Signals from a frozen row cannot relabel a peer's current failure.""" + result = execute_demo_episode(_InactiveStaleTruncationEnv()) + + assert result.terminal_reason == "failure" + assert result.terminal_reasons == ("success", "failure") + assert result.truncated == (False, False) + + +class _VectorValidatorEnv(_StaggeredVectorEnv): + def create_demo_segments(self): + return ( + DemoSegment( + actions=(1,), + name="validated", + validator=lambda: torch.tensor([True, False]), + ), + ) + + def step(self, action: int): + self.actions.append(action) + return ( + None, + torch.zeros(2), + torch.zeros(2, dtype=torch.bool), + torch.zeros(2, dtype=torch.bool), + {"success": torch.zeros(2, dtype=torch.bool)}, + ) + + +def test_validator_batch_abort_has_consistent_peer_status() -> None: + """A batch-aborted peer is not simultaneously marked segment-successful.""" + result = execute_demo_episode(_VectorValidatorEnv()) + + assert result.segments[0].successes == (False, False) + assert result.segments[0].failure_reasons == ( + "batch_aborted", + "segment_validation_failed", + ) + assert result.segments[0].failure_reason == "segment_validation_failed" + + +class _CancellationEnv(_ValidatedSegmentEnv): + def __init__(self) -> None: + super().__init__(validation=True) + self.validator_called = False + + def create_demo_segments(self): + yield DemoSegment( + actions=(1,), + name="first", + validator=self._validate, + ) + self.requested_second_segment = True + yield DemoSegment(actions=(3,), name="second") + + def _validate(self) -> torch.Tensor: + self.validator_called = True + return torch.ones(1, dtype=torch.bool) + + +def test_cancellation_after_last_action_does_not_advance_lazy_plan() -> None: + """Cancellation is observed before validation or requesting another segment.""" + env = _CancellationEnv() + + result = execute_demo_episode(env, should_stop=lambda: bool(env.actions)) + + assert env.actions == [1] + assert not env.validator_called + assert not env.requested_second_segment + assert result.terminal_reason == "interrupted" + assert result.terminal_reasons == ("interrupted",) + + class _LegacyEnv(_SegmentedEnv): create_demo_segments = EmbodiedEnv.create_demo_segments @@ -249,6 +522,191 @@ def test_expert_rollout_writer_uses_independent_per_env_lengths() -> None: assert env.current_rollout_step == 3 +def test_expert_rollout_writer_freezes_inactive_demo_row() -> None: + """Sticky terminal rows do not receive frames from later shared actions.""" + env = _RolloutWriterStub() + env._demo_active_mask = torch.tensor([False, True]) + obs = TensorDict({"state": torch.tensor([[1.0, 1.0], [2.0, 2.0]])}, batch_size=[2]) + + EmbodiedEnv._write_episode_rollout_step( + env, + obs=obs, + action=torch.tensor([[3.0, 3.0], [4.0, 4.0]]), + rewards=torch.tensor([0.5, 1.0]), + terminateds=torch.zeros(2, dtype=torch.bool), + truncateds=torch.zeros(2, dtype=torch.bool), + ) + + assert env.rollout_steps.tolist() == [0, 3] + assert not env.rollout_buffer["valid"][0].any() + assert env.rollout_buffer["valid"][1, 2] + + +class _RobotQposStub: + def get_qpos(self) -> torch.Tensor: + return torch.tensor([[0.1, 0.2], [0.3, 0.4]]) + + +class _ActionMaskStub: + num_envs = 2 + active_joint_ids = [0, 1] + robot = _RobotQposStub() + _demo_active_mask = torch.tensor([False, True]) + + +def test_processed_qpos_action_holds_inactive_row() -> None: + """Inactive qpos-controlled rows hold their measured joint position.""" + env = _ActionMaskStub() + + masked = EmbodiedEnv._mask_processed_demo_action( + env, torch.tensor([[9.0, 9.0], [8.0, 8.0]]) + ) + + assert torch.allclose(masked[0], torch.tensor([0.1, 0.2])) + assert torch.allclose(masked[1], torch.tensor([8.0, 8.0])) + + +def test_processed_qpos_action_supports_full_robot_layout() -> None: + """A full-DOF IK command holds measured full qpos for inactive rows.""" + env = _ActionMaskStub() + env.active_joint_ids = [0] + + masked = EmbodiedEnv._mask_processed_demo_action( + env, torch.tensor([[9.0, 9.0], [8.0, 8.0]]) + ) + + assert torch.allclose(masked[0], torch.tensor([0.1, 0.2])) + assert torch.allclose(masked[1], torch.tensor([8.0, 8.0])) + + +class _FullDofRobotStub: + def __init__(self) -> None: + self.command = None + self.joint_ids = None + + def get_qpos(self) -> torch.Tensor: + return torch.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]) + + def set_qpos(self, qpos: torch.Tensor, joint_ids: list[int]) -> None: + self.command = qpos + self.joint_ids = joint_ids + + +def test_full_dof_hold_is_sliced_before_active_joint_step() -> None: + """Masking and application agree when IK returns full robot qpos.""" + env = type( + "FullDofActionStub", + (), + { + "num_envs": 2, + "device": torch.device("cpu"), + "active_joint_ids": [0, 2], + "_demo_active_mask": torch.tensor([False, True]), + "robot": _FullDofRobotStub(), + }, + )() + action = TensorDict( + {"qpos": torch.tensor([[9.0, 9.0, 9.0], [8.0, 8.0, 8.0]])}, + batch_size=[2], + ) + + masked = EmbodiedEnv._mask_processed_demo_action(env, action) + applied = EmbodiedEnv._step_action(env, masked) + + assert torch.allclose( + env.robot.command, + torch.tensor([[0.1, 0.3], [8.0, 8.0]]), + ) + assert env.robot.joint_ids == [0, 2] + assert applied["qpos"].shape == (2, 2) + + +def test_processed_velocity_action_zeros_inactive_row() -> None: + """Inactive velocity-controlled rows receive a no-op command.""" + env = _ActionMaskStub() + action = TensorDict( + {"qvel": torch.tensor([[9.0, 9.0], [8.0, 8.0]])}, batch_size=[2] + ) + + masked = EmbodiedEnv._mask_processed_demo_action(env, action) + + assert torch.equal(masked["qvel"][0], torch.zeros(2)) + assert torch.equal(masked["qvel"][1], torch.tensor([8.0, 8.0])) + + +class _TrajectoryRobotStub: + def get_local_pose(self) -> torch.Tensor: + return torch.ones(2, 7) + + def get_qpos(self) -> torch.Tensor: + return torch.tensor([[0.1, 0.2], [0.3, 0.4]]) + + +class _TrajectoryWriterStub: + num_envs = 2 + device = torch.device("cpu") + + def __init__(self) -> None: + self._traj_buffer = TensorDict( + { + "states": { + "robot": { + "root_pose": torch.zeros(2, 4, 7), + "qpos": torch.zeros(2, 4, 2), + } + }, + "actions": torch.zeros(2, 4, 2), + }, + batch_size=[2, 4], + ) + self._traj_steps = torch.tensor([1, 1]) + self._traj_raw_action = torch.tensor([[9.0, 9.0], [8.0, 8.0]]) + self._demo_active_mask = torch.tensor([False, True]) + self.robot = _TrajectoryRobotStub() + self.sim = type( + "SimulationStub", + (), + {"_articulations": {}, "_rigid_objects": {}}, + )() + + +def test_trajectory_writer_freezes_inactive_demo_row() -> None: + """Trajectory cursors obey the same sticky activity mask as rollout rows.""" + env = _TrajectoryWriterStub() + + EmbodiedEnv._write_trajectory_step(env) + + assert env._traj_steps.tolist() == [1, 2] + assert torch.equal( + env._traj_buffer["states"]["robot"]["qpos"][0, 1], torch.zeros(2) + ) + assert torch.equal( + env._traj_buffer["states"]["robot"]["qpos"][1, 1], + torch.tensor([0.3, 0.4]), + ) + + +def test_success_status_freezes_after_staggered_demo_completion() -> None: + """Later stale done signals cannot clear an already successful frozen row.""" + env = type( + "SuccessStatusStub", + (), + { + "episode_success_status": torch.tensor([True, False]), + "_demo_no_auto_reset": True, + "_demo_active_mask": torch.tensor([False, True]), + }, + )() + + EmbodiedEnv._update_episode_success_status( + env, + {"success": torch.tensor([False, True])}, + torch.tensor([True, True]), + ) + + assert env.episode_success_status.tolist() == [True, True] + + def test_clear_expert_rows_preserves_unrelated_environment() -> None: """Clearing a completed row is an actual in-place selective mutation.""" env = _RolloutWriterStub() @@ -294,3 +752,61 @@ def test_legacy_metadata_reports_consistent_episode_success() -> None: assert metadata["success"] assert metadata["terminal_reason"] == "success" assert metadata["segments"][0]["success"] + + +class _CloseDependencyStub: + def __init__(self, *, error: Exception | None = None) -> None: + self.calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] + self.error = error + + def __call__(self, *args: Any, **kwargs: Any) -> None: + self.calls.append((args, kwargs)) + if self.error is not None: + raise self.error + + +class _CloseEnvStub: + def __init__(self, *, finalize_error: Exception | None = None) -> None: + self._closed = False + self._close_error = None + self._close_lock = threading.RLock() + self.discard = _CloseDependencyStub() + self.finalize = _CloseDependencyStub(error=finalize_error) + self.report = _CloseDependencyStub() + self.destroy = _CloseDependencyStub() + self.dataset_manager = type( + "DatasetManagerStub", (), {"finalize": self.finalize} + )() + self._profiler = type("ProfilerStub", (), {"report": self.report})() + self.sim = type("SimulationStub", (), {"destroy": self.destroy})() + + def _discard_pending_recordings(self) -> None: + self.discard() + + +def test_close_is_idempotent_and_never_commits_pending_data() -> None: + """Only the first close aborts pending data and finalizes committed writes.""" + env = _CloseEnvStub() + + EmbodiedEnv.close(env) + EmbodiedEnv.close(env) + + assert len(env.discard.calls) == 1 + assert len(env.finalize.calls) == 1 + assert len(env.report.calls) == 1 + assert env.destroy.calls == [((), {})] + + +def test_close_propagates_durability_failure_before_process_exit() -> None: + """Recorder failures are observable and simulator cleanup cannot hide them.""" + env = _CloseEnvStub(finalize_error=OSError("disk full")) + + with pytest.raises(RuntimeError, match="disk full"): + EmbodiedEnv.close(env) + + with pytest.raises(RuntimeError, match="disk full"): + EmbodiedEnv.close(env) + + assert env.destroy.calls == [((), {"exit_process": False})] + assert len(env.discard.calls) == 1 + assert len(env.finalize.calls) == 1 diff --git a/tests/gym/envs/test_replay.py b/tests/gym/envs/test_replay.py index be9441f92..80a64ef86 100644 --- a/tests/gym/envs/test_replay.py +++ b/tests/gym/envs/test_replay.py @@ -422,7 +422,7 @@ def test_auto_save_at_episode_end(tmp_path): assert data["meta"]["lengths"] == [4] -def test_auto_save_on_close(tmp_path): +def test_close_discards_uncommitted_trajectory(tmp_path): save_dir = tmp_path / "trajs" env = ReplayTestEnv(record_trajectory=True, num_envs=2, device="cpu") env.cfg.trajectory_save_dir = str(save_dir) @@ -436,7 +436,7 @@ def test_auto_save_on_close(tmp_path): gc.collect() files = list(save_dir.glob("*.pt")) - assert len(files) == 2 # one per in-flight env + assert files == [] def test_async_envs_do_not_corrupt_recording(): diff --git a/tests/lab/scripts/test_run_env.py b/tests/lab/scripts/test_run_env.py index 6084e16e1..2c3072c25 100644 --- a/tests/lab/scripts/test_run_env.py +++ b/tests/lab/scripts/test_run_env.py @@ -16,6 +16,7 @@ from __future__ import annotations +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -114,6 +115,37 @@ def test_run_env_preserves_configured_viser_image_fps() -> None: assert merged["visualization"]["sensor_image_fps"] == configured_fps +def test_replay_restores_wrapper_state_without_closing_caller_env(monkeypatch) -> None: + """Replay leaves the environment close to its CLI owner.""" + env = MagicMock() + env.unwrapped = env + replay_env = MagicMock() + monkeypatch.setattr( + run_env, + "load_trajectory", + lambda path: {"meta": {"num_envs": 1, "num_steps": 1}}, + ) + monkeypatch.setattr(run_env, "ReplayWrapper", lambda *args, **kwargs: replay_env) + monkeypatch.setattr(run_env, "replay_auto", lambda *args, **kwargs: None) + + run_env.replay(env, "trajectory.pt") + + replay_env.close.assert_not_called() + replay_env.env.sim.enable_physics.assert_called_once_with(True) + assert replay_env.env._replay_no_auto_reset is False + + +def test_preview_quit_returns_without_zero_exit(monkeypatch) -> None: + """Preview quit lets CLI cleanup failures determine the process status.""" + env = MagicMock() + env.reset.return_value = (None, {}) + monkeypatch.setattr("builtins.input", lambda: "q") + + run_env.preview(env) + + env.reset.assert_called_once_with() + + class _ResetTrackingEnv: def __init__(self) -> None: self.reset_options = [] @@ -123,6 +155,19 @@ def reset(self, options=None): return None, {} +class _LifecycleTrackingEnv(_ResetTrackingEnv): + def __init__(self) -> None: + super().__init__() + self.events = [] + + def reset(self, options=None): + self.events.append(("reset", options)) + return super().reset(options=options) + + def close(self, *, exit_process=None) -> None: + self.events.append(("close", exit_process)) + + def _episode_result(*, success: bool, reason: str) -> DemoEpisodeResult: return DemoEpisodeResult( episode_index=0, @@ -178,3 +223,169 @@ def test_generate_function_discards_partial_episode_on_exception(monkeypatch) -> generate_function(env, reset_before=False) assert env.reset_options == [{"save_data": False}] + + +@pytest.mark.parametrize( + "exit_type", + [KeyboardInterrupt, SystemExit, GeneratorExit], + ids=["keyboard-interrupt", "system-exit", "generator-exit"], +) +def test_generate_function_discards_partial_episode_on_early_exit( + monkeypatch, exit_type +) -> None: + """Non-Exception exits abort the pending episode before propagating.""" + env = _ResetTrackingEnv() + + def exit_during_episode(*args, **kwargs): + raise exit_type() + + monkeypatch.setattr( + "embodichain.lab.scripts.run_env.execute_demo_episode", + exit_during_episode, + ) + + with pytest.raises(exit_type): + generate_function(env, reset_before=False) + + assert env.reset_options == [{"save_data": False}] + + +def test_cli_aborts_before_closing_environment_once(monkeypatch) -> None: + """CLI owns finalization and closes only after discarding pending data.""" + env = _LifecycleTrackingEnv() + args = SimpleNamespace( + replay=False, + replay_mode="kinematic", + preview=False, + save_path="", + save_video=False, + debug_mode=False, + regenerate=False, + record_trajectory=False, + ) + parser = MagicMock() + parser.parse_args.return_value = args + gym_config = { + "id": GYM_ID, + "max_episodes": 1, + "demo_max_attempts": 1, + } + + monkeypatch.setattr(run_env, "_create_parser", lambda: parser) + monkeypatch.setattr(run_env, "discover_task_packages", lambda: None) + monkeypatch.setattr(run_env, "execute_init_hooks", lambda: None) + monkeypatch.setattr( + run_env, + "build_env_cfg_from_args", + lambda parsed_args: (object(), gym_config, {}), + ) + monkeypatch.setattr(run_env.gymnasium, "make", lambda **kwargs: env) + monkeypatch.setattr(run_env, "generate_function", lambda *args, **kwargs: True) + monkeypatch.setattr( + "embodichain.lab.sim.sim_manager.SimulationManager.flush_cleanup_queue", + lambda: None, + ) + + run_env.cli([]) + + abort_event = ("reset", {"save_data": False}) + assert env.events == [abort_event, abort_event, ("close", None)] + + +def test_close_durability_failure_is_not_swallowed() -> None: + """A failed recorder barrier makes the runner fail after aborting pending data.""" + env = _LifecycleTrackingEnv() + + def fail_close(*, exit_process=None) -> None: + env.events.append(("close", exit_process)) + raise OSError("dataset flush failed") + + env.close = fail_close + + with pytest.raises(OSError, match="dataset flush failed"): + run_env._abort_and_close_env(env) + + assert env.events == [ + ("reset", {"save_data": False}), + ("close", None), + ] + + +def test_cli_preserves_main_error_and_disables_fast_exit(monkeypatch) -> None: + """Failure cleanup returns normally so the original CLI error reaches the shell.""" + env = _LifecycleTrackingEnv() + args = SimpleNamespace( + replay=False, + replay_mode="kinematic", + preview=False, + save_path="", + save_video=False, + debug_mode=False, + regenerate=False, + record_trajectory=False, + ) + parser = MagicMock() + parser.parse_args.return_value = args + gym_config = {"id": GYM_ID} + + monkeypatch.setattr(run_env, "_create_parser", lambda: parser) + monkeypatch.setattr(run_env, "discover_task_packages", lambda: None) + monkeypatch.setattr(run_env, "execute_init_hooks", lambda: None) + monkeypatch.setattr( + run_env, + "build_env_cfg_from_args", + lambda parsed_args: (object(), gym_config, {}), + ) + monkeypatch.setattr(run_env.gymnasium, "make", lambda **kwargs: env) + monkeypatch.setattr( + run_env, + "main", + MagicMock(side_effect=RuntimeError("generation failed")), + ) + monkeypatch.setattr( + "embodichain.lab.sim.sim_manager.SimulationManager.flush_cleanup_queue", + lambda: None, + ) + + with pytest.raises(RuntimeError, match="generation failed"): + run_env.cli([]) + + assert env.events == [ + ("reset", {"save_data": False}), + ("close", False), + ] + + +def test_zero_system_exit_cannot_hide_cleanup_failure(monkeypatch) -> None: + """A nominal exit code is replaced by a failed durability barrier.""" + env = _LifecycleTrackingEnv() + args = SimpleNamespace( + replay=False, + replay_mode="kinematic", + preview=False, + replay_trajectory=None, + ) + parser = MagicMock() + parser.parse_args.return_value = args + monkeypatch.setattr(run_env, "_create_parser", lambda: parser) + monkeypatch.setattr(run_env, "discover_task_packages", lambda: None) + monkeypatch.setattr(run_env, "execute_init_hooks", lambda: None) + monkeypatch.setattr( + run_env, + "build_env_cfg_from_args", + lambda parsed_args: (object(), {"id": GYM_ID}, {}), + ) + monkeypatch.setattr(run_env.gymnasium, "make", lambda **kwargs: env) + monkeypatch.setattr(run_env, "main", MagicMock(side_effect=SystemExit(0))) + monkeypatch.setattr( + "embodichain.lab.sim.sim_manager.SimulationManager.flush_cleanup_queue", + lambda: None, + ) + + def fail_close(*, exit_process=None) -> None: + raise OSError("dataset flush failed") + + env.close = fail_close + + with pytest.raises(OSError, match="dataset flush failed"): + run_env.cli([]) From cd0e8d10e7a2f9ac0feb88a85b15a7342508eb4b Mon Sep 17 00:00:00 2001 From: yuecideng Date: Thu, 6 Aug 2026 15:45:23 +0000 Subject: [PATCH 3/3] test(data): tolerate spawned consumer startup latency --- tests/data_pipeline/test_online_data.py | 35 ++++++++++++++++++++++--- 1 file changed, 32 insertions(+), 3 deletions(-) diff --git a/tests/data_pipeline/test_online_data.py b/tests/data_pipeline/test_online_data.py index 70ef32d53..1a99b724e 100644 --- a/tests/data_pipeline/test_online_data.py +++ b/tests/data_pipeline/test_online_data.py @@ -32,6 +32,7 @@ import os import threading import unittest +from queue import Empty from unittest.mock import MagicMock import pytest @@ -65,6 +66,9 @@ STATE_DIM = 6 OBS_DIM = 10 ACTION_DIM = 4 +CONSUMER_PROCESS_TIMEOUT = 60.0 +RESULT_QUEUE_TIMEOUT = 5.0 +PROCESS_CLEANUP_TIMEOUT = 5.0 # --------------------------------------------------------------------------- @@ -436,10 +440,35 @@ def test_real_consumer_process_can_sample(self, start_method: str) -> None: ) process.start() - result = result_queue.get(timeout=10.0) - process.join(timeout=10.0) + try: + # A spawned consumer imports torch and the EmbodiChain package from + # scratch. Under xdist load this can exceed the old 10-second + # queue timeout even though the consumer is making progress. + process.join(timeout=CONSUMER_PROCESS_TIMEOUT) + if process.is_alive(): + pytest.fail( + "consumer process did not exit within " + f"{CONSUMER_PROCESS_TIMEOUT} seconds" + ) + if process.exitcode != 0: + pytest.fail( + "consumer process exited without sampling successfully " + f"(exit code {process.exitcode})" + ) + try: + result = result_queue.get(timeout=RESULT_QUEUE_TIMEOUT) + except Empty: + pytest.fail("consumer process exited without publishing a result") + finally: + if process.is_alive(): + process.terminate() + process.join(timeout=PROCESS_CLEANUP_TIMEOUT) + if process.is_alive(): + process.kill() + process.join(timeout=PROCESS_CLEANUP_TIMEOUT) + result_queue.close() + result_queue.join_thread() - assert process.exitcode == 0 assert result == ("ok", (1, 1)) assert not self.engine._close_signal.is_set()