Compare commits

...

14 Commits

Author SHA1 Message Date
Steven Palma 6c60f4651b update image 2026-07-29 16:12:39 +02:00
Steven Palma 117216b29b refactor(transforms): several updates 2026-07-29 15:43:13 +02:00
liyux 2422d43fb1 tune showcase: softer shadow, dropout, jitter intensity 2026-07-29 15:27:31 +02:00
liyux 53a5cffb4c tune showcase to balanced augmentation intensity 2026-07-29 15:27:31 +02:00
liyux 2a3f40c673 update showcase with better sample frame 2026-07-29 15:27:31 +02:00
liyux 2ddfb5a376 add augmentation showcase image for PR 2026-07-29 15:27:31 +02:00
Yuxian LI d8b0a6b17f feat(transforms): add 8 robotics-relevant image augmentations
Add GaussianNoise, MotionBlur, JPEGCompression, GaussianPatchBrightness,
RandomShadow, CoarseDropout, GammaCorrection, and PlanckianJitter.

Each transform addresses a real-world failure mode not covered by the
existing 6 defaults (sensor noise, motion blur, compression artifacts,
uneven lighting, cast shadows, partial occlusion, exposure variation,
color temperature shift).

All transforms are pure PyTorch, follow the make_params/transform
pattern, and integrate with ImageTransformConfig via a registry.
2026-07-29 15:27:31 +02:00
Steven Palma 7d615acf9a fix(robots): retry SO follower/leader bus reads on transient Feetech errors (#4207)
* fix(robots): retry SO follower/leader bus reads on transient Feetech errors

SO-100/SO-101 teleoperation aborts when a sync_read of Present_Position
returns a corrupted status packet ("Incorrect status packet!"), which the
Feetech bus emits intermittently under load. The read path already supports
a num_retry argument but the SO follower and leader never used it, so a single
transient failure crashed the control loop.

Add a max_read_retry config option (default 3) to SOFollowerConfig and
SOLeaderConfig and forward it to every Present_Position sync_read. Retries are
immediate and only happen on failure, so the steady-state read cost is
unchanged; set max_read_retry=0 to restore the previous behavior.

Fixes #3131

* chore(robots): change defaults

---------

Co-authored-by: isaka1022 <isaka1022@gmail.com>
2026-07-29 15:20:14 +02:00
Steven Palma 09572babee perf(docker): split dependency install from source copy for CI layer caching (#4208)
* perf(docker): split dep install from src copy for CI layer caching

Install third-party deps (torch + all extras) in a layer keyed only on
pyproject.toml + uv.lock via --no-install-project, then copy src and install
the local package. Editing src/ no longer busts the heavy dependency layer,
so BuildKit layer cache hits across CI builds.

Applied to both Dockerfile.user and Dockerfile.internal.

* chore(ci): less verbose comments + copy all files

---------

Co-authored-by: dongmao.zhang <dongmao.zhang@bytedance.com>
2026-07-29 15:04:46 +02:00
Predrag Cvetkovic 35339d31e5 fix(datasets): bound memory of augment_dataset_quantile_stats by sampling frames (#3749)
* fix(datasets): bound memory of augment_dataset_quantile_stats by sampling frames

Per-episode stats previously materialized every frame (and decoded up to 16
episodes in parallel), so peak memory scaled with episode length and OOM'd on
large datasets (#2889). Numeric features are now read in full from the table
(exact), while only image/video frames are sub-sampled per episode using the
existing sample_indices heuristic. Worker count is configurable via
LEROBOT_STATS_MAX_WORKERS; --no-sampling restores exact behavior.

* Update tests/datasets/test_augment_quantile_stats.py

Co-authored-by: Haoming Song <1847575517@qq.com>
Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com>

---------

Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com>
Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
Co-authored-by: Haoming Song <1847575517@qq.com>
2026-07-29 12:32:03 +02:00
Steven Palma f37be3edbe fix(eval): prevent eval_policy crash when start_seed is None and num_envs>1 (#4203)
* fix(eval): align seed list length with num_envs when unseeded

eval_policy appended a single None per batch to all_seeds on the unseeded path while the reward and success lists grew by num_envs. The per-episode zip(..., strict=True) then raised ValueError for num_envs > 1. Extend all_seeds by num_envs so the lists stay aligned.

* chore(tests): delete lerobot_eval test

---------

Co-authored-by: Devin Lai <markauto75@gmail.com>
2026-07-28 18:41:28 +02:00
Khalil Meftah 4d076845ac fix peft factory test mocking (#4201) 2026-07-28 17:54:58 +02:00
Steven Palma 413972c812 fix(env): eval env lifecycle (#4194)
Co-authored-by: itxaiohanglover <1531137510@qq.com>
Co-authored-by: nickndeng <nickndeng@gmail.com>
Co-authored-by: nickndeng <107904079+nickndeng@users.noreply.github.com>
Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>
2026-07-28 16:45:48 +02:00
Steven Palma 0449aa02f6 fix(utils): handle missing/unresponsive TTS on Linux (#4199)
* fix: handle missing/unresponsive TTS on Linux

spd-say may be installed but hang indefinitely when speech-dispatcher
is not running. Add a 5s timeout and catch TimeoutExpired alongside
FileNotFoundError so recording continues without audio.

* chore(utils): add log warning for say

---------

Co-authored-by: Jiwen Cai <jiwenc@nvidia.com>
2026-07-28 16:45:32 +02:00
19 changed files with 899 additions and 60 deletions
+4 -5
View File
@@ -68,17 +68,16 @@ ENV HOME=/home/user_lerobot \
# issues with MuJoCo and OpenGL drivers. # issues with MuJoCo and OpenGL drivers.
RUN uv venv --python python${PYTHON_VERSION} RUN uv venv --python python${PYTHON_VERSION}
# Install Python dependencies for caching # Install third-party dependencies separately for layer caching
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./ COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
COPY --chown=user_lerobot:user_lerobot src/ src/ RUN uv sync --locked --extra all --no-install-project --no-cache
RUN uv sync --locked --extra all --no-cache
RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas
# Copy the rest of the application source code # Copy the application source code and install the local project
# Make sure to have the git-LFS files for testing # Make sure to have the git-LFS files for testing
COPY --chown=user_lerobot:user_lerobot . . COPY --chown=user_lerobot:user_lerobot . .
RUN uv sync --locked --extra all --no-cache
# Set the default command # Set the default command
CMD ["/bin/bash"] CMD ["/bin/bash"]
+4 -5
View File
@@ -60,15 +60,14 @@ ENV HOME=/home/user_lerobot \
# run other Python projects in the same container without dependency conflicts. # run other Python projects in the same container without dependency conflicts.
RUN uv venv RUN uv venv
# Install Python dependencies for caching # Install third-party dependencies separately for layer caching
COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./ COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
COPY --chown=user_lerobot:user_lerobot src/ src/ RUN uv sync --locked --extra all --no-install-project --no-cache
RUN uv sync --locked --extra all --no-cache # Copy the application code and install the local project
# Copy the rest of the application code
# Make sure to have the git-LFS files for testing # Make sure to have the git-LFS files for testing
COPY --chown=user_lerobot:user_lerobot . . COPY --chown=user_lerobot:user_lerobot . .
RUN uv sync --locked --extra all --no-cache
# Set the default command # Set the default command
CMD ["/bin/bash"] CMD ["/bin/bash"]
Binary file not shown.

After

Width:  |  Height:  |  Size: 682 KiB

+6 -1
View File
@@ -384,7 +384,12 @@ class LiberoEnv(gym.Env):
def close(self): def close(self):
if self._env is not None: if self._env is not None:
self._env.close() try:
self._env.close()
finally:
# LIBERO deletes its inner env on close, so this wrapper must
# be recreated before the next reset.
self._env = None
def _make_env_fns( def _make_env_fns(
@@ -46,6 +46,12 @@ class SOFollowerConfig:
position_i_coefficient: int = 0 position_i_coefficient: int = 0
position_d_coefficient: int = 32 position_d_coefficient: int = 32
# Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
# return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
# at once, which otherwise aborts the control loop. Retries are immediate (no sleep) and only happen on
# failure, so the steady-state read cost is unchanged.
num_read_retries: int = 2
@RobotConfig.register_subclass("so101_follower") @RobotConfig.register_subclass("so101_follower")
@RobotConfig.register_subclass("so100_follower") @RobotConfig.register_subclass("so100_follower")
@@ -180,7 +180,7 @@ class SOFollower(Robot):
def get_observation(self) -> RobotObservation: def get_observation(self) -> RobotObservation:
# Read arm position # Read arm position
start = time.perf_counter() start = time.perf_counter()
obs_dict = self.bus.sync_read("Present_Position") obs_dict = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()} obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()}
dt_ms = (time.perf_counter() - start) * 1e3 dt_ms = (time.perf_counter() - start) * 1e3
logger.debug(f"{self} read state: {dt_ms:.1f}ms") logger.debug(f"{self} read state: {dt_ms:.1f}ms")
@@ -221,7 +221,7 @@ class SOFollower(Robot):
# Cap goal position when too far away from present position. # Cap goal position when too far away from present position.
# /!\ Slower fps expected due to reading from the follower. # /!\ Slower fps expected due to reading from the follower.
if self.config.max_relative_target is not None: if self.config.max_relative_target is not None:
present_pos = self.bus.sync_read("Present_Position") present_pos = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in goal_pos.items()} goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in goal_pos.items()}
goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target) goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target)
@@ -36,6 +36,7 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
import argparse import argparse
import concurrent.futures import concurrent.futures
import logging import logging
import os
from pathlib import Path from pathlib import Path
import numpy as np import numpy as np
@@ -52,6 +53,7 @@ from lerobot.datasets import (
get_feature_stats, get_feature_stats,
write_stats, write_stats,
) )
from lerobot.datasets.compute_stats import sample_indices
from lerobot.utils.utils import init_logging from lerobot.utils.utils import init_logging
@@ -77,12 +79,14 @@ def has_quantile_stats(stats: dict[str, dict] | None, quantile_list_keys: list[s
return False return False
def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict: def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampling: bool = True) -> dict:
"""Process a single episode and return its statistics. """Process a single episode and return its statistics.
Args: Args:
dataset: The LeRobot dataset dataset: The LeRobot dataset
episode_idx: Index of the episode to process episode_idx: Index of the episode to process
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
Returns: Returns:
Dictionary containing episode statistics Dictionary containing episode statistics
@@ -92,16 +96,31 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"] start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"] end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]
collected_data: dict[str, list] = {} episode_len = end_idx - start_idx
for idx in range(start_idx, end_idx):
item = dataset[idx]
for key, value in item.items():
if key not in dataset.features:
continue
if key not in collected_data: # Images/video are the memory hog, so sub-sample those frames per episode;
collected_data[key] = [] # numeric columns are cheap, so read them in full (exact).
collected_data[key].append(value) image_keys = [k for k in dataset.features if dataset.features[k]["dtype"] in ("image", "video")]
numeric_keys = [
k for k in dataset.features if dataset.features[k]["dtype"] not in ("image", "video", "string")
]
collected_data: dict[str, list] = {}
# Numeric features: every frame, read directly from the underlying table.
if numeric_keys:
numeric_cols = dataset.hf_dataset.select_columns(numeric_keys)[start_idx:end_idx]
for key in numeric_keys:
collected_data[key] = [torch.as_tensor(v) for v in numeric_cols[key]]
# Image/video features: decode only a sampled subset of frames.
if image_keys:
sampled_offsets = sample_indices(episode_len) if use_sampling else list(range(episode_len))
for offset in sampled_offsets:
item = dataset[start_idx + offset]
for key in image_keys:
if key in item:
collected_data.setdefault(key, []).append(item[key])
ep_stats = {} ep_stats = {}
for key, data_list in collected_data.items(): for key, data_list in collected_data.items():
@@ -131,11 +150,13 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
return ep_stats return ep_stats
def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dict]: def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]:
"""Compute quantile statistics for all episodes in the dataset. """Compute quantile statistics for all episodes in the dataset.
Args: Args:
dataset: The LeRobot dataset to compute statistics for dataset: The LeRobot dataset to compute statistics for
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
Returns: Returns:
Dictionary containing aggregated statistics with quantiles Dictionary containing aggregated statistics with quantiles
@@ -153,15 +174,15 @@ def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dic
if has_videos: if has_videos:
logging.info("Dataset contains video keys - using sequential processing for thread safety") logging.info("Dataset contains video keys - using sequential processing for thread safety")
for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"): for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
ep_stats = process_single_episode(dataset, episode_idx) ep_stats = process_single_episode(dataset, episode_idx, use_sampling)
episode_stats_list.append(ep_stats) episode_stats_list.append(ep_stats)
else: else:
logging.info("Dataset has no video keys - using parallel processing for better performance") logging.info("Dataset has no video keys - using parallel processing for better performance")
max_workers = min(dataset.num_episodes, 16) max_workers = min(dataset.num_episodes, int(os.environ.get("LEROBOT_STATS_MAX_WORKERS", 16)))
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
future_to_episode = { future_to_episode = {
executor.submit(process_single_episode, dataset, episode_idx): episode_idx executor.submit(process_single_episode, dataset, episode_idx, use_sampling): episode_idx
for episode_idx in range(dataset.num_episodes) for episode_idx in range(dataset.num_episodes)
} }
@@ -188,6 +209,7 @@ def augment_dataset_with_quantile_stats(
repo_id: str, repo_id: str,
root: str | Path | None = None, root: str | Path | None = None,
overwrite: bool = False, overwrite: bool = False,
use_sampling: bool = True,
) -> None: ) -> None:
"""Augment a dataset with quantile statistics if they are missing. """Augment a dataset with quantile statistics if they are missing.
@@ -195,6 +217,8 @@ def augment_dataset_with_quantile_stats(
repo_id: Repository ID of the dataset repo_id: Repository ID of the dataset
root: Local root directory for the dataset root: Local root directory for the dataset
overwrite: Overwrite existing quantile statistics if they already exist overwrite: Overwrite existing quantile statistics if they already exist
use_sampling: If True, sub-sample image/video frames per episode to bound
memory. If False, use every frame (exact, higher memory).
""" """
logging.info(f"Loading dataset: {repo_id}") logging.info(f"Loading dataset: {repo_id}")
dataset = LeRobotDataset( dataset = LeRobotDataset(
@@ -208,7 +232,7 @@ def augment_dataset_with_quantile_stats(
logging.info("Dataset does not contain quantile statistics. Computing them now...") logging.info("Dataset does not contain quantile statistics. Computing them now...")
new_stats = compute_quantile_stats_for_dataset(dataset) new_stats = compute_quantile_stats_for_dataset(dataset, use_sampling=use_sampling)
logging.info("Updating dataset metadata with new quantile statistics") logging.info("Updating dataset metadata with new quantile statistics")
dataset.meta.stats = new_stats dataset.meta.stats = new_stats
@@ -248,6 +272,14 @@ def main():
action="store_true", action="store_true",
help="Overwrite existing quantile statistics if they already exist", help="Overwrite existing quantile statistics if they already exist",
) )
parser.add_argument(
"--no-sampling",
action="store_true",
help=(
"Compute stats over every frame (exact, higher memory). By default, "
"image/video frames are sub-sampled per episode to bound memory."
),
)
args = parser.parse_args() args = parser.parse_args()
root = Path(args.root) if args.root else None root = Path(args.root) if args.root else None
@@ -258,6 +290,7 @@ def main():
repo_id=args.repo_id, repo_id=args.repo_id,
root=root, root=root,
overwrite=args.overwrite, overwrite=args.overwrite,
use_sampling=not args.no_sampling,
) )
+1 -1
View File
@@ -564,7 +564,7 @@ def eval_policy(
if seeds: if seeds:
all_seeds.extend(seeds) all_seeds.extend(seeds)
else: else:
all_seeds.append(None) all_seeds.extend([None] * env.num_envs)
# FIXME: episode_data is either None or it doesn't exist # FIXME: episode_data is either None or it doesn't exist
if return_episode_data: if return_episode_data:
+17 -13
View File
@@ -22,7 +22,8 @@ import dataclasses
import logging import logging
import sys import sys
import time import time
from contextlib import nullcontext from collections.abc import Iterator
from contextlib import contextmanager, nullcontext
from pprint import pformat from pprint import pformat
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any
@@ -76,6 +77,20 @@ else:
from .lerobot_eval import eval_policy_all from .lerobot_eval import eval_policy_all
@contextmanager
def _make_eval_envs(cfg: TrainPipelineConfig) -> Iterator[dict[str, dict[int, Any]]]:
"""Create evaluation environments for one run and always dispose of them."""
envs = make_env(
cfg.env,
n_envs=cfg.eval.batch_size,
use_async_envs=cfg.eval.use_async_envs,
)
try:
yield envs
finally:
close_envs(envs)
def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]: def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
"""Return worker-only DataLoader options, disabling them for single-process loading.""" """Return worker-only DataLoader options, disabling them for single-process loading."""
workers_enabled = cfg.num_workers > 0 workers_enabled = cfg.num_workers > 0
@@ -280,14 +295,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if not is_main_process: if not is_main_process:
dataset, eval_dataset = make_train_eval_datasets(cfg) dataset, eval_dataset = make_train_eval_datasets(cfg)
# Create environment used for evaluating checkpoints during training on simulation data.
# On real-world data, no need to create an environment as evaluations are done outside train.py,
# using the eval.py instead, with gym_dora environment and dora-rs.
eval_env = None
if cfg.env_eval_freq > 0 and cfg.env is not None and is_main_process:
logging.info("Creating env")
eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
if cfg.is_reward_model_training: if cfg.is_reward_model_training:
if is_main_process: if is_main_process:
logging.info("Creating reward model") logging.info("Creating reward model")
@@ -695,7 +702,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process: if is_main_process:
step_id = get_step_identifier(step, cfg.steps) step_id = get_step_identifier(step, cfg.steps)
logging.info(f"Eval policy at step {step}") logging.info(f"Eval policy at step {step}")
with torch.no_grad(), accelerator.autocast(): with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
eval_info = eval_policy_all( eval_info = eval_policy_all(
envs=eval_env, # dict[suite][task_id] -> vec_env envs=eval_env, # dict[suite][task_id] -> vec_env
policy=accelerator.unwrap_model(policy), policy=accelerator.unwrap_model(policy),
@@ -743,9 +750,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if is_main_process: if is_main_process:
progbar.close() progbar.close()
if eval_env:
close_envs(eval_env)
is_fsdp = accelerator.distributed_type == DistributedType.FSDP is_fsdp = accelerator.distributed_type == DistributedType.FSDP
model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None
if is_main_process: if is_main_process:
@@ -29,6 +29,12 @@ class SOLeaderConfig:
# Whether to use degrees for angles # Whether to use degrees for angles
use_degrees: bool = True use_degrees: bool = True
# Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
# return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
# at once, which otherwise aborts the teleoperation loop. Retries are immediate (no sleep) and only
# happen on failure, so the steady-state read cost is unchanged.
num_read_retries: int = 2
@TeleoperatorConfig.register_subclass("so101_leader") @TeleoperatorConfig.register_subclass("so101_leader")
@TeleoperatorConfig.register_subclass("so100_leader") @TeleoperatorConfig.register_subclass("so100_leader")
@@ -145,7 +145,7 @@ class SOLeader(Teleoperator):
@check_if_not_connected @check_if_not_connected
def get_action(self) -> dict[str, float]: def get_action(self) -> dict[str, float]:
start = time.perf_counter() start = time.perf_counter()
action = self.bus.sync_read("Present_Position") action = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
action = {f"{motor}.pos": val for motor, val in action.items()} action = {f"{motor}.pos": val for motor, val in action.items()}
dt_ms = (time.perf_counter() - start) * 1e3 dt_ms = (time.perf_counter() - start) * 1e3
logger.debug(f"{self} read action: {dt_ms:.1f}ms") logger.debug(f"{self} read action: {dt_ms:.1f}ms")
+16
View File
@@ -13,18 +13,34 @@
# limitations under the License. # limitations under the License.
from .transforms import ( from .transforms import (
CoarseDropout,
GammaCorrection,
GaussianNoise,
GaussianPatchBrightness,
ImageTransformConfig, ImageTransformConfig,
ImageTransforms, ImageTransforms,
ImageTransformsConfig, ImageTransformsConfig,
JPEGCompression,
MotionBlur,
PlanckianJitter,
RandomShadow,
RandomSubsetApply, RandomSubsetApply,
SharpnessJitter, SharpnessJitter,
make_transform_from_config, make_transform_from_config,
) )
__all__ = [ __all__ = [
"CoarseDropout",
"GammaCorrection",
"GaussianNoise",
"GaussianPatchBrightness",
"ImageTransformConfig", "ImageTransformConfig",
"ImageTransforms", "ImageTransforms",
"ImageTransformsConfig", "ImageTransformsConfig",
"JPEGCompression",
"MotionBlur",
"PlanckianJitter",
"RandomShadow",
"RandomSubsetApply", "RandomSubsetApply",
"SharpnessJitter", "SharpnessJitter",
"make_transform_from_config", "make_transform_from_config",
+471 -3
View File
@@ -14,11 +14,13 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import collections import collections
import math
from collections.abc import Callable, Sequence from collections.abc import Callable, Sequence
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Any from typing import Any
import torch import torch
from torchvision.io import decode_image, encode_jpeg
from torchvision.transforms import v2 from torchvision.transforms import v2
from torchvision.transforms.v2 import ( from torchvision.transforms.v2 import (
Transform, Transform,
@@ -144,6 +146,471 @@ class SharpnessJitter(Transform):
return self._call_kernel(F.adjust_sharpness, inpt, sharpness_factor=sharpness_factor) return self._call_kernel(F.adjust_sharpness, inpt, sharpness_factor=sharpness_factor)
class GaussianNoise(Transform):
"""Add Gaussian noise to simulate camera sensor noise.
Models readout noise from ADC quantization, which increases in low-light conditions.
Common in real-robot setups where wrist cameras operate in suboptimal lighting.
Args:
std: Range (min, max) for noise standard deviation in pixel-value scale (0-255).
"""
def __init__(self, std: float | Sequence[float] = (5.0, 25.0)) -> None:
super().__init__()
if isinstance(std, (int, float)):
self.std = (0.0, float(std))
elif isinstance(std, Sequence) and len(std) == 2:
self.std = (float(std[0]), float(std[1]))
else:
raise TypeError("std must be a number or a sequence with length 2.")
if not 0.0 <= self.std[0] <= self.std[1]:
raise ValueError(f"std must satisfy 0 <= min <= max, but got {self.std}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
return {
"std": torch.empty(1).uniform_(self.std[0], self.std[1]).item(),
"seed": torch.randint(0, torch.iinfo(torch.int64).max, ()).item(),
}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if isinstance(inpt, torch.Tensor) and inpt.is_floating_point():
generator = torch.Generator(device=inpt.device).manual_seed(params["seed"])
noise = torch.randn(inpt.shape, device=inpt.device, dtype=inpt.dtype, generator=generator)
return (inpt + noise * (params["std"] / 255.0)).clamp(0.0, 1.0)
return inpt
class MotionBlur(Transform):
"""Apply directional motion blur to simulate fast robot or object movement.
Generates a 1D averaging kernel along a random direction, applied via depthwise convolution.
Args:
kernel_size: An odd kernel size or a range containing at least one odd kernel size.
"""
def __init__(self, kernel_size: int | Sequence[int] = (3, 11)) -> None:
super().__init__()
if isinstance(kernel_size, int):
self.kernel_size = (kernel_size, kernel_size)
elif isinstance(kernel_size, Sequence) and len(kernel_size) == 2:
self.kernel_size = (int(kernel_size[0]), int(kernel_size[1]))
else:
raise TypeError("kernel_size must be an int or a sequence with length 2.")
if not 1 <= self.kernel_size[0] <= self.kernel_size[1]:
raise ValueError(f"kernel_size must satisfy 1 <= min <= max, but got {self.kernel_size}.")
self._first_odd_kernel_size = self.kernel_size[0] + (self.kernel_size[0] + 1) % 2
if self._first_odd_kernel_size > self.kernel_size[1]:
raise ValueError(f"kernel_size range must contain an odd value, but got {self.kernel_size}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
num_odd_sizes = (self.kernel_size[1] - self._first_odd_kernel_size) // 2 + 1
size_index = int(torch.randint(0, num_odd_sizes, ()).item())
ks = self._first_odd_kernel_size + 2 * size_index
angle = torch.empty(1).uniform_(0, 360).item()
return {"kernel_size": ks, "angle": angle}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt
if inpt.ndim < 3:
raise ValueError(f"MotionBlur expects [..., C, H, W] input, but got shape {inpt.shape}.")
kernel_size = params["kernel_size"]
radius = kernel_size // 2
angle = math.radians(params["angle"])
positions = torch.linspace(-radius, radius, kernel_size, device=inpt.device)
x_coords = (positions * math.cos(angle)).round().to(torch.long) + radius
y_coords = (positions * math.sin(angle)).round().to(torch.long) + radius
kernel = torch.zeros((kernel_size, kernel_size), device=inpt.device, dtype=inpt.dtype)
kernel[y_coords, x_coords] = 1
kernel /= kernel.sum()
channels, height, width = inpt.shape[-3:]
flat_input = inpt.reshape(-1, channels, height, width)
depthwise_kernel = kernel.expand(channels, 1, kernel_size, kernel_size)
padded = torch.nn.functional.pad(flat_input, (radius,) * 4, mode="replicate")
output = torch.nn.functional.conv2d(padded, depthwise_kernel, groups=channels)
return output.reshape(inpt.shape).clamp(0.0, 1.0)
class JPEGCompression(Transform):
"""Simulate JPEG compression artifacts (block artifacts, color banding).
Models quality degradation from video compression in network-streamed camera feeds.
Args:
quality: Range (min, max) for JPEG quality factor (lower = more artifacts).
"""
def __init__(self, quality: int | Sequence[int] = (15, 75)) -> None:
super().__init__()
if isinstance(quality, int):
self.quality = (quality, quality)
elif isinstance(quality, Sequence) and len(quality) == 2:
self.quality = (int(quality[0]), int(quality[1]))
else:
raise TypeError("quality must be an int or a sequence with length 2.")
if not 1 <= self.quality[0] <= self.quality[1] <= 100:
raise ValueError(f"quality must satisfy 1 <= min <= max <= 100, but got {self.quality}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
return {"quality": int(torch.randint(self.quality[0], self.quality[1] + 1, (1,)).item())}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt
if inpt.ndim < 3:
raise ValueError(f"JPEGCompression expects [..., C, H, W] input, but got shape {inpt.shape}.")
channels, height, width = inpt.shape[-3:]
if channels not in (1, 3):
raise ValueError(f"JPEGCompression expects 1 or 3 channels, but got {channels}.")
flat_input = inpt.reshape(-1, channels, height, width)
flat_uint8 = (flat_input.clamp(0.0, 1.0) * 255).round().to(torch.uint8).cpu()
decoded_frames = [
decode_image(encode_jpeg(frame, quality=params["quality"])) for frame in flat_uint8.unbind()
]
output = torch.stack(decoded_frames).to(device=inpt.device, dtype=inpt.dtype) / 255.0
return output.reshape(inpt.shape)
class GaussianPatchBrightness(Transform):
"""Apply spatially-varying brightness with Gaussian patches.
Simulates uneven overhead lighting, spotlights, and shadow patches commonly
encountered in real robot workspaces with multiple light sources.
Args:
num_patches: Range (min, max) for number of brightness patches.
sigma_range: Range for Gaussian sigma as fraction of image size.
factor_range: Range for brightness factor (< 1 darkens, > 1 brightens).
"""
def __init__(
self,
num_patches: int | Sequence[int] = (1, 4),
sigma_range: Sequence[float] = (0.05, 0.25),
factor_range: Sequence[float] = (0.4, 1.6),
) -> None:
super().__init__()
if isinstance(num_patches, int):
self.num_patches = (num_patches, num_patches)
elif isinstance(num_patches, Sequence) and len(num_patches) == 2:
self.num_patches = (int(num_patches[0]), int(num_patches[1]))
else:
raise TypeError("num_patches must be an int or a sequence with length 2.")
if not 1 <= self.num_patches[0] <= self.num_patches[1]:
raise ValueError(f"num_patches must satisfy 1 <= min <= max, but got {self.num_patches}.")
if not isinstance(sigma_range, Sequence) or len(sigma_range) != 2:
raise TypeError("sigma_range must be a sequence with length 2.")
self.sigma_range = (float(sigma_range[0]), float(sigma_range[1]))
if not 0.0 < self.sigma_range[0] <= self.sigma_range[1]:
raise ValueError(f"sigma_range must satisfy 0 < min <= max, but got {self.sigma_range}.")
if not isinstance(factor_range, Sequence) or len(factor_range) != 2:
raise TypeError("factor_range must be a sequence with length 2.")
self.factor_range = (float(factor_range[0]), float(factor_range[1]))
if not 0.0 <= self.factor_range[0] <= self.factor_range[1]:
raise ValueError(f"factor_range must satisfy 0 <= min <= max, but got {self.factor_range}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
n = int(torch.randint(self.num_patches[0], self.num_patches[1] + 1, (1,)).item())
return {
"centers": torch.rand(n, 2).tolist(),
"sigmas": torch.empty(n).uniform_(self.sigma_range[0], self.sigma_range[1]).tolist(),
"factors": torch.empty(n).uniform_(self.factor_range[0], self.factor_range[1]).tolist(),
}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt
h, w = inpt.shape[-2:]
mask = torch.ones(h, w, device=inpt.device, dtype=inpt.dtype)
grid_y = torch.linspace(0, 1, h, device=inpt.device, dtype=inpt.dtype)
grid_x = torch.linspace(0, 1, w, device=inpt.device, dtype=inpt.dtype)
yy, xx = torch.meshgrid(grid_y, grid_x, indexing="ij")
for (cy, cx), sigma, factor in zip(
params["centers"], params["sigmas"], params["factors"], strict=True
):
gauss = torch.exp(-((yy - cy) ** 2 + (xx - cx) ** 2) / (2 * sigma**2))
mask = mask * (1.0 + (factor - 1.0) * gauss)
broadcast_shape = (1,) * (inpt.ndim - 2) + (h, w)
return (inpt * mask.reshape(broadcast_shape)).clamp(0.0, 1.0)
class RandomShadow(Transform):
"""Add random vertical band shadow with smooth edges.
Simulates cast shadows from objects or people near the robot workspace.
Symmetric: randomly brightens or darkens to prevent BatchNorm stats shift.
Args:
opacity: Range (min, max) for shadow/highlight opacity.
"""
def __init__(self, opacity: float | Sequence[float] = (0.3, 0.6)) -> None:
super().__init__()
if isinstance(opacity, (int, float)):
self.opacity = (float(opacity), float(opacity))
elif isinstance(opacity, Sequence) and len(opacity) == 2:
self.opacity = (float(opacity[0]), float(opacity[1]))
else:
raise TypeError("opacity must be a number or a sequence with length 2.")
if not 0.0 <= self.opacity[0] <= self.opacity[1] <= 1.0:
raise ValueError(f"opacity must satisfy 0 <= min <= max <= 1, but got {self.opacity}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
return {
"opacity": torch.empty(1).uniform_(self.opacity[0], self.opacity[1]).item(),
"start": torch.rand(1).item(),
"width": torch.empty(1).uniform_(1 / 3, 2 / 3).item(),
"direction": -1.0 if torch.rand(1).item() < 0.5 else 1.0,
}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt
if inpt.ndim < 3:
raise ValueError(f"RandomShadow expects [..., C, H, W] input, but got shape {inpt.shape}.")
h, w = inpt.shape[-2:]
band_width = max(1, min(w, round(params["width"] * w)))
x_start = round(params["start"] * (w - band_width))
x_end = x_start + band_width
mask = torch.ones(h, w, device=inpt.device, dtype=inpt.dtype)
mask[:, x_start:x_end] = 1.0 + params["direction"] * params["opacity"]
smoothing_size = min(8, h, w)
if smoothing_size > 1:
batched_mask = mask[None, None]
small = torch.nn.functional.avg_pool2d(batched_mask, smoothing_size, stride=smoothing_size)
mask = torch.nn.functional.interpolate(small, size=(h, w), mode="bilinear", align_corners=False)[
0, 0
]
broadcast_shape = (1,) * (inpt.ndim - 2) + (h, w)
return (inpt * mask.reshape(broadcast_shape)).clamp(0.0, 1.0)
class CoarseDropout(Transform):
"""Drop random rectangular patches to simulate partial occlusion.
Models objects, hands, or cables passing through the camera field of view
during robot manipulation.
Args:
max_holes: Maximum number of rectangular patches to drop.
max_height_frac: Maximum patch height as fraction of image height.
max_width_frac: Maximum patch width as fraction of image width.
fill_value: Value to fill dropped regions with.
"""
def __init__(
self,
max_holes: int = 8,
max_height_frac: float = 0.07,
max_width_frac: float = 0.07,
fill_value: float = 0.0,
) -> None:
super().__init__()
if not isinstance(max_holes, int):
raise TypeError("max_holes must be an int.")
if max_holes < 1:
raise ValueError(f"max_holes must be at least 1, but got {max_holes}.")
if not 0.0 < max_height_frac <= 1.0:
raise ValueError(f"max_height_frac must be in (0, 1], but got {max_height_frac}.")
if not 0.0 < max_width_frac <= 1.0:
raise ValueError(f"max_width_frac must be in (0, 1], but got {max_width_frac}.")
if not 0.0 <= fill_value <= 1.0:
raise ValueError(f"fill_value must be in [0, 1], but got {fill_value}.")
self.max_holes = max_holes
self.max_height_frac = max_height_frac
self.max_width_frac = max_width_frac
self.fill_value = fill_value
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
n = int(torch.randint(1, self.max_holes + 1, (1,)).item())
sizes = torch.rand(n, 2)
sizes[:, 0] *= self.max_height_frac
sizes[:, 1] *= self.max_width_frac
return {"sizes": sizes.tolist(), "positions": torch.rand(n, 2).tolist()}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt
if inpt.ndim < 3:
raise ValueError(f"CoarseDropout expects [..., C, H, W] input, but got shape {inpt.shape}.")
h, w = inpt.shape[-2:]
result = inpt.clone()
for (height_frac, width_frac), (y_frac, x_frac) in zip(
params["sizes"], params["positions"], strict=True
):
hole_h = max(1, min(h, round(height_frac * h)))
hole_w = max(1, min(w, round(width_frac * w)))
y = round(y_frac * (h - hole_h))
x = round(x_frac * (w - hole_w))
result[..., y : y + hole_h, x : x + hole_w] = self.fill_value
return result
class GammaCorrection(Transform):
"""Apply random gamma correction to simulate exposure variation.
Models different camera auto-exposure settings and sensor response curves.
Uses log-symmetric sampling so brightening and darkening are equally likely,
preventing BatchNorm statistics shift.
Args:
gamma: Range (min, max) for gamma value. Values < 1 brighten, > 1 darken.
"""
def __init__(self, gamma: float | Sequence[float] = (0.5, 2.0)) -> None:
super().__init__()
if isinstance(gamma, (int, float)):
gamma = float(gamma)
if gamma <= 0:
raise ValueError(f"gamma must be positive, but got {gamma}.")
self.gamma = (min(gamma, 1.0 / gamma), max(gamma, 1.0 / gamma))
elif isinstance(gamma, Sequence) and len(gamma) == 2:
self.gamma = (float(gamma[0]), float(gamma[1]))
else:
raise TypeError("gamma must be a number or a sequence with length 2.")
if not 0.0 < self.gamma[0] <= self.gamma[1]:
raise ValueError(f"gamma must satisfy 0 < min <= max, but got {self.gamma}.")
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
log_lo = math.log(self.gamma[0])
log_hi = math.log(self.gamma[1])
gamma = math.exp(torch.empty(1).uniform_(log_lo, log_hi).item())
return {"gamma": gamma}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if isinstance(inpt, torch.Tensor) and inpt.is_floating_point():
return inpt.pow(params["gamma"]).clamp(0.0, 1.0)
return inpt
# From the paper authors' MIT-licensed reference implementation:
# https://github.com/TheZino/PlanckianJitter
_PLANCKIAN_BLACKBODY_COEFFICIENTS = (
(0.6743, 0.4029, 0.0013),
(0.6281, 0.4241, 0.1665),
(0.5919, 0.4372, 0.2513),
(0.5623, 0.4457, 0.3154),
(0.5376, 0.4515, 0.3672),
(0.5163, 0.4555, 0.4103),
(0.4979, 0.4584, 0.4468),
(0.4816, 0.4604, 0.4782),
(0.4672, 0.4619, 0.5053),
(0.4542, 0.4630, 0.5289),
(0.4426, 0.4638, 0.5497),
(0.4320, 0.4644, 0.5681),
(0.4223, 0.4648, 0.5844),
(0.4135, 0.4651, 0.5990),
(0.4054, 0.4653, 0.6121),
(0.3980, 0.4654, 0.6239),
(0.3911, 0.4655, 0.6346),
(0.3847, 0.4656, 0.6444),
(0.3787, 0.4656, 0.6532),
(0.3732, 0.4656, 0.6613),
(0.3680, 0.4655, 0.6688),
(0.3632, 0.4655, 0.6756),
(0.3586, 0.4655, 0.6820),
(0.3544, 0.4654, 0.6878),
(0.3503, 0.4653, 0.6933),
)
_PLANCKIAN_MIN_TEMPERATURE = 3_000
_PLANCKIAN_MAX_TEMPERATURE = 15_000
_PLANCKIAN_TEMPERATURE_STEP = 500
class PlanckianJitter(Transform):
"""Simulate color temperature shift along the Planckian locus.
Samples one black-body temperature and applies the corresponding correlated red
and blue channel scaling while preserving the green channel. Coefficients between
the tabulated 500 K intervals are linearly interpolated.
Reference: Zini et al., "Planckian Jitter", CVPR 2022 Workshop.
Args:
temperature: A fixed color temperature or range in Kelvin. Supported values
are between 3000 K and 15000 K.
"""
def __init__(self, temperature: int | Sequence[int] = (3_000, 15_000)) -> None:
super().__init__()
if isinstance(temperature, int):
self.temperature = (temperature, temperature)
elif isinstance(temperature, Sequence) and len(temperature) == 2:
self.temperature = (int(temperature[0]), int(temperature[1]))
else:
raise TypeError("temperature must be an int or a sequence with length 2.")
if not (
_PLANCKIAN_MIN_TEMPERATURE
<= self.temperature[0]
<= self.temperature[1]
<= _PLANCKIAN_MAX_TEMPERATURE
):
raise ValueError(
"temperature must satisfy "
f"{_PLANCKIAN_MIN_TEMPERATURE} <= min <= max <= {_PLANCKIAN_MAX_TEMPERATURE}, "
f"but got {self.temperature}."
)
def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
temperature = int(torch.randint(self.temperature[0], self.temperature[1] + 1, ()).item())
return {"temperature": temperature}
def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
return inpt
if inpt.ndim < 3 or inpt.shape[-3] != 3:
raise ValueError(f"PlanckianJitter expects [..., 3, H, W] input, but got shape {inpt.shape}.")
table_position = (params["temperature"] - _PLANCKIAN_MIN_TEMPERATURE) / _PLANCKIAN_TEMPERATURE_STEP
left_index = math.floor(table_position)
right_index = min(left_index + 1, len(_PLANCKIAN_BLACKBODY_COEFFICIENTS) - 1)
interpolation_weight = table_position - left_index
left = torch.tensor(
_PLANCKIAN_BLACKBODY_COEFFICIENTS[left_index],
device=inpt.device,
dtype=inpt.dtype,
)
right = torch.tensor(
_PLANCKIAN_BLACKBODY_COEFFICIENTS[right_index],
device=inpt.device,
dtype=inpt.dtype,
)
coefficients = torch.lerp(left, right, interpolation_weight)
scale = torch.stack(
(
coefficients[0] / coefficients[1],
coefficients.new_tensor(1.0),
coefficients[2] / coefficients[1],
)
)
broadcast_shape = (1,) * (inpt.ndim - 3) + (3, 1, 1)
return (inpt * scale.reshape(broadcast_shape)).clamp(0.0, 1.0)
_CUSTOM_TRANSFORMS: dict[str, type[Transform]] = {
"SharpnessJitter": SharpnessJitter,
"GaussianNoise": GaussianNoise,
"MotionBlur": MotionBlur,
"JPEGCompression": JPEGCompression,
"GaussianPatchBrightness": GaussianPatchBrightness,
"RandomShadow": RandomShadow,
"CoarseDropout": CoarseDropout,
"GammaCorrection": GammaCorrection,
"PlanckianJitter": PlanckianJitter,
}
@dataclass @dataclass
class ImageTransformConfig: class ImageTransformConfig:
""" """
@@ -216,16 +683,17 @@ class ImageTransformsConfig:
def make_transform_from_config(cfg: ImageTransformConfig) -> Transform: def make_transform_from_config(cfg: ImageTransformConfig) -> Transform:
if cfg.type == "SharpnessJitter": if cfg.type in _CUSTOM_TRANSFORMS:
return SharpnessJitter(**cfg.kwargs) return _CUSTOM_TRANSFORMS[cfg.type](**cfg.kwargs)
transform_cls = getattr(v2, cfg.type, None) transform_cls = getattr(v2, cfg.type, None)
if isinstance(transform_cls, type) and issubclass(transform_cls, Transform): if isinstance(transform_cls, type) and issubclass(transform_cls, Transform):
return transform_cls(**cfg.kwargs) return transform_cls(**cfg.kwargs)
valid_custom = ", ".join(sorted(_CUSTOM_TRANSFORMS.keys()))
raise ValueError( raise ValueError(
f"Transform '{cfg.type}' is not valid. It must be a class in " f"Transform '{cfg.type}' is not valid. It must be a class in "
f"torchvision.transforms.v2 or 'SharpnessJitter'." f"torchvision.transforms.v2 or one of: {valid_custom}."
) )
+7 -4
View File
@@ -133,10 +133,13 @@ def say(text: str, blocking: bool = False):
else: else:
raise RuntimeError("Unsupported operating system for text-to-speech.") raise RuntimeError("Unsupported operating system for text-to-speech.")
if blocking: try:
subprocess.run(cmd, check=True) if blocking:
else: subprocess.run(cmd, check=True, timeout=5)
subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0) else:
subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
except (FileNotFoundError, subprocess.TimeoutExpired) as e:
logging.warning("Text-to-speech command failed: %s | Error: %s", cmd, e)
def log_say(text: str, play_sounds: bool = True, blocking: bool = False): def log_say(text: str, play_sounds: bool = True, blocking: bool = False):
@@ -0,0 +1,104 @@
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# 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.
import numpy as np
import pytest
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
from lerobot.scripts.augment_dataset_quantile_stats import (
compute_quantile_stats_for_dataset,
has_quantile_stats,
)
def _numeric_keys(dataset):
return [k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string")]
def _image_keys(dataset):
return [k for k, v in dataset.features.items() if v["dtype"] in ("image", "video")]
def test_numeric_stats_are_unaffected_by_sampling(tmp_path, lerobot_dataset_factory):
"""Sampling only touches image/video frames; numeric features are read in
full either way, so their stats must be identical with and without sampling."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
numeric_keys = _numeric_keys(dataset)
assert numeric_keys, "fixture should expose numeric features"
for key in numeric_keys:
if key not in exact:
continue
for stat in ("mean", "std", "q01", "q50", "q99"):
if stat in exact[key]:
np.testing.assert_allclose(
sampled[key][stat],
exact[key][stat],
rtol=1e-6,
atol=1e-6,
err_msg=f"numeric feature '{key}' stat '{stat}' changed under sampling",
)
def test_image_sampling_reduces_data_but_keeps_stats_close(tmp_path, lerobot_dataset_factory):
"""For images, sampling should reduce the number of samples considered while
keeping the resulting statistics close to the exact ones."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
image_keys = _image_keys(dataset)
assert image_keys, "fixture should expose at least one image feature"
for key in image_keys:
# sampling actually looked at fewer pixels
assert sampled[key]["count"][0] < exact[key]["count"][0]
# but per-channel mean stays close
np.testing.assert_allclose(
sampled[key]["mean"],
exact[key]["mean"],
rtol=0.15,
err_msg=f"image feature '{key}' mean drifted too far under sampling",
)
def test_short_episodes_use_all_frames(tmp_path, lerobot_dataset_factory):
"""With episodes shorter than the sampling floor, sampling is a no-op and
must produce exactly the same stats as the exact path."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=40, use_videos=False
)
exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
for key in _image_keys(dataset):
assert sampled[key]["count"][0] == exact[key]["count"][0]
def test_quantile_stats_present_after_compute(tmp_path, lerobot_dataset_factory):
"""The computed stats should contain quantile keys for the dataset."""
dataset = lerobot_dataset_factory(
root=tmp_path / "ds", total_episodes=2, total_frames=200, use_videos=False
)
stats = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
assert has_quantile_stats(stats)
+158
View File
@@ -28,9 +28,17 @@ from lerobot.scripts.lerobot_imgtransform_viz import (
save_each_transform, save_each_transform,
) )
from lerobot.transforms import ( from lerobot.transforms import (
CoarseDropout,
GammaCorrection,
GaussianNoise,
GaussianPatchBrightness,
ImageTransformConfig, ImageTransformConfig,
ImageTransforms, ImageTransforms,
ImageTransformsConfig, ImageTransformsConfig,
JPEGCompression,
MotionBlur,
PlanckianJitter,
RandomShadow,
RandomSubsetApply, RandomSubsetApply,
SharpnessJitter, SharpnessJitter,
make_transform_from_config, make_transform_from_config,
@@ -455,3 +463,153 @@ def test_save_each_transform(img_tensor_factory, tmp_path):
assert (transform_dir / file_name).exists(), ( assert (transform_dir / file_name).exists(), (
f"{file_name} was not found in {transform} directory." f"{file_name} was not found in {transform} directory."
) )
# --- Tests for robotics-relevant augmentations ---
ROBOTICS_TRANSFORMS = [
("GaussianNoise", GaussianNoise, {"std": (5.0, 25.0)}),
("MotionBlur", MotionBlur, {"kernel_size": (3, 11)}),
("JPEGCompression", JPEGCompression, {"quality": (15, 75)}),
("GaussianPatchBrightness", GaussianPatchBrightness, {}),
("RandomShadow", RandomShadow, {"opacity": (0.3, 0.6)}),
("CoarseDropout", CoarseDropout, {"max_holes": 8}),
("GammaCorrection", GammaCorrection, {"gamma": (0.5, 2.0)}),
("PlanckianJitter", PlanckianJitter, {"temperature": (3_000, 15_000)}),
]
@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
def test_robotics_transform_shape_preserved(name, cls, kwargs, img_tensor_factory):
img = img_tensor_factory()
tf = cls(**kwargs)
out = tf(img)
assert out.shape == img.shape, f"{name} changed shape: {img.shape} -> {out.shape}"
@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
def test_robotics_transform_output_range(name, cls, kwargs, img_tensor_factory):
img = img_tensor_factory()
tf = cls(**kwargs)
out = tf(img)
assert out.min() >= -0.01, f"{name} min below range: {out.min():.4f}"
assert out.max() <= 1.01, f"{name} max above range: {out.max():.4f}"
@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
def test_robotics_transform_float_output(name, cls, kwargs, img_tensor_factory):
img = img_tensor_factory()
tf = cls(**kwargs)
out = tf(img)
assert out.is_floating_point(), f"{name} output dtype={out.dtype}"
@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
def test_robotics_transform_non_float_passthrough(name, cls, kwargs):
int_img = torch.randint(0, 255, (3, 32, 32), dtype=torch.uint8)
tf = cls(**kwargs)
out = tf(int_img)
assert torch.equal(out, int_img), f"{name} modified non-float input"
@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
def test_robotics_transform_via_config(name, cls, kwargs):
cfg = ImageTransformConfig(type=name, kwargs=kwargs)
tf = make_transform_from_config(cfg)
assert isinstance(tf, cls), f"Config produced {type(tf)}, expected {cls}"
def test_make_transform_error_message_includes_custom():
"""Error message should list all registered custom transforms."""
with pytest.raises(ValueError, match="GaussianNoise"):
make_transform_from_config(ImageTransformConfig(type="NonExistent"))
@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
@pytest.mark.parametrize("shape", [(4, 3, 32, 32), (2, 4, 3, 16, 16)])
def test_robotics_transform_supports_temporal_batches(name, cls, kwargs, shape):
img = torch.rand(shape)
out = cls(**kwargs)(img)
assert out.shape == img.shape, f"{name} changed shape: {img.shape} -> {out.shape}"
assert out.min() >= 0
assert out.max() <= 1
@pytest.mark.parametrize(
"cls,kwargs",
[
(GaussianNoise, {"std": (25.0, 25.0)}),
(MotionBlur, {"kernel_size": 5}),
(JPEGCompression, {"quality": 10}),
(
GaussianPatchBrightness,
{"num_patches": 1, "sigma_range": (0.2, 0.2), "factor_range": (0.5, 0.5)},
),
(RandomShadow, {"opacity": 0.5}),
(CoarseDropout, {"max_holes": 1, "fill_value": 0.0}),
(GammaCorrection, {"gamma": (2.0, 2.0)}),
(PlanckianJitter, {"temperature": 3_000}),
],
)
def test_robotics_transform_is_not_silent_noop(cls, kwargs):
img = torch.rand(3, 32, 32)
out = cls(**kwargs)(img)
assert not torch.equal(out, img)
@pytest.mark.parametrize(
"transform",
[
GaussianNoise(std=25),
RandomShadow(opacity=0.5),
CoarseDropout(max_holes=4),
],
)
def test_robotics_transform_random_params_are_reused(transform):
img = torch.rand(3, 32, 32)
params = transform.make_params([img])
torch.testing.assert_close(transform.transform(img, params), transform.transform(img, params))
def test_motion_blur_kernel_size_stays_in_configured_range():
transform = MotionBlur(kernel_size=(4, 10))
sampled_sizes = {transform.make_params([])["kernel_size"] for _ in range(100)}
assert sampled_sizes <= {5, 7, 9}
assert sampled_sizes
def test_gamma_correction_scalar_below_one_defines_symmetric_range():
transform = GammaCorrection(gamma=0.5)
assert transform.gamma == (0.5, 2.0)
assert transform(torch.rand(3, 8, 8)).shape == (3, 8, 8)
def test_planckian_jitter_uses_correlated_temperature_coefficients():
img = torch.full((2, 3, 8, 8), 0.25)
out = PlanckianJitter(temperature=3_000)(img)
torch.testing.assert_close(out[:, 1], img[:, 1])
assert torch.all(out[:, 0] > out[:, 1])
assert torch.all(out[:, 2] < out[:, 1])
def test_random_shadow_supports_small_images():
img = torch.rand(3, 7, 7)
assert RandomShadow()(img).shape == img.shape
@pytest.mark.parametrize(
"cls,kwargs",
[
(GaussianNoise, {"std": (-1.0, 1.0)}),
(MotionBlur, {"kernel_size": 4}),
(JPEGCompression, {"quality": (0, 75)}),
(GaussianPatchBrightness, {"sigma_range": (0.0, 0.25)}),
(RandomShadow, {"opacity": (0.3, 1.1)}),
(CoarseDropout, {"max_holes": 0}),
(GammaCorrection, {"gamma": 0.0}),
(PlanckianJitter, {"temperature": (2_000, 6_500)}),
],
)
def test_robotics_transform_rejects_invalid_config(cls, kwargs):
with pytest.raises(ValueError):
cls(**kwargs)
+13
View File
@@ -294,6 +294,19 @@ def test__sync_read(addr, length, ids_values, mock_motors, dummy_motors):
assert read_values == ids_values assert read_values == ids_values
def test__sync_read_retries_after_transient_failure(mock_motors, dummy_motors):
addr, length, ids_values = (10, 4, {1: 1337})
stub = mock_motors.build_sync_read_stub(addr, length, ids_values, num_invalid_try=1)
bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
bus.connect(handshake=False)
read_values, read_comm = bus._sync_read(addr, length, list(ids_values), num_retry=1)
assert read_comm == scs.COMM_SUCCESS
assert read_values == ids_values
assert mock_motors.stubs[stub].calls == 2
@pytest.mark.parametrize("raise_on_error", (True, False)) @pytest.mark.parametrize("raise_on_error", (True, False))
def test__sync_read_comm(raise_on_error, mock_motors, dummy_motors): def test__sync_read_comm(raise_on_error, mock_motors, dummy_motors):
addr, length, ids_values = (10, 4, {1: 1337}) addr, length, ids_values = (10, 4, {1: 1337})
+12 -8
View File
@@ -12,7 +12,6 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import sys
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import MagicMock from unittest.mock import MagicMock
@@ -47,18 +46,23 @@ def test_make_policy_keeps_peft_adapter_and_base_revisions_separate(monkeypatch)
peft_config_from_pretrained = MagicMock(return_value=peft_config) peft_config_from_pretrained = MagicMock(return_value=peft_config)
adapted_policy = torch.nn.Linear(1, 1) adapted_policy = torch.nn.Linear(1, 1)
peft_model_from_pretrained = MagicMock(return_value=adapted_policy) peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
monkeypatch.setitem( require_package = MagicMock()
sys.modules, monkeypatch.setattr(policy_factory, "require_package", require_package)
"peft", monkeypatch.setattr(
SimpleNamespace( policy_factory,
PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained), "PeftConfig",
PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained), SimpleNamespace(from_pretrained=peft_config_from_pretrained),
), )
monkeypatch.setattr(
policy_factory,
"PeftModel",
SimpleNamespace(from_pretrained=peft_model_from_pretrained),
) )
policy = policy_factory.make_policy(cfg, ds_meta=dataset_meta) policy = policy_factory.make_policy(cfg, ds_meta=dataset_meta)
assert policy is adapted_policy assert policy is adapted_policy
require_package.assert_called_once_with("peft", extra="peft")
peft_config_from_pretrained.assert_called_once_with( peft_config_from_pretrained.assert_called_once_with(
"user/adapter", "user/adapter",
revision="adapter-sha", revision="adapter-sha",
+23 -2
View File
@@ -49,7 +49,7 @@ def _make_bus_mock() -> MagicMock:
@pytest.fixture @pytest.fixture
def follower(): def follower(tmp_path):
bus_mock = _make_bus_mock() bus_mock = _make_bus_mock()
def _bus_side_effect(*_args, **kwargs): def _bus_side_effect(*_args, **kwargs):
@@ -71,7 +71,7 @@ def follower():
), ),
patch.object(SO100Follower, "configure", lambda self: None), patch.object(SO100Follower, "configure", lambda self: None),
): ):
cfg = SO100FollowerConfig(port="/dev/null") cfg = SO100FollowerConfig(port="/dev/null", calibration_dir=tmp_path)
robot = SO100Follower(cfg) robot = SO100Follower(cfg)
yield robot yield robot
if robot.is_connected: if robot.is_connected:
@@ -99,6 +99,27 @@ def test_get_observation(follower):
assert obs[f"{motor}.pos"] == idx assert obs[f"{motor}.pos"] == idx
def test_get_observation_uses_read_retries(follower):
# Feetech buses can intermittently fail a sync_read; the follower should forward the configured
# retry count so transient failures don't abort the control loop (see #3131).
follower.config.num_read_retries = 7
follower.connect()
follower.get_observation()
follower.bus.sync_read.assert_called_once_with("Present_Position", num_retry=7)
def test_send_action_uses_read_retries(follower):
follower.config.max_relative_target = 10.0
follower.config.num_read_retries = 7
follower.connect()
action = {f"{motor}.pos": value * 10 for value, motor in enumerate(follower.bus.motors, 1)}
follower.send_action(action)
follower.bus.sync_read.assert_called_once_with("Present_Position", num_retry=7)
def test_send_action(follower): def test_send_action(follower):
follower.connect() follower.connect()