mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 09:46:00 +00:00
Haoquan Fang
24017e960c
* add molmoact2 policy * add apache headers to molmoact2 files * simplify molmoact2 package imports * align molmoact2 feature validation with eo pattern * remove molmoact2 processor override from factory * guard molmoact2 transformers imports * guard molmoact2 processor transformers import * add scipy dependency to molmoact2 extra * use a single molmoact2 action queue * move molmoact2 config logic into config * fix molmoact2 hf image key resolution * load molmoact2 without remote code * lazy import molmoact2 scipy * format molmoact2 files * skip molmoact2 tests without optional deps * fix molmoact2 pre-commit checks * validate molmoact2 gripper range
238 lines
8.5 KiB
Python
238 lines
8.5 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2026 The Allen Institute for Artificial Intelligence and 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.
|
|
|
|
# ruff: noqa
|
|
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
from typing import ClassVar
|
|
|
|
import numpy as np
|
|
from tokenizers import ByteLevelBPETokenizer
|
|
from tokenizers.trainers import BpeTrainer
|
|
from huggingface_hub import snapshot_download
|
|
from transformers import PreTrainedTokenizerFast
|
|
from transformers.processing_utils import ProcessorMixin
|
|
|
|
|
|
def _hf_token() -> str | None:
|
|
return os.environ.get("HF_TOKEN") or os.environ.get("HF_ACCESS_TOKEN")
|
|
|
|
|
|
def _resolve_tokenizer_location(
|
|
tokenizer_path: str,
|
|
*,
|
|
revision: str | None = None,
|
|
force_download: bool = False,
|
|
) -> str:
|
|
local_path = Path(str(tokenizer_path)).expanduser()
|
|
if local_path.exists():
|
|
return str(local_path)
|
|
return snapshot_download(
|
|
repo_id=str(tokenizer_path),
|
|
repo_type="model",
|
|
revision=revision,
|
|
force_download=force_download,
|
|
ignore_patterns=["*.py", "*.pyc", "__pycache__/*"],
|
|
token=_hf_token(),
|
|
)
|
|
|
|
|
|
class UniversalActionProcessor(ProcessorMixin):
|
|
attributes: ClassVar[list[str]] = ["tokenizer"]
|
|
tokenizer_class: str = "AutoTokenizer"
|
|
|
|
def __init__(
|
|
self,
|
|
tokenizer: PreTrainedTokenizerFast,
|
|
scale: float = 10,
|
|
vocab_size: int = 1024,
|
|
min_token: int = 0,
|
|
*,
|
|
action_dim: int | None = None,
|
|
time_horizon: int | None = None,
|
|
):
|
|
self.scale = scale
|
|
self.vocab_size = vocab_size
|
|
self.min_token = min_token
|
|
|
|
# Action horizon and dimension needed during decoding. These can be specified
|
|
# in three ways (in order of priority):
|
|
# 1. passed in as kwargs to decode()
|
|
# 2. in the constructor
|
|
# 3. cached from the last time decode() was called
|
|
self.time_horizon = time_horizon
|
|
self.action_dim = action_dim
|
|
self.called_time_horizon = time_horizon
|
|
self.called_action_dim = action_dim
|
|
|
|
super().__init__(tokenizer)
|
|
self.bpe_tokenizer = self.tokenizer
|
|
|
|
def __call__(self, action_chunk: np.array) -> np.array:
|
|
from scipy.fft import dct
|
|
|
|
assert action_chunk.ndim <= 3, "Only 3 dimensions supported: [batch, timesteps, action_dim]"
|
|
if action_chunk.ndim == 2:
|
|
action_chunk = action_chunk[None, ...]
|
|
|
|
# Cache the time horizon and action dimension for decoding
|
|
self.called_time_horizon = action_chunk.shape[-2]
|
|
self.called_action_dim = action_chunk.shape[-1]
|
|
|
|
dct_coeff = dct(action_chunk, axis=1, norm="ortho")
|
|
dct_coeff = np.around(dct_coeff * self.scale)
|
|
tokens = []
|
|
for elem in dct_coeff:
|
|
token_str = "".join(map(chr, np.maximum(elem.flatten() - self.min_token, 0).astype(int)))
|
|
tokens.append(self.bpe_tokenizer(token_str)["input_ids"])
|
|
return tokens
|
|
|
|
def decode(
|
|
self,
|
|
tokens: list[list[int]],
|
|
*,
|
|
time_horizon: int | None = None,
|
|
action_dim: int | None = None,
|
|
) -> np.array:
|
|
from scipy.fft import idct
|
|
|
|
self.time_horizon = time_horizon or self.time_horizon or self.called_time_horizon
|
|
self.action_dim = action_dim or self.action_dim or self.called_action_dim
|
|
|
|
# Cache the time horizon and action dimension for the next call
|
|
self.called_time_horizon = self.time_horizon
|
|
self.called_action_dim = self.action_dim
|
|
|
|
assert self.time_horizon is not None and self.action_dim is not None, (
|
|
"Tokenizer not initialized, call encode() once or pass in time_horizon and action_dim."
|
|
)
|
|
|
|
decoded_actions = []
|
|
for token in tokens:
|
|
try:
|
|
decoded_tokens = self.bpe_tokenizer.decode(token)
|
|
decoded_dct_coeff = np.array(list(map(ord, decoded_tokens))) + self.min_token
|
|
decoded_dct_coeff = decoded_dct_coeff.reshape(-1, self.action_dim)
|
|
assert decoded_dct_coeff.shape == (
|
|
self.time_horizon,
|
|
self.action_dim,
|
|
), (
|
|
f"Decoded DCT coefficients have shape {decoded_dct_coeff.shape}, expected ({self.time_horizon}, {self.action_dim})"
|
|
)
|
|
except Exception as e:
|
|
print(f"Error decoding tokens: {e}")
|
|
print(f"Tokens: {token}")
|
|
decoded_dct_coeff = np.zeros((self.time_horizon, self.action_dim))
|
|
decoded_actions.append(idct(decoded_dct_coeff / self.scale, axis=0, norm="ortho"))
|
|
return np.stack(decoded_actions)
|
|
|
|
@classmethod
|
|
def fit(
|
|
cls,
|
|
action_data: list[np.array],
|
|
scale: float = 10,
|
|
vocab_size: int = 1024,
|
|
*,
|
|
time_horizon: int | None = None,
|
|
action_dim: int | None = None,
|
|
) -> "UniversalActionProcessor":
|
|
from scipy.fft import dct
|
|
|
|
# Run DCT over all inputs
|
|
dct_tokens = [dct(a, axis=0, norm="ortho").flatten() for a in action_data]
|
|
|
|
# Quantize and find min token
|
|
max_token = int(np.around(np.concatenate(dct_tokens) * scale).max())
|
|
min_token = int(np.around(np.concatenate(dct_tokens) * scale).min())
|
|
min_vocab_size = max_token - min_token
|
|
|
|
assert min_vocab_size <= vocab_size, (
|
|
f"Vocab size {vocab_size} is too small for the range of tokens {min_vocab_size}"
|
|
)
|
|
if min_vocab_size + 100 > vocab_size:
|
|
logging.warning(
|
|
f"Initial alphabet size {min_vocab_size} is almost as large as the vocab"
|
|
f"size {vocab_size}, consider increasing vocab size"
|
|
)
|
|
|
|
# Make token iterator for BPE training
|
|
def _token_iter():
|
|
for tokens in dct_tokens:
|
|
rounded_tokens = np.around(tokens * scale) - min_token
|
|
rounded_tokens = rounded_tokens.astype(int)
|
|
string = "".join(map(chr, rounded_tokens))
|
|
yield string
|
|
|
|
# Train BPE tokenizer
|
|
bpe = ByteLevelBPETokenizer()
|
|
|
|
# Set up the entire range of possible tokens as the initial alphabet
|
|
alphabet = [chr(i) for i in range(max_token - min_token + 1)]
|
|
trainer = BpeTrainer(
|
|
vocab_size=vocab_size,
|
|
min_frequency=2,
|
|
show_progress=True,
|
|
special_tokens=[],
|
|
initial_alphabet=alphabet,
|
|
max_token_length=10000,
|
|
)
|
|
|
|
# Train the inner tokenizer (don't use ByteLevelBPETokenizer.train_from_iterator()
|
|
# because it doesn't support custom alphabets)
|
|
bpe._tokenizer.train_from_iterator(_token_iter(), trainer=trainer)
|
|
|
|
return cls(
|
|
PreTrainedTokenizerFast(tokenizer_object=bpe, clean_up_tokenization_spaces=False),
|
|
scale=scale,
|
|
vocab_size=vocab_size,
|
|
min_token=min_token,
|
|
time_horizon=time_horizon,
|
|
action_dim=action_dim,
|
|
)
|
|
|
|
@classmethod
|
|
def from_pretrained_local(
|
|
cls,
|
|
pretrained_model_name_or_path: str,
|
|
*,
|
|
revision: str | None = None,
|
|
force_download: bool = False,
|
|
) -> "UniversalActionProcessor":
|
|
location = Path(
|
|
_resolve_tokenizer_location(
|
|
pretrained_model_name_or_path,
|
|
revision=revision,
|
|
force_download=force_download,
|
|
)
|
|
)
|
|
processor_config = {}
|
|
processor_config_path = location / "processor_config.json"
|
|
if processor_config_path.exists():
|
|
import json
|
|
|
|
processor_config = json.loads(processor_config_path.read_text())
|
|
tokenizer = PreTrainedTokenizerFast.from_pretrained(str(location))
|
|
return cls(
|
|
tokenizer,
|
|
scale=processor_config.get("scale", 10),
|
|
vocab_size=processor_config.get("vocab_size", 1024),
|
|
min_token=processor_config.get("min_token", 0),
|
|
action_dim=processor_config.get("action_dim"),
|
|
time_horizon=processor_config.get("time_horizon"),
|
|
)
|