Compare commits

..

3 Commits

Author SHA1 Message Date
Steven Palma 45243dcf7c chore(dataset): add check dataset shape 2026-07-27 18:36:15 +02:00
Steven Palma 034693f724 Merge branch 'main' into fix/save-episode-zero-width-feature 2026-07-27 15:52:12 +02:00
pranjalthebhatia bb3ef3537f fix(datasets): support features with a zero-width dimension (shape=(0,))
Declaring a numeric feature with `shape=(0,)` crashed `save_episode()` in two
distinct places, leaving `dtype: "string"` as the only (type-lossy) workaround:

  - `compute_episode_stats` -> `RunningQuantileStats.update` reshaped a size-0
    array, raising "ValueError: cannot reshape array of size 0 into shape (0)".
  - `get_hf_features_from_features` mapped it to a fixed-size Arrow list of
    length 0, which pyarrow rejects ("list_size needs to be a strict positive
    integer").

The issue only reported the first error; the second surfaces once the first is
fixed. This change handles both:

  - Skip zero-width features during episode stats, exactly as string/language
    features are already skipped.
  - Store 1-D zero-width features as a variable-length sequence (length=-1) so
    each per-frame value is simply an empty list.

Adds a unit test (stats layer) and an integration test that records, saves, and
reads back a zero-width feature, asserting it round-trips as an empty vector and
is excluded from stats.

Fixes #3654

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-04 17:42:55 -04:00
26 changed files with 127 additions and 489 deletions
-18
View File
@@ -14,7 +14,6 @@
import builtins
import datetime as dt
import json
import multiprocessing
import os
import tempfile
from dataclasses import dataclass, field
@@ -102,12 +101,6 @@ class TrainPipelineConfig(HubMixin):
batch_size: int = 8
prefetch_factor: int = 4
persistent_workers: bool = True
# DataLoader worker start method. "spawn" is safer than "fork" with
# non-fork-safe libs (PyAV / torchcodec / ffmpeg), but adds some
# worker-startup time per run since workers re-import modules instead
# of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
# when appropriate, or set it to `null` to use Python's platform default.
dataloader_multiprocessing_context: str | None = "spawn"
steps: int = 100_000
# Run policy in the simulation environment every N steps to measure reward/success (0 = disabled).
env_eval_freq: int = 20_000
@@ -219,17 +212,6 @@ class TrainPipelineConfig(HubMixin):
self.reward_model.pretrained_path = str(policy_dir)
def validate(self) -> None:
available_contexts = multiprocessing.get_all_start_methods()
if (
self.dataloader_multiprocessing_context is not None
and self.dataloader_multiprocessing_context not in available_contexts
):
raise ValueError(
"`dataloader_multiprocessing_context` must be None or one of "
f"{available_contexts} on this platform, got "
f"{self.dataloader_multiprocessing_context!r}."
)
self._resolve_pretrained_from_cli()
if self.policy is None and self.reward_model is None:
+7
View File
@@ -519,6 +519,13 @@ def compute_episode_stats(
if features[key]["dtype"] in {"string", "language"}:
continue
# Features with a zero-width dimension contain no statistics-bearing
# values. Skip them like strings instead of letting
# get_feature_stats -> RunningQuantileStats.update reshape a size-0 array,
# which raises "ValueError: cannot reshape array of size 0".
if any(dim == 0 for dim in features[key].get("shape", ())):
continue
if features[key]["dtype"] in ["image", "video"]:
ep_ft_array = sample_images(data)
axes_to_reduce = (0, 2, 3)
+11 -3
View File
@@ -64,12 +64,20 @@ def get_hf_features_from_features(features: dict) -> datasets.Features:
continue
elif ft["dtype"] == "image":
hf_features[key] = datasets.Image()
elif len(ft["shape"]) > 1 and any(dim == 0 for dim in ft["shape"]):
raise ValueError(
f"Multidimensional features with a zero-width dimension are not supported: "
f"'{key}' has shape {ft['shape']}. Only the one-dimensional shape (0,) is supported."
)
elif ft["shape"] == (1,):
hf_features[key] = datasets.Value(dtype=ft["dtype"])
elif len(ft["shape"]) == 1:
hf_features[key] = datasets.Sequence(
length=ft["shape"][0], feature=datasets.Value(dtype=ft["dtype"])
)
# A zero-width feature (shape=(0,)) has no fixed-size Arrow representation:
# pyarrow rejects a fixed-size list of length 0 ("list_size needs to be a
# strict positive integer"). Store it as a variable-length sequence
# (length=-1) so each per-frame value is simply an empty list.
seq_length = ft["shape"][0] if ft["shape"][0] > 0 else -1
hf_features[key] = datasets.Sequence(length=seq_length, feature=datasets.Value(dtype=ft["dtype"]))
elif len(ft["shape"]) == 2:
hf_features[key] = datasets.Array2D(shape=ft["shape"], dtype=ft["dtype"])
elif len(ft["shape"]) == 3:
-3
View File
@@ -155,7 +155,6 @@ class MetaworldEnv(gym.Env):
env.model.cam_pos[2] = [0.75, 0.075, 0.7]
env.reset()
env._freeze_rand_vec = False # otherwise no randomization
env.seeded_rand_vec = True # use seeded RNG so reset(seed=X) controls object positions
self._env = env
def render(self) -> np.ndarray:
@@ -221,8 +220,6 @@ class MetaworldEnv(gym.Env):
self._ensure_env()
super().reset(seed=seed)
if seed is not None:
self._env.seed(seed)
raw_obs, info = self._env.reset(seed=seed)
observation = self._format_raw_obs(raw_obs)
+1 -9
View File
@@ -44,19 +44,12 @@ from lerobot.utils.constants import (
POLICY_PREPROCESSOR_DEFAULT_NAME,
)
from lerobot.utils.feature_utils import dataset_to_policy_features
from lerobot.utils.import_utils import _peft_available, require_package
from .evo1.configuration_evo1 import Evo1Config
from .groot.configuration_groot import GrootConfig
from .pretrained import PreTrainedPolicy
from .utils import validate_visual_features_consistency
if TYPE_CHECKING or _peft_available:
from peft import PeftConfig, PeftModel
else:
PeftConfig = None
PeftModel = None
def _reconnect_relative_absolute_steps(
preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline
@@ -184,7 +177,6 @@ def make_pre_post_processors(
return make_groot_pre_post_processors_from_pretrained(
config=policy_cfg,
pretrained_path=pretrained_path,
revision=pretrained_revision,
dataset_stats=kwargs.get("dataset_stats"),
dataset_meta=kwargs.get("dataset_meta"),
preprocessor_overrides=kwargs.get("preprocessor_overrides"),
@@ -341,7 +333,7 @@ def make_policy(
# Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
# of the adapter and the adapter's config contains the path to the base policy. So we need the
# adapter config first, then load the correct policy and then apply PEFT.
require_package("peft", extra="peft")
from peft import PeftConfig, PeftModel
logging.info("Loading policy's PEFT adapter.")
@@ -37,19 +37,13 @@ def is_image_feature(key: str) -> bool:
@dataclass
class ConcurrencyConfig:
"""Configuration for the concurrency of the actor and learner.
Possible values are:
- "threads": Use threads for the actor and learner.
- "processes": Use processes for the actor and learner.
``multiprocessing_context`` selects the process-wide start method when
processes are used. Set it to ``None`` to preserve Python's default or a
method already selected by the embedding application.
"""
actor: str = "threads"
learner: str = "threads"
multiprocessing_context: str | None = "spawn"
@dataclass
@@ -475,7 +475,6 @@ def make_groot_pre_post_processors_from_pretrained(
config: GrootConfig,
pretrained_path: str,
*,
revision: str | None = None,
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
dataset_meta: Any | None = None,
preprocessor_overrides: dict[str, Any] | None = None,
@@ -512,7 +511,6 @@ def make_groot_pre_post_processors_from_pretrained(
preprocessor, postprocessor = _load_groot_processor_pipelines(
pretrained_path,
revision=revision,
preprocessor_overrides=preprocessor_overrides,
postprocessor_overrides=postprocessor_overrides,
preprocessor_config_filename=preprocessor_config_filename,
@@ -528,7 +526,6 @@ def make_groot_pre_post_processors_from_pretrained(
def _load_groot_processor_pipelines(
pretrained_path: str,
*,
revision: str | None,
preprocessor_overrides: dict[str, Any],
postprocessor_overrides: dict[str, Any],
preprocessor_config_filename: str,
@@ -543,7 +540,6 @@ def _load_groot_processor_pipelines(
preprocessor = PolicyProcessorPipeline.from_pretrained(
pretrained_model_name_or_path=pretrained_path,
config_filename=preprocessor_config_filename,
revision=revision,
overrides=preprocessor_overrides,
to_transition=batch_to_transition,
to_output=transition_to_batch,
@@ -551,7 +547,6 @@ def _load_groot_processor_pipelines(
postprocessor = PolicyProcessorPipeline.from_pretrained(
pretrained_model_name_or_path=pretrained_path,
config_filename=postprocessor_config_filename,
revision=revision,
overrides=postprocessor_overrides,
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
@@ -43,22 +43,11 @@ from torch.distributions import Beta
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.utils.constants import ACTION
from lerobot.utils.import_utils import (
_peft_available,
_scipy_available,
_transformers_available,
require_package,
)
from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
from ..rtc.modeling_rtc import RTCProcessor
from .configuration_molmoact2 import MolmoAct2Config
if TYPE_CHECKING or _peft_available:
from peft import LoraConfig, get_peft_model
else:
LoraConfig = None
get_peft_model = None
logger = logging.getLogger(__name__)
@@ -1742,11 +1731,13 @@ class MolmoAct2Policy(PreTrainedPolicy):
def _build_inner_lora_config(self):
require_package("peft", extra="molmoact2")
from peft import LoraConfig
return LoraConfig(**self._get_inner_peft_targets())
def _apply_lora_adapters(self) -> None:
require_package("peft", extra="molmoact2")
from peft import get_peft_model
peft_config = self._build_inner_lora_config()
self._validate_peft_config(peft_config)
+5 -13
View File
@@ -34,22 +34,14 @@ from lerobot.configs import PreTrainedConfig
from lerobot.configs.train import TrainPipelineConfig
from lerobot.utils.device_utils import resolve_safetensors_device
from lerobot.utils.hub import HubMixin
from lerobot.utils.import_utils import _peft_available, require_package
from .utils import log_model_loading_keys
if TYPE_CHECKING or _peft_available:
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType, get_peft_model
else:
PEFT_TYPE_TO_CONFIG_MAPPING = None
PeftType = None
get_peft_model = None
T = TypeVar("T", bound="PreTrainedPolicy")
if TYPE_CHECKING:
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
T = TypeVar("T", bound="PreTrainedPolicy")
def _build_card_context(
cfg: TrainPipelineConfig | None,
@@ -392,7 +384,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.)
These are merged with policy defaults to build the final config.
"""
require_package("peft", extra="peft")
from peft import get_peft_model
# If user provided a complete config, use it directly (with overrides)
if peft_config is not None:
@@ -463,7 +455,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
Returns:
Preprocessed dict with renamed keys and init_type mapped to method-specific key.
"""
require_package("peft", extra="peft")
from peft import PeftType
cli_overrides = cli_overrides.copy()
@@ -488,7 +480,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
def _build_peft_config(self, cli_overrides: dict):
"""Build a PEFT config from policy defaults and CLI overrides."""
require_package("peft", extra="peft")
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
# Determine PEFT method type (default to LORA)
method_type_str = cli_overrides.get("method_type") or "lora"
@@ -515,7 +507,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict):
"""Apply CLI overrides to an existing PEFT config."""
require_package("peft", extra="peft")
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
# Get method type from existing config or CLI override
method_type_str = cli_overrides.get("method_type")
@@ -132,20 +132,10 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
for axis in ["x", "y", "z"]:
for axis in ["x", "y", "z", "gripper"]:
features[PipelineFeatureType.ACTION].pop(f"delta_{axis}", None)
features[PipelineFeatureType.ACTION].pop("gripper", None)
for feat in [
"enabled",
"target_x",
"target_y",
"target_z",
"target_wx",
"target_wy",
"target_wz",
"gripper_vel",
]:
for feat in ["enabled", "target_x", "target_y", "target_z", "target_wx", "target_wy", "target_wz"]:
features[PipelineFeatureType.ACTION][f"{feat}"] = PolicyFeature(
type=FeatureType.ACTION, shape=(1,)
)
+4 -18
View File
@@ -713,8 +713,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
ProcessorMigrationError: If the model requires migration to processor format.
"""
model_id = str(pretrained_model_name_or_path)
model_path = Path(model_id)
is_local_source = model_path.is_dir() or model_path.is_file()
hub_download_kwargs = {
"force_download": force_download,
"resume_download": resume_download,
@@ -733,7 +731,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
# 3. Build steps with overrides
steps, validated_overrides = cls._build_steps_with_overrides(
loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs, is_local_source
loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs
)
# 4. Validate that all overrides were used
@@ -923,7 +921,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
model_id: str,
base_path: Path | None,
hub_download_kwargs: dict[str, Any],
is_local_source: bool = False,
) -> tuple[list[ProcessorStep], set[str]]:
"""Build all processor steps with overrides and state loading.
@@ -947,7 +944,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
3. **State Loading** (via _load_step_state):
- **If step has "state_file"**: Load tensor state from .safetensors
- **Local first**: Check base_path/state_file.safetensors
- **Hub fallback**: Download state file if the pipeline was loaded from the Hub
- **Hub fallback**: Download state file if not found locally
- **Optional**: Only load if step has load_state_dict method
4. **Override Tracking**:
@@ -965,7 +962,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
model_id: The model identifier (needed for Hub state file downloads)
base_path: Local directory path for finding state files
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
is_local_source: Whether model_id resolved to a local directory or config file.
Returns:
Tuple of (instantiated_steps_list, unused_override_keys)
@@ -979,9 +975,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides)
for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True):
cls._load_step_state(
step_instance, step_entry, model_id, base_path, hub_download_kwargs, is_local_source
)
cls._load_step_state(step_instance, step_entry, model_id, base_path, hub_download_kwargs)
return steps, remaining_override_keys
@@ -1145,7 +1139,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
model_id: str,
base_path: Path | None,
hub_download_kwargs: dict[str, Any],
is_local_source: bool = False,
) -> None:
"""Load state dictionary for a processor step if available.
@@ -1164,7 +1157,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
- **Use case**: Loading from local saved model directory
2. **Hub download fallback**: Download state file from repository
- **When triggered**: Local file not found and the pipeline source is a Hub repo
- **When triggered**: Local file not found or base_path is None
- **Process**: Use hf_hub_download with same parameters as config
- **Example**: Download "normalize_step_0.safetensors" from "user/repo"
- **Result**: Downloaded to local cache, path returned
@@ -1185,7 +1178,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
model_id: The model identifier (used for Hub downloads if needed)
base_path: Local directory path for finding state files (None for Hub-only)
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
is_local_source: Whether model_id resolved to a local directory or config file.
Note:
This method modifies step_instance in-place and returns None.
@@ -1199,12 +1191,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
# Try local file first
if base_path and (base_path / state_filename).exists():
state_path = str(base_path / state_filename)
elif is_local_source:
state_path = base_path / state_filename if base_path else Path(state_filename)
raise FileNotFoundError(
f"State file '{state_filename}' was not found for local processor pipeline "
f"'{model_id}' at '{state_path}'."
)
else:
# Download from Hub
state_path = hf_hub_download(
+4 -2
View File
@@ -91,7 +91,7 @@ from lerobot.robots import so_follower # noqa: F401
from lerobot.teleoperators import gamepad, so_leader # noqa: F401
from lerobot.teleoperators.utils import TeleopEvents
from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
from lerobot.utils.process import ProcessSignalHandler
from lerobot.utils.random_utils import set_seed
from lerobot.utils.robot_utils import precise_sleep
from lerobot.utils.transition import (
@@ -124,7 +124,9 @@ def actor_cli(cfg: TrainRLServerPipelineConfig):
cfg.validate()
display_pid = False
if not use_threads(cfg):
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
import torch.multiprocessing as mp
mp.set_start_method("spawn")
display_pid = True
# Create logs directory to ensure it exists
+4 -2
View File
@@ -102,7 +102,7 @@ from lerobot.utils.constants import (
)
from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.io_utils import load_json, write_json
from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
from lerobot.utils.process import ProcessSignalHandler
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
format_big_number,
@@ -123,7 +123,9 @@ def train_cli(cfg: TrainRLServerPipelineConfig):
# Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
require_package("grpcio", extra="hilserl", import_name="grpc")
if not use_threads(cfg):
ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
import torch.multiprocessing as mp
mp.set_start_method("spawn")
# Use the job_name from the config
train(
@@ -323,10 +323,6 @@ class LeKiwiClient(Robot):
np.ndarray: the action sent to the motors, potentially clipped.
"""
# Action values may be torch tensors (e.g. replayed from a dataset) or numpy
# scalars; json.dumps only serializes Python primitives, so coerce each value to a
# plain float before sending.
action = {key: float(value) for key, value in action.items()}
self.zmq_cmd_socket.send_string(json.dumps(action)) # action is in motor space
# TODO(Steven): Remove the np conversion when it is possible to record a non-numpy array value
+2 -11
View File
@@ -326,17 +326,8 @@ class RolloutConfig:
policy_path = parser.get_path_arg("policy")
if policy_path:
yaml_overrides = parser.get_yaml_overrides("policy")
cli_overrides = parser.get_cli_overrides("policy") or []
policy_overrides = yaml_overrides + cli_overrides
pretrained_revision = parser.parse_arg("pretrained_revision", cli_overrides)
if pretrained_revision is None:
pretrained_revision = parser.parse_arg("pretrained_revision", yaml_overrides)
self.policy = PreTrainedConfig.from_pretrained(
policy_path,
revision=pretrained_revision,
cli_overrides=policy_overrides,
)
cli_overrides = parser.get_cli_overrides("policy")
self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
self.policy.pretrained_path = policy_path
if self.policy is None:
raise ValueError("--policy.path is required for rollout")
+15 -52
View File
@@ -24,11 +24,10 @@ from __future__ import annotations
import logging
from dataclasses import dataclass, field
from threading import Event
from typing import TYPE_CHECKING
import torch
from lerobot.configs import FeatureType, PreTrainedConfig
from lerobot.configs import FeatureType
from lerobot.datasets import (
LeRobotDataset,
aggregate_pipeline_dataset_features,
@@ -48,7 +47,6 @@ from lerobot.processor.relative_action_processor import RelativeActionsProcessor
from lerobot.robots import make_robot_from_config
from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
from lerobot.utils.import_utils import _peft_available, require_package
from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
from .inference import (
@@ -59,12 +57,6 @@ from .inference import (
)
from .robot_wrapper import ThreadSafeRobot
if TYPE_CHECKING or _peft_available:
from peft import PeftConfig, PeftModel
else:
PeftConfig = None
PeftModel = None
logger = logging.getLogger(__name__)
@@ -167,35 +159,6 @@ class RolloutContext:
# ---------------------------------------------------------------------------
def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy:
"""Load policy weights, keeping adapter and base-model revisions independent."""
pretrained_revision = policy_config.pretrained_revision
policy_class = get_policy_class(policy_config.type)
if not policy_config.use_peft:
return policy_class.from_pretrained(
policy_config.pretrained_path,
config=policy_config,
revision=pretrained_revision,
)
require_package("peft", extra="peft")
peft_path = policy_config.pretrained_path
peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision)
policy = policy_class.from_pretrained(
pretrained_name_or_path=peft_config.base_model_name_or_path,
config=policy_config,
revision=peft_config.revision,
)
return PeftModel.from_pretrained(
policy,
peft_path,
config=peft_config,
revision=pretrained_revision,
)
def build_rollout_context(
cfg: RolloutConfig,
shutdown_event: Event,
@@ -213,6 +176,7 @@ def build_rollout_context(
# --- 1. Policy (heavy I/O, but no hardware yet) -------------------
logger.info("Loading policy from '%s'...", cfg.policy.pretrained_path)
policy_config = cfg.policy
policy_class = get_policy_class(policy_config.type)
if hasattr(policy_config, "compile_model"):
policy_config.compile_model = cfg.use_torch_compile
@@ -223,7 +187,17 @@ def build_rollout_context(
"Please use `cpu` or `cuda` backend."
)
policy = _load_pretrained_policy(policy_config)
if policy_config.use_peft:
from peft import PeftConfig, PeftModel
peft_path = policy_config.pretrained_path
peft_config = PeftConfig.from_pretrained(peft_path)
policy = policy_class.from_pretrained(
pretrained_name_or_path=peft_config.base_model_name_or_path, config=policy_config
)
policy = PeftModel.from_pretrained(policy, peft_path, config=peft_config)
else:
policy = policy_class.from_pretrained(policy_config.pretrained_path, config=policy_config)
if is_rtc:
policy.config.rtc_config = cfg.inference.rtc
@@ -302,22 +276,12 @@ def build_rollout_context(
# ``observation_features`` values are either a tuple (camera shape) or the
# ``float`` type itself used as a sentinel for scalar motor features —
# see ``dict[str, type | tuple]`` annotation on ``Robot.observation_features``.
# Keep cameras (tuple) plus both joint-position (.pos) and base-velocity (.vel)
# scalar state features. LeKiwi's observation.state is 9-dim (6 arm .pos +
# x/y/theta.vel) and the policy was trained/normalized on all 9; the old .pos-only
# filter fed a 6-dim state into a 9-dim normalizer → RuntimeError (size 6 vs 9).
# Pure-arm robots have no .vel state keys, so this is a no-op for them.
observation_features_hw = {
k: v
for k, v in all_obs_features.items()
if isinstance(v, tuple) or (v is float and k.endswith((".pos", ".vel")))
if isinstance(v, tuple) or (v is float and k.endswith(".pos"))
}
# Keep both joint-position (.pos) and base-velocity (.vel) action features so
# mobile manipulators command the base too (e.g. LeKiwi: 6 arm .pos +
# x/y/theta.vel = 9-dim action). Pure-arm robots have no .vel keys, so this is
# a no-op for them. Without the .vel keys the base velocities are silently
# dropped from dataset_features[ACTION]/ordered_action_keys and the base never moves.
action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith((".pos", ".vel"))}
action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith(".pos")}
# The action side is always needed: sync inference reads action names from
# ``dataset_features[ACTION]`` to map policy tensors back to robot actions.
@@ -428,7 +392,6 @@ def build_rollout_context(
preprocessor, postprocessor = make_pre_post_processors(
policy_cfg=policy_config,
pretrained_path=cfg.policy.pretrained_path,
pretrained_revision=policy_config.pretrained_revision,
dataset_stats=dataset_stats,
preprocessor_overrides={
"device_processor": {"device": cfg.device},
+9 -12
View File
@@ -62,7 +62,7 @@ from dataclasses import asdict
from functools import partial
from pathlib import Path
from pprint import pformat
from typing import TYPE_CHECKING, Any, TypedDict
from typing import Any, TypedDict
import einops
import gymnasium as gym
@@ -87,7 +87,7 @@ from lerobot.processor import PolicyProcessorPipeline
from lerobot.types import PolicyAction
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
from lerobot.utils.import_utils import register_third_party_plugins
from lerobot.utils.io_utils import write_video
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
@@ -95,11 +95,6 @@ from lerobot.utils.utils import (
inside_slurm,
)
if TYPE_CHECKING or _peft_available:
from peft import PeftModel
else:
PeftModel = None
def _env_features_to_dataset_features(env_features: dict) -> dict:
"""Convert EnvConfig.features to the dict format expected by LeRobotDataset.create()."""
@@ -449,11 +444,13 @@ def eval_policy(
exc = ValueError(
f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
)
if not _peft_available:
raise exc
require_package("peft", extra="peft")
if not isinstance(policy, PeftModel):
raise exc
try:
from peft import PeftModel
if not isinstance(policy, PeftModel):
raise exc
except ImportError:
raise exc from None
start = time.time()
# Preserve the mode for direct callers. eval_policy_all scopes the mode
-1
View File
@@ -61,7 +61,6 @@ from lerobot.robots import ( # noqa: F401
earthrover_mini_plus,
hope_jr,
koch_follower,
lekiwi,
make_robot_from_config,
omx_follower,
openarm_follower,
-1
View File
@@ -165,7 +165,6 @@ from lerobot.robots import ( # noqa: F401
earthrover_mini_plus,
hope_jr,
koch_follower,
lekiwi,
omx_follower,
openarm_follower,
reachy2,
+8 -19
View File
@@ -57,7 +57,7 @@ from lerobot.optim.factory import make_optimizer_and_scheduler
from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
from lerobot.rewards import make_reward_pre_post_processors
from lerobot.utils.collate import lerobot_collate_fn
from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
from lerobot.utils.import_utils import register_third_party_plugins
from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import (
@@ -68,24 +68,9 @@ from lerobot.utils.utils import (
inside_slurm,
)
if TYPE_CHECKING or _peft_available:
from peft import PeftModel
else:
PeftModel = None
from .lerobot_eval import eval_policy_all
def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
"""Return worker-only DataLoader options, disabling them for single-process loading."""
workers_enabled = cfg.num_workers > 0
return {
"prefetch_factor": cfg.prefetch_factor if workers_enabled else None,
"persistent_workers": cfg.persistent_workers and workers_enabled,
"multiprocessing_context": cfg.dataloader_multiprocessing_context if workers_enabled else None,
}
def update_policy(
train_metrics: MetricsTracker,
policy: PreTrainedPolicy,
@@ -212,6 +197,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if cfg.job.is_remote:
return submit_to_hf(cfg)
from lerobot.utils.import_utils import require_package
require_package("accelerate", extra="training")
from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs, DistributedType
@@ -315,7 +302,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if cfg.peft is not None:
if cfg.is_reward_model_training:
raise ValueError("PEFT is only supported for policy training. ")
require_package("peft", extra="peft")
from peft import PeftModel
if isinstance(policy, PeftModel):
logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.")
@@ -486,7 +473,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
pin_memory=device.type == "cuda",
drop_last=False,
collate_fn=collate_fn,
**_dataloader_worker_kwargs(cfg),
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
)
# Build eval dataloader if a held-out split exists
@@ -512,7 +500,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
pin_memory=device.type == "cuda",
drop_last=False,
collate_fn=eval_collate_fn,
**_dataloader_worker_kwargs(cfg),
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
)
# Prepare everything with accelerator
-28
View File
@@ -16,39 +16,11 @@
# limitations under the License.
import logging
import multiprocessing
import os
import signal
import sys
def ensure_multiprocessing_start_method(start_method: str | None) -> None:
"""Set a multiprocessing start method once, or verify the existing method matches.
Passing ``None`` leaves Python's process-wide default untouched. This is useful
when LeRobot is embedded in an application that owns multiprocessing setup.
"""
if start_method is None:
return
available_methods = multiprocessing.get_all_start_methods()
if start_method not in available_methods:
raise ValueError(
f"Multiprocessing start method must be one of {available_methods} on this platform, "
f"got {start_method!r}."
)
current_method = multiprocessing.get_start_method(allow_none=True)
if current_method is None:
multiprocessing.set_start_method(start_method)
elif current_method != start_method:
raise RuntimeError(
f"Multiprocessing start method is already {current_method!r}; cannot change it to "
f"{start_method!r}. Set the configured multiprocessing context to null to keep the "
"application's existing method, or launch LeRobot in a fresh process."
)
class ProcessSignalHandler:
"""Utility class to attach graceful shutdown signal handlers.
+20
View File
@@ -687,6 +687,26 @@ def test_compute_episode_stats_string_features_skipped():
assert "q01" in stats["action"]
@pytest.mark.parametrize("shape", [(0,), (0, 2), (2, 0), (1, 0, 2)])
def test_compute_episode_stats_zero_width_feature_skipped(shape):
"""Features with any zero-width dimension carry no values and are skipped."""
episode_data = {
"action": np.random.normal(0, 1, (100, 5)).astype(np.float32),
"target": np.zeros((100, *shape), dtype=np.float32),
}
features = {
"action": {"dtype": "float32", "shape": (5,)},
"target": {"dtype": "float32", "shape": shape},
}
stats = compute_episode_stats(episode_data, features)
# Zero-width features are skipped, just like strings; non-empty features are unaffected.
assert "target" not in stats
assert "action" in stats
assert "q01" in stats["action"]
def test_aggregate_feature_stats_with_quantiles():
"""Test aggregating feature stats that include quantiles."""
stats_ft_list = [
+31
View File
@@ -27,6 +27,7 @@ pytest.importorskip("datasets", reason="datasets is required (install lerobot[da
from lerobot.configs import VideoEncoderConfig
from lerobot.datasets.dataset_writer import _encode_video_worker
from lerobot.datasets.feature_utils import get_hf_features_from_features
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.utils import DEFAULT_IMAGE_PATH
from tests.fixtures.constants import DEFAULT_FPS, DUMMY_REPO_ID
@@ -189,6 +190,36 @@ def test_save_multiple_episodes(tmp_path):
assert dataset.meta.total_frames == total_frames
def test_save_episode_with_zero_width_feature(tmp_path):
"""A one-dimensional empty numeric feature round-trips and has no statistics."""
features = {
**SIMPLE_FEATURES,
"target": {"dtype": "float32", "shape": (0,), "names": None},
}
root = tmp_path / "ds"
dataset = LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=features, root=root)
for _ in range(4):
dataset.add_frame(_make_frame(features))
dataset.save_episode()
dataset.finalize()
assert dataset.meta.total_episodes == 1
assert dataset.meta.total_frames == 4
reloaded = LeRobotDataset(repo_id=DUMMY_REPO_ID, root=root)
target = np.asarray(reloaded[0]["target"])
assert target.shape == (0,)
assert "target" not in (reloaded.meta.stats or {})
@pytest.mark.parametrize("shape", [(0, 2), (2, 0), (1, 0, 2)])
def test_multidimensional_zero_width_feature_rejected(shape):
features = {"target": {"dtype": "float32", "shape": shape, "names": None}}
with pytest.raises(ValueError, match="Multidimensional features with a zero-width dimension"):
get_hf_features_from_features(features)
# ── clear / lifecycle ────────────────────────────────────────────────
@@ -113,7 +113,6 @@ def test_gaussian_actor_config_default_initialization():
# Concurrency configuration
assert config.concurrency.actor == "threads"
assert config.concurrency.learner == "threads"
assert config.concurrency.multiprocessing_context == "spawn"
assert isinstance(config.actor_network_kwargs, ActorNetworkConfig)
assert isinstance(config.policy_kwargs, PolicyConfig)
@@ -153,7 +152,6 @@ def test_concurrency_config():
config = ConcurrencyConfig()
assert config.actor == "threads"
assert config.learner == "threads"
assert config.multiprocessing_context == "spawn"
def test_gaussian_actor_config_custom_initialization():
@@ -26,17 +26,8 @@ import tempfile
from pathlib import Path
import pytest
import torch
from safetensors.torch import save_file
from lerobot.configs import PipelineFeatureType, PolicyFeature
from lerobot.processor.pipeline import (
DataProcessorPipeline,
ProcessorMigrationError,
ProcessorStep,
ProcessorStepRegistry,
)
from lerobot.types import EnvTransition
from lerobot.processor.pipeline import DataProcessorPipeline, ProcessorMigrationError
# Simplified Config Loading Tests
@@ -107,140 +98,6 @@ def test_load_config_nonexistent_path_tries_hub():
DataProcessorPipeline._load_config("nonexistent/path", "processor.json", {})
def test_from_pretrained_local_directory_missing_state_does_not_call_hub(monkeypatch):
"""Local processor dirs must fail locally when a state file is missing."""
@ProcessorStepRegistry.register("local_missing_state_step")
class LocalMissingStateStep(ProcessorStep):
def __call__(self, transition: EnvTransition) -> EnvTransition:
return transition
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return features
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
try:
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
config = {
"name": "LocalMissingStatePipeline",
"steps": [{"registry_name": "local_missing_state_step", "state_file": "missing.safetensors"}],
}
(tmp_path / "processor.json").write_text(json.dumps(config))
def fail_hub_download(*args, **kwargs):
pytest.fail("local missing processor state should not call hf_hub_download")
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fail_hub_download)
with pytest.raises(FileNotFoundError, match="missing.safetensors.*local processor pipeline"):
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="processor.json")
finally:
ProcessorStepRegistry.unregister("local_missing_state_step")
def test_from_pretrained_local_config_file_missing_state_does_not_call_hub(monkeypatch):
"""Local single-file processor configs must also keep missing state resolution local."""
@ProcessorStepRegistry.register("local_file_missing_state_step")
class LocalFileMissingStateStep(ProcessorStep):
def __call__(self, transition: EnvTransition) -> EnvTransition:
return transition
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return features
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
pass
try:
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
config_path = tmp_path / "processor.json"
config = {
"name": "LocalFileMissingStatePipeline",
"steps": [
{"registry_name": "local_file_missing_state_step", "state_file": "missing.safetensors"}
],
}
config_path.write_text(json.dumps(config))
def fail_hub_download(*args, **kwargs):
pytest.fail("local missing processor state should not call hf_hub_download")
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fail_hub_download)
with pytest.raises(FileNotFoundError, match="missing.safetensors.*local processor pipeline"):
DataProcessorPipeline.from_pretrained(config_path, config_filename="ignored.json")
finally:
ProcessorStepRegistry.unregister("local_file_missing_state_step")
def test_from_pretrained_hub_source_missing_local_state_still_calls_hub(monkeypatch, tmp_path):
"""Hub sources still fall back to hf_hub_download for state files."""
@ProcessorStepRegistry.register("hub_state_step")
class HubStateStep(ProcessorStep):
def __init__(self):
self.value = torch.tensor(0)
def __call__(self, transition: EnvTransition) -> EnvTransition:
return transition
def transform_features(
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
return features
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
self.value = state["value"]
try:
state_path = tmp_path / "downloaded.safetensors"
save_file({"value": torch.tensor(7)}, state_path)
loaded_config = {
"name": "HubStatePipeline",
"steps": [{"registry_name": "hub_state_step", "state_file": "hub_state.safetensors"}],
}
calls = []
def fake_load_config(cls, model_id, config_filename, hub_download_kwargs):
return loaded_config, tmp_path / "hub_cache"
def fake_hub_download(**kwargs):
calls.append(kwargs)
return str(state_path)
monkeypatch.setattr(DataProcessorPipeline, "_load_config", classmethod(fake_load_config))
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fake_hub_download)
pipeline = DataProcessorPipeline.from_pretrained("user/repo", config_filename="processor.json")
assert calls == [
{
"repo_id": "user/repo",
"filename": "hub_state.safetensors",
"repo_type": "model",
"force_download": False,
"resume_download": None,
"proxies": None,
"token": None,
"cache_dir": None,
"local_files_only": False,
"revision": None,
}
]
assert pipeline.steps[0].value.item() == 7
finally:
ProcessorStepRegistry.unregister("hub_state_step")
# Config Validation Tests
-112
View File
@@ -17,8 +17,6 @@
from __future__ import annotations
import dataclasses
import sys
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
@@ -108,116 +106,6 @@ def test_sentry_config_defaults():
assert cfg.target_video_file_size_mb is None
def test_rollout_config_passes_policy_pretrained_revision(monkeypatch):
from lerobot.configs import PreTrainedConfig, parser
from lerobot.rollout import RolloutConfig
from tests.mocks.mock_robot import MockRobotConfig
captured = {}
def fake_from_pretrained(cls, pretrained_name_or_path, **kwargs):
captured["pretrained_name_or_path"] = pretrained_name_or_path
captured.update(kwargs)
return SimpleNamespace(device="cpu", pretrained_revision=kwargs["revision"])
monkeypatch.setattr(parser, "get_yaml_overrides", lambda _: ["--pretrained_revision=yaml-sha"])
monkeypatch.setattr(
sys,
"argv",
["lerobot-rollout", "--policy.path=user/policy", "--policy.pretrained_revision=cli-sha"],
)
monkeypatch.setattr(PreTrainedConfig, "from_pretrained", classmethod(fake_from_pretrained))
cfg = RolloutConfig(robot=MockRobotConfig())
assert captured["pretrained_name_or_path"] == "user/policy"
assert captured["revision"] == "cli-sha"
assert captured["cli_overrides"] == [
"--pretrained_revision=yaml-sha",
"--pretrained_revision=cli-sha",
]
assert cfg.policy.pretrained_path == "user/policy"
assert cfg.policy.pretrained_revision == "cli-sha"
def test_load_pretrained_policy_passes_revision(monkeypatch):
import lerobot.rollout.context as rollout_context
policy_config = SimpleNamespace(
type="mock",
use_peft=False,
pretrained_path="user/policy",
pretrained_revision="policy-sha",
)
policy_class = MagicMock()
loaded_policy = MagicMock()
policy_class.from_pretrained.return_value = loaded_policy
monkeypatch.setattr(rollout_context, "get_policy_class", lambda _: policy_class)
policy = rollout_context._load_pretrained_policy(policy_config)
assert policy is loaded_policy
policy_class.from_pretrained.assert_called_once_with(
"user/policy",
config=policy_config,
revision="policy-sha",
)
def test_load_pretrained_peft_policy_keeps_adapter_and_base_revisions_separate(monkeypatch):
import lerobot.rollout.context as rollout_context
policy_config = SimpleNamespace(
type="mock",
use_peft=True,
pretrained_path="user/adapter",
pretrained_revision="adapter-sha",
)
policy_class = MagicMock()
base_policy = MagicMock()
policy_class.from_pretrained.return_value = base_policy
monkeypatch.setattr(rollout_context, "get_policy_class", lambda _: policy_class)
peft_config = SimpleNamespace(
base_model_name_or_path="user/base-policy",
revision="base-sha",
)
peft_config_from_pretrained = MagicMock(return_value=peft_config)
adapted_policy = MagicMock()
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
require_package = MagicMock()
monkeypatch.setattr(rollout_context, "require_package", require_package)
monkeypatch.setattr(
rollout_context,
"PeftConfig",
SimpleNamespace(from_pretrained=peft_config_from_pretrained),
raising=False,
)
monkeypatch.setattr(
rollout_context,
"PeftModel",
SimpleNamespace(from_pretrained=peft_model_from_pretrained),
raising=False,
)
policy = rollout_context._load_pretrained_policy(policy_config)
assert policy is adapted_policy
require_package.assert_called_once_with("peft", extra="peft")
peft_config_from_pretrained.assert_called_once_with("user/adapter", revision="adapter-sha")
policy_class.from_pretrained.assert_called_once_with(
pretrained_name_or_path="user/base-policy",
config=policy_config,
revision="base-sha",
)
peft_model_from_pretrained.assert_called_once_with(
base_policy,
"user/adapter",
config=peft_config,
revision="adapter-sha",
)
# ---------------------------------------------------------------------------
# RolloutRingBuffer
# ---------------------------------------------------------------------------