diff --git a/agent_context/topics/data-pipeline/data-pipeline.md b/agent_context/topics/data-pipeline/data-pipeline.md index 0325fa85f..5e58ddd84 100644 --- a/agent_context/topics/data-pipeline/data-pipeline.md +++ b/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 diff --git a/agent_context/topics/data-pipeline/persistence.md b/agent_context/topics/data-pipeline/persistence.md index a39898ef2..4d4569fdd 100644 --- a/agent_context/topics/data-pipeline/persistence.md +++ b/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. | diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index 7b266735b..aa4282628 100644 --- a/docs/source/api_reference/public_api.rst +++ b/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 -------------------------------- diff --git a/embodichain/lab/gym/envs/managers/async_datasets.py b/embodichain/lab/gym/envs/managers/async_datasets.py index 9a58acac0..159f0c37a 100644 --- a/embodichain/lab/gym/envs/managers/async_datasets.py +++ b/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,11 +95,12 @@ 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` @@ -105,11 +108,16 @@ class AsyncLeRobotRecorder(LeRobotRecorder): ``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. diff --git a/embodichain/lab/gym/envs/managers/datasets.py b/embodichain/lab/gym/envs/managers/datasets.py index 61b3a9221..810ad0980 100644 --- a/embodichain/lab/gym/envs/managers/datasets.py +++ b/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 diff --git a/embodichain/lab/task_program/integrations/configured.py b/embodichain/lab/task_program/integrations/configured.py index 666b15827..a161e067a 100644 --- a/embodichain/lab/task_program/integrations/configured.py +++ b/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 diff --git a/embodichain_tasks/configs/tasks/special/stay_still_save_3cam/catalog.yaml b/embodichain_tasks/configs/tasks/special/stay_still_save_3cam/catalog.yaml new file mode 100644 index 000000000..ea825249a --- /dev/null +++ b/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 diff --git a/embodichain_tasks/configs/tasks/special/stay_still_save_3cam/env_ur10_3cam.json b/embodichain_tasks/configs/tasks/special/stay_still_save_3cam/env_ur10_3cam.json new file mode 100644 index 000000000..a2f95a5f9 --- /dev/null +++ b/embodichain_tasks/configs/tasks/special/stay_still_save_3cam/env_ur10_3cam.json @@ -0,0 +1,180 @@ +{ + "id": "StayStillSave3Cam-v1", + "physics": "default", + "max_episodes": 1, + "env": { + "dataset": { + "lerobot": { + "func": "AsyncLeRobotRecorder", + "mode": "save", + "params": { + "robot_meta": { + "robot_type": "UR10" + }, + "instruction": { + "lang": "Hold still for 300 steps while recording three RGB cameras" + }, + "extra": { + "scene_type": "benchmark", + "task_description": "stay still three cameras 300 steps", + "data_type": "sim" + }, + "use_videos": false, + "image_writer_threads": 8, + "image_writer_processes": 0, + "image_compress_level": 1, + "async_queue_maxsize": 8 + }, + "save_failed_episodes": true + } + } + }, + "robot": { + "uid": "UR10", + "fpath": "UniversalRobots/UR10/UR10.urdf", + "init_pos": [ + 0.0, + 0.0, + 0.7775 + ], + "init_qpos": [ + 1.57079, + -1.57079, + 1.57079, + -1.57079, + -1.57079, + -3.14159 + ] + }, + "sensor": [ + { + "sensor_type": "Camera", + "uid": "cam_high", + "width": 640, + "height": 480, + "intrinsics": [ + 488.1665344238281, + 488.1665344238281, + 320.0, + 240.0 + ], + "extrinsics": { + "eye": [ + 1, + 0, + 3 + ], + "target": [ + 0, + 0, + 1 + ] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_left", + "width": 640, + "height": 480, + "intrinsics": [ + 488.1665344238281, + 488.1665344238281, + 320.0, + 240.0 + ], + "extrinsics": { + "eye": [ + 0, + 2.8, + 2.4 + ], + "target": [ + 0, + 0, + 1 + ] + } + }, + { + "sensor_type": "Camera", + "uid": "cam_right", + "width": 640, + "height": 480, + "intrinsics": [ + 488.1665344238281, + 488.1665344238281, + 320.0, + 240.0 + ], + "extrinsics": { + "eye": [ + 0, + -2.8, + 2.4 + ], + "target": [ + 0, + 0, + 1 + ] + } + } + ], + "light": { + "direct": [ + { + "uid": "light_1", + "light_type": "sun", + "color": [ + 1.0, + 1.0, + 1.0 + ], + "intensity": 5.0, + "direction": [ + 0.0, + 0.0, + -1.0 + ] + } + ] + }, + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "CircleTableSimple/circle_table_simple.ply", + "compute_uv": true + }, + "attrs": { + "mass_props": { + "mass": 10.0 + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + } + }, + "body_scale": [ + 1, + 1, + 1 + ], + "body_type": "kinematic", + "init_pos": [ + 0.8, + 0.0, + 0.825 + ], + "init_rot": [ + 0, + 90, + 0 + ] + } + ], + "rigid_object": [], + "max_episode_steps": 301 +} diff --git a/embodichain_tasks/embodichain_tasks/special/stay_still_save_3cam.py b/embodichain_tasks/embodichain_tasks/special/stay_still_save_3cam.py new file mode 100644 index 000000000..d4c22ff71 --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/special/stay_still_save_3cam.py @@ -0,0 +1,59 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Three-camera stationary demo used to benchmark long LeRobot episodes.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +from embodichain.lab.gym.envs import DemoSegment, EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.utils import logger + +__all__ = ["StayStillSave3CamEnv"] + + +@register_env("StayStillSave3Cam-v1", max_episode_steps=301) +class StayStillSave3CamEnv(EmbodiedEnv): + """Hold the UR10 pose for 300 steps while recording three RGB cameras.""" + + def __init__( + self, + cfg: EmbodiedEnvCfg | None = None, + **kwargs: Any, + ) -> None: + super().__init__(cfg, **kwargs) + + def create_demo_segments(self, **kwargs: Any) -> Iterable[DemoSegment]: + """Yield one deterministic 300-step stationary demonstration segment.""" + + del kwargs + init_pose = self.robot.get_qpos().clone() + + def actions() -> Iterable[Any]: + for _ in range(300): + yield init_pose.clone() + + logger.log_info("Generated 300 hold-still demo actions for three cameras.") + yield DemoSegment( + actions=actions(), + name="hold_still_300_steps", + instruction="Hold still while recording three RGB cameras", + metadata={"segment_index": 0, "segment_count": 1, "steps": 300}, + progress_total_steps=300, + ) diff --git a/tests/gym/envs/managers/test_async_dataset_functors.py b/tests/gym/envs/managers/test_async_dataset_functors.py index 531c73c69..60b5693af 100644 --- a/tests/gym/envs/managers/test_async_dataset_functors.py +++ b/tests/gym/envs/managers/test_async_dataset_functors.py @@ -18,6 +18,7 @@ from __future__ import annotations +import threading from unittest.mock import Mock, patch import pytest @@ -180,6 +181,7 @@ def _make_recorder( mock_dataset: _MockDataset, *, action_contract: dict[str, object] | None = None, + async_queue_maxsize: int | None = None, ) -> AsyncLeRobotRecorder: """Build an AsyncLeRobotRecorder with the LeRobotDataset.create patched out.""" from embodichain.lab.gym.envs.managers.cfg import DatasetFunctorCfg @@ -194,6 +196,8 @@ def _make_recorder( } if action_contract is not None: params["action_contract"] = action_contract + if async_queue_maxsize is not None: + params["async_queue_maxsize"] = async_queue_maxsize cfg = DatasetFunctorCfg(func=AsyncLeRobotRecorder, params=params) with patch("embodichain.lab.gym.envs.managers.datasets.LeRobotDataset") as mock_cls: mock_cls.create.return_value = mock_dataset @@ -204,6 +208,64 @@ def _make_recorder( class TestAsyncLeRobotRecorder: """Tests for AsyncLeRobotRecorder enqueue/drain behavior.""" + def test_configures_bounded_queue_and_reports_stats(self): + """A positive queue limit enables backpressure without changing drain order.""" + env = _MockEnv(num_envs=1, steps=1) + recorder = _make_recorder(env, _MockDataset(), async_queue_maxsize=1) + try: + assert recorder._save_queue.maxsize == 1 + stats = recorder.async_stats + assert stats["queue_maxsize"] == 1 + assert stats["queue_size"] == 0 + assert stats["queue_peak"] == 0 + assert stats["backpressure_events"] == 0 + finally: + recorder.finalize() + + def test_rejects_negative_queue_limit(self): + """Negative queue limits fail before starting a background worker.""" + env = _MockEnv(num_envs=1, steps=1) + with pytest.raises(ValueError, match="async_queue_maxsize"): + _make_recorder(env, _MockDataset(), async_queue_maxsize=-1) + + def test_bounded_queue_blocks_until_worker_catches_up(self): + """A full queue blocks the producer and records the backpressure time.""" + env = _MockEnv(num_envs=3, steps=1) + recorder = _make_recorder(env, _MockDataset(), async_queue_maxsize=1) + worker_entered = threading.Event() + release_worker = threading.Event() + + def persist_payload(*args, **kwargs): + worker_entered.set() + assert release_worker.wait(timeout=2.0) + return True + + recorder._persist_episode_payload = Mock(side_effect=persist_payload) + enqueue_done = threading.Event() + + def enqueue_payloads() -> None: + recorder(env, env_ids=torch.tensor([0, 1, 2])) + enqueue_done.set() + + enqueue_thread = threading.Thread(target=enqueue_payloads) + enqueue_thread.start() + + try: + assert worker_entered.wait(timeout=2.0) + assert not enqueue_done.wait(timeout=0.05) + finally: + release_worker.set() + enqueue_thread.join(timeout=2.0) + env.current_rollout_step = 0 + recorder.finalize() + + assert enqueue_done.is_set() + stats = recorder.async_stats + assert stats["queue_maxsize"] == 1 + assert stats["queue_peak"] == 1 + assert stats["backpressure_events"] >= 1 + assert stats["backpressure_s"] > 0.0 + def test_call_accepts_construction_only_kwargs(self): """``__call__`` must accept image_writer_threads/processes (regression). diff --git a/tests/gym/envs/managers/test_dataset_functors.py b/tests/gym/envs/managers/test_dataset_functors.py index 3c248052a..242c034aa 100644 --- a/tests/gym/envs/managers/test_dataset_functors.py +++ b/tests/gym/envs/managers/test_dataset_functors.py @@ -257,6 +257,111 @@ def test_initialization_with_defaults(self, mock_lerobot_dataset): assert recorder.dataset_fps == 30 assert mock_lerobot_dataset.create.call_args.kwargs["fps"] == 30 + @pytest.mark.parametrize( + ("value", "error_type"), + [ + (True, TypeError), + (1.0, TypeError), + ("3", TypeError), + (-1, ValueError), + (10, ValueError), + ], + ) + @patch("embodichain.lab.gym.envs.managers.datasets.LeRobotDataset") + def test_image_compress_level_rejects_invalid_values( + self, mock_lerobot_dataset, value, error_type + ): + """Invalid compression values fail before LeRobot dataset creation.""" + with pytest.raises(error_type, match="image_compress_level"): + LeRobotRecorder( + MockFunctorCfg(params={"image_compress_level": value}), + MockEnvForDataset(), + ) + + mock_lerobot_dataset.create.assert_not_called() + + @pytest.mark.parametrize(("threads", "processes"), [(0, 0), (4, 2)]) + @patch("embodichain.lab.gym.envs.managers.datasets._ImageCompressionLeRobotDataset") + def test_image_compress_level_uses_subclass_and_writer_options( + self, mock_dataset_cls, threads, processes + ): + """Explicit compression uses the private subclass and preserves writer args.""" + mock_dataset = Mock() + mock_dataset.meta = Mock() + mock_dataset.meta.info = {"fps": 30} + mock_dataset_cls.create.return_value = mock_dataset + recorder = LeRobotRecorder( + MockFunctorCfg( + params={ + "image_compress_level": 3, + "image_writer_threads": threads, + "image_writer_processes": processes, + } + ), + MockEnvForDataset(), + ) + + mock_dataset_cls.create.assert_called_once() + assert ( + mock_dataset_cls.create.call_args.kwargs["image_writer_threads"] == threads + ) + assert ( + mock_dataset_cls.create.call_args.kwargs["image_writer_processes"] + == processes + ) + assert recorder.image_compress_level == 3 + assert mock_dataset.image_compress_level == 3 + + @pytest.mark.parametrize("image_compress_level", [0, 1, 3, 6]) + def test_image_compress_level_preserves_png_pixels( + self, tmp_path, image_compress_level + ): + """Configured PNG levels decode to the exact source pixels.""" + from embodichain.lab.gym.envs.managers import datasets + + image = np.zeros((8, 10, 3), dtype=np.uint8) + image[2:6, 3:8] = (13, 101, 247) + image[0, 0] = (255, 7, 91) + dataset = datasets._ImageCompressionLeRobotDataset.__new__( + datasets._ImageCompressionLeRobotDataset + ) + dataset.image_writer = None + dataset.image_compress_level = image_compress_level + path = tmp_path / f"image_{image_compress_level}.png" + + dataset._save_image(image, path, compress_level=6) + + from PIL import Image + + with Image.open(path) as decoded: + np.testing.assert_array_equal(np.asarray(decoded), image) + + def test_image_compress_level_reaches_async_image_writer(self, tmp_path): + """The configured level reaches a real LeRobot image-writer thread.""" + from lerobot.datasets.image_writer import AsyncImageWriter + from embodichain.lab.gym.envs.managers import datasets + + image = np.zeros((8, 10, 3), dtype=np.uint8) + image[2:6, 3:8] = (13, 101, 247) + dataset = datasets._ImageCompressionLeRobotDataset.__new__( + datasets._ImageCompressionLeRobotDataset + ) + writer = AsyncImageWriter(num_processes=0, num_threads=2) + dataset.image_writer = writer + dataset.image_compress_level = 3 + path = tmp_path / "async_image.png" + + try: + dataset._save_image(image, path, compress_level=6) + writer.wait_until_done() + finally: + writer.stop() + + from PIL import Image + + with Image.open(path) as decoded: + np.testing.assert_array_equal(np.asarray(decoded), image) + @patch("embodichain.lab.gym.envs.managers.datasets.LeRobotDataset") def test_non_integer_environment_frequency_is_rejected(self, mock_lerobot_dataset): """LeRobot recording rejects cadence that its integer FPS cannot encode.""" diff --git a/tests/gym/envs/tasks/test_stay_still_save_3cam.py b/tests/gym/envs/tasks/test_stay_still_save_3cam.py new file mode 100644 index 000000000..ab7fdcca2 --- /dev/null +++ b/tests/gym/envs/tasks/test_stay_still_save_3cam.py @@ -0,0 +1,40 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from embodichain_tasks.special.stay_still_save_3cam import ( + StayStillSave3CamEnv, +) + + +def test_stay_still_three_camera_demo_has_300_steps() -> None: + """The benchmark task emits one deterministic segment of 300 actions.""" + + env = StayStillSave3CamEnv.__new__(StayStillSave3CamEnv) + env.robot = SimpleNamespace(get_qpos=lambda: torch.zeros(1, 6)) + + (segment,) = tuple(env.create_demo_segments()) + actions = list(segment.actions) + + assert segment.name == "hold_still_300_steps" + assert segment.progress_total_steps == 300 + assert len(actions) == 300 + assert all(action.shape == (1, 6) for action in actions) diff --git a/tests/gym/envs/test_official_task_layout.py b/tests/gym/envs/test_official_task_layout.py index fceae2c67..4c599fd8b 100644 --- a/tests/gym/envs/test_official_task_layout.py +++ b/tests/gym/envs/test_official_task_layout.py @@ -45,6 +45,7 @@ "StackBlocksTwo-v1": "embodichain_tasks.manipulation.tableware.stack_blocks_two", "StackCups-v1": "embodichain_tasks.manipulation.tableware.stack_cups", "StayStillSave-v1": "embodichain_tasks.special.stay_still_save", + "StayStillSave3Cam-v1": "embodichain_tasks.special.stay_still_save_3cam", } REMOVED_AGENT_ENV_IDS = {"PourWaterAgent-v3", "RearrangementAgent-v3"} CONFIG_DEFINED_TASK_PROGRAM_TASKS = {"pour_water"}