admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-24 02:06:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into feat/add-xvla

Browse Source
This commit is contained in:
Michel Aractingi
2025-11-28 10:54:42 +01:00
committed by GitHub
parent d4e6d60ec3 c55fbe1b3e
commit 8f59d93458
39 changed files with 1782 additions and 1817 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmarks/video/benchmark.py
-94
View File
@@ -1,94 +0,0 @@
#!/usr/bin/env python
# Copyright 2024 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 threading
import time
from contextlib import ContextDecorator
class TimeBenchmark(ContextDecorator):
"""
Measures execution time using a context manager or decorator.
This class supports both context manager and decorator usage, and is thread-safe for multithreaded
environments.
Args:
print: If True, prints the elapsed time upon exiting the context or completing the function. Defaults
to False.
Examples:
Using as a context manager:
>>> benchmark = TimeBenchmark()
>>> with benchmark:
... time.sleep(1)
>>> print(f"Block took {benchmark.result:.4f} seconds")
Block took approximately 1.0000 seconds
Using with multithreading:
```python
import threading
benchmark = TimeBenchmark()
def context_manager_example():
with benchmark:
time.sleep(0.01)
print(f"Block took {benchmark.result_ms:.2f} milliseconds")
threads = []
for _ in range(3):
t1 = threading.Thread(target=context_manager_example)
threads.append(t1)
for t in threads:
t.start()
for t in threads:
t.join()
```
Expected output:
Block took approximately 10.00 milliseconds
Block took approximately 10.00 milliseconds
Block took approximately 10.00 milliseconds
"""
def __init__(self, print=False):
self.local = threading.local()
self.print_time = print
def __enter__(self):
self.local.start_time = time.perf_counter()
return self
def __exit__(self, *exc):
self.local.end_time = time.perf_counter()
self.local.elapsed_time = self.local.end_time - self.local.start_time
if self.print_time:
print(f"Elapsed time: {self.local.elapsed_time:.4f} seconds")
return False
@property
def result(self):
return getattr(self.local, "elapsed_time", None)
@property
def result_ms(self):
return self.result * 1e3
benchmarks/video/capture_camera_feed.py
-102
View File
@@ -1,102 +0,0 @@
#!/usr/bin/env python
# Copyright 2024 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.
"""Capture video feed from a camera as raw images."""
import argparse
import datetime as dt
import os
import time
from pathlib import Path
import cv2
import rerun as rr
# see https://rerun.io/docs/howto/visualization/limit-ram
RERUN_MEMORY_LIMIT = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "5%")
def display_and_save_video_stream(output_dir: Path, fps: int, width: int, height: int, duration: int):
rr.init("lerobot_capture_camera_feed")
rr.spawn(memory_limit=RERUN_MEMORY_LIMIT)
now = dt.datetime.now()
capture_dir = output_dir / f"{now:%Y-%m-%d}" / f"{now:%H-%M-%S}"
if not capture_dir.exists():
capture_dir.mkdir(parents=True, exist_ok=True)
# Opens the default webcam
cap = cv2.VideoCapture(0)
if not cap.isOpened():
print("Error: Could not open video stream.")
return
cap.set(cv2.CAP_PROP_FPS, fps)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
frame_index = 0
start_time = time.time()
while time.time() - start_time < duration:
ret, frame = cap.read()
if not ret:
print("Error: Could not read frame.")
break
rr.log("video/stream", rr.Image(frame), static=True)
cv2.imwrite(str(capture_dir / f"frame_{frame_index:06d}.png"), frame)
frame_index += 1
# Release the capture
cap.release()
# TODO(Steven): Add a graceful shutdown via a close() method for the Viewer context, though not currently supported in the Rerun API.
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--output-dir",
type=Path,
default=Path("outputs/cam_capture/"),
help="Directory where the capture images are written. A subfolder named with the current date & time will be created inside it for each capture.",
)
parser.add_argument(
"--fps",
type=int,
default=30,
help="Frames Per Second of the capture.",
)
parser.add_argument(
"--width",
type=int,
default=1280,
help="Width of the captured images.",
)
parser.add_argument(
"--height",
type=int,
default=720,
help="Height of the captured images.",
)
parser.add_argument(
"--duration",
type=int,
default=20,
help="Duration in seconds for which the video stream should be captured.",
)
args = parser.parse_args()
display_and_save_video_stream(**vars(args))
benchmarks/video/run_video_benchmark.py
+28 -33
View File
@@ -21,11 +21,13 @@ See the provided README.md or run `python benchmark/video/run_video_benchmark.py
import argparse import argparse
import datetime as dt import datetime as dt
import itertools
import random import random
import shutil import shutil
from collections import OrderedDict from collections import OrderedDict
from concurrent.futures import ThreadPoolExecutor, as_completed from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path from pathlib import Path
from threading import Lock
import einops import einops
import numpy as np import numpy as np
@@ -35,13 +37,13 @@ import torch
from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity
from tqdm import tqdm from tqdm import tqdm
from benchmarks.video.benchmark import TimeBenchmark
from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.video_utils import ( from lerobot.datasets.video_utils import (
decode_video_frames_torchvision, decode_video_frames,
encode_video_frames, encode_video_frames,
) )
from lerobot.utils.constants import OBS_IMAGE from lerobot.utils.constants import OBS_IMAGE
from lerobot.utils.utils import TimerManager
BASE_ENCODING = OrderedDict( BASE_ENCODING = OrderedDict(
[ [
@@ -86,7 +88,7 @@ def load_original_frames(imgs_dir: Path, timestamps: list[float], fps: int) -> t
frames = [] frames = []
for ts in timestamps: for ts in timestamps:
idx = int(ts * fps) idx = int(ts * fps)
frame = PIL.Image.open(imgs_dir / f"frame_{idx:06d}.png") frame = PIL.Image.open(imgs_dir / f"frame-{idx:06d}.png")
frame = torch.from_numpy(np.array(frame)) frame = torch.from_numpy(np.array(frame))
frame = frame.type(torch.float32) / 255 frame = frame.type(torch.float32) / 255
frame = einops.rearrange(frame, "h w c -> c h w") frame = einops.rearrange(frame, "h w c -> c h w")
@@ -97,21 +99,21 @@ def load_original_frames(imgs_dir: Path, timestamps: list[float], fps: int) -> t
def save_decoded_frames( def save_decoded_frames(
imgs_dir: Path, save_dir: Path, frames: torch.Tensor, timestamps: list[float], fps: int imgs_dir: Path, save_dir: Path, frames: torch.Tensor, timestamps: list[float], fps: int
) -> None: ) -> None:
if save_dir.exists() and len(list(save_dir.glob("frame_*.png"))) == len(timestamps): if save_dir.exists() and len(list(save_dir.glob("frame-*.png"))) == len(timestamps):
return return
save_dir.mkdir(parents=True, exist_ok=True) save_dir.mkdir(parents=True, exist_ok=True)
for i, ts in enumerate(timestamps): for i, ts in enumerate(timestamps):
idx = int(ts * fps) idx = int(ts * fps)
frame_hwc = (frames[i].permute((1, 2, 0)) * 255).type(torch.uint8).cpu().numpy() frame_hwc = (frames[i].permute((1, 2, 0)) * 255).type(torch.uint8).cpu().numpy()
PIL.Image.fromarray(frame_hwc).save(save_dir / f"frame_{idx:06d}_decoded.png") PIL.Image.fromarray(frame_hwc).save(save_dir / f"frame-{idx:06d}_decoded.png")
shutil.copyfile(imgs_dir / f"frame_{idx:06d}.png", save_dir / f"frame_{idx:06d}_original.png") shutil.copyfile(imgs_dir / f"frame-{idx:06d}.png", save_dir / f"frame-{idx:06d}_original.png")
def save_first_episode(imgs_dir: Path, dataset: LeRobotDataset) -> None: def save_first_episode(imgs_dir: Path, dataset: LeRobotDataset) -> None:
episode_index = 0 episode_index = 0
ep_num_images = dataset.meta.episodes["length"][episode_index] ep_num_images = dataset.meta.episodes["length"][episode_index]
if imgs_dir.exists() and len(list(imgs_dir.glob("frame_*.png"))) == ep_num_images: if imgs_dir.exists() and len(list(imgs_dir.glob("frame-*.png"))) == ep_num_images:
return return
imgs_dir.mkdir(parents=True, exist_ok=True) imgs_dir.mkdir(parents=True, exist_ok=True)
@@ -125,7 +127,7 @@ def save_first_episode(imgs_dir: Path, dataset: LeRobotDataset) -> None:
tqdm(imgs_dataset, desc=f"saving {dataset.repo_id} first episode images", leave=False) tqdm(imgs_dataset, desc=f"saving {dataset.repo_id} first episode images", leave=False)
): ):
img = item[img_keys[0]] img = item[img_keys[0]]
img.save(str(imgs_dir / f"frame_{i:06d}.png"), quality=100) img.save(str(imgs_dir / f"frame-{i:06d}.png"), quality=100)
if i >= ep_num_images - 1: if i >= ep_num_images - 1:
break break
@@ -149,18 +151,6 @@ def sample_timestamps(timestamps_mode: str, ep_num_images: int, fps: int) -> lis
return [idx / fps for idx in frame_indexes] return [idx / fps for idx in frame_indexes]
def decode_video_frames(
video_path: str,
timestamps: list[float],
tolerance_s: float,
backend: str,
) -> torch.Tensor:
if backend in ["pyav", "video_reader"]:
return decode_video_frames_torchvision(video_path, timestamps, tolerance_s, backend)
else:
raise NotImplementedError(backend)
def benchmark_decoding( def benchmark_decoding(
imgs_dir: Path, imgs_dir: Path,
video_path: Path, video_path: Path,
@@ -172,8 +162,8 @@ def benchmark_decoding(
num_workers: int = 4, num_workers: int = 4,
save_frames: bool = False, save_frames: bool = False,
) -> dict: ) -> dict:
def process_sample(sample: int): def process_sample(sample: int, lock: Lock):
time_benchmark = TimeBenchmark() time_benchmark = TimerManager(log=False)
timestamps = sample_timestamps(timestamps_mode, ep_num_images, fps) timestamps = sample_timestamps(timestamps_mode, ep_num_images, fps)
num_frames = len(timestamps) num_frames = len(timestamps)
result = { result = {
@@ -182,13 +172,13 @@ def benchmark_decoding(
"mse_values": [], "mse_values": [],
} }
with time_benchmark: with time_benchmark, lock:
frames = decode_video_frames(video_path, timestamps=timestamps, tolerance_s=5e-1, backend=backend) frames = decode_video_frames(video_path, timestamps=timestamps, tolerance_s=5e-1, backend=backend)
result["load_time_video_ms"] = time_benchmark.result_ms / num_frames result["load_time_video_ms"] = (time_benchmark.last * 1000) / num_frames
with time_benchmark: with time_benchmark:
original_frames = load_original_frames(imgs_dir, timestamps, fps) original_frames = load_original_frames(imgs_dir, timestamps, fps)
result["load_time_images_ms"] = time_benchmark.result_ms / num_frames result["load_time_images_ms"] = (time_benchmark.last * 1000) / num_frames
frames_np, original_frames_np = frames.numpy(), original_frames.numpy() frames_np, original_frames_np = frames.numpy(), original_frames.numpy()
for i in range(num_frames): for i in range(num_frames):
@@ -215,8 +205,10 @@ def benchmark_decoding(
# A sample is a single set of decoded frames specified by timestamps_mode (e.g. a single frame, 2 frames, etc.). # A sample is a single set of decoded frames specified by timestamps_mode (e.g. a single frame, 2 frames, etc.).
# For each sample, we record metrics (loading time and quality metrics) which are then averaged over all samples. # For each sample, we record metrics (loading time and quality metrics) which are then averaged over all samples.
# As these samples are independent, we run them in parallel threads to speed up the benchmark. # As these samples are independent, we run them in parallel threads to speed up the benchmark.
# Use a single shared lock for all worker threads
shared_lock = Lock()
with ThreadPoolExecutor(max_workers=num_workers) as executor: with ThreadPoolExecutor(max_workers=num_workers) as executor:
futures = [executor.submit(process_sample, i) for i in range(num_samples)] futures = [executor.submit(process_sample, i, shared_lock) for i in range(num_samples)]
for future in tqdm(as_completed(futures), total=num_samples, desc="samples", leave=False): for future in tqdm(as_completed(futures), total=num_samples, desc="samples", leave=False):
result = future.result() result = future.result()
load_times_video_ms.append(result["load_time_video_ms"]) load_times_video_ms.append(result["load_time_video_ms"])
@@ -358,11 +350,14 @@ def main(
imgs_dir = output_dir / "images" / dataset.repo_id.replace("/", "_") imgs_dir = output_dir / "images" / dataset.repo_id.replace("/", "_")
# We only use the first episode # We only use the first episode
save_first_episode(imgs_dir, dataset) save_first_episode(imgs_dir, dataset)
for key, values in tqdm(encoding_benchmarks.items(), desc="encodings (g, crf)", leave=False): for duet in [
for value in tqdm(values, desc=f"encodings ({key})", leave=False): dict(zip(encoding_benchmarks.keys(), unique_combination, strict=False))
for unique_combination in itertools.product(*encoding_benchmarks.values())
]:
encoding_cfg = BASE_ENCODING.copy() encoding_cfg = BASE_ENCODING.copy()
encoding_cfg["vcodec"] = video_codec encoding_cfg["vcodec"] = video_codec
encoding_cfg["pix_fmt"] = pixel_format encoding_cfg["pix_fmt"] = pixel_format
for key, value in duet.items():
encoding_cfg[key] = value encoding_cfg[key] = value
args_path = Path("_".join(str(value) for value in encoding_cfg.values())) args_path = Path("_".join(str(value) for value in encoding_cfg.values()))
video_path = output_dir / "videos" / args_path / f"{repo_id.replace('/', '_')}.mp4" video_path = output_dir / "videos" / args_path / f"{repo_id.replace('/', '_')}.mp4"
@@ -409,9 +404,9 @@ if __name__ == "__main__":
nargs="*", nargs="*",
default=[ default=[
"lerobot/pusht_image", "lerobot/pusht_image",
"aliberts/aloha_mobile_shrimp_image", "lerobot/aloha_mobile_shrimp_image",
"aliberts/paris_street", "lerobot/paris_street",
"aliberts/kitchen", "lerobot/kitchen",
], ],
help="Datasets repo-ids to test against. First episodes only are used. Must be images.", help="Datasets repo-ids to test against. First episodes only are used. Must be images.",
) )
@@ -419,7 +414,7 @@ if __name__ == "__main__":
"--vcodec", "--vcodec",
type=str, type=str,
nargs="*", nargs="*",
default=["libx264", "hevc", "libsvtav1"], default=["h264", "hevc", "libsvtav1"],
help="Video codecs to be tested", help="Video codecs to be tested",
) )
parser.add_argument( parser.add_argument(
@@ -468,7 +463,7 @@ if __name__ == "__main__":
"--backends", "--backends",
type=str, type=str,
nargs="*", nargs="*",
default=["pyav", "video_reader"], default=["torchcodec", "pyav"],
help="Torchvision decoding backend to be tested.", help="Torchvision decoding backend to be tested.",
) )
parser.add_argument( parser.add_argument(
docs/source/async.mdx
+1 -1
View File
@@ -196,7 +196,7 @@ client_cfg = RobotClientConfig(
server_address="localhost:8080", server_address="localhost:8080",
policy_device="mps", policy_device="mps",
policy_type="smolvla", policy_type="smolvla",
pretrained_name_or_path="fracapuano/smolvla_async", pretrained_name_or_path="<user>/smolvla_async",
chunk_size_threshold=0.5, chunk_size_threshold=0.5,
actions_per_chunk=50, # make sure this is less than the max actions of the policy actions_per_chunk=50, # make sure this is less than the max actions of the policy
) )
docs/source/envhub_leisaac.mdx
+2 -2
View File
@@ -139,7 +139,7 @@ from lerobot.teleoperators import ( # noqa: F401
make_teleoperator_from_config, make_teleoperator_from_config,
so101_leader, so101_leader,
) )
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.utils import init_logging from lerobot.utils.utils import init_logging
from lerobot.envs.factory import make_env from lerobot.envs.factory import make_env
@@ -196,7 +196,7 @@ def teleop_loop(teleop: Teleoperator, env: gym.Env, fps: int):
obs, info = env.reset() obs, info = env.reset()
dt_s = time.perf_counter() - loop_start dt_s = time.perf_counter() - loop_start
busy_wait(1 / fps - dt_s) precise_sleep(1 / fps - dt_s)
loop_s = time.perf_counter() - loop_start loop_s = time.perf_counter() - loop_start
print(f"\ntime: {loop_s * 1e3:.2f}ms ({1 / loop_s:.0f} Hz)") print(f"\ntime: {loop_s * 1e3:.2f}ms ({1 / loop_s:.0f} Hz)")
docs/source/il_robots.mdx
+2 -2
View File
@@ -393,7 +393,7 @@ import time
from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.robots.so100_follower.so100_follower import SO100Follower
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.utils import log_say from lerobot.utils.utils import log_say
episode_idx = 0 episode_idx = 0
@@ -415,7 +415,7 @@ for idx in range(dataset.num_frames):
} }
robot.send_action(action) robot.send_action(action)
busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0)) precise_sleep(1.0 / dataset.fps - (time.perf_counter() - t0))
robot.disconnect() robot.disconnect()
``` ```
examples/backward_compatibility/replay.py
+2 -2
View File
@@ -45,7 +45,7 @@ from lerobot.robots import ( # noqa: F401
so101_follower, so101_follower,
) )
from lerobot.utils.constants import ACTION from lerobot.utils.constants import ACTION
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.utils import ( from lerobot.utils.utils import (
init_logging, init_logging,
log_say, log_say,
@@ -97,7 +97,7 @@ def replay(cfg: ReplayConfig):
robot.send_action(action) robot.send_action(action)
dt_s = time.perf_counter() - start_episode_t dt_s = time.perf_counter() - start_episode_t
busy_wait(1 / dataset.fps - dt_s) precise_sleep(1 / dataset.fps - dt_s)
robot.disconnect() robot.disconnect()
examples/dataset/load_lerobot_dataset.py
+80 -75
View File
@@ -34,105 +34,106 @@ from huggingface_hub import HfApi
import lerobot import lerobot
from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
# We ported a number of existing datasets ourselves, use this to see the list:
print("List of available datasets:")
pprint(lerobot.available_datasets)
# You can also browse through the datasets created/ported by the community on the hub using the hub api: def main():
hub_api = HfApi() # We ported a number of existing datasets ourselves, use this to see the list:
repo_ids = [info.id for info in hub_api.list_datasets(task_categories="robotics", tags=["LeRobot"])] print("List of available datasets:")
pprint(repo_ids) pprint(lerobot.available_datasets)
# Or simply explore them in your web browser directly at: # You can also browse through the datasets created/ported by the community on the hub using the hub api:
# https://huggingface.co/datasets?other=LeRobot hub_api = HfApi()
repo_ids = [info.id for info in hub_api.list_datasets(task_categories="robotics", tags=["LeRobot"])]
pprint(repo_ids)
# Let's take this one for this example # Or simply explore them in your web browser directly at:
repo_id = "lerobot/aloha_mobile_cabinet" # https://huggingface.co/datasets?other=LeRobot
# We can have a look and fetch its metadata to know more about it:
ds_meta = LeRobotDatasetMetadata(repo_id)
# By instantiating just this class, you can quickly access useful information about the content and the # Let's take this one for this example
# structure of the dataset without downloading the actual data yet (only metadata files — which are repo_id = "lerobot/aloha_mobile_cabinet"
# lightweight). # We can have a look and fetch its metadata to know more about it:
print(f"Total number of episodes: {ds_meta.total_episodes}") ds_meta = LeRobotDatasetMetadata(repo_id)
print(f"Average number of frames per episode: {ds_meta.total_frames / ds_meta.total_episodes:.3f}")
print(f"Frames per second used during data collection: {ds_meta.fps}")
print(f"Robot type: {ds_meta.robot_type}")
print(f"keys to access images from cameras: {ds_meta.camera_keys=}\n")
print("Tasks:") # By instantiating just this class, you can quickly access useful information about the content and the
print(ds_meta.tasks) # structure of the dataset without downloading the actual data yet (only metadata files — which are
print("Features:") # lightweight).
pprint(ds_meta.features) print(f"Total number of episodes: {ds_meta.total_episodes}")
print(f"Average number of frames per episode: {ds_meta.total_frames / ds_meta.total_episodes:.3f}")
print(f"Frames per second used during data collection: {ds_meta.fps}")
print(f"Robot type: {ds_meta.robot_type}")
print(f"keys to access images from cameras: {ds_meta.camera_keys=}\n")
# You can also get a short summary by simply printing the object: print("Tasks:")
print(ds_meta) print(ds_meta.tasks)
print("Features:")
pprint(ds_meta.features)
# You can then load the actual dataset from the hub. # You can also get a short summary by simply printing the object:
# Either load any subset of episodes: print(ds_meta)
dataset = LeRobotDataset(repo_id, episodes=[0, 10, 11, 23])
# And see how many frames you have: # You can then load the actual dataset from the hub.
print(f"Selected episodes: {dataset.episodes}") # Either load any subset of episodes:
print(f"Number of episodes selected: {dataset.num_episodes}") dataset = LeRobotDataset(repo_id, episodes=[0, 10, 11, 23])
print(f"Number of frames selected: {dataset.num_frames}")
# Or simply load the entire dataset: # And see how many frames you have:
dataset = LeRobotDataset(repo_id) print(f"Selected episodes: {dataset.episodes}")
print(f"Number of episodes selected: {dataset.num_episodes}") print(f"Number of episodes selected: {dataset.num_episodes}")
print(f"Number of frames selected: {dataset.num_frames}") print(f"Number of frames selected: {dataset.num_frames}")
# The previous metadata class is contained in the 'meta' attribute of the dataset: # Or simply load the entire dataset:
print(dataset.meta) dataset = LeRobotDataset(repo_id)
print(f"Number of episodes selected: {dataset.num_episodes}")
print(f"Number of frames selected: {dataset.num_frames}")
# LeRobotDataset actually wraps an underlying Hugging Face dataset # The previous metadata class is contained in the 'meta' attribute of the dataset:
# (see https://huggingface.co/docs/datasets for more information). print(dataset.meta)
print(dataset.hf_dataset)
# LeRobot datasets also subclasses PyTorch datasets so you can do everything you know and love from working # LeRobotDataset actually wraps an underlying Hugging Face dataset
# with the latter, like iterating through the dataset. # (see https://huggingface.co/docs/datasets for more information).
# The __getitem__ iterates over the frames of the dataset. Since our datasets are also structured by print(dataset.hf_dataset)
# episodes, you can access the frame indices of any episode using dataset.meta.episodes. Here, we access
# frame indices associated to the first episode:
episode_index = 0
from_idx = dataset.meta.episodes["dataset_from_index"][episode_index]
to_idx = dataset.meta.episodes["dataset_to_index"][episode_index]
# Then we grab all the image frames from the first camera: # LeRobot datasets also subclasses PyTorch datasets so you can do everything you know and love from working
camera_key = dataset.meta.camera_keys[0] # with the latter, like iterating through the dataset.
frames = [dataset[idx][camera_key] for idx in range(from_idx, to_idx)] # The __getitem__ iterates over the frames of the dataset. Since our datasets are also structured by
# episodes, you can access the frame indices of any episode using dataset.meta.episodes. Here, we access
# frame indices associated to the first episode:
episode_index = 0
from_idx = dataset.meta.episodes["dataset_from_index"][episode_index]
to_idx = dataset.meta.episodes["dataset_to_index"][episode_index]
# The objects returned by the dataset are all torch.Tensors # Then we grab all the image frames from the first camera:
print(type(frames[0])) camera_key = dataset.meta.camera_keys[0]
print(frames[0].shape) frames = [dataset[idx][camera_key] for idx in range(from_idx, to_idx)]
# Since we're using pytorch, the shape is in pytorch, channel-first convention (c, h, w). # The objects returned by the dataset are all torch.Tensors
# We can compare this shape with the information available for that feature print(type(frames[0]))
pprint(dataset.features[camera_key]) print(frames[0].shape)
# In particular:
print(dataset.features[camera_key]["shape"])
# The shape is in (h, w, c) which is a more universal format.
# For many machine learning applications we need to load the history of past observations or trajectories of # Since we're using pytorch, the shape is in pytorch, channel-first convention (c, h, w).
# future actions. Our datasets can load previous and future frames for each key/modality, using timestamps # We can compare this shape with the information available for that feature
# differences with the current loaded frame. For instance: pprint(dataset.features[camera_key])
delta_timestamps = { # In particular:
print(dataset.features[camera_key]["shape"])
# The shape is in (h, w, c) which is a more universal format.
# For many machine learning applications we need to load the history of past observations or trajectories of
# future actions. Our datasets can load previous and future frames for each key/modality, using timestamps
# differences with the current loaded frame. For instance:
delta_timestamps = {
# loads 4 images: 1 second before current frame, 500 ms before, 200 ms before, and current frame # loads 4 images: 1 second before current frame, 500 ms before, 200 ms before, and current frame
camera_key: [-1, -0.5, -0.20, 0], camera_key: [-1, -0.5, -0.20, 0],
# loads 6 state vectors: 1.5 seconds before, 1 second before, ... 200 ms, 100 ms, and current frame # loads 6 state vectors: 1.5 seconds before, 1 second before, ... 200 ms, 100 ms, and current frame
"observation.state": [-1.5, -1, -0.5, -0.20, -0.10, 0], "observation.state": [-1.5, -1, -0.5, -0.20, -0.10, 0],
# loads 64 action vectors: current frame, 1 frame in the future, 2 frames, ... 63 frames in the future # loads 64 action vectors: current frame, 1 frame in the future, 2 frames, ... 63 frames in the future
"action": [t / dataset.fps for t in range(64)], "action": [t / dataset.fps for t in range(64)],
} }
# Note that in any case, these delta_timestamps values need to be multiples of (1/fps) so that added to any # Note that in any case, these delta_timestamps values need to be multiples of (1/fps) so that added to any
# timestamp, you still get a valid timestamp. # timestamp, you still get a valid timestamp.
dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps) dataset = LeRobotDataset(repo_id, delta_timestamps=delta_timestamps)
print(f"\n{dataset[0][camera_key].shape=}") # (4, c, h, w) print(f"\n{dataset[0][camera_key].shape=}") # (4, c, h, w)
print(f"{dataset[0]['observation.state'].shape=}") # (6, c) print(f"{dataset[0]['observation.state'].shape=}") # (6, c)
print(f"{dataset[0]['action'].shape=}\n") # (64, c) print(f"{dataset[0]['action'].shape=}\n") # (64, c)
if __name__ == "__main__":
dataloader = torch.utils.data.DataLoader( dataloader = torch.utils.data.DataLoader(
dataset, dataset,
num_workers=4, num_workers=4,
@@ -144,3 +145,7 @@ if __name__ == "__main__":
print(f"{batch['observation.state'].shape=}") # (32, 6, c) print(f"{batch['observation.state'].shape=}") # (32, 6, c)
print(f"{batch['action'].shape=}") # (32, 64, c) print(f"{batch['action'].shape=}") # (32, 64, c)
break break
if __name__ == "__main__":
main()
examples/lekiwi/evaluate.py
+39 -33
View File
@@ -33,55 +33,57 @@ TASK_DESCRIPTION = "My task description"
HF_MODEL_ID = "<hf_username>/<model_repo_id>" HF_MODEL_ID = "<hf_username>/<model_repo_id>"
HF_DATASET_ID = "<hf_username>/<eval_dataset_repo_id>" HF_DATASET_ID = "<hf_username>/<eval_dataset_repo_id>"
# Create the robot configuration & robot
robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
robot = LeKiwiClient(robot_config) def main():
# Create the robot configuration & robot
robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
# Create policy robot = LeKiwiClient(robot_config)
policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
# Configure the dataset features # Create policy
action_features = hw_to_dataset_features(robot.action_features, ACTION) policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
dataset_features = {**action_features, **obs_features}
# Create the dataset # Configure the dataset features
dataset = LeRobotDataset.create( action_features = hw_to_dataset_features(robot.action_features, ACTION)
obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
dataset_features = {**action_features, **obs_features}
# Create the dataset
dataset = LeRobotDataset.create(
repo_id=HF_DATASET_ID, repo_id=HF_DATASET_ID,
fps=FPS, fps=FPS,
features=dataset_features, features=dataset_features,
robot_type=robot.name, robot_type=robot.name,
use_videos=True, use_videos=True,
image_writer_threads=4, image_writer_threads=4,
) )
# Build Policy Processors # Build Policy Processors
preprocessor, postprocessor = make_pre_post_processors( preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=policy, policy_cfg=policy,
pretrained_path=HF_MODEL_ID, pretrained_path=HF_MODEL_ID,
dataset_stats=dataset.meta.stats, dataset_stats=dataset.meta.stats,
# The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility. # The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}}, preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
) )
# Connect the robot # Connect the robot
# To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi` # To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
robot.connect() robot.connect()
# TODO(Steven): Update this example to use pipelines # TODO(Steven): Update this example to use pipelines
teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors() teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
# Initialize the keyboard listener and rerun visualization # Initialize the keyboard listener and rerun visualization
listener, events = init_keyboard_listener() listener, events = init_keyboard_listener()
init_rerun(session_name="lekiwi_evaluate") init_rerun(session_name="lekiwi_evaluate")
if not robot.is_connected: if not robot.is_connected:
raise ValueError("Robot is not connected!") raise ValueError("Robot is not connected!")
print("Starting evaluate loop...") print("Starting evaluate loop...")
recorded_episodes = 0 recorded_episodes = 0
while recorded_episodes < NUM_EPISODES and not events["stop_recording"]: while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
log_say(f"Running inference, recording eval episode {recorded_episodes} of {NUM_EPISODES}") log_say(f"Running inference, recording eval episode {recorded_episodes} of {NUM_EPISODES}")
# Main record loop # Main record loop
@@ -129,10 +131,14 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
dataset.save_episode() dataset.save_episode()
recorded_episodes += 1 recorded_episodes += 1
# Clean up # Clean up
log_say("Stop recording") log_say("Stop recording")
robot.disconnect() robot.disconnect()
listener.stop() listener.stop()
dataset.finalize() dataset.finalize()
dataset.push_to_hub() dataset.push_to_hub()
if __name__ == "__main__":
main()
examples/lekiwi/record.py
+43 -37
View File
@@ -34,50 +34,52 @@ RESET_TIME_SEC = 10
TASK_DESCRIPTION = "My task description" TASK_DESCRIPTION = "My task description"
HF_REPO_ID = "<hf_username>/<dataset_repo_id>" HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
# Create the robot and teleoperator configurations
robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
leader_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm")
keyboard_config = KeyboardTeleopConfig()
# Initialize the robot and teleoperator def main():
robot = LeKiwiClient(robot_config) # Create the robot and teleoperator configurations
leader_arm = SO100Leader(leader_arm_config) robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
keyboard = KeyboardTeleop(keyboard_config) leader_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm")
keyboard_config = KeyboardTeleopConfig()
# TODO(Steven): Update this example to use pipelines # Initialize the robot and teleoperator
teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors() robot = LeKiwiClient(robot_config)
leader_arm = SO100Leader(leader_arm_config)
keyboard = KeyboardTeleop(keyboard_config)
# Configure the dataset features # TODO(Steven): Update this example to use pipelines
action_features = hw_to_dataset_features(robot.action_features, ACTION) teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()
obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
dataset_features = {**action_features, **obs_features}
# Create the dataset # Configure the dataset features
dataset = LeRobotDataset.create( action_features = hw_to_dataset_features(robot.action_features, ACTION)
obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
dataset_features = {**action_features, **obs_features}
# Create the dataset
dataset = LeRobotDataset.create(
repo_id=HF_REPO_ID, repo_id=HF_REPO_ID,
fps=FPS, fps=FPS,
features=dataset_features, features=dataset_features,
robot_type=robot.name, robot_type=robot.name,
use_videos=True, use_videos=True,
image_writer_threads=4, image_writer_threads=4,
) )
# Connect the robot and teleoperator # Connect the robot and teleoperator
# To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi` # To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
robot.connect() robot.connect()
leader_arm.connect() leader_arm.connect()
keyboard.connect() keyboard.connect()
# Initialize the keyboard listener and rerun visualization # Initialize the keyboard listener and rerun visualization
listener, events = init_keyboard_listener() listener, events = init_keyboard_listener()
init_rerun(session_name="lekiwi_record") init_rerun(session_name="lekiwi_record")
if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected: if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
raise ValueError("Robot or teleop is not connected!") raise ValueError("Robot or teleop is not connected!")
print("Starting record loop...") print("Starting record loop...")
recorded_episodes = 0 recorded_episodes = 0
while recorded_episodes < NUM_EPISODES and not events["stop_recording"]: while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
log_say(f"Recording episode {recorded_episodes}") log_say(f"Recording episode {recorded_episodes}")
# Main record loop # Main record loop
@@ -124,12 +126,16 @@ while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
dataset.save_episode() dataset.save_episode()
recorded_episodes += 1 recorded_episodes += 1
# Clean up # Clean up
log_say("Stop recording") log_say("Stop recording")
robot.disconnect() robot.disconnect()
leader_arm.disconnect() leader_arm.disconnect()
keyboard.disconnect() keyboard.disconnect()
listener.stop() listener.stop()
dataset.finalize() dataset.finalize()
dataset.push_to_hub() dataset.push_to_hub()
if __name__ == "__main__":
main()
examples/lekiwi/replay.py
+24 -18
View File
@@ -20,32 +20,34 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient
from lerobot.utils.constants import ACTION from lerobot.utils.constants import ACTION
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.utils import log_say from lerobot.utils.utils import log_say
EPISODE_IDX = 0 EPISODE_IDX = 0
# Initialize the robot config
robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
# Initialize the robot def main():
robot = LeKiwiClient(robot_config) # Initialize the robot config
robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
# Fetch the dataset to replay # Initialize the robot
dataset = LeRobotDataset("<hf_username>/<dataset_repo_id>", episodes=[EPISODE_IDX]) robot = LeKiwiClient(robot_config)
# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
actions = episode_frames.select_columns(ACTION)
# Connect to the robot # Fetch the dataset to replay
robot.connect() dataset = LeRobotDataset("<hf_username>/<dataset_repo_id>", episodes=[EPISODE_IDX])
# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
actions = episode_frames.select_columns(ACTION)
if not robot.is_connected: # Connect to the robot
robot.connect()
if not robot.is_connected:
raise ValueError("Robot is not connected!") raise ValueError("Robot is not connected!")
print("Starting replay loop...") print("Starting replay loop...")
log_say(f"Replaying episode {EPISODE_IDX}") log_say(f"Replaying episode {EPISODE_IDX}")
for idx in range(len(episode_frames)): for idx in range(len(episode_frames)):
t0 = time.perf_counter() t0 = time.perf_counter()
# Get recorded action from dataset # Get recorded action from dataset
@@ -56,6 +58,10 @@ for idx in range(len(episode_frames)):
# Send action to robot # Send action to robot
_ = robot.send_action(action) _ = robot.send_action(action)
busy_wait(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0)) precise_sleep(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0))
robot.disconnect() robot.disconnect()
if __name__ == "__main__":
main()
examples/lekiwi/teleoperate.py
+26 -20
View File
@@ -19,35 +19,37 @@ import time
from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop, KeyboardTeleopConfig from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop, KeyboardTeleopConfig
from lerobot.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig from lerobot.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.visualization_utils import init_rerun, log_rerun_data from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
FPS = 30 FPS = 30
# Create the robot and teleoperator configurations
robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="my_lekiwi")
teleop_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm")
keyboard_config = KeyboardTeleopConfig(id="my_laptop_keyboard")
# Initialize the robot and teleoperator def main():
robot = LeKiwiClient(robot_config) # Create the robot and teleoperator configurations
leader_arm = SO100Leader(teleop_arm_config) robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="my_lekiwi")
keyboard = KeyboardTeleop(keyboard_config) teleop_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm")
keyboard_config = KeyboardTeleopConfig(id="my_laptop_keyboard")
# Connect to the robot and teleoperator # Initialize the robot and teleoperator
# To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi` robot = LeKiwiClient(robot_config)
robot.connect() leader_arm = SO100Leader(teleop_arm_config)
leader_arm.connect() keyboard = KeyboardTeleop(keyboard_config)
keyboard.connect()
# Init rerun viewer # Connect to the robot and teleoperator
init_rerun(session_name="lekiwi_teleop") # To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
robot.connect()
leader_arm.connect()
keyboard.connect()
if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected: # Init rerun viewer
init_rerun(session_name="lekiwi_teleop")
if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
raise ValueError("Robot or teleop is not connected!") raise ValueError("Robot or teleop is not connected!")
print("Starting teleop loop...") print("Starting teleop loop...")
while True: while True:
t0 = time.perf_counter() t0 = time.perf_counter()
# Get robot observation # Get robot observation
@@ -69,4 +71,8 @@ while True:
# Visualize # Visualize
log_rerun_data(observation=observation, action=action) log_rerun_data(observation=observation, action=action)
busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) precise_sleep(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
if __name__ == "__main__":
main()
examples/phone_to_so100/evaluate.py
+46 -38
View File
@@ -52,29 +52,31 @@ TASK_DESCRIPTION = "My task description"
HF_MODEL_ID = "<hf_username>/<model_repo_id>" HF_MODEL_ID = "<hf_username>/<model_repo_id>"
HF_DATASET_ID = "<hf_username>/<dataset_repo_id>" HF_DATASET_ID = "<hf_username>/<dataset_repo_id>"
# Create the robot configuration & robot
camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)} def main():
robot_config = SO100FollowerConfig( # Create the robot configuration & robot
camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
robot_config = SO100FollowerConfig(
port="/dev/tty.usbmodem58760434471", port="/dev/tty.usbmodem58760434471",
id="my_awesome_follower_arm", id="my_awesome_follower_arm",
cameras=camera_config, cameras=camera_config,
use_degrees=True, use_degrees=True,
) )
robot = SO100Follower(robot_config) robot = SO100Follower(robot_config)
# Create policy # Create policy
policy = ACTPolicy.from_pretrained(HF_MODEL_ID) policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
kinematics_solver = RobotKinematics( kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()), joint_names=list(robot.bus.motors.keys()),
) )
# Build pipeline to convert EE action to joints action # Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
steps=[ steps=[
InverseKinematicsEEToJoints( InverseKinematicsEEToJoints(
kinematics=kinematics_solver, kinematics=kinematics_solver,
@@ -84,19 +86,21 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Build pipeline to convert joints observation to EE observation # Build pipeline to convert joints observation to EE observation
robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation]( robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
steps=[ steps=[
ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) ForwardKinematicsJointsToEE(
kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())
)
], ],
to_transition=observation_to_transition, to_transition=observation_to_transition,
to_output=transition_to_observation, to_output=transition_to_observation,
) )
# Create the dataset # Create the dataset
dataset = LeRobotDataset.create( dataset = LeRobotDataset.create(
repo_id=HF_DATASET_ID, repo_id=HF_DATASET_ID,
fps=FPS, fps=FPS,
features=combine_feature_dicts( features=combine_feature_dicts(
@@ -121,30 +125,30 @@ dataset = LeRobotDataset.create(
robot_type=robot.name, robot_type=robot.name,
use_videos=True, use_videos=True,
image_writer_threads=4, image_writer_threads=4,
) )
# Build Policy Processors # Build Policy Processors
preprocessor, postprocessor = make_pre_post_processors( preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=policy, policy_cfg=policy,
pretrained_path=HF_MODEL_ID, pretrained_path=HF_MODEL_ID,
dataset_stats=dataset.meta.stats, dataset_stats=dataset.meta.stats,
# The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility. # The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}}, preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
) )
# Connect the robot # Connect the robot
robot.connect() robot.connect()
# Initialize the keyboard listener and rerun visualization # Initialize the keyboard listener and rerun visualization
listener, events = init_keyboard_listener() listener, events = init_keyboard_listener()
init_rerun(session_name="phone_so100_evaluate") init_rerun(session_name="phone_so100_evaluate")
if not robot.is_connected: if not robot.is_connected:
raise ValueError("Robot is not connected!") raise ValueError("Robot is not connected!")
print("Starting evaluate loop...") print("Starting evaluate loop...")
episode_idx = 0 episode_idx = 0
for episode_idx in range(NUM_EPISODES): for episode_idx in range(NUM_EPISODES):
log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}") log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
# Main record loop # Main record loop
@@ -190,10 +194,14 @@ for episode_idx in range(NUM_EPISODES):
dataset.save_episode() dataset.save_episode()
episode_idx += 1 episode_idx += 1
# Clean up # Clean up
log_say("Stop recording") log_say("Stop recording")
robot.disconnect() robot.disconnect()
listener.stop() listener.stop()
dataset.finalize() dataset.finalize()
dataset.push_to_hub() dataset.push_to_hub()
if __name__ == "__main__":
main()
examples/phone_to_so100/record.py
+51 -42
View File
@@ -50,29 +50,33 @@ RESET_TIME_SEC = 30
TASK_DESCRIPTION = "My task description" TASK_DESCRIPTION = "My task description"
HF_REPO_ID = "<hf_username>/<dataset_repo_id>" HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
# Create the robot and teleoperator configurations
camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)} def main():
robot_config = SO100FollowerConfig( # Create the robot and teleoperator configurations
camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
robot_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411", port="/dev/tty.usbmodem5A460814411",
id="my_awesome_follower_arm", id="my_awesome_follower_arm",
cameras=camera_config, cameras=camera_config,
use_degrees=True, use_degrees=True,
) )
teleop_config = PhoneConfig(phone_os=PhoneOS.IOS) # or PhoneOS.ANDROID teleop_config = PhoneConfig(phone_os=PhoneOS.IOS) # or PhoneOS.ANDROID
# Initialize the robot and teleoperator # Initialize the robot and teleoperator
robot = SO100Follower(robot_config) robot = SO100Follower(robot_config)
phone = Phone(teleop_config) phone = Phone(teleop_config)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
kinematics_solver = RobotKinematics( kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()), joint_names=list(robot.bus.motors.keys()),
) )
# Build pipeline to convert phone action to EE action # Build pipeline to convert phone action to EE action
phone_to_robot_ee_pose_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( phone_to_robot_ee_pose_processor = RobotProcessorPipeline[
tuple[RobotAction, RobotObservation], RobotAction
](
steps=[ steps=[
MapPhoneActionToRobotAction(platform=teleop_config.phone_os), MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
EEReferenceAndDelta( EEReferenceAndDelta(
@@ -89,10 +93,10 @@ phone_to_robot_ee_pose_processor = RobotProcessorPipeline[tuple[RobotAction, Rob
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Build pipeline to convert EE action to joints action # Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
steps=[ steps=[
InverseKinematicsEEToJoints( InverseKinematicsEEToJoints(
kinematics=kinematics_solver, kinematics=kinematics_solver,
@@ -102,19 +106,21 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Build pipeline to convert joint observation to EE observation # Build pipeline to convert joint observation to EE observation
robot_joints_to_ee_pose = RobotProcessorPipeline[RobotObservation, RobotObservation]( robot_joints_to_ee_pose = RobotProcessorPipeline[RobotObservation, RobotObservation](
steps=[ steps=[
ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) ForwardKinematicsJointsToEE(
kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())
)
], ],
to_transition=observation_to_transition, to_transition=observation_to_transition,
to_output=transition_to_observation, to_output=transition_to_observation,
) )
# Create the dataset # Create the dataset
dataset = LeRobotDataset.create( dataset = LeRobotDataset.create(
repo_id=HF_REPO_ID, repo_id=HF_REPO_ID,
fps=FPS, fps=FPS,
features=combine_feature_dicts( features=combine_feature_dicts(
@@ -134,23 +140,22 @@ dataset = LeRobotDataset.create(
robot_type=robot.name, robot_type=robot.name,
use_videos=True, use_videos=True,
image_writer_threads=4, image_writer_threads=4,
) )
# Connect the robot and teleoperator # Connect the robot and teleoperator
robot.connect() robot.connect()
phone.connect() phone.connect()
# Initialize the keyboard listener and rerun visualization # Initialize the keyboard listener and rerun visualization
listener, events = init_keyboard_listener() listener, events = init_keyboard_listener()
init_rerun(session_name="phone_so100_record") init_rerun(session_name="phone_so100_record")
if not robot.is_connected or not phone.is_connected: if not robot.is_connected or not phone.is_connected:
raise ValueError("Robot or teleop is not connected!") raise ValueError("Robot or teleop is not connected!")
print("Starting record loop. Move your phone to teleoperate the robot...")
print("Starting record loop. Move your phone to teleoperate the robot...") episode_idx = 0
episode_idx = 0 while episode_idx < NUM_EPISODES and not events["stop_recording"]:
while episode_idx < NUM_EPISODES and not events["stop_recording"]:
log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}") log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
# Main record loop # Main record loop
@@ -195,11 +200,15 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]:
dataset.save_episode() dataset.save_episode()
episode_idx += 1 episode_idx += 1
# Clean up # Clean up
log_say("Stop recording") log_say("Stop recording")
robot.disconnect() robot.disconnect()
phone.disconnect() phone.disconnect()
listener.stop() listener.stop()
dataset.finalize() dataset.finalize()
dataset.push_to_hub() dataset.push_to_hub()
if __name__ == "__main__":
main()
examples/phone_to_so100/replay.py
+32 -26
View File
@@ -29,29 +29,31 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
) )
from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.robots.so100_follower.so100_follower import SO100Follower
from lerobot.utils.constants import ACTION from lerobot.utils.constants import ACTION
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.utils import log_say from lerobot.utils.utils import log_say
EPISODE_IDX = 0 EPISODE_IDX = 0
HF_REPO_ID = "<hf_username>/<dataset_repo_id>" HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
# Initialize the robot config
robot_config = SO100FollowerConfig( def main():
# Initialize the robot config
robot_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
) )
# Initialize the robot # Initialize the robot
robot = SO100Follower(robot_config) robot = SO100Follower(robot_config)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
kinematics_solver = RobotKinematics( kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()), joint_names=list(robot.bus.motors.keys()),
) )
# Build pipeline to convert EE action to joints action # Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
steps=[ steps=[
InverseKinematicsEEToJoints( InverseKinematicsEEToJoints(
kinematics=kinematics_solver, kinematics=kinematics_solver,
@@ -61,23 +63,23 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Fetch the dataset to replay # Fetch the dataset to replay
dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX]) dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0 # Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX) episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
actions = episode_frames.select_columns(ACTION) actions = episode_frames.select_columns(ACTION)
# Connect to the robot # Connect to the robot
robot.connect() robot.connect()
if not robot.is_connected: if not robot.is_connected:
raise ValueError("Robot is not connected!") raise ValueError("Robot is not connected!")
print("Starting replay loop...") print("Starting replay loop...")
log_say(f"Replaying episode {EPISODE_IDX}") log_say(f"Replaying episode {EPISODE_IDX}")
for idx in range(len(episode_frames)): for idx in range(len(episode_frames)):
t0 = time.perf_counter() t0 = time.perf_counter()
# Get recorded action from dataset # Get recorded action from dataset
@@ -94,7 +96,11 @@ for idx in range(len(episode_frames)):
# Send action to robot # Send action to robot
_ = robot.send_action(joint_action) _ = robot.send_action(joint_action)
busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0)) precise_sleep(1.0 / dataset.fps - (time.perf_counter() - t0))
# Clean up # Clean up
robot.disconnect() robot.disconnect()
if __name__ == "__main__":
main()
examples/phone_to_so100/teleoperate.py
+31 -23
View File
@@ -32,30 +32,34 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower
from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
from lerobot.teleoperators.phone.teleop_phone import Phone from lerobot.teleoperators.phone.teleop_phone import Phone
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.visualization_utils import init_rerun, log_rerun_data from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
FPS = 30 FPS = 30
# Initialize the robot and teleoperator
robot_config = SO100FollowerConfig( def main():
# Initialize the robot and teleoperator
robot_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
) )
teleop_config = PhoneConfig(phone_os=PhoneOS.IOS) # or PhoneOS.ANDROID teleop_config = PhoneConfig(phone_os=PhoneOS.IOS) # or PhoneOS.ANDROID
# Initialize the robot and teleoperator # Initialize the robot and teleoperator
robot = SO100Follower(robot_config) robot = SO100Follower(robot_config)
teleop_device = Phone(teleop_config) teleop_device = Phone(teleop_config)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
kinematics_solver = RobotKinematics( kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()), joint_names=list(robot.bus.motors.keys()),
) )
# Build pipeline to convert phone action to ee pose action to joint action # Build pipeline to convert phone action to ee pose action to joint action
phone_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( phone_to_robot_joints_processor = RobotProcessorPipeline[
tuple[RobotAction, RobotObservation], RobotAction
](
steps=[ steps=[
MapPhoneActionToRobotAction(platform=teleop_config.phone_os), MapPhoneActionToRobotAction(platform=teleop_config.phone_os),
EEReferenceAndDelta( EEReferenceAndDelta(
@@ -79,20 +83,20 @@ phone_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, Robo
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Connect to the robot and teleoperator # Connect to the robot and teleoperator
robot.connect() robot.connect()
teleop_device.connect() teleop_device.connect()
# Init rerun viewer # Init rerun viewer
init_rerun(session_name="phone_so100_teleop") init_rerun(session_name="phone_so100_teleop")
if not robot.is_connected or not teleop_device.is_connected: if not robot.is_connected or not teleop_device.is_connected:
raise ValueError("Robot or teleop is not connected!") raise ValueError("Robot or teleop is not connected!")
print("Starting teleop loop. Move your phone to teleoperate the robot...") print("Starting teleop loop. Move your phone to teleoperate the robot...")
while True: while True:
t0 = time.perf_counter() t0 = time.perf_counter()
# Get robot observation # Get robot observation
@@ -110,4 +114,8 @@ while True:
# Visualize # Visualize
log_rerun_data(observation=phone_obs, action=joint_action) log_rerun_data(observation=phone_obs, action=joint_action)
busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) precise_sleep(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
if __name__ == "__main__":
main()
examples/so100_to_so100_EE/evaluate.py
+46 -39
View File
@@ -52,29 +52,31 @@ TASK_DESCRIPTION = "My task description"
HF_MODEL_ID = "<hf_username>/<model_repo_id>" HF_MODEL_ID = "<hf_username>/<model_repo_id>"
HF_DATASET_ID = "<hf_username>/<dataset_repo_id>" HF_DATASET_ID = "<hf_username>/<dataset_repo_id>"
# Create the robot configuration & robot
camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)} def main():
robot_config = SO100FollowerConfig( # Create the robot configuration & robot
camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
robot_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411", port="/dev/tty.usbmodem5A460814411",
id="my_awesome_follower_arm", id="my_awesome_follower_arm",
cameras=camera_config, cameras=camera_config,
use_degrees=True, use_degrees=True,
) )
robot = SO100Follower(robot_config) robot = SO100Follower(robot_config)
# Create policy # Create policy
policy = ACTPolicy.from_pretrained(HF_MODEL_ID) policy = ACTPolicy.from_pretrained(HF_MODEL_ID)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
kinematics_solver = RobotKinematics( kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()), joint_names=list(robot.bus.motors.keys()),
) )
# Build pipeline to convert EE action to joints action # Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
steps=[ steps=[
InverseKinematicsEEToJoints( InverseKinematicsEEToJoints(
kinematics=kinematics_solver, kinematics=kinematics_solver,
@@ -84,20 +86,21 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Build pipeline to convert joints observation to EE observation # Build pipeline to convert joints observation to EE observation
robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation]( robot_joints_to_ee_pose_processor = RobotProcessorPipeline[RobotObservation, RobotObservation](
steps=[ steps=[
ForwardKinematicsJointsToEE(kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())) ForwardKinematicsJointsToEE(
kinematics=kinematics_solver, motor_names=list(robot.bus.motors.keys())
)
], ],
to_transition=observation_to_transition, to_transition=observation_to_transition,
to_output=transition_to_observation, to_output=transition_to_observation,
) )
# Create the dataset
# Create the dataset dataset = LeRobotDataset.create(
dataset = LeRobotDataset.create(
repo_id=HF_DATASET_ID, repo_id=HF_DATASET_ID,
fps=FPS, fps=FPS,
features=combine_feature_dicts( features=combine_feature_dicts(
@@ -122,30 +125,30 @@ dataset = LeRobotDataset.create(
robot_type=robot.name, robot_type=robot.name,
use_videos=True, use_videos=True,
image_writer_threads=4, image_writer_threads=4,
) )
# Build Policy Processors # Build Policy Processors
preprocessor, postprocessor = make_pre_post_processors( preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=policy, policy_cfg=policy,
pretrained_path=HF_MODEL_ID, pretrained_path=HF_MODEL_ID,
dataset_stats=dataset.meta.stats, dataset_stats=dataset.meta.stats,
# The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility. # The inference device is automatically set to match the detected hardware, overriding any previous device settings from training to ensure compatibility.
preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}}, preprocessor_overrides={"device_processor": {"device": str(policy.config.device)}},
) )
# Connect the robot and teleoperator # Connect the robot and teleoperator
robot.connect() robot.connect()
# Initialize the keyboard listener and rerun visualization # Initialize the keyboard listener and rerun visualization
listener, events = init_keyboard_listener() listener, events = init_keyboard_listener()
init_rerun(session_name="so100_so100_evaluate") init_rerun(session_name="so100_so100_evaluate")
if not robot.is_connected: if not robot.is_connected:
raise ValueError("Robot is not connected!") raise ValueError("Robot is not connected!")
print("Starting evaluate loop...") print("Starting evaluate loop...")
episode_idx = 0 episode_idx = 0
for episode_idx in range(NUM_EPISODES): for episode_idx in range(NUM_EPISODES):
log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}") log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
# Main record loop # Main record loop
@@ -191,10 +194,14 @@ for episode_idx in range(NUM_EPISODES):
dataset.save_episode() dataset.save_episode()
episode_idx += 1 episode_idx += 1
# Clean up # Clean up
log_say("Stop recording") log_say("Stop recording")
robot.disconnect() robot.disconnect()
listener.stop() listener.stop()
dataset.finalize() dataset.finalize()
dataset.push_to_hub() dataset.push_to_hub()
if __name__ == "__main__":
main()
examples/so100_to_so100_EE/record.py
+53 -45
View File
@@ -48,33 +48,38 @@ RESET_TIME_SEC = 30
TASK_DESCRIPTION = "My task description" TASK_DESCRIPTION = "My task description"
HF_REPO_ID = "<hf_username>/<dataset_repo_id>" HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
# Create the robot and teleoperator configurations
camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
follower_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", cameras=camera_config, use_degrees=True
)
leader_config = SO100LeaderConfig(port="/dev/tty.usbmodem5A460819811", id="my_awesome_leader_arm")
# Initialize the robot and teleoperator def main():
follower = SO100Follower(follower_config) # Create the robot and teleoperator configurations
leader = SO100Leader(leader_config) camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
follower_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411",
id="my_awesome_follower_arm",
cameras=camera_config,
use_degrees=True,
)
leader_config = SO100LeaderConfig(port="/dev/tty.usbmodem5A460819811", id="my_awesome_leader_arm")
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf # Initialize the robot and teleoperator
follower_kinematics_solver = RobotKinematics( follower = SO100Follower(follower_config)
leader = SO100Leader(leader_config)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
follower_kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(follower.bus.motors.keys()), joint_names=list(follower.bus.motors.keys()),
) )
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
leader_kinematics_solver = RobotKinematics( leader_kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(leader.bus.motors.keys()), joint_names=list(leader.bus.motors.keys()),
) )
# Build pipeline to convert follower joints to EE observation # Build pipeline to convert follower joints to EE observation
follower_joints_to_ee = RobotProcessorPipeline[RobotObservation, RobotObservation]( follower_joints_to_ee = RobotProcessorPipeline[RobotObservation, RobotObservation](
steps=[ steps=[
ForwardKinematicsJointsToEE( ForwardKinematicsJointsToEE(
kinematics=follower_kinematics_solver, motor_names=list(follower.bus.motors.keys()) kinematics=follower_kinematics_solver, motor_names=list(follower.bus.motors.keys())
@@ -82,10 +87,10 @@ follower_joints_to_ee = RobotProcessorPipeline[RobotObservation, RobotObservatio
], ],
to_transition=observation_to_transition, to_transition=observation_to_transition,
to_output=transition_to_observation, to_output=transition_to_observation,
) )
# Build pipeline to convert leader joints to EE action # Build pipeline to convert leader joints to EE action
leader_joints_to_ee = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( leader_joints_to_ee = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
steps=[ steps=[
ForwardKinematicsJointsToEE( ForwardKinematicsJointsToEE(
kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys()) kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
@@ -93,10 +98,10 @@ leader_joints_to_ee = RobotProcessorPipeline[tuple[RobotAction, RobotObservation
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Build pipeline to convert EE action to follower joints # Build pipeline to convert EE action to follower joints
ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
[ [
EEBoundsAndSafety( EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]}, end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
@@ -110,10 +115,10 @@ ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservati
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Create the dataset # Create the dataset
dataset = LeRobotDataset.create( dataset = LeRobotDataset.create(
repo_id=HF_REPO_ID, repo_id=HF_REPO_ID,
fps=FPS, fps=FPS,
features=combine_feature_dicts( features=combine_feature_dicts(
@@ -133,23 +138,22 @@ dataset = LeRobotDataset.create(
robot_type=follower.name, robot_type=follower.name,
use_videos=True, use_videos=True,
image_writer_threads=4, image_writer_threads=4,
) )
# Connect the robot and teleoperator
leader.connect()
follower.connect()
# Connect the robot and teleoperator # Initialize the keyboard listener and rerun visualization
leader.connect() listener, events = init_keyboard_listener()
follower.connect() init_rerun(session_name="recording_phone")
# Initialize the keyboard listener and rerun visualization if not leader.is_connected or not follower.is_connected:
listener, events = init_keyboard_listener()
init_rerun(session_name="recording_phone")
if not leader.is_connected or not follower.is_connected:
raise ValueError("Robot or teleop is not connected!") raise ValueError("Robot or teleop is not connected!")
print("Starting record loop...") print("Starting record loop...")
episode_idx = 0 episode_idx = 0
while episode_idx < NUM_EPISODES and not events["stop_recording"]: while episode_idx < NUM_EPISODES and not events["stop_recording"]:
log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}") log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
# Main record loop # Main record loop
@@ -194,11 +198,15 @@ while episode_idx < NUM_EPISODES and not events["stop_recording"]:
dataset.save_episode() dataset.save_episode()
episode_idx += 1 episode_idx += 1
# Clean up # Clean up
log_say("Stop recording") log_say("Stop recording")
leader.disconnect() leader.disconnect()
follower.disconnect() follower.disconnect()
listener.stop() listener.stop()
dataset.finalize() dataset.finalize()
dataset.push_to_hub() dataset.push_to_hub()
if __name__ == "__main__":
main()
examples/so100_to_so100_EE/replay.py
+32 -26
View File
@@ -30,29 +30,31 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
) )
from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.robots.so100_follower.so100_follower import SO100Follower
from lerobot.utils.constants import ACTION from lerobot.utils.constants import ACTION
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.utils import log_say from lerobot.utils.utils import log_say
EPISODE_IDX = 0 EPISODE_IDX = 0
HF_REPO_ID = "<hf_username>/<dataset_repo_id>" HF_REPO_ID = "<hf_username>/<dataset_repo_id>"
# Initialize the robot config
robot_config = SO100FollowerConfig( def main():
# Initialize the robot config
robot_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
) )
# Initialize the robot # Initialize the robot
robot = SO100Follower(robot_config) robot = SO100Follower(robot_config)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
kinematics_solver = RobotKinematics( kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(robot.bus.motors.keys()), joint_names=list(robot.bus.motors.keys()),
) )
# Build pipeline to convert EE action to joints action # Build pipeline to convert EE action to joints action
robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
steps=[ steps=[
InverseKinematicsEEToJoints( InverseKinematicsEEToJoints(
kinematics=kinematics_solver, kinematics=kinematics_solver,
@@ -62,23 +64,23 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Fetch the dataset to replay # Fetch the dataset to replay
dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX]) dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
# Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0 # Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX) episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
actions = episode_frames.select_columns(ACTION) actions = episode_frames.select_columns(ACTION)
# Connect to the robot # Connect to the robot
robot.connect() robot.connect()
if not robot.is_connected: if not robot.is_connected:
raise ValueError("Robot is not connected!") raise ValueError("Robot is not connected!")
print("Starting replay loop...") print("Starting replay loop...")
log_say(f"Replaying episode {EPISODE_IDX}") log_say(f"Replaying episode {EPISODE_IDX}")
for idx in range(len(episode_frames)): for idx in range(len(episode_frames)):
t0 = time.perf_counter() t0 = time.perf_counter()
# Get recorded action from dataset # Get recorded action from dataset
@@ -95,7 +97,11 @@ for idx in range(len(episode_frames)):
# Send action to robot # Send action to robot
_ = robot.send_action(joint_action) _ = robot.send_action(joint_action)
busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0)) precise_sleep(1.0 / dataset.fps - (time.perf_counter() - t0))
# Clean up # Clean up
robot.disconnect() robot.disconnect()
if __name__ == "__main__":
main()
examples/so100_to_so100_EE/teleoperate.py
+34 -28
View File
@@ -32,37 +32,39 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.robots.so100_follower.so100_follower import SO100Follower
from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderConfig from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderConfig
from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.visualization_utils import init_rerun, log_rerun_data from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
FPS = 30 FPS = 30
# Initialize the robot and teleoperator config
follower_config = SO100FollowerConfig( def main():
# Initialize the robot and teleoperator config
follower_config = SO100FollowerConfig(
port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True port="/dev/tty.usbmodem5A460814411", id="my_awesome_follower_arm", use_degrees=True
) )
leader_config = SO100LeaderConfig(port="/dev/tty.usbmodem5A460819811", id="my_awesome_leader_arm") leader_config = SO100LeaderConfig(port="/dev/tty.usbmodem5A460819811", id="my_awesome_leader_arm")
# Initialize the robot and teleoperator # Initialize the robot and teleoperator
follower = SO100Follower(follower_config) follower = SO100Follower(follower_config)
leader = SO100Leader(leader_config) leader = SO100Leader(leader_config)
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
follower_kinematics_solver = RobotKinematics( follower_kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(follower.bus.motors.keys()), joint_names=list(follower.bus.motors.keys()),
) )
# NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo: https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
leader_kinematics_solver = RobotKinematics( leader_kinematics_solver = RobotKinematics(
urdf_path="./SO101/so101_new_calib.urdf", urdf_path="./SO101/so101_new_calib.urdf",
target_frame_name="gripper_frame_link", target_frame_name="gripper_frame_link",
joint_names=list(leader.bus.motors.keys()), joint_names=list(leader.bus.motors.keys()),
) )
# Build pipeline to convert teleop joints to EE action # Build pipeline to convert teleop joints to EE action
leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction]( leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
steps=[ steps=[
ForwardKinematicsJointsToEE( ForwardKinematicsJointsToEE(
kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys()) kinematics=leader_kinematics_solver, motor_names=list(leader.bus.motors.keys())
@@ -70,10 +72,10 @@ leader_to_ee = RobotProcessorPipeline[RobotAction, RobotAction](
], ],
to_transition=robot_action_to_transition, to_transition=robot_action_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# build pipeline to convert EE action to robot joints # build pipeline to convert EE action to robot joints
ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction]( ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservation], RobotAction](
[ [
EEBoundsAndSafety( EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]}, end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
@@ -87,17 +89,17 @@ ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservati
], ],
to_transition=robot_action_observation_to_transition, to_transition=robot_action_observation_to_transition,
to_output=transition_to_robot_action, to_output=transition_to_robot_action,
) )
# Connect to the robot and teleoperator # Connect to the robot and teleoperator
follower.connect() follower.connect()
leader.connect() leader.connect()
# Init rerun viewer # Init rerun viewer
init_rerun(session_name="so100_so100_EE_teleop") init_rerun(session_name="so100_so100_EE_teleop")
print("Starting teleop loop...") print("Starting teleop loop...")
while True: while True:
t0 = time.perf_counter() t0 = time.perf_counter()
# Get robot observation # Get robot observation
@@ -118,4 +120,8 @@ while True:
# Visualize # Visualize
log_rerun_data(observation=leader_ee_act, action=follower_joints_act) log_rerun_data(observation=leader_ee_act, action=follower_joints_act)
busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) precise_sleep(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
if __name__ == "__main__":
main()
examples/tutorial/act/act_training_example.py
+50 -44
View File
@@ -19,60 +19,62 @@ def make_delta_timestamps(delta_indices: list[int] | None, fps: int) -> list[flo
return [i / fps for i in delta_indices] return [i / fps for i in delta_indices]
output_directory = Path("outputs/robot_learning_tutorial/act") def main():
output_directory.mkdir(parents=True, exist_ok=True) output_directory = Path("outputs/robot_learning_tutorial/act")
output_directory.mkdir(parents=True, exist_ok=True)
# Select your device # Select your device
device = torch.device("mps") # or "cuda" or "cpu" device = torch.device("mps") # or "cuda" or "cpu"
dataset_id = "lerobot/svla_so101_pickplace" dataset_id = "lerobot/svla_so101_pickplace"
# This specifies the inputs the model will be expecting and the outputs it will produce # This specifies the inputs the model will be expecting and the outputs it will produce
dataset_metadata = LeRobotDatasetMetadata(dataset_id) dataset_metadata = LeRobotDatasetMetadata(dataset_id)
features = dataset_to_policy_features(dataset_metadata.features) features = dataset_to_policy_features(dataset_metadata.features)
output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION} output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
input_features = {key: ft for key, ft in features.items() if key not in output_features} input_features = {key: ft for key, ft in features.items() if key not in output_features}
cfg = ACTConfig(input_features=input_features, output_features=output_features) cfg = ACTConfig(input_features=input_features, output_features=output_features)
policy = ACTPolicy(cfg) policy = ACTPolicy(cfg)
preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats) preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
policy.train() policy.train()
policy.to(device) policy.to(device)
# To perform action chunking, ACT expects a given number of actions as targets # To perform action chunking, ACT expects a given number of actions as targets
delta_timestamps = { delta_timestamps = {
"action": make_delta_timestamps(cfg.action_delta_indices, dataset_metadata.fps), "action": make_delta_timestamps(cfg.action_delta_indices, dataset_metadata.fps),
} }
# add image features if they are present # add image features if they are present
delta_timestamps |= { delta_timestamps |= {
k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps) for k in cfg.image_features k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps)
} for k in cfg.image_features
}
# Instantiate the dataset # Instantiate the dataset
dataset = LeRobotDataset(dataset_id, delta_timestamps=delta_timestamps) dataset = LeRobotDataset(dataset_id, delta_timestamps=delta_timestamps)
# Create the optimizer and dataloader for offline training # Create the optimizer and dataloader for offline training
optimizer = cfg.get_optimizer_preset().build(policy.parameters()) optimizer = cfg.get_optimizer_preset().build(policy.parameters())
batch_size = 32 batch_size = 32
dataloader = torch.utils.data.DataLoader( dataloader = torch.utils.data.DataLoader(
dataset, dataset,
batch_size=batch_size, batch_size=batch_size,
shuffle=True, shuffle=True,
pin_memory=device.type != "cpu", pin_memory=device.type != "cpu",
drop_last=True, drop_last=True,
) )
# Number of training steps and logging frequency # Number of training steps and logging frequency
training_steps = 1 training_steps = 1
log_freq = 1 log_freq = 1
# Run training loop # Run training loop
step = 0 step = 0
done = False done = False
while not done: while not done:
for batch in dataloader: for batch in dataloader:
batch = preprocessor(batch) batch = preprocessor(batch)
loss, _ = policy.forward(batch) loss, _ = policy.forward(batch)
@@ -87,12 +89,16 @@ while not done:
done = True done = True
break break
# Save the policy checkpoint, alongside the pre/post processors # Save the policy checkpoint, alongside the pre/post processors
policy.save_pretrained(output_directory) policy.save_pretrained(output_directory)
preprocessor.save_pretrained(output_directory) preprocessor.save_pretrained(output_directory)
postprocessor.save_pretrained(output_directory) postprocessor.save_pretrained(output_directory)
# Save all assets to the Hub # Save all assets to the Hub
policy.push_to_hub("fracapuano/robot_learning_tutorial_act") policy.push_to_hub("<user>/robot_learning_tutorial_act")
preprocessor.push_to_hub("fracapuano/robot_learning_tutorial_act") preprocessor.push_to_hub("<user>/robot_learning_tutorial_act")
postprocessor.push_to_hub("fracapuano/robot_learning_tutorial_act") postprocessor.push_to_hub("<user>/robot_learning_tutorial_act")
if __name__ == "__main__":
main()
examples/tutorial/act/act_using_example.py
+30 -24
View File
@@ -8,37 +8,39 @@ from lerobot.policies.utils import build_inference_frame, make_robot_action
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.robots.so100_follower.so100_follower import SO100Follower
device = torch.device("mps") # or "cuda" or "cpu"
model_id = "fracapuano/robot_learning_tutorial_act"
model = ACTPolicy.from_pretrained(model_id)
dataset_id = "lerobot/svla_so101_pickplace"
# This only downloads the metadata for the dataset, ~10s of MB even for large-scale datasets
dataset_metadata = LeRobotDatasetMetadata(dataset_id)
preprocess, postprocess = make_pre_post_processors(model.config, dataset_stats=dataset_metadata.stats)
# # find ports using lerobot-find-port
follower_port = ... # something like "/dev/tty.usbmodem58760431631"
# # the robot ids are used the load the right calibration files
follower_id = ... # something like "follower_so100"
MAX_EPISODES = 5 MAX_EPISODES = 5
MAX_STEPS_PER_EPISODE = 20 MAX_STEPS_PER_EPISODE = 20
# Robot and environment configuration
# Camera keys must match the name and resolutions of the ones used for training! def main():
# You can check the camera keys expected by a model in the info.json card on the model card on the Hub device = torch.device("mps") # or "cuda" or "cpu"
camera_config = { model_id = "<user>/robot_learning_tutorial_act"
model = ACTPolicy.from_pretrained(model_id)
dataset_id = "lerobot/svla_so101_pickplace"
# This only downloads the metadata for the dataset, ~10s of MB even for large-scale datasets
dataset_metadata = LeRobotDatasetMetadata(dataset_id)
preprocess, postprocess = make_pre_post_processors(model.config, dataset_stats=dataset_metadata.stats)
# # find ports using lerobot-find-port
follower_port = ... # something like "/dev/tty.usbmodem58760431631"
# # the robot ids are used the load the right calibration files
follower_id = ... # something like "follower_so100"
# Robot and environment configuration
# Camera keys must match the name and resolutions of the ones used for training!
# You can check the camera keys expected by a model in the info.json card on the model card on the Hub
camera_config = {
"side": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30), "side": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
"up": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30), "up": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
} }
robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config) robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
robot = SO100Follower(robot_cfg) robot = SO100Follower(robot_cfg)
robot.connect() robot.connect()
for _ in range(MAX_EPISODES): for _ in range(MAX_EPISODES):
for _ in range(MAX_STEPS_PER_EPISODE): for _ in range(MAX_STEPS_PER_EPISODE):
obs = robot.get_observation() obs = robot.get_observation()
obs_frame = build_inference_frame( obs_frame = build_inference_frame(
@@ -55,3 +57,7 @@ for _ in range(MAX_EPISODES):
robot.send_action(action) robot.send_action(action)
print("Episode finished! Starting new episode...") print("Episode finished! Starting new episode...")
if __name__ == "__main__":
main()
examples/tutorial/async-inf/policy_server.py
+11 -5
View File
@@ -1,11 +1,17 @@
from lerobot.async_inference.configs import PolicyServerConfig from lerobot.async_inference.configs import PolicyServerConfig
from lerobot.async_inference.policy_server import serve from lerobot.async_inference.policy_server import serve
host = ... # something like "127.0.0.1" if you're exposing to localhost
port = ... # something like 8080
config = PolicyServerConfig( def main():
host = ... # something like "127.0.0.1" if you're exposing to localhost
port = ... # something like 8080
config = PolicyServerConfig(
host=host, host=host,
port=port, port=port,
) )
serve(config) serve(config)
if __name__ == "__main__":
main()
examples/tutorial/async-inf/robot_client.py
+25 -19
View File
@@ -6,41 +6,43 @@ from lerobot.async_inference.robot_client import RobotClient
from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
from lerobot.robots.so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower import SO100FollowerConfig
# these cameras must match the ones expected by the policy - find your cameras with lerobot-find-cameras
# check the config.json on the Hub for the policy you are using to see the expected camera specs def main():
camera_cfg = { # these cameras must match the ones expected by the policy - find your cameras with lerobot-find-cameras
# check the config.json on the Hub for the policy you are using to see the expected camera specs
camera_cfg = {
"up": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30), "up": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
"side": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30), "side": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
} }
# # find ports using lerobot-find-port # # find ports using lerobot-find-port
follower_port = ... # something like "/dev/tty.usbmodem58760431631" follower_port = ... # something like "/dev/tty.usbmodem58760431631"
# # the robot ids are used the load the right calibration files # # the robot ids are used the load the right calibration files
follower_id = ... # something like "follower_so100" follower_id = ... # something like "follower_so100"
robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_cfg) robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_cfg)
server_address = ... # something like "127.0.0.1:8080" if using localhost server_address = ... # something like "127.0.0.1:8080" if using localhost
# 3. Create client configuration # 3. Create client configuration
client_cfg = RobotClientConfig( client_cfg = RobotClientConfig(
robot=robot_cfg, robot=robot_cfg,
server_address=server_address, server_address=server_address,
policy_device="mps", policy_device="mps",
policy_type="act", policy_type="act",
pretrained_name_or_path="fracapuano/robot_learning_tutorial_act", pretrained_name_or_path="<user>/robot_learning_tutorial_act",
chunk_size_threshold=0.5, # g chunk_size_threshold=0.5, # g
actions_per_chunk=50, # make sure this is less than the max actions of the policy actions_per_chunk=50, # make sure this is less than the max actions of the policy
) )
# 4. Create and start client # 4. Create and start client
client = RobotClient(client_cfg) client = RobotClient(client_cfg)
# 5. Provide a textual description of the task # 5. Provide a textual description of the task
task = ... task = ...
if client.start(): if client.start():
# Start action receiver thread # Start action receiver thread
action_receiver_thread = threading.Thread(target=client.receive_actions, daemon=True) action_receiver_thread = threading.Thread(target=client.receive_actions, daemon=True)
action_receiver_thread.start() action_receiver_thread.start()
@@ -53,3 +55,7 @@ if client.start():
action_receiver_thread.join() action_receiver_thread.join()
# (Optionally) plot the action queue size # (Optionally) plot the action queue size
visualize_action_queue_size(client.action_queue_size) visualize_action_queue_size(client.action_queue_size)
if __name__ == "__main__":
main()
examples/tutorial/diffusion/diffusion_training_example.py
+50 -44
View File
@@ -19,61 +19,63 @@ def make_delta_timestamps(delta_indices: list[int] | None, fps: int) -> list[flo
return [i / fps for i in delta_indices] return [i / fps for i in delta_indices]
output_directory = Path("outputs/robot_learning_tutorial/diffusion") def main():
output_directory.mkdir(parents=True, exist_ok=True) output_directory = Path("outputs/robot_learning_tutorial/diffusion")
output_directory.mkdir(parents=True, exist_ok=True)
# Select your device # Select your device
device = torch.device("mps") # or "cuda" or "cpu" device = torch.device("mps") # or "cuda" or "cpu"
dataset_id = "lerobot/svla_so101_pickplace" dataset_id = "lerobot/svla_so101_pickplace"
# This specifies the inputs the model will be expecting and the outputs it will produce # This specifies the inputs the model will be expecting and the outputs it will produce
dataset_metadata = LeRobotDatasetMetadata(dataset_id) dataset_metadata = LeRobotDatasetMetadata(dataset_id)
features = dataset_to_policy_features(dataset_metadata.features) features = dataset_to_policy_features(dataset_metadata.features)
output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION} output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
input_features = {key: ft for key, ft in features.items() if key not in output_features} input_features = {key: ft for key, ft in features.items() if key not in output_features}
cfg = DiffusionConfig(input_features=input_features, output_features=output_features) cfg = DiffusionConfig(input_features=input_features, output_features=output_features)
policy = DiffusionPolicy(cfg) policy = DiffusionPolicy(cfg)
preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats) preprocessor, postprocessor = make_pre_post_processors(cfg, dataset_stats=dataset_metadata.stats)
policy.train() policy.train()
policy.to(device) policy.to(device)
# To perform action chunking, ACT expects a given number of actions as targets # To perform action chunking, ACT expects a given number of actions as targets
delta_timestamps = { delta_timestamps = {
"observation.state": make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps), "observation.state": make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps),
"action": make_delta_timestamps(cfg.action_delta_indices, dataset_metadata.fps), "action": make_delta_timestamps(cfg.action_delta_indices, dataset_metadata.fps),
} }
# add image features if they are present # add image features if they are present
delta_timestamps |= { delta_timestamps |= {
k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps) for k in cfg.image_features k: make_delta_timestamps(cfg.observation_delta_indices, dataset_metadata.fps)
} for k in cfg.image_features
}
# Instantiate the dataset # Instantiate the dataset
dataset = LeRobotDataset(dataset_id, delta_timestamps=delta_timestamps) dataset = LeRobotDataset(dataset_id, delta_timestamps=delta_timestamps)
# Create the optimizer and dataloader for offline training # Create the optimizer and dataloader for offline training
optimizer = cfg.get_optimizer_preset().build(policy.parameters()) optimizer = cfg.get_optimizer_preset().build(policy.parameters())
batch_size = 32 batch_size = 32
dataloader = torch.utils.data.DataLoader( dataloader = torch.utils.data.DataLoader(
dataset, dataset,
batch_size=batch_size, batch_size=batch_size,
shuffle=True, shuffle=True,
pin_memory=device.type != "cpu", pin_memory=device.type != "cpu",
drop_last=True, drop_last=True,
) )
# Number of training steps and logging frequency # Number of training steps and logging frequency
training_steps = 1 training_steps = 1
log_freq = 1 log_freq = 1
# Run training loop # Run training loop
step = 0 step = 0
done = False done = False
while not done: while not done:
for batch in dataloader: for batch in dataloader:
batch = preprocessor(batch) batch = preprocessor(batch)
loss, _ = policy.forward(batch) loss, _ = policy.forward(batch)
@@ -88,12 +90,16 @@ while not done:
done = True done = True
break break
# Save the policy checkpoint, alongside the pre/post processors # Save the policy checkpoint, alongside the pre/post processors
policy.save_pretrained(output_directory) policy.save_pretrained(output_directory)
preprocessor.save_pretrained(output_directory) preprocessor.save_pretrained(output_directory)
postprocessor.save_pretrained(output_directory) postprocessor.save_pretrained(output_directory)
# Save all assets to the Hub # Save all assets to the Hub
policy.push_to_hub("fracapuano/robot_learning_tutorial_diffusion") policy.push_to_hub("<user>/robot_learning_tutorial_diffusion")
preprocessor.push_to_hub("fracapuano/robot_learning_tutorial_diffusion") preprocessor.push_to_hub("<user>/robot_learning_tutorial_diffusion")
postprocessor.push_to_hub("fracapuano/robot_learning_tutorial_diffusion") postprocessor.push_to_hub("<user>/robot_learning_tutorial_diffusion")
if __name__ == "__main__":
main()
examples/tutorial/diffusion/diffusion_using_example.py
+30 -26
View File
@@ -8,42 +8,42 @@ from lerobot.policies.utils import build_inference_frame, make_robot_action
from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
from lerobot.robots.so100_follower.so100_follower import SO100Follower from lerobot.robots.so100_follower.so100_follower import SO100Follower
device = torch.device("mps") # or "cuda" or "cpu"
model_id = "fracapuano/robot_learning_tutorial_diffusion"
model = DiffusionPolicy.from_pretrained(model_id)
dataset_id = "lerobot/svla_so101_pickplace"
# This only downloads the metadata for the dataset, ~10s of MB even for large-scale datasets
dataset_metadata = LeRobotDatasetMetadata(dataset_id)
preprocess, postprocess = make_pre_post_processors(
model.config, model_id, dataset_stats=dataset_metadata.stats
)
MAX_EPISODES = 5 MAX_EPISODES = 5
MAX_STEPS_PER_EPISODE = 20 MAX_STEPS_PER_EPISODE = 20
# # find ports using lerobot-find-port def main():
follower_port = ... # something like "/dev/tty.usbmodem58760431631" device = torch.device("mps") # or "cuda" or "cpu"
model_id = "<user>/robot_learning_tutorial_diffusion"
# # the robot ids are used the load the right calibration files model = DiffusionPolicy.from_pretrained(model_id)
follower_id = ... # something like "follower_so100"
# Robot and environment configuration dataset_id = "lerobot/svla_so101_pickplace"
# Camera keys must match the name and resolutions of the ones used for training! # This only downloads the metadata for the dataset, ~10s of MB even for large-scale datasets
# You can check the camera keys expected by a model in the info.json card on the model card on the Hub dataset_metadata = LeRobotDatasetMetadata(dataset_id)
camera_config = { preprocess, postprocess = make_pre_post_processors(
model.config, model_id, dataset_stats=dataset_metadata.stats
)
# # find ports using lerobot-find-port
follower_port = ... # something like "/dev/tty.usbmodem58760431631"
# # the robot ids are used the load the right calibration files
follower_id = ... # something like "follower_so100"
# Robot and environment configuration
# Camera keys must match the name and resolutions of the ones used for training!
# You can check the camera keys expected by a model in the info.json card on the model card on the Hub
camera_config = {
"side": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30), "side": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
"up": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30), "up": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
} }
robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config) robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
robot = SO100Follower(robot_cfg) robot = SO100Follower(robot_cfg)
robot.connect() robot.connect()
for _ in range(MAX_EPISODES):
for _ in range(MAX_EPISODES):
for _ in range(MAX_STEPS_PER_EPISODE): for _ in range(MAX_STEPS_PER_EPISODE):
obs = robot.get_observation() obs = robot.get_observation()
obs_frame = build_inference_frame( obs_frame = build_inference_frame(
@@ -58,3 +58,7 @@ for _ in range(MAX_EPISODES):
robot.send_action(action) robot.send_action(action)
print("Episode finished! Starting new episode...") print("Episode finished! Starting new episode...")
if __name__ == "__main__":
main()
examples/tutorial/pi0/using_pi0_example.py
+30 -24
View File
@@ -11,46 +11,48 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower
MAX_EPISODES = 5 MAX_EPISODES = 5
MAX_STEPS_PER_EPISODE = 20 MAX_STEPS_PER_EPISODE = 20
device = torch.device("mps") # or "cuda" or "cpu"
model_id = "lerobot/pi0_base"
model = PI0Policy.from_pretrained(model_id) def main():
device = torch.device("mps") # or "cuda" or "cpu"
model_id = "lerobot/pi0_base"
preprocess, postprocess = make_pre_post_processors( model = PI0Policy.from_pretrained(model_id)
preprocess, postprocess = make_pre_post_processors(
model.config, model.config,
model_id, model_id,
# This overrides allows to run on MPS, otherwise defaults to CUDA (if available) # This overrides allows to run on MPS, otherwise defaults to CUDA (if available)
preprocessor_overrides={"device_processor": {"device": str(device)}}, preprocessor_overrides={"device_processor": {"device": str(device)}},
) )
# find ports using lerobot-find-port # find ports using lerobot-find-port
follower_port = ... # something like "/dev/tty.usbmodem58760431631" follower_port = ... # something like "/dev/tty.usbmodem58760431631"
# the robot ids are used the load the right calibration files # the robot ids are used the load the right calibration files
follower_id = ... # something like "follower_so100" follower_id = ... # something like "follower_so100"
# Robot and environment configuration # Robot and environment configuration
# Camera keys must match the name and resolutions of the ones used for training! # Camera keys must match the name and resolutions of the ones used for training!
# You can check the camera keys expected by a model in the info.json card on the model card on the Hub # You can check the camera keys expected by a model in the info.json card on the model card on the Hub
camera_config = { camera_config = {
"base_0_rgb": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30), "base_0_rgb": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
"left_wrist_0_rgb": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30), "left_wrist_0_rgb": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
"right_wrist_0_rgb": OpenCVCameraConfig(index_or_path=2, width=640, height=480, fps=30), "right_wrist_0_rgb": OpenCVCameraConfig(index_or_path=2, width=640, height=480, fps=30),
} }
robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config) robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
robot = SO100Follower(robot_cfg) robot = SO100Follower(robot_cfg)
robot.connect() robot.connect()
task = "" # something like "pick the red block" task = "" # something like "pick the red block"
robot_type = "" # something like "so100_follower" for multi-embodiment datasets robot_type = "" # something like "so100_follower" for multi-embodiment datasets
# This is used to match the raw observation keys to the keys expected by the policy # This is used to match the raw observation keys to the keys expected by the policy
action_features = hw_to_dataset_features(robot.action_features, "action") action_features = hw_to_dataset_features(robot.action_features, "action")
obs_features = hw_to_dataset_features(robot.observation_features, "observation") obs_features = hw_to_dataset_features(robot.observation_features, "observation")
dataset_features = {**action_features, **obs_features} dataset_features = {**action_features, **obs_features}
for _ in range(MAX_EPISODES): for _ in range(MAX_EPISODES):
for _ in range(MAX_STEPS_PER_EPISODE): for _ in range(MAX_STEPS_PER_EPISODE):
obs = robot.get_observation() obs = robot.get_observation()
obs_frame = build_inference_frame( obs_frame = build_inference_frame(
@@ -65,3 +67,7 @@ for _ in range(MAX_EPISODES):
robot.send_action(action) robot.send_action(action)
print("Episode finished! Starting new episode...") print("Episode finished! Starting new episode...")
if __name__ == "__main__":
main()
examples/tutorial/rl/hilserl_example.py
+61 -59
View File
@@ -20,6 +20,8 @@ from lerobot.teleoperators.utils import TeleopEvents
LOG_EVERY = 10 LOG_EVERY = 10
SEND_EVERY = 10 SEND_EVERY = 10
MAX_EPISODES = 5
MAX_STEPS_PER_EPISODE = 20
def run_learner( def run_learner(
@@ -223,80 +225,76 @@ def make_policy_obs(obs, device: torch.device = "cpu"):
} }
"""Main function - coordinates actor and learner processes.""" def main():
"""Main function - coordinates actor and learner processes."""
device = "mps" # or "cuda" or "cpu" device = "mps" # or "cuda" or "cpu"
output_directory = Path("outputs/robot_learning_tutorial/hil_serl") output_directory = Path("outputs/robot_learning_tutorial/hil_serl")
output_directory.mkdir(parents=True, exist_ok=True) output_directory.mkdir(parents=True, exist_ok=True)
# find ports using lerobot-find-port # find ports using lerobot-find-port
follower_port = ... follower_port = ...
leader_port = ... leader_port = ...
# the robot ids are used the load the right calibration files # the robot ids are used the load the right calibration files
follower_id = ... follower_id = ...
leader_id = ... leader_id = ...
# A pretrained model (to be used in-distribution!) # A pretrained model (to be used in-distribution!)
reward_classifier_id = "fracapuano/reward_classifier_hil_serl_example" reward_classifier_id = "<user>/reward_classifier_hil_serl_example"
reward_classifier = Classifier.from_pretrained(reward_classifier_id) reward_classifier = Classifier.from_pretrained(reward_classifier_id)
reward_classifier.to(device) reward_classifier.to(device)
reward_classifier.eval() reward_classifier.eval()
MAX_EPISODES = 5 # Robot and environment configuration
MAX_STEPS_PER_EPISODE = 20 robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id)
teleop_cfg = SO100LeaderConfig(port=leader_port, id=leader_id)
processor_cfg = HILSerlProcessorConfig(control_mode="leader")
# Robot and environment configuration env_cfg = HILSerlRobotEnvConfig(robot=robot_cfg, teleop=teleop_cfg, processor=processor_cfg)
robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id)
teleop_cfg = SO100LeaderConfig(port=leader_port, id=leader_id)
processor_cfg = HILSerlProcessorConfig(control_mode="leader")
env_cfg = HILSerlRobotEnvConfig(robot=robot_cfg, teleop=teleop_cfg, processor=processor_cfg) # Create robot environment
env, teleop_device = make_robot_env(env_cfg)
# Create robot environment obs_features = hw_to_dataset_features(env.robot.observation_features, "observation")
env, teleop_device = make_robot_env(env_cfg) action_features = hw_to_dataset_features(env.robot.action_features, "action")
obs_features = hw_to_dataset_features(env.robot.observation_features, "observation") # Create SAC policy for action selection
action_features = hw_to_dataset_features(env.robot.action_features, "action") policy_cfg = SACConfig(
# Create SAC policy for action selection
policy_cfg = SACConfig(
device=device, device=device,
input_features=obs_features, input_features=obs_features,
output_features=action_features, output_features=action_features,
) )
policy_actor = SACPolicy(policy_cfg) policy_actor = SACPolicy(policy_cfg)
policy_learner = SACPolicy(policy_cfg) policy_learner = SACPolicy(policy_cfg)
demonstrations_repo_id = "lerobot/example_hil_serl_dataset" demonstrations_repo_id = "lerobot/example_hil_serl_dataset"
offline_dataset = LeRobotDataset(repo_id=demonstrations_repo_id) offline_dataset = LeRobotDataset(repo_id=demonstrations_repo_id)
# Online buffer: initialized from scratch # Online buffer: initialized from scratch
online_replay_buffer = ReplayBuffer(device=device, state_keys=list(obs_features.keys())) online_replay_buffer = ReplayBuffer(device=device, state_keys=list(obs_features.keys()))
# Offline buffer: Created from dataset (pre-populated it with demonstrations) # Offline buffer: Created from dataset (pre-populated it with demonstrations)
offline_replay_buffer = ReplayBuffer.from_lerobot_dataset( offline_replay_buffer = ReplayBuffer.from_lerobot_dataset(
lerobot_dataset=offline_dataset, device=device, state_keys=list(obs_features.keys()) lerobot_dataset=offline_dataset, device=device, state_keys=list(obs_features.keys())
) )
# Create communication channels between learner and actor processes # Create communication channels between learner and actor processes
transitions_queue = mp.Queue(maxsize=10) transitions_queue = mp.Queue(maxsize=10)
parameters_queue = mp.Queue(maxsize=2) parameters_queue = mp.Queue(maxsize=2)
shutdown_event = mp.Event() shutdown_event = mp.Event()
# Signal handler for graceful shutdown
# Signal handler for graceful shutdown def signal_handler(sig):
def signal_handler(sig):
print(f"\nSignal {sig} received, shutting down...") print(f"\nSignal {sig} received, shutting down...")
shutdown_event.set() shutdown_event.set()
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)
signal.signal(signal.SIGINT, signal_handler) # Create processes
signal.signal(signal.SIGTERM, signal_handler) learner_process = mp.Process(
# Create processes
learner_process = mp.Process(
target=run_learner, target=run_learner,
args=( args=(
transitions_queue, transitions_queue,
@@ -307,9 +305,9 @@ learner_process = mp.Process(
offline_replay_buffer, offline_replay_buffer,
), ),
kwargs={"device": device}, # can run on accelerated hardware for training kwargs={"device": device}, # can run on accelerated hardware for training
) )
actor_process = mp.Process( actor_process = mp.Process(
target=run_actor, target=run_actor,
args=( args=(
transitions_queue, transitions_queue,
@@ -321,25 +319,29 @@ actor_process = mp.Process(
output_directory, output_directory,
), ),
kwargs={"device": "cpu"}, # actor is frozen, can run on CPU or accelerate for inference kwargs={"device": "cpu"}, # actor is frozen, can run on CPU or accelerate for inference
) )
learner_process.start() learner_process.start()
actor_process.start() actor_process.start()
try: try:
# Wait for actor to finish (it controls the episode loop) # Wait for actor to finish (it controls the episode loop)
actor_process.join() actor_process.join()
shutdown_event.set() shutdown_event.set()
learner_process.join(timeout=10) learner_process.join(timeout=10)
except KeyboardInterrupt: except KeyboardInterrupt:
print("Main process interrupted") print("Main process interrupted")
shutdown_event.set() shutdown_event.set()
actor_process.join(timeout=5) actor_process.join(timeout=5)
learner_process.join(timeout=10) learner_process.join(timeout=10)
finally: finally:
if learner_process.is_alive(): if learner_process.is_alive():
learner_process.terminate() learner_process.terminate()
if actor_process.is_alive(): if actor_process.is_alive():
actor_process.terminate() actor_process.terminate()
if __name__ == "__main__":
main()
examples/tutorial/rl/reward_classifier_example.py
+28 -23
View File
@@ -4,37 +4,38 @@ from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.policies.factory import make_policy, make_pre_post_processors from lerobot.policies.factory import make_policy, make_pre_post_processors
from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
# Device to use for training
device = "mps" # or "cuda", or "cpu"
# Load the dataset used for training def main():
repo_id = "lerobot/example_hil_serl_dataset" # Device to use for training
dataset = LeRobotDataset(repo_id) device = "mps" # or "cuda", or "cpu"
# Configure the policy to extract features from the image frames # Load the dataset used for training
camera_keys = dataset.meta.camera_keys repo_id = "lerobot/example_hil_serl_dataset"
dataset = LeRobotDataset(repo_id)
config = RewardClassifierConfig( # Configure the policy to extract features from the image frames
camera_keys = dataset.meta.camera_keys
config = RewardClassifierConfig(
num_cameras=len(camera_keys), num_cameras=len(camera_keys),
device=device, device=device,
# backbone model to extract features from the image frames # backbone model to extract features from the image frames
model_name="microsoft/resnet-18", model_name="microsoft/resnet-18",
) )
# Make policy, preprocessor, and optimizer # Make policy, preprocessor, and optimizer
policy = make_policy(config, ds_meta=dataset.meta) policy = make_policy(config, ds_meta=dataset.meta)
optimizer = config.get_optimizer_preset().build(policy.parameters()) optimizer = config.get_optimizer_preset().build(policy.parameters())
preprocessor, _ = make_pre_post_processors(policy_cfg=config, dataset_stats=dataset.meta.stats) preprocessor, _ = make_pre_post_processors(policy_cfg=config, dataset_stats=dataset.meta.stats)
classifier_id = "<user>/reward_classifier_hil_serl_example"
classifier_id = "fracapuano/reward_classifier_hil_serl_example" # Instantiate a dataloader
dataloader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True)
# Instantiate a dataloader # Training loop
dataloader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True) num_epochs = 5
for epoch in range(num_epochs):
# Training loop
num_epochs = 5
for epoch in range(num_epochs):
total_loss = 0 total_loss = 0
total_accuracy = 0 total_accuracy = 0
for batch in dataloader: for batch in dataloader:
@@ -56,7 +57,11 @@ for epoch in range(num_epochs):
avg_accuracy = total_accuracy / len(dataloader) avg_accuracy = total_accuracy / len(dataloader)
print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss:.4f}, Accuracy: {avg_accuracy:.2f}%") print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss:.4f}, Accuracy: {avg_accuracy:.2f}%")
print("Training finished!") print("Training finished!")
# You can now save the trained policy. # You can now save the trained policy.
policy.push_to_hub(classifier_id) policy.push_to_hub(classifier_id)
if __name__ == "__main__":
main()
examples/tutorial/smolvla/using_smolvla_example.py
+30 -24
View File
@@ -11,45 +11,47 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower
MAX_EPISODES = 5 MAX_EPISODES = 5
MAX_STEPS_PER_EPISODE = 20 MAX_STEPS_PER_EPISODE = 20
device = torch.device("mps") # or "cuda" or "cpu"
model_id = "lerobot/smolvla_base"
model = SmolVLAPolicy.from_pretrained(model_id) def main():
device = torch.device("mps") # or "cuda" or "cpu"
model_id = "lerobot/smolvla_base"
preprocess, postprocess = make_pre_post_processors( model = SmolVLAPolicy.from_pretrained(model_id)
preprocess, postprocess = make_pre_post_processors(
model.config, model.config,
model_id, model_id,
# This overrides allows to run on MPS, otherwise defaults to CUDA (if available) # This overrides allows to run on MPS, otherwise defaults to CUDA (if available)
preprocessor_overrides={"device_processor": {"device": str(device)}}, preprocessor_overrides={"device_processor": {"device": str(device)}},
) )
# find ports using lerobot-find-port # find ports using lerobot-find-port
follower_port = ... # something like "/dev/tty.usbmodem58760431631" follower_port = ... # something like "/dev/tty.usbmodem58760431631"
# the robot ids are used the load the right calibration files # the robot ids are used the load the right calibration files
follower_id = ... # something like "follower_so100" follower_id = ... # something like "follower_so100"
# Robot and environment configuration # Robot and environment configuration
# Camera keys must match the name and resolutions of the ones used for training! # Camera keys must match the name and resolutions of the ones used for training!
# You can check the camera keys expected by a model in the info.json card on the model card on the Hub # You can check the camera keys expected by a model in the info.json card on the model card on the Hub
camera_config = { camera_config = {
"camera1": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30), "camera1": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=30),
"camera2": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30), "camera2": OpenCVCameraConfig(index_or_path=1, width=640, height=480, fps=30),
} }
robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config) robot_cfg = SO100FollowerConfig(port=follower_port, id=follower_id, cameras=camera_config)
robot = SO100Follower(robot_cfg) robot = SO100Follower(robot_cfg)
robot.connect() robot.connect()
task = "" # something like "pick the red block" task = "" # something like "pick the red block"
robot_type = "" # something like "so100_follower" for multi-embodiment datasets robot_type = "" # something like "so100_follower" for multi-embodiment datasets
# This is used to match the raw observation keys to the keys expected by the policy # This is used to match the raw observation keys to the keys expected by the policy
action_features = hw_to_dataset_features(robot.action_features, "action") action_features = hw_to_dataset_features(robot.action_features, "action")
obs_features = hw_to_dataset_features(robot.observation_features, "observation") obs_features = hw_to_dataset_features(robot.observation_features, "observation")
dataset_features = {**action_features, **obs_features} dataset_features = {**action_features, **obs_features}
for _ in range(MAX_EPISODES): for _ in range(MAX_EPISODES):
for _ in range(MAX_STEPS_PER_EPISODE): for _ in range(MAX_STEPS_PER_EPISODE):
obs = robot.get_observation() obs = robot.get_observation()
obs_frame = build_inference_frame( obs_frame = build_inference_frame(
@@ -64,3 +66,7 @@ for _ in range(MAX_EPISODES):
robot.send_action(action) robot.send_action(action)
print("Episode finished! Starting new episode...") print("Episode finished! Starting new episode...")
if __name__ == "__main__":
main()
pyproject.toml
+1 -1
View File
@@ -25,7 +25,7 @@ discord = "https://discord.gg/s3KuuzsPFb"
[project] [project]
name = "lerobot" name = "lerobot"
version = "0.4.2" version = "0.4.3"
description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch" description = "🤗 LeRobot: State-of-the-art Machine Learning for Real-World Robotics in Pytorch"
readme = "README.md" readme = "README.md"
license = { text = "Apache-2.0" } license = { text = "Apache-2.0" }
src/lerobot/datasets/lerobot_dataset.py
+1 -2
View File
@@ -16,7 +16,6 @@
import concurrent.futures import concurrent.futures
import contextlib import contextlib
import logging import logging
import platform
import shutil import shutil
import tempfile import tempfile
from collections.abc import Callable from collections.abc import Callable
@@ -1149,7 +1148,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
def save_episode( def save_episode(
self, self,
episode_data: dict | None = None, episode_data: dict | None = None,
parallel_encoding: bool = platform.system() == "Linux", parallel_encoding: bool = True,
) -> None: ) -> None:
""" """
This will save to disk the current episode in self.episode_buffer. This will save to disk the current episode in self.episode_buffer.
src/lerobot/rl/actor.py
+2 -2
View File
@@ -78,7 +78,7 @@ from lerobot.transport.utils import (
transitions_to_bytes, transitions_to_bytes,
) )
from lerobot.utils.random_utils import set_seed from lerobot.utils.random_utils import set_seed
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.transition import ( from lerobot.utils.transition import (
Transition, Transition,
move_state_dict_to_device, move_state_dict_to_device,
@@ -398,7 +398,7 @@ def act_with_policy(
if cfg.env.fps is not None: if cfg.env.fps is not None:
dt_time = time.perf_counter() - start_time dt_time = time.perf_counter() - start_time
busy_wait(1 / cfg.env.fps - dt_time) precise_sleep(1 / cfg.env.fps - dt_time)
# Communication Functions - Group all gRPC/messaging functions # Communication Functions - Group all gRPC/messaging functions
src/lerobot/rl/gym_manipulator.py
+5 -5
View File
@@ -74,7 +74,7 @@ from lerobot.teleoperators import (
from lerobot.teleoperators.teleoperator import Teleoperator from lerobot.teleoperators.teleoperator import Teleoperator
from lerobot.teleoperators.utils import TeleopEvents from lerobot.teleoperators.utils import TeleopEvents
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, REWARD from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, REWARD
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.utils import log_say from lerobot.utils.utils import log_say
logging.basicConfig(level=logging.INFO) logging.basicConfig(level=logging.INFO)
@@ -114,7 +114,7 @@ def reset_follower_position(robot_arm: Robot, target_position: np.ndarray) -> No
for pose in trajectory: for pose in trajectory:
action_dict = dict(zip(current_position_dict, pose, strict=False)) action_dict = dict(zip(current_position_dict, pose, strict=False))
robot_arm.bus.sync_write("Goal_Position", action_dict) robot_arm.bus.sync_write("Goal_Position", action_dict)
busy_wait(0.015) precise_sleep(0.015)
class RobotEnv(gym.Env): class RobotEnv(gym.Env):
@@ -238,7 +238,7 @@ class RobotEnv(gym.Env):
reset_follower_position(self.robot, np.array(self.reset_pose)) reset_follower_position(self.robot, np.array(self.reset_pose))
log_say("Reset the environment done.", play_sounds=True) log_say("Reset the environment done.", play_sounds=True)
busy_wait(self.reset_time_s - (time.perf_counter() - start_time)) precise_sleep(self.reset_time_s - (time.perf_counter() - start_time))
super().reset(seed=seed, options=options) super().reset(seed=seed, options=options)
@@ -713,7 +713,7 @@ def control_loop(
transition = env_processor(transition) transition = env_processor(transition)
# Maintain fps timing # Maintain fps timing
busy_wait(dt - (time.perf_counter() - step_start_time)) precise_sleep(dt - (time.perf_counter() - step_start_time))
if dataset is not None and cfg.dataset.push_to_hub: if dataset is not None and cfg.dataset.push_to_hub:
logging.info("Pushing dataset to hub") logging.info("Pushing dataset to hub")
@@ -745,7 +745,7 @@ def replay_trajectory(
) )
transition = action_processor(transition) transition = action_processor(transition)
env.step(transition[TransitionKey.ACTION]) env.step(transition[TransitionKey.ACTION])
busy_wait(1 / cfg.env.fps - (time.perf_counter() - start_time)) precise_sleep(1 / cfg.env.fps - (time.perf_counter() - start_time))
@parser.wrap() @parser.wrap()
src/lerobot/scripts/lerobot_find_joint_limits.py
+2 -2
View File
@@ -50,7 +50,7 @@ from lerobot.teleoperators import ( # noqa: F401
make_teleoperator_from_config, make_teleoperator_from_config,
so100_leader, so100_leader,
) )
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
@dataclass @dataclass
@@ -114,7 +114,7 @@ def find_joint_and_ee_bounds(cfg: FindJointLimitsConfig):
print(f"Min joint pos position {np.round(min_pos, 4).tolist()}") print(f"Min joint pos position {np.round(min_pos, 4).tolist()}")
break break
busy_wait(0.01) precise_sleep(0.01)
def main(): def main():
src/lerobot/scripts/lerobot_record.py
+2 -2
View File
@@ -119,7 +119,7 @@ from lerobot.utils.control_utils import (
sanity_check_dataset_robot_compatibility, sanity_check_dataset_robot_compatibility,
) )
from lerobot.utils.import_utils import register_third_party_devices from lerobot.utils.import_utils import register_third_party_devices
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.utils import ( from lerobot.utils.utils import (
get_safe_torch_device, get_safe_torch_device,
init_logging, init_logging,
@@ -364,7 +364,7 @@ def record_loop(
log_rerun_data(observation=obs_processed, action=action_values) log_rerun_data(observation=obs_processed, action=action_values)
dt_s = time.perf_counter() - start_loop_t dt_s = time.perf_counter() - start_loop_t
busy_wait(1 / fps - dt_s) precise_sleep(1 / fps - dt_s)
timestamp = time.perf_counter() - start_episode_t timestamp = time.perf_counter() - start_episode_t
src/lerobot/scripts/lerobot_replay.py
+2 -2
View File
@@ -62,7 +62,7 @@ from lerobot.robots import ( # noqa: F401
) )
from lerobot.utils.constants import ACTION from lerobot.utils.constants import ACTION
from lerobot.utils.import_utils import register_third_party_devices from lerobot.utils.import_utils import register_third_party_devices
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.utils import ( from lerobot.utils.utils import (
init_logging, init_logging,
log_say, log_say,
@@ -121,7 +121,7 @@ def replay(cfg: ReplayConfig):
_ = robot.send_action(processed_action) _ = robot.send_action(processed_action)
dt_s = time.perf_counter() - start_episode_t dt_s = time.perf_counter() - start_episode_t
busy_wait(1 / dataset.fps - dt_s) precise_sleep(1 / dataset.fps - dt_s)
robot.disconnect() robot.disconnect()
src/lerobot/scripts/lerobot_teleoperate.py
+5 -4
View File
@@ -89,7 +89,7 @@ from lerobot.teleoperators import ( # noqa: F401
so101_leader, so101_leader,
) )
from lerobot.utils.import_utils import register_third_party_devices from lerobot.utils.import_utils import register_third_party_devices
from lerobot.utils.robot_utils import busy_wait from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.utils import init_logging, move_cursor_up from lerobot.utils.utils import init_logging, move_cursor_up
from lerobot.utils.visualization_utils import init_rerun, log_rerun_data from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
@@ -170,12 +170,13 @@ def teleop_loop(
# Display the final robot action that was sent # Display the final robot action that was sent
for motor, value in robot_action_to_send.items(): for motor, value in robot_action_to_send.items():
print(f"{motor:<{display_len}} | {value:>7.2f}") print(f"{motor:<{display_len}} | {value:>7.2f}")
move_cursor_up(len(robot_action_to_send) + 5) move_cursor_up(len(robot_action_to_send) + 3)
dt_s = time.perf_counter() - loop_start dt_s = time.perf_counter() - loop_start
busy_wait(1 / fps - dt_s) precise_sleep(1 / fps - dt_s)
loop_s = time.perf_counter() - loop_start loop_s = time.perf_counter() - loop_start
print(f"\ntime: {loop_s * 1e3:.2f}ms ({1 / loop_s:.0f} Hz)") print(f"Teleop loop time: {loop_s * 1e3:.2f}ms ({1 / loop_s:.0f} Hz)")
move_cursor_up(1)
if duration is not None and time.perf_counter() - start >= duration: if duration is not None and time.perf_counter() - start >= duration:
return return
src/lerobot/utils/robot_utils.py
+33 -7
View File
@@ -16,14 +16,40 @@ import platform
import time import time
def busy_wait(seconds): def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: float = 0.003):
if platform.system() == "Darwin" or platform.system() == "Windows": """
# On Mac and Windows, `time.sleep` is not accurate and we need to use this while loop trick, Wait for `seconds` with better precision than time.sleep alone at the expense of more CPU usage.
# but it consumes CPU cycles.
Parameters:
- seconds: duration to wait
- spin_threshold: if remaining <= spin_threshold -> spin; otherwise sleep (seconds). Default 10ms
- sleep_margin: when sleeping leave this much time before deadline to avoid oversleep. Default 3ms
Note:
The default parameters are chosen to prioritize timing accuracy over CPU usage for the common 30 FPS use case.
"""
if seconds <= 0:
return
system = platform.system()
# On macOS and Windows the scheduler / sleep granularity can make
# short sleeps inaccurate. Instead of burning CPU for the whole
# duration, sleep for most of the time and spin for the final few
# milliseconds to achieve good accuracy with much lower CPU usage.
if system in ("Darwin", "Windows"):
end_time = time.perf_counter() + seconds end_time = time.perf_counter() + seconds
while time.perf_counter() < end_time: while True:
remaining = end_time - time.perf_counter()
if remaining <= 0:
break
# If there's more than a couple milliseconds left, sleep most
# of the remaining time and leave a small margin for the final spin.
if remaining > spin_threshold:
# Sleep but avoid sleeping past the end by leaving a small margin.
time.sleep(max(remaining - sleep_margin, 0))
else:
# Final short spin to hit precise timing without long sleeps.
pass pass
else: else:
# On Linux time.sleep is accurate # On Linux time.sleep is accurate enough for most uses
if seconds > 0:
time.sleep(seconds) time.sleep(seconds)