mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-29 04:36:04 +00:00
Compare commits
22 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 9cade7db18 | |||
| 775bbc19f4 | |||
| 2b78c7e582 | |||
| b101ff0e56 | |||
| ec2dbc1c98 | |||
| d526785e47 | |||
| 4af7c70664 | |||
| a855570097 | |||
| 167e22ba51 | |||
| 00c25c65c2 | |||
| 23f6d5dabd | |||
| 9b25b7fe0a | |||
| c1b6ea85d6 | |||
| ffe25afb8f | |||
| 95211b98f1 | |||
| 95256d766d | |||
| fd53716688 | |||
| a96540a2c4 | |||
| acd42b4d85 | |||
| bbeacfe57d | |||
| 801346e18c | |||
| ab87fd9764 |
@@ -136,6 +136,10 @@ config = RealSenseCameraConfig(
|
|||||||
height=480,
|
height=480,
|
||||||
color_mode=ColorMode.RGB,
|
color_mode=ColorMode.RGB,
|
||||||
use_depth=True,
|
use_depth=True,
|
||||||
|
# Optional fixed color controls. Omit them to leave the current sensor settings unchanged.
|
||||||
|
exposure=120,
|
||||||
|
gain=64,
|
||||||
|
white_balance=4600,
|
||||||
rotation=Cv2Rotation.NO_ROTATION
|
rotation=Cv2Rotation.NO_ROTATION
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -154,6 +158,15 @@ finally:
|
|||||||
```
|
```
|
||||||
<!-- prettier-ignore-end -->
|
<!-- prettier-ignore-end -->
|
||||||
|
|
||||||
|
Manual color controls disable the corresponding automatic exposure or white-balance mode. Their
|
||||||
|
supported ranges vary by camera model; an invalid value raises an error at connection time that
|
||||||
|
includes the range reported by the sensor. Requesting an unsupported control also raises an error.
|
||||||
|
Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
|
||||||
|
require `use_rgb=True`.
|
||||||
|
|
||||||
|
On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual
|
||||||
|
exposure or gain also affects the depth stream.
|
||||||
|
|
||||||
</hfoption>
|
</hfoption>
|
||||||
</hfoptions>
|
</hfoptions>
|
||||||
|
|
||||||
|
|||||||
+143
-7
@@ -1,9 +1,8 @@
|
|||||||
<div class="flex justify-center">
|
<div class="flex justify-center">
|
||||||
<a target="_blank" href="https://huggingface.co/lerobot">
|
<a target="_blank" href="https://huggingface.co/lerobot">
|
||||||
<img
|
<img
|
||||||
alt="HuggingFace Expert Acceleration Program"
|
alt="LeRobot, Hugging Face Robotics Library"
|
||||||
src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/lerobot-logo-thumbnail.png"
|
src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/lerobot-logo-thumbnail.png"
|
||||||
style="width: 100%"
|
|
||||||
></img>
|
></img>
|
||||||
</a>
|
</a>
|
||||||
</div>
|
</div>
|
||||||
@@ -12,12 +11,149 @@
|
|||||||
|
|
||||||
**State-of-the-art machine learning for real-world robotics**
|
**State-of-the-art machine learning for real-world robotics**
|
||||||
|
|
||||||
🤗 LeRobot aims to provide models, datasets, and tools for real-world robotics in PyTorch. The goal is to lower the barrier for entry to robotics so that everyone can contribute and benefit from sharing datasets and pretrained models.
|
🤗 LeRobot provides a hardware-agnostic, Python-native interface for controlling real robots - from affordable arms like the SO-101 to full humanoids. Plus the tools to record, store, and share the datasets they generate. Every dataset uses the standardized **LeRobotDataset** format (synchronized video + action/state data) and can be streamed directly from the [Hugging Face Hub](https://huggingface.co/lerobot).
|
||||||
|
|
||||||
🤗 LeRobot contains state-of-the-art approaches that have been shown to transfer to the real-world with a focus on imitation learning and reinforcement learning.
|
🤗 On top of that data, LeRobot implements state-of-the-art policies - from lightweight imitation-learning models like ACT to large vision-language-action models like π₀ and SmolVLA - all trainable, shareable, and deployable with the same handful of CLI commands.
|
||||||
|
|
||||||
🤗 LeRobot already provides a set of pretrained models, datasets with human collected demonstrations, and simulated environments so that everyone can get started.
|
The goal: lower the barrier to entry for robotics, so that everyone can contribute to, and benefit from, shared datasets and pretrained models.
|
||||||
|
|
||||||
🤗 LeRobot hosts pretrained models and datasets on the LeRobot HuggingFace page.
|
<div align="center" style="display: flex; justify-content: center; gap: 8px; flex-wrap: wrap; margin: 20px 0;">
|
||||||
|
<a href="https://discord.gg/s3KuuzsPFb" target="_blank">
|
||||||
|
<img alt="Discord" src="https://img.shields.io/badge/Discord-Join_the_Community-5865F2?style=flat&logo=discord&logoColor=white">
|
||||||
|
</a>
|
||||||
|
<a href="https://x.com/LeRobotHF" target="_blank">
|
||||||
|
<img alt="X (Twitter)" src="https://img.shields.io/badge/X-Follow_%40LeRobotHF-black?style=flat&logo=x&logoColor=white">
|
||||||
|
</a>
|
||||||
|
<a href="https://huggingface.co/lerobot" target="_blank">
|
||||||
|
<img alt="Hugging Face Hub" src="https://img.shields.io/badge/HF_Hub-Models_%26_Datasets-FFD21E?style=flat">
|
||||||
|
</a>
|
||||||
|
</div>
|
||||||
|
|
||||||
Join the LeRobot community on [Discord](https://discord.gg/s3KuuzsPFb)
|
<div align="center">
|
||||||
|
<img src="../../media/readme/robots_control_video.webp" width="640px" alt="Reachy 2 Demo">
|
||||||
|
</div>
|
||||||
|
|
||||||
|
## Get Started
|
||||||
|
|
||||||
|
New here? [Install LeRobot](./installation), then pick your path:
|
||||||
|
|
||||||
|
<div class="grid grid-cols-1 md:grid-cols-3 gap-4 my-6">
|
||||||
|
<div class="border dark:border-gray-700 rounded-lg p-4 shadow">
|
||||||
|
<div class="text-lg font-semibold mb-2">🔧 I have a robot</div>
|
||||||
|
<p class="text-gray-700 dark:text-gray-300 text-sm">
|
||||||
|
LeRobot supports a wide range of arms and mobile robots. Popular picks:
|
||||||
|
</p>
|
||||||
|
<ul class="text-gray-700 dark:text-gray-300 text-sm list-disc pl-5 mb-2">
|
||||||
|
<li>
|
||||||
|
<a href="./so101">SO-101</a> - our flagship, low-cost arm
|
||||||
|
</li>
|
||||||
|
<li>
|
||||||
|
<a href="./lekiwi">LeKiwi</a> - a mobile base with an arm on top
|
||||||
|
</li>
|
||||||
|
<li>
|
||||||
|
<a href="./koch">Koch v1.1</a> - a long-time community favorite
|
||||||
|
</li>
|
||||||
|
<li>
|
||||||
|
or find yours under <strong>Robots</strong> in the sidebar
|
||||||
|
</li>
|
||||||
|
</ul>
|
||||||
|
<p class="text-gray-700 dark:text-gray-300 text-sm">
|
||||||
|
Once it's assembled and calibrated, record a dataset and train your first
|
||||||
|
policy with the <a href="./il_robots">imitation learning tutorial</a>.
|
||||||
|
</p>
|
||||||
|
</div>
|
||||||
|
<div class="border dark:border-gray-700 rounded-lg p-4 shadow">
|
||||||
|
<div class="text-lg font-semibold mb-2">💻 No hardware yet</div>
|
||||||
|
<p class="text-gray-700 dark:text-gray-300 text-sm">
|
||||||
|
You can still train and evaluate policies without owning a robot:
|
||||||
|
</p>
|
||||||
|
<ul class="text-gray-700 dark:text-gray-300 text-sm list-disc pl-5 mb-2">
|
||||||
|
<li>
|
||||||
|
train on an existing{" "}
|
||||||
|
<a href="https://huggingface.co/datasets?other=LeRobot">
|
||||||
|
LeRobot dataset
|
||||||
|
</a>{" "}
|
||||||
|
from the Hub
|
||||||
|
</li>
|
||||||
|
<li>
|
||||||
|
evaluate in <a href="./envhub">simulation</a>, against benchmarks like
|
||||||
|
LIBERO or Meta-World
|
||||||
|
</li>
|
||||||
|
<li>
|
||||||
|
try the free <a href="./notebooks">Colab notebooks</a> - nothing to
|
||||||
|
install
|
||||||
|
</li>
|
||||||
|
</ul>
|
||||||
|
</div>
|
||||||
|
<div class="border dark:border-gray-700 rounded-lg p-4 shadow">
|
||||||
|
<div class="text-lg font-semibold mb-2">🤝 I want to contribute</div>
|
||||||
|
<p class="text-gray-700 dark:text-gray-300 text-sm">
|
||||||
|
Start with the <a href="./contributing">Contributing guide</a>, then{" "}
|
||||||
|
<a href="./bring_your_own_policies">add a new policy</a> or{" "}
|
||||||
|
<a href="./integrate_hardware">bring your own hardware</a>.
|
||||||
|
</p>
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
## Explore the Docs
|
||||||
|
|
||||||
|
<div class="grid grid-cols-1 md:grid-cols-3 gap-4 my-6">
|
||||||
|
<a
|
||||||
|
class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
|
||||||
|
href="./cheat-sheet"
|
||||||
|
>
|
||||||
|
<div class="font-semibold mb-1">📋 Cheat Sheet</div>
|
||||||
|
<p class="text-gray-700 dark:text-gray-300 text-sm">
|
||||||
|
Every LeRobot CLI command, copy-paste ready.
|
||||||
|
</p>
|
||||||
|
</a>
|
||||||
|
<a
|
||||||
|
class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
|
||||||
|
href="./hardware_guide"
|
||||||
|
>
|
||||||
|
<div class="font-semibold mb-1">🖥️ Compute & Hardware Guide</div>
|
||||||
|
<p class="text-gray-700 dark:text-gray-300 text-sm">
|
||||||
|
Which policy fits your GPU, and how long training takes.
|
||||||
|
</p>
|
||||||
|
</a>
|
||||||
|
<a
|
||||||
|
class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
|
||||||
|
href="./lerobot-dataset-v3"
|
||||||
|
>
|
||||||
|
<div class="font-semibold mb-1">🗂️ LeRobotDataset</div>
|
||||||
|
<p class="text-gray-700 dark:text-gray-300 text-sm">
|
||||||
|
Load, stream, and visualize robot datasets from the Hub.
|
||||||
|
</p>
|
||||||
|
</a>
|
||||||
|
<a
|
||||||
|
class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
|
||||||
|
href="./lelab"
|
||||||
|
>
|
||||||
|
<div class="font-semibold mb-1">🖱️ LeLab</div>
|
||||||
|
<p class="text-gray-700 dark:text-gray-300 text-sm">
|
||||||
|
A browser GUI for calibrating, recording, and training - no CLI required.
|
||||||
|
</p>
|
||||||
|
</a>
|
||||||
|
<a
|
||||||
|
class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
|
||||||
|
href="./act"
|
||||||
|
>
|
||||||
|
<div class="font-semibold mb-1">🧠 Policies</div>
|
||||||
|
<p class="text-gray-700 dark:text-gray-300 text-sm">
|
||||||
|
Start with ACT, our recommended first policy - or browse SmolVLA, π₀, and
|
||||||
|
more in the sidebar.
|
||||||
|
</p>
|
||||||
|
</a>
|
||||||
|
<a
|
||||||
|
class="!no-underline border dark:border-gray-700 rounded-lg p-4 shadow hover:shadow-lg"
|
||||||
|
href="./envhub"
|
||||||
|
>
|
||||||
|
<div class="font-semibold mb-1">🎮 Simulation & Benchmarks</div>
|
||||||
|
<p class="text-gray-700 dark:text-gray-300 text-sm">
|
||||||
|
Train and evaluate in simulated environments before touching real
|
||||||
|
hardware.
|
||||||
|
</p>
|
||||||
|
</a>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<!-- ASSET TODO: small monochrome icons matching HF's style would read cleaner than emoji in the cards above, once available. -->
|
||||||
|
|||||||
@@ -120,14 +120,22 @@ class OpenCVCamera(Camera):
|
|||||||
self.rotation: int | None = get_cv2_rotation(config.rotation)
|
self.rotation: int | None = get_cv2_rotation(config.rotation)
|
||||||
self.backend: int = config.backend
|
self.backend: int = config.backend
|
||||||
|
|
||||||
if self.height and self.width:
|
self.capture_width: int | None = None
|
||||||
self.capture_width, self.capture_height = self.width, self.height
|
self.capture_height: int | None = None
|
||||||
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
|
self._reset_connection_settings()
|
||||||
self.capture_width, self.capture_height = self.height, self.width
|
|
||||||
|
|
||||||
def __str__(self) -> str:
|
def __str__(self) -> str:
|
||||||
return f"{self.__class__.__name__}({self.index_or_path})"
|
return f"{self.__class__.__name__}({self.index_or_path})"
|
||||||
|
|
||||||
|
def _reset_connection_settings(self) -> None:
|
||||||
|
"""Restore settings that may have been auto-detected during a failed connection."""
|
||||||
|
self.fps = self.config.fps
|
||||||
|
self.width = self.config.width
|
||||||
|
self.height = self.config.height
|
||||||
|
self.capture_width, self.capture_height = self.width, self.height
|
||||||
|
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
|
||||||
|
self.capture_width, self.capture_height = self.height, self.width
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def is_connected(self) -> bool:
|
def is_connected(self) -> bool:
|
||||||
"""Checks if the camera is currently connected and opened."""
|
"""Checks if the camera is currently connected and opened."""
|
||||||
@@ -164,17 +172,25 @@ class OpenCVCamera(Camera):
|
|||||||
f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
|
f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
|
||||||
)
|
)
|
||||||
|
|
||||||
self._configure_capture_settings()
|
try:
|
||||||
self._start_read_thread()
|
self._configure_capture_settings()
|
||||||
|
self._start_read_thread()
|
||||||
|
|
||||||
if warmup and self.warmup_s > 0:
|
if warmup and self.warmup_s > 0:
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
while time.time() - start_time < self.warmup_s:
|
while time.time() - start_time < self.warmup_s:
|
||||||
self.async_read(timeout_ms=self.warmup_s * 1000)
|
self.async_read(timeout_ms=self.warmup_s * 1000)
|
||||||
time.sleep(0.1)
|
time.sleep(0.1)
|
||||||
with self.frame_lock:
|
with self.frame_lock:
|
||||||
if self.latest_frame is None:
|
if self.latest_frame is None:
|
||||||
raise ConnectionError(f"{self} failed to capture frames during warmup.")
|
raise ConnectionError(f"{self} failed to capture frames during warmup.")
|
||||||
|
except BaseException:
|
||||||
|
try:
|
||||||
|
self._cleanup_resources()
|
||||||
|
except Exception:
|
||||||
|
logger.exception(f"Failed to fully clean up {self} after connect() failed.")
|
||||||
|
self._reset_connection_settings()
|
||||||
|
raise
|
||||||
|
|
||||||
logger.info(f"{self} connected.")
|
logger.info(f"{self} connected.")
|
||||||
|
|
||||||
@@ -312,32 +328,36 @@ class OpenCVCamera(Camera):
|
|||||||
|
|
||||||
for target in targets_to_scan:
|
for target in targets_to_scan:
|
||||||
camera = cv2.VideoCapture(target)
|
camera = cv2.VideoCapture(target)
|
||||||
if camera.isOpened():
|
try:
|
||||||
default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
|
if camera.isOpened():
|
||||||
default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||||
default_fps = camera.get(cv2.CAP_PROP_FPS)
|
default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||||
default_format = camera.get(cv2.CAP_PROP_FORMAT)
|
default_fps = camera.get(cv2.CAP_PROP_FPS)
|
||||||
|
default_format = camera.get(cv2.CAP_PROP_FORMAT)
|
||||||
|
|
||||||
# Get FOURCC code and convert to string
|
# Get FOURCC code and convert to string
|
||||||
default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC)
|
default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC)
|
||||||
default_fourcc_code_int = int(default_fourcc_code)
|
default_fourcc_code_int = int(default_fourcc_code)
|
||||||
default_fourcc = "".join([chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)])
|
default_fourcc = "".join(
|
||||||
|
[chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)]
|
||||||
|
)
|
||||||
|
|
||||||
camera_info = {
|
camera_info = {
|
||||||
"name": f"OpenCV Camera @ {target}",
|
"name": f"OpenCV Camera @ {target}",
|
||||||
"type": "OpenCV",
|
"type": "OpenCV",
|
||||||
"id": target,
|
"id": target,
|
||||||
"backend_api": camera.getBackendName(),
|
"backend_api": camera.getBackendName(),
|
||||||
"default_stream_profile": {
|
"default_stream_profile": {
|
||||||
"format": default_format,
|
"format": default_format,
|
||||||
"fourcc": default_fourcc,
|
"fourcc": default_fourcc,
|
||||||
"width": default_width,
|
"width": default_width,
|
||||||
"height": default_height,
|
"height": default_height,
|
||||||
"fps": default_fps,
|
"fps": default_fps,
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
found_cameras_info.append(camera_info)
|
found_cameras_info.append(camera_info)
|
||||||
|
finally:
|
||||||
camera.release()
|
camera.release()
|
||||||
|
|
||||||
return found_cameras_info
|
return found_cameras_info
|
||||||
@@ -496,6 +516,26 @@ class OpenCVCamera(Camera):
|
|||||||
self.latest_timestamp = None
|
self.latest_timestamp = None
|
||||||
self.new_frame_event.clear()
|
self.new_frame_event.clear()
|
||||||
|
|
||||||
|
def _cleanup_resources(self) -> None:
|
||||||
|
"""Stop background reads and release the capture, including after partial setup."""
|
||||||
|
read_thread = self.thread
|
||||||
|
videocapture = self.videocapture
|
||||||
|
|
||||||
|
try:
|
||||||
|
self._stop_read_thread()
|
||||||
|
finally:
|
||||||
|
self.videocapture = None
|
||||||
|
try:
|
||||||
|
if videocapture is not None:
|
||||||
|
videocapture.release()
|
||||||
|
finally:
|
||||||
|
# Releasing the device may unblock a hardware read that outlived
|
||||||
|
# the first bounded join in _stop_read_thread().
|
||||||
|
if read_thread is not None and read_thread.is_alive():
|
||||||
|
read_thread.join(timeout=2.0)
|
||||||
|
if read_thread.is_alive(): # pragma: no cover
|
||||||
|
logger.warning(f"{self} read thread remained alive after releasing the capture.")
|
||||||
|
|
||||||
@check_if_not_connected
|
@check_if_not_connected
|
||||||
def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
|
def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
|
||||||
"""
|
"""
|
||||||
@@ -586,16 +626,6 @@ class OpenCVCamera(Camera):
|
|||||||
if not self.is_connected and self.thread is None:
|
if not self.is_connected and self.thread is None:
|
||||||
raise DeviceNotConnectedError(f"{self} not connected.")
|
raise DeviceNotConnectedError(f"{self} not connected.")
|
||||||
|
|
||||||
if self.thread is not None:
|
self._cleanup_resources()
|
||||||
self._stop_read_thread()
|
|
||||||
|
|
||||||
if self.videocapture is not None:
|
|
||||||
self.videocapture.release()
|
|
||||||
self.videocapture = None
|
|
||||||
|
|
||||||
with self.frame_lock:
|
|
||||||
self.latest_frame = None
|
|
||||||
self.latest_timestamp = None
|
|
||||||
self.new_frame_event.clear()
|
|
||||||
|
|
||||||
logger.info(f"{self} disconnected.")
|
logger.info(f"{self} disconnected.")
|
||||||
|
|||||||
@@ -121,6 +121,9 @@ class RealSenseCamera(Camera):
|
|||||||
|
|
||||||
self.config = config
|
self.config = config
|
||||||
|
|
||||||
|
self.width: int | None = config.width
|
||||||
|
self.height: int | None = config.height
|
||||||
|
|
||||||
if config.serial_number_or_name.isdigit():
|
if config.serial_number_or_name.isdigit():
|
||||||
self.serial_number = config.serial_number_or_name
|
self.serial_number = config.serial_number_or_name
|
||||||
else:
|
else:
|
||||||
@@ -131,6 +134,9 @@ class RealSenseCamera(Camera):
|
|||||||
self.use_rgb = config.use_rgb
|
self.use_rgb = config.use_rgb
|
||||||
self.use_depth = config.use_depth
|
self.use_depth = config.use_depth
|
||||||
self.warmup_s = config.warmup_s
|
self.warmup_s = config.warmup_s
|
||||||
|
self.exposure: int | None = config.exposure
|
||||||
|
self.gain: int | None = config.gain
|
||||||
|
self.white_balance: int | None = config.white_balance
|
||||||
|
|
||||||
self.rs_pipeline: rs.pipeline | None = None
|
self.rs_pipeline: rs.pipeline | None = None
|
||||||
self.rs_profile: rs.pipeline_profile | None = None
|
self.rs_profile: rs.pipeline_profile | None = None
|
||||||
@@ -145,14 +151,23 @@ class RealSenseCamera(Camera):
|
|||||||
|
|
||||||
self.rotation: int | None = get_cv2_rotation(config.rotation)
|
self.rotation: int | None = get_cv2_rotation(config.rotation)
|
||||||
|
|
||||||
if self.height and self.width:
|
self.capture_width: int | None = None
|
||||||
self.capture_width, self.capture_height = self.width, self.height
|
self.capture_height: int | None = None
|
||||||
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
|
self._reset_connection_settings()
|
||||||
self.capture_width, self.capture_height = self.height, self.width
|
|
||||||
|
|
||||||
def __str__(self) -> str:
|
def __str__(self) -> str:
|
||||||
return f"{self.__class__.__name__}({self.serial_number})"
|
return f"{self.__class__.__name__}({self.serial_number})"
|
||||||
|
|
||||||
|
def _reset_connection_settings(self) -> None:
|
||||||
|
"""Restore settings that may have been auto-detected during a failed connection."""
|
||||||
|
self.fps = self.config.fps
|
||||||
|
self.width = self.config.width
|
||||||
|
self.height = self.config.height
|
||||||
|
self.warmup_s = self.config.warmup_s
|
||||||
|
self.capture_width, self.capture_height = self.width, self.height
|
||||||
|
if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
|
||||||
|
self.capture_width, self.capture_height = self.height, self.width
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def is_connected(self) -> bool:
|
def is_connected(self) -> bool:
|
||||||
"""Checks if the camera pipeline is started and streams are active."""
|
"""Checks if the camera pipeline is started and streams are active."""
|
||||||
@@ -172,7 +187,8 @@ class RealSenseCamera(Camera):
|
|||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
DeviceAlreadyConnectedError: If the camera is already connected.
|
DeviceAlreadyConnectedError: If the camera is already connected.
|
||||||
ValueError: If the configuration is invalid (e.g., missing serial/name, name not unique).
|
ValueError: If the configuration is invalid, a requested sensor option is unsupported,
|
||||||
|
or a requested sensor value is invalid.
|
||||||
ConnectionError: If the camera is found but fails to start the pipeline or no RealSense devices are detected at all.
|
ConnectionError: If the camera is found but fails to start the pipeline or no RealSense devices are detected at all.
|
||||||
RuntimeError: If the pipeline starts but fails to apply requested settings.
|
RuntimeError: If the pipeline starts but fails to apply requested settings.
|
||||||
"""
|
"""
|
||||||
@@ -190,22 +206,31 @@ class RealSenseCamera(Camera):
|
|||||||
f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
|
f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
|
||||||
) from e
|
) from e
|
||||||
|
|
||||||
self._configure_capture_settings()
|
try:
|
||||||
self._start_read_thread()
|
self._configure_capture_settings()
|
||||||
|
self._configure_sensor_options()
|
||||||
|
self._start_read_thread()
|
||||||
|
|
||||||
# NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise.
|
# NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise.
|
||||||
self.warmup_s = max(self.warmup_s, 1)
|
self.warmup_s = max(self.warmup_s, 1)
|
||||||
|
|
||||||
warmup_read = self.async_read if self.use_rgb else self.async_read_depth
|
warmup_read = self.async_read if self.use_rgb else self.async_read_depth
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
while time.time() - start_time < self.warmup_s:
|
while time.time() - start_time < self.warmup_s:
|
||||||
warmup_read(timeout_ms=self.warmup_s * 1000)
|
warmup_read(timeout_ms=self.warmup_s * 1000)
|
||||||
time.sleep(0.1)
|
time.sleep(0.1)
|
||||||
with self.frame_lock:
|
with self.frame_lock:
|
||||||
if (self.use_rgb and self.latest_color_frame is None) or (
|
if (self.use_rgb and self.latest_color_frame is None) or (
|
||||||
self.use_depth and self.latest_depth_frame is None
|
self.use_depth and self.latest_depth_frame is None
|
||||||
):
|
):
|
||||||
raise ConnectionError(f"{self} failed to capture frames during warmup.")
|
raise ConnectionError(f"{self} failed to capture frames during warmup.")
|
||||||
|
except BaseException:
|
||||||
|
try:
|
||||||
|
self._cleanup_resources()
|
||||||
|
except Exception:
|
||||||
|
logger.exception(f"Failed to fully clean up {self} after connect() failed.")
|
||||||
|
self._reset_connection_settings()
|
||||||
|
raise
|
||||||
|
|
||||||
logger.info(f"{self} connected.")
|
logger.info(f"{self} connected.")
|
||||||
|
|
||||||
@@ -339,6 +364,111 @@ class RealSenseCamera(Camera):
|
|||||||
self.new_frame_event.clear()
|
self.new_frame_event.clear()
|
||||||
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":
|
||||||
|
"""Returns the sensor that controls the color stream.
|
||||||
|
|
||||||
|
Most RealSense cameras expose "RGB Camera" for color. The D405 has no
|
||||||
|
separate RGB module — its color stream comes from "Stereo Module".
|
||||||
|
We try RGB Camera first, then fall back to Stereo Module.
|
||||||
|
"""
|
||||||
|
if self.rs_profile is None:
|
||||||
|
raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
|
||||||
|
|
||||||
|
device = self.rs_profile.get_device()
|
||||||
|
sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
|
||||||
|
|
||||||
|
for name in ("RGB Camera", "Stereo Module"):
|
||||||
|
if name in sensors:
|
||||||
|
return sensors[name]
|
||||||
|
|
||||||
|
available = list(sensors.keys())
|
||||||
|
raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}")
|
||||||
|
|
||||||
|
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."""
|
||||||
|
try:
|
||||||
|
sensor.set_option(option, value)
|
||||||
|
except Exception as e:
|
||||||
|
range_info = ""
|
||||||
|
try:
|
||||||
|
option_range = sensor.get_option_range(option)
|
||||||
|
range_info = (
|
||||||
|
f" (supported range: min={option_range.min}, max={option_range.max}, "
|
||||||
|
f"step={option_range.step}, default={option_range.default})"
|
||||||
|
)
|
||||||
|
except Exception:
|
||||||
|
range_info = " (option range unavailable)"
|
||||||
|
raise ValueError(
|
||||||
|
f"{self}: failed to set {label} to {value}{range_info}. Original error: {e}"
|
||||||
|
) from e
|
||||||
|
|
||||||
|
def _configure_sensor_options(self) -> None:
|
||||||
|
"""Applies manual sensor options (exposure, gain, white balance) to the color sensor.
|
||||||
|
|
||||||
|
When exposure or gain is set, auto-exposure is disabled first. When white_balance
|
||||||
|
is set, auto white balance is disabled first. An omitted option is left unchanged,
|
||||||
|
and configuration is skipped entirely if all options are omitted.
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If the sensor does not support a requested option or a requested
|
||||||
|
value is invalid. Invalid-value errors include the option name, requested
|
||||||
|
value, and supported range when available.
|
||||||
|
"""
|
||||||
|
if self.exposure is None and self.gain is None and self.white_balance is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
color_sensor = self._get_color_sensor()
|
||||||
|
|
||||||
|
requested_options = (
|
||||||
|
(rs.option.exposure, self.exposure, "exposure"),
|
||||||
|
(rs.option.gain, self.gain, "gain"),
|
||||||
|
(rs.option.white_balance, self.white_balance, "white balance"),
|
||||||
|
)
|
||||||
|
unsupported_options = [
|
||||||
|
label
|
||||||
|
for option, value, label in requested_options
|
||||||
|
if value is not None and not color_sensor.supports(option)
|
||||||
|
]
|
||||||
|
if unsupported_options:
|
||||||
|
raise ValueError(
|
||||||
|
f"{self}: color sensor does not support requested manual options: {unsupported_options}."
|
||||||
|
)
|
||||||
|
|
||||||
|
manual_exposure_requested = self.exposure is not None or self.gain is not None
|
||||||
|
if manual_exposure_requested:
|
||||||
|
if color_sensor.supports(rs.option.enable_auto_exposure):
|
||||||
|
self._set_sensor_option(color_sensor, rs.option.enable_auto_exposure, 0, "auto-exposure")
|
||||||
|
logger.info(f"{self} auto-exposure disabled.")
|
||||||
|
else:
|
||||||
|
logger.warning(
|
||||||
|
f"{self} sensor does not support disabling auto-exposure; "
|
||||||
|
"applying manual exposure/gain directly."
|
||||||
|
)
|
||||||
|
|
||||||
|
if self.exposure is not None:
|
||||||
|
self._set_sensor_option(color_sensor, rs.option.exposure, self.exposure, "exposure")
|
||||||
|
logger.info(f"{self} exposure set to {self.exposure}.")
|
||||||
|
|
||||||
|
if self.gain is not None:
|
||||||
|
self._set_sensor_option(color_sensor, rs.option.gain, self.gain, "gain")
|
||||||
|
logger.info(f"{self} gain set to {self.gain}.")
|
||||||
|
|
||||||
|
if self.white_balance is not None:
|
||||||
|
if color_sensor.supports(rs.option.enable_auto_white_balance):
|
||||||
|
self._set_sensor_option(
|
||||||
|
color_sensor, rs.option.enable_auto_white_balance, 0, "auto white balance"
|
||||||
|
)
|
||||||
|
logger.info(f"{self} auto white balance disabled.")
|
||||||
|
else:
|
||||||
|
logger.warning(
|
||||||
|
f"{self} sensor does not support disabling auto white balance; "
|
||||||
|
"applying manual white balance directly."
|
||||||
|
)
|
||||||
|
self._set_sensor_option(
|
||||||
|
color_sensor, rs.option.white_balance, self.white_balance, "white balance"
|
||||||
|
)
|
||||||
|
logger.info(f"{self} white balance set to {self.white_balance}.")
|
||||||
|
|
||||||
@check_if_not_connected
|
@check_if_not_connected
|
||||||
def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
|
def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
|
||||||
"""
|
"""
|
||||||
@@ -541,6 +671,27 @@ class RealSenseCamera(Camera):
|
|||||||
self.latest_timestamp = None
|
self.latest_timestamp = None
|
||||||
self.new_frame_event.clear()
|
self.new_frame_event.clear()
|
||||||
|
|
||||||
|
def _cleanup_resources(self) -> None:
|
||||||
|
"""Stop background reads and stop the pipeline, including after partial setup."""
|
||||||
|
read_thread = self.thread
|
||||||
|
rs_pipeline = self.rs_pipeline
|
||||||
|
|
||||||
|
try:
|
||||||
|
self._stop_read_thread()
|
||||||
|
finally:
|
||||||
|
self.rs_pipeline = None
|
||||||
|
self.rs_profile = None
|
||||||
|
try:
|
||||||
|
if rs_pipeline is not None:
|
||||||
|
rs_pipeline.stop()
|
||||||
|
finally:
|
||||||
|
# Stopping the pipeline may unblock a hardware read that outlived
|
||||||
|
# the first bounded join in _stop_read_thread().
|
||||||
|
if read_thread is not None and read_thread.is_alive():
|
||||||
|
read_thread.join(timeout=2.0)
|
||||||
|
if read_thread.is_alive(): # pragma: no cover
|
||||||
|
logger.warning(f"{self} read thread remained alive after stopping the pipeline.")
|
||||||
|
|
||||||
def _async_read(self, timeout_ms: float, read_depth: bool = False) -> NDArray[Any]:
|
def _async_read(self, timeout_ms: float, read_depth: bool = False) -> NDArray[Any]:
|
||||||
"""Shared helper for :meth:`async_read`/:meth:`async_read_depth`: return the latest buffered frame."""
|
"""Shared helper for :meth:`async_read`/:meth:`async_read_depth`: return the latest buffered frame."""
|
||||||
if self.thread is None or not self.thread.is_alive():
|
if self.thread is None or not self.thread.is_alive():
|
||||||
@@ -684,18 +835,5 @@ class RealSenseCamera(Camera):
|
|||||||
f"Attempted to disconnect {self}, but it appears already disconnected."
|
f"Attempted to disconnect {self}, but it appears already disconnected."
|
||||||
)
|
)
|
||||||
|
|
||||||
if self.thread is not None:
|
self._cleanup_resources()
|
||||||
self._stop_read_thread()
|
|
||||||
|
|
||||||
if self.rs_pipeline is not None:
|
|
||||||
self.rs_pipeline.stop()
|
|
||||||
self.rs_pipeline = None
|
|
||||||
self.rs_profile = None
|
|
||||||
|
|
||||||
with self.frame_lock:
|
|
||||||
self.latest_color_frame = None
|
|
||||||
self.latest_depth_frame = None
|
|
||||||
self.latest_timestamp = None
|
|
||||||
self.new_frame_event.clear()
|
|
||||||
|
|
||||||
logger.info(f"{self} disconnected.")
|
logger.info(f"{self} disconnected.")
|
||||||
|
|||||||
@@ -46,6 +46,17 @@ class RealSenseCameraConfig(CameraConfig):
|
|||||||
use_depth: Whether to enable depth stream. Defaults to False.
|
use_depth: Whether to enable depth stream. Defaults to False.
|
||||||
rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
|
rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
|
||||||
warmup_s: Time reading frames before returning from connect (in seconds)
|
warmup_s: Time reading frames before returning from connect (in seconds)
|
||||||
|
exposure: Manual exposure value for the color sensor. When set, auto-exposure is
|
||||||
|
disabled and this fixed value is used. Valid ranges are camera-model specific
|
||||||
|
and reported if the value is rejected. Defaults to None (leave unchanged).
|
||||||
|
gain: Manual gain value for the color sensor. When set, auto-exposure is disabled
|
||||||
|
and this fixed gain is used, which also freezes exposure at its current value
|
||||||
|
when no exposure is configured. Valid ranges are camera-model specific and
|
||||||
|
reported if the value is rejected. Defaults to None (leave unchanged).
|
||||||
|
white_balance: Manual white balance value for the color sensor. When set, auto
|
||||||
|
white balance is disabled and this fixed value is used. Valid ranges are
|
||||||
|
camera-model specific and reported if the value is rejected. Defaults to None
|
||||||
|
(leave unchanged).
|
||||||
|
|
||||||
Note:
|
Note:
|
||||||
- Either name or serial_number must be specified.
|
- Either name or serial_number must be specified.
|
||||||
@@ -61,6 +72,9 @@ class RealSenseCameraConfig(CameraConfig):
|
|||||||
use_depth: bool = False
|
use_depth: bool = False
|
||||||
rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
|
rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
|
||||||
warmup_s: int = 1
|
warmup_s: int = 1
|
||||||
|
exposure: int | None = None
|
||||||
|
gain: int | None = None
|
||||||
|
white_balance: int | None = None
|
||||||
|
|
||||||
def __post_init__(self) -> None:
|
def __post_init__(self) -> None:
|
||||||
self.color_mode = ColorMode(self.color_mode)
|
self.color_mode = ColorMode(self.color_mode)
|
||||||
@@ -69,6 +83,18 @@ class RealSenseCameraConfig(CameraConfig):
|
|||||||
if not self.use_rgb and not self.use_depth:
|
if not self.use_rgb and not self.use_depth:
|
||||||
raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.")
|
raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.")
|
||||||
|
|
||||||
|
manual_color_options = {
|
||||||
|
"exposure": self.exposure,
|
||||||
|
"gain": self.gain,
|
||||||
|
"white_balance": self.white_balance,
|
||||||
|
}
|
||||||
|
configured_color_options = [name for name, value in manual_color_options.items() if value is not None]
|
||||||
|
if configured_color_options and not self.use_rgb:
|
||||||
|
raise ValueError(
|
||||||
|
"Manual color sensor options require `use_rgb=True`. "
|
||||||
|
f"Configured options: {configured_color_options}."
|
||||||
|
)
|
||||||
|
|
||||||
values = (self.fps, self.width, self.height)
|
values = (self.fps, self.width, self.height)
|
||||||
if any(v is not None for v in values) and any(v is None for v in values):
|
if any(v is not None for v in values) and any(v is None for v in values):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
|
|||||||
@@ -71,13 +71,19 @@ class DatasetRecordConfig:
|
|||||||
# Number of threads per encoder instance. None = auto (codec default).
|
# Number of threads per encoder instance. None = auto (codec default).
|
||||||
# Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
|
# Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
|
||||||
encoder_threads: int | None = None
|
encoder_threads: int | None = None
|
||||||
|
# Skip appending the date-time tag to repo_id, keeping the user-provided name as-is
|
||||||
|
# (e.g. self-managed versioned names intended for a later `lerobot-edit-dataset merge`).
|
||||||
|
no_stamp: bool = False
|
||||||
|
|
||||||
def stamp_repo_id(self) -> None:
|
def stamp_repo_id(self) -> None:
|
||||||
"""Append a date-time tag to ``repo_id`` so each recording session gets a unique name.
|
"""Append a date-time tag to ``repo_id`` so each recording session gets a unique name.
|
||||||
|
|
||||||
Must be called explicitly at dataset *creation* time — not on resume,
|
Must be called explicitly at dataset *creation* time — not on resume,
|
||||||
where the existing ``repo_id`` (already stamped) must be preserved.
|
where the existing ``repo_id`` (already stamped) must be preserved.
|
||||||
|
No-op when ``no_stamp`` is set, preserving a user-managed ``repo_id``.
|
||||||
"""
|
"""
|
||||||
|
if self.no_stamp:
|
||||||
|
return
|
||||||
if self.repo_id:
|
if self.repo_id:
|
||||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||||
self.repo_id = f"{self.repo_id}_{timestamp}"
|
self.repo_id = f"{self.repo_id}_{timestamp}"
|
||||||
|
|||||||
@@ -14,6 +14,7 @@
|
|||||||
import builtins
|
import builtins
|
||||||
import datetime as dt
|
import datetime as dt
|
||||||
import json
|
import json
|
||||||
|
import multiprocessing
|
||||||
import os
|
import os
|
||||||
import tempfile
|
import tempfile
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
@@ -101,6 +102,12 @@ class TrainPipelineConfig(HubMixin):
|
|||||||
batch_size: int = 8
|
batch_size: int = 8
|
||||||
prefetch_factor: int = 4
|
prefetch_factor: int = 4
|
||||||
persistent_workers: bool = True
|
persistent_workers: bool = True
|
||||||
|
# DataLoader worker start method. "spawn" is safer than "fork" with
|
||||||
|
# non-fork-safe libs (PyAV / torchcodec / ffmpeg), but adds some
|
||||||
|
# worker-startup time per run since workers re-import modules instead
|
||||||
|
# of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
|
||||||
|
# when appropriate, or set it to `null` to use Python's platform default.
|
||||||
|
dataloader_multiprocessing_context: str | None = "spawn"
|
||||||
steps: int = 100_000
|
steps: int = 100_000
|
||||||
# Run policy in the simulation environment every N steps to measure reward/success (0 = disabled).
|
# Run policy in the simulation environment every N steps to measure reward/success (0 = disabled).
|
||||||
env_eval_freq: int = 20_000
|
env_eval_freq: int = 20_000
|
||||||
@@ -212,6 +219,17 @@ class TrainPipelineConfig(HubMixin):
|
|||||||
self.reward_model.pretrained_path = str(policy_dir)
|
self.reward_model.pretrained_path = str(policy_dir)
|
||||||
|
|
||||||
def validate(self) -> None:
|
def validate(self) -> None:
|
||||||
|
available_contexts = multiprocessing.get_all_start_methods()
|
||||||
|
if (
|
||||||
|
self.dataloader_multiprocessing_context is not None
|
||||||
|
and self.dataloader_multiprocessing_context not in available_contexts
|
||||||
|
):
|
||||||
|
raise ValueError(
|
||||||
|
"`dataloader_multiprocessing_context` must be None or one of "
|
||||||
|
f"{available_contexts} on this platform, got "
|
||||||
|
f"{self.dataloader_multiprocessing_context!r}."
|
||||||
|
)
|
||||||
|
|
||||||
self._resolve_pretrained_from_cli()
|
self._resolve_pretrained_from_cli()
|
||||||
|
|
||||||
if self.policy is None and self.reward_model is None:
|
if self.policy is None and self.reward_model is None:
|
||||||
|
|||||||
@@ -172,6 +172,23 @@ 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:
|
||||||
@@ -369,7 +386,9 @@ class DatasetWriter:
|
|||||||
self._episodes_since_last_encoding = 0
|
self._episodes_since_last_encoding = 0
|
||||||
|
|
||||||
if episode_data is None:
|
if episode_data is None:
|
||||||
self.clear_episode_buffer(delete_images=len(self._meta.image_keys) > 0)
|
if 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."""
|
||||||
@@ -561,10 +580,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 images.
|
"""Discard the current episode buffer and optionally delete temp camera frames.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
delete_images: If ``True``, remove temporary image directories
|
delete_images: If ``True``, remove temporary camera frame directories
|
||||||
written for the current episode.
|
written for the current episode.
|
||||||
"""
|
"""
|
||||||
# Cancel streaming encoder if active
|
# Cancel streaming encoder if active
|
||||||
@@ -572,17 +591,7 @@ class DatasetWriter:
|
|||||||
self._streaming_encoder.cancel_episode()
|
self._streaming_encoder.cancel_episode()
|
||||||
|
|
||||||
if delete_images:
|
if delete_images:
|
||||||
if self.image_writer is not None:
|
self._delete_camera_frame_dirs(self._meta.camera_keys)
|
||||||
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()
|
||||||
|
|
||||||
|
|||||||
@@ -155,6 +155,7 @@ class MetaworldEnv(gym.Env):
|
|||||||
env.model.cam_pos[2] = [0.75, 0.075, 0.7]
|
env.model.cam_pos[2] = [0.75, 0.075, 0.7]
|
||||||
env.reset()
|
env.reset()
|
||||||
env._freeze_rand_vec = False # otherwise no randomization
|
env._freeze_rand_vec = False # otherwise no randomization
|
||||||
|
env.seeded_rand_vec = True # use seeded RNG so reset(seed=X) controls object positions
|
||||||
self._env = env
|
self._env = env
|
||||||
|
|
||||||
def render(self) -> np.ndarray:
|
def render(self) -> np.ndarray:
|
||||||
@@ -220,6 +221,8 @@ class MetaworldEnv(gym.Env):
|
|||||||
self._ensure_env()
|
self._ensure_env()
|
||||||
super().reset(seed=seed)
|
super().reset(seed=seed)
|
||||||
|
|
||||||
|
if seed is not None:
|
||||||
|
self._env.seed(seed)
|
||||||
raw_obs, info = self._env.reset(seed=seed)
|
raw_obs, info = self._env.reset(seed=seed)
|
||||||
|
|
||||||
observation = self._format_raw_obs(raw_obs)
|
observation = self._format_raw_obs(raw_obs)
|
||||||
|
|||||||
@@ -122,6 +122,9 @@ MODEL_ENCODING_TABLE = {
|
|||||||
"xm430-w350": X_SERIES_ENCODINGS_TABLE,
|
"xm430-w350": X_SERIES_ENCODINGS_TABLE,
|
||||||
"xm540-w270": X_SERIES_ENCODINGS_TABLE,
|
"xm540-w270": X_SERIES_ENCODINGS_TABLE,
|
||||||
"xc430-w150": X_SERIES_ENCODINGS_TABLE,
|
"xc430-w150": X_SERIES_ENCODINGS_TABLE,
|
||||||
|
"xh540-w150": X_SERIES_ENCODINGS_TABLE,
|
||||||
|
"xc330-t288": X_SERIES_ENCODINGS_TABLE,
|
||||||
|
"xc330-t181": X_SERIES_ENCODINGS_TABLE,
|
||||||
}
|
}
|
||||||
|
|
||||||
# {model: model_resolution}
|
# {model: model_resolution}
|
||||||
@@ -134,6 +137,9 @@ MODEL_RESOLUTION = {
|
|||||||
"xm430-w350": 4096,
|
"xm430-w350": 4096,
|
||||||
"xm540-w270": 4096,
|
"xm540-w270": 4096,
|
||||||
"xc430-w150": 4096,
|
"xc430-w150": 4096,
|
||||||
|
"xh540-w150": 4096,
|
||||||
|
"xc330-t288": 4096,
|
||||||
|
"xc330-t181": 4096,
|
||||||
}
|
}
|
||||||
|
|
||||||
# {model: model_number}
|
# {model: model_number}
|
||||||
@@ -145,6 +151,9 @@ MODEL_NUMBER_TABLE = {
|
|||||||
"xm430-w350": 1020,
|
"xm430-w350": 1020,
|
||||||
"xm540-w270": 1120,
|
"xm540-w270": 1120,
|
||||||
"xc430-w150": 1070,
|
"xc430-w150": 1070,
|
||||||
|
"xh540-w150": 1110,
|
||||||
|
"xc330-t288": 1220,
|
||||||
|
"xc330-t181": 1210,
|
||||||
}
|
}
|
||||||
|
|
||||||
# {model: available_operating_modes}
|
# {model: available_operating_modes}
|
||||||
@@ -156,6 +165,9 @@ MODEL_OPERATING_MODES = {
|
|||||||
"xm430-w350": [0, 1, 3, 4, 5, 16],
|
"xm430-w350": [0, 1, 3, 4, 5, 16],
|
||||||
"xm540-w270": [0, 1, 3, 4, 5, 16],
|
"xm540-w270": [0, 1, 3, 4, 5, 16],
|
||||||
"xc430-w150": [1, 3, 4, 16],
|
"xc430-w150": [1, 3, 4, 16],
|
||||||
|
"xh540-w150": [0, 1, 3, 4, 5, 16],
|
||||||
|
"xc330-t288": [0, 1, 3, 4, 5, 16],
|
||||||
|
"xc330-t181": [0, 1, 3, 4, 5, 16],
|
||||||
}
|
}
|
||||||
|
|
||||||
MODEL_CONTROL_TABLE = {
|
MODEL_CONTROL_TABLE = {
|
||||||
@@ -166,6 +178,9 @@ MODEL_CONTROL_TABLE = {
|
|||||||
"xm430-w350": X_SERIES_CONTROL_TABLE,
|
"xm430-w350": X_SERIES_CONTROL_TABLE,
|
||||||
"xm540-w270": X_SERIES_CONTROL_TABLE,
|
"xm540-w270": X_SERIES_CONTROL_TABLE,
|
||||||
"xc430-w150": X_SERIES_CONTROL_TABLE,
|
"xc430-w150": X_SERIES_CONTROL_TABLE,
|
||||||
|
"xh540-w150": X_SERIES_CONTROL_TABLE,
|
||||||
|
"xc330-t288": X_SERIES_CONTROL_TABLE,
|
||||||
|
"xc330-t181": X_SERIES_CONTROL_TABLE,
|
||||||
}
|
}
|
||||||
|
|
||||||
MODEL_BAUDRATE_TABLE = {
|
MODEL_BAUDRATE_TABLE = {
|
||||||
@@ -176,6 +191,9 @@ MODEL_BAUDRATE_TABLE = {
|
|||||||
"xm430-w350": X_SERIES_BAUDRATE_TABLE,
|
"xm430-w350": X_SERIES_BAUDRATE_TABLE,
|
||||||
"xm540-w270": X_SERIES_BAUDRATE_TABLE,
|
"xm540-w270": X_SERIES_BAUDRATE_TABLE,
|
||||||
"xc430-w150": X_SERIES_BAUDRATE_TABLE,
|
"xc430-w150": X_SERIES_BAUDRATE_TABLE,
|
||||||
|
"xh540-w150": X_SERIES_BAUDRATE_TABLE,
|
||||||
|
"xc330-t288": X_SERIES_BAUDRATE_TABLE,
|
||||||
|
"xc330-t181": X_SERIES_BAUDRATE_TABLE,
|
||||||
}
|
}
|
||||||
|
|
||||||
AVAILABLE_BAUDRATES = [
|
AVAILABLE_BAUDRATES = [
|
||||||
|
|||||||
@@ -44,12 +44,19 @@ from lerobot.utils.constants import (
|
|||||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||||
)
|
)
|
||||||
from lerobot.utils.feature_utils import dataset_to_policy_features
|
from lerobot.utils.feature_utils import dataset_to_policy_features
|
||||||
|
from lerobot.utils.import_utils import _peft_available, require_package
|
||||||
|
|
||||||
from .evo1.configuration_evo1 import Evo1Config
|
from .evo1.configuration_evo1 import Evo1Config
|
||||||
from .groot.configuration_groot import GrootConfig
|
from .groot.configuration_groot import GrootConfig
|
||||||
from .pretrained import PreTrainedPolicy
|
from .pretrained import PreTrainedPolicy
|
||||||
from .utils import validate_visual_features_consistency
|
from .utils import validate_visual_features_consistency
|
||||||
|
|
||||||
|
if TYPE_CHECKING or _peft_available:
|
||||||
|
from peft import PeftConfig, PeftModel
|
||||||
|
else:
|
||||||
|
PeftConfig = None
|
||||||
|
PeftModel = None
|
||||||
|
|
||||||
|
|
||||||
def _reconnect_relative_absolute_steps(
|
def _reconnect_relative_absolute_steps(
|
||||||
preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline
|
preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline
|
||||||
@@ -177,6 +184,7 @@ def make_pre_post_processors(
|
|||||||
return make_groot_pre_post_processors_from_pretrained(
|
return make_groot_pre_post_processors_from_pretrained(
|
||||||
config=policy_cfg,
|
config=policy_cfg,
|
||||||
pretrained_path=pretrained_path,
|
pretrained_path=pretrained_path,
|
||||||
|
revision=pretrained_revision,
|
||||||
dataset_stats=kwargs.get("dataset_stats"),
|
dataset_stats=kwargs.get("dataset_stats"),
|
||||||
dataset_meta=kwargs.get("dataset_meta"),
|
dataset_meta=kwargs.get("dataset_meta"),
|
||||||
preprocessor_overrides=kwargs.get("preprocessor_overrides"),
|
preprocessor_overrides=kwargs.get("preprocessor_overrides"),
|
||||||
@@ -333,12 +341,15 @@ def make_policy(
|
|||||||
# Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
|
# Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
|
||||||
# of the adapter and the adapter's config contains the path to the base policy. So we need the
|
# of the adapter and the adapter's config contains the path to the base policy. So we need the
|
||||||
# adapter config first, then load the correct policy and then apply PEFT.
|
# adapter config first, then load the correct policy and then apply PEFT.
|
||||||
from peft import PeftConfig, PeftModel
|
require_package("peft", extra="peft")
|
||||||
|
|
||||||
logging.info("Loading policy's PEFT adapter.")
|
logging.info("Loading policy's PEFT adapter.")
|
||||||
|
|
||||||
peft_pretrained_path = str(cfg.pretrained_path)
|
peft_pretrained_path = str(cfg.pretrained_path)
|
||||||
peft_config = PeftConfig.from_pretrained(peft_pretrained_path)
|
peft_config = PeftConfig.from_pretrained(
|
||||||
|
peft_pretrained_path,
|
||||||
|
revision=cfg.pretrained_revision,
|
||||||
|
)
|
||||||
|
|
||||||
kwargs["pretrained_name_or_path"] = peft_config.base_model_name_or_path
|
kwargs["pretrained_name_or_path"] = peft_config.base_model_name_or_path
|
||||||
if not kwargs["pretrained_name_or_path"]:
|
if not kwargs["pretrained_name_or_path"]:
|
||||||
@@ -349,9 +360,14 @@ def make_policy(
|
|||||||
"the adapter was trained."
|
"the adapter was trained."
|
||||||
)
|
)
|
||||||
|
|
||||||
|
kwargs["revision"] = peft_config.revision
|
||||||
policy = policy_cls.from_pretrained(**kwargs)
|
policy = policy_cls.from_pretrained(**kwargs)
|
||||||
policy = PeftModel.from_pretrained(
|
policy = PeftModel.from_pretrained(
|
||||||
policy, peft_pretrained_path, config=peft_config, is_trainable=True
|
policy,
|
||||||
|
peft_pretrained_path,
|
||||||
|
config=peft_config,
|
||||||
|
revision=cfg.pretrained_revision,
|
||||||
|
is_trainable=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
else:
|
else:
|
||||||
|
|||||||
@@ -37,13 +37,19 @@ def is_image_feature(key: str) -> bool:
|
|||||||
@dataclass
|
@dataclass
|
||||||
class ConcurrencyConfig:
|
class ConcurrencyConfig:
|
||||||
"""Configuration for the concurrency of the actor and learner.
|
"""Configuration for the concurrency of the actor and learner.
|
||||||
|
|
||||||
Possible values are:
|
Possible values are:
|
||||||
- "threads": Use threads for the actor and learner.
|
- "threads": Use threads for the actor and learner.
|
||||||
- "processes": Use processes for the actor and learner.
|
- "processes": Use processes for the actor and learner.
|
||||||
|
|
||||||
|
``multiprocessing_context`` selects the process-wide start method when
|
||||||
|
processes are used. Set it to ``None`` to preserve Python's default or a
|
||||||
|
method already selected by the embedding application.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
actor: str = "threads"
|
actor: str = "threads"
|
||||||
learner: str = "threads"
|
learner: str = "threads"
|
||||||
|
multiprocessing_context: str | None = "spawn"
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
|
|||||||
@@ -475,6 +475,7 @@ def make_groot_pre_post_processors_from_pretrained(
|
|||||||
config: GrootConfig,
|
config: GrootConfig,
|
||||||
pretrained_path: str,
|
pretrained_path: str,
|
||||||
*,
|
*,
|
||||||
|
revision: str | None = None,
|
||||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||||
dataset_meta: Any | None = None,
|
dataset_meta: Any | None = None,
|
||||||
preprocessor_overrides: dict[str, Any] | None = None,
|
preprocessor_overrides: dict[str, Any] | None = None,
|
||||||
@@ -511,6 +512,7 @@ def make_groot_pre_post_processors_from_pretrained(
|
|||||||
|
|
||||||
preprocessor, postprocessor = _load_groot_processor_pipelines(
|
preprocessor, postprocessor = _load_groot_processor_pipelines(
|
||||||
pretrained_path,
|
pretrained_path,
|
||||||
|
revision=revision,
|
||||||
preprocessor_overrides=preprocessor_overrides,
|
preprocessor_overrides=preprocessor_overrides,
|
||||||
postprocessor_overrides=postprocessor_overrides,
|
postprocessor_overrides=postprocessor_overrides,
|
||||||
preprocessor_config_filename=preprocessor_config_filename,
|
preprocessor_config_filename=preprocessor_config_filename,
|
||||||
@@ -526,6 +528,7 @@ def make_groot_pre_post_processors_from_pretrained(
|
|||||||
def _load_groot_processor_pipelines(
|
def _load_groot_processor_pipelines(
|
||||||
pretrained_path: str,
|
pretrained_path: str,
|
||||||
*,
|
*,
|
||||||
|
revision: str | None,
|
||||||
preprocessor_overrides: dict[str, Any],
|
preprocessor_overrides: dict[str, Any],
|
||||||
postprocessor_overrides: dict[str, Any],
|
postprocessor_overrides: dict[str, Any],
|
||||||
preprocessor_config_filename: str,
|
preprocessor_config_filename: str,
|
||||||
@@ -540,6 +543,7 @@ def _load_groot_processor_pipelines(
|
|||||||
preprocessor = PolicyProcessorPipeline.from_pretrained(
|
preprocessor = PolicyProcessorPipeline.from_pretrained(
|
||||||
pretrained_model_name_or_path=pretrained_path,
|
pretrained_model_name_or_path=pretrained_path,
|
||||||
config_filename=preprocessor_config_filename,
|
config_filename=preprocessor_config_filename,
|
||||||
|
revision=revision,
|
||||||
overrides=preprocessor_overrides,
|
overrides=preprocessor_overrides,
|
||||||
to_transition=batch_to_transition,
|
to_transition=batch_to_transition,
|
||||||
to_output=transition_to_batch,
|
to_output=transition_to_batch,
|
||||||
@@ -547,6 +551,7 @@ def _load_groot_processor_pipelines(
|
|||||||
postprocessor = PolicyProcessorPipeline.from_pretrained(
|
postprocessor = PolicyProcessorPipeline.from_pretrained(
|
||||||
pretrained_model_name_or_path=pretrained_path,
|
pretrained_model_name_or_path=pretrained_path,
|
||||||
config_filename=postprocessor_config_filename,
|
config_filename=postprocessor_config_filename,
|
||||||
|
revision=revision,
|
||||||
overrides=postprocessor_overrides,
|
overrides=postprocessor_overrides,
|
||||||
to_transition=policy_action_to_transition,
|
to_transition=policy_action_to_transition,
|
||||||
to_output=transition_to_policy_action,
|
to_output=transition_to_policy_action,
|
||||||
|
|||||||
@@ -43,11 +43,22 @@ from torch.distributions import Beta
|
|||||||
|
|
||||||
from lerobot.policies.pretrained import PreTrainedPolicy
|
from lerobot.policies.pretrained import PreTrainedPolicy
|
||||||
from lerobot.utils.constants import ACTION
|
from lerobot.utils.constants import ACTION
|
||||||
from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
|
from lerobot.utils.import_utils import (
|
||||||
|
_peft_available,
|
||||||
|
_scipy_available,
|
||||||
|
_transformers_available,
|
||||||
|
require_package,
|
||||||
|
)
|
||||||
|
|
||||||
from ..rtc.modeling_rtc import RTCProcessor
|
from ..rtc.modeling_rtc import RTCProcessor
|
||||||
from .configuration_molmoact2 import MolmoAct2Config
|
from .configuration_molmoact2 import MolmoAct2Config
|
||||||
|
|
||||||
|
if TYPE_CHECKING or _peft_available:
|
||||||
|
from peft import LoraConfig, get_peft_model
|
||||||
|
else:
|
||||||
|
LoraConfig = None
|
||||||
|
get_peft_model = None
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@@ -1731,13 +1742,11 @@ class MolmoAct2Policy(PreTrainedPolicy):
|
|||||||
|
|
||||||
def _build_inner_lora_config(self):
|
def _build_inner_lora_config(self):
|
||||||
require_package("peft", extra="molmoact2")
|
require_package("peft", extra="molmoact2")
|
||||||
from peft import LoraConfig
|
|
||||||
|
|
||||||
return LoraConfig(**self._get_inner_peft_targets())
|
return LoraConfig(**self._get_inner_peft_targets())
|
||||||
|
|
||||||
def _apply_lora_adapters(self) -> None:
|
def _apply_lora_adapters(self) -> None:
|
||||||
require_package("peft", extra="molmoact2")
|
require_package("peft", extra="molmoact2")
|
||||||
from peft import get_peft_model
|
|
||||||
|
|
||||||
peft_config = self._build_inner_lora_config()
|
peft_config = self._build_inner_lora_config()
|
||||||
self._validate_peft_config(peft_config)
|
self._validate_peft_config(peft_config)
|
||||||
|
|||||||
@@ -34,14 +34,22 @@ from lerobot.configs import PreTrainedConfig
|
|||||||
from lerobot.configs.train import TrainPipelineConfig
|
from lerobot.configs.train import TrainPipelineConfig
|
||||||
from lerobot.utils.device_utils import resolve_safetensors_device
|
from lerobot.utils.device_utils import resolve_safetensors_device
|
||||||
from lerobot.utils.hub import HubMixin
|
from lerobot.utils.hub import HubMixin
|
||||||
|
from lerobot.utils.import_utils import _peft_available, require_package
|
||||||
|
|
||||||
from .utils import log_model_loading_keys
|
from .utils import log_model_loading_keys
|
||||||
|
|
||||||
T = TypeVar("T", bound="PreTrainedPolicy")
|
if TYPE_CHECKING or _peft_available:
|
||||||
|
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType, get_peft_model
|
||||||
|
else:
|
||||||
|
PEFT_TYPE_TO_CONFIG_MAPPING = None
|
||||||
|
PeftType = None
|
||||||
|
get_peft_model = None
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
||||||
|
|
||||||
|
T = TypeVar("T", bound="PreTrainedPolicy")
|
||||||
|
|
||||||
|
|
||||||
def _build_card_context(
|
def _build_card_context(
|
||||||
cfg: TrainPipelineConfig | None,
|
cfg: TrainPipelineConfig | None,
|
||||||
@@ -384,7 +392,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
|||||||
peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.)
|
peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.)
|
||||||
These are merged with policy defaults to build the final config.
|
These are merged with policy defaults to build the final config.
|
||||||
"""
|
"""
|
||||||
from peft import get_peft_model
|
require_package("peft", extra="peft")
|
||||||
|
|
||||||
# If user provided a complete config, use it directly (with overrides)
|
# If user provided a complete config, use it directly (with overrides)
|
||||||
if peft_config is not None:
|
if peft_config is not None:
|
||||||
@@ -455,7 +463,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
|||||||
Returns:
|
Returns:
|
||||||
Preprocessed dict with renamed keys and init_type mapped to method-specific key.
|
Preprocessed dict with renamed keys and init_type mapped to method-specific key.
|
||||||
"""
|
"""
|
||||||
from peft import PeftType
|
require_package("peft", extra="peft")
|
||||||
|
|
||||||
cli_overrides = cli_overrides.copy()
|
cli_overrides = cli_overrides.copy()
|
||||||
|
|
||||||
@@ -480,7 +488,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
|||||||
|
|
||||||
def _build_peft_config(self, cli_overrides: dict):
|
def _build_peft_config(self, cli_overrides: dict):
|
||||||
"""Build a PEFT config from policy defaults and CLI overrides."""
|
"""Build a PEFT config from policy defaults and CLI overrides."""
|
||||||
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
|
require_package("peft", extra="peft")
|
||||||
|
|
||||||
# Determine PEFT method type (default to LORA)
|
# Determine PEFT method type (default to LORA)
|
||||||
method_type_str = cli_overrides.get("method_type") or "lora"
|
method_type_str = cli_overrides.get("method_type") or "lora"
|
||||||
@@ -507,7 +515,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
|||||||
|
|
||||||
def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict):
|
def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict):
|
||||||
"""Apply CLI overrides to an existing PEFT config."""
|
"""Apply CLI overrides to an existing PEFT config."""
|
||||||
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
|
require_package("peft", extra="peft")
|
||||||
|
|
||||||
# Get method type from existing config or CLI override
|
# Get method type from existing config or CLI override
|
||||||
method_type_str = cli_overrides.get("method_type")
|
method_type_str = cli_overrides.get("method_type")
|
||||||
|
|||||||
@@ -132,10 +132,20 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
|
|||||||
def transform_features(
|
def transform_features(
|
||||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||||
for axis in ["x", "y", "z", "gripper"]:
|
for axis in ["x", "y", "z"]:
|
||||||
features[PipelineFeatureType.ACTION].pop(f"delta_{axis}", None)
|
features[PipelineFeatureType.ACTION].pop(f"delta_{axis}", None)
|
||||||
|
features[PipelineFeatureType.ACTION].pop("gripper", None)
|
||||||
|
|
||||||
for feat in ["enabled", "target_x", "target_y", "target_z", "target_wx", "target_wy", "target_wz"]:
|
for feat in [
|
||||||
|
"enabled",
|
||||||
|
"target_x",
|
||||||
|
"target_y",
|
||||||
|
"target_z",
|
||||||
|
"target_wx",
|
||||||
|
"target_wy",
|
||||||
|
"target_wz",
|
||||||
|
"gripper_vel",
|
||||||
|
]:
|
||||||
features[PipelineFeatureType.ACTION][f"{feat}"] = PolicyFeature(
|
features[PipelineFeatureType.ACTION][f"{feat}"] = PolicyFeature(
|
||||||
type=FeatureType.ACTION, shape=(1,)
|
type=FeatureType.ACTION, shape=(1,)
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -713,6 +713,8 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
ProcessorMigrationError: If the model requires migration to processor format.
|
ProcessorMigrationError: If the model requires migration to processor format.
|
||||||
"""
|
"""
|
||||||
model_id = str(pretrained_model_name_or_path)
|
model_id = str(pretrained_model_name_or_path)
|
||||||
|
model_path = Path(model_id)
|
||||||
|
is_local_source = model_path.is_dir() or model_path.is_file()
|
||||||
hub_download_kwargs = {
|
hub_download_kwargs = {
|
||||||
"force_download": force_download,
|
"force_download": force_download,
|
||||||
"resume_download": resume_download,
|
"resume_download": resume_download,
|
||||||
@@ -731,7 +733,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
|
|
||||||
# 3. Build steps with overrides
|
# 3. Build steps with overrides
|
||||||
steps, validated_overrides = cls._build_steps_with_overrides(
|
steps, validated_overrides = cls._build_steps_with_overrides(
|
||||||
loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs
|
loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs, is_local_source
|
||||||
)
|
)
|
||||||
|
|
||||||
# 4. Validate that all overrides were used
|
# 4. Validate that all overrides were used
|
||||||
@@ -921,6 +923,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
model_id: str,
|
model_id: str,
|
||||||
base_path: Path | None,
|
base_path: Path | None,
|
||||||
hub_download_kwargs: dict[str, Any],
|
hub_download_kwargs: dict[str, Any],
|
||||||
|
is_local_source: bool = False,
|
||||||
) -> tuple[list[ProcessorStep], set[str]]:
|
) -> tuple[list[ProcessorStep], set[str]]:
|
||||||
"""Build all processor steps with overrides and state loading.
|
"""Build all processor steps with overrides and state loading.
|
||||||
|
|
||||||
@@ -944,7 +947,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
3. **State Loading** (via _load_step_state):
|
3. **State Loading** (via _load_step_state):
|
||||||
- **If step has "state_file"**: Load tensor state from .safetensors
|
- **If step has "state_file"**: Load tensor state from .safetensors
|
||||||
- **Local first**: Check base_path/state_file.safetensors
|
- **Local first**: Check base_path/state_file.safetensors
|
||||||
- **Hub fallback**: Download state file if not found locally
|
- **Hub fallback**: Download state file if the pipeline was loaded from the Hub
|
||||||
- **Optional**: Only load if step has load_state_dict method
|
- **Optional**: Only load if step has load_state_dict method
|
||||||
|
|
||||||
4. **Override Tracking**:
|
4. **Override Tracking**:
|
||||||
@@ -962,6 +965,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
model_id: The model identifier (needed for Hub state file downloads)
|
model_id: The model identifier (needed for Hub state file downloads)
|
||||||
base_path: Local directory path for finding state files
|
base_path: Local directory path for finding state files
|
||||||
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
|
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
|
||||||
|
is_local_source: Whether model_id resolved to a local directory or config file.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Tuple of (instantiated_steps_list, unused_override_keys)
|
Tuple of (instantiated_steps_list, unused_override_keys)
|
||||||
@@ -975,7 +979,9 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides)
|
steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides)
|
||||||
|
|
||||||
for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True):
|
for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True):
|
||||||
cls._load_step_state(step_instance, step_entry, model_id, base_path, hub_download_kwargs)
|
cls._load_step_state(
|
||||||
|
step_instance, step_entry, model_id, base_path, hub_download_kwargs, is_local_source
|
||||||
|
)
|
||||||
|
|
||||||
return steps, remaining_override_keys
|
return steps, remaining_override_keys
|
||||||
|
|
||||||
@@ -1139,6 +1145,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
model_id: str,
|
model_id: str,
|
||||||
base_path: Path | None,
|
base_path: Path | None,
|
||||||
hub_download_kwargs: dict[str, Any],
|
hub_download_kwargs: dict[str, Any],
|
||||||
|
is_local_source: bool = False,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Load state dictionary for a processor step if available.
|
"""Load state dictionary for a processor step if available.
|
||||||
|
|
||||||
@@ -1157,7 +1164,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
- **Use case**: Loading from local saved model directory
|
- **Use case**: Loading from local saved model directory
|
||||||
|
|
||||||
2. **Hub download fallback**: Download state file from repository
|
2. **Hub download fallback**: Download state file from repository
|
||||||
- **When triggered**: Local file not found or base_path is None
|
- **When triggered**: Local file not found and the pipeline source is a Hub repo
|
||||||
- **Process**: Use hf_hub_download with same parameters as config
|
- **Process**: Use hf_hub_download with same parameters as config
|
||||||
- **Example**: Download "normalize_step_0.safetensors" from "user/repo"
|
- **Example**: Download "normalize_step_0.safetensors" from "user/repo"
|
||||||
- **Result**: Downloaded to local cache, path returned
|
- **Result**: Downloaded to local cache, path returned
|
||||||
@@ -1178,6 +1185,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
model_id: The model identifier (used for Hub downloads if needed)
|
model_id: The model identifier (used for Hub downloads if needed)
|
||||||
base_path: Local directory path for finding state files (None for Hub-only)
|
base_path: Local directory path for finding state files (None for Hub-only)
|
||||||
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
|
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
|
||||||
|
is_local_source: Whether model_id resolved to a local directory or config file.
|
||||||
|
|
||||||
Note:
|
Note:
|
||||||
This method modifies step_instance in-place and returns None.
|
This method modifies step_instance in-place and returns None.
|
||||||
@@ -1191,6 +1199,12 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
# Try local file first
|
# Try local file first
|
||||||
if base_path and (base_path / state_filename).exists():
|
if base_path and (base_path / state_filename).exists():
|
||||||
state_path = str(base_path / state_filename)
|
state_path = str(base_path / state_filename)
|
||||||
|
elif is_local_source:
|
||||||
|
state_path = base_path / state_filename if base_path else Path(state_filename)
|
||||||
|
raise FileNotFoundError(
|
||||||
|
f"State file '{state_filename}' was not found for local processor pipeline "
|
||||||
|
f"'{model_id}' at '{state_path}'."
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
# Download from Hub
|
# Download from Hub
|
||||||
state_path = hf_hub_download(
|
state_path = hf_hub_download(
|
||||||
|
|||||||
@@ -91,7 +91,7 @@ from lerobot.robots import so_follower # noqa: F401
|
|||||||
from lerobot.teleoperators import gamepad, so_leader # noqa: F401
|
from lerobot.teleoperators import gamepad, so_leader # noqa: F401
|
||||||
from lerobot.teleoperators.utils import TeleopEvents
|
from lerobot.teleoperators.utils import TeleopEvents
|
||||||
from lerobot.utils.device_utils import get_safe_torch_device
|
from lerobot.utils.device_utils import get_safe_torch_device
|
||||||
from lerobot.utils.process import ProcessSignalHandler
|
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
|
||||||
from lerobot.utils.random_utils import set_seed
|
from lerobot.utils.random_utils import set_seed
|
||||||
from lerobot.utils.robot_utils import precise_sleep
|
from lerobot.utils.robot_utils import precise_sleep
|
||||||
from lerobot.utils.transition import (
|
from lerobot.utils.transition import (
|
||||||
@@ -124,9 +124,7 @@ def actor_cli(cfg: TrainRLServerPipelineConfig):
|
|||||||
cfg.validate()
|
cfg.validate()
|
||||||
display_pid = False
|
display_pid = False
|
||||||
if not use_threads(cfg):
|
if not use_threads(cfg):
|
||||||
import torch.multiprocessing as mp
|
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
|
||||||
|
|
||||||
mp.set_start_method("spawn")
|
|
||||||
display_pid = True
|
display_pid = True
|
||||||
|
|
||||||
# Create logs directory to ensure it exists
|
# Create logs directory to ensure it exists
|
||||||
|
|||||||
@@ -102,7 +102,7 @@ from lerobot.utils.constants import (
|
|||||||
)
|
)
|
||||||
from lerobot.utils.device_utils import get_safe_torch_device
|
from lerobot.utils.device_utils import get_safe_torch_device
|
||||||
from lerobot.utils.io_utils import load_json, write_json
|
from lerobot.utils.io_utils import load_json, write_json
|
||||||
from lerobot.utils.process import ProcessSignalHandler
|
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
|
||||||
from lerobot.utils.random_utils import set_seed
|
from lerobot.utils.random_utils import set_seed
|
||||||
from lerobot.utils.utils import (
|
from lerobot.utils.utils import (
|
||||||
format_big_number,
|
format_big_number,
|
||||||
@@ -123,9 +123,7 @@ def train_cli(cfg: TrainRLServerPipelineConfig):
|
|||||||
# Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
|
# Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
|
||||||
require_package("grpcio", extra="hilserl", import_name="grpc")
|
require_package("grpcio", extra="hilserl", import_name="grpc")
|
||||||
if not use_threads(cfg):
|
if not use_threads(cfg):
|
||||||
import torch.multiprocessing as mp
|
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
|
||||||
|
|
||||||
mp.set_start_method("spawn")
|
|
||||||
|
|
||||||
# Use the job_name from the config
|
# Use the job_name from the config
|
||||||
train(
|
train(
|
||||||
|
|||||||
@@ -323,6 +323,10 @@ class LeKiwiClient(Robot):
|
|||||||
np.ndarray: the action sent to the motors, potentially clipped.
|
np.ndarray: the action sent to the motors, potentially clipped.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
# Action values may be torch tensors (e.g. replayed from a dataset) or numpy
|
||||||
|
# scalars; json.dumps only serializes Python primitives, so coerce each value to a
|
||||||
|
# plain float before sending.
|
||||||
|
action = {key: float(value) for key, value in action.items()}
|
||||||
self.zmq_cmd_socket.send_string(json.dumps(action)) # action is in motor space
|
self.zmq_cmd_socket.send_string(json.dumps(action)) # action is in motor space
|
||||||
|
|
||||||
# TODO(Steven): Remove the np conversion when it is possible to record a non-numpy array value
|
# TODO(Steven): Remove the np conversion when it is possible to record a non-numpy array value
|
||||||
|
|||||||
@@ -326,8 +326,17 @@ class RolloutConfig:
|
|||||||
|
|
||||||
policy_path = parser.get_path_arg("policy")
|
policy_path = parser.get_path_arg("policy")
|
||||||
if policy_path:
|
if policy_path:
|
||||||
cli_overrides = parser.get_cli_overrides("policy")
|
yaml_overrides = parser.get_yaml_overrides("policy")
|
||||||
self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
|
cli_overrides = parser.get_cli_overrides("policy") or []
|
||||||
|
policy_overrides = yaml_overrides + cli_overrides
|
||||||
|
pretrained_revision = parser.parse_arg("pretrained_revision", cli_overrides)
|
||||||
|
if pretrained_revision is None:
|
||||||
|
pretrained_revision = parser.parse_arg("pretrained_revision", yaml_overrides)
|
||||||
|
self.policy = PreTrainedConfig.from_pretrained(
|
||||||
|
policy_path,
|
||||||
|
revision=pretrained_revision,
|
||||||
|
cli_overrides=policy_overrides,
|
||||||
|
)
|
||||||
self.policy.pretrained_path = policy_path
|
self.policy.pretrained_path = policy_path
|
||||||
if self.policy is None:
|
if self.policy is None:
|
||||||
raise ValueError("--policy.path is required for rollout")
|
raise ValueError("--policy.path is required for rollout")
|
||||||
|
|||||||
@@ -24,10 +24,11 @@ from __future__ import annotations
|
|||||||
import logging
|
import logging
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
from threading import Event
|
from threading import Event
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from lerobot.configs import FeatureType
|
from lerobot.configs import FeatureType, PreTrainedConfig
|
||||||
from lerobot.datasets import (
|
from lerobot.datasets import (
|
||||||
LeRobotDataset,
|
LeRobotDataset,
|
||||||
aggregate_pipeline_dataset_features,
|
aggregate_pipeline_dataset_features,
|
||||||
@@ -47,6 +48,7 @@ from lerobot.processor.relative_action_processor import RelativeActionsProcessor
|
|||||||
from lerobot.robots import make_robot_from_config
|
from lerobot.robots import make_robot_from_config
|
||||||
from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
|
from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
|
||||||
from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
|
from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
|
||||||
|
from lerobot.utils.import_utils import _peft_available, require_package
|
||||||
|
|
||||||
from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
|
from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
|
||||||
from .inference import (
|
from .inference import (
|
||||||
@@ -57,6 +59,12 @@ from .inference import (
|
|||||||
)
|
)
|
||||||
from .robot_wrapper import ThreadSafeRobot
|
from .robot_wrapper import ThreadSafeRobot
|
||||||
|
|
||||||
|
if TYPE_CHECKING or _peft_available:
|
||||||
|
from peft import PeftConfig, PeftModel
|
||||||
|
else:
|
||||||
|
PeftConfig = None
|
||||||
|
PeftModel = None
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@@ -159,6 +167,35 @@ class RolloutContext:
|
|||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy:
|
||||||
|
"""Load policy weights, keeping adapter and base-model revisions independent."""
|
||||||
|
pretrained_revision = policy_config.pretrained_revision
|
||||||
|
policy_class = get_policy_class(policy_config.type)
|
||||||
|
|
||||||
|
if not policy_config.use_peft:
|
||||||
|
return policy_class.from_pretrained(
|
||||||
|
policy_config.pretrained_path,
|
||||||
|
config=policy_config,
|
||||||
|
revision=pretrained_revision,
|
||||||
|
)
|
||||||
|
|
||||||
|
require_package("peft", extra="peft")
|
||||||
|
|
||||||
|
peft_path = policy_config.pretrained_path
|
||||||
|
peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision)
|
||||||
|
policy = policy_class.from_pretrained(
|
||||||
|
pretrained_name_or_path=peft_config.base_model_name_or_path,
|
||||||
|
config=policy_config,
|
||||||
|
revision=peft_config.revision,
|
||||||
|
)
|
||||||
|
return PeftModel.from_pretrained(
|
||||||
|
policy,
|
||||||
|
peft_path,
|
||||||
|
config=peft_config,
|
||||||
|
revision=pretrained_revision,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def build_rollout_context(
|
def build_rollout_context(
|
||||||
cfg: RolloutConfig,
|
cfg: RolloutConfig,
|
||||||
shutdown_event: Event,
|
shutdown_event: Event,
|
||||||
@@ -176,7 +213,6 @@ def build_rollout_context(
|
|||||||
# --- 1. Policy (heavy I/O, but no hardware yet) -------------------
|
# --- 1. Policy (heavy I/O, but no hardware yet) -------------------
|
||||||
logger.info("Loading policy from '%s'...", cfg.policy.pretrained_path)
|
logger.info("Loading policy from '%s'...", cfg.policy.pretrained_path)
|
||||||
policy_config = cfg.policy
|
policy_config = cfg.policy
|
||||||
policy_class = get_policy_class(policy_config.type)
|
|
||||||
|
|
||||||
if hasattr(policy_config, "compile_model"):
|
if hasattr(policy_config, "compile_model"):
|
||||||
policy_config.compile_model = cfg.use_torch_compile
|
policy_config.compile_model = cfg.use_torch_compile
|
||||||
@@ -187,17 +223,7 @@ def build_rollout_context(
|
|||||||
"Please use `cpu` or `cuda` backend."
|
"Please use `cpu` or `cuda` backend."
|
||||||
)
|
)
|
||||||
|
|
||||||
if policy_config.use_peft:
|
policy = _load_pretrained_policy(policy_config)
|
||||||
from peft import PeftConfig, PeftModel
|
|
||||||
|
|
||||||
peft_path = policy_config.pretrained_path
|
|
||||||
peft_config = PeftConfig.from_pretrained(peft_path)
|
|
||||||
policy = policy_class.from_pretrained(
|
|
||||||
pretrained_name_or_path=peft_config.base_model_name_or_path, config=policy_config
|
|
||||||
)
|
|
||||||
policy = PeftModel.from_pretrained(policy, peft_path, config=peft_config)
|
|
||||||
else:
|
|
||||||
policy = policy_class.from_pretrained(policy_config.pretrained_path, config=policy_config)
|
|
||||||
|
|
||||||
if is_rtc:
|
if is_rtc:
|
||||||
policy.config.rtc_config = cfg.inference.rtc
|
policy.config.rtc_config = cfg.inference.rtc
|
||||||
@@ -276,12 +302,22 @@ def build_rollout_context(
|
|||||||
# ``observation_features`` values are either a tuple (camera shape) or the
|
# ``observation_features`` values are either a tuple (camera shape) or the
|
||||||
# ``float`` type itself used as a sentinel for scalar motor features —
|
# ``float`` type itself used as a sentinel for scalar motor features —
|
||||||
# see ``dict[str, type | tuple]`` annotation on ``Robot.observation_features``.
|
# see ``dict[str, type | tuple]`` annotation on ``Robot.observation_features``.
|
||||||
|
# Keep cameras (tuple) plus both joint-position (.pos) and base-velocity (.vel)
|
||||||
|
# scalar state features. LeKiwi's observation.state is 9-dim (6 arm .pos +
|
||||||
|
# x/y/theta.vel) and the policy was trained/normalized on all 9; the old .pos-only
|
||||||
|
# filter fed a 6-dim state into a 9-dim normalizer → RuntimeError (size 6 vs 9).
|
||||||
|
# Pure-arm robots have no .vel state keys, so this is a no-op for them.
|
||||||
observation_features_hw = {
|
observation_features_hw = {
|
||||||
k: v
|
k: v
|
||||||
for k, v in all_obs_features.items()
|
for k, v in all_obs_features.items()
|
||||||
if isinstance(v, tuple) or (v is float and k.endswith(".pos"))
|
if isinstance(v, tuple) or (v is float and k.endswith((".pos", ".vel")))
|
||||||
}
|
}
|
||||||
action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith(".pos")}
|
# Keep both joint-position (.pos) and base-velocity (.vel) action features so
|
||||||
|
# mobile manipulators command the base too (e.g. LeKiwi: 6 arm .pos +
|
||||||
|
# x/y/theta.vel = 9-dim action). Pure-arm robots have no .vel keys, so this is
|
||||||
|
# a no-op for them. Without the .vel keys the base velocities are silently
|
||||||
|
# dropped from dataset_features[ACTION]/ordered_action_keys and the base never moves.
|
||||||
|
action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith((".pos", ".vel"))}
|
||||||
|
|
||||||
# The action side is always needed: sync inference reads action names from
|
# The action side is always needed: sync inference reads action names from
|
||||||
# ``dataset_features[ACTION]`` to map policy tensors back to robot actions.
|
# ``dataset_features[ACTION]`` to map policy tensors back to robot actions.
|
||||||
@@ -392,6 +428,7 @@ def build_rollout_context(
|
|||||||
preprocessor, postprocessor = make_pre_post_processors(
|
preprocessor, postprocessor = make_pre_post_processors(
|
||||||
policy_cfg=policy_config,
|
policy_cfg=policy_config,
|
||||||
pretrained_path=cfg.policy.pretrained_path,
|
pretrained_path=cfg.policy.pretrained_path,
|
||||||
|
pretrained_revision=policy_config.pretrained_revision,
|
||||||
dataset_stats=dataset_stats,
|
dataset_stats=dataset_stats,
|
||||||
preprocessor_overrides={
|
preprocessor_overrides={
|
||||||
"device_processor": {"device": cfg.device},
|
"device_processor": {"device": cfg.device},
|
||||||
|
|||||||
@@ -94,6 +94,8 @@ from lerobot.datasets.video_utils import concatenate_video_files, get_video_dura
|
|||||||
from lerobot.utils.constants import HF_LEROBOT_HOME
|
from lerobot.utils.constants import HF_LEROBOT_HOME
|
||||||
from lerobot.utils.utils import flatten_dict, init_logging
|
from lerobot.utils.utils import flatten_dict, init_logging
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
V21 = "v2.1"
|
V21 = "v2.1"
|
||||||
V30 = "v3.0"
|
V30 = "v3.0"
|
||||||
|
|
||||||
@@ -476,11 +478,11 @@ def convert_dataset(
|
|||||||
# First check if the dataset already has a v3.0 version
|
# First check if the dataset already has a v3.0 version
|
||||||
if root is None and not force_conversion:
|
if root is None and not force_conversion:
|
||||||
try:
|
try:
|
||||||
print("Trying to download v3.0 version of the dataset from the hub...")
|
logger.info("Trying to download v3.0 version of the dataset from the hub...")
|
||||||
snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id)
|
snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id)
|
||||||
return
|
return
|
||||||
except Exception:
|
except Exception:
|
||||||
print("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
|
logger.info("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
|
||||||
|
|
||||||
# Set root based on whether local dataset path is provided
|
# Set root based on whether local dataset path is provided
|
||||||
use_local_dataset = False
|
use_local_dataset = False
|
||||||
@@ -488,7 +490,7 @@ def convert_dataset(
|
|||||||
if root.exists():
|
if root.exists():
|
||||||
validate_local_dataset_version(root)
|
validate_local_dataset_version(root)
|
||||||
use_local_dataset = True
|
use_local_dataset = True
|
||||||
print(f"Using local dataset at {root}")
|
logger.info(f"Using local dataset at {root}")
|
||||||
|
|
||||||
old_root = root.parent / f"{root.name}_old"
|
old_root = root.parent / f"{root.name}_old"
|
||||||
new_root = root.parent / f"{root.name}_v30"
|
new_root = root.parent / f"{root.name}_v30"
|
||||||
@@ -523,7 +525,7 @@ def convert_dataset(
|
|||||||
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, RevisionNotFoundError) as e:
|
||||||
print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
|
logger.warning(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
|
||||||
pass
|
pass
|
||||||
hub_api.delete_files(
|
hub_api.delete_files(
|
||||||
delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"],
|
delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"],
|
||||||
|
|||||||
@@ -154,14 +154,14 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
|||||||
repo_id = cfg.new_repo_id or cfg.repo_id
|
repo_id = cfg.new_repo_id or cfg.repo_id
|
||||||
commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
|
commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
|
||||||
api = HfApi()
|
api = HfApi()
|
||||||
print(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...", flush=True)
|
logger.info(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...")
|
||||||
api.create_repo(
|
api.create_repo(
|
||||||
repo_id=repo_id,
|
repo_id=repo_id,
|
||||||
repo_type="dataset",
|
repo_type="dataset",
|
||||||
private=cfg.push_private,
|
private=cfg.push_private,
|
||||||
exist_ok=True,
|
exist_ok=True,
|
||||||
)
|
)
|
||||||
print(f"[lerobot-annotate] uploading {root} -> {repo_id}...", flush=True)
|
logger.info(f"[lerobot-annotate] uploading {root} -> {repo_id}...")
|
||||||
commit_info = api.upload_folder(
|
commit_info = api.upload_folder(
|
||||||
folder_path=str(root),
|
folder_path=str(root),
|
||||||
repo_id=repo_id,
|
repo_id=repo_id,
|
||||||
@@ -172,7 +172,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
|||||||
# at the source dataset; a fresh card is generated below instead.
|
# at the source dataset; a fresh card is generated below instead.
|
||||||
ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
|
ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
|
||||||
)
|
)
|
||||||
print(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}", flush=True)
|
logger.info(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}")
|
||||||
|
|
||||||
dataset_info = load_info(root)
|
dataset_info = load_info(root)
|
||||||
card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
|
card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
|
||||||
@@ -200,14 +200,13 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
|
|||||||
with suppress(RevisionNotFoundError):
|
with suppress(RevisionNotFoundError):
|
||||||
api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
|
api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
|
||||||
api.create_tag(**tag_kwargs)
|
api.create_tag(**tag_kwargs)
|
||||||
print(f"[lerobot-annotate] tagged {repo_id} as {version_tag}", flush=True)
|
logger.info(f"[lerobot-annotate] tagged {repo_id} as {version_tag}")
|
||||||
except Exception as exc: # noqa: BLE001
|
except Exception as exc: # noqa: BLE001
|
||||||
print(
|
logger.warning(
|
||||||
f"[lerobot-annotate] WARNING: could not create tag {version_tag!r} on {repo_id}: {exc}. "
|
f"[lerobot-annotate] WARNING: could not create tag {version_tag!r} on {repo_id}: {exc}. "
|
||||||
"Dataset is uploaded but ``LeRobotDataset`` won't be able to load it until it's tagged. "
|
"Dataset is uploaded but ``LeRobotDataset`` won't be able to load it until it's tagged. "
|
||||||
"Run: from huggingface_hub import HfApi; "
|
"Run: from huggingface_hub import HfApi; "
|
||||||
f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)",
|
f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)"
|
||||||
flush=True,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -89,6 +89,8 @@ from lerobot.datasets import LeRobotDataset
|
|||||||
from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
|
from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
|
||||||
from lerobot.utils.utils import init_logging
|
from lerobot.utils.utils import init_logging
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
DEFAULT_FOXGLOVE_PORT = 8765
|
DEFAULT_FOXGLOVE_PORT = 8765
|
||||||
DEFAULT_RERUN_PORT = 9090
|
DEFAULT_RERUN_PORT = 9090
|
||||||
|
|
||||||
@@ -299,7 +301,7 @@ def visualize_dataset(
|
|||||||
while True:
|
while True:
|
||||||
time.sleep(1)
|
time.sleep(1)
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
print("Ctrl-C received. Exiting.")
|
logger.info("Ctrl-C received. Exiting.")
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
|
|||||||
@@ -62,7 +62,7 @@ from dataclasses import asdict
|
|||||||
from functools import partial
|
from functools import partial
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from pprint import pformat
|
from pprint import pformat
|
||||||
from typing import Any, TypedDict
|
from typing import TYPE_CHECKING, Any, TypedDict
|
||||||
|
|
||||||
import einops
|
import einops
|
||||||
import gymnasium as gym
|
import gymnasium as gym
|
||||||
@@ -87,7 +87,7 @@ from lerobot.processor import PolicyProcessorPipeline
|
|||||||
from lerobot.types import PolicyAction
|
from lerobot.types import PolicyAction
|
||||||
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
|
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
|
||||||
from lerobot.utils.device_utils import get_safe_torch_device
|
from lerobot.utils.device_utils import get_safe_torch_device
|
||||||
from lerobot.utils.import_utils import register_third_party_plugins
|
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
|
||||||
from lerobot.utils.io_utils import write_video
|
from lerobot.utils.io_utils import write_video
|
||||||
from lerobot.utils.random_utils import set_seed
|
from lerobot.utils.random_utils import set_seed
|
||||||
from lerobot.utils.utils import (
|
from lerobot.utils.utils import (
|
||||||
@@ -95,6 +95,14 @@ from lerobot.utils.utils import (
|
|||||||
inside_slurm,
|
inside_slurm,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING or _peft_available:
|
||||||
|
from peft import PeftModel
|
||||||
|
else:
|
||||||
|
PeftModel = None
|
||||||
|
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
def _env_features_to_dataset_features(env_features: dict) -> dict:
|
def _env_features_to_dataset_features(env_features: dict) -> dict:
|
||||||
"""Convert EnvConfig.features to the dict format expected by LeRobotDataset.create()."""
|
"""Convert EnvConfig.features to the dict format expected by LeRobotDataset.create()."""
|
||||||
@@ -444,15 +452,16 @@ def eval_policy(
|
|||||||
exc = ValueError(
|
exc = ValueError(
|
||||||
f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
|
f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
|
||||||
)
|
)
|
||||||
try:
|
if not _peft_available:
|
||||||
from peft import PeftModel
|
raise exc
|
||||||
|
require_package("peft", extra="peft")
|
||||||
if not isinstance(policy, PeftModel):
|
if not isinstance(policy, PeftModel):
|
||||||
raise exc
|
raise exc
|
||||||
except ImportError:
|
|
||||||
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
|
||||||
@@ -674,6 +683,8 @@ 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
|
||||||
|
|
||||||
|
|
||||||
@@ -791,13 +802,13 @@ def eval_main(cfg: EvalPipelineConfig):
|
|||||||
recording_repo_id=cfg.eval.recording_repo_id,
|
recording_repo_id=cfg.eval.recording_repo_id,
|
||||||
recording_private=cfg.eval.recording_private,
|
recording_private=cfg.eval.recording_private,
|
||||||
)
|
)
|
||||||
print("Overall Aggregated Metrics:")
|
logger.info("Overall Aggregated Metrics:")
|
||||||
print(info["overall"])
|
logger.info(info["overall"])
|
||||||
|
|
||||||
# Print per-suite stats
|
# Print per-suite stats
|
||||||
for task_group, task_group_info in info.items():
|
for task_group, task_group_info in info.items():
|
||||||
print(f"\nAggregated Metrics for {task_group}:")
|
logger.info(f"\nAggregated Metrics for {task_group}:")
|
||||||
print(task_group_info)
|
logger.info(task_group_info)
|
||||||
# Close all vec envs
|
# Close all vec envs
|
||||||
close_envs(envs)
|
close_envs(envs)
|
||||||
|
|
||||||
@@ -1010,40 +1021,48 @@ def eval_policy_all(
|
|||||||
recording_private=recording_private,
|
recording_private=recording_private,
|
||||||
)
|
)
|
||||||
|
|
||||||
if max_parallel_tasks <= 1:
|
# Set the shared policy's mode before launching any workers. Restoring it
|
||||||
prefetch_thread: threading.Thread | None = None
|
# inside individual tasks would let one task enable training mode while
|
||||||
for i, (task_group, task_id, env) in enumerate(tasks):
|
# another task is still evaluating.
|
||||||
if prefetch_thread is not None:
|
was_training = policy.training
|
||||||
prefetch_thread.join()
|
policy.eval()
|
||||||
prefetch_thread = None
|
try:
|
||||||
|
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 = fut.result()
|
tg, tid, metrics = task_runner(task_group, task_id, env)
|
||||||
_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):
|
||||||
|
|||||||
@@ -40,6 +40,7 @@ from PIL import Image
|
|||||||
from lerobot.cameras import ColorMode
|
from lerobot.cameras import ColorMode
|
||||||
from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
|
from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
|
||||||
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
|
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
|
||||||
|
from lerobot.utils.utils import init_logging
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -285,6 +286,8 @@ def save_images_from_all_cameras(
|
|||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
|
init_logging()
|
||||||
|
|
||||||
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."
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -453,9 +453,11 @@ def record(
|
|||||||
encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize,
|
encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize,
|
||||||
)
|
)
|
||||||
|
|
||||||
robot.connect()
|
# Connect the teleoperator before the robot so the robot isn't left idle (and possibly
|
||||||
|
# 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()
|
||||||
|
|
||||||
|
|||||||
@@ -61,6 +61,7 @@ from lerobot.robots import ( # noqa: F401
|
|||||||
earthrover_mini_plus,
|
earthrover_mini_plus,
|
||||||
hope_jr,
|
hope_jr,
|
||||||
koch_follower,
|
koch_follower,
|
||||||
|
lekiwi,
|
||||||
make_robot_from_config,
|
make_robot_from_config,
|
||||||
omx_follower,
|
omx_follower,
|
||||||
openarm_follower,
|
openarm_follower,
|
||||||
|
|||||||
@@ -165,6 +165,7 @@ from lerobot.robots import ( # noqa: F401
|
|||||||
earthrover_mini_plus,
|
earthrover_mini_plus,
|
||||||
hope_jr,
|
hope_jr,
|
||||||
koch_follower,
|
koch_follower,
|
||||||
|
lekiwi,
|
||||||
omx_follower,
|
omx_follower,
|
||||||
openarm_follower,
|
openarm_follower,
|
||||||
reachy2,
|
reachy2,
|
||||||
|
|||||||
@@ -57,7 +57,7 @@ from lerobot.optim.factory import make_optimizer_and_scheduler
|
|||||||
from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
|
from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
|
||||||
from lerobot.rewards import make_reward_pre_post_processors
|
from lerobot.rewards import make_reward_pre_post_processors
|
||||||
from lerobot.utils.collate import lerobot_collate_fn
|
from lerobot.utils.collate import lerobot_collate_fn
|
||||||
from lerobot.utils.import_utils import register_third_party_plugins
|
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
|
||||||
from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
|
from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
|
||||||
from lerobot.utils.random_utils import set_seed
|
from lerobot.utils.random_utils import set_seed
|
||||||
from lerobot.utils.utils import (
|
from lerobot.utils.utils import (
|
||||||
@@ -68,9 +68,24 @@ from lerobot.utils.utils import (
|
|||||||
inside_slurm,
|
inside_slurm,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING or _peft_available:
|
||||||
|
from peft import PeftModel
|
||||||
|
else:
|
||||||
|
PeftModel = None
|
||||||
|
|
||||||
from .lerobot_eval import eval_policy_all
|
from .lerobot_eval import eval_policy_all
|
||||||
|
|
||||||
|
|
||||||
|
def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
|
||||||
|
"""Return worker-only DataLoader options, disabling them for single-process loading."""
|
||||||
|
workers_enabled = cfg.num_workers > 0
|
||||||
|
return {
|
||||||
|
"prefetch_factor": cfg.prefetch_factor if workers_enabled else None,
|
||||||
|
"persistent_workers": cfg.persistent_workers and workers_enabled,
|
||||||
|
"multiprocessing_context": cfg.dataloader_multiprocessing_context if workers_enabled else None,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
def update_policy(
|
def update_policy(
|
||||||
train_metrics: MetricsTracker,
|
train_metrics: MetricsTracker,
|
||||||
policy: PreTrainedPolicy,
|
policy: PreTrainedPolicy,
|
||||||
@@ -197,8 +212,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
|||||||
if cfg.job.is_remote:
|
if cfg.job.is_remote:
|
||||||
return submit_to_hf(cfg)
|
return submit_to_hf(cfg)
|
||||||
|
|
||||||
from lerobot.utils.import_utils import require_package
|
|
||||||
|
|
||||||
require_package("accelerate", extra="training")
|
require_package("accelerate", extra="training")
|
||||||
from accelerate import Accelerator
|
from accelerate import Accelerator
|
||||||
from accelerate.utils import DistributedDataParallelKwargs, DistributedType
|
from accelerate.utils import DistributedDataParallelKwargs, DistributedType
|
||||||
@@ -302,7 +315,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
|||||||
if cfg.peft is not None:
|
if cfg.peft is not None:
|
||||||
if cfg.is_reward_model_training:
|
if cfg.is_reward_model_training:
|
||||||
raise ValueError("PEFT is only supported for policy training. ")
|
raise ValueError("PEFT is only supported for policy training. ")
|
||||||
from peft import PeftModel
|
require_package("peft", extra="peft")
|
||||||
|
|
||||||
if isinstance(policy, PeftModel):
|
if isinstance(policy, PeftModel):
|
||||||
logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.")
|
logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.")
|
||||||
@@ -473,8 +486,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
|||||||
pin_memory=device.type == "cuda",
|
pin_memory=device.type == "cuda",
|
||||||
drop_last=False,
|
drop_last=False,
|
||||||
collate_fn=collate_fn,
|
collate_fn=collate_fn,
|
||||||
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
|
**_dataloader_worker_kwargs(cfg),
|
||||||
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
# Build eval dataloader if a held-out split exists
|
# Build eval dataloader if a held-out split exists
|
||||||
@@ -500,8 +512,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
|||||||
pin_memory=device.type == "cuda",
|
pin_memory=device.type == "cuda",
|
||||||
drop_last=False,
|
drop_last=False,
|
||||||
collate_fn=eval_collate_fn,
|
collate_fn=eval_collate_fn,
|
||||||
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
|
**_dataloader_worker_kwargs(cfg),
|
||||||
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
# Prepare everything with accelerator
|
# Prepare everything with accelerator
|
||||||
|
|||||||
@@ -45,6 +45,7 @@ lerobot-train-tokenizer \
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
import json
|
import json
|
||||||
|
import logging
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
@@ -63,6 +64,9 @@ else:
|
|||||||
from lerobot.configs import NormalizationMode, parser
|
from lerobot.configs import NormalizationMode, parser
|
||||||
from lerobot.datasets import LeRobotDataset
|
from lerobot.datasets import LeRobotDataset
|
||||||
from lerobot.utils.constants import ACTION, OBS_STATE
|
from lerobot.utils.constants import ACTION, OBS_STATE
|
||||||
|
from lerobot.utils.utils import init_logging
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
@@ -274,11 +278,8 @@ def process_episode(args):
|
|||||||
|
|
||||||
return action_chunks
|
return action_chunks
|
||||||
|
|
||||||
except Exception as e:
|
except Exception:
|
||||||
print(f"Error processing episode {ep_idx}: {e}")
|
logger.exception("Error processing episode %s", ep_idx)
|
||||||
import traceback
|
|
||||||
|
|
||||||
traceback.print_exc()
|
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
@@ -300,10 +301,10 @@ def train_fast_tokenizer(
|
|||||||
Returns:
|
Returns:
|
||||||
Trained FAST tokenizer
|
Trained FAST tokenizer
|
||||||
"""
|
"""
|
||||||
print(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
|
logger.info(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
|
||||||
print(f"Action chunk shape: {action_chunks.shape}")
|
logger.info(f"Action chunk shape: {action_chunks.shape}")
|
||||||
print(f"Vocab size: {vocab_size}")
|
logger.info(f"Vocab size: {vocab_size}")
|
||||||
print(f"DCT scale: {scale}")
|
logger.info(f"DCT scale: {scale}")
|
||||||
|
|
||||||
# download the tokenizer source code (not pretrained weights)
|
# download the tokenizer source code (not pretrained weights)
|
||||||
# we'll train a new tokenizer on our own data
|
# we'll train a new tokenizer on our own data
|
||||||
@@ -314,7 +315,7 @@ def train_fast_tokenizer(
|
|||||||
|
|
||||||
# train the new tokenizer on our action data using .fit()
|
# train the new tokenizer on our action data using .fit()
|
||||||
# this trains the BPE tokenizer on DCT coefficients
|
# this trains the BPE tokenizer on DCT coefficients
|
||||||
print("Training new tokenizer (this may take a few minutes)...")
|
logger.info("Training new tokenizer (this may take a few minutes)...")
|
||||||
tokenizer = base_tokenizer.fit(
|
tokenizer = base_tokenizer.fit(
|
||||||
action_data_list,
|
action_data_list,
|
||||||
scale=scale,
|
scale=scale,
|
||||||
@@ -322,21 +323,21 @@ def train_fast_tokenizer(
|
|||||||
time_horizon=action_chunks.shape[1], # action_horizon
|
time_horizon=action_chunks.shape[1], # action_horizon
|
||||||
action_dim=action_chunks.shape[2], # encoded dimensions
|
action_dim=action_chunks.shape[2], # encoded dimensions
|
||||||
)
|
)
|
||||||
print("✓ Tokenizer training complete!")
|
logger.info("✓ Tokenizer training complete!")
|
||||||
|
|
||||||
# validate it works
|
# validate it works
|
||||||
sample_chunk = action_chunks[0]
|
sample_chunk = action_chunks[0]
|
||||||
encoded = tokenizer(sample_chunk[None])[0]
|
encoded = tokenizer(sample_chunk[None])[0]
|
||||||
if isinstance(encoded, list):
|
if isinstance(encoded, list):
|
||||||
encoded = np.array(encoded)
|
encoded = np.array(encoded)
|
||||||
print(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
|
logger.info(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
|
||||||
|
|
||||||
return tokenizer
|
return tokenizer
|
||||||
|
|
||||||
|
|
||||||
def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
|
def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
|
||||||
"""Compute compression statistics."""
|
"""Compute compression statistics."""
|
||||||
print("\nComputing compression statistics...")
|
logger.info("\nComputing compression statistics...")
|
||||||
|
|
||||||
# sample for stats (use max 1000 chunks for speed)
|
# sample for stats (use max 1000 chunks for speed)
|
||||||
sample_size = min(1000, len(action_chunks))
|
sample_size = min(1000, len(action_chunks))
|
||||||
@@ -366,12 +367,12 @@ def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
|
|||||||
"max_token_length": float(np.max(token_lengths)),
|
"max_token_length": float(np.max(token_lengths)),
|
||||||
}
|
}
|
||||||
|
|
||||||
print("Compression Statistics:")
|
logger.info("Compression Statistics:")
|
||||||
print(f" Average compression ratio: {stats['compression_ratio']:.2f}x")
|
logger.info(f" Average compression ratio: {stats['compression_ratio']:.2f}x")
|
||||||
print(f" Mean token length: {stats['mean_token_length']:.1f}")
|
logger.info(f" Mean token length: {stats['mean_token_length']:.1f}")
|
||||||
print(f" P99 token length: {stats['p99_token_length']:.0f}")
|
logger.info(f" P99 token length: {stats['p99_token_length']:.0f}")
|
||||||
print(f" Min token length: {stats['min_token_length']:.0f}")
|
logger.info(f" Min token length: {stats['min_token_length']:.0f}")
|
||||||
print(f" Max token length: {stats['max_token_length']:.0f}")
|
logger.info(f" Max token length: {stats['max_token_length']:.0f}")
|
||||||
|
|
||||||
return stats
|
return stats
|
||||||
|
|
||||||
@@ -385,9 +386,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
cfg: TokenizerTrainingConfig dataclass with all configuration parameters
|
cfg: TokenizerTrainingConfig dataclass with all configuration parameters
|
||||||
"""
|
"""
|
||||||
# load dataset
|
# load dataset
|
||||||
print(f"Loading dataset: {cfg.repo_id}")
|
logger.info(f"Loading dataset: {cfg.repo_id}")
|
||||||
dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root)
|
dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root)
|
||||||
print(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
|
logger.info(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
|
||||||
|
|
||||||
# parse normalization mode
|
# parse normalization mode
|
||||||
try:
|
try:
|
||||||
@@ -397,7 +398,7 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
f"Invalid normalization_mode: {cfg.normalization_mode}. "
|
f"Invalid normalization_mode: {cfg.normalization_mode}. "
|
||||||
f"Must be one of: {', '.join([m.value for m in NormalizationMode])}"
|
f"Must be one of: {', '.join([m.value for m in NormalizationMode])}"
|
||||||
) from err
|
) from err
|
||||||
print(f"Normalization mode: {norm_mode.value}")
|
logger.info(f"Normalization mode: {norm_mode.value}")
|
||||||
|
|
||||||
# parse encoded dimensions
|
# parse encoded dimensions
|
||||||
encoded_dim_ranges = []
|
encoded_dim_ranges = []
|
||||||
@@ -406,38 +407,38 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
encoded_dim_ranges.append((start, end))
|
encoded_dim_ranges.append((start, end))
|
||||||
|
|
||||||
total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges)
|
total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges)
|
||||||
print(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
|
logger.info(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
|
||||||
|
|
||||||
# parse relative dimensions
|
# parse relative dimensions
|
||||||
relative_dim_list = None
|
relative_dim_list = None
|
||||||
if cfg.relative_dims is not None and cfg.relative_dims.strip():
|
if cfg.relative_dims is not None and cfg.relative_dims.strip():
|
||||||
relative_dim_list = [int(d.strip()) for d in cfg.relative_dims.split(",")]
|
relative_dim_list = [int(d.strip()) for d in cfg.relative_dims.split(",")]
|
||||||
print(f"Relative dimensions: {relative_dim_list}")
|
logger.info(f"Relative dimensions: {relative_dim_list}")
|
||||||
else:
|
else:
|
||||||
print("No relative dimensions specified")
|
logger.info("No relative dimensions specified")
|
||||||
|
|
||||||
print(f"Use relative transform: {cfg.use_relative_transform}")
|
logger.info(f"Use relative transform: {cfg.use_relative_transform}")
|
||||||
if cfg.use_relative_transform and (relative_dim_list is None or len(relative_dim_list) == 0):
|
if cfg.use_relative_transform and (relative_dim_list is None or len(relative_dim_list) == 0):
|
||||||
print(
|
logger.warning(
|
||||||
"Warning: use_relative_transform=True but no relative_dims specified. "
|
"Warning: use_relative_transform=True but no relative_dims specified. "
|
||||||
"No relative transform will be applied."
|
"No relative transform will be applied."
|
||||||
)
|
)
|
||||||
|
|
||||||
print(f"Action horizon: {cfg.action_horizon}")
|
logger.info(f"Action horizon: {cfg.action_horizon}")
|
||||||
print(f"State key: {cfg.state_key}")
|
logger.info(f"State key: {cfg.state_key}")
|
||||||
|
|
||||||
# determine episodes to process
|
# determine episodes to process
|
||||||
num_episodes = dataset.num_episodes
|
num_episodes = dataset.num_episodes
|
||||||
if cfg.max_episodes is not None:
|
if cfg.max_episodes is not None:
|
||||||
num_episodes = min(cfg.max_episodes, num_episodes)
|
num_episodes = min(cfg.max_episodes, num_episodes)
|
||||||
|
|
||||||
print(f"Processing {num_episodes} episodes...")
|
logger.info(f"Processing {num_episodes} episodes...")
|
||||||
|
|
||||||
# process episodes sequentially (to avoid pickling issues with dataset)
|
# process episodes sequentially (to avoid pickling issues with dataset)
|
||||||
all_chunks = []
|
all_chunks = []
|
||||||
for ep_idx in range(num_episodes):
|
for ep_idx in range(num_episodes):
|
||||||
if ep_idx % 10 == 0:
|
if ep_idx % 10 == 0:
|
||||||
print(f" Processing episode {ep_idx}/{num_episodes}...")
|
logger.info(f" Processing episode {ep_idx}/{num_episodes}...")
|
||||||
|
|
||||||
chunks = process_episode(
|
chunks = process_episode(
|
||||||
(
|
(
|
||||||
@@ -455,19 +456,19 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
|
|
||||||
# concatenate all chunks
|
# concatenate all chunks
|
||||||
all_chunks = np.concatenate(all_chunks, axis=0)
|
all_chunks = np.concatenate(all_chunks, axis=0)
|
||||||
print(f"Collected {len(all_chunks)} action chunks")
|
logger.info(f"Collected {len(all_chunks)} action chunks")
|
||||||
|
|
||||||
# extract only encoded dimensions FIRST (before normalization)
|
# extract only encoded dimensions FIRST (before normalization)
|
||||||
encoded_chunks = []
|
encoded_chunks = []
|
||||||
for start, end in encoded_dim_ranges:
|
for start, end in encoded_dim_ranges:
|
||||||
encoded_chunks.append(all_chunks[:, :, start:end])
|
encoded_chunks.append(all_chunks[:, :, start:end])
|
||||||
encoded_chunks = np.concatenate(encoded_chunks, axis=-1) # [N, H, D_encoded]
|
encoded_chunks = np.concatenate(encoded_chunks, axis=-1) # [N, H, D_encoded]
|
||||||
print(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
|
logger.info(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
|
||||||
|
|
||||||
# apply normalization to encoded dimensions
|
# apply normalization to encoded dimensions
|
||||||
print("\nBefore normalization - overall stats:")
|
logger.info("\nBefore normalization - overall stats:")
|
||||||
print(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
|
logger.info(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
|
||||||
print(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
|
logger.info(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
|
||||||
|
|
||||||
# get normalization stats from dataset
|
# get normalization stats from dataset
|
||||||
norm_stats = dataset.meta.stats
|
norm_stats = dataset.meta.stats
|
||||||
@@ -489,9 +490,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
encoded_stats[stat_name] = stat_array[encoded_dim_indices]
|
encoded_stats[stat_name] = stat_array[encoded_dim_indices]
|
||||||
|
|
||||||
if encoded_stats:
|
if encoded_stats:
|
||||||
print(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
|
logger.info(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
|
||||||
for stat_name, stat_values in encoded_stats.items():
|
for stat_name, stat_values in encoded_stats.items():
|
||||||
print(
|
logger.info(
|
||||||
f" {stat_name}: shape={stat_values.shape}, "
|
f" {stat_name}: shape={stat_values.shape}, "
|
||||||
f"range=[{np.min(stat_values):.4f}, {np.max(stat_values):.4f}]"
|
f"range=[{np.min(stat_values):.4f}, {np.max(stat_values):.4f}]"
|
||||||
)
|
)
|
||||||
@@ -499,27 +500,27 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
# apply normalization based on mode
|
# apply normalization based on mode
|
||||||
try:
|
try:
|
||||||
encoded_chunks = apply_normalization(encoded_chunks, encoded_stats, norm_mode, eps=1e-8)
|
encoded_chunks = apply_normalization(encoded_chunks, encoded_stats, norm_mode, eps=1e-8)
|
||||||
print(f"\nApplied {norm_mode.value} normalization")
|
logger.info(f"\nApplied {norm_mode.value} normalization")
|
||||||
except ValueError as e:
|
except ValueError as e:
|
||||||
print(f"Warning: {e}. Using raw actions without normalization.")
|
logger.warning(f"Warning: {e}. Using raw actions without normalization.")
|
||||||
|
|
||||||
print("\nAfter normalization - overall stats:")
|
logger.info("\nAfter normalization - overall stats:")
|
||||||
print(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
|
logger.info(f" Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
|
||||||
print(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
|
logger.info(f" Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
|
||||||
|
|
||||||
print("\nPer-dimension stats (after normalization):")
|
logger.info("\nPer-dimension stats (after normalization):")
|
||||||
for d in range(encoded_chunks.shape[-1]):
|
for d in range(encoded_chunks.shape[-1]):
|
||||||
dim_data = encoded_chunks[:, :, d]
|
dim_data = encoded_chunks[:, :, d]
|
||||||
print(
|
logger.info(
|
||||||
f" Dim {d}: min={np.min(dim_data):7.4f}, max={np.max(dim_data):7.4f}, "
|
f" Dim {d}: min={np.min(dim_data):7.4f}, max={np.max(dim_data):7.4f}, "
|
||||||
f"mean={np.mean(dim_data):7.4f}, std={np.std(dim_data):7.4f}"
|
f"mean={np.mean(dim_data):7.4f}, std={np.std(dim_data):7.4f}"
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
print("Warning: Could not extract stats for encoded dimensions, using raw actions")
|
logger.warning("Warning: Could not extract stats for encoded dimensions, using raw actions")
|
||||||
else:
|
else:
|
||||||
print("Warning: No normalization stats found in dataset, using raw actions")
|
logger.warning("Warning: No normalization stats found in dataset, using raw actions")
|
||||||
|
|
||||||
print(f"Encoded chunks shape: {encoded_chunks.shape}")
|
logger.info(f"Encoded chunks shape: {encoded_chunks.shape}")
|
||||||
|
|
||||||
# train FAST tokenizer
|
# train FAST tokenizer
|
||||||
tokenizer = train_fast_tokenizer(
|
tokenizer = train_fast_tokenizer(
|
||||||
@@ -561,8 +562,8 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
with open(output_path / "metadata.json", "w") as f:
|
with open(output_path / "metadata.json", "w") as f:
|
||||||
json.dump(metadata, f, indent=2)
|
json.dump(metadata, f, indent=2)
|
||||||
|
|
||||||
print(f"\nSaved FAST tokenizer to {output_path}")
|
logger.info(f"\nSaved FAST tokenizer to {output_path}")
|
||||||
print(f"Metadata: {json.dumps(metadata, indent=2)}")
|
logger.info(f"Metadata: {json.dumps(metadata, indent=2)}")
|
||||||
|
|
||||||
# push to Hugging Face Hub if requested
|
# push to Hugging Face Hub if requested
|
||||||
if cfg.push_to_hub:
|
if cfg.push_to_hub:
|
||||||
@@ -570,10 +571,10 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
hub_repo_id = cfg.hub_repo_id
|
hub_repo_id = cfg.hub_repo_id
|
||||||
if hub_repo_id is None:
|
if hub_repo_id is None:
|
||||||
hub_repo_id = output_path.name
|
hub_repo_id = output_path.name
|
||||||
print(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
|
logger.info(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
|
||||||
|
|
||||||
print(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
|
logger.info(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
|
||||||
print(f" Private: {cfg.hub_private}")
|
logger.info(f" Private: {cfg.hub_private}")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# use the tokenizer's push_to_hub method
|
# use the tokenizer's push_to_hub method
|
||||||
@@ -593,14 +594,15 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
|
|||||||
commit_message="Upload tokenizer metadata",
|
commit_message="Upload tokenizer metadata",
|
||||||
)
|
)
|
||||||
|
|
||||||
print(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
|
logger.info(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"Error pushing to hub: {e}")
|
logger.error(f"Error pushing to hub: {e}")
|
||||||
print(" Make sure you're logged in with `huggingface-cli login`")
|
logger.error(" Make sure you're logged in with `huggingface-cli login`")
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
"""CLI entry point that parses arguments and runs the tokenizer training."""
|
"""CLI entry point that parses arguments and runs the tokenizer training."""
|
||||||
|
init_logging()
|
||||||
train_tokenizer()
|
train_tokenizer()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -16,11 +16,39 @@
|
|||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
|
import multiprocessing
|
||||||
import os
|
import os
|
||||||
import signal
|
import signal
|
||||||
import sys
|
import sys
|
||||||
|
|
||||||
|
|
||||||
|
def ensure_multiprocessing_start_method(start_method: str | None) -> None:
|
||||||
|
"""Set a multiprocessing start method once, or verify the existing method matches.
|
||||||
|
|
||||||
|
Passing ``None`` leaves Python's process-wide default untouched. This is useful
|
||||||
|
when LeRobot is embedded in an application that owns multiprocessing setup.
|
||||||
|
"""
|
||||||
|
if start_method is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
available_methods = multiprocessing.get_all_start_methods()
|
||||||
|
if start_method not in available_methods:
|
||||||
|
raise ValueError(
|
||||||
|
f"Multiprocessing start method must be one of {available_methods} on this platform, "
|
||||||
|
f"got {start_method!r}."
|
||||||
|
)
|
||||||
|
|
||||||
|
current_method = multiprocessing.get_start_method(allow_none=True)
|
||||||
|
if current_method is None:
|
||||||
|
multiprocessing.set_start_method(start_method)
|
||||||
|
elif current_method != start_method:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"Multiprocessing start method is already {current_method!r}; cannot change it to "
|
||||||
|
f"{start_method!r}. Set the configured multiprocessing context to null to keep the "
|
||||||
|
"application's existing method, or launch LeRobot in a fresh process."
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class ProcessSignalHandler:
|
class ProcessSignalHandler:
|
||||||
"""Utility class to attach graceful shutdown signal handlers.
|
"""Utility class to attach graceful shutdown signal handlers.
|
||||||
|
|
||||||
|
|||||||
@@ -20,7 +20,7 @@
|
|||||||
# ```
|
# ```
|
||||||
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from unittest.mock import patch
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
import cv2
|
import cv2
|
||||||
import numpy as np
|
import numpy as np
|
||||||
@@ -123,6 +123,73 @@ def test_invalid_width_connect():
|
|||||||
camera.connect(warmup=False)
|
camera.connect(warmup=False)
|
||||||
|
|
||||||
|
|
||||||
|
def test_connect_cleans_up_after_settings_failure_and_allows_retry():
|
||||||
|
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=0)
|
||||||
|
camera = OpenCVCamera(config)
|
||||||
|
opened_captures = []
|
||||||
|
|
||||||
|
def fail_settings():
|
||||||
|
opened_captures.append(camera.videocapture)
|
||||||
|
raise RuntimeError("settings failed")
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch.object(camera, "_configure_capture_settings", side_effect=fail_settings),
|
||||||
|
pytest.raises(RuntimeError, match="settings failed"),
|
||||||
|
):
|
||||||
|
camera.connect(warmup=False)
|
||||||
|
|
||||||
|
assert camera.videocapture is None
|
||||||
|
assert camera.thread is None
|
||||||
|
assert not camera.is_connected
|
||||||
|
assert opened_captures[0] is not None
|
||||||
|
assert not opened_captures[0].isOpened()
|
||||||
|
|
||||||
|
camera.connect(warmup=False)
|
||||||
|
assert camera.is_connected
|
||||||
|
camera.disconnect()
|
||||||
|
|
||||||
|
|
||||||
|
def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
|
||||||
|
config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=1)
|
||||||
|
camera = OpenCVCamera(config)
|
||||||
|
read_threads = []
|
||||||
|
|
||||||
|
def fail_warmup(*_args, **_kwargs):
|
||||||
|
read_threads.append(camera.thread)
|
||||||
|
raise TimeoutError("no frame")
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch.object(camera, "async_read", side_effect=fail_warmup),
|
||||||
|
pytest.raises(TimeoutError, match="no frame"),
|
||||||
|
):
|
||||||
|
camera.connect()
|
||||||
|
|
||||||
|
assert camera.videocapture is None
|
||||||
|
assert camera.thread is None
|
||||||
|
assert not camera.is_connected
|
||||||
|
assert read_threads[0] is not None
|
||||||
|
assert not read_threads[0].is_alive()
|
||||||
|
|
||||||
|
camera.connect(warmup=False)
|
||||||
|
assert camera.is_connected
|
||||||
|
camera.disconnect()
|
||||||
|
|
||||||
|
|
||||||
|
def test_find_cameras_releases_unopened_handles():
|
||||||
|
module_path = OpenCVCamera.__module__
|
||||||
|
unopened_capture = MagicMock()
|
||||||
|
unopened_capture.isOpened.return_value = False
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch(f"{module_path}.platform.system", return_value="Darwin"),
|
||||||
|
patch(f"{module_path}.MAX_OPENCV_INDEX", 1),
|
||||||
|
patch(f"{module_path}.cv2.VideoCapture", return_value=unopened_capture),
|
||||||
|
):
|
||||||
|
assert OpenCVCamera.find_cameras() == []
|
||||||
|
|
||||||
|
unopened_capture.release.assert_called_once_with()
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
|
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
|
||||||
def test_read(index_or_path):
|
def test_read(index_or_path):
|
||||||
config = OpenCVCameraConfig(index_or_path=index_or_path, warmup_s=0)
|
config = OpenCVCameraConfig(index_or_path=index_or_path, warmup_s=0)
|
||||||
|
|||||||
@@ -20,7 +20,7 @@
|
|||||||
# ```
|
# ```
|
||||||
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from unittest.mock import patch
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import pytest
|
import pytest
|
||||||
@@ -30,6 +30,8 @@ from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnected
|
|||||||
|
|
||||||
pytest.importorskip("pyrealsense2")
|
pytest.importorskip("pyrealsense2")
|
||||||
|
|
||||||
|
import pyrealsense2 as rs
|
||||||
|
|
||||||
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
|
from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
|
||||||
|
|
||||||
TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
|
TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
|
||||||
@@ -61,6 +63,17 @@ def test_abc_implementation():
|
|||||||
_ = RealSenseCamera(config)
|
_ = RealSenseCamera(config)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("option", ["exposure", "gain", "white_balance"])
|
||||||
|
def test_manual_color_option_requires_rgb(option):
|
||||||
|
with pytest.raises(ValueError, match="use_rgb=True"):
|
||||||
|
RealSenseCameraConfig(
|
||||||
|
serial_number_or_name="042",
|
||||||
|
use_rgb=False,
|
||||||
|
use_depth=True,
|
||||||
|
**{option: 100},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def test_connect():
|
def test_connect():
|
||||||
config = RealSenseCameraConfig(serial_number_or_name="042", warmup_s=0)
|
config = RealSenseCameraConfig(serial_number_or_name="042", warmup_s=0)
|
||||||
|
|
||||||
@@ -83,6 +96,27 @@ def test_connect_invalid_camera_path(patch_realsense):
|
|||||||
camera.connect(warmup=False)
|
camera.connect(warmup=False)
|
||||||
|
|
||||||
|
|
||||||
|
def test_connect_cleans_up_when_sensor_configuration_fails():
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
pipeline = MagicMock()
|
||||||
|
pipeline.start.return_value = MagicMock()
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("lerobot.cameras.realsense.camera_realsense.rs.pipeline", return_value=pipeline),
|
||||||
|
patch.object(camera, "_configure_rs_pipeline_config"),
|
||||||
|
patch.object(camera, "_configure_capture_settings"),
|
||||||
|
patch.object(camera, "_configure_sensor_options", side_effect=ValueError("invalid exposure")),
|
||||||
|
pytest.raises(ValueError, match="invalid exposure"),
|
||||||
|
):
|
||||||
|
camera.connect(warmup=False)
|
||||||
|
|
||||||
|
pipeline.stop.assert_called_once_with()
|
||||||
|
assert camera.rs_pipeline is None
|
||||||
|
assert camera.rs_profile is None
|
||||||
|
assert not camera.is_connected
|
||||||
|
|
||||||
|
|
||||||
def test_invalid_width_connect():
|
def test_invalid_width_connect():
|
||||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=99999, height=480, fps=30)
|
config = RealSenseCameraConfig(serial_number_or_name="042", width=99999, height=480, fps=30)
|
||||||
camera = RealSenseCamera(config)
|
camera = RealSenseCamera(config)
|
||||||
@@ -91,6 +125,33 @@ def test_invalid_width_connect():
|
|||||||
camera.connect(warmup=False)
|
camera.connect(warmup=False)
|
||||||
|
|
||||||
|
|
||||||
|
def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30)
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
read_threads = []
|
||||||
|
|
||||||
|
def fail_warmup(*_args, **_kwargs):
|
||||||
|
read_threads.append(camera.thread)
|
||||||
|
raise TimeoutError("no frame")
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch.object(camera, "async_read", side_effect=fail_warmup),
|
||||||
|
pytest.raises(TimeoutError, match="no frame"),
|
||||||
|
):
|
||||||
|
camera.connect()
|
||||||
|
|
||||||
|
assert camera.rs_pipeline is None
|
||||||
|
assert camera.rs_profile is None
|
||||||
|
assert camera.thread is None
|
||||||
|
assert not camera.is_connected
|
||||||
|
assert read_threads[0] is not None
|
||||||
|
assert not read_threads[0].is_alive()
|
||||||
|
|
||||||
|
camera.connect(warmup=False)
|
||||||
|
assert camera.is_connected
|
||||||
|
camera.disconnect()
|
||||||
|
|
||||||
|
|
||||||
def test_read():
|
def test_read():
|
||||||
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
|
config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
|
||||||
with RealSenseCamera(config) as camera:
|
with RealSenseCamera(config) as camera:
|
||||||
@@ -228,6 +289,203 @@ def test_read_latest_too_old():
|
|||||||
_ = camera.read_latest(max_age_ms=0) # immediately too old
|
_ = camera.read_latest(max_age_ms=0) # immediately too old
|
||||||
|
|
||||||
|
|
||||||
|
def _make_mock_sensor(name: str, supported_options: set | None = None) -> MagicMock:
|
||||||
|
"""Build a fake rs.sensor that reports a name and a configurable supported-options set."""
|
||||||
|
supported = supported_options if supported_options is not None else set()
|
||||||
|
sensor = MagicMock()
|
||||||
|
sensor.get_info.return_value = name
|
||||||
|
sensor.supports.side_effect = lambda opt: opt in supported
|
||||||
|
return sensor
|
||||||
|
|
||||||
|
|
||||||
|
def _attach_mock_color_sensor(camera: RealSenseCamera, sensor: MagicMock) -> None:
|
||||||
|
"""Wire camera.rs_profile so _get_color_sensor finds the given sensor."""
|
||||||
|
profile = MagicMock()
|
||||||
|
device = MagicMock()
|
||||||
|
device.query_sensors.return_value = [sensor]
|
||||||
|
profile.get_device.return_value = device
|
||||||
|
camera.rs_profile = profile
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_color_sensor_prefers_rgb_camera():
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
rgb = _make_mock_sensor("RGB Camera")
|
||||||
|
stereo = _make_mock_sensor("Stereo Module")
|
||||||
|
profile = MagicMock()
|
||||||
|
device = MagicMock()
|
||||||
|
device.query_sensors.return_value = [stereo, rgb]
|
||||||
|
profile.get_device.return_value = device
|
||||||
|
camera.rs_profile = profile
|
||||||
|
|
||||||
|
assert camera._get_color_sensor() is rgb
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_color_sensor_falls_back_to_stereo_module():
|
||||||
|
"""D405 has no separate RGB module; color comes from Stereo Module."""
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
stereo = _make_mock_sensor("Stereo Module")
|
||||||
|
_attach_mock_color_sensor(camera, stereo)
|
||||||
|
|
||||||
|
assert camera._get_color_sensor() is stereo
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_color_sensor_raises_with_available_sensors():
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
other = _make_mock_sensor("Motion Module")
|
||||||
|
_attach_mock_color_sensor(camera, other)
|
||||||
|
|
||||||
|
with pytest.raises(RuntimeError, match="Motion Module"):
|
||||||
|
camera._get_color_sensor()
|
||||||
|
|
||||||
|
|
||||||
|
def test_configure_sensor_options_skipped_when_none():
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042")
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
with patch.object(RealSenseCamera, "_get_color_sensor") as mock_get:
|
||||||
|
camera._configure_sensor_options()
|
||||||
|
mock_get.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
def test_configure_sensor_options_applies_all_values():
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120, gain=64, white_balance=4600)
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
sensor = _make_mock_sensor(
|
||||||
|
"RGB Camera",
|
||||||
|
supported_options={
|
||||||
|
rs.option.enable_auto_exposure,
|
||||||
|
rs.option.exposure,
|
||||||
|
rs.option.gain,
|
||||||
|
rs.option.enable_auto_white_balance,
|
||||||
|
rs.option.white_balance,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
_attach_mock_color_sensor(camera, sensor)
|
||||||
|
|
||||||
|
camera._configure_sensor_options()
|
||||||
|
|
||||||
|
sensor.set_option.assert_any_call(rs.option.enable_auto_exposure, 0)
|
||||||
|
sensor.set_option.assert_any_call(rs.option.exposure, 120)
|
||||||
|
sensor.set_option.assert_any_call(rs.option.gain, 64)
|
||||||
|
sensor.set_option.assert_any_call(rs.option.enable_auto_white_balance, 0)
|
||||||
|
sensor.set_option.assert_any_call(rs.option.white_balance, 4600)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("config_field", "option", "label"),
|
||||||
|
[
|
||||||
|
("exposure", rs.option.exposure, "exposure"),
|
||||||
|
("gain", rs.option.gain, "gain"),
|
||||||
|
("white_balance", rs.option.white_balance, "white balance"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_configure_sensor_options_raises_when_requested_option_is_unsupported(config_field, option, label):
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: 100})
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
sensor = _make_mock_sensor("RGB Camera", supported_options=set())
|
||||||
|
_attach_mock_color_sensor(camera, sensor)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match=label):
|
||||||
|
camera._configure_sensor_options()
|
||||||
|
|
||||||
|
sensor.supports.assert_any_call(option)
|
||||||
|
sensor.set_option.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("config_field", "option", "value"),
|
||||||
|
[
|
||||||
|
("exposure", rs.option.exposure, 120),
|
||||||
|
("gain", rs.option.gain, 64),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_configure_sensor_options_exposure_or_gain_disables_auto_exposure(config_field, option, value):
|
||||||
|
"""white_balance=None should not touch auto white balance."""
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: value})
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
sensor = _make_mock_sensor(
|
||||||
|
"RGB Camera",
|
||||||
|
supported_options={rs.option.enable_auto_exposure, option},
|
||||||
|
)
|
||||||
|
_attach_mock_color_sensor(camera, sensor)
|
||||||
|
|
||||||
|
camera._configure_sensor_options()
|
||||||
|
|
||||||
|
calls = [call.args for call in sensor.set_option.call_args_list]
|
||||||
|
assert (rs.option.enable_auto_exposure, 0) in calls
|
||||||
|
assert (option, value) in calls
|
||||||
|
for opt, _ in calls:
|
||||||
|
assert opt != rs.option.enable_auto_white_balance
|
||||||
|
assert opt != rs.option.white_balance
|
||||||
|
|
||||||
|
|
||||||
|
def test_configure_sensor_options_warns_when_auto_exposure_control_is_unsupported(caplog):
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.exposure})
|
||||||
|
_attach_mock_color_sensor(camera, sensor)
|
||||||
|
|
||||||
|
with caplog.at_level("WARNING"):
|
||||||
|
camera._configure_sensor_options()
|
||||||
|
|
||||||
|
sensor.set_option.assert_called_once_with(rs.option.exposure, 120)
|
||||||
|
assert "does not support disabling auto-exposure" in caplog.text
|
||||||
|
|
||||||
|
|
||||||
|
def test_configure_sensor_options_warns_when_auto_white_balance_control_is_unsupported(caplog):
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", white_balance=4600)
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.white_balance})
|
||||||
|
_attach_mock_color_sensor(camera, sensor)
|
||||||
|
|
||||||
|
with caplog.at_level("WARNING"):
|
||||||
|
camera._configure_sensor_options()
|
||||||
|
|
||||||
|
sensor.set_option.assert_called_once_with(rs.option.white_balance, 4600)
|
||||||
|
assert "does not support disabling auto white balance" in caplog.text
|
||||||
|
|
||||||
|
|
||||||
|
def test_configure_sensor_options_out_of_range_raises_value_error():
|
||||||
|
"""set_option errors should be re-raised as ValueError with range diagnostics."""
|
||||||
|
config = RealSenseCameraConfig(serial_number_or_name="042", exposure=999999)
|
||||||
|
camera = RealSenseCamera(config)
|
||||||
|
|
||||||
|
sensor = _make_mock_sensor(
|
||||||
|
"RGB Camera",
|
||||||
|
supported_options={rs.option.enable_auto_exposure, rs.option.exposure},
|
||||||
|
)
|
||||||
|
|
||||||
|
def fake_set_option(option, value):
|
||||||
|
if option == rs.option.exposure:
|
||||||
|
raise RuntimeError("value out of range")
|
||||||
|
|
||||||
|
sensor.set_option.side_effect = fake_set_option
|
||||||
|
|
||||||
|
option_range = MagicMock(min=1, max=10000, step=1, default=156)
|
||||||
|
sensor.get_option_range.return_value = option_range
|
||||||
|
|
||||||
|
_attach_mock_color_sensor(camera, sensor)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError, match="exposure") as exc_info:
|
||||||
|
camera._configure_sensor_options()
|
||||||
|
|
||||||
|
msg = str(exc_info.value)
|
||||||
|
assert "999999" in msg
|
||||||
|
assert "min=1" in msg
|
||||||
|
assert "max=10000" in msg
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"rotation",
|
"rotation",
|
||||||
[
|
[
|
||||||
|
|||||||
@@ -204,6 +204,38 @@ 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(
|
||||||
|
|||||||
@@ -0,0 +1,79 @@
|
|||||||
|
# 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 sys
|
||||||
|
from types import SimpleNamespace
|
||||||
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
import lerobot.policies.factory as policy_factory
|
||||||
|
|
||||||
|
|
||||||
|
def test_make_policy_keeps_peft_adapter_and_base_revisions_separate(monkeypatch):
|
||||||
|
cfg = SimpleNamespace(
|
||||||
|
type="mock",
|
||||||
|
device="cpu",
|
||||||
|
pretrained_path="user/adapter",
|
||||||
|
pretrained_revision="adapter-sha",
|
||||||
|
use_peft=True,
|
||||||
|
input_features={},
|
||||||
|
output_features={},
|
||||||
|
)
|
||||||
|
dataset_meta = SimpleNamespace(features={}, stats={})
|
||||||
|
|
||||||
|
base_policy = torch.nn.Linear(1, 1)
|
||||||
|
policy_from_pretrained = MagicMock(return_value=base_policy)
|
||||||
|
policy_class = SimpleNamespace(from_pretrained=policy_from_pretrained)
|
||||||
|
monkeypatch.setattr(policy_factory, "get_policy_class", lambda _: policy_class)
|
||||||
|
monkeypatch.setattr(policy_factory, "dataset_to_policy_features", lambda _: {})
|
||||||
|
monkeypatch.setattr(policy_factory, "validate_visual_features_consistency", lambda *args: None)
|
||||||
|
|
||||||
|
peft_config = SimpleNamespace(
|
||||||
|
base_model_name_or_path="user/base-policy",
|
||||||
|
revision="base-sha",
|
||||||
|
)
|
||||||
|
peft_config_from_pretrained = MagicMock(return_value=peft_config)
|
||||||
|
adapted_policy = torch.nn.Linear(1, 1)
|
||||||
|
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
|
||||||
|
monkeypatch.setitem(
|
||||||
|
sys.modules,
|
||||||
|
"peft",
|
||||||
|
SimpleNamespace(
|
||||||
|
PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained),
|
||||||
|
PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
policy = policy_factory.make_policy(cfg, ds_meta=dataset_meta)
|
||||||
|
|
||||||
|
assert policy is adapted_policy
|
||||||
|
peft_config_from_pretrained.assert_called_once_with(
|
||||||
|
"user/adapter",
|
||||||
|
revision="adapter-sha",
|
||||||
|
)
|
||||||
|
policy_from_pretrained.assert_called_once_with(
|
||||||
|
config=cfg,
|
||||||
|
dataset_stats=dataset_meta.stats,
|
||||||
|
dataset_meta=dataset_meta,
|
||||||
|
pretrained_name_or_path="user/base-policy",
|
||||||
|
revision="base-sha",
|
||||||
|
)
|
||||||
|
peft_model_from_pretrained.assert_called_once_with(
|
||||||
|
base_policy,
|
||||||
|
"user/adapter",
|
||||||
|
config=peft_config,
|
||||||
|
revision="adapter-sha",
|
||||||
|
is_trainable=True,
|
||||||
|
)
|
||||||
@@ -113,6 +113,7 @@ def test_gaussian_actor_config_default_initialization():
|
|||||||
# Concurrency configuration
|
# Concurrency configuration
|
||||||
assert config.concurrency.actor == "threads"
|
assert config.concurrency.actor == "threads"
|
||||||
assert config.concurrency.learner == "threads"
|
assert config.concurrency.learner == "threads"
|
||||||
|
assert config.concurrency.multiprocessing_context == "spawn"
|
||||||
|
|
||||||
assert isinstance(config.actor_network_kwargs, ActorNetworkConfig)
|
assert isinstance(config.actor_network_kwargs, ActorNetworkConfig)
|
||||||
assert isinstance(config.policy_kwargs, PolicyConfig)
|
assert isinstance(config.policy_kwargs, PolicyConfig)
|
||||||
@@ -152,6 +153,7 @@ def test_concurrency_config():
|
|||||||
config = ConcurrencyConfig()
|
config = ConcurrencyConfig()
|
||||||
assert config.actor == "threads"
|
assert config.actor == "threads"
|
||||||
assert config.learner == "threads"
|
assert config.learner == "threads"
|
||||||
|
assert config.multiprocessing_context == "spawn"
|
||||||
|
|
||||||
|
|
||||||
def test_gaussian_actor_config_custom_initialization():
|
def test_gaussian_actor_config_custom_initialization():
|
||||||
|
|||||||
@@ -26,8 +26,17 @@ import tempfile
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
import torch
|
||||||
|
from safetensors.torch import save_file
|
||||||
|
|
||||||
from lerobot.processor.pipeline import DataProcessorPipeline, ProcessorMigrationError
|
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||||
|
from lerobot.processor.pipeline import (
|
||||||
|
DataProcessorPipeline,
|
||||||
|
ProcessorMigrationError,
|
||||||
|
ProcessorStep,
|
||||||
|
ProcessorStepRegistry,
|
||||||
|
)
|
||||||
|
from lerobot.types import EnvTransition
|
||||||
|
|
||||||
# Simplified Config Loading Tests
|
# Simplified Config Loading Tests
|
||||||
|
|
||||||
@@ -98,6 +107,140 @@ def test_load_config_nonexistent_path_tries_hub():
|
|||||||
DataProcessorPipeline._load_config("nonexistent/path", "processor.json", {})
|
DataProcessorPipeline._load_config("nonexistent/path", "processor.json", {})
|
||||||
|
|
||||||
|
|
||||||
|
def test_from_pretrained_local_directory_missing_state_does_not_call_hub(monkeypatch):
|
||||||
|
"""Local processor dirs must fail locally when a state file is missing."""
|
||||||
|
|
||||||
|
@ProcessorStepRegistry.register("local_missing_state_step")
|
||||||
|
class LocalMissingStateStep(ProcessorStep):
|
||||||
|
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||||
|
return transition
|
||||||
|
|
||||||
|
def transform_features(
|
||||||
|
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||||
|
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||||
|
return features
|
||||||
|
|
||||||
|
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||||||
|
pass
|
||||||
|
|
||||||
|
try:
|
||||||
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||||
|
tmp_path = Path(tmp_dir)
|
||||||
|
config = {
|
||||||
|
"name": "LocalMissingStatePipeline",
|
||||||
|
"steps": [{"registry_name": "local_missing_state_step", "state_file": "missing.safetensors"}],
|
||||||
|
}
|
||||||
|
(tmp_path / "processor.json").write_text(json.dumps(config))
|
||||||
|
|
||||||
|
def fail_hub_download(*args, **kwargs):
|
||||||
|
pytest.fail("local missing processor state should not call hf_hub_download")
|
||||||
|
|
||||||
|
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fail_hub_download)
|
||||||
|
|
||||||
|
with pytest.raises(FileNotFoundError, match="missing.safetensors.*local processor pipeline"):
|
||||||
|
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="processor.json")
|
||||||
|
finally:
|
||||||
|
ProcessorStepRegistry.unregister("local_missing_state_step")
|
||||||
|
|
||||||
|
|
||||||
|
def test_from_pretrained_local_config_file_missing_state_does_not_call_hub(monkeypatch):
|
||||||
|
"""Local single-file processor configs must also keep missing state resolution local."""
|
||||||
|
|
||||||
|
@ProcessorStepRegistry.register("local_file_missing_state_step")
|
||||||
|
class LocalFileMissingStateStep(ProcessorStep):
|
||||||
|
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||||
|
return transition
|
||||||
|
|
||||||
|
def transform_features(
|
||||||
|
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||||
|
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||||
|
return features
|
||||||
|
|
||||||
|
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||||||
|
pass
|
||||||
|
|
||||||
|
try:
|
||||||
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||||
|
tmp_path = Path(tmp_dir)
|
||||||
|
config_path = tmp_path / "processor.json"
|
||||||
|
config = {
|
||||||
|
"name": "LocalFileMissingStatePipeline",
|
||||||
|
"steps": [
|
||||||
|
{"registry_name": "local_file_missing_state_step", "state_file": "missing.safetensors"}
|
||||||
|
],
|
||||||
|
}
|
||||||
|
config_path.write_text(json.dumps(config))
|
||||||
|
|
||||||
|
def fail_hub_download(*args, **kwargs):
|
||||||
|
pytest.fail("local missing processor state should not call hf_hub_download")
|
||||||
|
|
||||||
|
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fail_hub_download)
|
||||||
|
|
||||||
|
with pytest.raises(FileNotFoundError, match="missing.safetensors.*local processor pipeline"):
|
||||||
|
DataProcessorPipeline.from_pretrained(config_path, config_filename="ignored.json")
|
||||||
|
finally:
|
||||||
|
ProcessorStepRegistry.unregister("local_file_missing_state_step")
|
||||||
|
|
||||||
|
|
||||||
|
def test_from_pretrained_hub_source_missing_local_state_still_calls_hub(monkeypatch, tmp_path):
|
||||||
|
"""Hub sources still fall back to hf_hub_download for state files."""
|
||||||
|
|
||||||
|
@ProcessorStepRegistry.register("hub_state_step")
|
||||||
|
class HubStateStep(ProcessorStep):
|
||||||
|
def __init__(self):
|
||||||
|
self.value = torch.tensor(0)
|
||||||
|
|
||||||
|
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||||
|
return transition
|
||||||
|
|
||||||
|
def transform_features(
|
||||||
|
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||||
|
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||||
|
return features
|
||||||
|
|
||||||
|
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||||||
|
self.value = state["value"]
|
||||||
|
|
||||||
|
try:
|
||||||
|
state_path = tmp_path / "downloaded.safetensors"
|
||||||
|
save_file({"value": torch.tensor(7)}, state_path)
|
||||||
|
loaded_config = {
|
||||||
|
"name": "HubStatePipeline",
|
||||||
|
"steps": [{"registry_name": "hub_state_step", "state_file": "hub_state.safetensors"}],
|
||||||
|
}
|
||||||
|
calls = []
|
||||||
|
|
||||||
|
def fake_load_config(cls, model_id, config_filename, hub_download_kwargs):
|
||||||
|
return loaded_config, tmp_path / "hub_cache"
|
||||||
|
|
||||||
|
def fake_hub_download(**kwargs):
|
||||||
|
calls.append(kwargs)
|
||||||
|
return str(state_path)
|
||||||
|
|
||||||
|
monkeypatch.setattr(DataProcessorPipeline, "_load_config", classmethod(fake_load_config))
|
||||||
|
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fake_hub_download)
|
||||||
|
|
||||||
|
pipeline = DataProcessorPipeline.from_pretrained("user/repo", config_filename="processor.json")
|
||||||
|
|
||||||
|
assert calls == [
|
||||||
|
{
|
||||||
|
"repo_id": "user/repo",
|
||||||
|
"filename": "hub_state.safetensors",
|
||||||
|
"repo_type": "model",
|
||||||
|
"force_download": False,
|
||||||
|
"resume_download": None,
|
||||||
|
"proxies": None,
|
||||||
|
"token": None,
|
||||||
|
"cache_dir": None,
|
||||||
|
"local_files_only": False,
|
||||||
|
"revision": None,
|
||||||
|
}
|
||||||
|
]
|
||||||
|
assert pipeline.steps[0].value.item() == 7
|
||||||
|
finally:
|
||||||
|
ProcessorStepRegistry.unregister("hub_state_step")
|
||||||
|
|
||||||
|
|
||||||
# Config Validation Tests
|
# Config Validation Tests
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -17,6 +17,8 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import dataclasses
|
import dataclasses
|
||||||
|
import sys
|
||||||
|
from types import SimpleNamespace
|
||||||
from unittest.mock import MagicMock
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
@@ -106,6 +108,116 @@ def test_sentry_config_defaults():
|
|||||||
assert cfg.target_video_file_size_mb is None
|
assert cfg.target_video_file_size_mb is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_rollout_config_passes_policy_pretrained_revision(monkeypatch):
|
||||||
|
from lerobot.configs import PreTrainedConfig, parser
|
||||||
|
from lerobot.rollout import RolloutConfig
|
||||||
|
from tests.mocks.mock_robot import MockRobotConfig
|
||||||
|
|
||||||
|
captured = {}
|
||||||
|
|
||||||
|
def fake_from_pretrained(cls, pretrained_name_or_path, **kwargs):
|
||||||
|
captured["pretrained_name_or_path"] = pretrained_name_or_path
|
||||||
|
captured.update(kwargs)
|
||||||
|
return SimpleNamespace(device="cpu", pretrained_revision=kwargs["revision"])
|
||||||
|
|
||||||
|
monkeypatch.setattr(parser, "get_yaml_overrides", lambda _: ["--pretrained_revision=yaml-sha"])
|
||||||
|
monkeypatch.setattr(
|
||||||
|
sys,
|
||||||
|
"argv",
|
||||||
|
["lerobot-rollout", "--policy.path=user/policy", "--policy.pretrained_revision=cli-sha"],
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(PreTrainedConfig, "from_pretrained", classmethod(fake_from_pretrained))
|
||||||
|
|
||||||
|
cfg = RolloutConfig(robot=MockRobotConfig())
|
||||||
|
|
||||||
|
assert captured["pretrained_name_or_path"] == "user/policy"
|
||||||
|
assert captured["revision"] == "cli-sha"
|
||||||
|
assert captured["cli_overrides"] == [
|
||||||
|
"--pretrained_revision=yaml-sha",
|
||||||
|
"--pretrained_revision=cli-sha",
|
||||||
|
]
|
||||||
|
assert cfg.policy.pretrained_path == "user/policy"
|
||||||
|
assert cfg.policy.pretrained_revision == "cli-sha"
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_pretrained_policy_passes_revision(monkeypatch):
|
||||||
|
import lerobot.rollout.context as rollout_context
|
||||||
|
|
||||||
|
policy_config = SimpleNamespace(
|
||||||
|
type="mock",
|
||||||
|
use_peft=False,
|
||||||
|
pretrained_path="user/policy",
|
||||||
|
pretrained_revision="policy-sha",
|
||||||
|
)
|
||||||
|
policy_class = MagicMock()
|
||||||
|
loaded_policy = MagicMock()
|
||||||
|
policy_class.from_pretrained.return_value = loaded_policy
|
||||||
|
monkeypatch.setattr(rollout_context, "get_policy_class", lambda _: policy_class)
|
||||||
|
|
||||||
|
policy = rollout_context._load_pretrained_policy(policy_config)
|
||||||
|
|
||||||
|
assert policy is loaded_policy
|
||||||
|
policy_class.from_pretrained.assert_called_once_with(
|
||||||
|
"user/policy",
|
||||||
|
config=policy_config,
|
||||||
|
revision="policy-sha",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_pretrained_peft_policy_keeps_adapter_and_base_revisions_separate(monkeypatch):
|
||||||
|
import lerobot.rollout.context as rollout_context
|
||||||
|
|
||||||
|
policy_config = SimpleNamespace(
|
||||||
|
type="mock",
|
||||||
|
use_peft=True,
|
||||||
|
pretrained_path="user/adapter",
|
||||||
|
pretrained_revision="adapter-sha",
|
||||||
|
)
|
||||||
|
policy_class = MagicMock()
|
||||||
|
base_policy = MagicMock()
|
||||||
|
policy_class.from_pretrained.return_value = base_policy
|
||||||
|
monkeypatch.setattr(rollout_context, "get_policy_class", lambda _: policy_class)
|
||||||
|
|
||||||
|
peft_config = SimpleNamespace(
|
||||||
|
base_model_name_or_path="user/base-policy",
|
||||||
|
revision="base-sha",
|
||||||
|
)
|
||||||
|
peft_config_from_pretrained = MagicMock(return_value=peft_config)
|
||||||
|
adapted_policy = MagicMock()
|
||||||
|
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
|
||||||
|
require_package = MagicMock()
|
||||||
|
monkeypatch.setattr(rollout_context, "require_package", require_package)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
rollout_context,
|
||||||
|
"PeftConfig",
|
||||||
|
SimpleNamespace(from_pretrained=peft_config_from_pretrained),
|
||||||
|
raising=False,
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
rollout_context,
|
||||||
|
"PeftModel",
|
||||||
|
SimpleNamespace(from_pretrained=peft_model_from_pretrained),
|
||||||
|
raising=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
policy = rollout_context._load_pretrained_policy(policy_config)
|
||||||
|
|
||||||
|
assert policy is adapted_policy
|
||||||
|
require_package.assert_called_once_with("peft", extra="peft")
|
||||||
|
peft_config_from_pretrained.assert_called_once_with("user/adapter", revision="adapter-sha")
|
||||||
|
policy_class.from_pretrained.assert_called_once_with(
|
||||||
|
pretrained_name_or_path="user/base-policy",
|
||||||
|
config=policy_config,
|
||||||
|
revision="base-sha",
|
||||||
|
)
|
||||||
|
peft_model_from_pretrained.assert_called_once_with(
|
||||||
|
base_policy,
|
||||||
|
"user/adapter",
|
||||||
|
config=peft_config,
|
||||||
|
revision="adapter-sha",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# RolloutRingBuffer
|
# RolloutRingBuffer
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|||||||
Reference in New Issue
Block a user