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

big imports refactor

Browse Source
This commit is contained in:
Steven Palma
2026-04-11 15:03:24 +02:00
parent 964acd0151
commit d626964119
183 changed files with 892 additions and 558 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pyproject.toml
+11 -3
View File
@@ -98,7 +98,6 @@ training = [
"lerobot[dataset]",
"accelerate>=1.10.0,<2.0.0",
"wandb>=0.24.0,<0.25.0",
"diffusers>=0.27.2,<0.36.0",
]
hardware = [
"pynput>=1.7.8,<1.9.0",
@@ -116,7 +115,7 @@ build = [
# ── User-facing composite extras (map to CLI scripts) ─────
# lerobot-record, lerobot-replay, lerobot-calibrate, lerobot-teleoperate, etc.
robot = ["lerobot[dataset]", "lerobot[hardware]", "lerobot[viz]"]
# lerobot-eval
# lerobot-eval -- base evaluation framework. You also need the policy's extra (e.g., lerobot[pi]).
evaluation = ["lerobot[av-dep]"]
# lerobot-dataset-viz, lerobot-imgtransform-viz
dataset_viz = ["lerobot[dataset]", "lerobot[viz]"]
@@ -130,6 +129,7 @@ grpcio-dep = ["grpcio==1.73.1", "protobuf>=6.31.1,<6.32.0"]
can-dep = ["python-can>=4.2.0,<5.0.0"]
peft-dep = ["peft>=0.18.0,<1.0.0"]
scipy-dep = ["scipy>=1.14.0,<2.0.0"]
diffusers-dep = ["diffusers>=0.27.2,<0.36.0"]
qwen-vl-utils-dep = ["qwen-vl-utils>=0.0.11,<0.1.0"]
matplotlib-dep = ["matplotlib>=3.10.3,<4.0.0", "contourpy>=1.3.0,<2.0.0"] # NOTE: Explicitly listing contourpy helps the resolver converge faster.
@@ -162,6 +162,7 @@ intelrealsense = [
phone = ["hebi-py>=2.8.0,<2.12.0", "teleop>=0.1.0,<0.2.0", "fastapi<1.0", "lerobot[scipy-dep]"]
# Policies
diffusion = ["lerobot[diffusers-dep]"]
wallx = [
"lerobot[transformers-dep]",
"lerobot[peft]",
@@ -171,10 +172,11 @@ wallx = [
]
pi = ["lerobot[transformers-dep]", "lerobot[scipy-dep]"]
smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14,<0.6.0", "accelerate>=1.7.0,<2.0.0"]
multi_task_dit = ["lerobot[transformers-dep]"]
multi_task_dit = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]"]
groot = [
"lerobot[transformers-dep]",
"lerobot[peft]",
"lerobot[diffusers-dep]",
"dm-tree>=0.1.8,<1.0.0",
"timm>=1.0.0,<1.1.0",
"Pillow>=10.0.0,<13.0.0",
@@ -222,6 +224,7 @@ all = [
"lerobot[reachy2]",
"lerobot[kinematics]",
"lerobot[intelrealsense]",
"lerobot[diffusion]",
"lerobot[wallx]",
"lerobot[pi]",
"lerobot[smolvla]",
@@ -299,6 +302,11 @@ ignore = [
[tool.ruff.lint.per-file-ignores]
"__init__.py" = ["F401", "F403", "E402"]
# E402: require_package() guards must precede the imports they protect
"src/lerobot/policies/wall_x/modeling_wall_x.py" = ["E402"]
"src/lerobot/policies/smolvla/smolvlm_with_expert.py" = ["E402"]
"src/lerobot/policies/xvla/modeling_xvla.py" = ["E402"]
"src/lerobot/policies/sarm/processor_sarm.py" = ["E402"]
"src/lerobot/policies/wall_x/**" = ["N801", "N812", "SIM102", "SIM108", "SIM210", "SIM211", "B006", "B007", "SIM118"] # Supprese these as they are coming from original Qwen2_5_vl code TODO(pepijn): refactor original
[tool.ruff.lint.isort]
src/lerobot/async_inference/__init__.py
+6
View File
@@ -0,0 +1,6 @@
# Async inference server/client.
# Requires: lerobot[async]
from lerobot.utils.import_utils import require_package
require_package("grpcio", extra="async", import_name="grpc")
src/lerobot/async_inference/helpers.py
+1 -1
View File
@@ -22,7 +22,7 @@ from typing import Any
import torch
from lerobot.configs.types import PolicyFeature
from lerobot.configs import PolicyFeature
# NOTE: Configs need to be loaded for the client to be able to instantiate the policy config
from lerobot.policies import ( # noqa: F401
src/lerobot/async_inference/policy_server.py
+1 -1
View File
@@ -38,7 +38,7 @@ import draccus
import grpc
import torch
from lerobot.policies.factory import get_policy_class, make_pre_post_processors
from lerobot.policies import get_policy_class, make_pre_post_processors
from lerobot.processor import PolicyProcessorPipeline
from lerobot.transport import (
services_pb2, # type: ignore
src/lerobot/async_inference/robot_client.py
+2 -2
View File
@@ -47,8 +47,8 @@ import draccus
import grpc
import torch
from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig # noqa: F401
from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig # noqa: F401
from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401
from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401
from lerobot.robots import ( # noqa: F401
Robot,
RobotConfig,
src/lerobot/cameras/__init__.py
+6
View File
@@ -15,3 +15,9 @@
from .camera import Camera
from .configs import CameraConfig, ColorMode, Cv2Backends, Cv2Rotation
from .utils import make_cameras_from_configs
# NOTE: Camera submodule configs and implementations (OpenCVCameraConfig, RealSenseCamera, etc.)
# are intentionally NOT re-exported here to avoid pulling backend-specific dependencies.
# Import from submodules: ``from lerobot.cameras.opencv import OpenCVCameraConfig``
__all__ = ["Camera", "CameraConfig", "ColorMode", "Cv2Backends", "Cv2Rotation", "make_cameras_from_configs"]
src/lerobot/cameras/zmq/image_server.py
+2 -2
View File
@@ -31,8 +31,8 @@ import cv2
import numpy as np
import zmq
from lerobot.cameras.configs import ColorMode
from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
from ..configs import ColorMode
from ..opencv import OpenCVCamera, OpenCVCameraConfig
logger = logging.getLogger(__name__)
src/lerobot/configs/__init__.py
+47
View File
@@ -0,0 +1,47 @@
# 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.
"""
Public API for lerobot configuration types and base config classes.
NOTE: TrainPipelineConfig, EvalPipelineConfig, and TrainRLServerPipelineConfig
are intentionally NOT re-exported here to avoid circular dependencies
(they import lerobot.envs and lerobot.policies at module level).
Import them directly: ``from lerobot.configs.train import TrainPipelineConfig``
"""
from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
from .policies import PreTrainedConfig
from .types import (
FeatureType,
NormalizationMode,
PipelineFeatureType,
PolicyFeature,
RTCAttentionSchedule,
)
__all__ = [
# Types
"FeatureType",
"NormalizationMode",
"PipelineFeatureType",
"PolicyFeature",
"RTCAttentionSchedule",
# Config classes
"DatasetConfig",
"EvalConfig",
"PeftConfig",
"PreTrainedConfig",
"WandBConfig",
]
src/lerobot/configs/eval.py
+3 -2
View File
@@ -19,8 +19,9 @@ from pathlib import Path
from lerobot import envs, policies # noqa: F401
from lerobot.configs import parser
from lerobot.configs.default import EvalConfig
from lerobot.configs.policies import PreTrainedConfig
from .default import EvalConfig
from .policies import PreTrainedConfig
logger = getLogger(__name__)
src/lerobot/configs/policies.py
+3 -3
View File
@@ -26,13 +26,13 @@ from huggingface_hub import hf_hub_download
from huggingface_hub.constants import CONFIG_NAME
from huggingface_hub.errors import HfHubHTTPError
from lerobot.configs.types import FeatureType, PolicyFeature
from lerobot.optim.optimizers import OptimizerConfig
from lerobot.optim.schedulers import LRSchedulerConfig
from lerobot.optim import LRSchedulerConfig, OptimizerConfig
from lerobot.utils.constants import ACTION, OBS_STATE
from lerobot.utils.device_utils import auto_select_torch_device, is_amp_available, is_torch_device_available
from lerobot.utils.hub import HubMixin
from .types import FeatureType, PolicyFeature
T = TypeVar("T", bound="PreTrainedConfig")
logger = getLogger(__name__)
src/lerobot/configs/train.py
+4 -4
View File
@@ -24,12 +24,12 @@ from huggingface_hub.errors import HfHubHTTPError
from lerobot import envs
from lerobot.configs import parser
from lerobot.configs.default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
from lerobot.configs.policies import PreTrainedConfig
from lerobot.optim import OptimizerConfig
from lerobot.optim.schedulers import LRSchedulerConfig
from lerobot.optim import LRSchedulerConfig, OptimizerConfig
from lerobot.utils.hub import HubMixin
from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
from .policies import PreTrainedConfig
TRAIN_CONFIG_NAME = "train_config.json"
src/lerobot/data_processing/sarm_annotations/subtask_annotation.py
+1 -1
View File
@@ -76,7 +76,7 @@ import torch
from pydantic import BaseModel, Field
from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets import LeRobotDataset
# Pydantic Models for SARM Subtask Annotation
src/lerobot/datasets/__init__.py
+34 -10
View File
@@ -19,18 +19,33 @@ from lerobot.utils.import_utils import require_package
require_package("datasets", extra="dataset")
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.datasets.factory import make_dataset
from lerobot.datasets.image_writer import safe_stop_image_writer
from lerobot.datasets.io_utils import load_episodes, write_stats
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.multi_dataset import MultiLeRobotDataset
from lerobot.datasets.pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
from lerobot.datasets.sampler import EpisodeAwareSampler
from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
from lerobot.datasets.video_utils import VideoEncodingManager
from .compute_stats import aggregate_stats, get_feature_stats
from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
from .dataset_tools import (
convert_image_to_video_dataset,
delete_episodes,
merge_datasets,
modify_tasks,
recompute_stats,
remove_feature,
split_dataset,
)
from .factory import make_dataset
from .image_writer import safe_stop_image_writer
from .io_utils import load_episodes, write_stats
from .lerobot_dataset import LeRobotDataset
from .multi_dataset import MultiLeRobotDataset
from .pipeline_features import aggregate_pipeline_dataset_features, create_initial_features
from .sampler import EpisodeAwareSampler
from .streaming_dataset import StreamingLeRobotDataset
from .video_utils import VideoEncodingManager
# NOTE: Low-level I/O functions (cast_stats_to_numpy, get_parquet_file_size_in_mb, etc.)
# and legacy migration constants are intentionally NOT re-exported here.
# Import directly: ``from lerobot.datasets.io_utils import ...``
__all__ = [
"CODEBASE_VERSION",
"EpisodeAwareSampler",
"LeRobotDataset",
"LeRobotDatasetMetadata",
@@ -38,9 +53,18 @@ __all__ = [
"StreamingLeRobotDataset",
"VideoEncodingManager",
"aggregate_pipeline_dataset_features",
"aggregate_stats",
"convert_image_to_video_dataset",
"create_initial_features",
"delete_episodes",
"get_feature_stats",
"load_episodes",
"make_dataset",
"merge_datasets",
"modify_tasks",
"recompute_stats",
"remove_feature",
"safe_stop_image_writer",
"split_dataset",
"write_stats",
]
src/lerobot/datasets/aggregate.py
+6 -6
View File
@@ -23,10 +23,10 @@ import datasets
import pandas as pd
import tqdm
from lerobot.datasets.compute_stats import aggregate_stats
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.datasets.feature_utils import get_hf_features_from_features
from lerobot.datasets.io_utils import (
from .compute_stats import aggregate_stats
from .dataset_metadata import LeRobotDatasetMetadata
from .feature_utils import get_hf_features_from_features
from .io_utils import (
get_file_size_in_mb,
get_parquet_file_size_in_mb,
to_parquet_with_hf_images,
@@ -34,7 +34,7 @@ from lerobot.datasets.io_utils import (
write_stats,
write_tasks,
)
from lerobot.datasets.utils import (
from .utils import (
DEFAULT_CHUNK_SIZE,
DEFAULT_DATA_FILE_SIZE_IN_MB,
DEFAULT_DATA_PATH,
@@ -43,7 +43,7 @@ from lerobot.datasets.utils import (
DEFAULT_VIDEO_PATH,
update_chunk_file_indices,
)
from lerobot.datasets.video_utils import concatenate_video_files, get_video_duration_in_s
from .video_utils import concatenate_video_files, get_video_duration_in_s
def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
src/lerobot/datasets/compute_stats.py
+3 -3
View File
@@ -19,9 +19,11 @@ import logging
import numpy as np
from lerobot.datasets.io_utils import load_image_as_numpy
from lerobot.processor import RelativeActionsProcessorStep
from lerobot.utils.constants import ACTION, OBS_STATE
from .io_utils import load_image_as_numpy
DEFAULT_QUANTILES = [0.01, 0.10, 0.50, 0.90, 0.99]
@@ -696,8 +698,6 @@ def compute_relative_action_stats(
ValueError: If the dataset has fewer frames than ``chunk_size``.
RuntimeError: If no valid (single-episode) chunks are found.
"""
from lerobot.processor.relative_action_processor import RelativeActionsProcessorStep
if exclude_joints is None:
exclude_joints = []
src/lerobot/datasets/dataset_metadata.py
+8 -7
View File
@@ -23,9 +23,12 @@ import pyarrow as pa
import pyarrow.parquet as pq
from huggingface_hub import snapshot_download
from lerobot.datasets.compute_stats import aggregate_stats
from lerobot.datasets.feature_utils import create_empty_dataset_info
from lerobot.datasets.io_utils import (
from lerobot.utils.constants import DEFAULT_FEATURES, HF_LEROBOT_HOME, HF_LEROBOT_HUB_CACHE
from lerobot.utils.feature_utils import _validate_feature_names
from .compute_stats import aggregate_stats
from .feature_utils import create_empty_dataset_info
from .io_utils import (
get_file_size_in_mb,
load_episodes,
load_info,
@@ -37,7 +40,7 @@ from lerobot.datasets.io_utils import (
write_stats,
write_tasks,
)
from lerobot.datasets.utils import (
from .utils import (
DEFAULT_EPISODES_PATH,
INFO_PATH,
check_version_compatibility,
@@ -47,9 +50,7 @@ from lerobot.datasets.utils import (
is_valid_version,
update_chunk_file_indices,
)
from lerobot.datasets.video_utils import get_video_info
from lerobot.utils.constants import DEFAULT_FEATURES, HF_LEROBOT_HOME, HF_LEROBOT_HUB_CACHE
from lerobot.utils.feature_utils import _validate_feature_names
from .video_utils import get_video_info
CODEBASE_VERSION = "v3.0"
src/lerobot/datasets/dataset_reader.py
+4 -4
View File
@@ -21,17 +21,17 @@ from pathlib import Path
import datasets
import torch
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.datasets.feature_utils import (
from .dataset_metadata import LeRobotDatasetMetadata
from .feature_utils import (
check_delta_timestamps,
get_delta_indices,
get_hf_features_from_features,
)
from lerobot.datasets.io_utils import (
from .io_utils import (
hf_transform_to_torch,
load_nested_dataset,
)
from lerobot.datasets.video_utils import decode_video_frames
from .video_utils import decode_video_frames
class DatasetReader:
src/lerobot/datasets/dataset_tools.py
+13 -13
View File
@@ -36,22 +36,25 @@ import pyarrow.parquet as pq
import torch
from tqdm import tqdm
from lerobot.datasets.aggregate import aggregate_datasets
from lerobot.datasets.compute_stats import (
from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME, OBS_IMAGE, OBS_STATE
from lerobot.utils.utils import flatten_dict
from .aggregate import aggregate_datasets
from .compute_stats import (
aggregate_stats,
compute_episode_stats,
compute_relative_action_stats,
)
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.datasets.io_utils import (
from .dataset_metadata import LeRobotDatasetMetadata
from .io_utils import (
get_parquet_file_size_in_mb,
load_episodes,
write_info,
write_stats,
write_tasks,
)
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.utils import (
from .lerobot_dataset import LeRobotDataset
from .utils import (
DATA_DIR,
DEFAULT_CHUNK_SIZE,
DEFAULT_DATA_FILE_SIZE_IN_MB,
@@ -59,8 +62,7 @@ from lerobot.datasets.utils import (
DEFAULT_EPISODES_PATH,
update_chunk_file_indices,
)
from lerobot.datasets.video_utils import encode_video_frames, get_video_info
from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME, OBS_IMAGE, OBS_STATE
from .video_utils import encode_video_frames, get_video_info
def _load_episode_with_stats(src_dataset: LeRobotDataset, episode_idx: int) -> dict:
@@ -829,8 +831,6 @@ def _copy_and_reindex_episodes_metadata(
data_metadata: Dict mapping new episode index to its data file metadata
video_metadata: Optional dict mapping new episode index to its video metadata
"""
from lerobot.utils.utils import flatten_dict
if src_dataset.meta.episodes is None:
src_dataset.meta.episodes = load_episodes(src_dataset.meta.root)
@@ -922,8 +922,8 @@ def _write_parquet(df: pd.DataFrame, path: Path, meta: LeRobotDatasetMetadata) -
This ensures images are properly embedded and the file can be loaded correctly by HF datasets.
"""
from lerobot.datasets.feature_utils import get_hf_features_from_features
from lerobot.datasets.io_utils import embed_images
from .feature_utils import get_hf_features_from_features
from .io_utils import embed_images
hf_features = get_hf_features_from_features(meta.features)
ep_dataset = datasets.Dataset.from_dict(df.to_dict(orient="list"), features=hf_features, split="train")
@@ -1367,7 +1367,7 @@ def _copy_data_without_images(
episode_indices: Episodes to include
img_keys: Image keys to remove
"""
from lerobot.datasets.utils import DATA_DIR
from .utils import DATA_DIR
data_dir = src_dataset.root / DATA_DIR
parquet_files = sorted(data_dir.glob("*/*.parquet"))
src/lerobot/datasets/dataset_writer.py
+7 -7
View File
@@ -31,26 +31,26 @@ import PIL.Image
import pyarrow.parquet as pq
import torch
from lerobot.datasets.compute_stats import compute_episode_stats
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.datasets.feature_utils import (
from .compute_stats import compute_episode_stats
from .dataset_metadata import LeRobotDatasetMetadata
from .feature_utils import (
get_hf_features_from_features,
validate_episode_buffer,
validate_frame,
)
from lerobot.datasets.image_writer import AsyncImageWriter, write_image
from lerobot.datasets.io_utils import (
from .image_writer import AsyncImageWriter, write_image
from .io_utils import (
embed_images,
get_file_size_in_mb,
load_episodes,
write_info,
)
from lerobot.datasets.utils import (
from .utils import (
DEFAULT_EPISODES_PATH,
DEFAULT_IMAGE_PATH,
update_chunk_file_indices,
)
from lerobot.datasets.video_utils import (
from .video_utils import (
StreamingVideoEncoder,
concatenate_video_files,
encode_video_frames,
src/lerobot/datasets/factory.py
+6 -5
View File
@@ -18,15 +18,16 @@ from pprint import pformat
import torch
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs import PreTrainedConfig
from lerobot.configs.train import TrainPipelineConfig
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.multi_dataset import MultiLeRobotDataset
from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
from lerobot.transforms import ImageTransforms
from lerobot.utils.constants import ACTION, IMAGENET_STATS, OBS_PREFIX, REWARD
from .dataset_metadata import LeRobotDatasetMetadata
from .lerobot_dataset import LeRobotDataset
from .multi_dataset import MultiLeRobotDataset
from .streaming_dataset import StreamingLeRobotDataset
def resolve_delta_timestamps(
cfg: PreTrainedConfig, ds_meta: LeRobotDatasetMetadata
src/lerobot/datasets/feature_utils.py
+4 -3
View File
@@ -19,15 +19,16 @@ import datasets
import numpy as np
from PIL import Image as PILImage
from lerobot.datasets.utils import (
from lerobot.utils.constants import DEFAULT_FEATURES
from lerobot.utils.utils import is_valid_numpy_dtype_string
from .utils import (
DEFAULT_CHUNK_SIZE,
DEFAULT_DATA_FILE_SIZE_IN_MB,
DEFAULT_DATA_PATH,
DEFAULT_VIDEO_FILE_SIZE_IN_MB,
DEFAULT_VIDEO_PATH,
)
from lerobot.utils.constants import DEFAULT_FEATURES
from lerobot.utils.utils import is_valid_numpy_dtype_string
def get_hf_features_from_features(features: dict) -> datasets.Features:
src/lerobot/datasets/io_utils.py
+4 -3
View File
@@ -28,7 +28,10 @@ from datasets.table import embed_table_storage
from PIL import Image as PILImage
from torchvision import transforms
from lerobot.datasets.utils import (
from lerobot.utils.io_utils import load_json, write_json
from lerobot.utils.utils import SuppressProgressBars
from .utils import (
DEFAULT_DATA_FILE_SIZE_IN_MB,
DEFAULT_EPISODES_PATH,
DEFAULT_SUBTASKS_PATH,
@@ -40,8 +43,6 @@ from lerobot.datasets.utils import (
serialize_dict,
unflatten_dict,
)
from lerobot.utils.io_utils import load_json, write_json
from lerobot.utils.utils import SuppressProgressBars
def get_parquet_file_size_in_mb(parquet_path: str | Path) -> float:
src/lerobot/datasets/lerobot_dataset.py
+7 -6
View File
@@ -24,20 +24,21 @@ import torch.utils
from huggingface_hub import HfApi, snapshot_download
from huggingface_hub.errors import RevisionNotFoundError
from lerobot.datasets.dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
from lerobot.datasets.dataset_reader import DatasetReader
from lerobot.datasets.dataset_writer import DatasetWriter
from lerobot.datasets.utils import (
from lerobot.utils.constants import HF_LEROBOT_HUB_CACHE
from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
from .dataset_reader import DatasetReader
from .dataset_writer import DatasetWriter
from .utils import (
create_lerobot_dataset_card,
get_safe_version,
is_valid_version,
)
from lerobot.datasets.video_utils import (
from .video_utils import (
StreamingVideoEncoder,
get_safe_default_codec,
resolve_vcodec,
)
from lerobot.utils.constants import HF_LEROBOT_HUB_CACHE
logger = logging.getLogger(__name__)
src/lerobot/datasets/multi_dataset.py
+5 -4
View File
@@ -21,12 +21,13 @@ import datasets
import torch
import torch.utils
from lerobot.datasets.compute_stats import aggregate_stats
from lerobot.datasets.feature_utils import get_hf_features_from_features
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.video_utils import VideoFrame
from lerobot.utils.constants import HF_LEROBOT_HOME
from .compute_stats import aggregate_stats
from .feature_utils import get_hf_features_from_features
from .lerobot_dataset import LeRobotDataset
from .video_utils import VideoFrame
logger = logging.getLogger(__name__)
src/lerobot/datasets/pipeline_features.py
+1 -1
View File
@@ -16,7 +16,7 @@ import re
from collections.abc import Sequence
from typing import Any
from lerobot.configs.types import PipelineFeatureType
from lerobot.configs import PipelineFeatureType
from lerobot.processor import DataProcessorPipeline
from lerobot.types import RobotAction, RobotObservation
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR
src/lerobot/datasets/streaming_dataset.py
+7 -6
View File
@@ -22,20 +22,21 @@ import numpy as np
import torch
from datasets import load_dataset
from lerobot.datasets.dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
from lerobot.datasets.feature_utils import get_delta_indices
from lerobot.datasets.io_utils import item_to_torch
from lerobot.datasets.utils import (
from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE
from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
from .feature_utils import get_delta_indices
from .io_utils import item_to_torch
from .utils import (
check_version_compatibility,
find_float_index,
is_float_in_list,
safe_shard,
)
from lerobot.datasets.video_utils import (
from .video_utils import (
VideoDecoderCache,
decode_video_frames_torchcodec,
)
from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE
class LookBackError(Exception):
src/lerobot/datasets/video_utils.py
+1 -1
View File
@@ -598,7 +598,7 @@ class _CameraEncoderThread(threading.Thread):
self.encoder_threads = encoder_threads
def run(self) -> None:
from lerobot.datasets.compute_stats import RunningQuantileStats, auto_downsample_height_width
from .compute_stats import RunningQuantileStats, auto_downsample_height_width
container = None
output_stream = None
src/lerobot/envs/__init__.py
+18 -1
View File
@@ -12,4 +12,21 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from .configs import AlohaEnv, EnvConfig, HubEnvConfig, PushtEnv # noqa: F401
from .configs import AlohaEnv, EnvConfig, HILSerlRobotEnvConfig, HubEnvConfig, PushtEnv
from .factory import make_env, make_env_config, make_env_pre_post_processors
from .utils import check_env_attributes_and_types, close_envs, env_to_policy_features, preprocess_observation
__all__ = [
"AlohaEnv",
"EnvConfig",
"HILSerlRobotEnvConfig",
"HubEnvConfig",
"PushtEnv",
"check_env_attributes_and_types",
"close_envs",
"env_to_policy_features",
"make_env",
"make_env_config",
"make_env_pre_post_processors",
"preprocess_observation",
]
src/lerobot/envs/configs.py
+4 -11
View File
@@ -23,7 +23,8 @@ import draccus
import gymnasium as gym
from gymnasium.envs.registration import registry as gym_registry
from lerobot.configs.types import FeatureType, PolicyFeature
from lerobot.configs import FeatureType, PolicyFeature
from lerobot.processor import IsaaclabArenaProcessorStep, LiberoProcessorStep, PolicyProcessorPipeline
from lerobot.robots import RobotConfig
from lerobot.teleoperators.config import TeleoperatorConfig
from lerobot.utils.constants import (
@@ -124,8 +125,6 @@ class EnvConfig(draccus.ChoiceRegistry, abc.ABC):
def get_env_processors(self):
"""Return (preprocessor, postprocessor) for this env. Default: identity."""
from lerobot.processor.pipeline import PolicyProcessorPipeline
return PolicyProcessorPipeline(steps=[]), PolicyProcessorPipeline(steps=[])
@@ -418,7 +417,7 @@ class LiberoEnv(EnvConfig):
return kwargs
def create_envs(self, n_envs: int, use_async_envs: bool = False):
from lerobot.envs.libero import create_libero_envs
from .libero import create_libero_envs
if self.task is None:
raise ValueError("LiberoEnv requires a task to be specified")
@@ -436,9 +435,6 @@ class LiberoEnv(EnvConfig):
)
def get_env_processors(self):
from lerobot.processor.env_processor import LiberoProcessorStep
from lerobot.processor.pipeline import PolicyProcessorPipeline
return (
PolicyProcessorPipeline(steps=[LiberoProcessorStep()]),
PolicyProcessorPipeline(steps=[]),
@@ -487,7 +483,7 @@ class MetaworldEnv(EnvConfig):
}
def create_envs(self, n_envs: int, use_async_envs: bool = False):
from lerobot.envs.metaworld import create_metaworld_envs
from .metaworld import create_metaworld_envs
if self.task is None:
raise ValueError("MetaWorld requires a task to be specified")
@@ -568,9 +564,6 @@ class IsaaclabArenaEnv(HubEnvConfig):
return {}
def get_env_processors(self):
from lerobot.processor.env_processor import IsaaclabArenaProcessorStep
from lerobot.processor.pipeline import PolicyProcessorPipeline
state_keys = tuple(k.strip() for k in (self.state_keys or "").split(",") if k.strip())
camera_keys = tuple(k.strip() for k in (self.camera_keys or "").split(",") if k.strip())
if not state_keys and not camera_keys:
src/lerobot/envs/factory.py
+2 -2
View File
@@ -19,8 +19,8 @@ from typing import Any
import gymnasium as gym
from lerobot.envs.configs import EnvConfig, HubEnvConfig
from lerobot.envs.utils import _call_make_env, _download_hub_file, _import_hub_module, _normalize_hub_result
from .configs import EnvConfig, HubEnvConfig
from .utils import _call_make_env, _download_hub_file, _import_hub_module, _normalize_hub_result
def make_env_config(env_type: str, **kwargs) -> EnvConfig:
src/lerobot/envs/libero.py
+2 -1
View File
@@ -29,9 +29,10 @@ from gymnasium import spaces
from libero.libero import benchmark, get_libero_path
from libero.libero.envs import OffScreenRenderEnv
from lerobot.envs.utils import _LazyAsyncVectorEnv
from lerobot.types import RobotObservation
from .utils import _LazyAsyncVectorEnv
def _parse_camera_names(camera_name: str | Sequence[str]) -> list[str]:
"""Normalize camera_name into a non-empty list of strings."""
src/lerobot/envs/metaworld.py
+2 -1
View File
@@ -25,9 +25,10 @@ import metaworld.policies as policies
import numpy as np
from gymnasium import spaces
from lerobot.envs.utils import _LazyAsyncVectorEnv
from lerobot.types import RobotObservation
from .utils import _LazyAsyncVectorEnv
# ---- Load configuration data from the external JSON file ----
CONFIG_PATH = Path(__file__).parent / "metaworld_config.json"
try:
src/lerobot/envs/utils.py
+3 -2
View File
@@ -27,11 +27,12 @@ import torch
from huggingface_hub import hf_hub_download, snapshot_download
from torch import Tensor
from lerobot.configs.types import FeatureType, PolicyFeature
from lerobot.envs.configs import EnvConfig
from lerobot.configs import FeatureType, PolicyFeature
from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE, OBS_STR
from lerobot.utils.utils import get_channel_first_image_shape
from .configs import EnvConfig
def _convert_nested_dict(d):
result = {}
src/lerobot/model/__init__.py
+5
View File
@@ -0,0 +1,5 @@
# Kinematics utilities for robot modeling.
from .kinematics import RobotKinematics as RobotKinematics
__all__ = ["RobotKinematics"]
src/lerobot/motors/__init__.py
+2
View File
@@ -19,3 +19,5 @@ from .motors_bus import (
MotorCalibration,
MotorNormMode,
)
__all__ = ["Motor", "MotorCalibration", "MotorNormMode"]
src/lerobot/motors/dynamixel/__init__.py
+4
View File
@@ -14,5 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from lerobot.utils.import_utils import require_package
require_package("dynamixel-sdk", extra="dynamixel", import_name="dynamixel_sdk")
from .dynamixel import DriveMode, DynamixelMotorsBus, OperatingMode, TorqueMode
from .tables import *
src/lerobot/motors/feetech/__init__.py
+4
View File
@@ -14,5 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from lerobot.utils.import_utils import require_package
require_package("feetech-servo-sdk", extra="feetech", import_name="scservo_sdk")
from .feetech import DriveMode, FeetechMotorsBus, OperatingMode, TorqueMode
from .tables import *
src/lerobot/optim/__init__.py
+42 -1
View File
@@ -12,4 +12,45 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from .optimizers import OptimizerConfig as OptimizerConfig
from .optimizers import (
AdamConfig as AdamConfig,
AdamWConfig as AdamWConfig,
MultiAdamConfig as MultiAdamConfig,
OptimizerConfig as OptimizerConfig,
SGDConfig as SGDConfig,
XVLAAdamWConfig as XVLAAdamWConfig,
load_optimizer_state,
save_optimizer_state,
)
from .schedulers import (
CosineDecayWithWarmupSchedulerConfig as CosineDecayWithWarmupSchedulerConfig,
DiffuserSchedulerConfig as DiffuserSchedulerConfig,
LRSchedulerConfig as LRSchedulerConfig,
VQBeTSchedulerConfig as VQBeTSchedulerConfig,
load_scheduler_state,
save_scheduler_state,
)
# NOTE: make_optimizer_and_scheduler is intentionally NOT re-exported here
# to avoid circular dependencies (it imports lerobot.configs.train and lerobot.policies).
# Import directly: ``from lerobot.optim.factory import make_optimizer_and_scheduler``
__all__ = [
# Optimizer configs
"AdamConfig",
"AdamWConfig",
"MultiAdamConfig",
"OptimizerConfig",
"SGDConfig",
"XVLAAdamWConfig",
# Scheduler configs
"CosineDecayWithWarmupSchedulerConfig",
"DiffuserSchedulerConfig",
"LRSchedulerConfig",
"VQBeTSchedulerConfig",
# State management
"load_optimizer_state",
"load_scheduler_state",
"save_optimizer_state",
"save_scheduler_state",
]
src/lerobot/optim/factory.py
+1 -1
View File
@@ -19,7 +19,7 @@ from torch.optim import Optimizer
from torch.optim.lr_scheduler import LRScheduler
from lerobot.configs.train import TrainPipelineConfig
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies import PreTrainedPolicy
def make_optimizer_and_scheduler(
src/lerobot/optim/schedulers.py
+3
View File
@@ -47,6 +47,9 @@ class DiffuserSchedulerConfig(LRSchedulerConfig):
num_warmup_steps: int | None = None
def build(self, optimizer: Optimizer, num_training_steps: int) -> LambdaLR:
from lerobot.utils.import_utils import require_package
require_package("diffusers", extra="diffusion")
from diffusers.optimization import get_scheduler
kwargs = {**asdict(self), "num_training_steps": num_training_steps, "optimizer": optimizer}
src/lerobot/policies/__init__.py
+29 -5
View File
@@ -14,30 +14,54 @@
from .act.configuration_act import ACTConfig as ACTConfig
from .diffusion.configuration_diffusion import DiffusionConfig as DiffusionConfig
from .factory import get_policy_class, make_policy, make_policy_config, make_pre_post_processors
from .groot.configuration_groot import GrootConfig as GrootConfig
from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig as MultiTaskDiTConfig
from .pi0.configuration_pi0 import PI0Config as PI0Config
from .pi0_fast.configuration_pi0_fast import PI0FastConfig as PI0FastConfig
from .pi05.configuration_pi05 import PI05Config as PI05Config
from .pretrained import PreTrainedPolicy as PreTrainedPolicy
from .rtc import ActionInterpolator as ActionInterpolator
from .sarm.configuration_sarm import SARMConfig as SARMConfig
from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
from .smolvla.processor_smolvla import SmolVLANewLineProcessor
from .smolvla.processor_smolvla import (
SmolVLANewLineProcessor as SmolVLANewLineProcessor, # noqa: F401 - registers with ProcessorStepRegistry
)
from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
from .utils import make_robot_action, prepare_observation_for_inference
from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig
from .wall_x.configuration_wall_x import WallXConfig as WallXConfig
from .xvla.configuration_xvla import XVLAConfig as XVLAConfig
# NOTE: Policy modeling classes (e.g., SACPolicy) are intentionally NOT re-exported here.
# They have heavy optional dependencies and are loaded lazily via get_policy_class().
# Import directly: ``from lerobot.policies.sac.modeling_sac import SACPolicy``
__all__ = [
# Configuration classes
"ACTConfig",
"DiffusionConfig",
"GrootConfig",
"MultiTaskDiTConfig",
"PI0Config",
"PI05Config",
"PI0FastConfig",
"SmolVLAConfig",
"PI05Config",
"SARMConfig",
"SmolVLAConfig",
"TDMPCConfig",
"VQBeTConfig",
"GrootConfig",
"XVLAConfig",
"WallXConfig",
"XVLAConfig",
# Base class
"PreTrainedPolicy",
# RTC utilities
"ActionInterpolator",
# Utility functions
"make_robot_action",
"prepare_observation_for_inference",
# Factory functions
"get_policy_class",
"make_policy",
"make_policy_config",
"make_pre_post_processors",
]
src/lerobot/policies/act/configuration_act.py
+2 -3
View File
@@ -15,9 +15,8 @@
# limitations under the License.
from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import NormalizationMode
from lerobot.optim.optimizers import AdamWConfig
from lerobot.configs import NormalizationMode, PreTrainedConfig
from lerobot.optim import AdamWConfig
@PreTrainedConfig.register_subclass("act")
src/lerobot/policies/act/modeling_act.py
+3 -2
View File
@@ -33,10 +33,11 @@ from torch import Tensor, nn
from torchvision.models._utils import IntermediateLayerGetter
from torchvision.ops.misc import FrozenBatchNorm2d
from lerobot.policies.act.configuration_act import ACTConfig
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
from ..pretrained import PreTrainedPolicy
from .configuration_act import ACTConfig
class ACTPolicy(PreTrainedPolicy):
"""
src/lerobot/policies/act/processor_act.py
+4 -2
View File
@@ -17,7 +17,6 @@ from typing import Any
import torch
from lerobot.policies.act.configuration_act import ACTConfig
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
@@ -26,10 +25,13 @@ from lerobot.processor import (
PolicyProcessorPipeline,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_act import ACTConfig
def make_act_pre_post_processors(
config: ACTConfig,
src/lerobot/policies/diffusion/configuration_diffusion.py
+2 -4
View File
@@ -16,10 +16,8 @@
# limitations under the License.
from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import NormalizationMode
from lerobot.optim.optimizers import AdamConfig
from lerobot.optim.schedulers import DiffuserSchedulerConfig
from lerobot.configs import NormalizationMode, PreTrainedConfig
from lerobot.optim import AdamConfig, DiffuserSchedulerConfig
@PreTrainedConfig.register_subclass("diffusion")
src/lerobot/policies/diffusion/modeling_diffusion.py
+6 -5
View File
@@ -31,15 +31,16 @@ import torch.nn.functional as F # noqa: N812
import torchvision
from torch import Tensor, nn
from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies.utils import (
from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
from ..pretrained import PreTrainedPolicy
from ..utils import (
get_device_from_parameters,
get_dtype_from_parameters,
get_output_shape,
populate_queues,
)
from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
from .configuration_diffusion import DiffusionConfig
class DiffusionPolicy(PreTrainedPolicy):
@@ -156,7 +157,7 @@ def _make_noise_scheduler(name: str, **kwargs: dict):
"""
from lerobot.utils.import_utils import require_package
require_package("diffusers", extra="training")
require_package("diffusers", extra="diffusion")
from diffusers.schedulers.scheduling_ddim import DDIMScheduler
from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
src/lerobot/policies/diffusion/processor_diffusion.py
+4 -2
View File
@@ -18,7 +18,6 @@ from typing import Any
import torch
from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
@@ -27,10 +26,13 @@ from lerobot.processor import (
PolicyProcessorPipeline,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_diffusion import DiffusionConfig
def make_diffusion_pre_post_processors(
config: DiffusionConfig,
src/lerobot/policies/factory.py
+54 -59
View File
@@ -23,30 +23,14 @@ from typing import TYPE_CHECKING, Any, TypedDict, Unpack
import torch
if TYPE_CHECKING:
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.datasets import LeRobotDatasetMetadata
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType
from lerobot.envs.configs import EnvConfig
from lerobot.envs.utils import env_to_policy_features
from lerobot.policies.act.configuration_act import ACTConfig
from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
from lerobot.policies.groot.configuration_groot import GrootConfig
from lerobot.policies.multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
from lerobot.policies.pi0.configuration_pi0 import PI0Config
from lerobot.policies.pi05.configuration_pi05 import PI05Config
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies.sac.configuration_sac import SACConfig
from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
from lerobot.policies.sarm.configuration_sarm import SARMConfig
from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
from lerobot.policies.utils import validate_visual_features_consistency
from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
from lerobot.policies.wall_x.configuration_wall_x import WallXConfig
from lerobot.policies.xvla.configuration_xvla import XVLAConfig
from lerobot.processor import PolicyProcessorPipeline
from lerobot.processor.converters import (
from lerobot.configs import FeatureType, PreTrainedConfig
from lerobot.envs import EnvConfig, env_to_policy_features
from lerobot.processor import (
AbsoluteActionsProcessorStep,
PolicyProcessorPipeline,
RelativeActionsProcessorStep,
batch_to_transition,
policy_action_to_transition,
transition_to_batch,
@@ -58,6 +42,24 @@ from lerobot.utils.constants import (
POLICY_POSTPROCESSOR_DEFAULT_NAME,
POLICY_PREPROCESSOR_DEFAULT_NAME,
)
from lerobot.utils.feature_utils import dataset_to_policy_features
from .act.configuration_act import ACTConfig
from .diffusion.configuration_diffusion import DiffusionConfig
from .groot.configuration_groot import GrootConfig
from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
from .pi0.configuration_pi0 import PI0Config
from .pi05.configuration_pi05 import PI05Config
from .pretrained import PreTrainedPolicy
from .sac.configuration_sac import SACConfig
from .sac.reward_model.configuration_classifier import RewardClassifierConfig
from .sarm.configuration_sarm import SARMConfig
from .smolvla.configuration_smolvla import SmolVLAConfig
from .tdmpc.configuration_tdmpc import TDMPCConfig
from .utils import validate_visual_features_consistency
from .vqbet.configuration_vqbet import VQBeTConfig
from .wall_x.configuration_wall_x import WallXConfig
from .xvla.configuration_xvla import XVLAConfig
def _reconnect_relative_absolute_steps(
@@ -70,11 +72,6 @@ def _reconnect_relative_absolute_steps(
the RelativeActionsProcessorStep so it can read the cached state at inference time.
That reference is not serializable, so we re-establish it here after loading.
"""
from lerobot.processor.relative_action_processor import (
AbsoluteActionsProcessorStep,
RelativeActionsProcessorStep,
)
relative_step = next((s for s in preprocessor.steps if isinstance(s, RelativeActionsProcessorStep)), None)
if relative_step is None:
return
@@ -100,63 +97,63 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
NotImplementedError: If the policy name is not recognized.
"""
if name == "tdmpc":
from lerobot.policies.tdmpc.modeling_tdmpc import TDMPCPolicy
from .tdmpc.modeling_tdmpc import TDMPCPolicy
return TDMPCPolicy
elif name == "diffusion":
from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
from .diffusion.modeling_diffusion import DiffusionPolicy
return DiffusionPolicy
elif name == "act":
from lerobot.policies.act.modeling_act import ACTPolicy
from .act.modeling_act import ACTPolicy
return ACTPolicy
elif name == "multi_task_dit":
from lerobot.policies.multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy
from .multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy
return MultiTaskDiTPolicy
elif name == "vqbet":
from lerobot.policies.vqbet.modeling_vqbet import VQBeTPolicy
from .vqbet.modeling_vqbet import VQBeTPolicy
return VQBeTPolicy
elif name == "pi0":
from lerobot.policies.pi0.modeling_pi0 import PI0Policy
from .pi0.modeling_pi0 import PI0Policy
return PI0Policy
elif name == "pi0_fast":
from lerobot.policies.pi0_fast.modeling_pi0_fast import PI0FastPolicy
from .pi0_fast.modeling_pi0_fast import PI0FastPolicy
return PI0FastPolicy
elif name == "pi05":
from lerobot.policies.pi05.modeling_pi05 import PI05Policy
from .pi05.modeling_pi05 import PI05Policy
return PI05Policy
elif name == "sac":
from lerobot.policies.sac.modeling_sac import SACPolicy
from .sac.modeling_sac import SACPolicy
return SACPolicy
elif name == "reward_classifier":
from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
from .sac.reward_model.modeling_classifier import Classifier
return Classifier
elif name == "smolvla":
from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
from .smolvla.modeling_smolvla import SmolVLAPolicy
return SmolVLAPolicy
elif name == "sarm":
from lerobot.policies.sarm.modeling_sarm import SARMRewardModel
from .sarm.modeling_sarm import SARMRewardModel
return SARMRewardModel
elif name == "groot":
from lerobot.policies.groot.modeling_groot import GrootPolicy
from .groot.modeling_groot import GrootPolicy
return GrootPolicy
elif name == "xvla":
from lerobot.policies.xvla.modeling_xvla import XVLAPolicy
from .xvla.modeling_xvla import XVLAPolicy
return XVLAPolicy
elif name == "wall_x":
from lerobot.policies.wall_x.modeling_wall_x import WallXPolicy
from .wall_x.modeling_wall_x import WallXPolicy
return WallXPolicy
else:
@@ -316,7 +313,7 @@ def make_pre_post_processors(
# Create a new processor based on policy type
if isinstance(policy_cfg, TDMPCConfig):
from lerobot.policies.tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
from .tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
processors = make_tdmpc_pre_post_processors(
config=policy_cfg,
@@ -324,7 +321,7 @@ def make_pre_post_processors(
)
elif isinstance(policy_cfg, DiffusionConfig):
from lerobot.policies.diffusion.processor_diffusion import make_diffusion_pre_post_processors
from .diffusion.processor_diffusion import make_diffusion_pre_post_processors
processors = make_diffusion_pre_post_processors(
config=policy_cfg,
@@ -332,7 +329,7 @@ def make_pre_post_processors(
)
elif isinstance(policy_cfg, ACTConfig):
from lerobot.policies.act.processor_act import make_act_pre_post_processors
from .act.processor_act import make_act_pre_post_processors
processors = make_act_pre_post_processors(
config=policy_cfg,
@@ -340,7 +337,7 @@ def make_pre_post_processors(
)
elif isinstance(policy_cfg, MultiTaskDiTConfig):
from lerobot.policies.multi_task_dit.processor_multi_task_dit import (
from .multi_task_dit.processor_multi_task_dit import (
make_multi_task_dit_pre_post_processors,
)
@@ -350,7 +347,7 @@ def make_pre_post_processors(
)
elif isinstance(policy_cfg, VQBeTConfig):
from lerobot.policies.vqbet.processor_vqbet import make_vqbet_pre_post_processors
from .vqbet.processor_vqbet import make_vqbet_pre_post_processors
processors = make_vqbet_pre_post_processors(
config=policy_cfg,
@@ -358,7 +355,7 @@ def make_pre_post_processors(
)
elif isinstance(policy_cfg, PI0Config):
from lerobot.policies.pi0.processor_pi0 import make_pi0_pre_post_processors
from .pi0.processor_pi0 import make_pi0_pre_post_processors
processors = make_pi0_pre_post_processors(
config=policy_cfg,
@@ -366,7 +363,7 @@ def make_pre_post_processors(
)
elif isinstance(policy_cfg, PI05Config):
from lerobot.policies.pi05.processor_pi05 import make_pi05_pre_post_processors
from .pi05.processor_pi05 import make_pi05_pre_post_processors
processors = make_pi05_pre_post_processors(
config=policy_cfg,
@@ -374,7 +371,7 @@ def make_pre_post_processors(
)
elif isinstance(policy_cfg, SACConfig):
from lerobot.policies.sac.processor_sac import make_sac_pre_post_processors
from .sac.processor_sac import make_sac_pre_post_processors
processors = make_sac_pre_post_processors(
config=policy_cfg,
@@ -382,7 +379,7 @@ def make_pre_post_processors(
)
elif isinstance(policy_cfg, RewardClassifierConfig):
from lerobot.policies.sac.reward_model.processor_classifier import make_classifier_processor
from .sac.reward_model.processor_classifier import make_classifier_processor
processors = make_classifier_processor(
config=policy_cfg,
@@ -390,7 +387,7 @@ def make_pre_post_processors(
)
elif isinstance(policy_cfg, SmolVLAConfig):
from lerobot.policies.smolvla.processor_smolvla import make_smolvla_pre_post_processors
from .smolvla.processor_smolvla import make_smolvla_pre_post_processors
processors = make_smolvla_pre_post_processors(
config=policy_cfg,
@@ -398,7 +395,7 @@ def make_pre_post_processors(
)
elif isinstance(policy_cfg, SARMConfig):
from lerobot.policies.sarm.processor_sarm import make_sarm_pre_post_processors
from .sarm.processor_sarm import make_sarm_pre_post_processors
processors = make_sarm_pre_post_processors(
config=policy_cfg,
@@ -406,7 +403,7 @@ def make_pre_post_processors(
dataset_meta=kwargs.get("dataset_meta"),
)
elif isinstance(policy_cfg, GrootConfig):
from lerobot.policies.groot.processor_groot import make_groot_pre_post_processors
from .groot.processor_groot import make_groot_pre_post_processors
processors = make_groot_pre_post_processors(
config=policy_cfg,
@@ -414,7 +411,7 @@ def make_pre_post_processors(
)
elif isinstance(policy_cfg, XVLAConfig):
from lerobot.policies.xvla.processor_xvla import (
from .xvla.processor_xvla import (
make_xvla_pre_post_processors,
)
@@ -424,7 +421,7 @@ def make_pre_post_processors(
)
elif isinstance(policy_cfg, WallXConfig):
from lerobot.policies.wall_x.processor_wall_x import make_wall_x_pre_post_processors
from .wall_x.processor_wall_x import make_wall_x_pre_post_processors
processors = make_wall_x_pre_post_processors(
config=policy_cfg,
@@ -495,8 +492,6 @@ def make_policy(
kwargs = {}
if ds_meta is not None:
from lerobot.utils.feature_utils import dataset_to_policy_features
features = dataset_to_policy_features(ds_meta.features)
else:
if not cfg.pretrained_path:
src/lerobot/policies/groot/__init__.py
+4
View File
@@ -14,6 +14,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from lerobot.utils.import_utils import require_package
require_package("transformers", extra="groot")
from .configuration_groot import GrootConfig
from .modeling_groot import GrootPolicy
from .processor_groot import make_groot_pre_post_processors
src/lerobot/policies/groot/action_head/cross_attention_dit.py
+1 -1
View File
@@ -19,7 +19,7 @@ import torch.nn.functional as F # noqa: N812
from lerobot.utils.import_utils import require_package
require_package("diffusers", extra="training")
require_package("diffusers", extra="groot")
from diffusers import ConfigMixin, ModelMixin # noqa: E402
from diffusers.configuration_utils import register_to_config # noqa: E402
src/lerobot/policies/groot/action_head/flow_matching_action_head.py
+1 -2
View File
@@ -31,11 +31,10 @@ else:
PretrainedConfig = object
BatchFeature = None
from lerobot.policies.groot.action_head.action_encoder import (
from .action_encoder import (
SinusoidalPositionalEncoding,
swish,
)
from .cross_attention_dit import DiT, SelfAttentionTransformer
src/lerobot/policies/groot/configuration_groot.py
+2 -4
View File
@@ -16,10 +16,8 @@
from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
from lerobot.utils.constants import ACTION, OBS_STATE
src/lerobot/policies/groot/groot_n1.py
+4 -3
View File
@@ -41,12 +41,13 @@ try:
except ImportError:
tree = None
from lerobot.policies.groot.action_head.flow_matching_action_head import (
from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME
from .action_head.flow_matching_action_head import (
FlowmatchingActionHead,
FlowmatchingActionHeadConfig,
)
from lerobot.policies.groot.utils import ensure_eagle_cache_ready
from lerobot.utils.constants import ACTION, HF_LEROBOT_HOME
from .utils import ensure_eagle_cache_ready
DEFAULT_VENDOR_EAGLE_PATH = str((Path(__file__).resolve().parent / "eagle2_hg_model").resolve())
DEFAULT_TOKENIZER_ASSETS_REPO = "lerobot/eagle2hg-processor-groot-n1p5"
src/lerobot/policies/groot/modeling_groot.py
+5 -4
View File
@@ -41,12 +41,13 @@ from typing import TypeVar
import torch
from torch import Tensor
from lerobot.configs.types import FeatureType, PolicyFeature
from lerobot.policies.groot.configuration_groot import GrootConfig
from lerobot.policies.groot.groot_n1 import GR00TN15
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.configs import FeatureType, PolicyFeature
from lerobot.utils.constants import ACTION, OBS_IMAGES
from ..pretrained import PreTrainedPolicy
from .configuration_groot import GrootConfig
from .groot_n1 import GR00TN15
T = TypeVar("T", bound="GrootPolicy")
src/lerobot/policies/groot/processor_groot.py
+3 -4
View File
@@ -30,12 +30,11 @@ else:
AutoProcessor = None
ProcessorMixin = object
from lerobot.configs.types import (
from lerobot.configs import (
FeatureType,
NormalizationMode,
PolicyFeature,
)
from lerobot.policies.groot.configuration_groot import GrootConfig
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
@@ -44,8 +43,6 @@ from lerobot.processor import (
ProcessorStep,
ProcessorStepRegistry,
RenameObservationsProcessorStep,
)
from lerobot.processor.converters import (
policy_action_to_transition,
transition_to_policy_action,
)
@@ -60,6 +57,8 @@ from lerobot.utils.constants import (
POLICY_PREPROCESSOR_DEFAULT_NAME,
)
from .configuration_groot import GrootConfig
# Defaults for Eagle processor locations
DEFAULT_TOKENIZER_ASSETS_REPO = "lerobot/eagle2hg-processor-groot-n1p5"
src/lerobot/policies/multi_task_dit/configuration_multi_task_dit.py
+2 -4
View File
@@ -17,10 +17,8 @@
import logging
from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import NormalizationMode
from lerobot.optim.optimizers import AdamConfig
from lerobot.optim.schedulers import DiffuserSchedulerConfig
from lerobot.configs import NormalizationMode, PreTrainedConfig
from lerobot.optim import AdamConfig, DiffuserSchedulerConfig
@PreTrainedConfig.register_subclass("multi_task_dit")
src/lerobot/policies/multi_task_dit/modeling_multi_task_dit.py
+6 -4
View File
@@ -36,17 +36,16 @@ import torch.nn.functional as F # noqa: N812
import torchvision
from torch import Tensor
from lerobot.policies.multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
from lerobot.utils.import_utils import _transformers_available
from .configuration_multi_task_dit import MultiTaskDiTConfig
# Conditional import for type checking and lazy loading
if TYPE_CHECKING or _transformers_available:
from transformers import CLIPTextModel, CLIPVisionModel
else:
CLIPTextModel = None
CLIPVisionModel = None
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies.utils import populate_queues
from lerobot.utils.constants import (
ACTION,
OBS_IMAGES,
@@ -55,6 +54,9 @@ from lerobot.utils.constants import (
OBS_STATE,
)
from ..pretrained import PreTrainedPolicy
from ..utils import populate_queues
# -- Policy --
@@ -643,7 +645,7 @@ class DiffusionObjective(nn.Module):
from lerobot.utils.import_utils import require_package
require_package("diffusers", extra="training")
require_package("diffusers", extra="multi_task_dit")
from diffusers.schedulers.scheduling_ddim import DDIMScheduler
from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
src/lerobot/policies/multi_task_dit/processor_multi_task_dit.py
+4 -2
View File
@@ -18,7 +18,6 @@ from typing import Any
import torch
from lerobot.policies.multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
@@ -28,10 +27,13 @@ from lerobot.processor import (
RenameObservationsProcessorStep,
TokenizerProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_multi_task_dit import MultiTaskDiTConfig
def make_multi_task_dit_pre_post_processors(
config: MultiTaskDiTConfig,
src/lerobot/policies/pi0/configuration_pi0.py
+4 -5
View File
@@ -16,13 +16,12 @@
from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
from lerobot.policies.rtc.configuration_rtc import RTCConfig
from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
from ..rtc.configuration_rtc import RTCConfig
DEFAULT_IMAGE_SIZE = 224
src/lerobot/policies/pi0/modeling_pi0.py
+6 -5
View File
@@ -33,7 +33,7 @@ if TYPE_CHECKING or _transformers_available:
from transformers.models.auto import CONFIG_MAPPING
from transformers.models.gemma import modeling_gemma
from lerobot.policies.pi_gemma import (
from ..pi_gemma import (
PaliGemmaForConditionalGenerationWithPiGemma,
PiGemmaForCausalLM,
_gated_residual,
@@ -48,10 +48,7 @@ else:
PaliGemmaForConditionalGenerationWithPiGemma = None
from lerobot.configs.policies import PreTrainedConfig
from lerobot.policies.pi0.configuration_pi0 import DEFAULT_IMAGE_SIZE, PI0Config
from lerobot.policies.pretrained import PreTrainedPolicy, T
from lerobot.policies.rtc.modeling_rtc import RTCProcessor
from lerobot.configs import PreTrainedConfig
from lerobot.utils.constants import (
ACTION,
OBS_LANGUAGE_ATTENTION_MASK,
@@ -60,6 +57,10 @@ from lerobot.utils.constants import (
OPENPI_ATTENTION_MASK_VALUE,
)
from ..pretrained import PreTrainedPolicy, T
from ..rtc.modeling_rtc import RTCProcessor
from .configuration_pi0 import DEFAULT_IMAGE_SIZE, PI0Config
class ActionSelectKwargs(TypedDict, total=False):
inference_delay: int | None
src/lerobot/policies/pi0/processor_pi0.py
+5 -3
View File
@@ -18,8 +18,7 @@ from typing import Any
import torch
from lerobot.configs.types import PipelineFeatureType, PolicyFeature
from lerobot.policies.pi0.configuration_pi0 import PI0Config
from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor import (
AbsoluteActionsProcessorStep,
AddBatchDimensionProcessorStep,
@@ -34,10 +33,13 @@ from lerobot.processor import (
RenameObservationsProcessorStep,
TokenizerProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_pi0 import PI0Config
@ProcessorStepRegistry.register(name="pi0_new_line_processor")
class Pi0NewLineProcessor(ComplementaryDataProcessorStep):
src/lerobot/policies/pi05/configuration_pi05.py
+4 -5
View File
@@ -16,13 +16,12 @@
from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
from lerobot.policies.rtc.configuration_rtc import RTCConfig
from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
from ..rtc.configuration_rtc import RTCConfig
DEFAULT_IMAGE_SIZE = 224
src/lerobot/policies/pi05/modeling_pi05.py
+6 -5
View File
@@ -33,7 +33,7 @@ if TYPE_CHECKING or _transformers_available:
from transformers.models.auto import CONFIG_MAPPING
from transformers.models.gemma import modeling_gemma
from lerobot.policies.pi_gemma import (
from ..pi_gemma import (
PaliGemmaForConditionalGenerationWithPiGemma,
PiGemmaForCausalLM,
_gated_residual,
@@ -46,10 +46,7 @@ else:
_gated_residual = None
layernorm_forward = None
PaliGemmaForConditionalGenerationWithPiGemma = None
from lerobot.configs.policies import PreTrainedConfig
from lerobot.policies.pi05.configuration_pi05 import DEFAULT_IMAGE_SIZE, PI05Config
from lerobot.policies.pretrained import PreTrainedPolicy, T
from lerobot.policies.rtc.modeling_rtc import RTCProcessor
from lerobot.configs import PreTrainedConfig
from lerobot.utils.constants import (
ACTION,
OBS_LANGUAGE_ATTENTION_MASK,
@@ -57,6 +54,10 @@ from lerobot.utils.constants import (
OPENPI_ATTENTION_MASK_VALUE,
)
from ..pretrained import PreTrainedPolicy, T
from ..rtc.modeling_rtc import RTCProcessor
from .configuration_pi05 import DEFAULT_IMAGE_SIZE, PI05Config
class ActionSelectKwargs(TypedDict, total=False):
inference_delay: int | None
src/lerobot/policies/pi05/processor_pi05.py
+5 -3
View File
@@ -21,8 +21,7 @@ from typing import Any
import numpy as np
import torch
from lerobot.configs.types import PipelineFeatureType, PolicyFeature
from lerobot.policies.pi05.configuration_pi05 import PI05Config
from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor import (
AbsoluteActionsProcessorStep,
AddBatchDimensionProcessorStep,
@@ -36,8 +35,9 @@ from lerobot.processor import (
RenameObservationsProcessorStep,
TokenizerProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.types import EnvTransition, TransitionKey
from lerobot.utils.constants import (
OBS_STATE,
@@ -45,6 +45,8 @@ from lerobot.utils.constants import (
POLICY_PREPROCESSOR_DEFAULT_NAME,
)
from .configuration_pi05 import PI05Config
@ProcessorStepRegistry.register(name="pi05_prepare_state_tokenizer_processor_step")
@dataclass
src/lerobot/policies/pi0_fast/configuration_pi0_fast.py
+4 -5
View File
@@ -16,13 +16,12 @@
from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
from lerobot.policies.rtc.configuration_rtc import RTCConfig
from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
from ..rtc.configuration_rtc import RTCConfig
DEFAULT_IMAGE_SIZE = 224
src/lerobot/policies/pi0_fast/modeling_pi0_fast.py
+6 -5
View File
@@ -38,7 +38,7 @@ if TYPE_CHECKING or _transformers_available:
from transformers import AutoTokenizer
from transformers.models.auto import CONFIG_MAPPING
from lerobot.policies.pi_gemma import (
from ..pi_gemma import (
PaliGemmaForConditionalGenerationWithPiGemma,
PiGemmaModel,
)
@@ -48,10 +48,7 @@ else:
PiGemmaModel = None
PaliGemmaForConditionalGenerationWithPiGemma = None
from lerobot.configs.policies import PreTrainedConfig
from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig
from lerobot.policies.pretrained import PreTrainedPolicy, T
from lerobot.policies.rtc.modeling_rtc import RTCProcessor
from lerobot.configs import PreTrainedConfig
from lerobot.utils.constants import (
ACTION,
ACTION_TOKEN_MASK,
@@ -61,6 +58,10 @@ from lerobot.utils.constants import (
OPENPI_ATTENTION_MASK_VALUE,
)
from ..pretrained import PreTrainedPolicy, T
from ..rtc.modeling_rtc import RTCProcessor
from .configuration_pi0_fast import PI0FastConfig
class ActionSelectKwargs(TypedDict, total=False):
temperature: float | None
src/lerobot/policies/pi0_fast/processor_pi0_fast.py
+5 -3
View File
@@ -21,8 +21,7 @@ from typing import Any
import numpy as np
import torch
from lerobot.configs.types import PipelineFeatureType, PolicyFeature
from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig
from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor import (
AbsoluteActionsProcessorStep,
ActionTokenizerProcessorStep,
@@ -37,8 +36,9 @@ from lerobot.processor import (
RenameObservationsProcessorStep,
TokenizerProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.types import EnvTransition, TransitionKey
from lerobot.utils.constants import (
OBS_STATE,
@@ -46,6 +46,8 @@ from lerobot.utils.constants import (
POLICY_PREPROCESSOR_DEFAULT_NAME,
)
from .configuration_pi0_fast import PI0FastConfig
@ProcessorStepRegistry.register(name="pi0_fast_prepare_state_tokenizer_processor_step")
@dataclass
src/lerobot/policies/pretrained.py
+3 -2
View File
@@ -29,11 +29,12 @@ from huggingface_hub.errors import HfHubHTTPError
from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor
from torch import Tensor, nn
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs import PreTrainedConfig
from lerobot.configs.train import TrainPipelineConfig
from lerobot.policies.utils import log_model_loading_keys
from lerobot.utils.hub import HubMixin
from .utils import log_model_loading_keys
T = TypeVar("T", bound="PreTrainedPolicy")
src/lerobot/policies/rtc/__init__.py
+5 -5
View File
@@ -14,11 +14,11 @@
"""Real-Time Chunking (RTC) utilities for action-chunking policies."""
from lerobot.policies.rtc.action_interpolator import ActionInterpolator
from lerobot.policies.rtc.action_queue import ActionQueue
from lerobot.policies.rtc.configuration_rtc import RTCConfig
from lerobot.policies.rtc.latency_tracker import LatencyTracker
from lerobot.policies.rtc.modeling_rtc import RTCProcessor
from .action_interpolator import ActionInterpolator
from .action_queue import ActionQueue
from .configuration_rtc import RTCConfig
from .latency_tracker import LatencyTracker
from .modeling_rtc import RTCProcessor
__all__ = [
"ActionInterpolator",
src/lerobot/policies/rtc/action_queue.py
+1 -1
View File
@@ -27,7 +27,7 @@ from threading import Lock
import torch
from torch import Tensor
from lerobot.policies.rtc.configuration_rtc import RTCConfig
from .configuration_rtc import RTCConfig
logger = logging.getLogger(__name__)
src/lerobot/policies/rtc/configuration_rtc.py
+1 -1
View File
@@ -23,7 +23,7 @@ Based on:
from dataclasses import dataclass
from lerobot.configs.types import RTCAttentionSchedule
from lerobot.configs import RTCAttentionSchedule
@dataclass
src/lerobot/policies/rtc/modeling_rtc.py
+4 -3
View File
@@ -27,9 +27,10 @@ import math
import torch
from torch import Tensor
from lerobot.configs.types import RTCAttentionSchedule
from lerobot.policies.rtc.configuration_rtc import RTCConfig
from lerobot.policies.rtc.debug_tracker import Tracker
from lerobot.configs import RTCAttentionSchedule
from .configuration_rtc import RTCConfig
from .debug_tracker import Tracker
logger = logging.getLogger(__name__)
src/lerobot/policies/sac/configuration_sac.py
+2 -3
View File
@@ -17,9 +17,8 @@
from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import NormalizationMode
from lerobot.optim.optimizers import MultiAdamConfig
from lerobot.configs import NormalizationMode, PreTrainedConfig
from lerobot.optim import MultiAdamConfig
from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE
src/lerobot/policies/sac/modeling_sac.py
+4 -3
View File
@@ -28,11 +28,12 @@ import torch.nn.functional as F # noqa: N812
from torch import Tensor
from torch.distributions import MultivariateNormal, TanhTransform, Transform, TransformedDistribution
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies.sac.configuration_sac import SACConfig, is_image_feature
from lerobot.policies.utils import get_device_from_parameters
from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_STATE
from ..pretrained import PreTrainedPolicy
from ..utils import get_device_from_parameters
from .configuration_sac import SACConfig, is_image_feature
DISCRETE_DIMENSION_INDEX = -1 # Gripper is always the last dimension
src/lerobot/policies/sac/processor_sac.py
+4 -2
View File
@@ -19,7 +19,6 @@ from typing import Any
import torch
from lerobot.policies.sac.configuration_sac import SACConfig
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
@@ -28,10 +27,13 @@ from lerobot.processor import (
PolicyProcessorPipeline,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_sac import SACConfig
def make_sac_pre_post_processors(
config: SACConfig,
src/lerobot/policies/sac/reward_model/configuration_classifier.py
+2 -4
View File
@@ -15,10 +15,8 @@
# limitations under the License.
from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import NormalizationMode
from lerobot.optim.optimizers import AdamWConfig, OptimizerConfig
from lerobot.optim.schedulers import LRSchedulerConfig
from lerobot.configs import NormalizationMode, PreTrainedConfig
from lerobot.optim import AdamWConfig, LRSchedulerConfig, OptimizerConfig
from lerobot.utils.constants import OBS_IMAGE
src/lerobot/policies/sac/reward_model/modeling_classifier.py
+3 -2
View File
@@ -19,10 +19,11 @@ import logging
import torch
from torch import Tensor, nn
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
from lerobot.utils.constants import OBS_IMAGE, REWARD
from ...pretrained import PreTrainedPolicy
from .configuration_classifier import RewardClassifierConfig
class ClassifierOutput:
"""Wrapper for classifier outputs with additional metadata."""
src/lerobot/policies/sac/reward_model/processor_classifier.py
+4 -2
View File
@@ -18,15 +18,17 @@ from typing import Any
import torch
from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
from lerobot.processor import (
DeviceProcessorStep,
IdentityProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from .configuration_classifier import RewardClassifierConfig
def make_classifier_processor(
src/lerobot/policies/sarm/__init__.py
+3
View File
@@ -0,0 +1,3 @@
from .configuration_sarm import SARMConfig
__all__ = ["SARMConfig"]
src/lerobot/policies/sarm/compute_rabc_weights.py
+5 -4
View File
@@ -57,10 +57,11 @@ import pyarrow.parquet as pq
import torch
from tqdm import tqdm
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.policies.sarm.modeling_sarm import SARMRewardModel
from lerobot.policies.sarm.processor_sarm import make_sarm_pre_post_processors
from lerobot.policies.sarm.sarm_utils import normalize_stage_tau
from lerobot.datasets import LeRobotDataset
from .modeling_sarm import SARMRewardModel
from .processor_sarm import make_sarm_pre_post_processors
from .sarm_utils import normalize_stage_tau
def get_reward_model_path_from_parquet(parquet_path: Path) -> str | None:
src/lerobot/policies/sarm/configuration_sarm.py
+2 -4
View File
@@ -22,10 +22,8 @@ Paper: https://arxiv.org/abs/2509.25358
from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
from lerobot.utils.constants import OBS_IMAGES, OBS_STATE
src/lerobot/policies/sarm/modeling_sarm.py
+5 -4
View File
@@ -34,13 +34,14 @@ import torch.nn as nn
import torch.nn.functional as F # noqa: N812
from torch import Tensor
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies.sarm.configuration_sarm import SARMConfig
from lerobot.policies.sarm.sarm_utils import (
from lerobot.utils.constants import OBS_STR
from ..pretrained import PreTrainedPolicy
from .configuration_sarm import SARMConfig
from .sarm_utils import (
normalize_stage_tau,
pad_state_to_max_dim,
)
from lerobot.utils.constants import OBS_STR
class StageTransformer(nn.Module):
src/lerobot/policies/sarm/processor_sarm.py
+22 -20
View File
@@ -24,33 +24,35 @@ import pandas as pd
import torch
from faker import Faker
from PIL import Image
from lerobot.utils.import_utils import require_package
require_package("transformers", extra="sarm")
from transformers import CLIPModel, CLIPProcessor
from lerobot.configs.types import FeatureType, PolicyFeature
from lerobot.policies.sarm.configuration_sarm import SARMConfig
from lerobot.policies.sarm.sarm_utils import (
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyProcessorPipeline,
ProcessorStep,
RenameObservationsProcessorStep,
from_tensor_to_numpy,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.types import EnvTransition, PolicyAction, TransitionKey
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_sarm import SARMConfig
from .sarm_utils import (
apply_rewind_augmentation,
compute_absolute_indices,
find_stage_and_tau,
pad_state_to_max_dim,
)
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
NormalizerProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
ProcessorStep,
RenameObservationsProcessorStep,
)
from lerobot.processor.converters import (
from_tensor_to_numpy,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.processor.pipeline import PipelineFeatureType
from lerobot.types import EnvTransition, TransitionKey
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
class SARMEncodingProcessorStep(ProcessorStep):
src/lerobot/policies/smolvla/configuration_smolvla.py
+4 -7
View File
@@ -14,15 +14,12 @@
from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import (
CosineDecayWithWarmupSchedulerConfig,
)
from lerobot.policies.rtc.configuration_rtc import RTCConfig
from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
from lerobot.utils.constants import OBS_IMAGES
from ..rtc.configuration_rtc import RTCConfig
@PreTrainedConfig.register_subclass("smolvla")
@dataclass
src/lerobot/policies/smolvla/modeling_smolvla.py
+8 -7
View File
@@ -60,16 +60,17 @@ import torch
import torch.nn.functional as F # noqa: N812
from torch import Tensor, nn
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies.rtc.modeling_rtc import RTCProcessor
from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
from lerobot.policies.smolvla.smolvlm_with_expert import SmolVLMWithExpertModel
from lerobot.policies.utils import (
populate_queues,
)
from lerobot.utils.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
from lerobot.utils.device_utils import get_safe_dtype
from ..pretrained import PreTrainedPolicy
from ..rtc.modeling_rtc import RTCProcessor
from ..utils import (
populate_queues,
)
from .configuration_smolvla import SmolVLAConfig
from .smolvlm_with_expert import SmolVLMWithExpertModel
class ActionSelectKwargs(TypedDict, total=False):
inference_delay: int | None
src/lerobot/policies/smolvla/processor_smolvla.py
+5 -3
View File
@@ -18,8 +18,7 @@ from typing import Any
import torch
from lerobot.configs.types import PipelineFeatureType, PolicyFeature
from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor import (
AddBatchDimensionProcessorStep,
ComplementaryDataProcessorStep,
@@ -31,10 +30,13 @@ from lerobot.processor import (
RenameObservationsProcessorStep,
TokenizerProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_smolvla import SmolVLAConfig
def make_smolvla_pre_post_processors(
config: SmolVLAConfig,
src/lerobot/policies/smolvla/smolvlm_with_expert.py
+5
View File
@@ -16,6 +16,11 @@ import copy
import torch
from torch import nn
from lerobot.utils.import_utils import require_package
require_package("transformers", extra="smolvla")
from transformers import (
AutoConfig,
AutoModel,
src/lerobot/policies/tdmpc/configuration_tdmpc.py
+2 -3
View File
@@ -16,9 +16,8 @@
# limitations under the License.
from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import NormalizationMode
from lerobot.optim.optimizers import AdamConfig
from lerobot.configs import NormalizationMode, PreTrainedConfig
from lerobot.optim import AdamConfig
@PreTrainedConfig.register_subclass("tdmpc")
src/lerobot/policies/tdmpc/modeling_tdmpc.py
+4 -3
View File
@@ -35,11 +35,12 @@ import torch.nn as nn
import torch.nn.functional as F # noqa: N812
from torch import Tensor
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
from lerobot.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_PREFIX, OBS_STATE, OBS_STR, REWARD
from ..pretrained import PreTrainedPolicy
from ..utils import get_device_from_parameters, get_output_shape, populate_queues
from .configuration_tdmpc import TDMPCConfig
class TDMPCPolicy(PreTrainedPolicy):
"""Implementation of TD-MPC learning + inference.
src/lerobot/policies/tdmpc/processor_tdmpc.py
+4 -2
View File
@@ -18,7 +18,6 @@ from typing import Any
import torch
from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
@@ -27,10 +26,13 @@ from lerobot.processor import (
PolicyProcessorPipeline,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_tdmpc import TDMPCConfig
def make_tdmpc_pre_post_processors(
config: TDMPCConfig,
src/lerobot/policies/utils.py
+2 -4
View File
@@ -21,10 +21,10 @@ import numpy as np
import torch
from torch import nn
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType, PolicyFeature
from lerobot.configs import FeatureType, PolicyFeature, PreTrainedConfig
from lerobot.types import PolicyAction, RobotAction, RobotObservation
from lerobot.utils.constants import ACTION, OBS_STR
from lerobot.utils.feature_utils import build_dataset_frame
def populate_queues(
@@ -162,8 +162,6 @@ def build_inference_frame(
Returns:
A dictionary of preprocessed tensors ready for model inference.
"""
from lerobot.utils.feature_utils import build_dataset_frame
# Extracts the correct keys from the incoming raw observation
observation = build_dataset_frame(ds_features, observation, prefix=OBS_STR)
src/lerobot/policies/vqbet/configuration_vqbet.py
+2 -4
View File
@@ -18,10 +18,8 @@
from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import NormalizationMode
from lerobot.optim.optimizers import AdamConfig
from lerobot.optim.schedulers import VQBeTSchedulerConfig
from lerobot.configs import NormalizationMode, PreTrainedConfig
from lerobot.optim import AdamConfig, VQBeTSchedulerConfig
@PreTrainedConfig.register_subclass("vqbet")
src/lerobot/policies/vqbet/modeling_vqbet.py
+5 -4
View File
@@ -27,12 +27,13 @@ import torch.nn.functional as F # noqa: N812
import torchvision
from torch import Tensor, nn
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
from lerobot.policies.vqbet.vqbet_utils import GPT, ResidualVQ
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
from ..pretrained import PreTrainedPolicy
from ..utils import get_device_from_parameters, get_output_shape, populate_queues
from .configuration_vqbet import VQBeTConfig
from .vqbet_utils import GPT, ResidualVQ
# ruff: noqa: N806
src/lerobot/policies/vqbet/processor_vqbet.py
+4 -2
View File
@@ -19,7 +19,6 @@ from typing import Any
import torch
from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
from lerobot.processor import (
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
@@ -28,10 +27,13 @@ from lerobot.processor import (
PolicyProcessorPipeline,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_vqbet import VQBeTConfig
def make_vqbet_pre_post_processors(
config: VQBeTConfig,
src/lerobot/policies/vqbet/vqbet_utils.py
+1 -1
View File
@@ -30,7 +30,7 @@ from torch import einsum, nn
from torch.cuda.amp import autocast
from torch.optim import Optimizer
from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
from .configuration_vqbet import VQBeTConfig
# ruff: noqa: N806
src/lerobot/policies/wall_x/__init__.py
+1 -1
View File
@@ -16,4 +16,4 @@
from .configuration_wall_x import WallXConfig
__all__ = ["WallXConfig", "WallXPolicy", "make_wall_x_pre_post_processors"]
__all__ = ["WallXConfig"]
src/lerobot/policies/wall_x/configuration_wall_x.py
+2 -4
View File
@@ -14,10 +14,8 @@
from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
from lerobot.optim import AdamWConfig, CosineDecayWithWarmupSchedulerConfig
from lerobot.utils.constants import ACTION, OBS_STATE
src/lerobot/policies/wall_x/modeling_wall_x.py
+17 -8
View File
@@ -43,6 +43,14 @@ import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from lerobot.utils.import_utils import require_package
require_package("transformers", extra="wallx")
require_package("peft", extra="wallx")
require_package("torchdiffeq", extra="wallx")
require_package("qwen-vl-utils", extra="wallx", import_name="qwen_vl_utils")
from peft import LoraConfig, get_peft_model
from PIL import Image
from qwen_vl_utils.vision_process import smart_resize
@@ -59,10 +67,12 @@ from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
)
from transformers.utils import is_torchdynamo_compiling, logging
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies.utils import populate_queues
from lerobot.policies.wall_x.configuration_wall_x import WallXConfig
from lerobot.policies.wall_x.constant import (
from lerobot.utils.constants import ACTION, OBS_STATE
from ..pretrained import PreTrainedPolicy
from ..utils import populate_queues
from .configuration_wall_x import WallXConfig
from .constant import (
GENERATE_SUBTASK_RATIO,
IMAGE_FACTOR,
MAX_PIXELS,
@@ -72,19 +82,18 @@ from lerobot.policies.wall_x.constant import (
RESOLUTION,
TOKENIZER_MAX_LENGTH,
)
from lerobot.policies.wall_x.qwen_model.configuration_qwen2_5_vl import Qwen2_5_VLConfig
from lerobot.policies.wall_x.qwen_model.qwen2_5_vl_moe import (
from .qwen_model.configuration_qwen2_5_vl import Qwen2_5_VLConfig
from .qwen_model.qwen2_5_vl_moe import (
Qwen2_5_VisionTransformerPretrainedModel,
Qwen2_5_VLACausalLMOutputWithPast,
Qwen2_5_VLMoEModel,
)
from lerobot.policies.wall_x.utils import (
from .utils import (
get_wallx_normal_text,
preprocesser_call,
process_grounding_points,
replace_action_token,
)
from lerobot.utils.constants import ACTION, OBS_STATE
logger = logging.get_logger(__name__)
src/lerobot/policies/wall_x/processor_wall_x.py
+5 -3
View File
@@ -18,8 +18,7 @@ from typing import Any
import torch
from lerobot.configs.types import PipelineFeatureType, PolicyFeature
from lerobot.policies.wall_x.configuration_wall_x import WallXConfig
from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor import (
AddBatchDimensionProcessorStep,
ComplementaryDataProcessorStep,
@@ -30,10 +29,13 @@ from lerobot.processor import (
ProcessorStepRegistry,
RenameObservationsProcessorStep,
UnnormalizerProcessorStep,
policy_action_to_transition,
transition_to_policy_action,
)
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
from .configuration_wall_x import WallXConfig
def make_wall_x_pre_post_processors(
config: WallXConfig,
src/lerobot/policies/wall_x/utils.py
+3 -2
View File
@@ -30,10 +30,11 @@ from typing import Any
import torch
from transformers import BatchFeature
from lerobot.policies.wall_x.constant import (
from lerobot.utils.constants import OBS_IMAGES
from .constant import (
CAMERA_NAME_MAPPING,
)
from lerobot.utils.constants import OBS_IMAGES
@dataclass

Some files were not shown because too many files have changed in this diff Show More