mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-29 20:49:42 +00:00
Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 0d788abd85 | |||
| 7b78e751a6 |
+2
-3
@@ -321,11 +321,10 @@ SmolVLA ships with `freeze_vision_encoder=True`. Unfreezing usually **improves p
|
|||||||
|
|
||||||
```bash
|
```bash
|
||||||
lerobot-train ... --policy.type=smolvla \
|
lerobot-train ... --policy.type=smolvla \
|
||||||
--policy.fine_tune_vision_encoder=true
|
--policy.freeze_vision_encoder=false \
|
||||||
|
--policy.train_expert_only=false
|
||||||
```
|
```
|
||||||
|
|
||||||
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
|
### 7.7 Signals to stop / keep going
|
||||||
|
|
||||||
- Train loss plateaus → stop, save a Hub checkpoint.
|
- Train loss plateaus → stop, save a Hub checkpoint.
|
||||||
|
|||||||
@@ -70,19 +70,6 @@ cd lerobot && lerobot-train \
|
|||||||
GPU allows it, as long as loading times remain short.
|
GPU allows it, as long as loading times remain short.
|
||||||
</Tip>
|
</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
|
Fine-tuning is an art. For a complete overview of the options for finetuning, run
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
|
|||||||
@@ -67,8 +67,6 @@ class SmolVLAConfig(PreTrainedConfig):
|
|||||||
|
|
||||||
# Finetuning settings
|
# Finetuning settings
|
||||||
freeze_vision_encoder: bool = True
|
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_expert_only: bool = True
|
||||||
train_state_proj: bool = True
|
train_state_proj: bool = True
|
||||||
|
|
||||||
@@ -112,12 +110,6 @@ class SmolVLAConfig(PreTrainedConfig):
|
|||||||
super().__post_init__()
|
super().__post_init__()
|
||||||
|
|
||||||
"""Input validation (not exhaustive)."""
|
"""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:
|
if self.n_action_steps > self.chunk_size:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"The chunk size is the upper bound for the number of action steps per model invocation. Got "
|
f"The chunk size is the upper bound for the number of action steps per model invocation. Got "
|
||||||
|
|||||||
@@ -186,28 +186,9 @@ class SmolVLAPolicy(PreTrainedPolicy):
|
|||||||
if model_value is not None:
|
if model_value is not None:
|
||||||
model_value.rtc_processor = self.rtc_processor
|
model_value.rtc_processor = self.rtc_processor
|
||||||
|
|
||||||
def get_optim_params(self):
|
def get_optim_params(self) -> dict:
|
||||||
if not self.config.fine_tune_vision_encoder:
|
|
||||||
return self.parameters()
|
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(
|
def _get_action_chunk(
|
||||||
self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs: Unpack[ActionSelectKwargs]
|
self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs: Unpack[ActionSelectKwargs]
|
||||||
) -> Tensor:
|
) -> Tensor:
|
||||||
@@ -512,7 +493,6 @@ class VLAFlowMatching(nn.Module):
|
|||||||
self.vlm_with_expert = SmolVLMWithExpertModel(
|
self.vlm_with_expert = SmolVLMWithExpertModel(
|
||||||
model_id=self.config.vlm_model_name,
|
model_id=self.config.vlm_model_name,
|
||||||
freeze_vision_encoder=self.config.freeze_vision_encoder,
|
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,
|
train_expert_only=self.config.train_expert_only,
|
||||||
load_vlm_weights=self.config.load_vlm_weights,
|
load_vlm_weights=self.config.load_vlm_weights,
|
||||||
attention_mode=self.config.attention_mode,
|
attention_mode=self.config.attention_mode,
|
||||||
|
|||||||
@@ -78,7 +78,6 @@ class SmolVLMWithExpertModel(nn.Module):
|
|||||||
load_vlm_weights: bool = True,
|
load_vlm_weights: bool = True,
|
||||||
train_expert_only: bool = True,
|
train_expert_only: bool = True,
|
||||||
freeze_vision_encoder: bool = False,
|
freeze_vision_encoder: bool = False,
|
||||||
fine_tune_vision_encoder: bool = False,
|
|
||||||
attention_mode: str = "self_attn",
|
attention_mode: str = "self_attn",
|
||||||
num_expert_layers: int = -1,
|
num_expert_layers: int = -1,
|
||||||
num_vlm_layers: int = -1,
|
num_vlm_layers: int = -1,
|
||||||
@@ -142,7 +141,6 @@ class SmolVLMWithExpertModel(nn.Module):
|
|||||||
self.num_key_value_heads = self.config.text_config.num_key_value_heads
|
self.num_key_value_heads = self.config.text_config.num_key_value_heads
|
||||||
|
|
||||||
self.freeze_vision_encoder = freeze_vision_encoder
|
self.freeze_vision_encoder = freeze_vision_encoder
|
||||||
self.fine_tune_vision_encoder = fine_tune_vision_encoder
|
|
||||||
self.train_expert_only = train_expert_only
|
self.train_expert_only = train_expert_only
|
||||||
self.attention_mode = attention_mode
|
self.attention_mode = attention_mode
|
||||||
self.expert_hidden_size = lm_expert_config.hidden_size
|
self.expert_hidden_size = lm_expert_config.hidden_size
|
||||||
@@ -152,6 +150,10 @@ class SmolVLMWithExpertModel(nn.Module):
|
|||||||
return self.vlm.model
|
return self.vlm.model
|
||||||
|
|
||||||
def set_requires_grad(self):
|
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:
|
if self.train_expert_only:
|
||||||
self.vlm.eval()
|
self.vlm.eval()
|
||||||
for params in self.vlm.parameters():
|
for params in self.vlm.parameters():
|
||||||
@@ -174,18 +176,6 @@ class SmolVLMWithExpertModel(nn.Module):
|
|||||||
for name, params in self.vlm.named_parameters():
|
for name, params in self.vlm.named_parameters():
|
||||||
if any(k in name for k in frozen_layers):
|
if any(k in name for k in frozen_layers):
|
||||||
params.requires_grad = False
|
params.requires_grad = False
|
||||||
|
|
||||||
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
|
# To avoid unused params issue with distributed training
|
||||||
for name, params in self.lm_expert.named_parameters():
|
for name, params in self.lm_expert.named_parameters():
|
||||||
if "lm_head" in name:
|
if "lm_head" in name:
|
||||||
@@ -194,15 +184,11 @@ class SmolVLMWithExpertModel(nn.Module):
|
|||||||
def train(self, mode: bool = True):
|
def train(self, mode: bool = True):
|
||||||
super().train(mode)
|
super().train(mode)
|
||||||
|
|
||||||
if self.train_expert_only:
|
|
||||||
self.vlm.eval()
|
|
||||||
|
|
||||||
if self.freeze_vision_encoder:
|
if self.freeze_vision_encoder:
|
||||||
self.get_vlm_model().vision_model.eval()
|
self.get_vlm_model().vision_model.eval()
|
||||||
|
|
||||||
if self.fine_tune_vision_encoder:
|
if self.train_expert_only:
|
||||||
self.get_vlm_model().vision_model.train(mode)
|
self.vlm.eval()
|
||||||
self.get_vlm_model().connector.train(mode)
|
|
||||||
|
|
||||||
def embed_image(self, image: torch.Tensor):
|
def embed_image(self, image: torch.Tensor):
|
||||||
patch_attention_mask = None
|
patch_attention_mask = None
|
||||||
|
|||||||
@@ -510,10 +510,10 @@ class ForwardKinematicsJointsToEEAction(RobotActionProcessorStep):
|
|||||||
# We only use the ee pose in the dataset, so we don't need the joint positions
|
# We only use the ee pose in the dataset, so we don't need the joint positions
|
||||||
for n in self.motor_names:
|
for n in self.motor_names:
|
||||||
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
|
features[PipelineFeatureType.ACTION].pop(f"{n}.pos", None)
|
||||||
# Store end-effector features as actions in the dataset schema
|
# We specify the dataset features of this step that we want to be stored in the dataset
|
||||||
for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
|
for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]:
|
||||||
features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
|
features[PipelineFeatureType.ACTION][f"ee.{k}"] = PolicyFeature(
|
||||||
type=FeatureType.ACTION, shape=(1,)
|
type=FeatureType.STATE, shape=(1,)
|
||||||
)
|
)
|
||||||
return features
|
return features
|
||||||
|
|
||||||
|
|||||||
@@ -142,11 +142,25 @@ class SOFollower(Robot):
|
|||||||
range_mins[full_turn_motor] = 0
|
range_mins[full_turn_motor] = 0
|
||||||
range_maxes[full_turn_motor] = 4095
|
range_maxes[full_turn_motor] = 4095
|
||||||
|
|
||||||
|
drive_modes = dict.fromkeys(self.bus.motors, 0)
|
||||||
|
input(f"Fully close the gripper of {self} and press ENTER....")
|
||||||
|
gripper_closed_pos = self.bus.read(
|
||||||
|
"Present_Position", "gripper", normalize=False, num_retry=self.config.num_read_retries
|
||||||
|
)
|
||||||
|
distance_to_min = abs(gripper_closed_pos - range_mins["gripper"])
|
||||||
|
distance_to_max = abs(gripper_closed_pos - range_maxes["gripper"])
|
||||||
|
if min(distance_to_min, distance_to_max) > (range_maxes["gripper"] - range_mins["gripper"]) * 0.2:
|
||||||
|
raise ValueError("Gripper is not fully closed. Run calibration again.")
|
||||||
|
|
||||||
|
drive_modes["gripper"] = int(distance_to_max < distance_to_min)
|
||||||
|
if drive_modes["gripper"]:
|
||||||
|
logger.info("Gripper motor is inverted, setting drive_mode=1 to compensate.")
|
||||||
|
|
||||||
self.calibration = {}
|
self.calibration = {}
|
||||||
for motor, m in self.bus.motors.items():
|
for motor, m in self.bus.motors.items():
|
||||||
self.calibration[motor] = MotorCalibration(
|
self.calibration[motor] = MotorCalibration(
|
||||||
id=m.id,
|
id=m.id,
|
||||||
drive_mode=0,
|
drive_mode=drive_modes[motor],
|
||||||
homing_offset=homing_offsets[motor],
|
homing_offset=homing_offsets[motor],
|
||||||
range_min=range_mins[motor],
|
range_min=range_mins[motor],
|
||||||
range_max=range_maxes[motor],
|
range_max=range_maxes[motor],
|
||||||
|
|||||||
@@ -110,11 +110,25 @@ class SOLeader(Teleoperator):
|
|||||||
range_mins[full_turn_motor] = 0
|
range_mins[full_turn_motor] = 0
|
||||||
range_maxes[full_turn_motor] = 4095
|
range_maxes[full_turn_motor] = 4095
|
||||||
|
|
||||||
|
drive_modes = dict.fromkeys(self.bus.motors, 0)
|
||||||
|
input(f"Fully close the gripper of {self} and press ENTER....")
|
||||||
|
gripper_closed_pos = self.bus.read(
|
||||||
|
"Present_Position", "gripper", normalize=False, num_retry=self.config.num_read_retries
|
||||||
|
)
|
||||||
|
distance_to_min = abs(gripper_closed_pos - range_mins["gripper"])
|
||||||
|
distance_to_max = abs(gripper_closed_pos - range_maxes["gripper"])
|
||||||
|
if min(distance_to_min, distance_to_max) > (range_maxes["gripper"] - range_mins["gripper"]) * 0.2:
|
||||||
|
raise ValueError("Gripper is not fully closed. Run calibration again.")
|
||||||
|
|
||||||
|
drive_modes["gripper"] = int(distance_to_max < distance_to_min)
|
||||||
|
if drive_modes["gripper"]:
|
||||||
|
logger.info("Gripper motor is inverted, setting drive_mode=1 to compensate.")
|
||||||
|
|
||||||
self.calibration = {}
|
self.calibration = {}
|
||||||
for motor, m in self.bus.motors.items():
|
for motor, m in self.bus.motors.items():
|
||||||
self.calibration[motor] = MotorCalibration(
|
self.calibration[motor] = MotorCalibration(
|
||||||
id=m.id,
|
id=m.id,
|
||||||
drive_mode=0,
|
drive_mode=drive_modes[motor],
|
||||||
homing_offset=homing_offsets[motor],
|
homing_offset=homing_offsets[motor],
|
||||||
range_min=range_mins[motor],
|
range_min=range_mins[motor],
|
||||||
range_max=range_maxes[motor],
|
range_max=range_maxes[motor],
|
||||||
|
|||||||
@@ -1,45 +0,0 @@
|
|||||||
#!/usr/bin/env python
|
|
||||||
|
|
||||||
# Copyright 2025 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 pytest
|
|
||||||
|
|
||||||
from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
|
|
||||||
from lerobot.robots.so_follower.robot_kinematic_processor import (
|
|
||||||
ForwardKinematicsJointsToEEAction,
|
|
||||||
ForwardKinematicsJointsToEEObservation,
|
|
||||||
)
|
|
||||||
|
|
||||||
MOTOR_NAMES = ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"]
|
|
||||||
EE_KEYS = {f"ee.{k}" for k in ["x", "y", "z", "wx", "wy", "wz", "gripper_pos"]}
|
|
||||||
|
|
||||||
|
|
||||||
def _joint_bucket(feature_type: FeatureType) -> dict[str, PolicyFeature]:
|
|
||||||
return {f"{n}.pos": PolicyFeature(type=feature_type, shape=(1,)) for n in MOTOR_NAMES}
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
("step_cls", "bucket", "feature_type"),
|
|
||||||
[
|
|
||||||
(ForwardKinematicsJointsToEEAction, PipelineFeatureType.ACTION, FeatureType.ACTION),
|
|
||||||
(ForwardKinematicsJointsToEEObservation, PipelineFeatureType.OBSERVATION, FeatureType.STATE),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
def test_fk_feature_schema(step_cls, bucket, feature_type):
|
|
||||||
features = {PipelineFeatureType.ACTION: {}, PipelineFeatureType.OBSERVATION: {}}
|
|
||||||
features[bucket] = _joint_bucket(feature_type)
|
|
||||||
out = step_cls(kinematics=None, motor_names=MOTOR_NAMES).transform_features(features)[bucket]
|
|
||||||
assert set(out) == EE_KEYS
|
|
||||||
assert {feature.type for feature in out.values()} == {feature_type}
|
|
||||||
@@ -149,3 +149,51 @@ def test_configure_writes_position_pid_coefficients():
|
|||||||
bus_mock.write.assert_any_call("P_Coefficient", "shoulder_pan", 32)
|
bus_mock.write.assert_any_call("P_Coefficient", "shoulder_pan", 32)
|
||||||
bus_mock.write.assert_any_call("I_Coefficient", "shoulder_pan", 1)
|
bus_mock.write.assert_any_call("I_Coefficient", "shoulder_pan", 1)
|
||||||
bus_mock.write.assert_any_call("D_Coefficient", "shoulder_pan", 16)
|
bus_mock.write.assert_any_call("D_Coefficient", "shoulder_pan", 16)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
"gripper_closed_pos, expected_drive_mode",
|
||||||
|
[
|
||||||
|
(2035, 0), # closed position at range_min -> raw increases when opening -> not inverted
|
||||||
|
(3528, 1), # closed position at range_max -> raw increases when closing -> inverted
|
||||||
|
(2781, None), # not near either end stop -> unsafe to infer
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_calibrate_detects_gripper_drive_mode(follower, gripper_closed_pos, expected_drive_mode):
|
||||||
|
"""Regression test for #3942: the follower gripper can be mounted mirrored with respect to the
|
||||||
|
leader's, in which case its raw position increases when closing. Calibration must detect this
|
||||||
|
and set drive_mode=1 so that normalized values follow the 0=closed/100=open convention."""
|
||||||
|
follower.connect()
|
||||||
|
|
||||||
|
motors = list(follower.bus.motors)
|
||||||
|
follower.bus.set_half_turn_homings.return_value = dict.fromkeys(motors, 0)
|
||||||
|
follower.bus.record_ranges_of_motion.return_value = (
|
||||||
|
dict.fromkeys(motors, 2035),
|
||||||
|
dict.fromkeys(motors, 3528),
|
||||||
|
)
|
||||||
|
follower.bus.read.return_value = gripper_closed_pos
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("builtins.input", return_value=""),
|
||||||
|
patch.object(type(follower), "_save_calibration", lambda self: None),
|
||||||
|
):
|
||||||
|
follower.calibration = {}
|
||||||
|
if expected_drive_mode is None:
|
||||||
|
with pytest.raises(ValueError, match="Gripper is not fully closed"):
|
||||||
|
follower.calibrate()
|
||||||
|
else:
|
||||||
|
follower.calibrate()
|
||||||
|
|
||||||
|
follower.bus.read.assert_called_with(
|
||||||
|
"Present_Position",
|
||||||
|
"gripper",
|
||||||
|
normalize=False,
|
||||||
|
num_retry=follower.config.num_read_retries,
|
||||||
|
)
|
||||||
|
if expected_drive_mode is None:
|
||||||
|
return
|
||||||
|
|
||||||
|
assert follower.calibration["gripper"].drive_mode == expected_drive_mode
|
||||||
|
for motor in motors:
|
||||||
|
if motor != "gripper":
|
||||||
|
assert follower.calibration[motor].drive_mode == 0
|
||||||
|
|||||||
Reference in New Issue
Block a user