fix(processor): serialize FAST token mapping

This commit is contained in:
Pepijn
2026-07-15 16:04:26 +02:00
parent a09715121e
commit eddf75616e
2 changed files with 21 additions and 1 deletions
@@ -565,6 +565,8 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
config = { config = {
"trust_remote_code": self.trust_remote_code, "trust_remote_code": self.trust_remote_code,
"max_action_tokens": self.max_action_tokens, "max_action_tokens": self.max_action_tokens,
"fast_skip_tokens": self.fast_skip_tokens,
"paligemma_tokenizer_name": self.paligemma_tokenizer_name,
} }
# Only save tokenizer_name if it was used to create the tokenizer # Only save tokenizer_name if it was used to create the tokenizer
+19 -1
View File
@@ -25,7 +25,7 @@ import pytest
import torch import torch
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep from lerobot.processor import ActionTokenizerProcessorStep, DataProcessorPipeline, TokenizerProcessorStep
from lerobot.processor.converters import create_transition, identity_transition from lerobot.processor.converters import create_transition, identity_transition
from lerobot.types import TransitionKey from lerobot.types import TransitionKey
from lerobot.utils.constants import ( from lerobot.utils.constants import (
@@ -88,6 +88,24 @@ class MockTokenizer:
return result return result
def test_action_tokenizer_config_preserves_token_mapping():
processor = object.__new__(ActionTokenizerProcessorStep)
processor.trust_remote_code = True
processor.max_action_tokens = 384
processor.fast_skip_tokens = 64
processor.paligemma_tokenizer_name = "custom/paligemma"
processor.action_tokenizer_name = "custom/fast"
processor.action_tokenizer_input_object = None
assert processor.get_config() == {
"trust_remote_code": True,
"max_action_tokens": 384,
"fast_skip_tokens": 64,
"paligemma_tokenizer_name": "custom/paligemma",
"action_tokenizer_name": "custom/fast",
}
@pytest.fixture @pytest.fixture
def mock_tokenizer(): def mock_tokenizer():
"""Provide a mock tokenizer for testing.""" """Provide a mock tokenizer for testing."""