admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-06-17 16:27:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into feat/add_pi

Browse Source 
Signed-off-by: Steven Palma <imstevenpmwork@ieee.org>
This commit is contained in:
Steven Palma
2025-09-29 11:51:06 +02:00
committed by GitHub
parent 57c6469b1f 90684a9690
commit 57bcf201ad
55 changed files with 153 additions and 440 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/hilserl.mdx
-3
View File
@@ -95,7 +95,6 @@ class HILSerlProcessorConfig:
class ObservationConfig:
add_joint_velocity_to_observation: bool = False # Add joint velocities to state
add_current_to_observation: bool = False # Add motor currents to state
add_ee_pose_to_observation: bool = False # Add end-effector pose to state
display_cameras: bool = False # Display camera feeds during execution
class ImagePreprocessingConfig:
@@ -105,7 +104,6 @@ class ImagePreprocessingConfig:
class GripperConfig:
use_gripper: bool = True # Enable gripper control
gripper_penalty: float = 0.0 # Penalty for inappropriate gripper usage
gripper_penalty_in_reward: bool = False # Include gripper penalty in reward
class ResetConfig:
fixed_reset_joint_positions: Any | None = None # Joint positions for reset
@@ -288,7 +286,6 @@ You can enable multiple observation processing features simultaneously:
"observation": {
"add_joint_velocity_to_observation": true,
"add_current_to_observation": true,
"add_ee_pose_to_observation": false,
"display_cameras": false
}
}
docs/source/libero.mdx
+1
View File
@@ -106,6 +106,7 @@ For reference, here is the **original dataset** published by Physical Intelligen
lerobot-train \
--policy.type=smolvla \
--policy.repo_id=${HF_USER}/libero-test \
--policy.load_vlm_weights=true \
--dataset.repo_id=HuggingFaceVLA/libero \
--env.type=libero \
--env.task=libero_10 \
docs/source/phone_teleop.mdx
+1 -2
View File
@@ -136,13 +136,12 @@ Additionally you can customize mapping or safety limits by editing the processor
),
```
- The `EEBoundsAndSafety` step clamps EE motion to a workspace and checks for large ee step jumps to ensure safety. The `end_effector_bounds` are the bounds for the EE pose and can be modified to change the workspace. The `max_ee_step_m` and `max_ee_twist_step_rad` are the step limits for the EE pose and can be modified to change the safety limits.
- The `EEBoundsAndSafety` step clamps EE motion to a workspace and checks for large ee step jumps to ensure safety. The `end_effector_bounds` are the bounds for the EE pose and can be modified to change the workspace. The `max_ee_step_m` are the step limits for the EE pose and can be modified to change the safety limits.
```examples/phone_to_so100/teleoperate.py
EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
max_ee_step_m=0.10,
max_ee_twist_step_rad=0.50,
)
```
docs/source/processors_robots_teleop.mdx
+1 -1
View File
@@ -38,7 +38,7 @@ phone_to_robot_ee_pose_processor = RobotProcessorPipeline[RobotAction, RobotActi
kinematics=kinematics_solver, end_effector_step_sizes={"x": 0.5, "y": 0.5, "z": 0.5}, motor_names=list(robot.bus.motors.keys()),
),
EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]}, max_ee_step_m=0.20, max_ee_twist_step_rad=0.50,
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]}, max_ee_step_m=0.20,
),
GripperVelocityToJoint(),
],
examples/phone_to_so100/record.py
-1
View File
@@ -84,7 +84,6 @@ phone_to_robot_ee_pose_processor = RobotProcessorPipeline[tuple[RobotAction, Rob
EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
max_ee_step_m=0.20,
max_ee_twist_step_rad=0.50,
),
GripperVelocityToJoint(speed_factor=20.0),
],
examples/phone_to_so100/teleoperate.py
-1
View File
@@ -67,7 +67,6 @@ phone_to_robot_joints_processor = RobotProcessorPipeline[tuple[RobotAction, Robo
EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
max_ee_step_m=0.10,
max_ee_twist_step_rad=0.50,
),
GripperVelocityToJoint(
speed_factor=20.0,
examples/so100_to_so100_EE/record.py
-1
View File
@@ -101,7 +101,6 @@ ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservati
EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
max_ee_step_m=0.10,
max_ee_twist_step_rad=0.50,
),
InverseKinematicsEEToJoints(
kinematics=follower_kinematics_solver,
examples/so100_to_so100_EE/teleoperate.py
-1
View File
@@ -78,7 +78,6 @@ ee_to_follower_joints = RobotProcessorPipeline[tuple[RobotAction, RobotObservati
EEBoundsAndSafety(
end_effector_bounds={"min": [-1.0, -1.0, -1.0], "max": [1.0, 1.0, 1.0]},
max_ee_step_m=0.10,
max_ee_twist_step_rad=0.50,
),
InverseKinematicsEEToJoints(
kinematics=follower_kinematics_solver,
src/lerobot/async_inference/helpers.py
+1 -2
View File
@@ -38,7 +38,6 @@ from lerobot.utils.constants import OBS_IMAGES, OBS_STATE, OBS_STR
from lerobot.utils.utils import init_logging
Action = torch.Tensor
ActionChunk = torch.Tensor
# observation as received from the robot
RawObservation = dict[str, torch.Tensor]
@@ -53,7 +52,7 @@ Observation = dict[str, torch.Tensor]
def visualize_action_queue_size(action_queue_size: list[int]) -> None:
import matplotlib.pyplot as plt
fig, ax = plt.subplots()
_, ax = plt.subplots()
ax.set_title("Action Queue Size Over Time")
ax.set_xlabel("Environment steps")
ax.set_ylabel("Action Queue Size")
src/lerobot/cameras/utils.py
-4
View File
@@ -15,14 +15,10 @@
# limitations under the License.
import platform
from pathlib import Path
from typing import TypeAlias
from .camera import Camera
from .configs import CameraConfig, Cv2Rotation
IndexOrPath: TypeAlias = int | Path
def make_cameras_from_configs(camera_configs: dict[str, CameraConfig]) -> dict[str, Camera]:
cameras = {}
src/lerobot/configs/types.py
-5
View File
@@ -15,7 +15,6 @@
# https://stackoverflow.com/questions/24481852/serialising-an-enum-member-to-json
from dataclasses import dataclass
from enum import Enum
from typing import Any, Protocol
class FeatureType(str, Enum):
@@ -40,10 +39,6 @@ class NormalizationMode(str, Enum):
QUANTILE10 = "QUANTILE10"
class DictLike(Protocol):
def __getitem__(self, key: Any) -> Any: ...
@dataclass
class PolicyFeature:
type: FeatureType
src/lerobot/datasets/aggregate.py
+19 -26
View File
@@ -93,14 +93,13 @@ def update_data_df(df, src_meta, dst_meta):
pd.DataFrame: Updated DataFrame with adjusted indices.
"""
def _update(row):
row["episode_index"] = row["episode_index"] + dst_meta.info["total_episodes"]
row["index"] = row["index"] + dst_meta.info["total_frames"]
task = src_meta.tasks.iloc[row["task_index"]].name
row["task_index"] = dst_meta.tasks.loc[task].task_index.item()
return row
df["episode_index"] = df["episode_index"] + dst_meta.info["total_episodes"]
df["index"] = df["index"] + dst_meta.info["total_frames"]
return df.apply(_update, axis=1)
src_task_names = src_meta.tasks.index.take(df["task_index"].to_numpy())
df["task_index"] = dst_meta.tasks.loc[src_task_names, "task_index"].to_numpy()
return df
def update_meta_data(
@@ -126,27 +125,21 @@ def update_meta_data(
pd.DataFrame: Updated DataFrame with adjusted indices and timestamps.
"""
def _update(row):
row["meta/episodes/chunk_index"] = row["meta/episodes/chunk_index"] + meta_idx["chunk"]
row["meta/episodes/file_index"] = row["meta/episodes/file_index"] + meta_idx["file"]
row["data/chunk_index"] = row["data/chunk_index"] + data_idx["chunk"]
row["data/file_index"] = row["data/file_index"] + data_idx["file"]
for key, video_idx in videos_idx.items():
row[f"videos/{key}/chunk_index"] = row[f"videos/{key}/chunk_index"] + video_idx["chunk"]
row[f"videos/{key}/file_index"] = row[f"videos/{key}/file_index"] + video_idx["file"]
row[f"videos/{key}/from_timestamp"] = (
row[f"videos/{key}/from_timestamp"] + video_idx["latest_duration"]
)
row[f"videos/{key}/to_timestamp"] = (
row[f"videos/{key}/to_timestamp"] + video_idx["latest_duration"]
)
df["meta/episodes/chunk_index"] = df["meta/episodes/chunk_index"] + meta_idx["chunk"]
df["meta/episodes/file_index"] = df["meta/episodes/file_index"] + meta_idx["file"]
df["data/chunk_index"] = df["data/chunk_index"] + data_idx["chunk"]
df["data/file_index"] = df["data/file_index"] + data_idx["file"]
for key, video_idx in videos_idx.items():
df[f"videos/{key}/chunk_index"] = df[f"videos/{key}/chunk_index"] + video_idx["chunk"]
df[f"videos/{key}/file_index"] = df[f"videos/{key}/file_index"] + video_idx["file"]
df[f"videos/{key}/from_timestamp"] = df[f"videos/{key}/from_timestamp"] + video_idx["latest_duration"]
df[f"videos/{key}/to_timestamp"] = df[f"videos/{key}/to_timestamp"] + video_idx["latest_duration"]
row["dataset_from_index"] = row["dataset_from_index"] + dst_meta.info["total_frames"]
row["dataset_to_index"] = row["dataset_to_index"] + dst_meta.info["total_frames"]
row["episode_index"] = row["episode_index"] + dst_meta.info["total_episodes"]
return row
df["dataset_from_index"] = df["dataset_from_index"] + dst_meta.info["total_frames"]
df["dataset_to_index"] = df["dataset_to_index"] + dst_meta.info["total_frames"]
df["episode_index"] = df["episode_index"] + dst_meta.info["total_episodes"]
return df.apply(_update, axis=1)
return df
def aggregate_datasets(
src/lerobot/datasets/factory.py
+2 -2
View File
@@ -27,7 +27,7 @@ from lerobot.datasets.lerobot_dataset import (
)
from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
from lerobot.datasets.transforms import ImageTransforms
from lerobot.utils.constants import ACTION, OBS_PREFIX
from lerobot.utils.constants import ACTION, OBS_PREFIX, REWARD
IMAGENET_STATS = {
"mean": [[[0.485]], [[0.456]], [[0.406]]], # (c,1,1)
@@ -55,7 +55,7 @@ def resolve_delta_timestamps(
"""
delta_timestamps = {}
for key in ds_meta.features:
if key == "next.reward" and cfg.reward_delta_indices is not None:
if key == REWARD and cfg.reward_delta_indices is not None:
delta_timestamps[key] = [i / ds_meta.fps for i in cfg.reward_delta_indices]
if key == ACTION and cfg.action_delta_indices is not None:
delta_timestamps[key] = [i / ds_meta.fps for i in cfg.action_delta_indices]
src/lerobot/datasets/lerobot_dataset.py
+3 -13
View File
@@ -848,11 +848,6 @@ class LeRobotDataset(torch.utils.data.Dataset):
return item
def _add_padding_keys(self, item: dict, padding: dict[str, list[bool]]) -> dict:
for key, val in padding.items():
item[key] = torch.BoolTensor(val)
return item
def __len__(self):
return self.num_frames
@@ -1032,7 +1027,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
# Reset episode buffer and clean up temporary images (if not already deleted during video encoding)
self.clear_episode_buffer(delete_images=len(self.meta.image_keys) > 0)
def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None):
def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None) -> None:
"""
Batch save videos for multiple episodes.
@@ -1158,7 +1153,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
}
return metadata
def _save_episode_video(self, video_key: str, episode_index: int):
def _save_episode_video(self, video_key: str, episode_index: int) -> dict:
# Encode episode frames into a temporary video
ep_path = self._encode_temporary_episode_video(video_key, episode_index)
ep_size_in_mb = get_video_size_in_mb(ep_path)
@@ -1263,7 +1258,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
if self.image_writer is not None:
self.image_writer.wait_until_done()
def _encode_temporary_episode_video(self, video_key: str, episode_index: int) -> dict:
def _encode_temporary_episode_video(self, video_key: str, episode_index: int) -> Path:
"""
Use ffmpeg to convert frames stored as png into mp4 videos.
Note: `encode_video_frames` is a blocking call. Making it asynchronous shouldn't speedup encoding,
@@ -1396,11 +1391,6 @@ class MultiLeRobotDataset(torch.utils.data.Dataset):
"""
return {repo_id: i for i, repo_id in enumerate(self.repo_ids)}
@property
def repo_index_to_id(self):
"""Return the inverse mapping if repo_id_to_index."""
return {v: k for k, v in self.repo_id_to_index}
@property
def fps(self) -> int:
"""Frames per second used during data collection.
src/lerobot/datasets/push_dataset_to_hub/utils.py
-57
View File
@@ -13,67 +13,10 @@
# 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 inspect
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import datasets
import numpy
import PIL
import torch
from lerobot.datasets.video_utils import encode_video_frames
def concatenate_episodes(ep_dicts):
data_dict = {}
keys = ep_dicts[0].keys()
for key in keys:
if torch.is_tensor(ep_dicts[0][key][0]):
data_dict[key] = torch.cat([ep_dict[key] for ep_dict in ep_dicts])
else:
if key not in data_dict:
data_dict[key] = []
for ep_dict in ep_dicts:
for x in ep_dict[key]:
data_dict[key].append(x)
total_frames = data_dict["frame_index"].shape[0]
data_dict["index"] = torch.arange(0, total_frames, 1)
return data_dict
def save_images_concurrently(imgs_array: numpy.array, out_dir: Path, max_workers: int = 4):
out_dir = Path(out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
def save_image(img_array, i, out_dir):
img = PIL.Image.fromarray(img_array)
img.save(str(out_dir / f"frame_{i:06d}.png"), quality=100)
num_images = len(imgs_array)
with ThreadPoolExecutor(max_workers=max_workers) as executor:
[executor.submit(save_image, imgs_array[i], i, out_dir) for i in range(num_images)]
def get_default_encoding() -> dict:
"""Returns the default ffmpeg encoding parameters used by `encode_video_frames`."""
signature = inspect.signature(encode_video_frames)
return {
k: v.default
for k, v in signature.parameters.items()
if v.default is not inspect.Parameter.empty and k in ["vcodec", "pix_fmt", "g", "crf"]
}
def check_repo_id(repo_id: str) -> None:
if len(repo_id.split("/")) != 2:
raise ValueError(
f"""`repo_id` is expected to contain a community or user id `/` the name of the dataset
(e.g. 'lerobot/pusht'), but contains '{repo_id}'."""
)
# TODO(aliberts): remove
def calculate_episode_data_index(hf_dataset: datasets.Dataset) -> dict[str, torch.Tensor]:
src/lerobot/datasets/streaming_dataset.py
+1 -3
View File
@@ -298,9 +298,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
return padding_mask
def make_frame(
self, dataset_iterator: Backtrackable, previous_dataset_iterator: Backtrackable | None = None
) -> Generator:
def make_frame(self, dataset_iterator: Backtrackable) -> Generator:
"""Makes a frame starting from a dataset iterator"""
item = next(dataset_iterator)
item = item_to_torch(item)
src/lerobot/datasets/utils.py
-49
View File
@@ -67,18 +67,6 @@ DEFAULT_IMAGE_PATH = "images/{image_key}/episode-{episode_index:06d}/frame-{fram
LEGACY_EPISODES_PATH = "meta/episodes.jsonl"
LEGACY_EPISODES_STATS_PATH = "meta/episodes_stats.jsonl"
LEGACY_TASKS_PATH = "meta/tasks.jsonl"
LEGACY_DEFAULT_VIDEO_PATH = "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4"
LEGACY_DEFAULT_PARQUET_PATH = "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet"
DATASET_CARD_TEMPLATE = """
---
# Metadata will go there
---
This dataset was created using [LeRobot](https://github.com/huggingface/lerobot).
## {}
"""
DEFAULT_FEATURES = {
"timestamp": {"dtype": "float32", "shape": (1,), "names": None},
@@ -383,12 +371,6 @@ def load_episodes(local_dir: Path) -> datasets.Dataset:
return episodes
def backward_compatible_episodes_stats(
stats: dict[str, dict[str, np.ndarray]], episodes: list[int]
) -> dict[int, dict[str, dict[str, np.ndarray]]]:
return dict.fromkeys(episodes, stats)
def load_image_as_numpy(
fpath: str | Path, dtype: np.dtype = np.float32, channel_first: bool = True
) -> np.ndarray:
@@ -1346,12 +1328,6 @@ class Backtrackable(Generic[T]):
# When cursor<0, slice so the order remains chronological
return list(self._back_buf)[: self._cursor or None]
def lookahead_buffer(self) -> list[T]:
"""
Return a copy of the current lookahead buffer.
"""
return list(self._ahead_buf)
def can_peek_back(self, steps: int = 1) -> bool:
"""
Check if we can go back `steps` items without raising an IndexError.
@@ -1377,31 +1353,6 @@ class Backtrackable(Generic[T]):
except StopIteration:
return False
def reset_cursor(self) -> None:
"""
Reset cursor to the most recent position (equivalent to calling next()
until you're back to the latest item).
"""
self._cursor = 0
def clear_ahead_buffer(self) -> None:
"""
Clear the ahead buffer, discarding any pre-fetched items.
"""
self._ahead_buf.clear()
def switch_source_iterable(self, new_source: Iterable[T]) -> None:
"""
Switch the source of the backtrackable to a new iterable, keeping the history.
This is useful when iterating over a sequence of datasets. The history from the
previous source is kept, but the lookahead buffer is cleared. The cursor is reset
to the present.
"""
self._source = iter(new_source)
self.clear_ahead_buffer()
self.reset_cursor()
def safe_shard(dataset: datasets.IterableDataset, index: int, num_shards: int) -> datasets.Dataset:
"""
src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py
+15 -3
View File
@@ -34,6 +34,7 @@ python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \
"""
import argparse
import logging
import shutil
from pathlib import Path
from typing import Any
@@ -71,6 +72,7 @@ from lerobot.datasets.utils import (
)
from lerobot.datasets.video_utils import concatenate_video_files, get_video_duration_in_s
from lerobot.utils.constants import HF_LEROBOT_HOME
from lerobot.utils.utils import init_logging
V21 = "v2.1"
@@ -144,6 +146,7 @@ def legacy_load_tasks(local_dir: Path) -> tuple[dict, dict]:
def convert_tasks(root, new_root):
logging.info(f"Converting tasks from {root} to {new_root}")
tasks, _ = legacy_load_tasks(root)
task_indices = tasks.keys()
task_strings = tasks.values()
@@ -185,7 +188,10 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
num_frames = 0
paths_to_cat = []
episodes_metadata = []
for ep_path in ep_paths:
logging.info(f"Converting data files from {len(ep_paths)} episodes")
for ep_path in tqdm.tqdm(ep_paths, desc="convert data files"):
ep_size_in_mb = get_parquet_file_size_in_mb(ep_path)
ep_num_frames = get_parquet_num_frames(ep_path)
ep_metadata = {
@@ -209,7 +215,6 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
# Reset for the next file
size_in_mb = ep_size_in_mb
num_frames = ep_num_frames
paths_to_cat = [ep_path]
chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
@@ -236,6 +241,8 @@ def get_image_keys(root):
def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int):
logging.info(f"Converting videos from {root} to {new_root}")
video_keys = get_video_keys(root)
if len(video_keys) == 0:
return None
@@ -254,7 +261,7 @@ def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int):
episods_metadata = []
num_cameras = len(video_keys)
num_episodes = num_eps_per_cam[0]
for ep_idx in range(num_episodes):
for ep_idx in tqdm.tqdm(range(num_episodes), desc="convert videos"):
# Sanity check
ep_ids = [eps_metadata_per_cam[cam_idx][ep_idx]["episode_index"] for cam_idx in range(num_cameras)]
ep_ids += [ep_idx]
@@ -281,6 +288,7 @@ def convert_videos_of_camera(root: Path, new_root: Path, video_key: str, video_f
duration_in_s = 0.0
paths_to_cat = []
episodes_metadata = []
for ep_path in tqdm.tqdm(ep_paths, desc=f"convert videos of {video_key}"):
ep_size_in_mb = get_video_size_in_mb(ep_path)
ep_duration_in_s = get_video_duration_in_s(ep_path)
@@ -374,6 +382,8 @@ def generate_episode_metadata_dict(
def convert_episodes_metadata(root, new_root, episodes_metadata, episodes_video_metadata=None):
logging.info(f"Converting episodes metadata from {root} to {new_root}")
episodes_legacy_metadata = legacy_load_episodes(root)
episodes_stats = legacy_load_episodes_stats(root)
@@ -405,6 +415,7 @@ def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb):
info["data_path"] = DEFAULT_DATA_PATH
info["video_path"] = DEFAULT_VIDEO_PATH
info["fps"] = int(info["fps"])
logging.info(f"Converting info from {root} to {new_root}")
for key in info["features"]:
if info["features"][key]["dtype"] == "video":
# already has fps in video_info
@@ -469,6 +480,7 @@ def convert_dataset(
if __name__ == "__main__":
init_logging()
parser = argparse.ArgumentParser()
parser.add_argument(
"--repo-id",
src/lerobot/datasets/video_utils.py
+1 -14
View File
@@ -428,7 +428,7 @@ def concatenate_video_files(
with tempfile.NamedTemporaryFile(mode="w", suffix=".ffconcat", delete=False) as tmp_concatenate_file:
tmp_concatenate_file.write("ffconcat version 1.0\n")
for input_path in input_video_paths:
tmp_concatenate_file.write(f"file '{str(input_path)}'\n")
tmp_concatenate_file.write(f"file '{str(input_path.resolve())}'\n")
tmp_concatenate_file.flush()
tmp_concatenate_path = tmp_concatenate_file.name
@@ -585,19 +585,6 @@ def get_video_pixel_channels(pix_fmt: str) -> int:
raise ValueError("Unknown format")
def get_image_pixel_channels(image: Image):
if image.mode == "L":
return 1 # Grayscale
elif image.mode == "LA":
return 2 # Grayscale + Alpha
elif image.mode == "RGB":
return 3 # RGB
elif image.mode == "RGBA":
return 4 # RGBA
else:
raise ValueError("Unknown format")
def get_video_duration_in_s(video_path: Path | str) -> float:
"""
Get the duration of a video file in seconds using PyAV.
src/lerobot/envs/configs.py
-2
View File
@@ -193,7 +193,6 @@ class ObservationConfig:
add_joint_velocity_to_observation: bool = False
add_current_to_observation: bool = False
add_ee_pose_to_observation: bool = False
display_cameras: bool = False
@@ -203,7 +202,6 @@ class GripperConfig:
use_gripper: bool = True
gripper_penalty: float = 0.0
gripper_penalty_in_reward: bool = False
@dataclass
src/lerobot/motors/motors_bus.py
-6
View File
@@ -99,12 +99,6 @@ class Motor:
norm_mode: MotorNormMode
class JointOutOfRangeError(Exception):
def __init__(self, message="Joint is out of range"):
self.message = message
super().__init__(self.message)
class PortHandler(Protocol):
def __init__(self, port_name):
self.is_open: bool
src/lerobot/policies/sac/configuration_sac.py
-2
View File
@@ -139,8 +139,6 @@ class SACConfig(PreTrainedConfig):
# Training parameter
# Number of steps for online training
online_steps: int = 1000000
# Seed for the online environment
online_env_seed: int = 10000
# Capacity of the online replay buffer
online_buffer_capacity: int = 100000
# Capacity of the offline replay buffer
src/lerobot/policies/sac/modeling_sac.py
-12
View File
@@ -1061,15 +1061,3 @@ class TanhMultivariateNormalDiag(TransformedDistribution):
x = transform(x)
return x
def _convert_normalization_params_to_tensor(normalization_params: dict) -> dict:
converted_params = {}
for outer_key, inner_dict in normalization_params.items():
converted_params[outer_key] = {}
for key, value in inner_dict.items():
converted_params[outer_key][key] = torch.tensor(value)
if "image" in outer_key:
converted_params[outer_key][key] = converted_params[outer_key][key].view(3, 1, 1)
return converted_params
src/lerobot/policies/vqbet/configuration_vqbet.py
-2
View File
@@ -82,7 +82,6 @@ class VQBeTConfig(PreTrainedConfig):
gpt_n_head: Number of headers of GPT
gpt_hidden_dim: Size of hidden dimensions of GPT
dropout: Dropout rate for GPT
mlp_hidden_dim: Size of hidden dimensions of offset header / bin prediction headers parts of VQ-BeT
offset_loss_weight: A constant that is multiplied to the offset loss
primary_code_loss_weight: A constant that is multiplied to the primary code prediction loss
secondary_code_loss_weight: A constant that is multiplied to the secondary code prediction loss
@@ -125,7 +124,6 @@ class VQBeTConfig(PreTrainedConfig):
gpt_n_head: int = 8
gpt_hidden_dim: int = 512
dropout: float = 0.1
mlp_hidden_dim: int = 1024
offset_loss_weight: float = 10000.0
primary_code_loss_weight: float = 5.0
secondary_code_loss_weight: float = 0.5
src/lerobot/policies/vqbet/vqbet_utils.py
-10
View File
@@ -231,16 +231,6 @@ class GPT(nn.Module):
torch.nn.init.zeros_(module.bias)
torch.nn.init.ones_(module.weight)
def crop_block_size(self, gpt_block_size):
# model surgery to decrease the block size if necessary
# e.g. we may load the GPT2 pretrained model checkpoint (block size 1024)
# but want to use a smaller block size for some smaller, simpler model
assert gpt_block_size <= self.config.gpt_block_size
self.config.gpt_block_size = gpt_block_size
self.transformer.wpe.weight = nn.Parameter(self.transformer.wpe.weight[:gpt_block_size])
for block in self.transformer.h:
block.attn.bias = block.attn.bias[:, :, :gpt_block_size, :gpt_block_size]
def configure_parameters(self):
"""
This long function is unfortunately doing something very simple and is being very defensive:
src/lerobot/processor/converters.py
+7 -7
View File
@@ -23,7 +23,7 @@ from typing import Any
import numpy as np
import torch
from lerobot.utils.constants import ACTION, OBS_PREFIX
from lerobot.utils.constants import ACTION, DONE, OBS_PREFIX, REWARD, TRUNCATED
from .core import EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey
@@ -355,9 +355,9 @@ def batch_to_transition(batch: dict[str, Any]) -> EnvTransition:
return create_transition(
observation=observation_keys if observation_keys else None,
action=batch.get(ACTION),
reward=batch.get("next.reward", 0.0),
done=batch.get("next.done", False),
truncated=batch.get("next.truncated", False),
reward=batch.get(REWARD, 0.0),
done=batch.get(DONE, False),
truncated=batch.get(TRUNCATED, False),
info=batch.get("info", {}),
complementary_data=complementary_data if complementary_data else None,
)
@@ -380,9 +380,9 @@ def transition_to_batch(transition: EnvTransition) -> dict[str, Any]:
batch = {
ACTION: transition.get(TransitionKey.ACTION),
"next.reward": transition.get(TransitionKey.REWARD, 0.0),
"next.done": transition.get(TransitionKey.DONE, False),
"next.truncated": transition.get(TransitionKey.TRUNCATED, False),
REWARD: transition.get(TransitionKey.REWARD, 0.0),
DONE: transition.get(TransitionKey.DONE, False),
TRUNCATED: transition.get(TransitionKey.TRUNCATED, False),
"info": transition.get(TransitionKey.INFO, {}),
}
src/lerobot/processor/delta_action_processor.py
-2
View File
@@ -83,14 +83,12 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
Attributes:
position_scale: A factor to scale the delta position inputs.
rotation_scale: A factor to scale the delta rotation inputs (currently unused).
noise_threshold: The magnitude below which delta inputs are considered noise
and do not trigger an "enabled" state.
"""
# Scale factors for delta movements
position_scale: float = 1.0
rotation_scale: float = 0.0 # No rotation deltas for gamepad/keyboard
noise_threshold: float = 1e-3 # 1 mm threshold to filter out noise
def action(self, action: RobotAction) -> RobotAction:
src/lerobot/rl/actor.py
-2
View File
@@ -97,8 +97,6 @@ from .gym_manipulator import (
step_env_and_process_transition,
)
ACTOR_SHUTDOWN_TIMEOUT = 30
# Main entry point
src/lerobot/rl/buffer.py
+8 -8
View File
@@ -24,7 +24,7 @@ import torch.nn.functional as F # noqa: N812
from tqdm import tqdm
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.utils.constants import ACTION, OBS_IMAGE
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, REWARD
from lerobot.utils.transition import Transition
@@ -534,8 +534,8 @@ class ReplayBuffer:
features[ACTION] = act_info
# Add "reward" and "done"
features["next.reward"] = {"dtype": "float32", "shape": (1,)}
features["next.done"] = {"dtype": "bool", "shape": (1,)}
features[REWARD] = {"dtype": "float32", "shape": (1,)}
features[DONE] = {"dtype": "bool", "shape": (1,)}
# Add state keys
for key in self.states:
@@ -578,8 +578,8 @@ class ReplayBuffer:
# Fill action, reward, done
frame_dict[ACTION] = self.actions[actual_idx].cpu()
frame_dict["next.reward"] = torch.tensor([self.rewards[actual_idx]], dtype=torch.float32).cpu()
frame_dict["next.done"] = torch.tensor([self.dones[actual_idx]], dtype=torch.bool).cpu()
frame_dict[REWARD] = torch.tensor([self.rewards[actual_idx]], dtype=torch.float32).cpu()
frame_dict[DONE] = torch.tensor([self.dones[actual_idx]], dtype=torch.bool).cpu()
frame_dict["task"] = task_name
# Add complementary_info if available
@@ -648,7 +648,7 @@ class ReplayBuffer:
# Check if the dataset has "next.done" key
sample = dataset[0]
has_done_key = "next.done" in sample
has_done_key = DONE in sample
# Check for complementary_info keys
complementary_info_keys = [key for key in sample if key.startswith("complementary_info.")]
@@ -671,11 +671,11 @@ class ReplayBuffer:
action = current_sample[ACTION].unsqueeze(0) # Add batch dimension
# ----- 3) Reward and done -----
reward = float(current_sample["next.reward"].item()) # ensure float
reward = float(current_sample[REWARD].item()) # ensure float
# Determine done flag - use next.done if available, otherwise infer from episode boundaries
if has_done_key:
done = bool(current_sample["next.done"].item()) # ensure bool
done = bool(current_sample[DONE].item()) # ensure bool
else:
# If this is the last frame or if next frame is in a different episode, mark as done
done = False
src/lerobot/rl/crop_dataset_roi.py
+19 -6
View File
@@ -25,6 +25,7 @@ import torchvision.transforms.functional as F # type: ignore # noqa: N812
from tqdm import tqdm # type: ignore
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.utils.constants import DONE, REWARD
def select_rect_roi(img):
@@ -159,7 +160,7 @@ def get_image_from_lerobot_dataset(dataset: LeRobotDataset):
return image_dict
def convert_lerobot_dataset_to_cropper_lerobot_dataset(
def convert_lerobot_dataset_to_cropped_lerobot_dataset(
original_dataset: LeRobotDataset,
crop_params_dict: dict[str, tuple[int, int, int, int]],
new_repo_id: str,
@@ -189,7 +190,7 @@ def convert_lerobot_dataset_to_cropper_lerobot_dataset(
# 1. Create a new (empty) LeRobotDataset for writing.
new_dataset = LeRobotDataset.create(
repo_id=new_repo_id,
fps=original_dataset.fps,
fps=int(original_dataset.fps),
root=new_dataset_root,
robot_type=original_dataset.meta.robot_type,
features=original_dataset.meta.info["features"],
@@ -212,7 +213,7 @@ def convert_lerobot_dataset_to_cropper_lerobot_dataset(
for key, value in frame.items():
if key in ("task_index", "timestamp", "episode_index", "frame_index", "index", "task"):
continue
if key in ("next.done", "next.reward"):
if key in (DONE, REWARD):
# if not isinstance(value, str) and len(value.shape) == 0:
value = value.unsqueeze(0)
@@ -274,6 +275,12 @@ if __name__ == "__main__":
default="",
help="The natural language task to describe the dataset.",
)
parser.add_argument(
"--new-repo-id",
type=str,
default=None,
help="The repository id for the new cropped and resized dataset. If not provided, it defaults to `repo_id` + '_cropped_resized'.",
)
args = parser.parse_args()
dataset = LeRobotDataset(repo_id=args.repo_id, root=args.root)
@@ -293,10 +300,16 @@ if __name__ == "__main__":
for key, roi in rois.items():
print(f"{key}: {roi}")
new_repo_id = args.repo_id + "_cropped_resized"
new_dataset_root = Path(str(dataset.root) + "_cropped_resized")
new_repo_id = args.new_repo_id if args.new_repo_id else args.repo_id + "_cropped_resized"
cropped_resized_dataset = convert_lerobot_dataset_to_cropper_lerobot_dataset(
if args.new_repo_id:
new_dataset_name = args.new_repo_id.split("/")[-1]
# Parent 1: HF user, Parent 2: HF LeRobot Home
new_dataset_root = dataset.root.parent.parent / new_dataset_name
else:
new_dataset_root = Path(str(dataset.root) + "_cropped_resized")
cropped_resized_dataset = convert_lerobot_dataset_to_cropped_lerobot_dataset(
original_dataset=dataset,
crop_params_dict=rois,
new_repo_id=new_repo_id,
src/lerobot/rl/gym_manipulator.py
+5 -5
View File
@@ -73,7 +73,7 @@ from lerobot.teleoperators import (
)
from lerobot.teleoperators.teleoperator import Teleoperator
from lerobot.teleoperators.utils import TeleopEvents
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, REWARD
from lerobot.utils.robot_utils import busy_wait
from lerobot.utils.utils import log_say
@@ -602,8 +602,8 @@ def control_loop(
action_features = teleop_device.action_features
features = {
ACTION: action_features,
"next.reward": {"dtype": "float32", "shape": (1,), "names": None},
"next.done": {"dtype": "bool", "shape": (1,), "names": None},
REWARD: {"dtype": "float32", "shape": (1,), "names": None},
DONE: {"dtype": "bool", "shape": (1,), "names": None},
}
if use_gripper:
features["complementary_info.discrete_penalty"] = {
@@ -673,8 +673,8 @@ def control_loop(
frame = {
**observations,
ACTION: action_to_record.cpu(),
"next.reward": np.array([transition[TransitionKey.REWARD]], dtype=np.float32),
"next.done": np.array([terminated or truncated], dtype=bool),
REWARD: np.array([transition[TransitionKey.REWARD]], dtype=np.float32),
DONE: np.array([terminated or truncated], dtype=bool),
}
if use_gripper:
discrete_penalty = transition[TransitionKey.COMPLEMENTARY_DATA].get("discrete_penalty", 0.0)
src/lerobot/rl/learner.py
-2
View File
@@ -102,8 +102,6 @@ from lerobot.utils.utils import (
from .learner_service import MAX_WORKERS, SHUTDOWN_TIMEOUT, LearnerService
LOG_PREFIX = "[LEARNER]"
@parser.wrap()
def train_cli(cfg: TrainRLServerPipelineConfig):
src/lerobot/robots/hope_jr/hope_jr_arm.py
+1 -1
View File
@@ -105,7 +105,7 @@ class HopeJrArm(Robot):
def is_calibrated(self) -> bool:
return self.bus.is_calibrated
def calibrate(self, limb_name: str = None) -> None:
def calibrate(self) -> None:
groups = {
"all": list(self.bus.motors.keys()),
"shoulder": ["shoulder_pitch", "shoulder_yaw", "shoulder_roll"],
src/lerobot/robots/so100_follower/robot_kinematic_processor.py
-6
View File
@@ -193,16 +193,12 @@ class EEBoundsAndSafety(RobotActionProcessorStep):
Attributes:
end_effector_bounds: A dictionary with "min" and "max" keys for position clipping.
max_ee_step_m: The maximum allowed change in position (in meters) between steps.
max_ee_twist_step_rad: The maximum allowed change in orientation (in radians) between steps.
_last_pos: Internal state storing the last commanded position.
_last_twist: Internal state storing the last commanded orientation.
"""
end_effector_bounds: dict
max_ee_step_m: float = 0.05
max_ee_twist_step_rad: float = 0.20
_last_pos: np.ndarray | None = field(default=None, init=False, repr=False)
_last_twist: np.ndarray | None = field(default=None, init=False, repr=False)
def action(self, action: RobotAction) -> RobotAction:
x = action["ee.x"]
@@ -233,7 +229,6 @@ class EEBoundsAndSafety(RobotActionProcessorStep):
raise ValueError(f"EE jump {n:.3f}m > {self.max_ee_step_m}m")
self._last_pos = pos
self._last_twist = twist
action["ee.x"] = float(pos[0])
action["ee.y"] = float(pos[1])
@@ -246,7 +241,6 @@ class EEBoundsAndSafety(RobotActionProcessorStep):
def reset(self):
"""Resets the last known position and orientation."""
self._last_pos = None
self._last_twist = None
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
src/lerobot/robots/stretch3/configuration_stretch3.py
-2
View File
@@ -49,5 +49,3 @@ class Stretch3RobotConfig(RobotConfig):
),
}
)
mock: bool = False
src/lerobot/robots/stretch3/robot_stretch3.py
-4
View File
@@ -164,10 +164,6 @@ class Stretch3Robot(Robot):
# TODO(aliberts): return action_sent when motion is limited
return action
def print_logs(self) -> None:
pass
# TODO(aliberts): move robot-specific logs logic here
def teleop_safety_stop(self) -> None:
if self.teleop is not None:
self.teleop._safety_stop(robot=self)
src/lerobot/scripts/lerobot_dataset_viz.py
+5 -5
View File
@@ -75,7 +75,7 @@ import torch.utils.data
import tqdm
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.utils.constants import ACTION, OBS_STATE
from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD
class EpisodeSampler(torch.utils.data.Sampler):
@@ -166,11 +166,11 @@ def visualize_dataset(
for dim_idx, val in enumerate(batch[OBS_STATE][i]):
rr.log(f"state/{dim_idx}", rr.Scalar(val.item()))
if "next.done" in batch:
rr.log("next.done", rr.Scalar(batch["next.done"][i].item()))
if DONE in batch:
rr.log(DONE, rr.Scalar(batch[DONE][i].item()))
if "next.reward" in batch:
rr.log("next.reward", rr.Scalar(batch["next.reward"][i].item()))
if REWARD in batch:
rr.log(REWARD, rr.Scalar(batch[REWARD][i].item()))
if "next.success" in batch:
rr.log("next.success", rr.Scalar(batch["next.success"][i].item()))
src/lerobot/scripts/lerobot_eval.py
+3 -3
View File
@@ -81,7 +81,7 @@ from lerobot.envs.utils import (
from lerobot.policies.factory import make_policy, make_pre_post_processors
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.processor import PolicyAction, PolicyProcessorPipeline
from lerobot.utils.constants import ACTION, OBS_STR
from lerobot.utils.constants import ACTION, DONE, OBS_STR, REWARD
from lerobot.utils.io_utils import write_video
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
@@ -451,9 +451,9 @@ def _compile_episode_data(
"episode_index": torch.tensor([start_episode_index + ep_ix] * (num_frames - 1)),
"frame_index": torch.arange(0, num_frames - 1, 1),
"timestamp": torch.arange(0, num_frames - 1, 1) / fps,
"next.done": rollout_data["done"][ep_ix, : num_frames - 1],
DONE: rollout_data["done"][ep_ix, : num_frames - 1],
"next.success": rollout_data["success"][ep_ix, : num_frames - 1],
"next.reward": rollout_data["reward"][ep_ix, : num_frames - 1].type(torch.float32),
REWARD: rollout_data["reward"][ep_ix, : num_frames - 1].type(torch.float32),
}
# For the last observation frame, all other keys will just be copy padded.
src/lerobot/scripts/lerobot_train.py
+2 -2
View File
@@ -190,8 +190,8 @@ def train(cfg: TrainPipelineConfig):
"device_processor": {"device": device.type},
"normalizer_processor": {"stats": dataset.meta.stats},
}
postprocessor_kwargs["postprocessor_overrides"] = {
"unnormalizer_processor": {"stats": dataset.meta.stats}
processor_kwargs["postprocessor_overrides"] = {
"unnormalizer_processor": {"stats": dataset.meta.stats},
}
preprocessor, postprocessor = make_pre_post_processors(
src/lerobot/teleoperators/gamepad/gamepad_utils.py
-22
View File
@@ -52,10 +52,6 @@ class InputController:
"""Get the current movement deltas (dx, dy, dz) in meters."""
return 0.0, 0.0, 0.0
def should_quit(self):
"""Return True if the user has requested to quit."""
return not self.running
def update(self):
"""Update controller state - call this once per frame."""
pass
@@ -198,14 +194,6 @@ class KeyboardController(InputController):
return delta_x, delta_y, delta_z
def should_quit(self):
"""Return True if ESC was pressed."""
return self.key_states["quit"]
def should_save(self):
"""Return True if Enter was pressed (save episode)."""
return self.key_states["success"] or self.key_states["failure"]
class GamepadController(InputController):
"""Generate motion deltas from gamepad input."""
@@ -351,8 +339,6 @@ class GamepadControllerHID(InputController):
# Button states
self.buttons = {}
self.quit_requested = False
self.save_requested = False
def find_device(self):
"""Look for the gamepad device by vendor and product ID."""
@@ -472,11 +458,3 @@ class GamepadControllerHID(InputController):
delta_z = -self.right_y * self.z_step_size # Up/down
return delta_x, delta_y, delta_z
def should_quit(self):
"""Return True if quit button was pressed."""
return self.quit_requested
def should_save(self):
"""Return True if save button was pressed."""
return self.save_requested
src/lerobot/teleoperators/keyboard/configuration_keyboard.py
+2 -1
View File
@@ -22,8 +22,9 @@ from ..config import TeleoperatorConfig
@TeleoperatorConfig.register_subclass("keyboard")
@dataclass
class KeyboardTeleopConfig(TeleoperatorConfig):
"""KeyboardTeleopConfig"""
# TODO(Steven): Consider setting in here the keys that we want to capture/listen
mock: bool = False
@TeleoperatorConfig.register_subclass("keyboard_ee")
src/lerobot/teleoperators/stretch3_gamepad/configuration_stretch3.py
+1 -1
View File
@@ -22,4 +22,4 @@ from ..config import TeleoperatorConfig
@TeleoperatorConfig.register_subclass("stretch3")
@dataclass
class Stretch3GamePadConfig(TeleoperatorConfig):
mock: bool = False
"""Stretch3GamePadConfig"""
src/lerobot/teleoperators/stretch3_gamepad/stretch3_gamepad.py
-4
View File
@@ -112,10 +112,6 @@ class Stretch3GamePad(Teleoperator):
def send_feedback(self, feedback: np.ndarray) -> None:
pass
def print_logs(self) -> None:
pass
# TODO(aliberts): move robot-specific logs logic here
def disconnect(self) -> None:
self.api.stop()
self.is_connected = False
src/lerobot/utils/constants.py
-1
View File
@@ -33,7 +33,6 @@ TRUNCATED = "next.truncated"
DONE = "next.done"
ROBOTS = "robots"
ROBOT_TYPE = "robot_type"
TELEOPERATORS = "teleoperators"
# files & directories
src/lerobot/utils/control_utils.py
-59
View File
@@ -27,7 +27,6 @@ from typing import Any
import numpy as np
import torch
from deepdiff import DeepDiff
from termcolor import colored
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.utils import DEFAULT_FEATURES
@@ -36,64 +35,6 @@ from lerobot.processor import PolicyAction, PolicyProcessorPipeline
from lerobot.robots import Robot
def log_control_info(robot: Robot, dt_s, episode_index=None, frame_index=None, fps=None):
"""
Logs performance metrics for a single step of the robot control loop.
This function formats and prints a single line of log information, including episode/frame counters,
total loop time (dt), and detailed timings for various robot and camera operations. It can also
highlight performance drops in yellow if the actual FPS is lower than the target FPS.
Args:
robot: The `Robot` instance, used to access its internal logs for detailed timings.
dt_s: The total duration of the control loop step in seconds.
episode_index: The index of the current episode.
frame_index: The index of the current frame within the episode.
fps: The target frames per second, used to check for performance degradation.
"""
log_items = []
if episode_index is not None:
log_items.append(f"ep:{episode_index}")
if frame_index is not None:
log_items.append(f"frame:{frame_index}")
def log_dt(shortname, dt_val_s):
nonlocal log_items, fps
info_str = f"{shortname}:{dt_val_s * 1000:5.2f} ({1 / dt_val_s:3.1f}hz)"
if fps is not None:
actual_fps = 1 / dt_val_s
if actual_fps < fps - 1:
info_str = colored(info_str, "yellow")
log_items.append(info_str)
# total step time displayed in milliseconds and its frequency
log_dt("dt", dt_s)
# TODO(aliberts): move robot-specific logs logic in robot.print_logs()
if not robot.robot_type.startswith("stretch"):
for name in robot.leader_arms:
key = f"read_leader_{name}_pos_dt_s"
if key in robot.logs:
log_dt("dtRlead", robot.logs[key])
for name in robot.follower_arms:
key = f"write_follower_{name}_goal_pos_dt_s"
if key in robot.logs:
log_dt("dtWfoll", robot.logs[key])
key = f"read_follower_{name}_pos_dt_s"
if key in robot.logs:
log_dt("dtRfoll", robot.logs[key])
for name in robot.cameras:
key = f"read_camera_{name}_dt_s"
if key in robot.logs:
log_dt(f"dtR{name}", robot.logs[key])
info_str = " ".join(log_items)
logging.info(info_str)
@cache
def is_headless():
"""
src/lerobot/utils/errors.py
-11
View File
@@ -30,14 +30,3 @@ class DeviceAlreadyConnectedError(ConnectionError):
):
self.message = message
super().__init__(self.message)
class InvalidActionError(ValueError):
"""Exception raised when an action is already invalid."""
def __init__(
self,
message="The action is invalid. Check the value follows what it is expected from the action space.",
):
self.message = message
super().__init__(self.message)
src/lerobot/utils/import_utils.py
-4
View File
@@ -57,8 +57,4 @@ def is_package_available(pkg_name: str, return_version: bool = False) -> tuple[b
return package_exists
_torch_available, _torch_version = is_package_available("torch", return_version=True)
_transformers_available = is_package_available("transformers")
_gym_xarm_available = is_package_available("gym_xarm")
_gym_aloha_available = is_package_available("gym_aloha")
_gym_pusht_available = is_package_available("gym_pusht")
src/lerobot/utils/utils.py
-4
View File
@@ -330,10 +330,6 @@ class TimerManager:
def history(self) -> list[float]:
return deepcopy(self._history)
@property
def fps_history(self) -> list[float]:
return [1.0 / t for t in self._history]
@property
def fps_last(self) -> float:
return 0.0 if self.last == 0 else 1.0 / self.last
tests/datasets/test_datasets.py
+3 -3
View File
@@ -46,7 +46,7 @@ from lerobot.datasets.utils import (
from lerobot.envs.factory import make_env_config
from lerobot.policies.factory import make_policy_config
from lerobot.robots import make_robot_from_config
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, OBS_STR, REWARD
from tests.fixtures.constants import DUMMY_CHW, DUMMY_HWC, DUMMY_REPO_ID
from tests.mocks.mock_robot import MockRobotConfig
from tests.utils import require_x86_64_kernel
@@ -399,8 +399,8 @@ def test_factory(env_name, repo_id, policy_name):
("timestamp", 0, True),
# TODO(rcadene): should we rename it agent_pos?
(OBS_STATE, 1, True),
("next.reward", 0, False),
("next.done", 0, False),
(REWARD, 0, False),
(DONE, 0, False),
]
# test number of dimensions
tests/policies/hilserl/test_modeling_classifier.py
+5 -5
View File
@@ -19,7 +19,7 @@ import torch
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
from lerobot.policies.sac.reward_model.modeling_classifier import ClassifierOutput
from lerobot.utils.constants import OBS_IMAGE
from lerobot.utils.constants import OBS_IMAGE, REWARD
from tests.utils import require_package
@@ -45,7 +45,7 @@ def test_binary_classifier_with_default_params():
OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
}
config.output_features = {
"next.reward": PolicyFeature(type=FeatureType.REWARD, shape=(1,)),
REWARD: PolicyFeature(type=FeatureType.REWARD, shape=(1,)),
}
config.normalization_mapping = {
"VISUAL": NormalizationMode.IDENTITY,
@@ -58,7 +58,7 @@ def test_binary_classifier_with_default_params():
input = {
OBS_IMAGE: torch.rand((batch_size, 3, 128, 128)),
"next.reward": torch.randint(low=0, high=2, size=(batch_size,)).float(),
REWARD: torch.randint(low=0, high=2, size=(batch_size,)).float(),
}
images, labels = classifier.extract_images_and_labels(input)
@@ -87,7 +87,7 @@ def test_multiclass_classifier():
OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
}
config.output_features = {
"next.reward": PolicyFeature(type=FeatureType.REWARD, shape=(num_classes,)),
REWARD: PolicyFeature(type=FeatureType.REWARD, shape=(num_classes,)),
}
config.num_cameras = 1
config.num_classes = num_classes
@@ -97,7 +97,7 @@ def test_multiclass_classifier():
input = {
OBS_IMAGE: torch.rand((batch_size, 3, 128, 128)),
"next.reward": torch.rand((batch_size, num_classes)),
REWARD: torch.rand((batch_size, num_classes)),
}
images, labels = classifier.extract_images_and_labels(input)
tests/policies/test_sac_config.py
-1
View File
@@ -69,7 +69,6 @@ def test_sac_config_default_initialization():
# Training parameters
assert config.online_steps == 1000000
assert config.online_env_seed == 10000
assert config.online_buffer_capacity == 100000
assert config.offline_buffer_capacity == 100000
assert config.async_prefetch is False
tests/processor/test_batch_conversion.py
+34 -34
View File
@@ -2,7 +2,7 @@ import torch
from lerobot.processor import DataProcessorPipeline, TransitionKey
from lerobot.processor.converters import batch_to_transition, transition_to_batch
from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_PREFIX, OBS_STATE
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_PREFIX, OBS_STATE, REWARD, TRUNCATED
def _dummy_batch():
@@ -12,9 +12,9 @@ def _dummy_batch():
f"{OBS_IMAGE}.right": torch.randn(1, 3, 128, 128),
OBS_STATE: torch.tensor([[0.1, 0.2, 0.3, 0.4]]),
ACTION: torch.tensor([[0.5]]),
"next.reward": 1.0,
"next.done": False,
"next.truncated": False,
REWARD: 1.0,
DONE: False,
TRUNCATED: False,
"info": {"key": "value"},
}
@@ -38,9 +38,9 @@ def test_observation_grouping_roundtrip():
# Check other fields
assert torch.allclose(batch_out[ACTION], batch_in[ACTION])
assert batch_out["next.reward"] == batch_in["next.reward"]
assert batch_out["next.done"] == batch_in["next.done"]
assert batch_out["next.truncated"] == batch_in["next.truncated"]
assert batch_out[REWARD] == batch_in[REWARD]
assert batch_out[DONE] == batch_in[DONE]
assert batch_out[TRUNCATED] == batch_in[TRUNCATED]
assert batch_out["info"] == batch_in["info"]
@@ -51,9 +51,9 @@ def test_batch_to_transition_observation_grouping():
f"{OBS_IMAGE}.left": torch.randn(1, 3, 128, 128),
OBS_STATE: [1, 2, 3, 4],
ACTION: torch.tensor([0.1, 0.2, 0.3, 0.4]),
"next.reward": 1.5,
"next.done": True,
"next.truncated": False,
REWARD: 1.5,
DONE: True,
TRUNCATED: False,
"info": {"episode": 42},
}
@@ -115,9 +115,9 @@ def test_transition_to_batch_observation_flattening():
# Check other fields are mapped to next.* format
assert batch[ACTION] == "action_data"
assert batch["next.reward"] == 1.5
assert batch["next.done"]
assert not batch["next.truncated"]
assert batch[REWARD] == 1.5
assert batch[DONE]
assert not batch[TRUNCATED]
assert batch["info"] == {"episode": 42}
@@ -125,9 +125,9 @@ def test_no_observation_keys():
"""Test behavior when there are no observation.* keys."""
batch = {
ACTION: torch.tensor([1.0, 2.0]),
"next.reward": 2.0,
"next.done": False,
"next.truncated": True,
REWARD: 2.0,
DONE: False,
TRUNCATED: True,
"info": {"test": "no_obs"},
}
@@ -146,9 +146,9 @@ def test_no_observation_keys():
# Round trip should work
reconstructed_batch = transition_to_batch(transition)
assert torch.allclose(reconstructed_batch[ACTION], torch.tensor([1.0, 2.0]))
assert reconstructed_batch["next.reward"] == 2.0
assert not reconstructed_batch["next.done"]
assert reconstructed_batch["next.truncated"]
assert reconstructed_batch[REWARD] == 2.0
assert not reconstructed_batch[DONE]
assert reconstructed_batch[TRUNCATED]
assert reconstructed_batch["info"] == {"test": "no_obs"}
@@ -173,9 +173,9 @@ def test_minimal_batch():
reconstructed_batch = transition_to_batch(transition)
assert reconstructed_batch[OBS_STATE] == "minimal_state"
assert torch.allclose(reconstructed_batch[ACTION], torch.tensor([0.5]))
assert reconstructed_batch["next.reward"] == 0.0
assert not reconstructed_batch["next.done"]
assert not reconstructed_batch["next.truncated"]
assert reconstructed_batch[REWARD] == 0.0
assert not reconstructed_batch[DONE]
assert not reconstructed_batch[TRUNCATED]
assert reconstructed_batch["info"] == {}
@@ -197,9 +197,9 @@ def test_empty_batch():
# Round trip
reconstructed_batch = transition_to_batch(transition)
assert reconstructed_batch[ACTION] is None
assert reconstructed_batch["next.reward"] == 0.0
assert not reconstructed_batch["next.done"]
assert not reconstructed_batch["next.truncated"]
assert reconstructed_batch[REWARD] == 0.0
assert not reconstructed_batch[DONE]
assert not reconstructed_batch[TRUNCATED]
assert reconstructed_batch["info"] == {}
@@ -210,9 +210,9 @@ def test_complex_nested_observation():
f"{OBS_IMAGE}.left": {"image": torch.randn(1, 3, 128, 128), "timestamp": 1234567891},
OBS_STATE: torch.randn(7),
ACTION: torch.randn(8),
"next.reward": 3.14,
"next.done": False,
"next.truncated": True,
REWARD: 3.14,
DONE: False,
TRUNCATED: True,
"info": {"episode_length": 200, "success": True},
}
@@ -240,9 +240,9 @@ def test_complex_nested_observation():
assert torch.allclose(batch[ACTION], reconstructed_batch[ACTION])
# Check other fields
assert batch["next.reward"] == reconstructed_batch["next.reward"]
assert batch["next.done"] == reconstructed_batch["next.done"]
assert batch["next.truncated"] == reconstructed_batch["next.truncated"]
assert batch[REWARD] == reconstructed_batch[REWARD]
assert batch[DONE] == reconstructed_batch[DONE]
assert batch[TRUNCATED] == reconstructed_batch[TRUNCATED]
assert batch["info"] == reconstructed_batch["info"]
@@ -267,13 +267,13 @@ def test_custom_converter():
batch = {
OBS_STATE: torch.randn(1, 4),
ACTION: torch.randn(1, 2),
"next.reward": 1.0,
"next.done": False,
REWARD: 1.0,
DONE: False,
}
result = processor(batch)
# Check the reward was doubled by our custom converter
assert result["next.reward"] == 2.0
assert result[REWARD] == 2.0
assert torch.allclose(result[OBS_STATE], batch[OBS_STATE])
assert torch.allclose(result[ACTION], batch[ACTION])
tests/processor/test_converters.py
+3 -3
View File
@@ -9,7 +9,7 @@ from lerobot.processor.converters import (
to_tensor,
transition_to_batch,
)
from lerobot.utils.constants import ACTION, OBS_STATE, OBS_STR
from lerobot.utils.constants import ACTION, DONE, OBS_STATE, OBS_STR, REWARD
# Tests for the unified to_tensor function
@@ -201,8 +201,8 @@ def test_batch_to_transition_with_index_fields():
batch = {
OBS_STATE: torch.randn(1, 7),
ACTION: torch.randn(1, 4),
"next.reward": 1.5,
"next.done": False,
REWARD: 1.5,
DONE: False,
"task": ["pick_cube"],
"index": torch.tensor([42], dtype=torch.int64),
"task_index": torch.tensor([3], dtype=torch.int64),
tests/processor/test_pipeline.py
+7 -7
View File
@@ -35,7 +35,7 @@ from lerobot.processor import (
TransitionKey,
)
from lerobot.processor.converters import create_transition, identity_transition
from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_IMAGES, OBS_STATE
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STATE, REWARD, TRUNCATED
from tests.conftest import assert_contract_is_typed
@@ -258,9 +258,9 @@ def test_step_through_with_dict():
batch = {
OBS_IMAGE: None,
ACTION: None,
"next.reward": 0.0,
"next.done": False,
"next.truncated": False,
REWARD: 0.0,
DONE: False,
TRUNCATED: False,
"info": {},
}
@@ -1843,9 +1843,9 @@ def test_save_load_with_custom_converter_functions():
batch = {
OBS_IMAGE: torch.randn(1, 3, 32, 32),
ACTION: torch.randn(1, 7),
"next.reward": torch.tensor([1.0]),
"next.done": torch.tensor([False]),
"next.truncated": torch.tensor([False]),
REWARD: torch.tensor([1.0]),
DONE: torch.tensor([False]),
TRUNCATED: torch.tensor([False]),
"info": {},
}
tests/utils/test_replay_buffer.py
+3 -3
View File
@@ -22,7 +22,7 @@ import torch
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.rl.buffer import BatchTransition, ReplayBuffer, random_crop_vectorized
from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE, OBS_STR
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_STATE, OBS_STR, REWARD
from tests.fixtures.constants import DUMMY_REPO_ID
@@ -380,9 +380,9 @@ def test_to_lerobot_dataset(tmp_path):
for feature, value in ds[i].items():
if feature == ACTION:
assert torch.equal(value, buffer.actions[i])
elif feature == "next.reward":
elif feature == REWARD:
assert torch.equal(value, buffer.rewards[i])
elif feature == "next.done":
elif feature == DONE:
assert torch.equal(value, buffer.dones[i])
elif feature == OBS_IMAGE:
# Tensor -> numpy is not precise, so we have some diff there