mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-29 20:49:42 +00:00
Compare commits
14 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 6c60f4651b | |||
| 117216b29b | |||
| 2422d43fb1 | |||
| 53a5cffb4c | |||
| 2a3f40c673 | |||
| 2ddfb5a376 | |||
| d8b0a6b17f | |||
| 7d615acf9a | |||
| 09572babee | |||
| 35339d31e5 | |||
| f37be3edbe | |||
| 4d076845ac | |||
| 413972c812 | |||
| 0449aa02f6 |
@@ -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"]
|
||||||
|
|||||||
@@ -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 |
@@ -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:
|
||||||
|
try:
|
||||||
self._env.close()
|
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,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
@@ -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",
|
||||||
|
|||||||
@@ -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}."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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.")
|
||||||
|
|
||||||
|
try:
|
||||||
if blocking:
|
if blocking:
|
||||||
subprocess.run(cmd, check=True)
|
subprocess.run(cmd, check=True, timeout=5)
|
||||||
else:
|
else:
|
||||||
subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
|
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)
|
||||||
@@ -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)
|
||||||
|
|||||||
@@ -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,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",
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user