Improve async LeRobot persistence throughput and memory bounds

agent_context/topics/data-pipeline/data-pipeline.md

@@ -75,8 +75,10 @@ deduplication and sticky partial-commit errors; they are not a crash-recovery jo
 
 Async persistence clones tensor payloads to CPU and copies metadata before
 enqueue. One FIFO worker owns LeRobot access. Finalize rejects new work, drains,
-joins and surfaces errors. The queue is unbounded: writer lag can grow memory.
-Choose synchronous persistence or throttle production when memory bounds matter.
+joins and surfaces errors. Set ``async_queue_maxsize`` to a positive payload
+count to apply producer backpressure when memory bounds matter; the default
+``0`` preserves the historical unbounded queue. Queue depth and backpressure
+time are available through the recorder's async stats.
 
 ## Change sites and focused validation
 

agent_context/topics/data-pipeline/persistence.md

@@ -45,8 +45,9 @@
    processing and are checked again before sync writes or async enqueue.
 9. The async worker records per-payload errors but continues draining later items. Finalize
    rejects new work, queues a sentinel after all existing payloads, joins the worker, calls
-   base finalization, and aggregates background plus storage errors. Its queue is unbounded,
-   so sustained writer lag can grow RAM.
+   base finalization, and aggregates background plus storage errors. A positive
+   ``async_queue_maxsize`` bounds pending payloads and blocks the producer when the writer
+   lags; ``0`` keeps the historical unbounded queue.
 10. When sidecars are enabled, depth uses one HEVC MP4 per sensor per episode. Each writer
     creates a same-directory temp file and replaces the final file on close; failed/empty
     episodes clean the temp. The manager attempts every sensor, records only successfully
@@ -70,5 +71,5 @@
 | Episode indices reorder or LeRobot races | More than one persistence worker or non-FIFO access; retain single-worker ownership. |
 | Retry writes a duplicate episode | A post-commit failure bypassed sticky fragment state or fragment ID changed. Keep `save_episode()` as the explicit irreversible point. |
 | Depth file exists without matching episode, or vice versa | Determine whether failure was before or after LeRobot commit; only pre-commit work can be aborted transactionally. |
-| Async collection consumes growing RAM | Writer cannot match producer throughput; queue is intentionally unbounded. Throttle production or select sync persistence. |
+| Async collection consumes growing RAM | Writer cannot match producer throughput; set a positive ``async_queue_maxsize`` or throttle production. The default ``0`` remains unbounded for compatibility. |
 | Finalize succeeds while work remains | Sentinel/join ordering or error aggregation changed; queued items must precede sentinel and base finalization. |

docs/source/api_reference/public_api.rst

@@ -2394,6 +2394,15 @@ embodichain_tasks.special.stay_still_save
 
    StayStillSaveEnv
 
+embodichain_tasks.special.stay_still_save_3cam
+----------------------------------------------
+
+.. currentmodule:: embodichain_tasks.special.stay_still_save_3cam
+
+.. autosummary::
+
+   StayStillSave3CamEnv
+
 embodichain_tasks.utils.importer
 --------------------------------
 

embodichain/lab/gym/envs/managers/async_datasets.py

@@ -29,6 +29,7 @@
 import copy
 import queue
 import threading
+import time
 from typing import TYPE_CHECKING, Dict, Optional, Union
 
 import torch
@@ -41,6 +42,7 @@
 
 if TYPE_CHECKING:
     from embodichain.lab.gym.envs import EmbodiedEnv
+    from .cfg import DatasetFunctorCfg
 
 try:
     from lerobot.datasets.lerobot_dataset import LeRobotDataset  # noqa: F401
@@ -93,23 +95,29 @@ class AsyncLeRobotRecorder(LeRobotRecorder):
       flushes the image writer and finalizes the dataset.
 
     .. note::
-        The clone copies each episode's camera frames into host RAM. Memory
-        use is bounded by how far the worker falls behind (typically it keeps
-        up, since per-frame PNG write is the only heavy step and can itself be
-        offloaded via ``image_writer_threads``). For very high resolutions or
-        many envs, monitor RSS.
+        The clone copies each episode's camera frames into host RAM. The
+        default queue is unbounded. A positive ``async_queue_maxsize`` limits
+        waiting payloads; the active worker payload, caller's current clone,
+        rollout buffer and image-writer queue consume additional memory.
+        When the queue is full, committing an episode blocks until space is
+        available, allowing storage to catch up without dropping data.
 
     Args:
         cfg: :class:`~embodichain.lab.gym.envs.managers.cfg.DatasetFunctorCfg`
             with the same ``params`` as ``LeRobotRecorder``. The
             ``image_writer_threads`` / ``image_writer_processes`` params are
             honored and combine with the background worker (two levels of
             async: episode conversion off the sim thread, PNG writes off the
-            worker thread).
+            worker thread). ``async_queue_maxsize`` optionally bounds the
+            number of committed payloads waiting for the worker; zero keeps
+            the historical unbounded queue behavior.
         env: The environment instance.
     """
 
-    def __init__(self, cfg, env: EmbodiedEnv):
+    def __init__(self, cfg: DatasetFunctorCfg, env: EmbodiedEnv) -> None:
+        queue_maxsize = cfg.params.get("async_queue_maxsize", 0)
+        if type(queue_maxsize) is not int or queue_maxsize < 0:
+            raise ValueError("async_queue_maxsize must be a non-negative integer")
         if not LEROBOT_AVAILABLE:
             logger.log_error(
                 "LeRobot is not installed. Please install it with: pip install lerobot"
@@ -123,8 +131,17 @@ def __init__(self, cfg, env: EmbodiedEnv):
 
         # Single-worker queue. A single worker guarantees LeRobotDataset is
         # 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()
+        # episode ordering deterministic. A bounded queue provides optional
+        # backpressure so high-resolution payload clones cannot grow memory
+        # without limit while the writer falls behind.
+        self.async_queue_maxsize = queue_maxsize
+        self._save_queue: "queue.Queue[Optional[tuple]]" = queue.Queue(
+            maxsize=self.async_queue_maxsize
+        )
+        self._queue_metrics_lock = threading.Lock()
+        self._max_queue_depth = 0
+        self._backpressure_events = 0
+        self._backpressure_s = 0.0
         self._background_error_lock = threading.Lock()
         self._background_errors: list[tuple[int, str]] = []
         self._async_finalize_lock = threading.Lock()
@@ -272,7 +289,53 @@ def __call__(
                         "eligible segment spans."
                     )
                 for payload in payloads:
-                    self._save_queue.put(payload)
+                    self._enqueue_payload(payload)
+
+    def _enqueue_payload(self, payload: tuple) -> None:
+        """Enqueue one payload and account for bounded-queue backpressure."""
+
+        enqueue_elapsed = 0.0
+        queue_was_full = False
+        try:
+            self._save_queue.put_nowait(payload)
+        except queue.Full:
+            queue_was_full = True
+            enqueue_start = time.perf_counter()
+            self._save_queue.put(payload)
+            enqueue_elapsed = time.perf_counter() - enqueue_start
+        with self._queue_metrics_lock:
+            self._max_queue_depth = max(
+                self._max_queue_depth,
+                (
+                    self.async_queue_maxsize
+                    if queue_was_full
+                    else self._save_queue.qsize()
+                ),
+            )
+            if queue_was_full:
+                self._backpressure_events += 1
+                self._backpressure_s += enqueue_elapsed
+
+    @property
+    def async_stats(self) -> dict[str, float | int]:
+        """Return a snapshot of queue depth and producer waiting time.
+
+        Queue sizes and their sampled peak are approximate. Backpressure
+        counts attempts that found the queue full and measures only time
+        waiting to enqueue, excluding payload cloning and finalization.
+
+        Returns:
+            Queue size/limit/peak, blocked enqueue count and waiting seconds.
+        """
+
+        with self._queue_metrics_lock:
+            return {
+                "queue_size": self._save_queue.qsize(),
+                "queue_maxsize": self.async_queue_maxsize,
+                "queue_peak": self._max_queue_depth,
+                "backpressure_events": self._backpressure_events,
+                "backpressure_s": self._backpressure_s,
+            }
 
     def finalize(self) -> Optional[str]:
         """Drain committed writes, finalize storage, and surface all failures.

embodichain/lab/gym/envs/managers/datasets.py

@@ -109,6 +109,25 @@
     __all__ = []
 
 
+if LEROBOT_AVAILABLE:
+
+    class _ImageCompressionLeRobotDataset(LeRobotDataset):
+        """LeRobot dataset that overrides image PNG compression on request."""
+
+        def _save_image(self, image: Any, fpath: Path, compress_level: int = 1) -> None:
+            # LeRobot passes level 6 for image features and level 1 for the
+            # temporary PNGs used by video features.  Keep the latter at its
+            # upstream default so ``image_compress_level`` only tunes image
+            # feature writes.
+            image_compress_level = getattr(self, "image_compress_level", None)
+            if image_compress_level is not None and compress_level == 6:
+                compress_level = image_compress_level
+            super()._save_image(image, fpath, compress_level)
+
+else:
+    _ImageCompressionLeRobotDataset = None
+
+
 class LeRobotRecorder(Functor):
     """Functor for recording episodes in LeRobot format.
 
@@ -135,6 +154,9 @@ def __init__(self, cfg: DatasetFunctorCfg, env: EmbodiedEnv):
                 - use_videos: Whether to save videos
                 - image_writer_threads: Number of threads for image writing
                 - image_writer_processes: Number of processes for image writing
+                - image_compress_level: Optional PNG compression level (0-9)
+                  for image features; video temporary PNGs keep LeRobot's
+                  default level 1
             env: The environment instance
         """
         if not LEROBOT_AVAILABLE:
@@ -184,6 +206,9 @@ def __init__(self, cfg: DatasetFunctorCfg, env: EmbodiedEnv):
 
         # Experimental parameters for extra episode info saving.
         self.use_videos = params.get("use_videos", False)
+        self.image_compress_level = self._parse_image_compress_level(
+            params.get("image_compress_level")
+        )
 
         # Async image writing (lerobot official AsyncImageWriter).
         # When > 0, per-frame PNG writes are offloaded to a thread/process pool
@@ -227,6 +252,31 @@ def __init__(self, cfg: DatasetFunctorCfg, env: EmbodiedEnv):
         # Initialize dataset
         self._initialize_dataset()
 
+    @staticmethod
+    def _parse_image_compress_level(value: Any) -> int | None:
+        """Validate an optional lossless PNG compression level.
+
+        Args:
+            value: ``None`` to preserve LeRobot defaults, or an integer from
+                0 through 9.
+
+        Returns:
+            The validated compression level, or ``None``.
+
+        Raises:
+            TypeError: If ``value`` is not ``None`` or an integer.
+            ValueError: If the integer is outside the PNG level range.
+        """
+        if value is None:
+            return None
+        if type(value) is not int:
+            raise TypeError(
+                "image_compress_level must be None or an integer from 0 to 9."
+            )
+        if not 0 <= value <= 9:
+            raise ValueError("image_compress_level must be between 0 and 9.")
+        return value
+
     @staticmethod
     def _parse_action_contract(value: Any) -> dict[str, Any] | None:
         """Validate and normalize the optional action contract."""
@@ -1464,7 +1514,12 @@ def _initialize_dataset(self) -> None:
 
         features = self._build_features()
 
-        self.dataset = LeRobotDataset.create(
+        dataset_cls = (
+            _ImageCompressionLeRobotDataset
+            if self.image_compress_level is not None
+            else LeRobotDataset
+        )
+        self.dataset = dataset_cls.create(
             repo_id=dataset_name,
             fps=self.dataset_fps,
             root=str(self.dataset_full_path),
@@ -1475,6 +1530,11 @@ def _initialize_dataset(self) -> None:
             image_writer_processes=self.image_writer_processes,
             image_writer_threads=self.image_writer_threads,
         )
+        if self.image_compress_level is not None:
+            # ``LeRobotDataset.create`` constructs the object with
+            # ``cls.__new__(cls)``. Set the override after create, before the
+            # first frame reaches ``_save_image``.
+            self.dataset.image_compress_level = self.image_compress_level
         logger.log_info(f"Created LeRobot dataset at: {self.dataset_full_path}")
 
         # Set up the depth sidecar manager now that the dataset root and fps are

embodichain/lab/task_program/integrations/configured.py

@@ -19,6 +19,7 @@
 from __future__ import annotations
 
 from collections.abc import Mapping, Sequence
+from copy import deepcopy
 from dataclasses import asdict, dataclass
 from enum import Enum
 import hashlib

embodichain_tasks/configs/tasks/special/stay_still_save_3cam/catalog.yaml

@@ -0,0 +1,11 @@
+task_key: stay_still_save_3cam
+title: Stay Still Save Three Cameras
+summary: Collect 300-step stationary UR10 episodes from three RGB cameras.
+tags:
+- special
+- data_collection
+- benchmark
+default_deployment: default
+deployments:
+  default:
+    config: env_ur10_3cam.json