chore: update docs + remove legacy codepaths

This commit is contained in:
Steven Palma
2026-07-02 14:56:42 +02:00
parent f5ac58adb9
commit edc01c3b94
6 changed files with 29 additions and 139 deletions
@@ -100,9 +100,6 @@ class Evo1Config(PreTrainedConfig):
optimizer_grad_clip_norm: float = 1.0
scheduler_warmup_steps: int = 300
# Deprecated, has no effect. Kept only so configs serialized by earlier EVO1 checkpoints
# (which stored this field) can still be parsed; draccus rejects unknown fields.
drop_last: bool = True
def __post_init__(self):
super().__post_init__()
+10 -11
View File
@@ -43,11 +43,11 @@ logger = logging.getLogger(__name__)
def _batched_resize_01(images: torch.Tensor, image_size: int) -> torch.Tensor:
"""Resize a batch of ``[0, 1]`` images to ``(image_size, image_size)`` on-device.
Numerically mirrors InternVL3's per-image PIL preprocessing
Numerically mirrors InternVL3's reference PIL preprocessing
(``to_pil_image`` -> ``Image.resize`` -> ``to_tensor``): the float input is quantized to uint8
exactly as ``to_pil_image`` does, then resized with bicubic interpolation and antialiasing,
which matches PIL's default resampler. This runs as a single batched op instead of a per-image
Python loop with a GPU->CPU->PIL->GPU round-trip.
which matches PIL's default resampler. Matching the reference pixel-for-pixel keeps the policy
interchangeable with checkpoints produced by the upstream EVO1 preprocessing.
Args:
images: float tensor of shape ``(N, C, H, W)`` with values in ``[0, 1]``.
@@ -75,14 +75,14 @@ def _batched_pixel_values(
) -> torch.Tensor:
"""Build InternVL3 ``pixel_values`` from per-camera ``[0, 1]`` image batches without leaving the device.
Equivalent to running the old per-sample/per-image PIL path (resize -> to_tensor -> ImageNet
normalize, a single tile per image) but batched across the whole minibatch. Absent views (fewer
cameras than ``max_views``) are zero-padded to reproduce the previous ``torch.zeros_like``
padding; those views are masked out downstream via the attention mask.
Each image is resized, converted to ``dtype``, and ImageNet-normalized (a single tile per
image), batched across the whole minibatch. Absent views (fewer cameras than ``max_views``)
are filled with zero images; their placeholder tokens are masked out of attention downstream
via ``_mask_absent_image_tokens``.
Returns:
``pixel_values`` of shape ``(B * max_views, C, image_size, image_size)``, ordered row-major
over ``(sample, view)`` to match the old preprocessing.
over ``(sample, view)`` to line up with the per-view image placeholders in the prompt.
"""
resized: list[torch.Tensor] = []
for image in camera_images:
@@ -273,10 +273,9 @@ class InternVL3Embedder(nn.Module):
image_masks: torch.Tensor,
batch_num_tiles_list: list[list[int]],
) -> torch.Tensor:
"""Zero attention over the image-context tokens of absent views, fully vectorized.
"""Zero attention over the image-context tokens of absent (zero-padded) views.
Reproduces the previous per-sample/per-image Python loop, which called ``.item()`` once per
image and forced a device->host sync each time, without any host<->device synchronization.
Fully vectorized: runs without any host<->device synchronization.
"""
# A single tile per image (max_num=1), so every image occupies the same number of
# context tokens.
+1 -1
View File
@@ -359,7 +359,7 @@ class Evo1Policy(PreTrainedPolicy):
# Keep each present camera as a batched (B, C, H, W) tensor on its current (GPU) device.
# Resizing/normalization and zero-padding of absent views happen batched inside the
# embedder, so images never leave the device here (no per-sample .cpu() round-trip).
# embedder, so images never leave the device here.
camera_images: list[Tensor] = []
for camera_key in present_keys:
image = batch[camera_key]
+10 -69
View File
@@ -15,7 +15,7 @@
from __future__ import annotations
from copy import deepcopy
from dataclasses import dataclass, replace
from dataclasses import dataclass
from typing import Any
import torch
@@ -302,74 +302,24 @@ def _pad_evo1_stats(
return padded_stats
def _refresh_evo1_normalization_steps(
config: Evo1Config,
preprocessor: PolicyProcessorPipeline,
postprocessor: PolicyProcessorPipeline,
) -> None:
normalization_features = _evo1_normalization_features(config)
action_features = _evo1_action_features(config)
for step in preprocessor.steps:
if isinstance(step, NormalizerProcessorStep):
step.features = normalization_features
step.stats = _pad_evo1_stats(config, step.stats)
step.to(device=step.device, dtype=step.dtype)
for step in postprocessor.steps:
if isinstance(step, UnnormalizerProcessorStep):
step.features = action_features
step.stats = _pad_evo1_stats(config, step.stats)
step.to(device=step.device, dtype=step.dtype)
def ensure_evo1_processor_steps(
def reconcile_evo1_processors(
config: Evo1Config,
preprocessor: PolicyProcessorPipeline,
postprocessor: PolicyProcessorPipeline,
) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
"""Reconcile checkpoint-loaded pipelines with the current EVO1 config.
Adds the EVO1 steps when loading older checkpoints that do not serialize them, restores the
EVO1 batch converter (converters are not serialized), and refreshes the config-driven step
parameters (padding widths, action cropping, gripper binarization) so CLI overrides at
load/eval time take effect on checkpoints that already serialize these steps.
Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
(converters are plain functions and are never serialized), and eval-time CLI overrides of the
action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This
restores the converter and rebuilds the action step from the current config so those overrides
take effect.
"""
# Pipelines reloaded from a checkpoint come back with the default batch converter, which drops
# non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
preprocessor.to_transition = evo1_batch_to_transition
has_state_padding = any(isinstance(step, Evo1PadStateProcessorStep) for step in preprocessor.steps)
if not has_state_padding:
steps = list(preprocessor.steps)
insert_idx = next(
(idx for idx, step in enumerate(steps) if isinstance(step, NormalizerProcessorStep)),
len(steps),
)
steps.insert(insert_idx, Evo1PadStateProcessorStep(max_state_dim=config.max_state_dim))
preprocessor.steps = steps
has_action_padding = any(isinstance(step, Evo1PadActionProcessorStep) for step in preprocessor.steps)
if not has_action_padding:
steps = list(preprocessor.steps)
insert_idx = next(
(idx for idx, step in enumerate(steps) if isinstance(step, NormalizerProcessorStep)),
len(steps),
)
steps.insert(insert_idx, Evo1PadActionProcessorStep(max_action_dim=config.max_action_dim))
preprocessor.steps = steps
preprocessor.steps = [
replace(step, max_state_dim=config.max_state_dim)
if isinstance(step, Evo1PadStateProcessorStep)
else replace(step, max_action_dim=config.max_action_dim)
if isinstance(step, Evo1PadActionProcessorStep)
else step
for step in preprocessor.steps
]
current_action_step = Evo1ActionProcessorStep(
action_step = Evo1ActionProcessorStep(
action_dim=_evo1_action_dim(config),
binarize_gripper=config.binarize_gripper,
gripper_index=config.gripper_index,
@@ -386,20 +336,11 @@ def ensure_evo1_processor_steps(
(idx + 1 for idx, step in enumerate(steps) if isinstance(step, UnnormalizerProcessorStep)),
0,
)
steps.insert(insert_idx, current_action_step)
steps.insert(insert_idx, action_step)
else:
steps[action_step_idx] = current_action_step
# Actions must leave the postprocessor as float32 (numpy cannot represent bf16); older
# checkpoints serialized the device step without a float_dtype.
steps = [
replace(step, float_dtype="float32")
if isinstance(step, DeviceProcessorStep) and step.float_dtype is None
else step
for step in steps
]
steps[action_step_idx] = action_step
postprocessor.steps = steps
_refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
return preprocessor, postprocessor
+2 -2
View File
@@ -338,9 +338,9 @@ def make_pre_post_processors(
)
_reconnect_relative_absolute_steps(preprocessor, postprocessor)
if isinstance(policy_cfg, Evo1Config):
from .evo1.processor_evo1 import ensure_evo1_processor_steps
from .evo1.processor_evo1 import reconcile_evo1_processors
preprocessor, postprocessor = ensure_evo1_processor_steps(
preprocessor, postprocessor = reconcile_evo1_processors(
policy_cfg,
preprocessor,
postprocessor,
+6 -53
View File
@@ -34,9 +34,9 @@ from lerobot.policies.evo1.processor_evo1 import (
Evo1ActionProcessorStep,
Evo1PadActionProcessorStep,
Evo1PadStateProcessorStep,
ensure_evo1_processor_steps,
evo1_batch_to_transition,
make_evo1_pre_post_processors,
reconcile_evo1_processors,
)
from lerobot.policies.factory import get_policy_class, make_policy_config
from lerobot.processor import (
@@ -444,55 +444,6 @@ def test_evo1_postprocessor_returns_float32_for_bf16_actions():
assert processed.dtype == torch.float32
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 preprocessor.to_transition is evo1_batch_to_transition
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_processor_save_load_round_trip_applies_config_overrides(tmp_path):
train_config = make_config()
preprocessor, postprocessor = make_evo1_pre_post_processors(train_config, dataset_stats=make_stats())
@@ -512,14 +463,16 @@ def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path):
to_output=transition_to_policy_action,
)
# Simulate eval-time CLI overrides on a checkpoint that already serializes the EVO1 steps.
# Simulate eval-time CLI overrides applied on top of the loaded pipelines.
eval_config = make_config(binarize_gripper=True, postprocess_action_dim=ACTION_DIM)
loaded_pre, loaded_post = ensure_evo1_processor_steps(eval_config, loaded_pre, loaded_post)
loaded_pre, loaded_post = reconcile_evo1_processors(eval_config, loaded_pre, loaded_post)
assert loaded_pre.to_transition is evo1_batch_to_transition
assert sum(isinstance(step, Evo1ActionProcessorStep) for step in loaded_post.steps) == 1
action_step = next(step for step in loaded_post.steps if isinstance(step, Evo1ActionProcessorStep))
assert action_step.binarize_gripper is True
assert action_step.action_dim == ACTION_DIM
# The float32 output dtype is part of the serialized pipeline itself.
device_step = next(step for step in loaded_post.steps if isinstance(step, DeviceProcessorStep))
assert device_step.float_dtype == "float32"
@@ -798,7 +751,7 @@ def test_evo1_batched_pixel_values_shape_and_zero_padding():
assert pixel_values.shape == (batch_size * max_views, 3, image_size, image_size)
grouped = pixel_values.reshape(batch_size, max_views, 3, image_size, image_size)
# Absent views (indices 1, 2) are zero images normalized to -mean/std, matching the old padding.
# Absent views (indices 1, 2) are zero images, normalized to the constant -mean/std.
expected_pad = (-mean / std).view(1, 3, 1, 1)
for view in (1, 2):
assert torch.allclose(