admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-10 19:41:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(evo1): move LIBERO padding into policy processors

Browse Source
This commit is contained in:
javadcc_mac
2026-06-21 15:58:38 +08:00
parent 4cfa762da8
commit 25556ceefe
16 changed files with 637 additions and 252 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/envs/test_dispatch.py
+10 -61
View File
@@ -13,7 +13,7 @@ from gymnasium.envs.registration import register, registry as gym_registry
from lerobot.configs.types import PolicyFeature
from lerobot.envs.configs import EnvConfig, LiberoEnv
from lerobot.envs.factory import make_env, make_env_config, make_env_pre_post_processors
from lerobot.processor import LiberoActionProcessorStep, LiberoProcessorStep
from lerobot.processor import LiberoProcessorStep
from lerobot.utils.constants import OBS_PREFIX, OBS_STATE
logger = logging.getLogger(__name__)
@@ -86,38 +86,18 @@ def test_processors_delegation_supports_legacy_override_signature():
assert isinstance(post, DataProcessorPipeline)
def test_libero_evo1_processors_use_padded_state_and_env_action_dim():
"""EVO1 uses padded LIBERO state features while env actions stay executable."""
class _Evo1Config:
type = "evo1"
max_state_dim = 24
def test_libero_processors_are_policy_agnostic():
cfg = LiberoEnv()
pre, post = make_env_pre_post_processors(cfg, policy_cfg=_Evo1Config())
pre, post = make_env_pre_post_processors(cfg, policy_cfg=object())
assert isinstance(pre.steps[0], LiberoProcessorStep)
assert pre.steps[0].max_state_dim == 24
assert isinstance(post.steps[0], LiberoActionProcessorStep)
assert post.steps[0].action_dim == cfg.features["action"].shape[0] == 7
assert post.steps[0].binarize_gripper is True
class _OtherConfig:
type = "other"
pre_other, post_other = make_env_pre_post_processors(cfg, policy_cfg=_OtherConfig())
assert pre_other.steps[0].max_state_dim is None
assert post_other.steps[0].binarize_gripper is False
cfg.binarize_gripper = False
_, post_disabled = make_env_pre_post_processors(cfg, policy_cfg=_Evo1Config())
assert post_disabled.steps[0].binarize_gripper is False
assert len(post.steps) == 0
def test_libero_processor_pads_state_to_max_dim():
step = LiberoProcessorStep(max_state_dim=24)
def test_libero_processor_flattens_state_to_raw_8_dim():
step = LiberoProcessorStep()
observation = {
OBS_PREFIX
+ "robot_state": {
OBS_PREFIX + "robot_state": {
"eef": {
"pos": torch.tensor([[1.0, 2.0, 3.0]]),
"quat": torch.tensor([[0.0, 0.0, 0.0, 1.0]]),
@@ -127,39 +107,8 @@ def test_libero_processor_pads_state_to_max_dim():
}
state = step.observation(observation)[OBS_STATE]
assert state.shape == (1, 24)
assert torch.allclose(state[:, :8], torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 4.0, 5.0]]))
assert torch.count_nonzero(state[:, 8:]).item() == 0
def test_libero_action_processor_slices_padded_action():
step = LiberoActionProcessorStep(action_dim=7)
action = torch.arange(2 * 3 * 24, dtype=torch.float32).reshape(2, 3, 24)
sliced = step.action(action)
assert sliced.shape == (2, 3, 7)
assert torch.equal(sliced, action[..., :7])
with pytest.raises(ValueError, match="smaller than action_dim=7"):
step.action(torch.zeros(2, 6))
def test_libero_action_processor_can_binarize_gripper():
step = LiberoActionProcessorStep(action_dim=7, binarize_gripper=True)
action = torch.tensor(
[
[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 0.5, 7.0],
[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 0.6, 7.0],
],
dtype=torch.float32,
)
processed = step.action(action)
assert processed.shape == (2, 7)
assert torch.equal(processed[:, :6], action[:, :6])
assert torch.equal(processed[:, 6], torch.tensor([1.0, -1.0]))
assert torch.equal(action[:, 6], torch.tensor([0.5, 0.6]))
assert state.shape == (1, 8)
assert torch.allclose(state, torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 4.0, 5.0]]))
def test_base_create_envs():
tests/policies/evo1/test_evo1.py
+155 -4
View File
@@ -16,6 +16,7 @@
from __future__ import annotations
import pytest
import torch
from torch import nn
@@ -23,7 +24,15 @@ import lerobot.policies.evo1.modeling_evo1 as modeling_evo1
from lerobot.configs.types import FeatureType, PolicyFeature
from lerobot.policies.evo1.configuration_evo1 import Evo1Config
from lerobot.policies.evo1.flow_matching import FlowmatchingActionHead
from lerobot.policies.evo1.processor_evo1 import (
Evo1ActionProcessorStep,
Evo1PadActionProcessorStep,
Evo1PadStateProcessorStep,
ensure_evo1_processor_steps,
make_evo1_pre_post_processors,
)
from lerobot.policies.factory import get_policy_class, make_policy_config
from lerobot.processor import NormalizerProcessorStep, PolicyProcessorPipeline, UnnormalizerProcessorStep
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
STATE_DIM = 4
@@ -108,6 +117,19 @@ def make_batch(include_action=True):
return batch
def make_stats(state_dim=STATE_DIM, action_dim=ACTION_DIM):
return {
OBS_STATE: {
"min": torch.full((state_dim,), -2.0),
"max": torch.full((state_dim,), 2.0),
},
ACTION: {
"min": torch.full((action_dim,), -1.0),
"max": torch.full((action_dim,), 1.0),
},
}
def test_evo1_factory_registration():
cfg = make_policy_config(
"evo1",
@@ -191,22 +213,151 @@ def test_evo1_stage_defaults_and_consistency():
raise AssertionError("Expected inconsistent finetune config to raise ValueError")
def test_evo1_rejects_non_square_image_resolution():
with pytest.raises(ValueError, match="square image_resolution"):
make_config(image_resolution=(448, 320))
def test_evo1_build_model_config_uses_image_resolution_and_trainable_checkpointing():
stage1 = make_config(training_stage="stage1", image_resolution=(224, 224))
stage1_model_config = modeling_evo1.EVO1Policy._build_model_config(stage1)
assert stage1_model_config["image_size"] == 224
assert stage1_model_config["enable_gradient_checkpointing"] is False
stage2 = make_config(training_stage="stage2", image_resolution=(224, 224))
stage2_model_config = modeling_evo1.EVO1Policy._build_model_config(stage2)
assert stage2_model_config["enable_gradient_checkpointing"] is True
def test_evo1_policy_processors_pad_state_crop_action_and_binarize_gripper():
libero_action_dim = 7
config = make_config(
max_state_dim=MAX_STATE_DIM,
max_action_dim=8,
postprocess_action_dim=libero_action_dim,
binarize_gripper=True,
output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(libero_action_dim,))},
)
stats = make_stats(action_dim=libero_action_dim)
preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=stats)
assert isinstance(preprocessor.steps[2], Evo1PadStateProcessorStep)
assert isinstance(preprocessor.steps[3], Evo1PadActionProcessorStep)
assert isinstance(preprocessor.steps[4], NormalizerProcessorStep)
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
assert isinstance(postprocessor.steps[1], Evo1ActionProcessorStep)
normalizer = preprocessor.steps[4]
assert normalizer.features[OBS_STATE].shape == (MAX_STATE_DIM,)
assert normalizer.features[ACTION].shape == (8,)
assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
assert normalizer._tensor_stats[ACTION]["min"].shape == (8,)
processed_batch = preprocessor(
{
"task": "pick the block",
OBS_STATE: torch.zeros(STATE_DIM),
ACTION: torch.zeros(libero_action_dim),
f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
}
)
processed_state = processed_batch[OBS_STATE]
assert processed_state.shape == (1, MAX_STATE_DIM)
assert torch.allclose(processed_state, torch.zeros_like(processed_state))
assert processed_batch[ACTION].shape == (1, 8)
assert torch.allclose(processed_batch[ACTION], torch.zeros_like(processed_batch[ACTION]))
assert processed_batch["action_mask"].shape == (1, 8)
assert processed_batch["action_mask"][:, :libero_action_dim].all()
assert not processed_batch["action_mask"][:, libero_action_dim:].any()
action = torch.tensor(
[
[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.5, 0.7],
[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7],
],
dtype=torch.float32,
)
processed = postprocessor(action)
assert processed.shape == (2, 7)
assert torch.allclose(processed[:, :6], action[:, :6])
assert torch.equal(processed[:, 6], torch.tensor([1.0, -1.0]))
def test_evo1_legacy_processors_are_completed_before_normalization():
config = make_config(
max_state_dim=MAX_STATE_DIM,
max_action_dim=8,
postprocess_action_dim=7,
binarize_gripper=True,
)
stats = make_stats(action_dim=7)
legacy_pre = PolicyProcessorPipeline(
steps=[
NormalizerProcessorStep(
features={**config.input_features, **config.output_features},
norm_map=config.normalization_mapping,
stats=stats,
)
]
)
legacy_post = PolicyProcessorPipeline(
steps=[
UnnormalizerProcessorStep(
features=config.output_features,
norm_map=config.normalization_mapping,
stats=stats,
)
]
)
preprocessor, postprocessor = ensure_evo1_processor_steps(config, legacy_pre, legacy_post)
assert isinstance(preprocessor.steps[0], Evo1PadStateProcessorStep)
assert isinstance(preprocessor.steps[1], Evo1PadActionProcessorStep)
assert isinstance(preprocessor.steps[2], NormalizerProcessorStep)
assert isinstance(postprocessor.steps[0], UnnormalizerProcessorStep)
assert isinstance(postprocessor.steps[1], Evo1ActionProcessorStep)
assert postprocessor.steps[1].action_dim == 7
assert postprocessor.steps[1].binarize_gripper is True
assert preprocessor.steps[2].features[OBS_STATE].shape == (MAX_STATE_DIM,)
assert preprocessor.steps[2]._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
assert preprocessor.steps[2]._tensor_stats[ACTION]["min"].shape == (8,)
assert postprocessor.steps[0].features[ACTION].shape == (8,)
assert postprocessor.steps[0]._tensor_stats[ACTION]["min"].shape == (8,)
preprocessor, postprocessor = ensure_evo1_processor_steps(config, preprocessor, postprocessor)
assert sum(isinstance(step, Evo1PadStateProcessorStep) for step in preprocessor.steps) == 1
assert sum(isinstance(step, Evo1PadActionProcessorStep) for step in preprocessor.steps) == 1
assert sum(isinstance(step, Evo1ActionProcessorStep) for step in postprocessor.steps) == 1
def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
monkeypatch.setattr(modeling_evo1, "EVO1", DummyEVO1)
policy = modeling_evo1.EVO1Policy(make_config())
preprocessor, _postprocessor = make_evo1_pre_post_processors(policy.config, dataset_stats=make_stats())
training_batch = preprocessor(make_batch(include_action=True))
loss, metrics = policy.forward(make_batch(include_action=True))
assert training_batch[ACTION].shape == (2, CHUNK_SIZE, MAX_ACTION_DIM)
assert training_batch["action_mask"].shape == (2, CHUNK_SIZE, MAX_ACTION_DIM)
assert training_batch["action_mask"][:, :, :ACTION_DIM].all()
assert not training_batch["action_mask"][:, :, ACTION_DIM:].any()
loss, metrics = policy.forward(training_batch)
assert loss.ndim == 0
assert torch.isfinite(loss)
assert metrics["active_action_dims"] == ACTION_DIM * CHUNK_SIZE
assert policy.model.get_vl_embeddings_calls == 1
action_chunk = policy.predict_action_chunk(make_batch(include_action=False))
assert action_chunk.shape == (2, CHUNK_SIZE, ACTION_DIM)
assert action_chunk.shape == (2, CHUNK_SIZE, MAX_ACTION_DIM)
policy.reset()
selected = policy.select_action(make_batch(include_action=False))
assert selected.shape == (2, ACTION_DIM)
assert selected.shape == (2, MAX_ACTION_DIM)
def test_stage1_frozen_vlm_embeddings_do_not_track_gradients(monkeypatch):
@@ -220,7 +371,7 @@ def test_stage1_frozen_vlm_embeddings_do_not_track_gradients(monkeypatch):
assert policy.model.grad_enabled_calls == [False]
assert policy.model.embedder_training_calls == [False]
assert not fused_tokens.requires_grad
assert policy.model.embedder.training is True
assert policy.model.embedder.training is False
def test_stage2_vlm_embeddings_track_gradients(monkeypatch):