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
3.5 KiB
Python
93 lines
3.5 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 pytest
|
|
|
|
from lerobot.configs.policies import PreTrainedConfig
|
|
from lerobot.configs.types import FeatureType, PolicyFeature
|
|
from lerobot.policies.lingbot_va.configuration_lingbot_va import LingBotVAConfig
|
|
from lerobot.utils.constants import ACTION, OBS_IMAGES
|
|
|
|
|
|
def make_config(**overrides) -> LingBotVAConfig:
|
|
kwargs = {"device": "cpu"}
|
|
kwargs.update(overrides)
|
|
return LingBotVAConfig(**kwargs)
|
|
|
|
|
|
def test_registered_in_choice_registry() -> None:
|
|
assert "lingbot_va" in PreTrainedConfig.get_known_choices()
|
|
assert PreTrainedConfig.get_choice_class("lingbot_va") is LingBotVAConfig
|
|
|
|
|
|
def test_type_property() -> None:
|
|
assert make_config().type == "lingbot_va"
|
|
|
|
|
|
def test_chunk_size_and_action_steps() -> None:
|
|
cfg = make_config(frame_chunk_size=4, action_per_frame=4)
|
|
assert cfg.chunk_size == 16
|
|
assert cfg.n_action_steps == 16
|
|
assert cfg.action_delta_indices == list(range(16))
|
|
# 4 * (frame_chunk_size - 1) + 1 frames: exactly what AutoencoderKLWan._encode consumes to
|
|
# produce frame_chunk_size latent frames. Loading 16 decoded 3 frames the VAE never read.
|
|
assert cfg.observation_delta_indices == list(range(13))
|
|
assert cfg.reward_delta_indices is None
|
|
|
|
|
|
def test_observation_delta_indices_are_all_consumed_by_the_vae() -> None:
|
|
"""Every requested frame must survive ``iter_ = 1 + (n - 1) // 4``, and F must be preserved."""
|
|
for frame_chunk_size, action_per_frame in ((1, 4), (2, 16), (4, 4), (3, 8)):
|
|
cfg = make_config(frame_chunk_size=frame_chunk_size, action_per_frame=action_per_frame)
|
|
deltas = cfg.observation_delta_indices
|
|
num_latent_frames = 1 + (len(deltas) - 1) // 4
|
|
assert num_latent_frames == frame_chunk_size, (frame_chunk_size, deltas)
|
|
assert len(deltas) == 4 * (num_latent_frames - 1) + 1, (frame_chunk_size, deltas)
|
|
assert deltas[0] == 0
|
|
assert len(set(deltas)) == len(deltas)
|
|
|
|
|
|
def test_optimizer_and_scheduler_presets() -> None:
|
|
cfg = make_config()
|
|
opt = cfg.get_optimizer_preset()
|
|
assert opt.lr == cfg.optimizer_lr
|
|
sched = cfg.get_scheduler_preset()
|
|
assert sched.num_warmup_steps == cfg.scheduler_warmup_steps
|
|
|
|
|
|
def test_validate_features_sets_action_feature() -> None:
|
|
cfg = make_config()
|
|
cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128))}
|
|
cfg.output_features = {}
|
|
cfg.validate_features()
|
|
assert ACTION in cfg.output_features
|
|
assert cfg.output_features[ACTION].shape == (len(cfg.used_action_channel_ids),)
|
|
|
|
|
|
def test_validate_features_no_visual_raises() -> None:
|
|
cfg = make_config()
|
|
cfg.input_features = {}
|
|
cfg.output_features = {}
|
|
with pytest.raises(ValueError, match="at least one visual input feature"):
|
|
cfg.validate_features()
|
|
|
|
|
|
def test_invalid_attn_mode_raises() -> None:
|
|
with pytest.raises(ValueError, match="attn_mode"):
|
|
make_config(attn_mode="banana")
|