mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 13:09:40 +00:00
Compare commits
10 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 6cd4f1839b | |||
| dc0570586b | |||
| 36b8face98 | |||
| cd8984cc0a | |||
| b9ded9e761 | |||
| 185f3e1708 | |||
| e36783253a | |||
| 289e577fc7 | |||
| 9c32722eb9 | |||
| b49cb50e01 |
+14
-11
@@ -61,15 +61,19 @@ Full details in [`docs/source/so101.mdx`](./docs/source/so101.mdx) and [`docs/so
|
|||||||
**4.1 Install**
|
**4.1 Install**
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
pip install 'lerobot[feetech]' # SO-100/SO-101 motor stack
|
# uv (recommended — see AGENTS.md and CLAUDE.md)
|
||||||
# pip install 'lerobot[all]' # everything
|
uv sync --locked --extra feetech # SO-100/SO-101 motor stack
|
||||||
# pip install 'lerobot[aloha,pusht]' # specific features
|
# uv sync --locked --extra all # everything
|
||||||
# pip install 'lerobot[smolvla]' # add SmolVLA deps
|
# uv sync --locked --extra smolvla # add SmolVLA deps
|
||||||
git lfs install && git lfs pull
|
|
||||||
hf auth login # required to push datasets/policies
|
|
||||||
```
|
|
||||||
|
|
||||||
Contributors can alternatively use `uv sync --locked --extra feetech` (see `AGENTS.md`).
|
# pip (alternative, e.g. when not working from source)
|
||||||
|
# pip install 'lerobot[feetech]'
|
||||||
|
# pip install 'lerobot[all]'
|
||||||
|
# pip install 'lerobot[smolvla]'
|
||||||
|
|
||||||
|
git lfs install && git lfs pull
|
||||||
|
hf auth login # required to push datasets/policies
|
||||||
|
```
|
||||||
|
|
||||||
**4.2 Find USB ports** — run once per arm, unplug when prompted.
|
**4.2 Find USB ports** — run once per arm, unplug when prompted.
|
||||||
|
|
||||||
@@ -321,11 +325,10 @@ SmolVLA ships with `freeze_vision_encoder=True`. Unfreezing usually **improves p
|
|||||||
|
|
||||||
```bash
|
```bash
|
||||||
lerobot-train ... --policy.type=smolvla \
|
lerobot-train ... --policy.type=smolvla \
|
||||||
--policy.fine_tune_vision_encoder=true
|
--policy.freeze_vision_encoder=false \
|
||||||
|
--policy.train_expert_only=false
|
||||||
```
|
```
|
||||||
|
|
||||||
This selectively trains the vision encoder and connector while leaving the language model frozen. Their learning rate defaults to `0.1 × optimizer_lr`; adjust it with `--policy.vision_encoder_lr_multiplier` if needed.
|
|
||||||
|
|
||||||
### 7.7 Signals to stop / keep going
|
### 7.7 Signals to stop / keep going
|
||||||
|
|
||||||
- Train loss plateaus → stop, save a Hub checkpoint.
|
- Train loss plateaus → stop, save a Hub checkpoint.
|
||||||
|
|||||||
@@ -164,8 +164,8 @@ includes the range reported by the sensor. Requesting an unsupported control als
|
|||||||
Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
|
Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
|
||||||
require `use_rgb=True`.
|
require `use_rgb=True`.
|
||||||
|
|
||||||
On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual
|
Manual color controls require a dedicated RGB module. Cameras without one, such as the RealSense
|
||||||
exposure or gain also affects the depth stream.
|
D405, do not support them and raise an error at connection time.
|
||||||
|
|
||||||
</hfoption>
|
</hfoption>
|
||||||
</hfoptions>
|
</hfoptions>
|
||||||
|
|||||||
@@ -70,19 +70,6 @@ cd lerobot && lerobot-train \
|
|||||||
GPU allows it, as long as loading times remain short.
|
GPU allows it, as long as loading times remain short.
|
||||||
</Tip>
|
</Tip>
|
||||||
|
|
||||||
For tasks that require adapting visual features, such as distinguishing new colors or shapes, selectively
|
|
||||||
fine-tune the vision encoder and its connector:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
lerobot-train ... \
|
|
||||||
--policy.path=lerobot/smolvla_base \
|
|
||||||
--policy.fine_tune_vision_encoder=true
|
|
||||||
```
|
|
||||||
|
|
||||||
This keeps the language model frozen with the default `train_expert_only=true` setting and trains the vision
|
|
||||||
path at `0.1` times the main learning rate by default. Fine-tuning the vision encoder increases memory use and
|
|
||||||
can reduce the model's general visual knowledge, so enable it only when the frozen encoder is insufficient.
|
|
||||||
|
|
||||||
Fine-tuning is an art. For a complete overview of the options for finetuning, run
|
Fine-tuning is an art. For a complete overview of the options for finetuning, run
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
|
|||||||
@@ -365,11 +365,12 @@ class RealSenseCamera(Camera):
|
|||||||
return self._async_read(timeout_ms=10000, read_depth=read_depth)
|
return self._async_read(timeout_ms=10000, read_depth=read_depth)
|
||||||
|
|
||||||
def _get_color_sensor(self) -> "rs.sensor":
|
def _get_color_sensor(self) -> "rs.sensor":
|
||||||
"""Returns the sensor that controls the color stream.
|
"""Returns the dedicated "RGB Camera" sensor that controls the color stream.
|
||||||
|
|
||||||
Most RealSense cameras expose "RGB Camera" for color. The D405 has no
|
Manual color controls are only applied to a dedicated RGB module. Cameras
|
||||||
separate RGB module — its color stream comes from "Stereo Module".
|
without one (e.g. the D405, whose color stream comes from the shared
|
||||||
We try RGB Camera first, then fall back to Stereo Module.
|
"Stereo Module") are unsupported, so we never fall back to another sensor
|
||||||
|
to avoid altering the depth stream.
|
||||||
"""
|
"""
|
||||||
if self.rs_profile is None:
|
if self.rs_profile is None:
|
||||||
raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
|
raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
|
||||||
@@ -377,12 +378,14 @@ class RealSenseCamera(Camera):
|
|||||||
device = self.rs_profile.get_device()
|
device = self.rs_profile.get_device()
|
||||||
sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
|
sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
|
||||||
|
|
||||||
for name in ("RGB Camera", "Stereo Module"):
|
if "RGB Camera" in sensors:
|
||||||
if name in sensors:
|
return sensors["RGB Camera"]
|
||||||
return sensors[name]
|
|
||||||
|
|
||||||
available = list(sensors.keys())
|
available = list(sensors.keys())
|
||||||
raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}")
|
raise RuntimeError(
|
||||||
|
f"{self}: manual color controls require a dedicated 'RGB Camera' module, which this camera does not have. ",
|
||||||
|
f"Available sensors: {available}.",
|
||||||
|
)
|
||||||
|
|
||||||
def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
|
def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
|
||||||
"""Sets a sensor option, re-raising range errors with actionable diagnostics."""
|
"""Sets a sensor option, re-raising range errors with actionable diagnostics."""
|
||||||
|
|||||||
@@ -67,8 +67,6 @@ class SmolVLAConfig(PreTrainedConfig):
|
|||||||
|
|
||||||
# Finetuning settings
|
# Finetuning settings
|
||||||
freeze_vision_encoder: bool = True
|
freeze_vision_encoder: bool = True
|
||||||
fine_tune_vision_encoder: bool = False # Fine-tune vision + connector; takes priority over freezing.
|
|
||||||
vision_encoder_lr_multiplier: float = 0.1
|
|
||||||
train_expert_only: bool = True
|
train_expert_only: bool = True
|
||||||
train_state_proj: bool = True
|
train_state_proj: bool = True
|
||||||
|
|
||||||
@@ -112,12 +110,6 @@ class SmolVLAConfig(PreTrainedConfig):
|
|||||||
super().__post_init__()
|
super().__post_init__()
|
||||||
|
|
||||||
"""Input validation (not exhaustive)."""
|
"""Input validation (not exhaustive)."""
|
||||||
if self.fine_tune_vision_encoder:
|
|
||||||
self.freeze_vision_encoder = False
|
|
||||||
if self.vision_encoder_lr_multiplier <= 0:
|
|
||||||
raise ValueError(
|
|
||||||
f"`vision_encoder_lr_multiplier` must be positive, got {self.vision_encoder_lr_multiplier}."
|
|
||||||
)
|
|
||||||
if self.n_action_steps > self.chunk_size:
|
if self.n_action_steps > self.chunk_size:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"The chunk size is the upper bound for the number of action steps per model invocation. Got "
|
f"The chunk size is the upper bound for the number of action steps per model invocation. Got "
|
||||||
|
|||||||
@@ -186,27 +186,8 @@ class SmolVLAPolicy(PreTrainedPolicy):
|
|||||||
if model_value is not None:
|
if model_value is not None:
|
||||||
model_value.rtc_processor = self.rtc_processor
|
model_value.rtc_processor = self.rtc_processor
|
||||||
|
|
||||||
def get_optim_params(self):
|
def get_optim_params(self) -> dict:
|
||||||
if not self.config.fine_tune_vision_encoder:
|
return self.parameters()
|
||||||
return self.parameters()
|
|
||||||
|
|
||||||
vision_params = []
|
|
||||||
other_params = []
|
|
||||||
for name, param in self.named_parameters():
|
|
||||||
if not param.requires_grad:
|
|
||||||
continue
|
|
||||||
if ".vision_model." in name or ".connector." in name:
|
|
||||||
vision_params.append(param)
|
|
||||||
else:
|
|
||||||
other_params.append(param)
|
|
||||||
|
|
||||||
return [
|
|
||||||
{"params": other_params},
|
|
||||||
{
|
|
||||||
"params": vision_params,
|
|
||||||
"lr": self.config.optimizer_lr * self.config.vision_encoder_lr_multiplier,
|
|
||||||
},
|
|
||||||
]
|
|
||||||
|
|
||||||
def _get_action_chunk(
|
def _get_action_chunk(
|
||||||
self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs: Unpack[ActionSelectKwargs]
|
self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs: Unpack[ActionSelectKwargs]
|
||||||
@@ -512,7 +493,6 @@ class VLAFlowMatching(nn.Module):
|
|||||||
self.vlm_with_expert = SmolVLMWithExpertModel(
|
self.vlm_with_expert = SmolVLMWithExpertModel(
|
||||||
model_id=self.config.vlm_model_name,
|
model_id=self.config.vlm_model_name,
|
||||||
freeze_vision_encoder=self.config.freeze_vision_encoder,
|
freeze_vision_encoder=self.config.freeze_vision_encoder,
|
||||||
fine_tune_vision_encoder=self.config.fine_tune_vision_encoder,
|
|
||||||
train_expert_only=self.config.train_expert_only,
|
train_expert_only=self.config.train_expert_only,
|
||||||
load_vlm_weights=self.config.load_vlm_weights,
|
load_vlm_weights=self.config.load_vlm_weights,
|
||||||
attention_mode=self.config.attention_mode,
|
attention_mode=self.config.attention_mode,
|
||||||
|
|||||||
@@ -78,7 +78,6 @@ class SmolVLMWithExpertModel(nn.Module):
|
|||||||
load_vlm_weights: bool = True,
|
load_vlm_weights: bool = True,
|
||||||
train_expert_only: bool = True,
|
train_expert_only: bool = True,
|
||||||
freeze_vision_encoder: bool = False,
|
freeze_vision_encoder: bool = False,
|
||||||
fine_tune_vision_encoder: bool = False,
|
|
||||||
attention_mode: str = "self_attn",
|
attention_mode: str = "self_attn",
|
||||||
num_expert_layers: int = -1,
|
num_expert_layers: int = -1,
|
||||||
num_vlm_layers: int = -1,
|
num_vlm_layers: int = -1,
|
||||||
@@ -142,7 +141,6 @@ class SmolVLMWithExpertModel(nn.Module):
|
|||||||
self.num_key_value_heads = self.config.text_config.num_key_value_heads
|
self.num_key_value_heads = self.config.text_config.num_key_value_heads
|
||||||
|
|
||||||
self.freeze_vision_encoder = freeze_vision_encoder
|
self.freeze_vision_encoder = freeze_vision_encoder
|
||||||
self.fine_tune_vision_encoder = fine_tune_vision_encoder
|
|
||||||
self.train_expert_only = train_expert_only
|
self.train_expert_only = train_expert_only
|
||||||
self.attention_mode = attention_mode
|
self.attention_mode = attention_mode
|
||||||
self.expert_hidden_size = lm_expert_config.hidden_size
|
self.expert_hidden_size = lm_expert_config.hidden_size
|
||||||
@@ -152,6 +150,10 @@ class SmolVLMWithExpertModel(nn.Module):
|
|||||||
return self.vlm.model
|
return self.vlm.model
|
||||||
|
|
||||||
def set_requires_grad(self):
|
def set_requires_grad(self):
|
||||||
|
if self.freeze_vision_encoder:
|
||||||
|
self.get_vlm_model().vision_model.eval()
|
||||||
|
for params in self.get_vlm_model().vision_model.parameters():
|
||||||
|
params.requires_grad = False
|
||||||
if self.train_expert_only:
|
if self.train_expert_only:
|
||||||
self.vlm.eval()
|
self.vlm.eval()
|
||||||
for params in self.vlm.parameters():
|
for params in self.vlm.parameters():
|
||||||
@@ -174,18 +176,6 @@ class SmolVLMWithExpertModel(nn.Module):
|
|||||||
for name, params in self.vlm.named_parameters():
|
for name, params in self.vlm.named_parameters():
|
||||||
if any(k in name for k in frozen_layers):
|
if any(k in name for k in frozen_layers):
|
||||||
params.requires_grad = False
|
params.requires_grad = False
|
||||||
|
|
||||||
if self.freeze_vision_encoder:
|
|
||||||
self.get_vlm_model().vision_model.eval()
|
|
||||||
for params in self.get_vlm_model().vision_model.parameters():
|
|
||||||
params.requires_grad = False
|
|
||||||
|
|
||||||
if self.fine_tune_vision_encoder:
|
|
||||||
for params in self.get_vlm_model().vision_model.parameters():
|
|
||||||
params.requires_grad = True
|
|
||||||
for params in self.get_vlm_model().connector.parameters():
|
|
||||||
params.requires_grad = True
|
|
||||||
|
|
||||||
# To avoid unused params issue with distributed training
|
# To avoid unused params issue with distributed training
|
||||||
for name, params in self.lm_expert.named_parameters():
|
for name, params in self.lm_expert.named_parameters():
|
||||||
if "lm_head" in name:
|
if "lm_head" in name:
|
||||||
@@ -194,15 +184,11 @@ class SmolVLMWithExpertModel(nn.Module):
|
|||||||
def train(self, mode: bool = True):
|
def train(self, mode: bool = True):
|
||||||
super().train(mode)
|
super().train(mode)
|
||||||
|
|
||||||
if self.train_expert_only:
|
|
||||||
self.vlm.eval()
|
|
||||||
|
|
||||||
if self.freeze_vision_encoder:
|
if self.freeze_vision_encoder:
|
||||||
self.get_vlm_model().vision_model.eval()
|
self.get_vlm_model().vision_model.eval()
|
||||||
|
|
||||||
if self.fine_tune_vision_encoder:
|
if self.train_expert_only:
|
||||||
self.get_vlm_model().vision_model.train(mode)
|
self.vlm.eval()
|
||||||
self.get_vlm_model().connector.train(mode)
|
|
||||||
|
|
||||||
def embed_image(self, image: torch.Tensor):
|
def embed_image(self, image: torch.Tensor):
|
||||||
patch_attention_mask = None
|
patch_attention_mask = None
|
||||||
|
|||||||
@@ -18,7 +18,7 @@ import functools
|
|||||||
import threading
|
import threading
|
||||||
from collections.abc import Callable, Sequence
|
from collections.abc import Callable, Sequence
|
||||||
from contextlib import suppress
|
from contextlib import suppress
|
||||||
from typing import TypedDict
|
from typing import NotRequired, TypedDict
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
import torch.nn.functional as F # noqa: N812
|
import torch.nn.functional as F # noqa: N812
|
||||||
@@ -36,7 +36,7 @@ class BatchTransition(TypedDict):
|
|||||||
next_state: dict[str, torch.Tensor]
|
next_state: dict[str, torch.Tensor]
|
||||||
done: torch.Tensor
|
done: torch.Tensor
|
||||||
truncated: torch.Tensor
|
truncated: torch.Tensor
|
||||||
complementary_info: dict[str, torch.Tensor | float | int] | None = None
|
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
|
||||||
|
|
||||||
|
|
||||||
def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
|
def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
|
||||||
|
|||||||
@@ -28,7 +28,6 @@ lerobot-find-cameras
|
|||||||
# NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.
|
# NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
import concurrent.futures
|
|
||||||
import logging
|
import logging
|
||||||
import time
|
import time
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -133,7 +132,7 @@ def save_image(
|
|||||||
camera_identifier: str | int,
|
camera_identifier: str | int,
|
||||||
images_dir: Path,
|
images_dir: Path,
|
||||||
camera_type: str,
|
camera_type: str,
|
||||||
):
|
) -> None:
|
||||||
"""
|
"""
|
||||||
Saves a single image to disk using Pillow. Handles color conversion if necessary.
|
Saves a single image to disk using Pillow. Handles color conversion if necessary.
|
||||||
"""
|
"""
|
||||||
@@ -152,7 +151,7 @@ def save_image(
|
|||||||
logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}")
|
logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}")
|
||||||
|
|
||||||
|
|
||||||
def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
|
def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> dict[str, Any] | None:
|
||||||
"""Create and connect to a camera instance based on metadata."""
|
"""Create and connect to a camera instance based on metadata."""
|
||||||
cam_type = cam_meta.get("type")
|
cam_type = cam_meta.get("type")
|
||||||
cam_id = cam_meta.get("id")
|
cam_id = cam_meta.get("id")
|
||||||
@@ -165,12 +164,14 @@ def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
|
|||||||
cv_config = OpenCVCameraConfig(
|
cv_config = OpenCVCameraConfig(
|
||||||
index_or_path=cam_id,
|
index_or_path=cam_id,
|
||||||
color_mode=ColorMode.RGB,
|
color_mode=ColorMode.RGB,
|
||||||
|
warmup_s=warmup_s,
|
||||||
)
|
)
|
||||||
instance = OpenCVCamera(cv_config)
|
instance = OpenCVCamera(cv_config)
|
||||||
elif cam_type == "RealSense":
|
elif cam_type == "RealSense":
|
||||||
rs_config = RealSenseCameraConfig(
|
rs_config = RealSenseCameraConfig(
|
||||||
serial_number_or_name=cam_id,
|
serial_number_or_name=cam_id,
|
||||||
color_mode=ColorMode.RGB,
|
color_mode=ColorMode.RGB,
|
||||||
|
warmup_s=warmup_s,
|
||||||
)
|
)
|
||||||
instance = RealSenseCamera(rs_config)
|
instance = RealSenseCamera(rs_config)
|
||||||
else:
|
else:
|
||||||
@@ -188,9 +189,7 @@ def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
def process_camera_image(
|
def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_time: float) -> None:
|
||||||
cam_dict: dict[str, Any], output_dir: Path, current_time: float
|
|
||||||
) -> concurrent.futures.Future | None:
|
|
||||||
"""Capture and process an image from a single camera."""
|
"""Capture and process an image from a single camera."""
|
||||||
cam = cam_dict["instance"]
|
cam = cam_dict["instance"]
|
||||||
meta = cam_dict["meta"]
|
meta = cam_dict["meta"]
|
||||||
@@ -200,7 +199,7 @@ def process_camera_image(
|
|||||||
try:
|
try:
|
||||||
image_data = cam.read()
|
image_data = cam.read()
|
||||||
|
|
||||||
return save_image(
|
save_image(
|
||||||
image_data,
|
image_data,
|
||||||
cam_id_str,
|
cam_id_str,
|
||||||
output_dir,
|
output_dir,
|
||||||
@@ -215,21 +214,21 @@ def process_camera_image(
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
def cleanup_cameras(cameras_to_use: list[dict[str, Any]]):
|
def cleanup_camera(cam_dict: dict[str, Any]) -> None:
|
||||||
"""Disconnect all cameras."""
|
"""Disconnect all cameras."""
|
||||||
logger.info(f"Disconnecting {len(cameras_to_use)} cameras...")
|
logger.info(f"Disconnecting camera with ID {cam_dict['meta'].get('id')}...")
|
||||||
for cam_dict in cameras_to_use:
|
try:
|
||||||
try:
|
if cam_dict["instance"] and cam_dict["instance"].is_connected:
|
||||||
if cam_dict["instance"] and cam_dict["instance"].is_connected:
|
cam_dict["instance"].disconnect()
|
||||||
cam_dict["instance"].disconnect()
|
except Exception as e:
|
||||||
except Exception as e:
|
logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")
|
||||||
logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")
|
|
||||||
|
|
||||||
|
|
||||||
def save_images_from_all_cameras(
|
def save_images_from_all_cameras(
|
||||||
output_dir: Path,
|
output_dir: Path,
|
||||||
record_time_s: float = 2.0,
|
record_time_s: float = 2.0,
|
||||||
camera_type: str | None = None,
|
camera_type: str | None = None,
|
||||||
|
warmup_s: int = 1,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
Connects to detected cameras (optionally filtered by type) and saves images from each.
|
Connects to detected cameras (optionally filtered by type) and saves images from each.
|
||||||
@@ -240,6 +239,7 @@ def save_images_from_all_cameras(
|
|||||||
record_time_s: Duration in seconds to record images.
|
record_time_s: Duration in seconds to record images.
|
||||||
camera_type: Optional string to filter cameras ("realsense" or "opencv").
|
camera_type: Optional string to filter cameras ("realsense" or "opencv").
|
||||||
If None, uses all detected cameras.
|
If None, uses all detected cameras.
|
||||||
|
warmup_s: Duration in seconds to warmup camera before recording images.
|
||||||
"""
|
"""
|
||||||
output_dir.mkdir(parents=True, exist_ok=True)
|
output_dir.mkdir(parents=True, exist_ok=True)
|
||||||
logger.info(f"Saving images to {output_dir}")
|
logger.info(f"Saving images to {output_dir}")
|
||||||
@@ -249,40 +249,24 @@ def save_images_from_all_cameras(
|
|||||||
logger.warning("No cameras detected matching the criteria. Cannot save images.")
|
logger.warning("No cameras detected matching the criteria. Cannot save images.")
|
||||||
return
|
return
|
||||||
|
|
||||||
cameras_to_use = []
|
logger.info(
|
||||||
for cam_meta in all_camera_metadata:
|
f"Starting image capture for {record_time_s} seconds from {len(all_camera_metadata)} cameras."
|
||||||
camera_instance = create_camera_instance(cam_meta)
|
)
|
||||||
if camera_instance:
|
|
||||||
cameras_to_use.append(camera_instance)
|
|
||||||
|
|
||||||
if not cameras_to_use:
|
try:
|
||||||
logger.warning("No cameras could be connected. Aborting image save.")
|
for cam_meta in all_camera_metadata:
|
||||||
return
|
cam_dict = create_camera_instance(cam_meta, warmup_s=warmup_s)
|
||||||
|
if cam_dict is None:
|
||||||
logger.info(f"Starting image capture for {record_time_s} seconds from {len(cameras_to_use)} cameras.")
|
continue
|
||||||
start_time = time.perf_counter()
|
start_time = time.perf_counter()
|
||||||
|
|
||||||
with concurrent.futures.ThreadPoolExecutor(max_workers=len(cameras_to_use) * 2) as executor:
|
|
||||||
try:
|
|
||||||
while time.perf_counter() - start_time < record_time_s:
|
while time.perf_counter() - start_time < record_time_s:
|
||||||
futures = []
|
|
||||||
current_capture_time = time.perf_counter()
|
current_capture_time = time.perf_counter()
|
||||||
|
process_camera_image(cam_dict, output_dir, current_capture_time)
|
||||||
for cam_dict in cameras_to_use:
|
cleanup_camera(cam_dict)
|
||||||
future = process_camera_image(cam_dict, output_dir, current_capture_time)
|
except KeyboardInterrupt:
|
||||||
if future:
|
logger.info("Capture interrupted by user.")
|
||||||
futures.append(future)
|
finally:
|
||||||
|
print(f"Image capture finished. Images saved to {output_dir}")
|
||||||
if futures:
|
|
||||||
concurrent.futures.wait(futures)
|
|
||||||
|
|
||||||
except KeyboardInterrupt:
|
|
||||||
logger.info("Capture interrupted by user.")
|
|
||||||
finally:
|
|
||||||
print("\nFinalizing image saving...")
|
|
||||||
executor.shutdown(wait=True)
|
|
||||||
cleanup_cameras(cameras_to_use)
|
|
||||||
print(f"Image capture finished. Images saved to {output_dir}")
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
@@ -291,7 +275,6 @@ def main():
|
|||||||
parser = argparse.ArgumentParser(
|
parser = argparse.ArgumentParser(
|
||||||
description="Unified camera utility script for listing cameras and capturing images."
|
description="Unified camera utility script for listing cameras and capturing images."
|
||||||
)
|
)
|
||||||
|
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"camera_type",
|
"camera_type",
|
||||||
type=str,
|
type=str,
|
||||||
@@ -309,8 +292,14 @@ def main():
|
|||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--record-time-s",
|
"--record-time-s",
|
||||||
type=float,
|
type=float,
|
||||||
default=6.0,
|
default=2.0,
|
||||||
help="Time duration to attempt capturing frames. Default: 6 seconds.",
|
help="Time duration to attempt capturing frames. Default: 2 seconds.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--warmup-s",
|
||||||
|
type=int,
|
||||||
|
default=1,
|
||||||
|
help="Time duration to warmup camera before attempting to capture frames. Default: 1 second.",
|
||||||
)
|
)
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
save_images_from_all_cameras(**vars(args))
|
save_images_from_all_cameras(**vars(args))
|
||||||
|
|||||||
@@ -171,7 +171,13 @@ class IOSPhone(BasePhone, Teleoperator):
|
|||||||
# HEBI provides orientation in w, x, y, z format.
|
# HEBI provides orientation in w, x, y, z format.
|
||||||
# Scipy's Rotation expects x, y, z, w.
|
# Scipy's Rotation expects x, y, z, w.
|
||||||
quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]])) # wxyz to xyzw
|
quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]])) # wxyz to xyzw
|
||||||
rot = Rotation.from_quat(quat_xyzw)
|
# ARKit can emit zero/NaN quaternions before tracking is ready or on a
|
||||||
|
# dropped packet. Rotation.from_quat now rejects those; degrade the same
|
||||||
|
# way as a missing pose so teleop stays alive mid-session.
|
||||||
|
try:
|
||||||
|
rot = Rotation.from_quat(quat_xyzw)
|
||||||
|
except ValueError:
|
||||||
|
return False, None, None, None
|
||||||
pos = ar_pos - rot.apply(self.config.camera_offset)
|
pos = ar_pos - rot.apply(self.config.camera_offset)
|
||||||
return True, pos, rot, pose
|
return True, pos, rot, pose
|
||||||
|
|
||||||
|
|||||||
@@ -37,16 +37,25 @@ def auto_select_torch_device() -> torch.device:
|
|||||||
|
|
||||||
# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
|
# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
|
||||||
def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
|
def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
|
||||||
"""Given a string, return a torch.device with checks on whether the device is available."""
|
"""Given a string, return a torch.device with checks on whether the device is available.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If the requested device family is known but not available on
|
||||||
|
this machine (``AssertionError`` was previously used and is easy to
|
||||||
|
mistake for a programmer bug under ``python -O`` where asserts vanish).
|
||||||
|
"""
|
||||||
try_device = str(try_device)
|
try_device = str(try_device)
|
||||||
if try_device.startswith("cuda"):
|
if try_device.startswith("cuda"):
|
||||||
assert torch.cuda.is_available()
|
if not torch.cuda.is_available():
|
||||||
|
raise ValueError(f"Requested device {try_device!r} but CUDA is not available.")
|
||||||
device = torch.device(try_device)
|
device = torch.device(try_device)
|
||||||
elif try_device == "mps":
|
elif try_device == "mps":
|
||||||
assert torch.backends.mps.is_available()
|
if not torch.backends.mps.is_available():
|
||||||
|
raise ValueError("Requested device 'mps' but MPS is not available.")
|
||||||
device = torch.device("mps")
|
device = torch.device("mps")
|
||||||
elif try_device == "xpu":
|
elif try_device == "xpu":
|
||||||
assert torch.xpu.is_available()
|
if not torch.xpu.is_available():
|
||||||
|
raise ValueError("Requested device 'xpu' but XPU is not available.")
|
||||||
device = torch.device("xpu")
|
device = torch.device("xpu")
|
||||||
elif try_device == "cpu":
|
elif try_device == "cpu":
|
||||||
device = torch.device("cpu")
|
device = torch.device("cpu")
|
||||||
|
|||||||
@@ -32,21 +32,21 @@ def load_json(fpath: Path) -> Any:
|
|||||||
Returns:
|
Returns:
|
||||||
Any: The data loaded from the JSON file.
|
Any: The data loaded from the JSON file.
|
||||||
"""
|
"""
|
||||||
with open(fpath) as f:
|
with open(fpath, encoding="utf-8") as f:
|
||||||
return json.load(f)
|
return json.load(f)
|
||||||
|
|
||||||
|
|
||||||
def write_json(data: dict, fpath: Path) -> None:
|
def write_json(data: JsonLike, fpath: Path) -> None:
|
||||||
"""Write data to a JSON file.
|
"""Write JSON-serializable data to a file.
|
||||||
|
|
||||||
Creates parent directories if they don't exist.
|
Creates parent directories if they don't exist.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
data (dict): The dictionary to write.
|
data: JSON-serializable data to write.
|
||||||
fpath (Path): The path to the output JSON file.
|
fpath (Path): The path to the output JSON file.
|
||||||
"""
|
"""
|
||||||
fpath.parent.mkdir(exist_ok=True, parents=True)
|
fpath.parent.mkdir(exist_ok=True, parents=True)
|
||||||
with open(fpath, "w") as f:
|
with open(fpath, "w", encoding="utf-8") as f:
|
||||||
json.dump(data, f, indent=4, ensure_ascii=False)
|
json.dump(data, f, indent=4, ensure_ascii=False)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -30,6 +30,10 @@ def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: f
|
|||||||
"""
|
"""
|
||||||
if seconds <= 0:
|
if seconds <= 0:
|
||||||
return
|
return
|
||||||
|
if spin_threshold < 0:
|
||||||
|
raise ValueError(f"spin_threshold must be >= 0, got {spin_threshold}")
|
||||||
|
if sleep_margin < 0:
|
||||||
|
raise ValueError(f"sleep_margin must be >= 0, got {sleep_margin}")
|
||||||
|
|
||||||
system = platform.system()
|
system = platform.system()
|
||||||
# On macOS and Windows the scheduler / sleep granularity can make
|
# On macOS and Windows the scheduler / sleep granularity can make
|
||||||
|
|||||||
@@ -29,10 +29,13 @@ class Rotation:
|
|||||||
def __init__(self, quat: np.ndarray) -> None:
|
def __init__(self, quat: np.ndarray) -> None:
|
||||||
"""Initialize rotation from quaternion [x, y, z, w]."""
|
"""Initialize rotation from quaternion [x, y, z, w]."""
|
||||||
self._quat = np.asarray(quat, dtype=float)
|
self._quat = np.asarray(quat, dtype=float)
|
||||||
# Normalize quaternion
|
if self._quat.shape != (4,):
|
||||||
|
raise ValueError(f"Quaternion must have shape (4,), got {self._quat.shape}")
|
||||||
|
# Normalize quaternion. Reject the zero vector — it has no orientation.
|
||||||
norm = np.linalg.norm(self._quat)
|
norm = np.linalg.norm(self._quat)
|
||||||
if norm > 0:
|
if norm <= 0.0 or not np.isfinite(norm):
|
||||||
self._quat = self._quat / norm
|
raise ValueError(f"Quaternion must be a non-zero finite vector; got {self._quat} (norm={norm})")
|
||||||
|
self._quat = self._quat / norm
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_rotvec(cls, rotvec: np.ndarray) -> "Rotation":
|
def from_rotvec(cls, rotvec: np.ndarray) -> "Rotation":
|
||||||
|
|||||||
@@ -14,7 +14,7 @@
|
|||||||
# 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.
|
||||||
|
|
||||||
from typing import TypedDict
|
from typing import NotRequired, TypedDict
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
@@ -28,7 +28,7 @@ class Transition(TypedDict):
|
|||||||
next_state: dict[str, torch.Tensor]
|
next_state: dict[str, torch.Tensor]
|
||||||
done: bool
|
done: bool
|
||||||
truncated: bool
|
truncated: bool
|
||||||
complementary_info: dict[str, torch.Tensor | float | int] | None = None
|
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
|
||||||
|
|
||||||
|
|
||||||
def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition:
|
def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition:
|
||||||
|
|||||||
@@ -24,7 +24,6 @@ import sys
|
|||||||
import time
|
import time
|
||||||
from collections.abc import Iterator
|
from collections.abc import Iterator
|
||||||
from copy import copy, deepcopy
|
from copy import copy, deepcopy
|
||||||
from datetime import datetime
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from statistics import mean
|
from statistics import mean
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any
|
||||||
@@ -61,14 +60,16 @@ def init_logging(
|
|||||||
accelerator: Optional Accelerator instance (for multi-GPU detection)
|
accelerator: Optional Accelerator instance (for multi-GPU detection)
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def custom_format(record: logging.LogRecord) -> str:
|
class LeRobotFormatter(logging.Formatter):
|
||||||
dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
def format(self, record: logging.LogRecord) -> str:
|
||||||
fnameline = f"{record.pathname}:{record.lineno}"
|
record.lerobot_location = f"{record.pathname}:{record.lineno}"[-15:]
|
||||||
pid_str = f"[PID: {os.getpid()}] " if display_pid else ""
|
record.lerobot_pid = f"[PID: {os.getpid()}] " if display_pid else ""
|
||||||
return f"{record.levelname} {pid_str}{dt} {fnameline[-15:]:>15} {record.getMessage()}"
|
return super().format(record)
|
||||||
|
|
||||||
formatter = logging.Formatter()
|
formatter = LeRobotFormatter(
|
||||||
formatter.format = custom_format
|
"%(levelname)s %(lerobot_pid)s%(asctime)s %(lerobot_location)15s %(message)s",
|
||||||
|
datefmt="%Y-%m-%d %H:%M:%S",
|
||||||
|
)
|
||||||
|
|
||||||
logger = logging.getLogger()
|
logger = logging.getLogger()
|
||||||
logger.setLevel(logging.NOTSET)
|
logger.setLevel(logging.NOTSET)
|
||||||
|
|||||||
@@ -322,15 +322,16 @@ def test_get_color_sensor_prefers_rgb_camera():
|
|||||||
assert camera._get_color_sensor() is rgb
|
assert camera._get_color_sensor() is rgb
|
||||||
|
|
||||||
|
|
||||||
def test_get_color_sensor_falls_back_to_stereo_module():
|
def test_get_color_sensor_raises_without_dedicated_rgb_module():
|
||||||
"""D405 has no separate RGB module; color comes from Stereo Module."""
|
"""D405 has no separate RGB module; we refuse to touch the shared Stereo Module."""
|
||||||
config = RealSenseCameraConfig(serial_number_or_name="042")
|
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||||
camera = RealSenseCamera(config)
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
stereo = _make_mock_sensor("Stereo Module")
|
stereo = _make_mock_sensor("Stereo Module")
|
||||||
_attach_mock_color_sensor(camera, stereo)
|
_attach_mock_color_sensor(camera, stereo)
|
||||||
|
|
||||||
assert camera._get_color_sensor() is stereo
|
with pytest.raises(RuntimeError, match="dedicated 'RGB Camera' module"):
|
||||||
|
camera._get_color_sensor()
|
||||||
|
|
||||||
|
|
||||||
def test_get_color_sensor_raises_with_available_sensors():
|
def test_get_color_sensor_raises_with_available_sensors():
|
||||||
|
|||||||
@@ -0,0 +1,36 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
# Copyright 2026 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.
|
||||||
|
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from lerobot.utils.device_utils import get_safe_torch_device, is_torch_device_available
|
||||||
|
|
||||||
|
|
||||||
|
def test_cpu_always_available():
|
||||||
|
assert get_safe_torch_device("cpu") == torch.device("cpu")
|
||||||
|
assert is_torch_device_available("cpu")
|
||||||
|
|
||||||
|
|
||||||
|
def test_missing_cuda_raises_valueerror():
|
||||||
|
with patch("torch.cuda.is_available", return_value=False), pytest.raises(ValueError, match="CUDA"):
|
||||||
|
get_safe_torch_device("cuda")
|
||||||
|
|
||||||
|
|
||||||
|
def test_missing_mps_raises_valueerror():
|
||||||
|
with patch("torch.backends.mps.is_available", return_value=False), pytest.raises(ValueError, match="MPS"):
|
||||||
|
get_safe_torch_device("mps")
|
||||||
@@ -0,0 +1,46 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
# Copyright 2026 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
|
||||||
|
|
||||||
|
from lerobot.utils.rotation import Rotation
|
||||||
|
|
||||||
|
|
||||||
|
def test_zero_quaternion_rejected():
|
||||||
|
with pytest.raises(ValueError, match="non-zero"):
|
||||||
|
Rotation(np.zeros(4))
|
||||||
|
|
||||||
|
|
||||||
|
def test_non_finite_quaternion_rejected():
|
||||||
|
with pytest.raises(ValueError, match="non-zero|finite"):
|
||||||
|
Rotation(np.array([np.nan, 0.0, 0.0, 1.0]))
|
||||||
|
|
||||||
|
|
||||||
|
def test_wrong_shape_rejected():
|
||||||
|
with pytest.raises(ValueError, match="shape"):
|
||||||
|
Rotation(np.array([1.0, 0.0, 0.0]))
|
||||||
|
|
||||||
|
|
||||||
|
def test_identity_roundtrip():
|
||||||
|
r = Rotation.from_rotvec(np.zeros(3))
|
||||||
|
assert np.allclose(r.as_rotvec(), 0.0)
|
||||||
|
assert np.allclose(r.as_matrix(), np.eye(3))
|
||||||
|
|
||||||
|
|
||||||
|
def test_rotvec_roundtrip():
|
||||||
|
rotvec = np.array([0.1, -0.2, 0.3])
|
||||||
|
r = Rotation.from_rotvec(rotvec)
|
||||||
|
assert np.allclose(r.as_rotvec(), rotvec, atol=1e-6)
|
||||||
Reference in New Issue
Block a user