Compare commits

..

2 Commits

Author SHA1 Message Date
Steven Palma 0371e99117 feat(policiy): fine tune vision encoder refactor completness 2026-07-29 18:35:24 +02:00
Functionhx 9e30807eeb feat: add fine_tune_vision_encoder flag to SmolVLA
Fixes #1774
2026-07-29 18:22:01 +02:00
47 changed files with 223 additions and 554 deletions
+10 -13
View File
@@ -61,20 +61,16 @@ Full details in [`docs/source/so101.mdx`](./docs/source/so101.mdx) and [`docs/so
**4.1 Install**
```bash
# uv (recommended — see AGENTS.md and CLAUDE.md)
uv sync --locked --extra feetech # SO-100/SO-101 motor stack
# uv sync --locked --extra all # everything
# uv sync --locked --extra smolvla # add SmolVLA deps
# pip (alternative, e.g. when not working from source)
# pip install 'lerobot[feetech]'
# pip install 'lerobot[all]'
# pip install 'lerobot[smolvla]'
pip install 'lerobot[feetech]' # SO-100/SO-101 motor stack
# pip install 'lerobot[all]' # everything
# pip install 'lerobot[aloha,pusht]' # specific features
# pip install 'lerobot[smolvla]' # add SmolVLA deps
git lfs install && git lfs pull
hf auth login # required to push datasets/policies
hf auth login # required to push datasets/policies
```
Contributors can alternatively use `uv sync --locked --extra feetech` (see `AGENTS.md`).
**4.2 Find USB ports** — run once per arm, unplug when prompted.
```bash
@@ -325,10 +321,11 @@ SmolVLA ships with `freeze_vision_encoder=True`. Unfreezing usually **improves p
```bash
lerobot-train ... --policy.type=smolvla \
--policy.freeze_vision_encoder=false \
--policy.train_expert_only=false
--policy.fine_tune_vision_encoder=true
```
This selectively trains the vision encoder and connector while leaving the language model frozen. Their learning rate defaults to `0.1 × optimizer_lr`; adjust it with `--policy.vision_encoder_lr_multiplier` if needed.
### 7.7 Signals to stop / keep going
- Train loss plateaus → stop, save a Hub checkpoint.
+15 -2
View File
@@ -88,6 +88,20 @@ policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)
The same policy can work with different environment processors, and the same environment processor can work with different policies:
````python
# Use SmolVLA policy with LIBERO environment
# Use SmolVLA policy with LIBERO environment
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
env_cfg=libero_cfg,
policy_cfg=smolvla_cfg,
)
smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
# Or use ACT policy with the same LIBERO environment
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
env_cfg=libero_cfg,
policy_cfg=act_cfg,
)
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
```python
# Use SmolVLA policy with LIBERO environment
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
@@ -102,7 +116,6 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
policy_cfg=act_cfg,
)
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
```
### 3. **Easier Experimentation**
@@ -132,7 +145,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
state = torch.cat([eef_pos, eef_axisangle, eef_vel,
gripper_pos, gripper_vel], dim=-1) # 14D
return state
```
````
### 4. **Cleaner Environment Code**
-1
View File
@@ -59,7 +59,6 @@ The `lerobot-rollout --strategy.type=dagger` mode requires **teleoperators with
- `bi_openarm_mini` - Bimanual OpenArm Mini
- `so_leader` - SO100 / SO101 leader arm
- `bi_so_leader` - Bimanual SO100 / SO101 leader arms
> [!IMPORTANT]
> The provided commands default to `bi_openarm_follower` + `bi_openarm_mini`.
+1 -1
View File
@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.
### Record
This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data).
This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
```bash
lerobot-record \
+1 -1
View File
@@ -18,7 +18,7 @@ If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus inter
- [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) for controlling Dynamixel servos
Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so_follower.py)
For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so101_follower/so101_follower.py)
Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):
+1 -1
View File
@@ -51,7 +51,7 @@ In addition to these instructions, you need to install the Feetech SDK & ZeroMQ
pip install -e ".[lekiwi]"
```
Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖.
Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:.
Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
# Step-by-Step Assembly Instructions
+1 -1
View File
@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs
## Reproducing π₀Fast results
We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40kk steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
The finetuned model can be found here:
+14 -1
View File
@@ -70,6 +70,19 @@ cd lerobot && lerobot-train \
GPU allows it, as long as loading times remain short.
</Tip>
For tasks that require adapting visual features, such as distinguishing new colors or shapes, selectively
fine-tune the vision encoder and its connector:
```bash
lerobot-train ... \
--policy.path=lerobot/smolvla_base \
--policy.fine_tune_vision_encoder=true
```
This keeps the language model frozen with the default `train_expert_only=true` setting and trains the vision
path at `0.1` times the main learning rate by default. Fine-tuning the vision encoder increases memory use and
can reduce the model's general visual knowledge, so enable it only when the frozen encoder is insufficient.
Fine-tuning is an art. For a complete overview of the options for finetuning, run
```bash
@@ -93,7 +106,7 @@ lerobot-train --help
## Evaluate the finetuned model and run it in real-time
Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots#record-a-dataset).
Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
Once you are logged in, you can run inference in your setup by doing:
```bash
+1 -1
View File
@@ -338,7 +338,7 @@ It is advisable to install one 3-pin cable in the motor after placing them befor
<hfoption id="Leader">
- Mount the leader holder onto the wrist and secure it with 4 M3x6mm screws.
- Attach the handle to the leader holder using 1 M2x6mm screw.
- Attach the handle to motor 5 using 1 M2x6mm screw.
- Insert the gripper motor, secure it with 2 M2x6mm screws on each side, attach a motor horn using a M3x6mm horn screw.
- Attach the follower trigger with 4 M3x6mm screws.
+3 -51
View File
@@ -11,10 +11,9 @@ LeRobot provides several utilities for manipulating datasets:
3. **Merge Datasets** - Combine multiple datasets into one. The datasets must have identical features, and episodes are concatenated in the order specified in `repo_ids`
4. **Add Features** - Add new features to a dataset
5. **Remove Features** - Remove features from a dataset
6. **Modify Tasks** - Change the natural-language task descriptions associated with episodes
7. **Convert to Video** - Convert image-based datasets to video format for efficient storage (RGB and depth cameras are encoded with separate encoders)
8. **Re-encode Videos** - Re-encode an existing video dataset's RGB and/or depth streams with new encoder settings
9. **Show the Info of Datasets** - Show the summary of datasets information such as number of episode etc.
6. **Convert to Video** - Convert image-based datasets to video format for efficient storage (RGB and depth cameras are encoded with separate encoders)
7. **Re-encode Videos** - Re-encode an existing video dataset's RGB and/or depth streams with new encoder settings
8. **Show the Info of Datasets** - Show the summary of datasets information such as number of episode etc.
The core implementation is in `lerobot.datasets.dataset_tools`.
An example script detailing how to use the tools API is available in `examples/dataset/use_dataset_tools.py`.
@@ -90,53 +89,6 @@ lerobot-edit-dataset \
--operation.feature_names "['observation.images.top']"
```
#### Modify Tasks
Change the natural-language task descriptions attached to episodes. This is useful for fixing typos, standardizing wording, or re-labeling episodes.
> [!WARNING]
> `modify_tasks` modifies the dataset **in-place** (updating `meta/tasks.parquet`, the `task_index` column in the data files, the `tasks` column in the episode metadata, and `total_tasks` in `meta/info.json`). The `--new_repo_id` and `--new_root` parameters are ignored for this operation.
```bash
# Set a single task for all episodes
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.new_task "Pick up the cube and place it"
# Set different tasks for specific episodes
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.episode_tasks '{"0": "Task A", "1": "Task B", "2": "Task A"}'
# Replace existing task strings wherever they appear
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}'
# Combine modes in a single run
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.new_task "Default task" \
--operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}' \
--operation.episode_tasks '{"5": "Special task for episode 5"}'
```
**Parameters:**
- `new_task`: A single task string used as the default for episodes not otherwise covered.
- `episode_tasks`: Mapping from episode index to task string.
- `task_replacements`: Mapping from existing task strings to their replacements, applied to episodes whose current task matches a key. Every key must be an existing task in the dataset.
The modes can be combined in a single run. Per episode, the task is resolved with the following precedence:
`episode_tasks` > `task_replacements` > `new_task` > original task
At least one of `new_task`, `episode_tasks`, or `task_replacements` must be specified. An episode that ends up with no task raises an error.
#### Convert to Video
Convert an image-based dataset to video format, creating a new LeRobotDataset where images are stored as videos. This is useful for reducing storage requirements and improving data loading performance. The new dataset will have the exact same structure as the original, but with images encoded as MP4 videos in the proper LeRobot format.
+1 -1
View File
@@ -67,7 +67,7 @@ dependencies = [
"einops>=0.8.0,<0.9.0",
# Config & Hub
"draccus>=0.11.6,<0.12.0",
"draccus==0.10.0", # TODO: Relax version constraint
"huggingface-hub>=1.0.0,<2.0.0",
"requests>=2.32.0,<3.0.0",
+2 -4
View File
@@ -163,10 +163,8 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
return None
def _save_pretrained(self, save_directory: Path) -> None:
# Encode against the base class so draccus includes the choice "type" key,
# which `from_pretrained` needs to resolve the concrete subclass.
with open(save_directory / CONFIG_NAME, "w") as f:
json.dump(draccus.encode(self, PreTrainedConfig), f, indent=4)
with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
draccus.dump(self, f, indent=4)
@classmethod
def from_pretrained(
+2 -4
View File
@@ -103,10 +103,8 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
pass
def _save_pretrained(self, save_directory: Path) -> None:
# Encode against the base class so draccus includes the choice "type" key,
# which `from_pretrained` needs to resolve the concrete subclass.
with open(save_directory / CONFIG_NAME, "w") as f:
json.dump(draccus.encode(self, RewardModelConfig), f, indent=4)
with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
draccus.dump(self, f, indent=4)
@classmethod
def from_pretrained(
+1 -5
View File
@@ -194,11 +194,7 @@ class TrainPipelineConfig(HubMixin):
)
if Path(config_path).resolve().exists():
# `config_path` may point at the checkpoint's train_config.json or at its
# pretrained_model/ directory (both documented above) — resolve either to
# the pretrained_model/ directory.
config_path_obj = Path(config_path)
policy_dir = config_path_obj.parent if config_path_obj.is_file() else config_path_obj
policy_dir = Path(config_path).parent
self.checkpoint_path = policy_dir.parent
elif self.job.is_remote:
return
+15 -37
View File
@@ -1435,18 +1435,15 @@ def modify_tasks(
dataset: LeRobotDataset,
new_task: str | None = None,
episode_tasks: dict[int, str] | None = None,
task_replacements: dict[str, str] | None = None,
) -> LeRobotDataset:
"""Modify tasks in a LeRobotDataset.
This function allows you to either:
1. Set a single task for the entire dataset (using `new_task`)
2. Set specific tasks for specific episodes (using `episode_tasks`)
3. Replace existing task strings wherever they appear (using `task_replacements`)
Per episode, the task is resolved with precedence:
`episode_tasks` > `task_replacements` > `new_task` > original task. An episode that ends
up with no task (none of the above apply and it had no original task) raises an error.
You can combine both: `new_task` sets the default, and `episode_tasks` overrides
specific episodes.
The dataset is modified in-place, updating only the task-related files:
- meta/tasks.parquet
@@ -1456,14 +1453,11 @@ def modify_tasks(
Args:
dataset: The source LeRobotDataset to modify.
new_task: Default task applied to any episode not covered by `episode_tasks` or a
matching `task_replacements` entry.
episode_tasks: Optional dict mapping episode indices to task strings. Takes precedence
over both `task_replacements` and `new_task`.
task_replacements: Optional dict mapping existing task strings to new ones. Applied to
episodes whose current task matches a key. Every key must be an existing task.
new_task: A single task string to apply to all episodes. If None and episode_tasks
is also None, raises an error.
episode_tasks: Optional dict mapping episode indices to their task strings.
Overrides `new_task` for specific episodes.
At least one of `new_task`, `episode_tasks`, or `task_replacements` must be provided.
Examples:
Set a single task for all episodes:
@@ -1481,17 +1475,11 @@ def modify_tasks(
new_task="Default task",
episode_tasks={5: "Special task for episode 5"}
)
Replace existing task strings in-place:
dataset = modify_tasks(
dataset,
task_replacements={"Pick up the cube": "Lift the cube"}
)
"""
if not new_task and not episode_tasks and not task_replacements:
raise ValueError("Must specify at least one of new_task, episode_tasks, or task_replacements")
if new_task is None and episode_tasks is None:
raise ValueError("Must specify at least one of new_task or episode_tasks")
if episode_tasks:
if episode_tasks is not None:
valid_indices = set(range(dataset.meta.total_episodes))
invalid = set(episode_tasks.keys()) - valid_indices
if invalid:
@@ -1501,29 +1489,19 @@ def modify_tasks(
if dataset.meta.episodes is None:
dataset.meta.episodes = load_episodes(dataset.root)
if task_replacements:
current_tasks = set(dataset.meta.tasks.index)
invalid_tasks = set(task_replacements) - current_tasks
if invalid_tasks:
raise ValueError(f"Task replacements reference unknown tasks: {sorted(invalid_tasks)}")
# Build the mapping from episode index to task string
episode_to_task: dict[int, str] = {}
for ep_idx in range(dataset.meta.total_episodes):
original_tasks = dataset.meta.episodes[ep_idx]["tasks"]
original_task = original_tasks[0] if original_tasks else None
if episode_tasks and ep_idx in episode_tasks:
episode_to_task[ep_idx] = episode_tasks[ep_idx]
elif task_replacements and original_task in task_replacements:
episode_to_task[ep_idx] = task_replacements[original_task]
elif new_task:
elif new_task is not None:
episode_to_task[ep_idx] = new_task
elif original_task:
# Keep original task if not overridden and no default provided
episode_to_task[ep_idx] = original_task
else:
raise ValueError(f"Episode {ep_idx} has no task; provide new_task or episode_tasks")
# Keep original task if not overridden and no default provided
original_tasks = dataset.meta.episodes[ep_idx]["tasks"]
if not original_tasks:
raise ValueError(f"Episode {ep_idx} has no tasks and no default task was provided")
episode_to_task[ep_idx] = original_tasks[0]
# Collect all unique tasks and create new task mapping
unique_tasks = sorted(set(episode_to_task.values()))
@@ -42,9 +42,6 @@ class Evo1Policy(PreTrainedPolicy):
config_class = Evo1Config
name = "evo1"
def supports_rtc(self) -> bool:
return True
def __init__(self, config: Evo1Config, *, vlm_hub_kwargs: dict | None = None, **kwargs):
super().__init__(config)
config.validate_features()
@@ -68,9 +68,6 @@ class GrootPolicy(PreTrainedPolicy):
name = "groot"
config_class = GrootConfig
def supports_rtc(self) -> bool:
return True
def __init__(self, config: GrootConfig, **kwargs):
"""Initialize Groot policy wrapper."""
require_package("transformers", extra="groot")
@@ -520,9 +520,6 @@ class MolmoAct2Policy(PreTrainedPolicy):
config_class = MolmoAct2Config
name = "molmoact2"
def supports_rtc(self) -> bool:
return self.config.inference_action_mode == "continuous"
def __init__(
self,
config: MolmoAct2Config,
-3
View File
@@ -749,9 +749,6 @@ class PI0Policy(PreTrainedPolicy):
config_class = PI0Config
name = "pi0"
def supports_rtc(self) -> bool:
return True
def __init__(
self,
config: PI0Config,
@@ -714,9 +714,6 @@ class PI05Policy(PreTrainedPolicy):
config_class = PI05Config
name = "pi05"
def supports_rtc(self) -> bool:
return True
def __init__(
self,
config: PI05Config,
-4
View File
@@ -249,10 +249,6 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
"""
raise NotImplementedError
def supports_rtc(self) -> bool:
"""Whether this policy implements Real-Time Chunking inference semantics."""
return False
# TODO(aliberts, rcadene): split into 'forward' and 'compute_loss'?
@abc.abstractmethod
def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
@@ -67,6 +67,8 @@ class SmolVLAConfig(PreTrainedConfig):
# Finetuning settings
freeze_vision_encoder: bool = True
fine_tune_vision_encoder: bool = False # Fine-tune vision + connector; takes priority over freezing.
vision_encoder_lr_multiplier: float = 0.1
train_expert_only: bool = True
train_state_proj: bool = True
@@ -110,6 +112,12 @@ class SmolVLAConfig(PreTrainedConfig):
super().__post_init__()
"""Input validation (not exhaustive)."""
if self.fine_tune_vision_encoder:
self.freeze_vision_encoder = False
if self.vision_encoder_lr_multiplier <= 0:
raise ValueError(
f"`vision_encoder_lr_multiplier` must be positive, got {self.vision_encoder_lr_multiplier}."
)
if self.n_action_steps > self.chunk_size:
raise ValueError(
f"The chunk size is the upper bound for the number of action steps per model invocation. Got "
@@ -145,9 +145,6 @@ class SmolVLAPolicy(PreTrainedPolicy):
config_class = SmolVLAConfig
name = "smolvla"
def supports_rtc(self) -> bool:
return True
def __init__(
self,
config: SmolVLAConfig,
@@ -189,8 +186,27 @@ class SmolVLAPolicy(PreTrainedPolicy):
if model_value is not None:
model_value.rtc_processor = self.rtc_processor
def get_optim_params(self) -> dict:
return self.parameters()
def get_optim_params(self):
if not self.config.fine_tune_vision_encoder:
return self.parameters()
vision_params = []
other_params = []
for name, param in self.named_parameters():
if not param.requires_grad:
continue
if ".vision_model." in name or ".connector." in name:
vision_params.append(param)
else:
other_params.append(param)
return [
{"params": other_params},
{
"params": vision_params,
"lr": self.config.optimizer_lr * self.config.vision_encoder_lr_multiplier,
},
]
def _get_action_chunk(
self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs: Unpack[ActionSelectKwargs]
@@ -496,6 +512,7 @@ class VLAFlowMatching(nn.Module):
self.vlm_with_expert = SmolVLMWithExpertModel(
model_id=self.config.vlm_model_name,
freeze_vision_encoder=self.config.freeze_vision_encoder,
fine_tune_vision_encoder=self.config.fine_tune_vision_encoder,
train_expert_only=self.config.train_expert_only,
load_vlm_weights=self.config.load_vlm_weights,
attention_mode=self.config.attention_mode,
@@ -78,6 +78,7 @@ class SmolVLMWithExpertModel(nn.Module):
load_vlm_weights: bool = True,
train_expert_only: bool = True,
freeze_vision_encoder: bool = False,
fine_tune_vision_encoder: bool = False,
attention_mode: str = "self_attn",
num_expert_layers: int = -1,
num_vlm_layers: int = -1,
@@ -141,6 +142,7 @@ class SmolVLMWithExpertModel(nn.Module):
self.num_key_value_heads = self.config.text_config.num_key_value_heads
self.freeze_vision_encoder = freeze_vision_encoder
self.fine_tune_vision_encoder = fine_tune_vision_encoder
self.train_expert_only = train_expert_only
self.attention_mode = attention_mode
self.expert_hidden_size = lm_expert_config.hidden_size
@@ -150,10 +152,6 @@ class SmolVLMWithExpertModel(nn.Module):
return self.vlm.model
def set_requires_grad(self):
if self.freeze_vision_encoder:
self.get_vlm_model().vision_model.eval()
for params in self.get_vlm_model().vision_model.parameters():
params.requires_grad = False
if self.train_expert_only:
self.vlm.eval()
for params in self.vlm.parameters():
@@ -168,23 +166,26 @@ class SmolVLMWithExpertModel(nn.Module):
last_layers.append(self.num_vlm_layers - 2)
frozen_layers = [
"lm_head",
"text_model.norm.weight",
"text_model.model.norm.weight",
]
for layer in last_layers:
frozen_layers.append(f"text_model.layers.{layer}.")
frozen_layers.append(f"text_model.model.layers.{layer}.")
unmatched_patterns = set(frozen_layers)
for name, params in self.vlm.named_parameters():
matched_patterns = [k for k in frozen_layers if k in name]
if matched_patterns:
if any(k in name for k in frozen_layers):
params.requires_grad = False
unmatched_patterns.difference_update(matched_patterns)
if unmatched_patterns:
raise RuntimeError(
"Some frozen layer patterns matched no VLM parameters, so the corresponding layers "
"would silently remain trainable (parameter naming may have changed in transformers): "
f"{sorted(unmatched_patterns)}"
)
if self.freeze_vision_encoder:
self.get_vlm_model().vision_model.eval()
for params in self.get_vlm_model().vision_model.parameters():
params.requires_grad = False
if self.fine_tune_vision_encoder:
for params in self.get_vlm_model().vision_model.parameters():
params.requires_grad = True
for params in self.get_vlm_model().connector.parameters():
params.requires_grad = True
# To avoid unused params issue with distributed training
for name, params in self.lm_expert.named_parameters():
if "lm_head" in name:
@@ -193,11 +194,15 @@ class SmolVLMWithExpertModel(nn.Module):
def train(self, mode: bool = True):
super().train(mode)
if self.train_expert_only:
self.vlm.eval()
if self.freeze_vision_encoder:
self.get_vlm_model().vision_model.eval()
if self.train_expert_only:
self.vlm.eval()
if self.fine_tune_vision_encoder:
self.get_vlm_model().vision_model.train(mode)
self.get_vlm_model().connector.train(mode)
def embed_image(self, image: torch.Tensor):
patch_attention_mask = None
+2 -5
View File
@@ -16,7 +16,6 @@ from __future__ import annotations
import abc
import builtins
import json
import logging
import os
from dataclasses import dataclass, field
@@ -79,10 +78,8 @@ class RLAlgorithmConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
def _save_pretrained(self, save_directory: Path) -> None:
"""Serialize this config as ``config.json`` inside ``save_directory``."""
# Encode against the base class so draccus includes the choice "type" key,
# which `from_pretrained` needs to resolve the concrete subclass.
with open(save_directory / CONFIG_NAME, "w") as f:
json.dump(draccus.encode(self, RLAlgorithmConfig), f, indent=4)
with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
draccus.dump(self, f, indent=4)
@classmethod
def from_pretrained(
+2 -2
View File
@@ -18,7 +18,7 @@ import functools
import threading
from collections.abc import Callable, Sequence
from contextlib import suppress
from typing import NotRequired, TypedDict
from typing import TypedDict
import torch
import torch.nn.functional as F # noqa: N812
@@ -36,7 +36,7 @@ class BatchTransition(TypedDict):
next_state: dict[str, torch.Tensor]
done: torch.Tensor
truncated: torch.Tensor
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
complementary_info: dict[str, torch.Tensor | float | int] | None = None
def random_crop_vectorized(images: torch.Tensor, output_size: tuple) -> torch.Tensor:
-7
View File
@@ -57,7 +57,6 @@ from .inference import (
SyncInferenceConfig,
create_inference_engine,
)
from .inference.rtc import supports_rtc_inference
from .robot_wrapper import ThreadSafeRobot
if TYPE_CHECKING or _peft_available:
@@ -227,12 +226,6 @@ def build_rollout_context(
policy = _load_pretrained_policy(policy_config)
if is_rtc:
if not supports_rtc_inference(policy):
raise ValueError(
f"RTC inference is not supported by policy type '{policy_config.type}': "
"the policy must implement RTC semantics and predict_action_chunk must accept "
"inference_delay and prev_chunk_left_over. Use '--inference.type=sync' instead."
)
policy.config.rtc_config = cfg.inference.rtc
if hasattr(policy, "init_rtc_processor"):
policy.init_rtc_processor()
-18
View File
@@ -22,7 +22,6 @@ way via ``notify_observation``.
from __future__ import annotations
import inspect
import logging
import math
import time
@@ -63,23 +62,6 @@ _RTC_JOIN_TIMEOUT_S: float = 3.0
# ---------------------------------------------------------------------------
def supports_rtc_inference(policy: PreTrainedPolicy) -> bool:
"""Whether a policy declares RTC support and accepts the RTC call shape."""
supports_rtc = getattr(policy, "supports_rtc", None)
if not callable(supports_rtc) or not supports_rtc():
return False
try:
inspect.signature(policy.predict_action_chunk).bind(
object(),
inference_delay=0,
prev_chunk_left_over=None,
)
except (TypeError, ValueError):
return False
return True
def _normalize_prev_actions_length(prev_actions: torch.Tensor, target_steps: int) -> torch.Tensor:
"""Pad or truncate RTC prefix actions to a fixed length for stable compiled inference."""
if prev_actions.ndim != 2:
+2 -15
View File
@@ -127,12 +127,6 @@ Modify tasks - set default task with overrides for specific episodes (WARNING: m
--operation.new_task "Default task" \
--operation.episode_tasks '{"5": "Special task for episode 5"}'
Modify tasks - replace existing task strings in-place (WARNING: modifies in-place):
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type modify_tasks \
--operation.task_replacements '{"Pick up the red cube": "Lift the red cube"}'
Convert image dataset to video format and save locally:
lerobot-edit-dataset \
--repo_id lerobot/pusht_image \
@@ -309,7 +303,6 @@ class RemoveFeatureConfig(OperationConfig):
class ModifyTasksConfig(OperationConfig):
new_task: str | None = None
episode_tasks: dict[str, str] | None = None
task_replacements: dict[str, str] | None = None
@OperationConfig.register_subclass("convert_image_to_video")
@@ -558,12 +551,9 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
new_task = cfg.operation.new_task
episode_tasks_raw = cfg.operation.episode_tasks
task_replacements = cfg.operation.task_replacements
if new_task is None and episode_tasks_raw is None and task_replacements is None:
raise ValueError(
"Must specify at least one of new_task, episode_tasks, or task_replacements for modify_tasks operation"
)
if new_task is None and episode_tasks_raw is None:
raise ValueError("Must specify at least one of new_task or episode_tasks for modify_tasks operation")
if cfg.new_repo_id is not None or cfg.new_root is not None:
logging.warning(
@@ -583,14 +573,11 @@ def handle_modify_tasks(cfg: EditDatasetConfig) -> None:
logging.info(f" Default task: '{new_task}'")
if episode_tasks:
logging.info(f" Episode-specific tasks: {episode_tasks}")
if task_replacements:
logging.info(f" Task replacements: {task_replacements}")
modified_dataset = modify_tasks(
dataset,
new_task=new_task,
episode_tasks=episode_tasks,
task_replacements=task_replacements,
)
logging.info(f"Dataset modified at {dataset.root}")
+49 -38
View File
@@ -28,6 +28,7 @@ lerobot-find-cameras
# NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.
import argparse
import concurrent.futures
import logging
import time
from pathlib import Path
@@ -132,7 +133,7 @@ def save_image(
camera_identifier: str | int,
images_dir: Path,
camera_type: str,
) -> None:
):
"""
Saves a single image to disk using Pillow. Handles color conversion if necessary.
"""
@@ -151,7 +152,7 @@ def save_image(
logger.error(f"Failed to save image for camera {camera_identifier} (type {camera_type}): {e}")
def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> dict[str, Any] | None:
def create_camera_instance(cam_meta: dict[str, Any]) -> dict[str, Any] | None:
"""Create and connect to a camera instance based on metadata."""
cam_type = cam_meta.get("type")
cam_id = cam_meta.get("id")
@@ -164,14 +165,12 @@ def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> di
cv_config = OpenCVCameraConfig(
index_or_path=cam_id,
color_mode=ColorMode.RGB,
warmup_s=warmup_s,
)
instance = OpenCVCamera(cv_config)
elif cam_type == "RealSense":
rs_config = RealSenseCameraConfig(
serial_number_or_name=cam_id,
color_mode=ColorMode.RGB,
warmup_s=warmup_s,
)
instance = RealSenseCamera(rs_config)
else:
@@ -189,7 +188,9 @@ def create_camera_instance(cam_meta: dict[str, Any], *, warmup_s: int = 1) -> di
return None
def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_time: float) -> None:
def process_camera_image(
cam_dict: dict[str, Any], output_dir: Path, current_time: float
) -> concurrent.futures.Future | None:
"""Capture and process an image from a single camera."""
cam = cam_dict["instance"]
meta = cam_dict["meta"]
@@ -199,7 +200,7 @@ def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_tim
try:
image_data = cam.read()
save_image(
return save_image(
image_data,
cam_id_str,
output_dir,
@@ -214,21 +215,21 @@ def process_camera_image(cam_dict: dict[str, Any], output_dir: Path, current_tim
return None
def cleanup_camera(cam_dict: dict[str, Any]) -> None:
def cleanup_cameras(cameras_to_use: list[dict[str, Any]]):
"""Disconnect all cameras."""
logger.info(f"Disconnecting camera with ID {cam_dict['meta'].get('id')}...")
try:
if cam_dict["instance"] and cam_dict["instance"].is_connected:
cam_dict["instance"].disconnect()
except Exception as e:
logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")
logger.info(f"Disconnecting {len(cameras_to_use)} cameras...")
for cam_dict in cameras_to_use:
try:
if cam_dict["instance"] and cam_dict["instance"].is_connected:
cam_dict["instance"].disconnect()
except Exception as e:
logger.error(f"Error disconnecting camera {cam_dict['meta'].get('id')}: {e}")
def save_images_from_all_cameras(
output_dir: Path,
record_time_s: float = 2.0,
camera_type: str | None = None,
warmup_s: int = 1,
):
"""
Connects to detected cameras (optionally filtered by type) and saves images from each.
@@ -239,7 +240,6 @@ def save_images_from_all_cameras(
record_time_s: Duration in seconds to record images.
camera_type: Optional string to filter cameras ("realsense" or "opencv").
If None, uses all detected cameras.
warmup_s: Duration in seconds to warmup camera before recording images.
"""
output_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Saving images to {output_dir}")
@@ -249,24 +249,40 @@ def save_images_from_all_cameras(
logger.warning("No cameras detected matching the criteria. Cannot save images.")
return
logger.info(
f"Starting image capture for {record_time_s} seconds from {len(all_camera_metadata)} cameras."
)
cameras_to_use = []
for cam_meta in all_camera_metadata:
camera_instance = create_camera_instance(cam_meta)
if camera_instance:
cameras_to_use.append(camera_instance)
try:
for cam_meta in all_camera_metadata:
cam_dict = create_camera_instance(cam_meta, warmup_s=warmup_s)
if cam_dict is None:
continue
start_time = time.perf_counter()
if not cameras_to_use:
logger.warning("No cameras could be connected. Aborting image save.")
return
logger.info(f"Starting image capture for {record_time_s} seconds from {len(cameras_to_use)} cameras.")
start_time = time.perf_counter()
with concurrent.futures.ThreadPoolExecutor(max_workers=len(cameras_to_use) * 2) as executor:
try:
while time.perf_counter() - start_time < record_time_s:
futures = []
current_capture_time = time.perf_counter()
process_camera_image(cam_dict, output_dir, current_capture_time)
cleanup_camera(cam_dict)
except KeyboardInterrupt:
logger.info("Capture interrupted by user.")
finally:
print(f"Image capture finished. Images saved to {output_dir}")
for cam_dict in cameras_to_use:
future = process_camera_image(cam_dict, output_dir, current_capture_time)
if future:
futures.append(future)
if futures:
concurrent.futures.wait(futures)
except KeyboardInterrupt:
logger.info("Capture interrupted by user.")
finally:
print("\nFinalizing image saving...")
executor.shutdown(wait=True)
cleanup_cameras(cameras_to_use)
print(f"Image capture finished. Images saved to {output_dir}")
def main():
@@ -275,6 +291,7 @@ def main():
parser = argparse.ArgumentParser(
description="Unified camera utility script for listing cameras and capturing images."
)
parser.add_argument(
"camera_type",
type=str,
@@ -292,14 +309,8 @@ def main():
parser.add_argument(
"--record-time-s",
type=float,
default=2.0,
help="Time duration to attempt capturing frames. Default: 2 seconds.",
)
parser.add_argument(
"--warmup-s",
type=int,
default=1,
help="Time duration to warmup camera before attempting to capture frames. Default: 1 second.",
default=6.0,
help="Time duration to attempt capturing frames. Default: 6 seconds.",
)
args = parser.parse_args()
save_images_from_all_cameras(**vars(args))
+2 -7
View File
@@ -348,6 +348,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
preprocessor_overrides = {
"device_processor": {"device": device.type},
"normalizer_processor": {
"stats": dataset.meta.stats,
"features": {**policy.config.input_features, **policy.config.output_features},
"norm_map": policy.config.normalization_mapping,
},
@@ -355,17 +356,11 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
}
postprocessor_overrides = {
"unnormalizer_processor": {
"stats": dataset.meta.stats,
"features": policy.config.output_features,
"norm_map": policy.config.normalization_mapping,
},
}
# On resume, the checkpoint's saved processor stats are authoritative: they may have
# been adapted by the policy (e.g. EVO1 pads state/action stats to max_state_dim),
# and force-feeding raw dataset stats over them crashes normalization (#4006).
# This mirrors the `dataset_stats` kwarg above, which is also skipped on resume.
if not cfg.resume:
preprocessor_overrides["normalizer_processor"]["stats"] = dataset.meta.stats
postprocessor_overrides["unnormalizer_processor"]["stats"] = dataset.meta.stats
if getattr(active_cfg, "use_relative_actions", False):
preprocessor_overrides["relative_actions_processor"] = {
"enabled": True,
@@ -67,15 +67,7 @@ class BiSOLeader(BimanualMixin, Teleoperator):
@cached_property
def feedback_features(self) -> dict[str, type]:
# Bimanual teleop has feedback (can be actuated for handover).
# Return the same structure as action_features for consistency with left/right arms.
left_arm_features = self.left_arm.feedback_features
right_arm_features = self.right_arm.feedback_features
return {
**{f"left_{k}": v for k, v in left_arm_features.items()},
**{f"right_{k}": v for k, v in right_arm_features.items()},
}
return {}
def setup_motors(self) -> None:
self.left_arm.setup_motors()
@@ -95,43 +87,6 @@ class BiSOLeader(BimanualMixin, Teleoperator):
return action_dict
def enable_torque(self) -> None:
"""Enable torque on both leader arms for smooth handover."""
self.left_arm.enable_torque()
self.right_arm.enable_torque()
def disable_torque(self) -> None:
"""Disable torque on both leader arms to allow human control."""
self.left_arm.disable_torque()
self.right_arm.disable_torque()
@check_if_not_connected
def send_feedback(self, feedback: dict[str, float]) -> None:
"""Route bimanual feedback to left and right arms with proper prefix stripping.
Receives feedback dict with keys like: left_shoulder_pan.pos, right_shoulder_pan.pos, ...
Splits and routes to each arm by removing the prefix.
This enables DAgger smooth handover: when transitioning from policy control to human
intervention, both leader arms are commanded to the follower's current pose to avoid
discontinuities.
"""
# Split feedback by arm prefix
left_feedback = {}
right_feedback = {}
for key, value in feedback.items():
if key.startswith("left_"):
# Strip "left_" prefix and pass to left arm
stripped_key = key[5:] # len("left_") == 5
left_feedback[stripped_key] = value
elif key.startswith("right_"):
# Strip "right_" prefix and pass to right arm
stripped_key = key[6:] # len("right_") == 6
right_feedback[stripped_key] = value
# Send to each arm
if left_feedback:
self.left_arm.send_feedback(left_feedback)
if right_feedback:
self.right_arm.send_feedback(right_feedback)
# TODO: Implement force feedback
raise NotImplementedError
@@ -171,13 +171,7 @@ class IOSPhone(BasePhone, Teleoperator):
# HEBI provides orientation in w, x, y, z format.
# Scipy's Rotation expects x, y, z, w.
quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]])) # wxyz to xyzw
# ARKit can emit zero/NaN quaternions before tracking is ready or on a
# dropped packet. Rotation.from_quat now rejects those; degrade the same
# way as a missing pose so teleop stays alive mid-session.
try:
rot = Rotation.from_quat(quat_xyzw)
except ValueError:
return False, None, None, None
rot = Rotation.from_quat(quat_xyzw)
pos = ar_pos - rot.apply(self.config.camera_offset)
return True, pos, rot, pose
+4 -13
View File
@@ -37,25 +37,16 @@ def auto_select_torch_device() -> torch.device:
# TODO(Steven): Remove log. log shouldn't be an argument, this should be handled by the logger level
def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
"""Given a string, return a torch.device with checks on whether the device is available.
Raises:
ValueError: If the requested device family is known but not available on
this machine (``AssertionError`` was previously used and is easy to
mistake for a programmer bug under ``python -O`` where asserts vanish).
"""
"""Given a string, return a torch.device with checks on whether the device is available."""
try_device = str(try_device)
if try_device.startswith("cuda"):
if not torch.cuda.is_available():
raise ValueError(f"Requested device {try_device!r} but CUDA is not available.")
assert torch.cuda.is_available()
device = torch.device(try_device)
elif try_device == "mps":
if not torch.backends.mps.is_available():
raise ValueError("Requested device 'mps' but MPS is not available.")
assert torch.backends.mps.is_available()
device = torch.device("mps")
elif try_device == "xpu":
if not torch.xpu.is_available():
raise ValueError("Requested device 'xpu' but XPU is not available.")
assert torch.xpu.is_available()
device = torch.device("xpu")
elif try_device == "cpu":
device = torch.device("cpu")
+5 -5
View File
@@ -32,21 +32,21 @@ def load_json(fpath: Path) -> Any:
Returns:
Any: The data loaded from the JSON file.
"""
with open(fpath, encoding="utf-8") as f:
with open(fpath) as f:
return json.load(f)
def write_json(data: JsonLike, fpath: Path) -> None:
"""Write JSON-serializable data to a file.
def write_json(data: dict, fpath: Path) -> None:
"""Write data to a JSON file.
Creates parent directories if they don't exist.
Args:
data: JSON-serializable data to write.
data (dict): The dictionary to write.
fpath (Path): The path to the output JSON file.
"""
fpath.parent.mkdir(exist_ok=True, parents=True)
with open(fpath, "w", encoding="utf-8") as f:
with open(fpath, "w") as f:
json.dump(data, f, indent=4, ensure_ascii=False)
-4
View File
@@ -85,8 +85,6 @@ def serialize_torch_rng_state() -> dict[str, torch.Tensor]:
torch_rng_state_dict = {"torch_rng_state": torch.get_rng_state()}
if torch.cuda.is_available():
torch_rng_state_dict["torch_cuda_rng_state"] = torch.cuda.get_rng_state()
if torch.backends.mps.is_available():
torch_rng_state_dict["torch_mps_rng_state"] = torch.mps.get_rng_state()
return torch_rng_state_dict
@@ -97,8 +95,6 @@ def deserialize_torch_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None
torch.set_rng_state(rng_state_dict["torch_rng_state"])
if torch.cuda.is_available() and "torch_cuda_rng_state" in rng_state_dict:
torch.cuda.set_rng_state(rng_state_dict["torch_cuda_rng_state"])
if torch.backends.mps.is_available() and "torch_mps_rng_state" in rng_state_dict:
torch.mps.set_rng_state(rng_state_dict["torch_mps_rng_state"])
def serialize_rng_state() -> dict[str, torch.Tensor]:
-4
View File
@@ -30,10 +30,6 @@ def precise_sleep(seconds: float, spin_threshold: float = 0.010, sleep_margin: f
"""
if seconds <= 0:
return
if spin_threshold < 0:
raise ValueError(f"spin_threshold must be >= 0, got {spin_threshold}")
if sleep_margin < 0:
raise ValueError(f"sleep_margin must be >= 0, got {sleep_margin}")
system = platform.system()
# On macOS and Windows the scheduler / sleep granularity can make
+3 -6
View File
@@ -29,13 +29,10 @@ class Rotation:
def __init__(self, quat: np.ndarray) -> None:
"""Initialize rotation from quaternion [x, y, z, w]."""
self._quat = np.asarray(quat, dtype=float)
if self._quat.shape != (4,):
raise ValueError(f"Quaternion must have shape (4,), got {self._quat.shape}")
# Normalize quaternion. Reject the zero vector — it has no orientation.
# Normalize quaternion
norm = np.linalg.norm(self._quat)
if norm <= 0.0 or not np.isfinite(norm):
raise ValueError(f"Quaternion must be a non-zero finite vector; got {self._quat} (norm={norm})")
self._quat = self._quat / norm
if norm > 0:
self._quat = self._quat / norm
@classmethod
def from_rotvec(cls, rotvec: np.ndarray) -> "Rotation":
+2 -2
View File
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import NotRequired, TypedDict
from typing import TypedDict
import torch
@@ -28,7 +28,7 @@ class Transition(TypedDict):
next_state: dict[str, torch.Tensor]
done: bool
truncated: bool
complementary_info: NotRequired[dict[str, torch.Tensor | float | int] | None]
complementary_info: dict[str, torch.Tensor | float | int] | None = None
def move_transition_to_device(transition: Transition, device: str = "cpu") -> Transition:
+8 -9
View File
@@ -24,6 +24,7 @@ import sys
import time
from collections.abc import Iterator
from copy import copy, deepcopy
from datetime import datetime
from pathlib import Path
from statistics import mean
from typing import TYPE_CHECKING, Any
@@ -60,16 +61,14 @@ def init_logging(
accelerator: Optional Accelerator instance (for multi-GPU detection)
"""
class LeRobotFormatter(logging.Formatter):
def format(self, record: logging.LogRecord) -> str:
record.lerobot_location = f"{record.pathname}:{record.lineno}"[-15:]
record.lerobot_pid = f"[PID: {os.getpid()}] " if display_pid else ""
return super().format(record)
def custom_format(record: logging.LogRecord) -> str:
dt = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
fnameline = f"{record.pathname}:{record.lineno}"
pid_str = f"[PID: {os.getpid()}] " if display_pid else ""
return f"{record.levelname} {pid_str}{dt} {fnameline[-15:]:>15} {record.getMessage()}"
formatter = LeRobotFormatter(
"%(levelname)s %(lerobot_pid)s%(asctime)s %(lerobot_location)15s %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
formatter = logging.Formatter()
formatter.format = custom_format
logger = logging.getLogger()
logger.setLevel(logging.NOTSET)
-24
View File
@@ -66,27 +66,3 @@ def test_from_pretrained_raises_when_no_root_config_and_no_checkpoints(monkeypat
with pytest.raises(FileNotFoundError, match="train_config.json not found"):
TrainPipelineConfig.from_pretrained("user/empty-repo")
@pytest.mark.parametrize("pass_dir", [False, True])
def test_resolve_resume_checkpoint_accepts_file_or_pretrained_model_dir(tmp_path, monkeypatch, pass_dir):
"""`--config_path` may point at the checkpoint's train_config.json or at its
pretrained_model/ directory; both must resolve `policy.pretrained_path` to the
pretrained_model/ directory (regression test for the directory case, which
previously resolved one level too high and failed on model.safetensors)."""
pretrained_dir = tmp_path / "checkpoints" / "000002" / "pretrained_model"
pretrained_dir.mkdir(parents=True)
(pretrained_dir / "train_config.json").touch()
target = pretrained_dir if pass_dir else pretrained_dir / "train_config.json"
from lerobot.policies.act.configuration_act import ACTConfig
cfg = tc.draccus.parse(TrainPipelineConfig, args=["--dataset.repo_id", "u/d"])
cfg.policy = ACTConfig()
cfg.resume = True
monkeypatch.setattr(tc.parser, "parse_arg", lambda name: str(target) if name == "config_path" else None)
cfg._resolve_resume_checkpoint()
assert cfg.policy.pretrained_path == pretrained_dir
assert cfg.checkpoint_path == pretrained_dir.parent
+1 -59
View File
@@ -1125,61 +1125,9 @@ def test_modify_tasks_default_with_overrides(sample_dataset):
assert ep_data["tasks"][0] == default_task
def test_modify_tasks_replacements(sample_dataset):
"""Test replacing task strings based on their current values."""
modified_dataset = modify_tasks(
sample_dataset,
task_replacements={
"task_0": "Pick the cube",
"task_1": "Place the cube",
},
)
assert len(modified_dataset.meta.tasks) == 2
assert "Pick the cube" in modified_dataset.meta.tasks.index
assert "Place the cube" in modified_dataset.meta.tasks.index
for ep_idx in range(5):
expected_task = "Pick the cube" if ep_idx % 2 == 0 else "Place the cube"
assert modified_dataset.meta.episodes[ep_idx]["tasks"][0] == expected_task
def test_modify_tasks_replacements_with_episode_overrides(sample_dataset):
"""Test that explicit episode overrides take precedence over replacements."""
modified_dataset = modify_tasks(
sample_dataset,
task_replacements={
"task_0": "Pick the cube",
"task_1": "Place the cube",
},
episode_tasks={1: "Inspect the cube"},
)
assert modified_dataset.meta.episodes[0]["tasks"][0] == "Pick the cube"
assert modified_dataset.meta.episodes[1]["tasks"][0] == "Inspect the cube"
assert modified_dataset.meta.episodes[3]["tasks"][0] == "Place the cube"
assert len(modified_dataset.meta.tasks) == 3
def test_modify_tasks_default_task_and_replacements(sample_dataset):
"""Test that new_task acts as the default for episodes not matched by task_replacements."""
modified_dataset = modify_tasks(
sample_dataset,
new_task="Default task",
task_replacements={"task_0": "Pick the cube"},
)
for ep_idx in range(5):
expected_task = "Pick the cube" if ep_idx % 2 == 0 else "Default task"
assert modified_dataset.meta.episodes[ep_idx]["tasks"][0] == expected_task
assert len(modified_dataset.meta.tasks) == 2
def test_modify_tasks_no_task_specified(sample_dataset):
"""Test error when no task is specified."""
with pytest.raises(
ValueError, match="Must specify at least one of new_task, episode_tasks, or task_replacements"
):
with pytest.raises(ValueError, match="Must specify at least one of new_task or episode_tasks"):
modify_tasks(sample_dataset)
@@ -1189,12 +1137,6 @@ def test_modify_tasks_invalid_episode_indices(sample_dataset):
modify_tasks(sample_dataset, episode_tasks={10: "Task", 20: "Task"})
def test_modify_tasks_invalid_task_replacements(sample_dataset):
"""Test error when task replacements refer to unknown task strings."""
with pytest.raises(ValueError, match="Task replacements reference unknown tasks"):
modify_tasks(sample_dataset, task_replacements={"missing_task": "New task"})
def test_modify_tasks_updates_info_json(sample_dataset):
"""Test that total_tasks is updated in info.json."""
episode_tasks = {0: "Task A", 1: "Task B", 2: "Task C", 3: "Task A", 4: "Task B"}
@@ -105,23 +105,6 @@ class TestOperationTypeParsing:
resolved_name = OperationConfig.get_choice_name(type(cfg.operation))
assert resolved_name == type_name
def test_modify_tasks_replacements_args_parse(self):
cfg = parse_cfg(
[
"--repo_id",
"test/repo",
"--operation.type",
"modify_tasks",
"--operation.task_replacements",
'{"task_0": "pick cube", "task_1": "place cube"}',
]
)
assert isinstance(cfg.operation, ModifyTasksConfig)
assert cfg.operation.task_replacements == {
"task_0": "pick cube",
"task_1": "place cube",
}
class TestDepthEncoderParsing:
"""Test that the depth encoder is exposed and parsed for video operations."""
-36
View File
@@ -1,36 +0,0 @@
#!/usr/bin/env python
# Copyright 2026 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.
from unittest.mock import patch
import pytest
import torch
from lerobot.utils.device_utils import get_safe_torch_device, is_torch_device_available
def test_cpu_always_available():
assert get_safe_torch_device("cpu") == torch.device("cpu")
assert is_torch_device_available("cpu")
def test_missing_cuda_raises_valueerror():
with patch("torch.cuda.is_available", return_value=False), pytest.raises(ValueError, match="CUDA"):
get_safe_torch_device("cuda")
def test_missing_mps_raises_valueerror():
with patch("torch.backends.mps.is_available", return_value=False), pytest.raises(ValueError, match="MPS"):
get_safe_torch_device("mps")
-11
View File
@@ -73,17 +73,6 @@ def test_serialize_deserialize_torch_rng(fixed_seed):
assert val2 == val3
@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
def test_serialize_deserialize_torch_rng_mps(fixed_seed):
_ = torch.rand(1, device="mps").item()
st = serialize_torch_rng_state()
assert "torch_mps_rng_state" in st
val2 = torch.rand(1, device="mps").item()
deserialize_torch_rng_state(st)
val3 = torch.rand(1, device="mps").item()
assert val2 == val3
def test_serialize_deserialize_rng(fixed_seed):
# Generate one from each library
_ = random.random()
-46
View File
@@ -1,46 +0,0 @@
#!/usr/bin/env python
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
import pytest
from lerobot.utils.rotation import Rotation
def test_zero_quaternion_rejected():
with pytest.raises(ValueError, match="non-zero"):
Rotation(np.zeros(4))
def test_non_finite_quaternion_rejected():
with pytest.raises(ValueError, match="non-zero|finite"):
Rotation(np.array([np.nan, 0.0, 0.0, 1.0]))
def test_wrong_shape_rejected():
with pytest.raises(ValueError, match="shape"):
Rotation(np.array([1.0, 0.0, 0.0]))
def test_identity_roundtrip():
r = Rotation.from_rotvec(np.zeros(3))
assert np.allclose(r.as_rotvec(), 0.0)
assert np.allclose(r.as_matrix(), np.eye(3))
def test_rotvec_roundtrip():
rotvec = np.array([0.1, -0.2, 0.3])
r = Rotation.from_rotvec(rotvec)
assert np.allclose(r.as_rotvec(), rotvec, atol=1e-6)
Generated
+18 -5
View File
@@ -1,5 +1,5 @@
version = 1
revision = 2
revision = 3
requires-python = ">=3.12"
resolution-markers = [
"(python_full_version >= '3.15' and platform_machine == 'AMD64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux')",
@@ -1359,17 +1359,18 @@ sdist = { url = "https://files.pythonhosted.org/packages/a2/55/8f8cab2afd404cf57
[[package]]
name = "draccus"
version = "0.11.6"
version = "0.10.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "mergedeep" },
{ name = "pyyaml" },
{ name = "pyyaml-include" },
{ name = "toml" },
{ name = "typing-inspect" },
]
sdist = { url = "https://files.pythonhosted.org/packages/df/b0/cc399719e0cf6fea451f155798e19ea8c5e9c838b701d1565f29ea626c18/draccus-0.11.6.tar.gz", hash = "sha256:d134f576a1f4febd93c6b200df7f92e5febffe9efdc0a5f381bf50c2e0568a39", size = 67940, upload-time = "2026-06-12T13:21:19.999Z" }
sdist = { url = "https://files.pythonhosted.org/packages/4e/e2/f5012fda17ee5d1eaf3481b6ca3e11dffa5348e5e08ab745538fdc8041bb/draccus-0.10.0.tar.gz", hash = "sha256:8dd08304219becdcd66cd16058ba98e9c3e6b7bfe48ccb9579dae39f8d37ae19", size = 62243, upload-time = "2025-02-05T07:27:48.182Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/18/f1/d56bef4563d1cfaf55dd58de298a456587b9fe0a85dabb55768d44ace8f7/draccus-0.11.6-py3-none-any.whl", hash = "sha256:1cf3f37c64766e0f17d757099b388e762f1d22c772cc2376b4e366b515fb5737", size = 85449, upload-time = "2026-06-12T13:21:18.83Z" },
{ url = "https://files.pythonhosted.org/packages/c4/9a/a83083b230d352ee5d205757b74006dbe084448ca45e3bc5ca99215b1e55/draccus-0.10.0-py3-none-any.whl", hash = "sha256:90243418ae0e9271c390a59cafb6acfd37001193696ed36fcc8525f791a83282", size = 71783, upload-time = "2025-02-05T07:27:46.1Z" },
]
[[package]]
@@ -3288,7 +3289,7 @@ requires-dist = [
{ name = "deepdiff", marker = "extra == 'deepdiff-dep'", specifier = ">=7.0.1,<9.0.0" },
{ name = "diffusers", marker = "extra == 'diffusers-dep'", specifier = ">=0.38.0,<0.40.0" },
{ name = "dm-tree", marker = "extra == 'groot'", specifier = ">=0.1.8,<1.0.0" },
{ name = "draccus", specifier = ">=0.11.6,<0.12.0" },
{ name = "draccus", specifier = "==0.10.0" },
{ name = "dynamixel-sdk", marker = "extra == 'dynamixel'", specifier = ">=3.7.31,<3.9.0" },
{ name = "einops", specifier = ">=0.8.0,<0.9.0" },
{ name = "faker", marker = "extra == 'sarm'", specifier = ">=33.0.0,<35.0.0" },
@@ -5635,6 +5636,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" },
]
[[package]]
name = "pyyaml-include"
version = "1.4.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pyyaml" },
]
sdist = { url = "https://files.pythonhosted.org/packages/7f/be/2d07ad85e3d593d69640876a8686eae2c533db8cb7bf298d25c421b4d2d5/pyyaml-include-1.4.1.tar.gz", hash = "sha256:1a96e33a99a3e56235f5221273832464025f02ff3d8539309a3bf00dec624471", size = 20592, upload-time = "2024-03-25T14:56:43.748Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/d5/ca/6a2cc3a73170d10b5af1f1613baa2ed1f8f46f62dd0bfab2bffd2c2fe260/pyyaml_include-1.4.1-py3-none-any.whl", hash = "sha256:323c7f3a19c82fbc4d73abbaab7ef4f793e146a13383866831631b26ccc7fb00", size = 19079, upload-time = "2024-03-25T14:56:41.274Z" },
]
[[package]]
name = "pyzmq"
version = "27.1.0"