mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-29 04:36:04 +00:00
Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 2de042690e | |||
| 124d03608c |
@@ -1,11 +0,0 @@
|
|||||||
version: 2
|
|
||||||
updates:
|
|
||||||
- package-ecosystem: "github-actions"
|
|
||||||
directory: "/"
|
|
||||||
schedule:
|
|
||||||
interval: "weekly"
|
|
||||||
cooldown:
|
|
||||||
default-days: 7
|
|
||||||
groups:
|
|
||||||
actions:
|
|
||||||
patterns: ["*"]
|
|
||||||
@@ -149,13 +149,14 @@ lerobot-rollout \
|
|||||||
|
|
||||||
Foot pedal input is also supported via `--strategy.input_device=pedal`. Configure pedal codes with `--strategy.pedal.*` flags.
|
Foot pedal input is also supported via `--strategy.input_device=pedal`. Configure pedal codes with `--strategy.pedal.*` flags.
|
||||||
|
|
||||||
| Flag | Description |
|
| Flag | Description |
|
||||||
| ------------------------------------ | ------------------------------------------------------- |
|
| ------------------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||||
| `--strategy.num_episodes` | Number of correction episodes to record (default: 10) |
|
| `--strategy.num_episodes` | Number of correction episodes to record (default: 10) |
|
||||||
| `--strategy.record_autonomous` | Record autonomous frames too (default: false) |
|
| `--strategy.record_autonomous` | Record autonomous frames too (default: false) |
|
||||||
| `--strategy.upload_every_n_episodes` | Push to Hub every N episodes (default: 5) |
|
| `--strategy.upload_every_n_episodes` | Push to Hub every N episodes (default: 5) |
|
||||||
| `--strategy.input_device` | Input device: `keyboard` or `pedal` (default: keyboard) |
|
| `--strategy.input_device` | Input device: `keyboard` or `pedal` (default: keyboard) |
|
||||||
| `--teleop.type` | **Required.** Teleoperator type |
|
| `--strategy.smooth_handover` | Smoothly hand control over at pause / correction start (default: true). Disable for clutch-style teleops that re-reference at the current robot pose on engage |
|
||||||
|
| `--teleop.type` | **Required.** Teleoperator type |
|
||||||
|
|
||||||
### Episodic (`--strategy.type=episodic`)
|
### Episodic (`--strategy.type=episodic`)
|
||||||
|
|
||||||
|
|||||||
@@ -519,13 +519,6 @@ def compute_episode_stats(
|
|||||||
if features[key]["dtype"] in {"string", "language"}:
|
if features[key]["dtype"] in {"string", "language"}:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# Features with a zero-width dimension contain no statistics-bearing
|
|
||||||
# values. Skip them like strings instead of letting
|
|
||||||
# get_feature_stats -> RunningQuantileStats.update reshape a size-0 array,
|
|
||||||
# which raises "ValueError: cannot reshape array of size 0".
|
|
||||||
if any(dim == 0 for dim in features[key].get("shape", ())):
|
|
||||||
continue
|
|
||||||
|
|
||||||
if features[key]["dtype"] in ["image", "video"]:
|
if features[key]["dtype"] in ["image", "video"]:
|
||||||
ep_ft_array = sample_images(data)
|
ep_ft_array = sample_images(data)
|
||||||
axes_to_reduce = (0, 2, 3)
|
axes_to_reduce = (0, 2, 3)
|
||||||
|
|||||||
@@ -172,23 +172,6 @@ class DatasetWriter:
|
|||||||
def _get_image_file_dir(self, episode_index: int, image_key: str) -> Path:
|
def _get_image_file_dir(self, episode_index: int, image_key: str) -> Path:
|
||||||
return self._get_image_file_path(episode_index, image_key, frame_index=0).parent
|
return self._get_image_file_path(episode_index, image_key, frame_index=0).parent
|
||||||
|
|
||||||
def _get_episode_buffer_index(self) -> int:
|
|
||||||
episode_index = self.episode_buffer["episode_index"]
|
|
||||||
# episode_index is `int` when freshly created, but becomes `np.ndarray` after
|
|
||||||
# save_episode() mutates the buffer. Handle both types here.
|
|
||||||
if isinstance(episode_index, np.ndarray):
|
|
||||||
episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
|
|
||||||
return int(episode_index)
|
|
||||||
|
|
||||||
def _delete_camera_frame_dirs(self, camera_keys: list[str]) -> None:
|
|
||||||
if self.image_writer is not None:
|
|
||||||
self._wait_image_writer()
|
|
||||||
episode_index = self._get_episode_buffer_index()
|
|
||||||
for camera_key in camera_keys:
|
|
||||||
img_dir = self._get_image_file_dir(episode_index, camera_key)
|
|
||||||
if img_dir.is_dir():
|
|
||||||
shutil.rmtree(img_dir)
|
|
||||||
|
|
||||||
def _save_image(
|
def _save_image(
|
||||||
self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1
|
self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1
|
||||||
) -> None:
|
) -> None:
|
||||||
@@ -386,9 +369,7 @@ class DatasetWriter:
|
|||||||
self._episodes_since_last_encoding = 0
|
self._episodes_since_last_encoding = 0
|
||||||
|
|
||||||
if episode_data is None:
|
if episode_data is None:
|
||||||
if len(self._meta.image_keys) > 0:
|
self.clear_episode_buffer(delete_images=len(self._meta.image_keys) > 0)
|
||||||
self._delete_camera_frame_dirs(self._meta.image_keys)
|
|
||||||
self.episode_buffer = self._create_episode_buffer()
|
|
||||||
|
|
||||||
def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None) -> None:
|
def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None) -> None:
|
||||||
"""Batch save videos for multiple episodes."""
|
"""Batch save videos for multiple episodes."""
|
||||||
@@ -580,10 +561,10 @@ class DatasetWriter:
|
|||||||
return metadata
|
return metadata
|
||||||
|
|
||||||
def clear_episode_buffer(self, delete_images: bool = True) -> None:
|
def clear_episode_buffer(self, delete_images: bool = True) -> None:
|
||||||
"""Discard the current episode buffer and optionally delete temp camera frames.
|
"""Discard the current episode buffer and optionally delete temp images.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
delete_images: If ``True``, remove temporary camera frame directories
|
delete_images: If ``True``, remove temporary image directories
|
||||||
written for the current episode.
|
written for the current episode.
|
||||||
"""
|
"""
|
||||||
# Cancel streaming encoder if active
|
# Cancel streaming encoder if active
|
||||||
@@ -591,7 +572,17 @@ class DatasetWriter:
|
|||||||
self._streaming_encoder.cancel_episode()
|
self._streaming_encoder.cancel_episode()
|
||||||
|
|
||||||
if delete_images:
|
if delete_images:
|
||||||
self._delete_camera_frame_dirs(self._meta.camera_keys)
|
if self.image_writer is not None:
|
||||||
|
self._wait_image_writer()
|
||||||
|
episode_index = self.episode_buffer["episode_index"]
|
||||||
|
# episode_index is `int` when freshly created, but becomes `np.ndarray` after
|
||||||
|
# save_episode() mutates the buffer. Handle both types here.
|
||||||
|
if isinstance(episode_index, np.ndarray):
|
||||||
|
episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
|
||||||
|
for cam_key in self._meta.image_keys:
|
||||||
|
img_dir = self._get_image_file_dir(episode_index, cam_key)
|
||||||
|
if img_dir.is_dir():
|
||||||
|
shutil.rmtree(img_dir)
|
||||||
|
|
||||||
self.episode_buffer = self._create_episode_buffer()
|
self.episode_buffer = self._create_episode_buffer()
|
||||||
|
|
||||||
|
|||||||
@@ -64,20 +64,12 @@ def get_hf_features_from_features(features: dict) -> datasets.Features:
|
|||||||
continue
|
continue
|
||||||
elif ft["dtype"] == "image":
|
elif ft["dtype"] == "image":
|
||||||
hf_features[key] = datasets.Image()
|
hf_features[key] = datasets.Image()
|
||||||
elif len(ft["shape"]) > 1 and any(dim == 0 for dim in ft["shape"]):
|
|
||||||
raise ValueError(
|
|
||||||
f"Multidimensional features with a zero-width dimension are not supported: "
|
|
||||||
f"'{key}' has shape {ft['shape']}. Only the one-dimensional shape (0,) is supported."
|
|
||||||
)
|
|
||||||
elif ft["shape"] == (1,):
|
elif ft["shape"] == (1,):
|
||||||
hf_features[key] = datasets.Value(dtype=ft["dtype"])
|
hf_features[key] = datasets.Value(dtype=ft["dtype"])
|
||||||
elif len(ft["shape"]) == 1:
|
elif len(ft["shape"]) == 1:
|
||||||
# A zero-width feature (shape=(0,)) has no fixed-size Arrow representation:
|
hf_features[key] = datasets.Sequence(
|
||||||
# pyarrow rejects a fixed-size list of length 0 ("list_size needs to be a
|
length=ft["shape"][0], feature=datasets.Value(dtype=ft["dtype"])
|
||||||
# strict positive integer"). Store it as a variable-length sequence
|
)
|
||||||
# (length=-1) so each per-frame value is simply an empty list.
|
|
||||||
seq_length = ft["shape"][0] if ft["shape"][0] > 0 else -1
|
|
||||||
hf_features[key] = datasets.Sequence(length=seq_length, feature=datasets.Value(dtype=ft["dtype"]))
|
|
||||||
elif len(ft["shape"]) == 2:
|
elif len(ft["shape"]) == 2:
|
||||||
hf_features[key] = datasets.Array2D(shape=ft["shape"], dtype=ft["dtype"])
|
hf_features[key] = datasets.Array2D(shape=ft["shape"], dtype=ft["dtype"])
|
||||||
elif len(ft["shape"]) == 3:
|
elif len(ft["shape"]) == 3:
|
||||||
|
|||||||
@@ -180,6 +180,14 @@ class DAggerStrategyConfig(RolloutStrategyConfig):
|
|||||||
# Target video file size in MB for episode rotation (record_autonomous
|
# Target video file size in MB for episode rotation (record_autonomous
|
||||||
# mode only). Defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB when None.
|
# mode only). Defaults to DEFAULT_VIDEO_FILE_SIZE_IN_MB when None.
|
||||||
target_video_file_size_mb: int | None = None
|
target_video_file_size_mb: int | None = None
|
||||||
|
# Whether to turn on or off the smooth handover behavior at phase transitions:
|
||||||
|
# the leader is driven to the follower position on pause (teleops with
|
||||||
|
# `send_feedback` capability), and the follower is slid to the teleop pose when
|
||||||
|
# a correction starts (non-actuated teleops). Disable for clutch-style
|
||||||
|
# teleoperators (e.g. VR controllers) that re-reference at the current robot
|
||||||
|
# pose on engage: the handover is already continuous there, and the blocking
|
||||||
|
# interpolation only delays the start of the correction.
|
||||||
|
smooth_handover: bool = True
|
||||||
input_device: str = "keyboard"
|
input_device: str = "keyboard"
|
||||||
keyboard: DAggerKeyboardConfig = field(default_factory=DAggerKeyboardConfig)
|
keyboard: DAggerKeyboardConfig = field(default_factory=DAggerKeyboardConfig)
|
||||||
pedal: DAggerPedalConfig = field(default_factory=DAggerPedalConfig)
|
pedal: DAggerPedalConfig = field(default_factory=DAggerPedalConfig)
|
||||||
|
|||||||
@@ -623,8 +623,8 @@ class DAggerStrategy(RolloutStrategy):
|
|||||||
# State-machine transition side-effects
|
# State-machine transition side-effects
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _apply_transition(
|
def _apply_transition(
|
||||||
|
self,
|
||||||
old_phase: DAggerPhase,
|
old_phase: DAggerPhase,
|
||||||
new_phase: DAggerPhase,
|
new_phase: DAggerPhase,
|
||||||
engine,
|
engine,
|
||||||
@@ -634,6 +634,10 @@ class DAggerStrategy(RolloutStrategy):
|
|||||||
) -> None:
|
) -> None:
|
||||||
"""Execute side-effects for a validated phase transition, including smooth handovers.
|
"""Execute side-effects for a validated phase transition, including smooth handovers.
|
||||||
|
|
||||||
|
The smooth handovers below can be disabled with
|
||||||
|
``--strategy.smooth_handover=false`` (useful for clutch-style teleops
|
||||||
|
that re-reference at the current robot pose on engage).
|
||||||
|
|
||||||
AUTONOMOUS -> PAUSED (actuated teleop):
|
AUTONOMOUS -> PAUSED (actuated teleop):
|
||||||
Pause the engine, then drive the leader arm to the follower's last
|
Pause the engine, then drive the leader arm to the follower's last
|
||||||
commanded position so the operator takes over without a jerk.
|
commanded position so the operator takes over without a jerk.
|
||||||
@@ -657,7 +661,7 @@ class DAggerStrategy(RolloutStrategy):
|
|||||||
logger.info("Pausing engine - robot holds position")
|
logger.info("Pausing engine - robot holds position")
|
||||||
engine.pause()
|
engine.pause()
|
||||||
|
|
||||||
if teleop_supports_feedback(teleop) and prev_action is not None:
|
if self.config.smooth_handover and teleop_supports_feedback(teleop) and prev_action is not None:
|
||||||
# TODO(Maxime): prev_action is in robot action key space (output of robot_action_processor).
|
# TODO(Maxime): prev_action is in robot action key space (output of robot_action_processor).
|
||||||
# send_feedback expects teleop feedback key space. For homogeneous setups (e.g. SO-101
|
# send_feedback expects teleop feedback key space. For homogeneous setups (e.g. SO-101
|
||||||
# leader + SO-101 follower) the keys are identical so this works. If the processor pipeline
|
# leader + SO-101 follower) the keys are identical so this works. If the processor pipeline
|
||||||
@@ -668,7 +672,11 @@ class DAggerStrategy(RolloutStrategy):
|
|||||||
|
|
||||||
elif old_phase == DAggerPhase.PAUSED and new_phase == DAggerPhase.CORRECTING:
|
elif old_phase == DAggerPhase.PAUSED and new_phase == DAggerPhase.CORRECTING:
|
||||||
logger.info("Entering correction mode - human teleop control")
|
logger.info("Entering correction mode - human teleop control")
|
||||||
if not teleop_supports_feedback(teleop) and prev_action is not None:
|
if (
|
||||||
|
self.config.smooth_handover
|
||||||
|
and not teleop_supports_feedback(teleop)
|
||||||
|
and prev_action is not None
|
||||||
|
):
|
||||||
logger.info("Smooth handover: sliding follower to teleop position")
|
logger.info("Smooth handover: sliding follower to teleop position")
|
||||||
obs = robot.get_observation()
|
obs = robot.get_observation()
|
||||||
teleop_action = teleop.get_action()
|
teleop_action = teleop.get_action()
|
||||||
|
|||||||
@@ -61,7 +61,6 @@ import pyarrow as pa
|
|||||||
import tqdm
|
import tqdm
|
||||||
from datasets import Dataset, Features, Image
|
from datasets import Dataset, Features, Image
|
||||||
from huggingface_hub import HfApi, snapshot_download
|
from huggingface_hub import HfApi, snapshot_download
|
||||||
from huggingface_hub.errors import RevisionNotFoundError
|
|
||||||
from requests import HTTPError
|
from requests import HTTPError
|
||||||
|
|
||||||
from lerobot.datasets import CODEBASE_VERSION, LeRobotDataset, aggregate_stats
|
from lerobot.datasets import CODEBASE_VERSION, LeRobotDataset, aggregate_stats
|
||||||
@@ -522,7 +521,7 @@ def convert_dataset(
|
|||||||
hub_api = HfApi()
|
hub_api = HfApi()
|
||||||
try:
|
try:
|
||||||
hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
|
hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
|
||||||
except (HTTPError, RevisionNotFoundError) as e:
|
except HTTPError as e:
|
||||||
print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
|
print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
|
||||||
pass
|
pass
|
||||||
hub_api.delete_files(
|
hub_api.delete_files(
|
||||||
|
|||||||
@@ -453,9 +453,6 @@ def eval_policy(
|
|||||||
raise exc from None
|
raise exc from None
|
||||||
|
|
||||||
start = time.time()
|
start = time.time()
|
||||||
# Preserve the mode for direct callers. eval_policy_all scopes the mode
|
|
||||||
# around all tasks so parallel evaluations cannot race with each other.
|
|
||||||
was_training = policy.training
|
|
||||||
policy.eval()
|
policy.eval()
|
||||||
|
|
||||||
# Determine how many batched rollouts we need to get n_episodes. Note that if n_episodes is not evenly
|
# Determine how many batched rollouts we need to get n_episodes. Note that if n_episodes is not evenly
|
||||||
@@ -677,8 +674,6 @@ def eval_policy(
|
|||||||
if save_predicted_video:
|
if save_predicted_video:
|
||||||
info["predicted_video_paths"] = predicted_video_paths
|
info["predicted_video_paths"] = predicted_video_paths
|
||||||
|
|
||||||
policy.train(was_training)
|
|
||||||
|
|
||||||
return info
|
return info
|
||||||
|
|
||||||
|
|
||||||
@@ -1015,48 +1010,40 @@ def eval_policy_all(
|
|||||||
recording_private=recording_private,
|
recording_private=recording_private,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Set the shared policy's mode before launching any workers. Restoring it
|
if max_parallel_tasks <= 1:
|
||||||
# inside individual tasks would let one task enable training mode while
|
prefetch_thread: threading.Thread | None = None
|
||||||
# another task is still evaluating.
|
for i, (task_group, task_id, env) in enumerate(tasks):
|
||||||
was_training = policy.training
|
if prefetch_thread is not None:
|
||||||
policy.eval()
|
prefetch_thread.join()
|
||||||
try:
|
prefetch_thread = None
|
||||||
if max_parallel_tasks <= 1:
|
|
||||||
prefetch_thread: threading.Thread | None = None
|
|
||||||
for i, (task_group, task_id, env) in enumerate(tasks):
|
|
||||||
if prefetch_thread is not None:
|
|
||||||
prefetch_thread.join()
|
|
||||||
prefetch_thread = None
|
|
||||||
|
|
||||||
|
try:
|
||||||
|
tg, tid, metrics = task_runner(task_group, task_id, env)
|
||||||
|
_accumulate_to(tg, metrics)
|
||||||
|
per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics})
|
||||||
|
finally:
|
||||||
|
env.close()
|
||||||
|
# Prefetch next task's workers *after* closing current env to prevent
|
||||||
|
# GPU memory overlap between consecutive tasks.
|
||||||
|
if i + 1 < len(tasks):
|
||||||
|
next_env = tasks[i + 1][2]
|
||||||
|
if hasattr(next_env, "_ensure"):
|
||||||
|
prefetch_thread = threading.Thread(target=next_env._ensure, daemon=True)
|
||||||
|
prefetch_thread.start()
|
||||||
|
else:
|
||||||
|
with cf.ThreadPoolExecutor(max_workers=max_parallel_tasks) as executor:
|
||||||
|
fut2meta = {}
|
||||||
|
for task_group, task_id, env in tasks:
|
||||||
|
fut = executor.submit(task_runner, task_group, task_id, env)
|
||||||
|
fut2meta[fut] = (task_group, task_id, env)
|
||||||
|
for fut in cf.as_completed(fut2meta):
|
||||||
|
tg, tid, env = fut2meta[fut]
|
||||||
try:
|
try:
|
||||||
tg, tid, metrics = task_runner(task_group, task_id, env)
|
tg, tid, metrics = fut.result()
|
||||||
_accumulate_to(tg, metrics)
|
_accumulate_to(tg, metrics)
|
||||||
per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics})
|
per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics})
|
||||||
finally:
|
finally:
|
||||||
env.close()
|
env.close()
|
||||||
# Prefetch next task's workers *after* closing current env to prevent
|
|
||||||
# GPU memory overlap between consecutive tasks.
|
|
||||||
if i + 1 < len(tasks):
|
|
||||||
next_env = tasks[i + 1][2]
|
|
||||||
if hasattr(next_env, "_ensure"):
|
|
||||||
prefetch_thread = threading.Thread(target=next_env._ensure, daemon=True)
|
|
||||||
prefetch_thread.start()
|
|
||||||
else:
|
|
||||||
with cf.ThreadPoolExecutor(max_workers=max_parallel_tasks) as executor:
|
|
||||||
fut2meta = {}
|
|
||||||
for task_group, task_id, env in tasks:
|
|
||||||
fut = executor.submit(task_runner, task_group, task_id, env)
|
|
||||||
fut2meta[fut] = (task_group, task_id, env)
|
|
||||||
for fut in cf.as_completed(fut2meta):
|
|
||||||
tg, tid, env = fut2meta[fut]
|
|
||||||
try:
|
|
||||||
tg, tid, metrics = fut.result()
|
|
||||||
_accumulate_to(tg, metrics)
|
|
||||||
per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics})
|
|
||||||
finally:
|
|
||||||
env.close()
|
|
||||||
finally:
|
|
||||||
policy.train(was_training)
|
|
||||||
|
|
||||||
# compute aggregated metrics helper (robust to lists/scalars)
|
# compute aggregated metrics helper (robust to lists/scalars)
|
||||||
def _agg_from_list(xs):
|
def _agg_from_list(xs):
|
||||||
|
|||||||
@@ -453,11 +453,9 @@ def record(
|
|||||||
encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize,
|
encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Connect the teleoperator before the robot so the robot isn't left idle (and possibly
|
robot.connect()
|
||||||
# tripping a firmware watchdog) during teleop init. Matches lerobot_teleoperate.py.
|
|
||||||
if teleop is not None:
|
if teleop is not None:
|
||||||
teleop.connect()
|
teleop.connect()
|
||||||
robot.connect()
|
|
||||||
|
|
||||||
listener, events = init_keyboard_listener()
|
listener, events = init_keyboard_listener()
|
||||||
|
|
||||||
|
|||||||
@@ -687,26 +687,6 @@ def test_compute_episode_stats_string_features_skipped():
|
|||||||
assert "q01" in stats["action"]
|
assert "q01" in stats["action"]
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("shape", [(0,), (0, 2), (2, 0), (1, 0, 2)])
|
|
||||||
def test_compute_episode_stats_zero_width_feature_skipped(shape):
|
|
||||||
"""Features with any zero-width dimension carry no values and are skipped."""
|
|
||||||
episode_data = {
|
|
||||||
"action": np.random.normal(0, 1, (100, 5)).astype(np.float32),
|
|
||||||
"target": np.zeros((100, *shape), dtype=np.float32),
|
|
||||||
}
|
|
||||||
features = {
|
|
||||||
"action": {"dtype": "float32", "shape": (5,)},
|
|
||||||
"target": {"dtype": "float32", "shape": shape},
|
|
||||||
}
|
|
||||||
|
|
||||||
stats = compute_episode_stats(episode_data, features)
|
|
||||||
|
|
||||||
# Zero-width features are skipped, just like strings; non-empty features are unaffected.
|
|
||||||
assert "target" not in stats
|
|
||||||
assert "action" in stats
|
|
||||||
assert "q01" in stats["action"]
|
|
||||||
|
|
||||||
|
|
||||||
def test_aggregate_feature_stats_with_quantiles():
|
def test_aggregate_feature_stats_with_quantiles():
|
||||||
"""Test aggregating feature stats that include quantiles."""
|
"""Test aggregating feature stats that include quantiles."""
|
||||||
stats_ft_list = [
|
stats_ft_list = [
|
||||||
|
|||||||
@@ -27,7 +27,6 @@ pytest.importorskip("datasets", reason="datasets is required (install lerobot[da
|
|||||||
|
|
||||||
from lerobot.configs import VideoEncoderConfig
|
from lerobot.configs import VideoEncoderConfig
|
||||||
from lerobot.datasets.dataset_writer import _encode_video_worker
|
from lerobot.datasets.dataset_writer import _encode_video_worker
|
||||||
from lerobot.datasets.feature_utils import get_hf_features_from_features
|
|
||||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||||
from lerobot.datasets.utils import DEFAULT_IMAGE_PATH
|
from lerobot.datasets.utils import DEFAULT_IMAGE_PATH
|
||||||
from tests.fixtures.constants import DEFAULT_FPS, DUMMY_REPO_ID
|
from tests.fixtures.constants import DEFAULT_FPS, DUMMY_REPO_ID
|
||||||
@@ -190,36 +189,6 @@ def test_save_multiple_episodes(tmp_path):
|
|||||||
assert dataset.meta.total_frames == total_frames
|
assert dataset.meta.total_frames == total_frames
|
||||||
|
|
||||||
|
|
||||||
def test_save_episode_with_zero_width_feature(tmp_path):
|
|
||||||
"""A one-dimensional empty numeric feature round-trips and has no statistics."""
|
|
||||||
features = {
|
|
||||||
**SIMPLE_FEATURES,
|
|
||||||
"target": {"dtype": "float32", "shape": (0,), "names": None},
|
|
||||||
}
|
|
||||||
root = tmp_path / "ds"
|
|
||||||
dataset = LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=features, root=root)
|
|
||||||
for _ in range(4):
|
|
||||||
dataset.add_frame(_make_frame(features))
|
|
||||||
dataset.save_episode()
|
|
||||||
dataset.finalize()
|
|
||||||
|
|
||||||
assert dataset.meta.total_episodes == 1
|
|
||||||
assert dataset.meta.total_frames == 4
|
|
||||||
|
|
||||||
reloaded = LeRobotDataset(repo_id=DUMMY_REPO_ID, root=root)
|
|
||||||
target = np.asarray(reloaded[0]["target"])
|
|
||||||
assert target.shape == (0,)
|
|
||||||
assert "target" not in (reloaded.meta.stats or {})
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("shape", [(0, 2), (2, 0), (1, 0, 2)])
|
|
||||||
def test_multidimensional_zero_width_feature_rejected(shape):
|
|
||||||
features = {"target": {"dtype": "float32", "shape": shape, "names": None}}
|
|
||||||
|
|
||||||
with pytest.raises(ValueError, match="Multidimensional features with a zero-width dimension"):
|
|
||||||
get_hf_features_from_features(features)
|
|
||||||
|
|
||||||
|
|
||||||
# ── clear / lifecycle ────────────────────────────────────────────────
|
# ── clear / lifecycle ────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
@@ -235,38 +204,6 @@ def test_clear_resets_buffer(tmp_path):
|
|||||||
assert dataset.writer.episode_buffer["size"] == 0
|
assert dataset.writer.episode_buffer["size"] == 0
|
||||||
|
|
||||||
|
|
||||||
def test_clear_removes_video_frame_staging_dir(tmp_path):
|
|
||||||
"""clear_episode_buffer() removes PNG staging dirs for video features."""
|
|
||||||
video_key = "observation.images.cam"
|
|
||||||
features = {
|
|
||||||
video_key: {
|
|
||||||
"dtype": "video",
|
|
||||||
"shape": (64, 96, 3),
|
|
||||||
"names": ["height", "width", "channels"],
|
|
||||||
},
|
|
||||||
"action": {"dtype": "float32", "shape": (2,), "names": None},
|
|
||||||
}
|
|
||||||
dataset = LeRobotDataset.create(
|
|
||||||
repo_id=DUMMY_REPO_ID,
|
|
||||||
fps=DEFAULT_FPS,
|
|
||||||
features=features,
|
|
||||||
root=tmp_path / "ds",
|
|
||||||
use_videos=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
dataset.add_frame(_make_frame(features))
|
|
||||||
video_staging_dir = (
|
|
||||||
dataset.root
|
|
||||||
/ Path(DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)).parent
|
|
||||||
)
|
|
||||||
assert video_staging_dir.is_dir()
|
|
||||||
|
|
||||||
dataset.clear_episode_buffer()
|
|
||||||
|
|
||||||
assert dataset.writer.episode_buffer["size"] == 0
|
|
||||||
assert not video_staging_dir.exists()
|
|
||||||
|
|
||||||
|
|
||||||
def test_finalize_is_idempotent(tmp_path):
|
def test_finalize_is_idempotent(tmp_path):
|
||||||
"""Calling finalize() twice does not raise."""
|
"""Calling finalize() twice does not raise."""
|
||||||
dataset = LeRobotDataset.create(
|
dataset = LeRobotDataset.create(
|
||||||
|
|||||||
Reference in New Issue
Block a user