mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-28 12:15:59 +00:00
Merge remote-tracking branch 'origin/main' into codex/episode-video-streaming-byte-cache
# Conflicts: # src/lerobot/scripts/lerobot_train.py
This commit is contained in:
@@ -0,0 +1,11 @@
|
|||||||
|
version: 2
|
||||||
|
updates:
|
||||||
|
- package-ecosystem: "github-actions"
|
||||||
|
directory: "/"
|
||||||
|
schedule:
|
||||||
|
interval: "weekly"
|
||||||
|
cooldown:
|
||||||
|
default-days: 7
|
||||||
|
groups:
|
||||||
|
actions:
|
||||||
|
patterns: ["*"]
|
||||||
@@ -14,6 +14,7 @@
|
|||||||
import builtins
|
import builtins
|
||||||
import datetime as dt
|
import datetime as dt
|
||||||
import json
|
import json
|
||||||
|
import multiprocessing
|
||||||
import os
|
import os
|
||||||
import tempfile
|
import tempfile
|
||||||
from dataclasses import dataclass, field
|
from dataclasses import dataclass, field
|
||||||
@@ -101,6 +102,12 @@ class TrainPipelineConfig(HubMixin):
|
|||||||
batch_size: int = 8
|
batch_size: int = 8
|
||||||
prefetch_factor: int = 4
|
prefetch_factor: int = 4
|
||||||
persistent_workers: bool = True
|
persistent_workers: bool = True
|
||||||
|
# DataLoader worker start method. "spawn" is safer than "fork" with
|
||||||
|
# non-fork-safe libs (PyAV / torchcodec / ffmpeg), but adds some
|
||||||
|
# worker-startup time per run since workers re-import modules instead
|
||||||
|
# of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
|
||||||
|
# when appropriate, or set it to `null` to use Python's platform default.
|
||||||
|
dataloader_multiprocessing_context: str | None = "spawn"
|
||||||
steps: int = 100_000
|
steps: int = 100_000
|
||||||
# Run policy in the simulation environment every N steps to measure reward/success (0 = disabled).
|
# Run policy in the simulation environment every N steps to measure reward/success (0 = disabled).
|
||||||
env_eval_freq: int = 20_000
|
env_eval_freq: int = 20_000
|
||||||
@@ -212,6 +219,17 @@ class TrainPipelineConfig(HubMixin):
|
|||||||
self.reward_model.pretrained_path = str(policy_dir)
|
self.reward_model.pretrained_path = str(policy_dir)
|
||||||
|
|
||||||
def validate(self) -> None:
|
def validate(self) -> None:
|
||||||
|
available_contexts = multiprocessing.get_all_start_methods()
|
||||||
|
if (
|
||||||
|
self.dataloader_multiprocessing_context is not None
|
||||||
|
and self.dataloader_multiprocessing_context not in available_contexts
|
||||||
|
):
|
||||||
|
raise ValueError(
|
||||||
|
"`dataloader_multiprocessing_context` must be None or one of "
|
||||||
|
f"{available_contexts} on this platform, got "
|
||||||
|
f"{self.dataloader_multiprocessing_context!r}."
|
||||||
|
)
|
||||||
|
|
||||||
self._resolve_pretrained_from_cli()
|
self._resolve_pretrained_from_cli()
|
||||||
|
|
||||||
if self.policy is None and self.reward_model is None:
|
if self.policy is None and self.reward_model is None:
|
||||||
|
|||||||
@@ -172,6 +172,23 @@ class DatasetWriter:
|
|||||||
def _get_image_file_dir(self, episode_index: int, image_key: str) -> Path:
|
def _get_image_file_dir(self, episode_index: int, image_key: str) -> Path:
|
||||||
return self._get_image_file_path(episode_index, image_key, frame_index=0).parent
|
return self._get_image_file_path(episode_index, image_key, frame_index=0).parent
|
||||||
|
|
||||||
|
def _get_episode_buffer_index(self) -> int:
|
||||||
|
episode_index = self.episode_buffer["episode_index"]
|
||||||
|
# episode_index is `int` when freshly created, but becomes `np.ndarray` after
|
||||||
|
# save_episode() mutates the buffer. Handle both types here.
|
||||||
|
if isinstance(episode_index, np.ndarray):
|
||||||
|
episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
|
||||||
|
return int(episode_index)
|
||||||
|
|
||||||
|
def _delete_camera_frame_dirs(self, camera_keys: list[str]) -> None:
|
||||||
|
if self.image_writer is not None:
|
||||||
|
self._wait_image_writer()
|
||||||
|
episode_index = self._get_episode_buffer_index()
|
||||||
|
for camera_key in camera_keys:
|
||||||
|
img_dir = self._get_image_file_dir(episode_index, camera_key)
|
||||||
|
if img_dir.is_dir():
|
||||||
|
shutil.rmtree(img_dir)
|
||||||
|
|
||||||
def _save_image(
|
def _save_image(
|
||||||
self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1
|
self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path, compress_level: int = 1
|
||||||
) -> None:
|
) -> None:
|
||||||
@@ -369,7 +386,9 @@ class DatasetWriter:
|
|||||||
self._episodes_since_last_encoding = 0
|
self._episodes_since_last_encoding = 0
|
||||||
|
|
||||||
if episode_data is None:
|
if episode_data is None:
|
||||||
self.clear_episode_buffer(delete_images=len(self._meta.image_keys) > 0)
|
if len(self._meta.image_keys) > 0:
|
||||||
|
self._delete_camera_frame_dirs(self._meta.image_keys)
|
||||||
|
self.episode_buffer = self._create_episode_buffer()
|
||||||
|
|
||||||
def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None) -> None:
|
def _batch_save_episode_video(self, start_episode: int, end_episode: int | None = None) -> None:
|
||||||
"""Batch save videos for multiple episodes."""
|
"""Batch save videos for multiple episodes."""
|
||||||
@@ -561,10 +580,10 @@ class DatasetWriter:
|
|||||||
return metadata
|
return metadata
|
||||||
|
|
||||||
def clear_episode_buffer(self, delete_images: bool = True) -> None:
|
def clear_episode_buffer(self, delete_images: bool = True) -> None:
|
||||||
"""Discard the current episode buffer and optionally delete temp images.
|
"""Discard the current episode buffer and optionally delete temp camera frames.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
delete_images: If ``True``, remove temporary image directories
|
delete_images: If ``True``, remove temporary camera frame directories
|
||||||
written for the current episode.
|
written for the current episode.
|
||||||
"""
|
"""
|
||||||
# Cancel streaming encoder if active
|
# Cancel streaming encoder if active
|
||||||
@@ -572,17 +591,7 @@ class DatasetWriter:
|
|||||||
self._streaming_encoder.cancel_episode()
|
self._streaming_encoder.cancel_episode()
|
||||||
|
|
||||||
if delete_images:
|
if delete_images:
|
||||||
if self.image_writer is not None:
|
self._delete_camera_frame_dirs(self._meta.camera_keys)
|
||||||
self._wait_image_writer()
|
|
||||||
episode_index = self.episode_buffer["episode_index"]
|
|
||||||
# episode_index is `int` when freshly created, but becomes `np.ndarray` after
|
|
||||||
# save_episode() mutates the buffer. Handle both types here.
|
|
||||||
if isinstance(episode_index, np.ndarray):
|
|
||||||
episode_index = episode_index.item() if episode_index.size == 1 else episode_index[0]
|
|
||||||
for cam_key in self._meta.image_keys:
|
|
||||||
img_dir = self._get_image_file_dir(episode_index, cam_key)
|
|
||||||
if img_dir.is_dir():
|
|
||||||
shutil.rmtree(img_dir)
|
|
||||||
|
|
||||||
self.episode_buffer = self._create_episode_buffer()
|
self.episode_buffer = self._create_episode_buffer()
|
||||||
|
|
||||||
|
|||||||
@@ -155,6 +155,7 @@ class MetaworldEnv(gym.Env):
|
|||||||
env.model.cam_pos[2] = [0.75, 0.075, 0.7]
|
env.model.cam_pos[2] = [0.75, 0.075, 0.7]
|
||||||
env.reset()
|
env.reset()
|
||||||
env._freeze_rand_vec = False # otherwise no randomization
|
env._freeze_rand_vec = False # otherwise no randomization
|
||||||
|
env.seeded_rand_vec = True # use seeded RNG so reset(seed=X) controls object positions
|
||||||
self._env = env
|
self._env = env
|
||||||
|
|
||||||
def render(self) -> np.ndarray:
|
def render(self) -> np.ndarray:
|
||||||
@@ -220,6 +221,8 @@ class MetaworldEnv(gym.Env):
|
|||||||
self._ensure_env()
|
self._ensure_env()
|
||||||
super().reset(seed=seed)
|
super().reset(seed=seed)
|
||||||
|
|
||||||
|
if seed is not None:
|
||||||
|
self._env.seed(seed)
|
||||||
raw_obs, info = self._env.reset(seed=seed)
|
raw_obs, info = self._env.reset(seed=seed)
|
||||||
|
|
||||||
observation = self._format_raw_obs(raw_obs)
|
observation = self._format_raw_obs(raw_obs)
|
||||||
|
|||||||
@@ -177,6 +177,7 @@ def make_pre_post_processors(
|
|||||||
return make_groot_pre_post_processors_from_pretrained(
|
return make_groot_pre_post_processors_from_pretrained(
|
||||||
config=policy_cfg,
|
config=policy_cfg,
|
||||||
pretrained_path=pretrained_path,
|
pretrained_path=pretrained_path,
|
||||||
|
revision=pretrained_revision,
|
||||||
dataset_stats=kwargs.get("dataset_stats"),
|
dataset_stats=kwargs.get("dataset_stats"),
|
||||||
dataset_meta=kwargs.get("dataset_meta"),
|
dataset_meta=kwargs.get("dataset_meta"),
|
||||||
preprocessor_overrides=kwargs.get("preprocessor_overrides"),
|
preprocessor_overrides=kwargs.get("preprocessor_overrides"),
|
||||||
|
|||||||
@@ -475,6 +475,7 @@ def make_groot_pre_post_processors_from_pretrained(
|
|||||||
config: GrootConfig,
|
config: GrootConfig,
|
||||||
pretrained_path: str,
|
pretrained_path: str,
|
||||||
*,
|
*,
|
||||||
|
revision: str | None = None,
|
||||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||||
dataset_meta: Any | None = None,
|
dataset_meta: Any | None = None,
|
||||||
preprocessor_overrides: dict[str, Any] | None = None,
|
preprocessor_overrides: dict[str, Any] | None = None,
|
||||||
@@ -511,6 +512,7 @@ def make_groot_pre_post_processors_from_pretrained(
|
|||||||
|
|
||||||
preprocessor, postprocessor = _load_groot_processor_pipelines(
|
preprocessor, postprocessor = _load_groot_processor_pipelines(
|
||||||
pretrained_path,
|
pretrained_path,
|
||||||
|
revision=revision,
|
||||||
preprocessor_overrides=preprocessor_overrides,
|
preprocessor_overrides=preprocessor_overrides,
|
||||||
postprocessor_overrides=postprocessor_overrides,
|
postprocessor_overrides=postprocessor_overrides,
|
||||||
preprocessor_config_filename=preprocessor_config_filename,
|
preprocessor_config_filename=preprocessor_config_filename,
|
||||||
@@ -526,6 +528,7 @@ def make_groot_pre_post_processors_from_pretrained(
|
|||||||
def _load_groot_processor_pipelines(
|
def _load_groot_processor_pipelines(
|
||||||
pretrained_path: str,
|
pretrained_path: str,
|
||||||
*,
|
*,
|
||||||
|
revision: str | None,
|
||||||
preprocessor_overrides: dict[str, Any],
|
preprocessor_overrides: dict[str, Any],
|
||||||
postprocessor_overrides: dict[str, Any],
|
postprocessor_overrides: dict[str, Any],
|
||||||
preprocessor_config_filename: str,
|
preprocessor_config_filename: str,
|
||||||
@@ -540,6 +543,7 @@ def _load_groot_processor_pipelines(
|
|||||||
preprocessor = PolicyProcessorPipeline.from_pretrained(
|
preprocessor = PolicyProcessorPipeline.from_pretrained(
|
||||||
pretrained_model_name_or_path=pretrained_path,
|
pretrained_model_name_or_path=pretrained_path,
|
||||||
config_filename=preprocessor_config_filename,
|
config_filename=preprocessor_config_filename,
|
||||||
|
revision=revision,
|
||||||
overrides=preprocessor_overrides,
|
overrides=preprocessor_overrides,
|
||||||
to_transition=batch_to_transition,
|
to_transition=batch_to_transition,
|
||||||
to_output=transition_to_batch,
|
to_output=transition_to_batch,
|
||||||
@@ -547,6 +551,7 @@ def _load_groot_processor_pipelines(
|
|||||||
postprocessor = PolicyProcessorPipeline.from_pretrained(
|
postprocessor = PolicyProcessorPipeline.from_pretrained(
|
||||||
pretrained_model_name_or_path=pretrained_path,
|
pretrained_model_name_or_path=pretrained_path,
|
||||||
config_filename=postprocessor_config_filename,
|
config_filename=postprocessor_config_filename,
|
||||||
|
revision=revision,
|
||||||
overrides=postprocessor_overrides,
|
overrides=postprocessor_overrides,
|
||||||
to_transition=policy_action_to_transition,
|
to_transition=policy_action_to_transition,
|
||||||
to_output=transition_to_policy_action,
|
to_output=transition_to_policy_action,
|
||||||
|
|||||||
@@ -713,6 +713,8 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
ProcessorMigrationError: If the model requires migration to processor format.
|
ProcessorMigrationError: If the model requires migration to processor format.
|
||||||
"""
|
"""
|
||||||
model_id = str(pretrained_model_name_or_path)
|
model_id = str(pretrained_model_name_or_path)
|
||||||
|
model_path = Path(model_id)
|
||||||
|
is_local_source = model_path.is_dir() or model_path.is_file()
|
||||||
hub_download_kwargs = {
|
hub_download_kwargs = {
|
||||||
"force_download": force_download,
|
"force_download": force_download,
|
||||||
"resume_download": resume_download,
|
"resume_download": resume_download,
|
||||||
@@ -731,7 +733,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
|
|
||||||
# 3. Build steps with overrides
|
# 3. Build steps with overrides
|
||||||
steps, validated_overrides = cls._build_steps_with_overrides(
|
steps, validated_overrides = cls._build_steps_with_overrides(
|
||||||
loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs
|
loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs, is_local_source
|
||||||
)
|
)
|
||||||
|
|
||||||
# 4. Validate that all overrides were used
|
# 4. Validate that all overrides were used
|
||||||
@@ -921,6 +923,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
model_id: str,
|
model_id: str,
|
||||||
base_path: Path | None,
|
base_path: Path | None,
|
||||||
hub_download_kwargs: dict[str, Any],
|
hub_download_kwargs: dict[str, Any],
|
||||||
|
is_local_source: bool = False,
|
||||||
) -> tuple[list[ProcessorStep], set[str]]:
|
) -> tuple[list[ProcessorStep], set[str]]:
|
||||||
"""Build all processor steps with overrides and state loading.
|
"""Build all processor steps with overrides and state loading.
|
||||||
|
|
||||||
@@ -944,7 +947,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
3. **State Loading** (via _load_step_state):
|
3. **State Loading** (via _load_step_state):
|
||||||
- **If step has "state_file"**: Load tensor state from .safetensors
|
- **If step has "state_file"**: Load tensor state from .safetensors
|
||||||
- **Local first**: Check base_path/state_file.safetensors
|
- **Local first**: Check base_path/state_file.safetensors
|
||||||
- **Hub fallback**: Download state file if not found locally
|
- **Hub fallback**: Download state file if the pipeline was loaded from the Hub
|
||||||
- **Optional**: Only load if step has load_state_dict method
|
- **Optional**: Only load if step has load_state_dict method
|
||||||
|
|
||||||
4. **Override Tracking**:
|
4. **Override Tracking**:
|
||||||
@@ -962,6 +965,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
model_id: The model identifier (needed for Hub state file downloads)
|
model_id: The model identifier (needed for Hub state file downloads)
|
||||||
base_path: Local directory path for finding state files
|
base_path: Local directory path for finding state files
|
||||||
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
|
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
|
||||||
|
is_local_source: Whether model_id resolved to a local directory or config file.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
Tuple of (instantiated_steps_list, unused_override_keys)
|
Tuple of (instantiated_steps_list, unused_override_keys)
|
||||||
@@ -975,7 +979,9 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides)
|
steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides)
|
||||||
|
|
||||||
for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True):
|
for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True):
|
||||||
cls._load_step_state(step_instance, step_entry, model_id, base_path, hub_download_kwargs)
|
cls._load_step_state(
|
||||||
|
step_instance, step_entry, model_id, base_path, hub_download_kwargs, is_local_source
|
||||||
|
)
|
||||||
|
|
||||||
return steps, remaining_override_keys
|
return steps, remaining_override_keys
|
||||||
|
|
||||||
@@ -1139,6 +1145,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
model_id: str,
|
model_id: str,
|
||||||
base_path: Path | None,
|
base_path: Path | None,
|
||||||
hub_download_kwargs: dict[str, Any],
|
hub_download_kwargs: dict[str, Any],
|
||||||
|
is_local_source: bool = False,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Load state dictionary for a processor step if available.
|
"""Load state dictionary for a processor step if available.
|
||||||
|
|
||||||
@@ -1157,7 +1164,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
- **Use case**: Loading from local saved model directory
|
- **Use case**: Loading from local saved model directory
|
||||||
|
|
||||||
2. **Hub download fallback**: Download state file from repository
|
2. **Hub download fallback**: Download state file from repository
|
||||||
- **When triggered**: Local file not found or base_path is None
|
- **When triggered**: Local file not found and the pipeline source is a Hub repo
|
||||||
- **Process**: Use hf_hub_download with same parameters as config
|
- **Process**: Use hf_hub_download with same parameters as config
|
||||||
- **Example**: Download "normalize_step_0.safetensors" from "user/repo"
|
- **Example**: Download "normalize_step_0.safetensors" from "user/repo"
|
||||||
- **Result**: Downloaded to local cache, path returned
|
- **Result**: Downloaded to local cache, path returned
|
||||||
@@ -1178,6 +1185,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
model_id: The model identifier (used for Hub downloads if needed)
|
model_id: The model identifier (used for Hub downloads if needed)
|
||||||
base_path: Local directory path for finding state files (None for Hub-only)
|
base_path: Local directory path for finding state files (None for Hub-only)
|
||||||
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
|
hub_download_kwargs: Parameters for hf_hub_download (tokens, cache, etc.)
|
||||||
|
is_local_source: Whether model_id resolved to a local directory or config file.
|
||||||
|
|
||||||
Note:
|
Note:
|
||||||
This method modifies step_instance in-place and returns None.
|
This method modifies step_instance in-place and returns None.
|
||||||
@@ -1191,6 +1199,12 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
|
|||||||
# Try local file first
|
# Try local file first
|
||||||
if base_path and (base_path / state_filename).exists():
|
if base_path and (base_path / state_filename).exists():
|
||||||
state_path = str(base_path / state_filename)
|
state_path = str(base_path / state_filename)
|
||||||
|
elif is_local_source:
|
||||||
|
state_path = base_path / state_filename if base_path else Path(state_filename)
|
||||||
|
raise FileNotFoundError(
|
||||||
|
f"State file '{state_filename}' was not found for local processor pipeline "
|
||||||
|
f"'{model_id}' at '{state_path}'."
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
# Download from Hub
|
# Download from Hub
|
||||||
state_path = hf_hub_download(
|
state_path = hf_hub_download(
|
||||||
|
|||||||
@@ -323,6 +323,10 @@ class LeKiwiClient(Robot):
|
|||||||
np.ndarray: the action sent to the motors, potentially clipped.
|
np.ndarray: the action sent to the motors, potentially clipped.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
# Action values may be torch tensors (e.g. replayed from a dataset) or numpy
|
||||||
|
# scalars; json.dumps only serializes Python primitives, so coerce each value to a
|
||||||
|
# plain float before sending.
|
||||||
|
action = {key: float(value) for key, value in action.items()}
|
||||||
self.zmq_cmd_socket.send_string(json.dumps(action)) # action is in motor space
|
self.zmq_cmd_socket.send_string(json.dumps(action)) # action is in motor space
|
||||||
|
|
||||||
# TODO(Steven): Remove the np conversion when it is possible to record a non-numpy array value
|
# TODO(Steven): Remove the np conversion when it is possible to record a non-numpy array value
|
||||||
|
|||||||
@@ -326,8 +326,17 @@ class RolloutConfig:
|
|||||||
|
|
||||||
policy_path = parser.get_path_arg("policy")
|
policy_path = parser.get_path_arg("policy")
|
||||||
if policy_path:
|
if policy_path:
|
||||||
cli_overrides = parser.get_cli_overrides("policy")
|
yaml_overrides = parser.get_yaml_overrides("policy")
|
||||||
self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
|
cli_overrides = parser.get_cli_overrides("policy") or []
|
||||||
|
policy_overrides = yaml_overrides + cli_overrides
|
||||||
|
pretrained_revision = parser.parse_arg("pretrained_revision", cli_overrides)
|
||||||
|
if pretrained_revision is None:
|
||||||
|
pretrained_revision = parser.parse_arg("pretrained_revision", yaml_overrides)
|
||||||
|
self.policy = PreTrainedConfig.from_pretrained(
|
||||||
|
policy_path,
|
||||||
|
revision=pretrained_revision,
|
||||||
|
cli_overrides=policy_overrides,
|
||||||
|
)
|
||||||
self.policy.pretrained_path = policy_path
|
self.policy.pretrained_path = policy_path
|
||||||
if self.policy is None:
|
if self.policy is None:
|
||||||
raise ValueError("--policy.path is required for rollout")
|
raise ValueError("--policy.path is required for rollout")
|
||||||
|
|||||||
@@ -27,7 +27,7 @@ from threading import Event
|
|||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from lerobot.configs import FeatureType
|
from lerobot.configs import FeatureType, PreTrainedConfig
|
||||||
from lerobot.datasets import (
|
from lerobot.datasets import (
|
||||||
LeRobotDataset,
|
LeRobotDataset,
|
||||||
aggregate_pipeline_dataset_features,
|
aggregate_pipeline_dataset_features,
|
||||||
@@ -159,6 +159,35 @@ class RolloutContext:
|
|||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy:
|
||||||
|
"""Load policy weights, keeping adapter and base-model revisions independent."""
|
||||||
|
pretrained_revision = policy_config.pretrained_revision
|
||||||
|
policy_class = get_policy_class(policy_config.type)
|
||||||
|
|
||||||
|
if not policy_config.use_peft:
|
||||||
|
return policy_class.from_pretrained(
|
||||||
|
policy_config.pretrained_path,
|
||||||
|
config=policy_config,
|
||||||
|
revision=pretrained_revision,
|
||||||
|
)
|
||||||
|
|
||||||
|
from peft import PeftConfig, PeftModel
|
||||||
|
|
||||||
|
peft_path = policy_config.pretrained_path
|
||||||
|
peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision)
|
||||||
|
policy = policy_class.from_pretrained(
|
||||||
|
pretrained_name_or_path=peft_config.base_model_name_or_path,
|
||||||
|
config=policy_config,
|
||||||
|
revision=peft_config.revision,
|
||||||
|
)
|
||||||
|
return PeftModel.from_pretrained(
|
||||||
|
policy,
|
||||||
|
peft_path,
|
||||||
|
config=peft_config,
|
||||||
|
revision=pretrained_revision,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def build_rollout_context(
|
def build_rollout_context(
|
||||||
cfg: RolloutConfig,
|
cfg: RolloutConfig,
|
||||||
shutdown_event: Event,
|
shutdown_event: Event,
|
||||||
@@ -176,7 +205,6 @@ def build_rollout_context(
|
|||||||
# --- 1. Policy (heavy I/O, but no hardware yet) -------------------
|
# --- 1. Policy (heavy I/O, but no hardware yet) -------------------
|
||||||
logger.info("Loading policy from '%s'...", cfg.policy.pretrained_path)
|
logger.info("Loading policy from '%s'...", cfg.policy.pretrained_path)
|
||||||
policy_config = cfg.policy
|
policy_config = cfg.policy
|
||||||
policy_class = get_policy_class(policy_config.type)
|
|
||||||
|
|
||||||
if hasattr(policy_config, "compile_model"):
|
if hasattr(policy_config, "compile_model"):
|
||||||
policy_config.compile_model = cfg.use_torch_compile
|
policy_config.compile_model = cfg.use_torch_compile
|
||||||
@@ -187,17 +215,7 @@ def build_rollout_context(
|
|||||||
"Please use `cpu` or `cuda` backend."
|
"Please use `cpu` or `cuda` backend."
|
||||||
)
|
)
|
||||||
|
|
||||||
if policy_config.use_peft:
|
policy = _load_pretrained_policy(policy_config)
|
||||||
from peft import PeftConfig, PeftModel
|
|
||||||
|
|
||||||
peft_path = policy_config.pretrained_path
|
|
||||||
peft_config = PeftConfig.from_pretrained(peft_path)
|
|
||||||
policy = policy_class.from_pretrained(
|
|
||||||
pretrained_name_or_path=peft_config.base_model_name_or_path, config=policy_config
|
|
||||||
)
|
|
||||||
policy = PeftModel.from_pretrained(policy, peft_path, config=peft_config)
|
|
||||||
else:
|
|
||||||
policy = policy_class.from_pretrained(policy_config.pretrained_path, config=policy_config)
|
|
||||||
|
|
||||||
if is_rtc:
|
if is_rtc:
|
||||||
policy.config.rtc_config = cfg.inference.rtc
|
policy.config.rtc_config = cfg.inference.rtc
|
||||||
@@ -392,6 +410,7 @@ def build_rollout_context(
|
|||||||
preprocessor, postprocessor = make_pre_post_processors(
|
preprocessor, postprocessor = make_pre_post_processors(
|
||||||
policy_cfg=policy_config,
|
policy_cfg=policy_config,
|
||||||
pretrained_path=cfg.policy.pretrained_path,
|
pretrained_path=cfg.policy.pretrained_path,
|
||||||
|
pretrained_revision=policy_config.pretrained_revision,
|
||||||
dataset_stats=dataset_stats,
|
dataset_stats=dataset_stats,
|
||||||
preprocessor_overrides={
|
preprocessor_overrides={
|
||||||
"device_processor": {"device": cfg.device},
|
"device_processor": {"device": cfg.device},
|
||||||
|
|||||||
@@ -61,6 +61,7 @@ import pyarrow as pa
|
|||||||
import tqdm
|
import tqdm
|
||||||
from datasets import Dataset, Features, Image
|
from datasets import Dataset, Features, Image
|
||||||
from huggingface_hub import HfApi, snapshot_download
|
from huggingface_hub import HfApi, snapshot_download
|
||||||
|
from huggingface_hub.errors import RevisionNotFoundError
|
||||||
from requests import HTTPError
|
from requests import HTTPError
|
||||||
|
|
||||||
from lerobot.datasets import CODEBASE_VERSION, LeRobotDataset, aggregate_stats
|
from lerobot.datasets import CODEBASE_VERSION, LeRobotDataset, aggregate_stats
|
||||||
@@ -521,7 +522,7 @@ def convert_dataset(
|
|||||||
hub_api = HfApi()
|
hub_api = HfApi()
|
||||||
try:
|
try:
|
||||||
hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
|
hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
|
||||||
except HTTPError as e:
|
except (HTTPError, RevisionNotFoundError) as e:
|
||||||
print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
|
print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
|
||||||
pass
|
pass
|
||||||
hub_api.delete_files(
|
hub_api.delete_files(
|
||||||
|
|||||||
@@ -453,6 +453,9 @@ def eval_policy(
|
|||||||
raise exc from None
|
raise exc from None
|
||||||
|
|
||||||
start = time.time()
|
start = time.time()
|
||||||
|
# Preserve the mode for direct callers. eval_policy_all scopes the mode
|
||||||
|
# around all tasks so parallel evaluations cannot race with each other.
|
||||||
|
was_training = policy.training
|
||||||
policy.eval()
|
policy.eval()
|
||||||
|
|
||||||
# Determine how many batched rollouts we need to get n_episodes. Note that if n_episodes is not evenly
|
# Determine how many batched rollouts we need to get n_episodes. Note that if n_episodes is not evenly
|
||||||
@@ -674,6 +677,8 @@ def eval_policy(
|
|||||||
if save_predicted_video:
|
if save_predicted_video:
|
||||||
info["predicted_video_paths"] = predicted_video_paths
|
info["predicted_video_paths"] = predicted_video_paths
|
||||||
|
|
||||||
|
policy.train(was_training)
|
||||||
|
|
||||||
return info
|
return info
|
||||||
|
|
||||||
|
|
||||||
@@ -1010,6 +1015,12 @@ def eval_policy_all(
|
|||||||
recording_private=recording_private,
|
recording_private=recording_private,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Set the shared policy's mode before launching any workers. Restoring it
|
||||||
|
# inside individual tasks would let one task enable training mode while
|
||||||
|
# another task is still evaluating.
|
||||||
|
was_training = policy.training
|
||||||
|
policy.eval()
|
||||||
|
try:
|
||||||
if max_parallel_tasks <= 1:
|
if max_parallel_tasks <= 1:
|
||||||
prefetch_thread: threading.Thread | None = None
|
prefetch_thread: threading.Thread | None = None
|
||||||
for i, (task_group, task_id, env) in enumerate(tasks):
|
for i, (task_group, task_id, env) in enumerate(tasks):
|
||||||
@@ -1044,6 +1055,8 @@ def eval_policy_all(
|
|||||||
per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics})
|
per_task_infos.append({"task_group": tg, "task_id": tid, "metrics": metrics})
|
||||||
finally:
|
finally:
|
||||||
env.close()
|
env.close()
|
||||||
|
finally:
|
||||||
|
policy.train(was_training)
|
||||||
|
|
||||||
# compute aggregated metrics helper (robust to lists/scalars)
|
# compute aggregated metrics helper (robust to lists/scalars)
|
||||||
def _agg_from_list(xs):
|
def _agg_from_list(xs):
|
||||||
|
|||||||
@@ -453,9 +453,11 @@ def record(
|
|||||||
encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize,
|
encoder_queue_maxsize=cfg.dataset.encoder_queue_maxsize,
|
||||||
)
|
)
|
||||||
|
|
||||||
robot.connect()
|
# Connect the teleoperator before the robot so the robot isn't left idle (and possibly
|
||||||
|
# tripping a firmware watchdog) during teleop init. Matches lerobot_teleoperate.py.
|
||||||
if teleop is not None:
|
if teleop is not None:
|
||||||
teleop.connect()
|
teleop.connect()
|
||||||
|
robot.connect()
|
||||||
|
|
||||||
listener, events = init_keyboard_listener()
|
listener, events = init_keyboard_listener()
|
||||||
|
|
||||||
|
|||||||
@@ -61,6 +61,7 @@ from lerobot.robots import ( # noqa: F401
|
|||||||
earthrover_mini_plus,
|
earthrover_mini_plus,
|
||||||
hope_jr,
|
hope_jr,
|
||||||
koch_follower,
|
koch_follower,
|
||||||
|
lekiwi,
|
||||||
make_robot_from_config,
|
make_robot_from_config,
|
||||||
omx_follower,
|
omx_follower,
|
||||||
openarm_follower,
|
openarm_follower,
|
||||||
|
|||||||
@@ -72,6 +72,20 @@ from lerobot.utils.utils import (
|
|||||||
from .lerobot_eval import eval_policy_all
|
from .lerobot_eval import eval_policy_all
|
||||||
|
|
||||||
|
|
||||||
|
def _dataloader_worker_kwargs(
|
||||||
|
cfg: TrainPipelineConfig,
|
||||||
|
*,
|
||||||
|
num_workers: int | None = None,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""Return worker-only DataLoader options, disabling them for single-process loading."""
|
||||||
|
workers_enabled = (cfg.num_workers if num_workers is None else num_workers) > 0
|
||||||
|
return {
|
||||||
|
"prefetch_factor": cfg.prefetch_factor if workers_enabled else None,
|
||||||
|
"persistent_workers": cfg.persistent_workers and workers_enabled,
|
||||||
|
"multiprocessing_context": cfg.dataloader_multiprocessing_context if workers_enabled else None,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
def update_policy(
|
def update_policy(
|
||||||
train_metrics: MetricsTracker,
|
train_metrics: MetricsTracker,
|
||||||
policy: PreTrainedPolicy,
|
policy: PreTrainedPolicy,
|
||||||
@@ -521,13 +535,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
|||||||
pin_memory=device.type == "cuda",
|
pin_memory=device.type == "cuda",
|
||||||
drop_last=False,
|
drop_last=False,
|
||||||
collate_fn=collate_fn,
|
collate_fn=collate_fn,
|
||||||
prefetch_factor=(
|
**_dataloader_worker_kwargs(
|
||||||
cfg.prefetch_factor
|
cfg,
|
||||||
if (train_num_workers if cfg.dataset.streaming else cfg.num_workers) > 0
|
num_workers=train_num_workers if cfg.dataset.streaming else cfg.num_workers,
|
||||||
else None
|
|
||||||
),
|
),
|
||||||
persistent_workers=cfg.persistent_workers
|
|
||||||
and (train_num_workers if cfg.dataset.streaming else cfg.num_workers) > 0,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
# Build eval dataloader if a held-out split exists
|
# Build eval dataloader if a held-out split exists
|
||||||
@@ -553,8 +564,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
|||||||
pin_memory=device.type == "cuda",
|
pin_memory=device.type == "cuda",
|
||||||
drop_last=False,
|
drop_last=False,
|
||||||
collate_fn=eval_collate_fn,
|
collate_fn=eval_collate_fn,
|
||||||
prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
|
**_dataloader_worker_kwargs(cfg),
|
||||||
persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
|
|
||||||
)
|
)
|
||||||
|
|
||||||
# Prepare everything with accelerator
|
# Prepare everything with accelerator
|
||||||
|
|||||||
@@ -204,6 +204,38 @@ def test_clear_resets_buffer(tmp_path):
|
|||||||
assert dataset.writer.episode_buffer["size"] == 0
|
assert dataset.writer.episode_buffer["size"] == 0
|
||||||
|
|
||||||
|
|
||||||
|
def test_clear_removes_video_frame_staging_dir(tmp_path):
|
||||||
|
"""clear_episode_buffer() removes PNG staging dirs for video features."""
|
||||||
|
video_key = "observation.images.cam"
|
||||||
|
features = {
|
||||||
|
video_key: {
|
||||||
|
"dtype": "video",
|
||||||
|
"shape": (64, 96, 3),
|
||||||
|
"names": ["height", "width", "channels"],
|
||||||
|
},
|
||||||
|
"action": {"dtype": "float32", "shape": (2,), "names": None},
|
||||||
|
}
|
||||||
|
dataset = LeRobotDataset.create(
|
||||||
|
repo_id=DUMMY_REPO_ID,
|
||||||
|
fps=DEFAULT_FPS,
|
||||||
|
features=features,
|
||||||
|
root=tmp_path / "ds",
|
||||||
|
use_videos=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
dataset.add_frame(_make_frame(features))
|
||||||
|
video_staging_dir = (
|
||||||
|
dataset.root
|
||||||
|
/ Path(DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)).parent
|
||||||
|
)
|
||||||
|
assert video_staging_dir.is_dir()
|
||||||
|
|
||||||
|
dataset.clear_episode_buffer()
|
||||||
|
|
||||||
|
assert dataset.writer.episode_buffer["size"] == 0
|
||||||
|
assert not video_staging_dir.exists()
|
||||||
|
|
||||||
|
|
||||||
def test_finalize_is_idempotent(tmp_path):
|
def test_finalize_is_idempotent(tmp_path):
|
||||||
"""Calling finalize() twice does not raise."""
|
"""Calling finalize() twice does not raise."""
|
||||||
dataset = LeRobotDataset.create(
|
dataset = LeRobotDataset.create(
|
||||||
|
|||||||
@@ -26,8 +26,17 @@ import tempfile
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
import torch
|
||||||
|
from safetensors.torch import save_file
|
||||||
|
|
||||||
from lerobot.processor.pipeline import DataProcessorPipeline, ProcessorMigrationError
|
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||||
|
from lerobot.processor.pipeline import (
|
||||||
|
DataProcessorPipeline,
|
||||||
|
ProcessorMigrationError,
|
||||||
|
ProcessorStep,
|
||||||
|
ProcessorStepRegistry,
|
||||||
|
)
|
||||||
|
from lerobot.types import EnvTransition
|
||||||
|
|
||||||
# Simplified Config Loading Tests
|
# Simplified Config Loading Tests
|
||||||
|
|
||||||
@@ -98,6 +107,140 @@ def test_load_config_nonexistent_path_tries_hub():
|
|||||||
DataProcessorPipeline._load_config("nonexistent/path", "processor.json", {})
|
DataProcessorPipeline._load_config("nonexistent/path", "processor.json", {})
|
||||||
|
|
||||||
|
|
||||||
|
def test_from_pretrained_local_directory_missing_state_does_not_call_hub(monkeypatch):
|
||||||
|
"""Local processor dirs must fail locally when a state file is missing."""
|
||||||
|
|
||||||
|
@ProcessorStepRegistry.register("local_missing_state_step")
|
||||||
|
class LocalMissingStateStep(ProcessorStep):
|
||||||
|
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||||
|
return transition
|
||||||
|
|
||||||
|
def transform_features(
|
||||||
|
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||||
|
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||||
|
return features
|
||||||
|
|
||||||
|
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||||||
|
pass
|
||||||
|
|
||||||
|
try:
|
||||||
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||||
|
tmp_path = Path(tmp_dir)
|
||||||
|
config = {
|
||||||
|
"name": "LocalMissingStatePipeline",
|
||||||
|
"steps": [{"registry_name": "local_missing_state_step", "state_file": "missing.safetensors"}],
|
||||||
|
}
|
||||||
|
(tmp_path / "processor.json").write_text(json.dumps(config))
|
||||||
|
|
||||||
|
def fail_hub_download(*args, **kwargs):
|
||||||
|
pytest.fail("local missing processor state should not call hf_hub_download")
|
||||||
|
|
||||||
|
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fail_hub_download)
|
||||||
|
|
||||||
|
with pytest.raises(FileNotFoundError, match="missing.safetensors.*local processor pipeline"):
|
||||||
|
DataProcessorPipeline.from_pretrained(tmp_path, config_filename="processor.json")
|
||||||
|
finally:
|
||||||
|
ProcessorStepRegistry.unregister("local_missing_state_step")
|
||||||
|
|
||||||
|
|
||||||
|
def test_from_pretrained_local_config_file_missing_state_does_not_call_hub(monkeypatch):
|
||||||
|
"""Local single-file processor configs must also keep missing state resolution local."""
|
||||||
|
|
||||||
|
@ProcessorStepRegistry.register("local_file_missing_state_step")
|
||||||
|
class LocalFileMissingStateStep(ProcessorStep):
|
||||||
|
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||||
|
return transition
|
||||||
|
|
||||||
|
def transform_features(
|
||||||
|
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||||
|
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||||
|
return features
|
||||||
|
|
||||||
|
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||||||
|
pass
|
||||||
|
|
||||||
|
try:
|
||||||
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||||
|
tmp_path = Path(tmp_dir)
|
||||||
|
config_path = tmp_path / "processor.json"
|
||||||
|
config = {
|
||||||
|
"name": "LocalFileMissingStatePipeline",
|
||||||
|
"steps": [
|
||||||
|
{"registry_name": "local_file_missing_state_step", "state_file": "missing.safetensors"}
|
||||||
|
],
|
||||||
|
}
|
||||||
|
config_path.write_text(json.dumps(config))
|
||||||
|
|
||||||
|
def fail_hub_download(*args, **kwargs):
|
||||||
|
pytest.fail("local missing processor state should not call hf_hub_download")
|
||||||
|
|
||||||
|
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fail_hub_download)
|
||||||
|
|
||||||
|
with pytest.raises(FileNotFoundError, match="missing.safetensors.*local processor pipeline"):
|
||||||
|
DataProcessorPipeline.from_pretrained(config_path, config_filename="ignored.json")
|
||||||
|
finally:
|
||||||
|
ProcessorStepRegistry.unregister("local_file_missing_state_step")
|
||||||
|
|
||||||
|
|
||||||
|
def test_from_pretrained_hub_source_missing_local_state_still_calls_hub(monkeypatch, tmp_path):
|
||||||
|
"""Hub sources still fall back to hf_hub_download for state files."""
|
||||||
|
|
||||||
|
@ProcessorStepRegistry.register("hub_state_step")
|
||||||
|
class HubStateStep(ProcessorStep):
|
||||||
|
def __init__(self):
|
||||||
|
self.value = torch.tensor(0)
|
||||||
|
|
||||||
|
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||||
|
return transition
|
||||||
|
|
||||||
|
def transform_features(
|
||||||
|
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||||
|
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||||
|
return features
|
||||||
|
|
||||||
|
def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
|
||||||
|
self.value = state["value"]
|
||||||
|
|
||||||
|
try:
|
||||||
|
state_path = tmp_path / "downloaded.safetensors"
|
||||||
|
save_file({"value": torch.tensor(7)}, state_path)
|
||||||
|
loaded_config = {
|
||||||
|
"name": "HubStatePipeline",
|
||||||
|
"steps": [{"registry_name": "hub_state_step", "state_file": "hub_state.safetensors"}],
|
||||||
|
}
|
||||||
|
calls = []
|
||||||
|
|
||||||
|
def fake_load_config(cls, model_id, config_filename, hub_download_kwargs):
|
||||||
|
return loaded_config, tmp_path / "hub_cache"
|
||||||
|
|
||||||
|
def fake_hub_download(**kwargs):
|
||||||
|
calls.append(kwargs)
|
||||||
|
return str(state_path)
|
||||||
|
|
||||||
|
monkeypatch.setattr(DataProcessorPipeline, "_load_config", classmethod(fake_load_config))
|
||||||
|
monkeypatch.setattr("lerobot.processor.pipeline.hf_hub_download", fake_hub_download)
|
||||||
|
|
||||||
|
pipeline = DataProcessorPipeline.from_pretrained("user/repo", config_filename="processor.json")
|
||||||
|
|
||||||
|
assert calls == [
|
||||||
|
{
|
||||||
|
"repo_id": "user/repo",
|
||||||
|
"filename": "hub_state.safetensors",
|
||||||
|
"repo_type": "model",
|
||||||
|
"force_download": False,
|
||||||
|
"resume_download": None,
|
||||||
|
"proxies": None,
|
||||||
|
"token": None,
|
||||||
|
"cache_dir": None,
|
||||||
|
"local_files_only": False,
|
||||||
|
"revision": None,
|
||||||
|
}
|
||||||
|
]
|
||||||
|
assert pipeline.steps[0].value.item() == 7
|
||||||
|
finally:
|
||||||
|
ProcessorStepRegistry.unregister("hub_state_step")
|
||||||
|
|
||||||
|
|
||||||
# Config Validation Tests
|
# Config Validation Tests
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -17,6 +17,8 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import dataclasses
|
import dataclasses
|
||||||
|
import sys
|
||||||
|
from types import SimpleNamespace
|
||||||
from unittest.mock import MagicMock
|
from unittest.mock import MagicMock
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
@@ -106,6 +108,109 @@ def test_sentry_config_defaults():
|
|||||||
assert cfg.target_video_file_size_mb is None
|
assert cfg.target_video_file_size_mb is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_rollout_config_passes_policy_pretrained_revision(monkeypatch):
|
||||||
|
from lerobot.configs import PreTrainedConfig, parser
|
||||||
|
from lerobot.rollout import RolloutConfig
|
||||||
|
from tests.mocks.mock_robot import MockRobotConfig
|
||||||
|
|
||||||
|
captured = {}
|
||||||
|
|
||||||
|
def fake_from_pretrained(cls, pretrained_name_or_path, **kwargs):
|
||||||
|
captured["pretrained_name_or_path"] = pretrained_name_or_path
|
||||||
|
captured.update(kwargs)
|
||||||
|
return SimpleNamespace(device="cpu", pretrained_revision=kwargs["revision"])
|
||||||
|
|
||||||
|
monkeypatch.setattr(parser, "get_yaml_overrides", lambda _: ["--pretrained_revision=yaml-sha"])
|
||||||
|
monkeypatch.setattr(
|
||||||
|
sys,
|
||||||
|
"argv",
|
||||||
|
["lerobot-rollout", "--policy.path=user/policy", "--policy.pretrained_revision=cli-sha"],
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(PreTrainedConfig, "from_pretrained", classmethod(fake_from_pretrained))
|
||||||
|
|
||||||
|
cfg = RolloutConfig(robot=MockRobotConfig())
|
||||||
|
|
||||||
|
assert captured["pretrained_name_or_path"] == "user/policy"
|
||||||
|
assert captured["revision"] == "cli-sha"
|
||||||
|
assert captured["cli_overrides"] == [
|
||||||
|
"--pretrained_revision=yaml-sha",
|
||||||
|
"--pretrained_revision=cli-sha",
|
||||||
|
]
|
||||||
|
assert cfg.policy.pretrained_path == "user/policy"
|
||||||
|
assert cfg.policy.pretrained_revision == "cli-sha"
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_pretrained_policy_passes_revision(monkeypatch):
|
||||||
|
import lerobot.rollout.context as rollout_context
|
||||||
|
|
||||||
|
policy_config = SimpleNamespace(
|
||||||
|
type="mock",
|
||||||
|
use_peft=False,
|
||||||
|
pretrained_path="user/policy",
|
||||||
|
pretrained_revision="policy-sha",
|
||||||
|
)
|
||||||
|
policy_class = MagicMock()
|
||||||
|
loaded_policy = MagicMock()
|
||||||
|
policy_class.from_pretrained.return_value = loaded_policy
|
||||||
|
monkeypatch.setattr(rollout_context, "get_policy_class", lambda _: policy_class)
|
||||||
|
|
||||||
|
policy = rollout_context._load_pretrained_policy(policy_config)
|
||||||
|
|
||||||
|
assert policy is loaded_policy
|
||||||
|
policy_class.from_pretrained.assert_called_once_with(
|
||||||
|
"user/policy",
|
||||||
|
config=policy_config,
|
||||||
|
revision="policy-sha",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_pretrained_peft_policy_keeps_adapter_and_base_revisions_separate(monkeypatch):
|
||||||
|
import lerobot.rollout.context as rollout_context
|
||||||
|
|
||||||
|
policy_config = SimpleNamespace(
|
||||||
|
type="mock",
|
||||||
|
use_peft=True,
|
||||||
|
pretrained_path="user/adapter",
|
||||||
|
pretrained_revision="adapter-sha",
|
||||||
|
)
|
||||||
|
policy_class = MagicMock()
|
||||||
|
base_policy = MagicMock()
|
||||||
|
policy_class.from_pretrained.return_value = base_policy
|
||||||
|
monkeypatch.setattr(rollout_context, "get_policy_class", lambda _: policy_class)
|
||||||
|
|
||||||
|
peft_config = SimpleNamespace(
|
||||||
|
base_model_name_or_path="user/base-policy",
|
||||||
|
revision="base-sha",
|
||||||
|
)
|
||||||
|
peft_config_from_pretrained = MagicMock(return_value=peft_config)
|
||||||
|
adapted_policy = MagicMock()
|
||||||
|
peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
|
||||||
|
monkeypatch.setitem(
|
||||||
|
sys.modules,
|
||||||
|
"peft",
|
||||||
|
SimpleNamespace(
|
||||||
|
PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained),
|
||||||
|
PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained),
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
policy = rollout_context._load_pretrained_policy(policy_config)
|
||||||
|
|
||||||
|
assert policy is adapted_policy
|
||||||
|
peft_config_from_pretrained.assert_called_once_with("user/adapter", revision="adapter-sha")
|
||||||
|
policy_class.from_pretrained.assert_called_once_with(
|
||||||
|
pretrained_name_or_path="user/base-policy",
|
||||||
|
config=policy_config,
|
||||||
|
revision="base-sha",
|
||||||
|
)
|
||||||
|
peft_model_from_pretrained.assert_called_once_with(
|
||||||
|
base_policy,
|
||||||
|
"user/adapter",
|
||||||
|
config=peft_config,
|
||||||
|
revision="adapter-sha",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# RolloutRingBuffer
|
# RolloutRingBuffer
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|||||||
Reference in New Issue
Block a user