# ----------------------------------------------------------------------------
# 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.
# ----------------------------------------------------------------------------
"""Asynchronous LeRobot recorder for parallel environments.
This module provides :class:`AsyncLeRobotRecorder`, which decouples episode
saving from the simulation loop. It is the recommended recorder when running
many parallel environments that complete episodes together: instead of
blocking ``env.reset()`` while episodes are converted and flushed to disk, the
completed episode buffers are cloned and handed to a background worker thread,
so the simulator can keep stepping.
"""
from __future__ import annotations
import copy
import queue
import threading
import time
from typing import TYPE_CHECKING, Dict, Optional, Union
import torch
from embodichain.utils import logger
from embodichain.lab.gym.envs.demo import DEMO_ANNOTATION_KEYS
from .datasets import LeRobotRecorder
__all__ = ["AsyncLeRobotRecorder"]
if TYPE_CHECKING:
from embodichain.lab.gym.envs import EmbodiedEnv
from .cfg import DatasetFunctorCfg
try:
from lerobot.datasets.lerobot_dataset import LeRobotDataset # noqa: F401
LEROBOT_AVAILABLE = True
__all__ = ["AsyncLeRobotRecorder"]
except ImportError:
LEROBOT_AVAILABLE = False
__all__ = []
[docs]
class AsyncLeRobotRecorder(LeRobotRecorder):
"""LeRobot recorder that saves episodes on a background thread.
Drop-in replacement for :class:`LeRobotRecorder` selected via the dataset
config ``"func": "AsyncLeRobotRecorder"``. It shares the same on-disk
format and feature-building logic; only the *timing* of the save differs.
Why this helps parallel environments
-------------------------------------
In the synchronous recorder, :meth:`LeRobotRecorder.__call__` runs inside
``env.reset()`` (via ``DatasetManager.apply``) and blocks the simulator
while it iterates ``add_frame`` + ``save_episode`` for every finished env.
With ``num_envs=N`` all finishing at once, the sim stalls for the *sum* of
all episodes' save time every reset.
This recorder instead, on each ``apply``:
1. Reads each env's rollout-buffer slice (``obs``/``actions``).
2. **Clones** the slice to CPU (detached from the live buffer).
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._persist_episode_payload` on each cloned payload.
Correctness
-----------
* **No concurrent dataset access.** ``LeRobotDataset`` is not thread-safe;
only the worker thread ever calls ``add_frame`` / ``save_episode`` /
mutates ``curr_episode``. The main thread only enqueues and, at close,
drains.
* **No buffer race.** The slice is cloned in the caller thread *before*
the buffer is cleared on reset, so the worker never reads memory that
the sim is overwriting.
* **Ordering.** A single worker preserves FIFO episode order, so
``episode_index`` assignment is deterministic.
* **Drain on close.** :meth:`finalize` joins the worker before the parent
flushes the image writer and finalizes the dataset.
.. note::
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). ``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.
"""
[docs]
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"
)
super().__init__(cfg, env)
if self.record_eef_observation:
raise ValueError(
"AsyncLeRobotRecorder cannot record EEF observations safely; "
"use LeRobotRecorder so FK runs on the simulation thread."
)
# 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. 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()
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",
daemon=True,
)
self._worker.start()
logger.log_info(
"[AsyncLeRobotRecorder] Background save worker started; "
"episode saves will not block env.reset()."
)
def _worker_loop(self) -> None:
"""Consume cloned episodes from the queue and persist them."""
while True:
item = self._save_queue.get()
if item is None:
# Sentinel: finalize() is draining. Exit the worker.
break
env_id, obs_clone, action_clone, annotations, episode_metadata = item
fragment_id = None
is_fragment = bool(
episode_metadata is not None and episode_metadata.get("fragment", False)
)
try:
fragment_id = self._fragment_id_from_metadata(episode_metadata)
saved = self._persist_episode_payload(
env_id,
obs_clone,
action_clone,
annotations=annotations,
episode_metadata=episode_metadata,
)
if not saved:
self._record_background_error(env_id, "episode save returned False")
except BaseException as error: # noqa: BLE001 - worker must not die
label = (
f"fragment {fragment_id!r}"
if fragment_id is not None
else ("fragment payload" if is_fragment else "episode")
)
self._record_background_error(env_id, f"{label}: {error}")
logger.log_warning(
f"[AsyncLeRobotRecorder] Background worker failed on "
f"env {env_id} {label}: {error}"
)
def _record_background_error(self, env_id: int, message: str) -> None:
"""Remember one failed committed payload for the durability barrier."""
with self._background_error_lock:
self._background_errors.append((env_id, message))
def __call__(
self,
env: EmbodiedEnv,
env_ids: Union[torch.Tensor, None],
save_path: Optional[str] = None,
robot_meta: Optional[Dict] = None,
instruction: Optional[str] = None,
extra: Optional[Dict] = None,
use_videos: bool = False,
**kwargs,
) -> None:
"""Enqueue completed episodes for background saving.
Reads each env's rollout-buffer slice, clones it to CPU (so the worker
is immune to buffer reuse on reset), and pushes it onto the queue. The
actual conversion and disk I/O happen asynchronously on the worker
thread. Returns immediately so ``env.reset()`` is not blocked.
Args:
env: The environment instance.
env_ids: Environment IDs to save. If None, enqueues all envs.
save_path: Unused at call time (honored at construction).
robot_meta: Unused at call time (honored at construction).
instruction: Unused at call time (honored at construction).
extra: Unused at call time (honored at construction).
use_videos: Unused at call time (honored at construction).
**kwargs: Construction-only params (e.g. ``image_writer_threads``)
passed through by ``DatasetManager.apply``; honored in
:meth:`__init__`, ignored here.
"""
context = getattr(env, "episode_context", None)
if context is not None:
env_ids = context.filter_recordable(env_ids)
if len(env_ids) == 0:
return
# 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)
# Initial reset may reach the recorder before the first
# transition. Do not enqueue an empty payload and later
# report it as a failed committed episode.
if step <= 0:
continue
obs_view = env.rollout_buffer["obs"][env_id, :step]
stored_actions = env.rollout_buffer["actions"][env_id, :step]
action_view = self._contract_action_list(env_id, step, stored_actions)
# 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
)
payloads = list(
self._episode_payloads(
env_id,
obs_clone,
action_clone,
annotations,
episode_metadata,
)
)
if (
episode_metadata is not None
and episode_metadata.get("output_mode") == "segment_fragments"
and not payloads
):
raise RuntimeError(
f"Committed fragment collection for env {env_id} had no "
"eligible segment spans."
)
for payload in payloads:
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,
}
[docs]
def finalize(self) -> Optional[str]:
"""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.
Returns:
The finalized dataset path.
Raises:
RuntimeError: If any queued episode or dataset resource failed.
"""
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