mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-18 15:31:47 +00:00
refactor pi052 to reuse pi05
This commit is contained in:
Pepijn
@@ -15,6 +15,7 @@
|
||||
# limitations under the License.
|
||||
|
||||
import builtins
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
from collections import deque
|
||||
@@ -23,6 +24,7 @@ from typing import TYPE_CHECKING, Literal, TypedDict, Unpack
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F # noqa: N812
|
||||
from safetensors.torch import load_file
|
||||
from torch import Tensor, nn
|
||||
|
||||
from lerobot.utils.import_utils import _transformers_available, require_package
|
||||
@@ -32,6 +34,7 @@ if TYPE_CHECKING or _transformers_available:
|
||||
from transformers.cache_utils import DynamicCache
|
||||
from transformers.models.auto import CONFIG_MAPPING
|
||||
from transformers.models.gemma import modeling_gemma
|
||||
from transformers.utils import cached_file
|
||||
|
||||
from ..pi_gemma import (
|
||||
PaliGemmaForConditionalGenerationWithPiGemma,
|
||||
@@ -47,6 +50,7 @@ else:
|
||||
_gated_residual = None
|
||||
layernorm_forward = None
|
||||
PaliGemmaForConditionalGenerationWithPiGemma = None
|
||||
cached_file = None
|
||||
from lerobot.configs import PreTrainedConfig
|
||||
from lerobot.utils.constants import (
|
||||
ACTION,
|
||||
@@ -66,6 +70,84 @@ class ActionSelectKwargs(TypedDict, total=False):
|
||||
execution_horizon: int | None
|
||||
|
||||
|
||||
_SAFETENSORS_FILE = "model.safetensors"
|
||||
_SAFETENSORS_INDEX = "model.safetensors.index.json"
|
||||
|
||||
|
||||
def _resolve_weight_files(
|
||||
pretrained_name_or_path: str | Path,
|
||||
*,
|
||||
force_download: bool,
|
||||
resume_download: bool | None,
|
||||
proxies: dict | None,
|
||||
token: str | bool | None,
|
||||
cache_dir: str | Path | None,
|
||||
local_files_only: bool,
|
||||
revision: str | None,
|
||||
) -> list[Path]:
|
||||
model_id = str(pretrained_name_or_path)
|
||||
local_dir = Path(model_id)
|
||||
load_kwargs = {
|
||||
"revision": revision,
|
||||
"cache_dir": cache_dir,
|
||||
"force_download": force_download,
|
||||
"resume_download": resume_download,
|
||||
"proxies": proxies,
|
||||
"token": token,
|
||||
"local_files_only": local_files_only,
|
||||
}
|
||||
|
||||
if local_dir.is_dir():
|
||||
index_path = local_dir / _SAFETENSORS_INDEX
|
||||
single_path = local_dir / _SAFETENSORS_FILE
|
||||
else:
|
||||
resolved_index = cached_file(
|
||||
model_id,
|
||||
_SAFETENSORS_INDEX,
|
||||
_raise_exceptions_for_missing_entries=False,
|
||||
**load_kwargs,
|
||||
)
|
||||
index_path = Path(resolved_index) if resolved_index is not None else None
|
||||
single_path = None
|
||||
if index_path is None:
|
||||
resolved_file = cached_file(model_id, _SAFETENSORS_FILE, **load_kwargs)
|
||||
single_path = Path(resolved_file) if resolved_file is not None else None
|
||||
|
||||
if index_path is None or not index_path.is_file():
|
||||
if single_path is None or not single_path.is_file():
|
||||
raise FileNotFoundError(f"No {_SAFETENSORS_FILE} found in {model_id!r}.")
|
||||
return [single_path]
|
||||
|
||||
index = json.loads(index_path.read_text())
|
||||
shard_names = sorted(set(index.get("weight_map", {}).values()))
|
||||
if not shard_names:
|
||||
raise ValueError(f"Invalid safetensors index without a weight_map: {index_path}")
|
||||
if local_dir.is_dir():
|
||||
files = [local_dir / name for name in shard_names]
|
||||
else:
|
||||
files = []
|
||||
for name in shard_names:
|
||||
resolved_file = cached_file(model_id, name, **load_kwargs)
|
||||
if resolved_file is None:
|
||||
raise FileNotFoundError(f"Checkpoint shard {name!r} not found in {model_id!r}.")
|
||||
files.append(Path(resolved_file))
|
||||
missing = [str(path) for path in files if not path.is_file()]
|
||||
if missing:
|
||||
raise FileNotFoundError(f"Missing checkpoint shards: {missing}")
|
||||
return files
|
||||
|
||||
|
||||
def _load_weight_files(files: list[Path]) -> dict[str, Tensor]:
|
||||
state_dict: dict[str, Tensor] = {}
|
||||
for path in files:
|
||||
shard = load_file(path)
|
||||
overlap = state_dict.keys() & shard.keys()
|
||||
if overlap:
|
||||
raise ValueError(f"Duplicate checkpoint keys in {path}: {sorted(overlap)[:5]}")
|
||||
state_dict.update(shard)
|
||||
return state_dict
|
||||
|
||||
|
||||
def get_safe_dtype(target_dtype, device_type):
|
||||
"""Get a safe dtype for the given device type."""
|
||||
if device_type == "mps" and target_dtype == torch.float64:
|
||||
@@ -563,6 +645,12 @@ class PaliGemmaWithExpertModel(
|
||||
class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
"""Core PI05 PyTorch model."""
|
||||
|
||||
use_hf_vision_checkpointing_api = False
|
||||
checkpoint_vision_embeddings = True
|
||||
use_typed_attention_masks = False
|
||||
use_on_device_suffix_mask = False
|
||||
precompute_denoise_times = False
|
||||
|
||||
def __init__(self, config: PI05Config, rtc_processor: RTCProcessor | None = None):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
@@ -606,7 +694,11 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
"""Enable gradient checkpointing for memory optimization."""
|
||||
self.gradient_checkpointing_enabled = True
|
||||
self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = True
|
||||
self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing = True
|
||||
vision_tower = self.paligemma_with_expert.paligemma.model.vision_tower
|
||||
if self.use_hf_vision_checkpointing_api:
|
||||
vision_tower.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
|
||||
else:
|
||||
vision_tower.gradient_checkpointing = True
|
||||
self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = True
|
||||
logging.info("Enabled gradient checkpointing for PI05Pytorch model")
|
||||
|
||||
@@ -614,7 +706,11 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
"""Disable gradient checkpointing."""
|
||||
self.gradient_checkpointing_enabled = False
|
||||
self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = False
|
||||
self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing = False
|
||||
vision_tower = self.paligemma_with_expert.paligemma.model.vision_tower
|
||||
if self.use_hf_vision_checkpointing_api:
|
||||
vision_tower.gradient_checkpointing_disable()
|
||||
else:
|
||||
vision_tower.gradient_checkpointing = False
|
||||
self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = False
|
||||
logging.info("Disabled gradient checkpointing for PI05Pytorch model")
|
||||
|
||||
@@ -629,10 +725,13 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
)
|
||||
return func(*args, **kwargs)
|
||||
|
||||
def _prepare_attention_masks_4d(self, att_2d_masks):
|
||||
def _prepare_attention_masks_4d(self, att_2d_masks, dtype=None):
|
||||
"""Helper method to prepare 4D attention masks for transformer."""
|
||||
att_2d_masks_4d = att_2d_masks[:, None, :, :]
|
||||
return torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
|
||||
result = torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
|
||||
if dtype is not None:
|
||||
result = result.to(dtype=dtype)
|
||||
return result
|
||||
|
||||
def sample_noise(self, shape, device):
|
||||
return torch.normal(
|
||||
@@ -658,13 +757,16 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
pad_masks = []
|
||||
att_masks = []
|
||||
|
||||
# Process images
|
||||
for img, img_mask in zip(images, img_masks, strict=True):
|
||||
if self.checkpoint_vision_embeddings:
|
||||
|
||||
def image_embed_func(img):
|
||||
return self.paligemma_with_expert.embed_image(img)
|
||||
def embed_image(img):
|
||||
return self._apply_checkpoint(self.paligemma_with_expert.embed_image, img)
|
||||
|
||||
img_emb = self._apply_checkpoint(image_embed_func, img)
|
||||
img_embs = [embed_image(img) for img in images]
|
||||
else:
|
||||
img_embs = [self.paligemma_with_expert.embed_image(img) for img in images]
|
||||
|
||||
for img_emb, img_mask in zip(img_embs, img_masks, strict=True):
|
||||
bsize, num_img_embs = img_emb.shape[:2]
|
||||
|
||||
embs.append(img_emb)
|
||||
@@ -734,8 +836,14 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
|
||||
embs = torch.cat(embs, dim=1)
|
||||
pad_masks = torch.cat(pad_masks, dim=1)
|
||||
att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)
|
||||
att_masks = att_masks[None, :].expand(bsize, len(att_masks))
|
||||
if self.use_on_device_suffix_mask:
|
||||
n = len(att_masks)
|
||||
att_masks = torch.zeros(n, dtype=embs.dtype, device=embs.device)
|
||||
att_masks[0] = 1
|
||||
att_masks = att_masks[None, :].expand(bsize, n)
|
||||
else:
|
||||
att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)
|
||||
att_masks = att_masks[None, :].expand(bsize, len(att_masks))
|
||||
|
||||
return embs, pad_masks, att_masks, adarms_cond
|
||||
|
||||
@@ -819,7 +927,8 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
|
||||
prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
|
||||
|
||||
prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks)
|
||||
mask_dtype = prefix_embs.dtype if self.use_typed_attention_masks else None
|
||||
prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks, dtype=mask_dtype)
|
||||
self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager" # noqa: SLF001
|
||||
|
||||
_, past_key_values = self.paligemma_with_expert.forward(
|
||||
@@ -832,10 +941,19 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch`
|
||||
|
||||
dt = -1.0 / num_steps
|
||||
|
||||
times = None
|
||||
if self.precompute_denoise_times:
|
||||
times = torch.tensor(
|
||||
[1.0 + step * dt for step in range(num_steps)], dtype=torch.float32, device=device
|
||||
)
|
||||
|
||||
x_t = noise
|
||||
for step in range(num_steps):
|
||||
time = 1.0 + step * dt
|
||||
time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
|
||||
if times is None:
|
||||
time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
|
||||
else:
|
||||
time_tensor = times[step].expand(bsize)
|
||||
|
||||
def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
|
||||
return self.denoise_step(
|
||||
@@ -913,6 +1031,9 @@ class PI05Policy(PreTrainedPolicy):
|
||||
|
||||
config_class = PI05Config
|
||||
name = "pi05"
|
||||
model_class = PI05Pytorch
|
||||
eval_after_pretrained_load = False
|
||||
show_openpi_disclaimer = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -930,7 +1051,7 @@ class PI05Policy(PreTrainedPolicy):
|
||||
|
||||
# Initialize the core PI05 model
|
||||
self.init_rtc_processor()
|
||||
self.model = PI05Pytorch(config, rtc_processor=self.rtc_processor)
|
||||
self.model = self.model_class(config, rtc_processor=self.rtc_processor)
|
||||
|
||||
# Enable gradient checkpointing if requested
|
||||
if config.gradient_checkpointing:
|
||||
@@ -956,16 +1077,16 @@ class PI05Policy(PreTrainedPolicy):
|
||||
strict: bool = True,
|
||||
**kwargs,
|
||||
) -> T:
|
||||
"""Override the from_pretrained method to handle key remapping and display important disclaimer."""
|
||||
print(
|
||||
"The PI05 model is a direct port of the OpenPI implementation. \n"
|
||||
"This implementation follows the original OpenPI structure for compatibility. \n"
|
||||
"Original implementation: https://github.com/Physical-Intelligence/openpi"
|
||||
)
|
||||
"""Load PI05-compatible single-file or sharded safetensors checkpoints."""
|
||||
if cls.show_openpi_disclaimer:
|
||||
print(
|
||||
"The PI05 model is a direct port of the OpenPI implementation. \n"
|
||||
"This implementation follows the original OpenPI structure for compatibility. \n"
|
||||
"Original implementation: https://github.com/Physical-Intelligence/openpi"
|
||||
)
|
||||
if pretrained_name_or_path is None:
|
||||
raise ValueError("pretrained_name_or_path is required")
|
||||
|
||||
# Use provided config if available, otherwise create default config
|
||||
if config is None:
|
||||
config = PreTrainedConfig.from_pretrained(
|
||||
pretrained_name_or_path=pretrained_name_or_path,
|
||||
@@ -979,85 +1100,35 @@ class PI05Policy(PreTrainedPolicy):
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Initialize model without loading weights
|
||||
# Check if dataset_stats were provided in kwargs
|
||||
model = cls(config, **kwargs)
|
||||
|
||||
# Load state dict (expects keys with "model." prefix)
|
||||
try:
|
||||
print(f"Loading model from: {pretrained_name_or_path}")
|
||||
try:
|
||||
from transformers.utils import cached_file
|
||||
|
||||
resolved_file = cached_file(
|
||||
pretrained_name_or_path,
|
||||
"model.safetensors",
|
||||
cache_dir=kwargs.get("cache_dir"),
|
||||
force_download=kwargs.get("force_download", False),
|
||||
resume_download=kwargs.get("resume_download"),
|
||||
proxies=kwargs.get("proxies"),
|
||||
token=kwargs.get("token"),
|
||||
revision=kwargs.get("revision"),
|
||||
local_files_only=kwargs.get("local_files_only", False),
|
||||
)
|
||||
from safetensors.torch import load_file
|
||||
|
||||
original_state_dict = load_file(resolved_file)
|
||||
print("✓ Loaded state dict from model.safetensors")
|
||||
except Exception as e:
|
||||
print(f"Could not load state dict from remote files: {e}")
|
||||
print("Returning model without loading pretrained weights")
|
||||
return model
|
||||
|
||||
# First, fix any key differences (see openpi model.py, _fix_pytorch_state_dict_keys)
|
||||
fixed_state_dict = model._fix_pytorch_state_dict_keys(original_state_dict, model.config)
|
||||
|
||||
# Then add "model." prefix for all keys that don't already have it
|
||||
remapped_state_dict = {}
|
||||
remap_count = 0
|
||||
|
||||
for key, value in fixed_state_dict.items():
|
||||
if not key.startswith("model."):
|
||||
new_key = f"model.{key}"
|
||||
remapped_state_dict[new_key] = value
|
||||
remap_count += 1
|
||||
else:
|
||||
remapped_state_dict[key] = value
|
||||
|
||||
if remap_count > 0:
|
||||
print(f"Remapped {remap_count} state dict keys")
|
||||
|
||||
# Load the remapped state dict into the model
|
||||
missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)
|
||||
|
||||
if missing_keys:
|
||||
print(f"Missing keys when loading state dict: {len(missing_keys)} keys")
|
||||
if len(missing_keys) <= 5:
|
||||
for key in missing_keys:
|
||||
print(f" - {key}")
|
||||
else:
|
||||
for key in missing_keys[:5]:
|
||||
print(f" - {key}")
|
||||
print(f" ... and {len(missing_keys) - 5} more")
|
||||
|
||||
if unexpected_keys:
|
||||
print(f"Unexpected keys when loading state dict: {len(unexpected_keys)} keys")
|
||||
if len(unexpected_keys) <= 5:
|
||||
for key in unexpected_keys:
|
||||
print(f" - {key}")
|
||||
else:
|
||||
for key in unexpected_keys[:5]:
|
||||
print(f" - {key}")
|
||||
print(f" ... and {len(unexpected_keys) - 5} more")
|
||||
|
||||
if not missing_keys and not unexpected_keys:
|
||||
print("All keys loaded successfully!")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Warning: Could not load state dict: {e}")
|
||||
|
||||
files = _resolve_weight_files(
|
||||
pretrained_name_or_path,
|
||||
force_download=force_download,
|
||||
resume_download=resume_download,
|
||||
proxies=proxies,
|
||||
token=token,
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
revision=revision,
|
||||
)
|
||||
fixed_state_dict = model._fix_pytorch_state_dict_keys(_load_weight_files(files), model.config)
|
||||
remapped_state_dict = {
|
||||
key if key.startswith("model.") else f"model.{key}": value
|
||||
for key, value in fixed_state_dict.items()
|
||||
}
|
||||
remapped_state_dict = model._prepare_pretrained_state_dict(remapped_state_dict)
|
||||
missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)
|
||||
if missing_keys:
|
||||
logging.warning("Missing %s checkpoint keys: %s", cls.name, missing_keys)
|
||||
if unexpected_keys:
|
||||
logging.warning("Unexpected %s checkpoint keys: %s", cls.name, unexpected_keys)
|
||||
if model.eval_after_pretrained_load:
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
def _prepare_pretrained_state_dict(self, state_dict: dict[str, Tensor]) -> dict[str, Tensor]:
|
||||
return state_dict
|
||||
|
||||
def _fix_pytorch_state_dict_keys(
|
||||
self, state_dict, model_config
|
||||
): # see openpi `BaseModelConfig, _fix_pytorch_state_dict_keys`
|
||||
@@ -1228,12 +1299,16 @@ class PI05Policy(PreTrainedPolicy):
|
||||
|
||||
# Action queue logic for n_action_steps > 1
|
||||
if len(self._action_queue) == 0:
|
||||
actions = self.predict_action_chunk(batch)[:, : self.config.n_action_steps]
|
||||
action_batch = self._prepare_action_batch(batch)
|
||||
actions = self.predict_action_chunk(action_batch)[:, : self.config.n_action_steps]
|
||||
# Transpose to get shape (n_action_steps, batch_size, action_dim)
|
||||
self._action_queue.extend(actions.transpose(0, 1))
|
||||
|
||||
return self._action_queue.popleft()
|
||||
|
||||
def _prepare_action_batch(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
|
||||
return batch
|
||||
|
||||
@torch.no_grad()
|
||||
def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
|
||||
"""Predict a chunk of actions given environment observations."""
|
||||
|
||||
@@ -16,237 +16,41 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import types
|
||||
from collections import deque
|
||||
from contextlib import nullcontext
|
||||
from pathlib import Path
|
||||
from typing import Any, Unpack
|
||||
|
||||
import torch
|
||||
from safetensors.torch import load_file
|
||||
from torch import Tensor
|
||||
from torch.nn import functional
|
||||
from transformers.utils import cached_file
|
||||
|
||||
from lerobot.configs import PreTrainedConfig
|
||||
from lerobot.utils.constants import (
|
||||
ACTION,
|
||||
OBS_LANGUAGE_ATTENTION_MASK,
|
||||
OBS_LANGUAGE_TOKENS,
|
||||
OBS_STATE,
|
||||
OPENPI_ATTENTION_MASK_VALUE,
|
||||
)
|
||||
from lerobot.utils.import_utils import require_package
|
||||
|
||||
from ..pi05.modeling_pi05 import (
|
||||
ActionSelectKwargs,
|
||||
PI05Policy,
|
||||
PI05Pytorch as PI05PytorchBase,
|
||||
create_sinusoidal_pos_embedding,
|
||||
make_att_2d_masks,
|
||||
)
|
||||
from ..pretrained import PreTrainedPolicy, T
|
||||
from .configuration_pi052 import PI052Config
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_SAFETENSORS_FILE = "model.safetensors"
|
||||
_SAFETENSORS_INDEX = "model.safetensors.index.json"
|
||||
|
||||
|
||||
def _resolve_weight_files(
|
||||
pretrained_name_or_path: str | Path,
|
||||
*,
|
||||
force_download: bool,
|
||||
resume_download: bool | None,
|
||||
proxies: dict | None,
|
||||
token: str | bool | None,
|
||||
cache_dir: str | Path | None,
|
||||
local_files_only: bool,
|
||||
revision: str | None,
|
||||
) -> list[Path]:
|
||||
model_id = str(pretrained_name_or_path)
|
||||
local_dir = Path(model_id)
|
||||
load_kwargs = {
|
||||
"revision": revision,
|
||||
"cache_dir": cache_dir,
|
||||
"force_download": force_download,
|
||||
"resume_download": resume_download,
|
||||
"proxies": proxies,
|
||||
"token": token,
|
||||
"local_files_only": local_files_only,
|
||||
}
|
||||
|
||||
if local_dir.is_dir():
|
||||
index_path = local_dir / _SAFETENSORS_INDEX
|
||||
single_path = local_dir / _SAFETENSORS_FILE
|
||||
else:
|
||||
resolved_index = cached_file(
|
||||
model_id,
|
||||
_SAFETENSORS_INDEX,
|
||||
_raise_exceptions_for_missing_entries=False,
|
||||
**load_kwargs,
|
||||
)
|
||||
index_path = Path(resolved_index) if resolved_index is not None else None
|
||||
single_path = None
|
||||
if index_path is None:
|
||||
resolved_file = cached_file(model_id, _SAFETENSORS_FILE, **load_kwargs)
|
||||
single_path = Path(resolved_file) if resolved_file is not None else None
|
||||
|
||||
if index_path is None or not index_path.is_file():
|
||||
if single_path is None or not single_path.is_file():
|
||||
raise FileNotFoundError(f"No {_SAFETENSORS_FILE} found in {model_id!r}.")
|
||||
return [single_path]
|
||||
|
||||
index = json.loads(index_path.read_text())
|
||||
shard_names = sorted(set(index.get("weight_map", {}).values()))
|
||||
if not shard_names:
|
||||
raise ValueError(f"Invalid safetensors index without a weight_map: {index_path}")
|
||||
if local_dir.is_dir():
|
||||
files = [local_dir / name for name in shard_names]
|
||||
else:
|
||||
files = []
|
||||
for name in shard_names:
|
||||
resolved_file = cached_file(model_id, name, **load_kwargs)
|
||||
if resolved_file is None:
|
||||
raise FileNotFoundError(f"Checkpoint shard {name!r} not found in {model_id!r}.")
|
||||
files.append(Path(resolved_file))
|
||||
missing = [str(path) for path in files if not path.is_file()]
|
||||
if missing:
|
||||
raise FileNotFoundError(f"Missing checkpoint shards: {missing}")
|
||||
return files
|
||||
|
||||
|
||||
def _load_weight_files(files: list[Path]) -> dict[str, Tensor]:
|
||||
state_dict: dict[str, Tensor] = {}
|
||||
for path in files:
|
||||
shard = load_file(path)
|
||||
overlap = state_dict.keys() & shard.keys()
|
||||
if overlap:
|
||||
raise ValueError(f"Duplicate checkpoint keys in {path}: {sorted(overlap)[:5]}")
|
||||
state_dict.update(shard)
|
||||
return state_dict
|
||||
|
||||
|
||||
class PI05Pytorch(PI05PytorchBase): # see openpi `PI0Pytorch`
|
||||
"""Core PI05 PyTorch model."""
|
||||
|
||||
def gradient_checkpointing_enable(self):
|
||||
"""Enable gradient checkpointing for memory optimization."""
|
||||
self.gradient_checkpointing_enabled = True
|
||||
self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = True
|
||||
self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing_enable(
|
||||
gradient_checkpointing_kwargs={"use_reentrant": False}
|
||||
)
|
||||
self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = True
|
||||
logging.info("Enabled gradient checkpointing for PI05Pytorch model")
|
||||
|
||||
def gradient_checkpointing_disable(self):
|
||||
"""Disable gradient checkpointing."""
|
||||
self.gradient_checkpointing_enabled = False
|
||||
self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = False
|
||||
self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing_disable()
|
||||
self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = False
|
||||
logging.info("Disabled gradient checkpointing for PI05Pytorch model")
|
||||
|
||||
def _prepare_attention_masks_4d(self, att_2d_masks, dtype=None):
|
||||
"""Helper method to prepare 4D attention masks for transformer."""
|
||||
att_2d_masks_4d = att_2d_masks[:, None, :, :]
|
||||
result = torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
|
||||
if dtype is not None:
|
||||
result = result.to(dtype=dtype)
|
||||
return result
|
||||
|
||||
def embed_prefix(
|
||||
self, images, img_masks, tokens, masks
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Embed images with SigLIP and language tokens with embedding layer."""
|
||||
embs = []
|
||||
pad_masks = []
|
||||
att_masks = []
|
||||
|
||||
# SigLIP checkpoints its encoder layers internally. An outer tower
|
||||
# checkpoint would recreate every layer activation at once in backward.
|
||||
img_embs = [self.paligemma_with_expert.embed_image(img) for img in images]
|
||||
|
||||
for img_emb, img_mask in zip(img_embs, img_masks, strict=True):
|
||||
bsize, num_img_embs = img_emb.shape[:2]
|
||||
embs.append(img_emb)
|
||||
pad_masks.append(img_mask[:, None].expand(bsize, num_img_embs))
|
||||
att_masks += [0] * num_img_embs
|
||||
|
||||
# Process language tokens
|
||||
def lang_embed_func(tokens):
|
||||
# GemmaTextScaledWordEmbedding already applies sqrt(hidden_size); do not scale twice.
|
||||
return self.paligemma_with_expert.embed_language_tokens(tokens)
|
||||
|
||||
lang_emb = self._apply_checkpoint(lang_embed_func, tokens)
|
||||
embs.append(lang_emb)
|
||||
pad_masks.append(masks)
|
||||
|
||||
num_lang_embs = lang_emb.shape[1]
|
||||
att_masks += [0] * num_lang_embs
|
||||
|
||||
embs = torch.cat(embs, dim=1)
|
||||
pad_masks = torch.cat(pad_masks, dim=1)
|
||||
att_masks = torch.tensor(att_masks, dtype=torch.bool, device=pad_masks.device)
|
||||
|
||||
bsize = pad_masks.shape[0]
|
||||
att_masks = att_masks[None, :].expand(bsize, len(att_masks))
|
||||
|
||||
return embs, pad_masks, att_masks
|
||||
|
||||
def embed_suffix(self, noisy_actions, timestep):
|
||||
"""Embed noisy_actions, timestep to prepare for Expert Gemma processing."""
|
||||
embs = []
|
||||
pad_masks = []
|
||||
att_masks = []
|
||||
|
||||
# Embed timestep using sine-cosine positional encoding
|
||||
time_emb = create_sinusoidal_pos_embedding(
|
||||
timestep,
|
||||
self.action_in_proj.out_features,
|
||||
min_period=self.config.min_period,
|
||||
max_period=self.config.max_period,
|
||||
device=timestep.device,
|
||||
)
|
||||
time_emb = time_emb.type(dtype=timestep.dtype)
|
||||
|
||||
# Fuse timestep + action information using an MLP
|
||||
def action_proj_func(noisy_actions):
|
||||
return self.action_in_proj(noisy_actions)
|
||||
|
||||
action_emb = self._apply_checkpoint(action_proj_func, noisy_actions)
|
||||
|
||||
def time_mlp_func(time_emb):
|
||||
x = self.time_mlp_in(time_emb)
|
||||
x = functional.silu(x)
|
||||
x = self.time_mlp_out(x)
|
||||
return functional.silu(x)
|
||||
|
||||
time_emb = self._apply_checkpoint(time_mlp_func, time_emb)
|
||||
action_time_emb = action_emb
|
||||
adarms_cond = time_emb
|
||||
|
||||
embs.append(action_time_emb)
|
||||
bsize, action_time_dim = action_time_emb.shape[:2]
|
||||
action_time_mask = torch.ones(bsize, action_time_dim, dtype=torch.bool, device=timestep.device)
|
||||
pad_masks.append(action_time_mask)
|
||||
|
||||
# Set attention masks so that image, language and state inputs do not attend to action tokens
|
||||
att_masks += [1] + ([0] * (self.config.chunk_size - 1))
|
||||
|
||||
embs = torch.cat(embs, dim=1)
|
||||
pad_masks = torch.cat(pad_masks, dim=1)
|
||||
# Build the constant suffix mask on-device to avoid a per-step host sync.
|
||||
n = len(att_masks)
|
||||
att_masks = torch.zeros(n, dtype=embs.dtype, device=embs.device)
|
||||
att_masks[0] = 1
|
||||
att_masks = att_masks[None, :].expand(bsize, n)
|
||||
|
||||
return embs, pad_masks, att_masks, adarms_cond
|
||||
use_hf_vision_checkpointing_api = True
|
||||
checkpoint_vision_embeddings = False
|
||||
use_typed_attention_masks = True
|
||||
use_on_device_suffix_mask = True
|
||||
precompute_denoise_times = True
|
||||
|
||||
def forward(self, images, img_masks, tokens, masks, actions, noise, time) -> Tensor:
|
||||
"""Do a full training forward pass and compute the loss."""
|
||||
@@ -299,91 +103,6 @@ class PI05Pytorch(PI05PytorchBase): # see openpi `PI0Pytorch`
|
||||
|
||||
return functional.mse_loss(u_t, v_t, reduction="none")
|
||||
|
||||
@torch.no_grad() # see openpi `sample_actions` (slightly adapted)
|
||||
def sample_actions(
|
||||
self,
|
||||
images,
|
||||
img_masks,
|
||||
tokens,
|
||||
masks,
|
||||
noise=None,
|
||||
num_steps=None,
|
||||
**kwargs: Unpack[ActionSelectKwargs],
|
||||
) -> Tensor:
|
||||
"""Do a full inference forward and compute the action."""
|
||||
if num_steps is None:
|
||||
num_steps = self.config.num_inference_steps
|
||||
|
||||
bsize = tokens.shape[0]
|
||||
device = tokens.device
|
||||
|
||||
if noise is None:
|
||||
# Sample noise with padded dimension as expected by action_in_proj
|
||||
actions_shape = (
|
||||
bsize,
|
||||
self.config.chunk_size,
|
||||
self.config.max_action_dim,
|
||||
) # Use config max_action_dim for internal processing
|
||||
noise = self.sample_noise(actions_shape, device)
|
||||
|
||||
prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, tokens, masks)
|
||||
prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
|
||||
prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
|
||||
|
||||
prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(
|
||||
prefix_att_2d_masks, dtype=prefix_embs.dtype
|
||||
)
|
||||
self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager" # noqa: SLF001
|
||||
|
||||
_, past_key_values = self.paligemma_with_expert.forward(
|
||||
attention_mask=prefix_att_2d_masks_4d,
|
||||
position_ids=prefix_position_ids,
|
||||
past_key_values=None,
|
||||
inputs_embeds=[prefix_embs, None],
|
||||
use_cache=True,
|
||||
)
|
||||
|
||||
dt = -1.0 / num_steps
|
||||
|
||||
# Precompute timesteps on-device to avoid a host sync per denoising step.
|
||||
times = torch.tensor([1.0 + s * dt for s in range(num_steps)], dtype=torch.float32, device=device)
|
||||
|
||||
x_t = noise
|
||||
for step in range(num_steps):
|
||||
time = 1.0 + step * dt # Python float kept for the RTC branch below
|
||||
time_tensor = times[step].expand(bsize)
|
||||
|
||||
def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
|
||||
return self.denoise_step(
|
||||
prefix_pad_masks=prefix_pad_masks,
|
||||
past_key_values=past_key_values,
|
||||
x_t=input_x_t,
|
||||
timestep=current_timestep,
|
||||
)
|
||||
|
||||
if self._rtc_enabled():
|
||||
inference_delay = kwargs.get("inference_delay")
|
||||
prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
|
||||
execution_horizon = kwargs.get("execution_horizon")
|
||||
|
||||
v_t = self.rtc_processor.denoise_step(
|
||||
x_t=x_t,
|
||||
prev_chunk_left_over=prev_chunk_left_over,
|
||||
inference_delay=inference_delay,
|
||||
time=time,
|
||||
original_denoise_step_partial=denoise_step_partial_call,
|
||||
execution_horizon=execution_horizon,
|
||||
)
|
||||
else:
|
||||
v_t = denoise_step_partial_call(x_t)
|
||||
|
||||
x_t = x_t + dt * v_t
|
||||
|
||||
if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled():
|
||||
self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
|
||||
|
||||
return x_t
|
||||
|
||||
def denoise_step(
|
||||
self,
|
||||
prefix_pad_masks,
|
||||
@@ -1094,21 +813,14 @@ class PI052Policy(PI05Policy):
|
||||
|
||||
config_class = PI052Config
|
||||
name = "pi052"
|
||||
model_class = PI05Pytorch
|
||||
eval_after_pretrained_load = True
|
||||
show_openpi_disclaimer = False
|
||||
|
||||
def __init__(self, config: PI052Config, **kwargs: Any) -> None:
|
||||
# Patch before constructing Gemma/SigLIP layers; the operation is optional and idempotent.
|
||||
_enable_hf_kernels()
|
||||
|
||||
require_package("transformers", extra="pi")
|
||||
PreTrainedPolicy.__init__(self, config)
|
||||
config.validate_features()
|
||||
self.config = config
|
||||
self.init_rtc_processor()
|
||||
self.model = PI05Pytorch(config, rtc_processor=self.rtc_processor)
|
||||
if config.gradient_checkpointing:
|
||||
self.model.gradient_checkpointing_enable()
|
||||
self.model.to(config.device)
|
||||
self.reset()
|
||||
super().__init__(config, **kwargs)
|
||||
|
||||
# Re-enable layers PI0.5 freezes when text supervision is requested.
|
||||
if config.text_loss_weight > 0 and config.unfreeze_lm_head:
|
||||
@@ -1155,19 +867,9 @@ class PI052Policy(PI05Policy):
|
||||
persistent=False,
|
||||
)
|
||||
|
||||
# Size per-environment inference state lazily.
|
||||
self.last_subtasks: list[str] | None = None
|
||||
self.last_subtasks_raw: list[str] | None = None
|
||||
self.last_subtasks_source: list[str] | None = None
|
||||
self._last_good_subtasks: list[str | None] | None = None
|
||||
|
||||
def reset(self):
|
||||
"""Reset action and high-level inference state."""
|
||||
# inlined PI05Policy.reset
|
||||
self._action_queue = deque(maxlen=self.config.n_action_steps)
|
||||
self._queues = {
|
||||
ACTION: deque(maxlen=self.config.n_action_steps),
|
||||
}
|
||||
super().reset()
|
||||
self.last_subtasks = None
|
||||
self.last_subtasks_raw = None
|
||||
self.last_subtasks_source = None
|
||||
@@ -1227,7 +929,7 @@ class PI052Policy(PI05Policy):
|
||||
and predict_actions_t is None
|
||||
and not getattr(self.config, "enable_fast_action_loss", False)
|
||||
):
|
||||
return self._pi05_flow_forward(batch, reduction=reduction)
|
||||
return super().forward(batch, reduction=reduction)
|
||||
|
||||
# Compute the host-side action-routing decision once for both flow and FAST.
|
||||
predict_any = predict_actions_t is None or bool(predict_actions_t.any().item())
|
||||
@@ -1413,7 +1115,6 @@ class PI052Policy(PI05Policy):
|
||||
reduction: str = "mean",
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
"""Run the single-repeat combined prefix and action path."""
|
||||
from lerobot.utils.constants import ACTION # noqa: PLC0415
|
||||
|
||||
noise = self.model.sample_noise(actions.shape, actions.device)
|
||||
time = self.model.sample_time(actions.shape[0], actions.device)
|
||||
@@ -1499,7 +1200,6 @@ class PI052Policy(PI05Policy):
|
||||
reduction: str = "mean",
|
||||
) -> tuple[Tensor, Tensor]:
|
||||
"""Run K independent action draws against one shared VLM prefix."""
|
||||
from lerobot.utils.constants import ACTION # noqa: PLC0415
|
||||
|
||||
model = self.model
|
||||
k = num_repeats
|
||||
@@ -1892,30 +1592,7 @@ class PI052Policy(PI05Policy):
|
||||
self._last_select_message_debug = ""
|
||||
return decoded
|
||||
|
||||
@torch.no_grad()
|
||||
def select_action(self, batch: dict[str, Tensor]) -> Tensor:
|
||||
"""Select an action via PI052's high-level → low-level inference path.
|
||||
|
||||
At action-chunk boundaries, first generate a low-level subtask from
|
||||
the high-level task prompt. Then retokenize that subtask as the
|
||||
low-level action prompt before sampling the action chunk. This keeps
|
||||
the public policy API identical to PI05 (`Tensor` action out), while
|
||||
matching the PI052 training/runtime conditioning more closely.
|
||||
"""
|
||||
assert not self._rtc_enabled(), (
|
||||
"RTC is not supported for select_action, use it with predict_action_chunk"
|
||||
)
|
||||
|
||||
self.eval()
|
||||
|
||||
if len(self._action_queue) == 0:
|
||||
action_batch = self._with_low_level_subtask_prompt(batch)
|
||||
actions = self.predict_action_chunk(action_batch)[:, : self.config.n_action_steps]
|
||||
self._action_queue.extend(actions.transpose(0, 1))
|
||||
|
||||
return self._action_queue.popleft()
|
||||
|
||||
def _with_low_level_subtask_prompt(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
|
||||
def _prepare_action_batch(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
|
||||
from .inference.pi052_adapter import _build_text_batch # noqa: PLC0415
|
||||
from .text_processor_pi052 import discretize_state_str # noqa: PLC0415
|
||||
|
||||
@@ -2126,73 +1803,14 @@ class PI052Policy(PI05Policy):
|
||||
return sorted_ix.gather(-1, choice).squeeze(-1)
|
||||
return torch.multinomial(probs, num_samples=1).squeeze(-1)
|
||||
|
||||
# PI0.5 flow-only fallback for unannotated batches.
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls: type[T],
|
||||
pretrained_name_or_path: str | Path,
|
||||
*,
|
||||
config: PreTrainedConfig | None = None,
|
||||
force_download: bool = False,
|
||||
resume_download: bool | None = None,
|
||||
proxies: dict | None = None,
|
||||
token: str | bool | None = None,
|
||||
cache_dir: str | Path | None = None,
|
||||
local_files_only: bool = False,
|
||||
revision: str | None = None,
|
||||
strict: bool = True,
|
||||
**kwargs,
|
||||
) -> T:
|
||||
"""Load a PI05/PI052 checkpoint, including sharded safetensors checkpoints."""
|
||||
if pretrained_name_or_path is None:
|
||||
raise ValueError("pretrained_name_or_path is required")
|
||||
|
||||
if config is None:
|
||||
config = PreTrainedConfig.from_pretrained(
|
||||
pretrained_name_or_path=pretrained_name_or_path,
|
||||
force_download=force_download,
|
||||
resume_download=resume_download,
|
||||
proxies=proxies,
|
||||
token=token,
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
revision=revision,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
model = cls(config, **kwargs)
|
||||
files = _resolve_weight_files(
|
||||
pretrained_name_or_path,
|
||||
force_download=force_download,
|
||||
resume_download=resume_download,
|
||||
proxies=proxies,
|
||||
token=token,
|
||||
cache_dir=cache_dir,
|
||||
local_files_only=local_files_only,
|
||||
revision=revision,
|
||||
)
|
||||
fixed_state_dict = model._fix_pytorch_state_dict_keys(_load_weight_files(files), model.config)
|
||||
remapped_state_dict = {
|
||||
key if key.startswith("model.") else f"model.{key}": value
|
||||
for key, value in fixed_state_dict.items()
|
||||
}
|
||||
|
||||
def _prepare_pretrained_state_dict(self, remapped_state_dict: dict[str, Tensor]) -> dict[str, Tensor]:
|
||||
lm_head_key = "model.paligemma_with_expert.paligemma.lm_head.weight"
|
||||
embed_tokens_key = "model.paligemma_with_expert.paligemma.model.language_model.embed_tokens.weight"
|
||||
if lm_head_key not in remapped_state_dict and embed_tokens_key in remapped_state_dict:
|
||||
remapped_state_dict[lm_head_key] = remapped_state_dict[embed_tokens_key].clone().float()
|
||||
elif lm_head_key in remapped_state_dict:
|
||||
remapped_state_dict[lm_head_key] = remapped_state_dict[lm_head_key].float()
|
||||
|
||||
missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)
|
||||
if not strict:
|
||||
if missing_keys:
|
||||
logger.warning("Missing PI052 checkpoint keys: %s", missing_keys)
|
||||
if unexpected_keys:
|
||||
logger.warning("Unexpected PI052 checkpoint keys: %s", unexpected_keys)
|
||||
model.to(config.device)
|
||||
model.eval()
|
||||
return model
|
||||
return remapped_state_dict
|
||||
|
||||
def get_optim_params(self):
|
||||
"""Return policy parameters, optionally split into LR-scaled groups.
|
||||
@@ -2264,63 +1882,8 @@ class PI052Policy(PI05Policy):
|
||||
|
||||
@torch.no_grad()
|
||||
def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
|
||||
"""Predict a chunk of actions given environment observations."""
|
||||
self.eval()
|
||||
|
||||
# Guard before first-observation FP8 calibration to prevent recursive prediction.
|
||||
if self.config.use_flashrt_fp8_mlp and not getattr(self, "_fp8_applied", False):
|
||||
self._fp8_applied = True
|
||||
self.apply_flashrt_fp8_mlp(batch)
|
||||
|
||||
# Prepare inputs
|
||||
images, img_masks = self._preprocess_images(batch)
|
||||
tokens, masks = batch[f"{OBS_LANGUAGE_TOKENS}"], batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
|
||||
|
||||
# Sample actions using the model (pass through RTC kwargs, no separate state needed for PI05)
|
||||
actions = self.model.sample_actions(images, img_masks, tokens, masks, **kwargs)
|
||||
|
||||
# Unpad actions to actual action dimension
|
||||
original_action_dim = self.config.output_features[ACTION].shape[0]
|
||||
actions = actions[:, :, :original_action_dim]
|
||||
|
||||
return actions
|
||||
|
||||
def _pi05_flow_forward(self, batch: dict[str, Tensor], reduction: str = "mean") -> tuple[Tensor, dict]:
|
||||
"""Run the batch through the model and compute the loss for training.
|
||||
|
||||
Args:
|
||||
batch: Training batch containing observations and actions.
|
||||
reduction: How to reduce the loss. Options:
|
||||
- "mean": Return scalar mean loss (default, backward compatible)
|
||||
- "none": Return per-sample losses of shape (batch_size,) for RA-BC weighting
|
||||
"""
|
||||
# Prepare inputs
|
||||
images, img_masks = self._preprocess_images(batch)
|
||||
tokens, masks = batch[f"{OBS_LANGUAGE_TOKENS}"], batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
|
||||
|
||||
actions = self.prepare_action(batch)
|
||||
|
||||
noise = self.model.sample_noise(actions.shape, actions.device)
|
||||
time = self.model.sample_time(actions.shape[0], actions.device)
|
||||
|
||||
# Compute loss (no separate state needed for PI05)
|
||||
losses = self.model.forward(images, img_masks, tokens, masks, actions, noise, time)
|
||||
|
||||
# Truncate losses to actual action dimensions
|
||||
original_action_dim = self.config.output_features[ACTION].shape[0]
|
||||
losses = losses[:, :, :original_action_dim]
|
||||
|
||||
loss_dict = {
|
||||
"loss_per_dim": losses.mean(dim=[0, 1]).detach().cpu().numpy().tolist(),
|
||||
}
|
||||
|
||||
if reduction == "none":
|
||||
# Return per-sample losses (B,) by averaging over time and action dims
|
||||
per_sample_loss = losses.mean(dim=(1, 2))
|
||||
loss_dict["loss"] = per_sample_loss.mean().item()
|
||||
return per_sample_loss, loss_dict
|
||||
else:
|
||||
# Default: return scalar mean loss
|
||||
loss = losses.mean()
|
||||
loss_dict["loss"] = loss.item()
|
||||
return loss, loss_dict
|
||||
return super().predict_action_chunk(batch, **kwargs)
|
||||
|
||||
@@ -37,7 +37,7 @@ def test_shifted_ce_none_retains_distinct_per_sample_losses():
|
||||
|
||||
|
||||
def test_checkpoint_resolution_forwards_explicit_hub_options(monkeypatch, tmp_path):
|
||||
import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
|
||||
import lerobot.policies.pi05.modeling_pi05 as modeling_pi05
|
||||
|
||||
checkpoint = tmp_path / "model.safetensors"
|
||||
checkpoint.touch()
|
||||
@@ -47,8 +47,8 @@ def test_checkpoint_resolution_forwards_explicit_hub_options(monkeypatch, tmp_pa
|
||||
calls.append((model_id, filename, kwargs))
|
||||
return None if filename.endswith("index.json") else str(checkpoint)
|
||||
|
||||
monkeypatch.setattr(modeling_pi052, "cached_file", fake_cached_file)
|
||||
files = modeling_pi052._resolve_weight_files(
|
||||
monkeypatch.setattr(modeling_pi05, "cached_file", fake_cached_file)
|
||||
files = modeling_pi05._resolve_weight_files(
|
||||
"org/model",
|
||||
force_download=True,
|
||||
resume_download=True,
|
||||
@@ -71,10 +71,10 @@ def test_checkpoint_resolution_forwards_explicit_hub_options(monkeypatch, tmp_pa
|
||||
|
||||
|
||||
def test_checkpoint_resolution_rejects_local_directory_without_weights(tmp_path):
|
||||
import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
|
||||
import lerobot.policies.pi05.modeling_pi05 as modeling_pi05
|
||||
|
||||
with pytest.raises(FileNotFoundError, match="model.safetensors"):
|
||||
modeling_pi052._resolve_weight_files(
|
||||
modeling_pi05._resolve_weight_files(
|
||||
tmp_path,
|
||||
force_download=False,
|
||||
resume_download=None,
|
||||
|
||||
@@ -41,14 +41,12 @@ def _checkpoint_model():
|
||||
tower = _MockVisionTower()
|
||||
language_model = SimpleNamespace(gradient_checkpointing=False)
|
||||
expert_model = SimpleNamespace(gradient_checkpointing=False)
|
||||
model = SimpleNamespace(
|
||||
gradient_checkpointing_enabled=False,
|
||||
paligemma_with_expert=SimpleNamespace(
|
||||
paligemma=SimpleNamespace(
|
||||
model=SimpleNamespace(language_model=language_model, vision_tower=tower)
|
||||
),
|
||||
gemma_expert=SimpleNamespace(model=expert_model),
|
||||
),
|
||||
model = PI05Pytorch.__new__(PI05Pytorch)
|
||||
nn.Module.__init__(model)
|
||||
model.gradient_checkpointing_enabled = False
|
||||
model.paligemma_with_expert = SimpleNamespace(
|
||||
paligemma=SimpleNamespace(model=SimpleNamespace(language_model=language_model, vision_tower=tower)),
|
||||
gemma_expert=SimpleNamespace(model=expert_model),
|
||||
)
|
||||
return model, tower, language_model, expert_model
|
||||
|
||||
|
||||
@@ -16,8 +16,12 @@
|
||||
|
||||
"""Test script to verify PI0.5 (pi05) support in PI0 policy"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from safetensors.torch import save_file
|
||||
from torch import nn
|
||||
|
||||
pytest.importorskip("transformers")
|
||||
|
||||
@@ -31,6 +35,26 @@ from lerobot.utils.random_utils import set_seed
|
||||
from tests.utils import require_cuda, require_hf_token # noqa: E402
|
||||
|
||||
|
||||
class _CheckpointPolicy(PI05Policy):
|
||||
def __init__(self, config, **kwargs):
|
||||
nn.Module.__init__(self)
|
||||
self.config = config
|
||||
self.loaded_state_dict = None
|
||||
|
||||
def load_state_dict(self, state_dict, strict=True, assign=False):
|
||||
self.loaded_state_dict = state_dict
|
||||
return [], []
|
||||
|
||||
|
||||
def test_from_pretrained_loads_existing_single_file_checkpoint(tmp_path):
|
||||
save_file({"weight": torch.tensor([1.0])}, tmp_path / "model.safetensors")
|
||||
|
||||
policy = _CheckpointPolicy.from_pretrained(tmp_path, config=SimpleNamespace())
|
||||
|
||||
assert policy.loaded_state_dict is not None
|
||||
torch.testing.assert_close(policy.loaded_state_dict["model.weight"], torch.tensor([1.0]))
|
||||
|
||||
|
||||
@require_cuda
|
||||
@require_hf_token
|
||||
def test_policy_instantiation():
|
||||
|
||||
Reference in New Issue
Block a user