mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-28 20:26:05 +00:00
fix(pi0_fast): align FAST semantics with OpenPI
Use OpenPI-compatible normalization, token boundaries, balanced loss, and strict full-chunk decoding so training and inference share one sequence contract. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
pepijn
@@ -55,14 +55,7 @@ MODEL_PATH_LEROBOT = "lerobot/pi0fast-base"
|
||||
# Expected action token shape: (batch_size, max_decoding_steps)
|
||||
EXPECTED_ACTION_TOKENS_SHAPE = (1, 2)
|
||||
|
||||
# Expected first 5 action tokens (for reproducibility check)
|
||||
EXPECTED_ACTION_TOKENS_FIRST_5 = torch.tensor([255020, 255589])
|
||||
|
||||
# Expected actions after detokenization
|
||||
EXPECTED_ACTIONS_SHAPE = (1, 2, 32) # (batch_size, n_action_steps, action_dim)
|
||||
EXPECTED_ACTIONS_MEAN = 0.046403881162405014
|
||||
EXPECTED_ACTIONS_STD = 0.2607129216194153
|
||||
EXPECTED_ACTIONS_FIRST_5 = torch.tensor([0.0000, 0.3536, 0.0707, 0.0000, 0.0000])
|
||||
|
||||
|
||||
@require_cuda
|
||||
@@ -99,9 +92,8 @@ def instantiate_lerobot_pi0_fast(
|
||||
pretrained_name_or_path=model_path,
|
||||
strict=True,
|
||||
)
|
||||
policy.config.validate_action_token_prefix = False
|
||||
policy.config.max_action_tokens = 2
|
||||
policy.config.max_decoding_steps = 2
|
||||
policy.config.max_decoding_steps = 256
|
||||
policy.config.chunk_size = 2
|
||||
policy.config.n_action_steps = 2
|
||||
else:
|
||||
@@ -110,9 +102,8 @@ def instantiate_lerobot_pi0_fast(
|
||||
max_action_dim=DUMMY_ACTION_DIM,
|
||||
max_state_dim=DUMMY_STATE_DIM,
|
||||
device=DEVICE,
|
||||
validate_action_token_prefix=False,
|
||||
max_action_tokens=2,
|
||||
max_decoding_steps=2,
|
||||
max_decoding_steps=256,
|
||||
chunk_size=2,
|
||||
)
|
||||
policy = PI0FastPolicy(config)
|
||||
@@ -262,56 +253,11 @@ def test_pi0_fast_action_generation(policy, preprocessor):
|
||||
print(f"LeRobot actions std: {lerobot_actions.std().item():.6f}")
|
||||
print(f"LeRobot actions first 5: {lerobot_actions[0, 0, :5]}")
|
||||
|
||||
print("\nExpected values (from original PI0Fast):")
|
||||
print(f"Expected actions shape: {EXPECTED_ACTIONS_SHAPE}")
|
||||
print(f"Expected actions mean: {EXPECTED_ACTIONS_MEAN:.6f}")
|
||||
print(f"Expected actions std: {EXPECTED_ACTIONS_STD:.6f}")
|
||||
print(f"Expected actions first 5: {EXPECTED_ACTIONS_FIRST_5}")
|
||||
|
||||
print("\nAction Comparison:")
|
||||
print("-" * 80)
|
||||
|
||||
# Compare shapes
|
||||
actual_shape = tuple(lerobot_actions.shape)
|
||||
print(f"Actual shape: {actual_shape}")
|
||||
|
||||
assert actual_shape == EXPECTED_ACTIONS_SHAPE, (
|
||||
f"Shape mismatch: {actual_shape} vs {EXPECTED_ACTIONS_SHAPE}"
|
||||
)
|
||||
print(f"Shape matches: {actual_shape}")
|
||||
|
||||
# Compare statistics
|
||||
actual_mean = lerobot_actions.mean().item()
|
||||
actual_std = lerobot_actions.std().item()
|
||||
|
||||
print(f"\nMean: {actual_mean:.6f} (expected: {EXPECTED_ACTIONS_MEAN:.6f})")
|
||||
print(f"Std: {actual_std:.6f} (expected: {EXPECTED_ACTIONS_STD:.6f})")
|
||||
|
||||
# Compare first 5 actions
|
||||
actual_first_5 = lerobot_actions[0, 0, :5]
|
||||
print("\nFirst 5 actions comparison:")
|
||||
print(f" Actual: {actual_first_5}")
|
||||
print(f" Expected: {EXPECTED_ACTIONS_FIRST_5}")
|
||||
|
||||
first_5_diff = torch.abs(actual_first_5 - EXPECTED_ACTIONS_FIRST_5)
|
||||
print(f" Max diff: {first_5_diff.max().item():.6e}")
|
||||
print(f" Mean diff: {first_5_diff.mean().item():.6e}")
|
||||
|
||||
# Check with different tolerances
|
||||
tolerances = [1e-5, 1e-4, 1e-3, 1e-2]
|
||||
for tol in tolerances:
|
||||
is_close = torch.allclose(actual_first_5, EXPECTED_ACTIONS_FIRST_5, atol=tol)
|
||||
status = "Success" if is_close else "Failure"
|
||||
print(f"{status}: First 5 actions close (atol={tol}): {is_close}")
|
||||
|
||||
# Assert with reasonable tolerance
|
||||
tolerance = 1e-3
|
||||
assert torch.allclose(actual_first_5, EXPECTED_ACTIONS_FIRST_5, atol=tolerance), (
|
||||
f"First 5 actions differ by more than tolerance ({tolerance})"
|
||||
)
|
||||
print(f"\nSuccess: Actions match expected values within tolerance ({tolerance})!")
|
||||
|
||||
print("\nAction generation test completed (values printed for reference)!")
|
||||
assert torch.isfinite(lerobot_actions).all()
|
||||
|
||||
|
||||
@require_cuda
|
||||
@@ -442,10 +388,6 @@ def test_pi0_fast_action_token_sampling(policy, preprocessor):
|
||||
print(f"Action tokens shape: {action_tokens.shape}")
|
||||
print(f"Action tokens first 10: {action_tokens[0, :10].tolist()}")
|
||||
|
||||
print("\nExpected values (from original PI0Fast):")
|
||||
print(f"Expected shape: {EXPECTED_ACTION_TOKENS_SHAPE}")
|
||||
print(f"Expected first 5: {EXPECTED_ACTION_TOKENS_FIRST_5.tolist()}")
|
||||
|
||||
# Verify shape
|
||||
actual_shape = tuple(action_tokens.shape)
|
||||
print(f"\nActual shape: {actual_shape}")
|
||||
@@ -454,11 +396,8 @@ def test_pi0_fast_action_token_sampling(policy, preprocessor):
|
||||
f"Shape mismatch: {actual_shape} vs {EXPECTED_ACTION_TOKENS_SHAPE}"
|
||||
)
|
||||
|
||||
# Compare first 5 tokens
|
||||
actual_first_5 = action_tokens[0, :5].cpu()
|
||||
assert torch.equal(actual_first_5, EXPECTED_ACTION_TOKENS_FIRST_5), (
|
||||
f"First 5 tokens mismatch: {actual_first_5} vs {EXPECTED_ACTION_TOKENS_FIRST_5}"
|
||||
)
|
||||
action_prefix = policy._paligemma_tokenizer.encode("Action: ", add_special_tokens=False)
|
||||
assert action_tokens[0, : len(action_prefix)].tolist() == action_prefix
|
||||
|
||||
print("\nAction token sampling test completed!")
|
||||
|
||||
@@ -500,19 +439,11 @@ def test_pi0_fast_detokenization(policy, preprocessor):
|
||||
|
||||
# Detokenize
|
||||
print("\n[LeRobot] Detokenizing action tokens...")
|
||||
action_horizon = policy.config.n_action_steps
|
||||
action_horizon = policy.config.chunk_size
|
||||
action_dim = policy.config.output_features["action"].shape[0]
|
||||
|
||||
try:
|
||||
continuous_actions = policy.detokenize_actions(
|
||||
action_tokens, action_horizon=action_horizon, action_dim=action_dim
|
||||
)
|
||||
print(f"Continuous actions shape: {continuous_actions.shape}")
|
||||
print(f"Continuous actions mean: {continuous_actions.mean().item():.6f}")
|
||||
print(f"Continuous actions std: {continuous_actions.std().item():.6f}")
|
||||
print(f"Continuous actions first 5: {continuous_actions[0, 0, :5]}")
|
||||
print("\nDetokenization successful!")
|
||||
except Exception as e:
|
||||
print(f"\nDetokenization failed with error: {e}")
|
||||
print("This may be expected if the action tokens are not valid FAST tokens.")
|
||||
print("The test will pass as long as the sampling works correctly.")
|
||||
continuous_actions = policy.detokenize_actions(
|
||||
action_tokens, action_horizon=action_horizon, action_dim=action_dim
|
||||
)
|
||||
assert continuous_actions.shape == (1, action_horizon, action_dim)
|
||||
assert torch.isfinite(continuous_actions).all()
|
||||
|
||||
@@ -0,0 +1,265 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 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.
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
pytest.importorskip("transformers")
|
||||
pytest.importorskip("scipy")
|
||||
|
||||
from lerobot.configs import NormalizationMode # noqa: E402
|
||||
from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig # noqa: E402
|
||||
from lerobot.policies.pi0_fast.modeling_pi0_fast import ( # noqa: E402
|
||||
PI0FastPolicy,
|
||||
PI0FastPytorch,
|
||||
_gather_last_valid_language_hidden,
|
||||
_reduce_fast_token_loss,
|
||||
)
|
||||
from lerobot.policies.pi0_fast.processor_pi0_fast import ( # noqa: E402
|
||||
Pi0FastPrepareStateAndLanguageTokenizerProcessorStep,
|
||||
)
|
||||
from lerobot.processor.tokenizer_processor import ActionTokenizerProcessorStep # noqa: E402
|
||||
from lerobot.types import TransitionKey # noqa: E402
|
||||
from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS # noqa: E402
|
||||
|
||||
|
||||
class _FakePaliGemmaTokenizer:
|
||||
vocab_size = 1000
|
||||
bos_token_id = 2
|
||||
eos_token_id = 1
|
||||
|
||||
def encode(self, text, add_special_tokens=True):
|
||||
if text == "Action: ":
|
||||
return [10, 11]
|
||||
if text == "|":
|
||||
return [12, self.eos_token_id] if add_special_tokens else [12]
|
||||
return [900, 901]
|
||||
|
||||
|
||||
def test_pi0_fast_uses_openpi_quantile_normalization_by_default():
|
||||
config = PI0FastConfig()
|
||||
|
||||
assert config.normalization_mapping == {
|
||||
"VISUAL": NormalizationMode.IDENTITY,
|
||||
"STATE": NormalizationMode.QUANTILES,
|
||||
"ACTION": NormalizationMode.QUANTILES,
|
||||
}
|
||||
|
||||
|
||||
def test_pi0_fast_action_tokens_have_no_second_bos():
|
||||
step = ActionTokenizerProcessorStep.__new__(ActionTokenizerProcessorStep)
|
||||
step.max_action_tokens = 8
|
||||
step.fast_skip_tokens = 128
|
||||
step.prepend_bos = False
|
||||
step.action_tokenizer = lambda _: torch.tensor([4, 9])
|
||||
step._paligemma_tokenizer = _FakePaliGemmaTokenizer()
|
||||
|
||||
tokens, mask = step._tokenize_action(torch.zeros(1, 2, 1))
|
||||
|
||||
mapped = [1000 - 1 - 128 - token for token in (4, 9)]
|
||||
assert tokens[0, :6].tolist() == [10, 11, *mapped, 12, 1]
|
||||
assert _FakePaliGemmaTokenizer.bos_token_id not in tokens[0, mask[0]].tolist()
|
||||
|
||||
|
||||
def test_action_tokenizer_serializes_bos_layout():
|
||||
step = ActionTokenizerProcessorStep.__new__(ActionTokenizerProcessorStep)
|
||||
step.trust_remote_code = True
|
||||
step.max_action_tokens = 8
|
||||
step.fast_skip_tokens = 128
|
||||
step.paligemma_tokenizer_name = "paligemma"
|
||||
step.prepend_bos = False
|
||||
step.action_tokenizer_name = "fast"
|
||||
step.action_tokenizer_input_object = None
|
||||
|
||||
assert step.get_config() == {
|
||||
"trust_remote_code": True,
|
||||
"max_action_tokens": 8,
|
||||
"fast_skip_tokens": 128,
|
||||
"paligemma_tokenizer_name": "paligemma",
|
||||
"prepend_bos": False,
|
||||
"action_tokenizer_name": "fast",
|
||||
}
|
||||
|
||||
|
||||
def test_pi0_fast_prompt_canonicalizes_task_to_lowercase():
|
||||
step = Pi0FastPrepareStateAndLanguageTokenizerProcessorStep()
|
||||
transition = {
|
||||
TransitionKey.OBSERVATION: {"observation.state": torch.zeros(1, 2)},
|
||||
TransitionKey.COMPLEMENTARY_DATA: {"task": [" Pick_UP\nCube "]},
|
||||
}
|
||||
|
||||
result = step(transition)
|
||||
|
||||
assert result[TransitionKey.COMPLEMENTARY_DATA]["task"] == ["Task: pick up cube, State: 128 128;\n"]
|
||||
|
||||
|
||||
def test_last_language_hidden_uses_attention_mask_not_padding():
|
||||
hidden = torch.arange(2 * 7, dtype=torch.float32).reshape(2, 7, 1)
|
||||
language_mask = torch.tensor([[True, True, False, False], [True, True, True, False]])
|
||||
|
||||
gathered = _gather_last_valid_language_hidden(hidden, language_mask, image_token_count=3)
|
||||
|
||||
assert gathered[:, 0].tolist() == [hidden[0, 4, 0].item(), hidden[1, 5, 0].item()]
|
||||
|
||||
|
||||
def test_fast_ce_averages_each_sample_before_the_batch():
|
||||
token_loss = torch.tensor([[2.0, 99.0, 99.0], [6.0, 6.0, 6.0]])
|
||||
mask = torch.tensor([[True, False, False], [True, True, True]])
|
||||
|
||||
loss = _reduce_fast_token_loss(token_loss, mask)
|
||||
|
||||
assert loss.item() == pytest.approx(4.0)
|
||||
|
||||
|
||||
class _TokenFromHiddenHead(nn.Module):
|
||||
def forward(self, hidden):
|
||||
logits = torch.full((*hidden.shape[:-1], 10), -100.0)
|
||||
logits.scatter_(-1, hidden.long(), 100.0)
|
||||
return logits
|
||||
|
||||
|
||||
class _ScriptedPaliGemma:
|
||||
def __init__(self):
|
||||
q_proj = SimpleNamespace(weight=torch.empty(1))
|
||||
language_model = SimpleNamespace(layers=[SimpleNamespace(self_attn=SimpleNamespace(q_proj=q_proj))])
|
||||
self.paligemma = SimpleNamespace(
|
||||
lm_head=_TokenFromHiddenHead(),
|
||||
model=SimpleNamespace(language_model=language_model),
|
||||
)
|
||||
self.calls = 0
|
||||
|
||||
def forward(self, inputs_embeds, **_kwargs):
|
||||
scripted_tokens = ([5, 6], [1, 7], [9, 1])
|
||||
token_ids = torch.tensor(scripted_tokens[self.calls], dtype=torch.float32)
|
||||
self.calls += 1
|
||||
hidden = token_ids[:, None, None].expand(-1, inputs_embeds[0].shape[1], -1).clone()
|
||||
return (hidden, None), object()
|
||||
|
||||
@staticmethod
|
||||
def embed_language_tokens(tokens):
|
||||
return tokens.to(dtype=torch.float32).unsqueeze(-1)
|
||||
|
||||
|
||||
def _make_scripted_generation_model(captured):
|
||||
model = PI0FastPytorch.__new__(PI0FastPytorch)
|
||||
nn.Module.__init__(model)
|
||||
model.config = SimpleNamespace(max_action_tokens=4)
|
||||
model._paligemma_tokenizer = SimpleNamespace(eos_token_id=1)
|
||||
model.paligemma_with_expert = _ScriptedPaliGemma()
|
||||
model._prepare_attention_masks_4d = lambda masks, dtype: masks
|
||||
|
||||
def embed_prefix(_images, _img_masks, tokens, masks, **_kwargs):
|
||||
captured.append(tokens.clone())
|
||||
image_masks = torch.ones(tokens.shape[0], 1, dtype=torch.bool)
|
||||
pad_masks = torch.cat([image_masks, masks], dim=1)
|
||||
embeddings = torch.zeros(tokens.shape[0], pad_masks.shape[1], 1)
|
||||
attention = torch.ones(tokens.shape[0], pad_masks.shape[1], pad_masks.shape[1], dtype=torch.bool)
|
||||
return embeddings, pad_masks, attention, 1, 0
|
||||
|
||||
model.embed_prefix_fast = embed_prefix
|
||||
return model
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sampler_name", ["sample_actions_fast", "sample_actions_fast_kv_cache"])
|
||||
def test_fast_generators_stop_each_sample_at_eos_without_boundary_bos(sampler_name):
|
||||
captured = []
|
||||
model = _make_scripted_generation_model(captured)
|
||||
tokens = torch.tensor([[2, 3, 0, 0], [2, 3, 4, 0]])
|
||||
masks = torch.tensor([[True, True, False, False], [True, True, True, False]])
|
||||
|
||||
generated = getattr(model, sampler_name)([], [], tokens, masks, max_decoding_steps=4, temperature=0.0)
|
||||
|
||||
assert torch.equal(generated, torch.tensor([[5, 1, 0, 0], [6, 7, 1, 0]]))
|
||||
assert torch.equal(captured[0], tokens)
|
||||
|
||||
|
||||
def test_predict_action_chunk_decodes_full_chunk(monkeypatch):
|
||||
policy = PI0FastPolicy.__new__(PI0FastPolicy)
|
||||
nn.Module.__init__(policy)
|
||||
policy.config = SimpleNamespace(
|
||||
chunk_size=8,
|
||||
n_action_steps=3,
|
||||
output_features={"action": SimpleNamespace(shape=(4,))},
|
||||
temperature=0.0,
|
||||
max_decoding_steps=16,
|
||||
use_kv_cache=False,
|
||||
)
|
||||
policy.model = SimpleNamespace(
|
||||
sample_actions_fast=lambda *args, **kwargs: torch.ones(1, 4, dtype=torch.long)
|
||||
)
|
||||
monkeypatch.setattr(policy, "_preprocess_images", lambda batch: ([], []))
|
||||
captured = {}
|
||||
|
||||
def detokenize(tokens, action_horizon, action_dim):
|
||||
captured["shape"] = (action_horizon, action_dim)
|
||||
return torch.zeros(1, action_horizon, action_dim)
|
||||
|
||||
monkeypatch.setattr(policy, "detokenize_actions", detokenize)
|
||||
batch = {
|
||||
OBS_LANGUAGE_TOKENS: torch.ones(1, 2, dtype=torch.long),
|
||||
OBS_LANGUAGE_ATTENTION_MASK: torch.ones(1, 2, dtype=torch.bool),
|
||||
}
|
||||
|
||||
actions = policy.predict_action_chunk(batch)
|
||||
|
||||
assert captured["shape"] == (8, 4)
|
||||
assert actions.shape == (1, 8, 4)
|
||||
|
||||
|
||||
class _FakeActionTokenizer:
|
||||
@staticmethod
|
||||
def decode(tokens, time_horizon, action_dim):
|
||||
if tokens != [[-119]]:
|
||||
raise ValueError("invalid FAST token")
|
||||
return [np.arange(time_horizon * action_dim).reshape(time_horizon, action_dim)]
|
||||
|
||||
|
||||
class _FakeDecodeTokenizer(_FakePaliGemmaTokenizer):
|
||||
decoded_text = ""
|
||||
|
||||
def decode(self, _tokens):
|
||||
return self.decoded_text
|
||||
|
||||
def encode(self, text, add_special_tokens=True):
|
||||
if text == "codes":
|
||||
return [990]
|
||||
return super().encode(text, add_special_tokens=add_special_tokens)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("decoded_text", ["", "not an action", "Action: bad"])
|
||||
def test_malformed_fast_generation_returns_zero_actions(decoded_text):
|
||||
policy = PI0FastPolicy.__new__(PI0FastPolicy)
|
||||
nn.Module.__init__(policy)
|
||||
tokenizer = _FakeDecodeTokenizer()
|
||||
tokenizer.decoded_text = decoded_text
|
||||
policy._paligemma_tokenizer = tokenizer
|
||||
policy.action_tokenizer = _FakeActionTokenizer()
|
||||
policy.config = SimpleNamespace(fast_skip_tokens=128)
|
||||
|
||||
actions = policy.detokenize_actions(torch.tensor([[7, 1, 0]]), action_horizon=2, action_dim=2)
|
||||
|
||||
assert actions.shape == (1, 2, 2)
|
||||
assert torch.equal(actions, torch.zeros_like(actions))
|
||||
|
||||
|
||||
def test_valid_fast_generation_decodes_exact_shape():
|
||||
policy = PI0FastPolicy.__new__(PI0FastPolicy)
|
||||
nn.Module.__init__(policy)
|
||||
tokenizer = _FakeDecodeTokenizer()
|
||||
tokenizer.decoded_text = "Action: codes|"
|
||||
policy._paligemma_tokenizer = tokenizer
|
||||
policy.action_tokenizer = _FakeActionTokenizer()
|
||||
policy.config = SimpleNamespace(fast_skip_tokens=128)
|
||||
|
||||
actions = policy.detokenize_actions(torch.tensor([[7, 1, 0]]), action_horizon=2, action_dim=2)
|
||||
|
||||
assert actions.shape == (1, 2, 2)
|
||||
assert np.isfinite(actions.numpy()).all()
|
||||
Reference in New Issue
Block a user