mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-02 06:29:50 +00:00
Maxime Ellerbach
c15f6acb7e
LingBot-VA's action stream doubles as its memory: _compute_kv_cache feeds the previous chunk's predicted actions back as clean context. With use_relative_actions each chunk is anchored on the state at its own start, so the same physical motion encodes differently depending on which chunk it fell in, and the cached stream sawtooths back to zero displacement at every boundary (mean 0.22, p95 0.91 normalized units against a total output range of 2.0). Training never contains such a reset -- one sample is exactly one chunk, hence one anchor -- and the model's only cross-frame supervision is "continue from the previous clean action frame", so it continues rather than restarting and the postprocessor double-counts the displacement: ~7.3 deg/chunk on the folding data, matching the observed monotonic drift after 4-5 chunks. First-difference the conditioning instead. An increment carries no anchor, so the cached history stays consistent across chunks; measured boundary inconsistency drops 7.30 -> 1.78 deg, the latter being the controller's tracking error, which is generic to chunked control rather than an artifact of the action encoding. The conversion touches ONLY the conditioning -- at inference in _preprocess_action_state, and on the matching clean copy in _add_noise_stream so training agrees. Predicted actions keep the cumulative-offset encoding, so the processors, the checkpoint pipeline and the executed commands are untouched. That is also what makes it safe to do here: normalization is affine with a non-zero constant, so cumsum does not commute with unnormalization -- inverting an accumulated stream would drift by (n-1)*(q99+q01)/2, which is 0 on a symmetric channel but 118 deg over one chunk on right_joint_6. Nothing inverts the conditioning, so that cannot arise. Gated on use_relative_actions only: an absolute action stream is already globally consistent, so those checkpoints keep their exact behaviour. A relative checkpoint is only valid for the encoding it was trained with. Also fixes observation_delta_indices: AutoencoderKLWan._encode consumes only the first 4 * (iter_ - 1) + 1 frames of an n-frame clip, so asking for frame_chunk_size * 4 decoded 3 frames per sample that never reached the encoder (verified by ablation -- scrambling them left the latents bit-identical). Same latent frame count, less video decode. The observation window stays shorter than the action window, so sample validity is unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
93 lines
4.3 KiB
Python
93 lines
4.3 KiB
Python
#!/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.
|
|
# 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.
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
|
|
from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig
|
|
from lerobot.policies.lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors
|
|
from lerobot.processor import PolicyProcessorPipeline, UnnormalizerProcessorStep
|
|
from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
|
|
from lerobot.utils.constants import (
|
|
ACTION,
|
|
OBS_IMAGES,
|
|
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
|
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
|
)
|
|
|
|
|
|
def _make_config(*, action_norm: NormalizationMode = NormalizationMode.QUANTILES) -> LingBotVAConfig:
|
|
# LingBotVAConfig defaults ACTION to IDENTITY (the LIBERO delta-EEF recipe), so the quantile
|
|
# tests below have to set a stats-based mode explicitly -- relying on the default is what left
|
|
# test_postprocessor_quantile_unnormalization asserting against an identity unnormalizer.
|
|
cfg = LingBotVAConfig(device="cpu")
|
|
cfg.normalization_mapping = {**cfg.normalization_mapping, "ACTION": action_norm}
|
|
cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128))}
|
|
cfg.output_features = {}
|
|
cfg.validate_features()
|
|
return cfg
|
|
|
|
|
|
def test_make_pre_post_processors_names_and_steps() -> None:
|
|
cfg = _make_config()
|
|
pre, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=None)
|
|
assert pre.name == POLICY_PREPROCESSOR_DEFAULT_NAME
|
|
assert post.name == POLICY_POSTPROCESSOR_DEFAULT_NAME
|
|
# Actions are unnormalized by the standard built-in quantile unnormalizer.
|
|
assert any(isinstance(s, UnnormalizerProcessorStep) for s in post.steps)
|
|
|
|
|
|
def test_freshly_built_postprocessor_is_identity() -> None:
|
|
# Without action stats the quantile unnormalizer is a no-op (identity passthrough): the real
|
|
# per-benchmark q01/q99 are restored from the saved checkpoint on load, not hardcoded here.
|
|
cfg = _make_config()
|
|
_, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=None)
|
|
normed = torch.tensor([[0.3, -0.5, 1.0, -1.0, 0.0, 0.7, -0.2]])
|
|
assert torch.allclose(post(normed), normed, atol=1e-6)
|
|
|
|
|
|
def test_postprocessor_quantile_unnormalization() -> None:
|
|
# QUANTILES unnormalize maps [-1, 1] -> [q01, q99]: -1 -> q01, +1 -> q99.
|
|
cfg = _make_config()
|
|
q01 = [-1.0, -0.5, 0.0, -1.0, -1.0, -1.0, -1.0]
|
|
q99 = [1.0, 0.5, 2.0, 1.0, 1.0, 1.0, 1.0]
|
|
stats = {ACTION: {"q01": q01, "q99": q99}}
|
|
_, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats=stats)
|
|
out_lo = post(torch.full((1, 7), -1.0))
|
|
out_hi = post(torch.full((1, 7), 1.0))
|
|
assert torch.allclose(out_lo, torch.tensor(q01).unsqueeze(0), atol=1e-4)
|
|
assert torch.allclose(out_hi, torch.tensor(q99).unsqueeze(0), atol=1e-4)
|
|
|
|
|
|
def test_postprocessor_stats_survive_save_load(tmp_path) -> None:
|
|
# Regression guard for the Hub mechanism: the q01/q99 stats live in the saved post-processor
|
|
# state and must round-trip through save_pretrained / from_pretrained.
|
|
cfg = _make_config()
|
|
q01 = [-0.6, -0.8, -0.9, -0.1, -0.15, -0.25, -1.0]
|
|
q99 = [0.9, 0.85, 0.9, 0.17, 0.18, 0.34, 1.0]
|
|
_, post = make_lingbot_va_pre_post_processors(cfg, dataset_stats={ACTION: {"q01": q01, "q99": q99}})
|
|
post.save_pretrained(tmp_path)
|
|
loaded = PolicyProcessorPipeline.from_pretrained(
|
|
tmp_path,
|
|
config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
|
|
to_transition=policy_action_to_transition,
|
|
to_output=transition_to_policy_action,
|
|
)
|
|
out = loaded(torch.full((1, 7), -1.0))
|
|
assert torch.allclose(out, torch.tensor(q01).unsqueeze(0), atol=1e-4)
|