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

Merge branch 'feat/add_pi' into feat/validate_pi_libero

Browse Source
This commit is contained in:
Pepijn
2025-09-13 17:54:36 +02:00
parent f0b969ae48 d0d714be47
commit de42da8225
5 changed files with 44 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/configs/policies.py
+3 -1
View File
@@ -72,9 +72,11 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
tags: list[str] | None = None tags: list[str] | None = None
# Add tags to your policy on the hub. # Add tags to your policy on the hub.
license: str | None = None license: str | None = None
# Either the repo ID of a model hosted on the Hub or a path to a directory containing weights
# saved using `Policy.save_pretrained`. If not provided, the policy is initialized from scratch.
pretrained_path: str | None = None
def __post_init__(self): def __post_init__(self):
self.pretrained_path = None
if not self.device or not is_torch_device_available(self.device): if not self.device or not is_torch_device_available(self.device):
auto_device = auto_select_torch_device() auto_device = auto_select_torch_device()
logging.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.") logging.warning(f"Device '{self.device}' is not available. Switching to '{auto_device}'.")
src/lerobot/policies/pi05_openpi/configuration_pi05openpi.py
+11 -28
View File
@@ -16,7 +16,7 @@
from dataclasses import dataclass, field from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.configs.types import NormalizationMode
from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
@@ -36,23 +36,20 @@ class PI05OpenPIConfig(PreTrainedConfig):
n_obs_steps: int = 1 n_obs_steps: int = 1
chunk_size: int = 50 # Number of action steps to predict, in openpi called "action_horizon" chunk_size: int = 50 # Number of action steps to predict, in openpi called "action_horizon"
n_action_steps: int = 50 # Number of action steps to execute n_action_steps: int = 50 # Number of action steps to execute
action_dim: int = 32 # Action dimension (will be padded to 32)
state_dim: int = 32 # State dimension (will be padded to 32) # Shorter state and action vectors will be padded to these dimensions
max_state_dim: int = 32 # State dimension (will be padded to 32)
max_action_dim: int = 32 # Action dimension (will be padded to 32)
# Flow matching parameters: see openpi `PI0Pytorch` # Flow matching parameters: see openpi `PI0Pytorch`
num_inference_steps: int = 10 # Number of denoising steps during inference num_inference_steps: int = 10 # Number of denoising steps during inference
time_sampling_beta_alpha: float = 1.5 # Beta distribution alpha parameter for time sampling time_sampling_beta_alpha: float = 1.5 # Beta distribution alpha parameter for time sampling
time_sampling_beta_beta: float = 1.0 # Beta distribution beta parameter for time sampling time_sampling_beta_beta: float = 1.0 # Beta distribution beta parameter for time sampling
min_period: float = 4e-3 # Min period for sinusoidal positional encoding min_period: float = 4e-3 # Min period for sinusoidal positional encoding
max_period: float = 4.0 # Max period for sinusoidal positional encodingis my max_period: float = 4.0 # Max period for sinusoidal positional encoding
# Image preprocessing # Image preprocessing
image_resolution: tuple[int, int] = (224, 224) # see openpi `preprocessing_pytorch.py` image_resolution: tuple[int, int] = (224, 224) # see openpi `preprocessing_pytorch.py`
image_keys: tuple[str, ...] = (
"observation.images.base_0_rgb",
"observation.images.left_wrist_0_rgb",
"observation.images.right_wrist_0_rgb",
)
# Normalization # Normalization
normalization_mapping: dict[str, NormalizationMode] = field( normalization_mapping: dict[str, NormalizationMode] = field(
@@ -103,26 +100,12 @@ class PI05OpenPIConfig(PreTrainedConfig):
def validate_features(self) -> None: def validate_features(self) -> None:
"""Validate and set up input/output features.""" """Validate and set up input/output features."""
# Add image features # Image features are now handled dynamically through dataset configuration
for key in self.image_keys: # No need to auto-add hardcoded image keys
if key not in self.input_features:
self.input_features[key] = PolicyFeature(
type=FeatureType.VISUAL,
shape=(3, 224, 224), # Default shape, will be resized
)
# Ensure state and action features exist # State and action features are also handled dynamically through dataset configuration
if "observation.state" not in self.input_features: # The actual dimensions come from the feature shapes, max dimensions are used for padding only
self.input_features["observation.state"] = PolicyFeature( pass
type=FeatureType.STATE,
shape=(self.state_dim,),
)
if "action" not in self.output_features:
self.output_features["action"] = PolicyFeature(
type=FeatureType.ACTION,
shape=(self.action_dim,),
)
def get_optimizer_preset(self) -> AdamWConfig: def get_optimizer_preset(self) -> AdamWConfig:
return AdamWConfig( return AdamWConfig(
src/lerobot/policies/pi05_openpi/modeling_pi05openpi.py
+10 -10
View File
@@ -503,8 +503,8 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
precision=config.dtype, precision=config.dtype,
) )
self.action_in_proj = nn.Linear(config.action_dim, action_expert_config.width) self.action_in_proj = nn.Linear(config.max_action_dim, action_expert_config.width)
self.action_out_proj = nn.Linear(action_expert_config.width, config.action_dim) self.action_out_proj = nn.Linear(action_expert_config.width, config.max_action_dim)
self.time_mlp_in = nn.Linear(action_expert_config.width, action_expert_config.width) self.time_mlp_in = nn.Linear(action_expert_config.width, action_expert_config.width)
self.time_mlp_out = nn.Linear(action_expert_config.width, action_expert_config.width) self.time_mlp_out = nn.Linear(action_expert_config.width, action_expert_config.width)
@@ -739,8 +739,8 @@ $(python -c "import transformers, os; print(os.path.dirname(transformers.__file_
actions_shape = ( actions_shape = (
bsize, bsize,
self.config.chunk_size, self.config.chunk_size,
self.config.action_dim, self.config.max_action_dim,
) # Use config action_dim for internal processing ) # Use config max_action_dim for internal processing
noise = self.sample_noise(actions_shape, device) noise = self.sample_noise(actions_shape, device)
prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix( prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
@@ -875,8 +875,8 @@ class PI05OpenPIPolicy(PreTrainedPolicy):
if pretrained_name_or_path is None: if pretrained_name_or_path is None:
raise ValueError("pretrained_name_or_path is required") raise ValueError("pretrained_name_or_path is required")
# Create default config # Use provided config if available, otherwise create default config
config = cls.config_class() config = kwargs.get("config", cls.config_class())
# Initialize model without loading weights # Initialize model without loading weights
# Check if dataset_stats were provided in kwargs # Check if dataset_stats were provided in kwargs
@@ -1236,12 +1236,12 @@ class PI05OpenPIPolicy(PreTrainedPolicy):
def prepare_state(self, batch): # see lerobot pi0 `prepare_state` (exact copy) def prepare_state(self, batch): # see lerobot pi0 `prepare_state` (exact copy)
"""Pad state""" """Pad state"""
state = pad_vector(batch[OBS_STATE], self.config.state_dim) state = pad_vector(batch[OBS_STATE], self.config.max_state_dim)
return state return state
def prepare_action(self, batch): # see lerobot pi0 `prepare_action` (exact copy) def prepare_action(self, batch): # see lerobot pi0 `prepare_action` (exact copy)
"""Pad action""" """Pad action"""
actions = pad_vector(batch[ACTION], self.config.action_dim) actions = pad_vector(batch[ACTION], self.config.max_action_dim)
return actions return actions
@torch.no_grad() @torch.no_grad()
@@ -1295,8 +1295,8 @@ class PI05OpenPIPolicy(PreTrainedPolicy):
losses = self.model.forward(images, img_masks, lang_tokens, lang_masks, state, actions) losses = self.model.forward(images, img_masks, lang_tokens, lang_masks, state, actions)
# Truncate losses to actual action dimensions # Truncate losses to actual action dimensions
if self.config.action_dim < 32: original_action_dim = self.config.output_features[ACTION].shape[0]
losses = losses[:, :, : self.config.action_dim] losses = losses[:, :, :original_action_dim]
loss = losses.mean() loss = losses.mean()
src/lerobot/policies/pi0_openpi/configuration_pi0openpi.py
+11 -28
View File
@@ -16,7 +16,7 @@
from dataclasses import dataclass, field from dataclasses import dataclass, field
from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.policies import PreTrainedConfig
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.configs.types import NormalizationMode
from lerobot.optim.optimizers import AdamWConfig from lerobot.optim.optimizers import AdamWConfig
from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig from lerobot.optim.schedulers import CosineDecayWithWarmupSchedulerConfig
@@ -33,23 +33,20 @@ class PI0OpenPIConfig(PreTrainedConfig):
n_obs_steps: int = 1 n_obs_steps: int = 1
chunk_size: int = 50 # Number of action steps to predict, in openpi called "action_horizon" chunk_size: int = 50 # Number of action steps to predict, in openpi called "action_horizon"
n_action_steps: int = 50 # Number of action steps to execute n_action_steps: int = 50 # Number of action steps to execute
action_dim: int = 32 # Action dimension (will be padded to 32)
state_dim: int = 32 # State dimension (will be padded to 32) # Shorter state and action vectors will be padded to these dimensions
max_state_dim: int = 32 # State dimension (will be padded to 32)
max_action_dim: int = 32 # Action dimension (will be padded to 32)
# Flow matching parameters: see openpi `PI0Pytorch` # Flow matching parameters: see openpi `PI0Pytorch`
num_inference_steps: int = 10 # Number of denoising steps during inference num_inference_steps: int = 10 # Number of denoising steps during inference
time_sampling_beta_alpha: float = 1.5 # Beta distribution alpha parameter for time sampling time_sampling_beta_alpha: float = 1.5 # Beta distribution alpha parameter for time sampling
time_sampling_beta_beta: float = 1.0 # Beta distribution beta parameter for time sampling time_sampling_beta_beta: float = 1.0 # Beta distribution beta parameter for time sampling
min_period: float = 4e-3 # Min period for sinusoidal positional encoding min_period: float = 4e-3 # Min period for sinusoidal positional encoding
max_period: float = 4.0 # Max period for sinusoidal positional encodingis my max_period: float = 4.0 # Max period for sinusoidal positional encoding
# Image preprocessing # Image preprocessing
image_resolution: tuple[int, int] = (224, 224) # see openpi `preprocessing_pytorch.py` image_resolution: tuple[int, int] = (224, 224) # see openpi `preprocessing_pytorch.py`
image_keys: tuple[str, ...] = (
"observation.images.base_0_rgb",
"observation.images.left_wrist_0_rgb",
"observation.images.right_wrist_0_rgb",
)
# Normalization # Normalization
normalization_mapping: dict[str, NormalizationMode] = field( normalization_mapping: dict[str, NormalizationMode] = field(
@@ -100,26 +97,12 @@ class PI0OpenPIConfig(PreTrainedConfig):
def validate_features(self) -> None: def validate_features(self) -> None:
"""Validate and set up input/output features.""" """Validate and set up input/output features."""
# Add image features # Image features are now handled dynamically through dataset configuration
for key in self.image_keys: # No need to auto-add hardcoded image keys
if key not in self.input_features:
self.input_features[key] = PolicyFeature(
type=FeatureType.VISUAL,
shape=(3, 224, 224), # Default shape, will be resized
)
# Ensure state and action features exist # State and action features are also handled dynamically through dataset configuration
if "observation.state" not in self.input_features: # The actual dimensions come from the feature shapes, max dimensions are used for padding only
self.input_features["observation.state"] = PolicyFeature( pass
type=FeatureType.STATE,
shape=(self.state_dim,),
)
if "action" not in self.output_features:
self.output_features["action"] = PolicyFeature(
type=FeatureType.ACTION,
shape=(self.action_dim,),
)
def get_optimizer_preset(self) -> AdamWConfig: def get_optimizer_preset(self) -> AdamWConfig:
return AdamWConfig( return AdamWConfig(
src/lerobot/policies/pi0_openpi/modeling_pi0openpi.py
+9 -9
View File
@@ -503,10 +503,10 @@ class PI0Pytorch(nn.Module): # see openpi `PI0Pytorch`
precision=config.dtype, precision=config.dtype,
) )
self.action_in_proj = nn.Linear(config.action_dim, action_expert_config.width) self.action_in_proj = nn.Linear(config.max_action_dim, action_expert_config.width)
self.action_out_proj = nn.Linear(action_expert_config.width, config.action_dim) self.action_out_proj = nn.Linear(action_expert_config.width, config.max_action_dim)
self.state_proj = nn.Linear(config.state_dim, action_expert_config.width) self.state_proj = nn.Linear(config.max_state_dim, action_expert_config.width)
self.action_time_mlp_in = nn.Linear(2 * action_expert_config.width, action_expert_config.width) self.action_time_mlp_in = nn.Linear(2 * action_expert_config.width, action_expert_config.width)
self.action_time_mlp_out = nn.Linear(action_expert_config.width, action_expert_config.width) self.action_time_mlp_out = nn.Linear(action_expert_config.width, action_expert_config.width)
@@ -758,8 +758,8 @@ $(python -c "import transformers, os; print(os.path.dirname(transformers.__file_
actions_shape = ( actions_shape = (
bsize, bsize,
self.config.chunk_size, self.config.chunk_size,
self.config.action_dim, self.config.max_action_dim,
) # Use config action_dim for internal processing ) # Use config max_action_dim for internal processing
noise = self.sample_noise(actions_shape, device) noise = self.sample_noise(actions_shape, device)
prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix( prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
@@ -894,8 +894,8 @@ class PI0OpenPIPolicy(PreTrainedPolicy):
if pretrained_name_or_path is None: if pretrained_name_or_path is None:
raise ValueError("pretrained_name_or_path is required") raise ValueError("pretrained_name_or_path is required")
# Create default config # Use provided config if available, otherwise create default config
config = cls.config_class() config = kwargs.get("config", cls.config_class())
# Initialize model without loading weights # Initialize model without loading weights
# Check if dataset_stats were provided in kwargs # Check if dataset_stats were provided in kwargs
@@ -1251,12 +1251,12 @@ class PI0OpenPIPolicy(PreTrainedPolicy):
def prepare_state(self, batch): # see lerobot pi0 `prepare_state` (exact copy) def prepare_state(self, batch): # see lerobot pi0 `prepare_state` (exact copy)
"""Pad state""" """Pad state"""
state = pad_vector(batch[OBS_STATE], self.config.state_dim) state = pad_vector(batch[OBS_STATE], self.config.max_state_dim)
return state return state
def prepare_action(self, batch): # see lerobot pi0 `prepare_action` (exact copy) def prepare_action(self, batch): # see lerobot pi0 `prepare_action` (exact copy)
"""Pad action""" """Pad action"""
actions = pad_vector(batch[ACTION], self.config.action_dim) actions = pad_vector(batch[ACTION], self.config.max_action_dim)
return actions return actions
@torch.no_grad() @torch.no_grad()