diff --git a/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py b/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py index 9837b9ef4..bf6d0480d 100644 --- a/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py +++ b/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py @@ -182,9 +182,24 @@ class LingBotVAConfig(PreTrainedConfig): @property def observation_delta_indices(self) -> list[int]: + """Observation frame deltas for the training clip, sized to what the VAE actually reads. + + ``diffusers``' ``AutoencoderKLWan._encode`` runs ``iter_ = 1 + (n - 1) // 4`` passes over + ``x[:, :, :1]`` then ``x[:, :, 1 + 4*(i-1) : 1 + 4*i]``, so it only ever consumes the first + ``4 * (iter_ - 1) + 1`` frames of an ``n``-frame clip. Asking for ``frame_chunk_size * 4`` + frames (the previous formula) therefore decoded 3 frames per sample that never reached the + encoder: at ``frame_chunk_size=2`` the deltas were ``[0, 4, ..., 28]`` and only + ``[0, 4, 8, 12, 16]`` were used -- verified by ablation, scrambling the tail left the latents + bit-identical. + + Requesting exactly ``4 * (frame_chunk_size - 1) + 1`` frames yields the same + ``frame_chunk_size`` latent frames with every loaded frame used, and drops the wasted video + decode. The stride is unchanged, so the frames that do reach the model are the same ones. + """ temporal_downsample = 4 stride = max(1, self.action_per_frame // temporal_downsample) - return list(range(0, self.frame_chunk_size * temporal_downsample * stride, stride)) + num_frames = temporal_downsample * (self.frame_chunk_size - 1) + 1 + return [i * stride for i in range(num_frames)] @property def action_delta_indices(self) -> list[int]: diff --git a/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py index cb93a5565..39b57c5ff 100644 --- a/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py +++ b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py @@ -215,6 +215,13 @@ class LingBotVAPolicy(PreTrainedPolicy): noisy_latents = scheduler.add_noise(latent, noise, timesteps, t_dim=2) targets = scheduler.training_target(latent, noise, timesteps) + # The clean copy below is the *conditioning* the model reads for earlier frames -- at + # inference it is what `_compute_kv_cache` caches. Convert it the same way there, so training + # and rollout agree. `noisy_latents` / `targets` are untouched: the model still predicts + # cumulative offsets, which is what the postprocessor knows how to invert. + if action_mode and self.config.use_relative_actions: + latent = self._to_action_increments(latent) + grid_id = ( get_mesh_id( latent.shape[-3] // patch_f, @@ -436,7 +443,9 @@ class LingBotVAPolicy(PreTrainedPolicy): # First call: this observation conditions the first chunk (it is *not* a keyframe). self._started = True actions, predictions = unpack_action_output( - self.predict_action_chunk(batch, return_intermediate_predictions=return_intermediate_predictions) + self.predict_action_chunk( + batch, return_intermediate_predictions=return_intermediate_predictions + ) ) # [B, chunk_size, n_used] self._action_queue.extend(actions.transpose(0, 1)) # [chunk_size, B, n_used] self._obs_buffer = [] @@ -743,13 +752,41 @@ class LingBotVAPolicy(PreTrainedPolicy): return mask # Action conditioning (executed action history) (de)normalization + def _to_action_increments(self, action_norm: Tensor) -> Tensor: + """First-difference an action chunk along its horizon, for use as conditioning only. + + ``action_norm`` is ``[B, action_dim, F, action_per_frame, 1]``; the horizon runs frame-major + across ``F * action_per_frame``. Element 0 is left as-is. + + Under ``use_relative_actions`` the values are cumulative offsets from the chunk's own anchor, + so the same physical motion encodes differently depending on which chunk it landed in — and + the memory sawtooths back to zero at every boundary. An increment carries no anchor, so the + cached history stays consistent across chunks. + + Applied ONLY to the conditioning: the predicted actions keep the cumulative-offset encoding + the processors expect. That matters, because normalization is affine with a non-zero constant + and ``cumsum`` does not commute with it — 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 the issue cannot arise here. + """ + b, c, f, n, _ = action_norm.shape + flat = action_norm.reshape(b, c, f * n) + increments = torch.cat([flat[..., :1], flat[..., 1:] - flat[..., :-1]], dim=-1) + return increments.reshape(b, c, f, n, 1) + def _preprocess_action_state(self, action_norm: Tensor) -> Tensor: """Build the action-conditioning tensor from the already-normalized executed actions. ``action_norm`` is the model-space action chunk ``[B, action_dim, F, action_per_frame, 1]``. - Upstream re-derives the conditioning from the raw executed action via quantile norm; here - the executed actions are already in the model-normalized space, so we pass them through. + Upstream re-derives the conditioning from the raw executed action via quantile norm; here the + executed actions are already in the model-normalized space, so only the optional increment + conversion is applied (mirrored on the training conditioning in ``_add_noise_stream``). + + Absolute-action runs are left alone: their stream is already globally consistent, so + there is nothing to re-encode and existing checkpoints keep their exact behaviour. """ + if self.config.use_relative_actions: + action_norm = self._to_action_increments(action_norm) return action_norm.to(self.config.device, self.dtype) def _compute_kv_cache(self, obs_buffer, executed_actions): diff --git a/tests/policies/lingbot_va/test_action_memory.py b/tests/policies/lingbot_va/test_action_memory.py new file mode 100644 index 000000000..a14b2e6c8 --- /dev/null +++ b/tests/policies/lingbot_va/test_action_memory.py @@ -0,0 +1,157 @@ +#!/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. + +"""Increment-encoded conditioning for the action KV cache under relative actions. + +LingBot-VA's action stream doubles as its memory -- `_compute_kv_cache` feeds the previous chunk's +actions back as clean context. Under `use_relative_actions` those values are cumulative offsets from +the chunk's own anchor, so the cached stream resets to zero displacement at every boundary; the +model, whose only cross-frame supervision is "continue from the previous clean action frame", +continues instead of restarting and the postprocessor double-counts the displacement. + +The conversion applies to the conditioning ONLY. The predicted actions stay cumulative offsets, so +the processors are untouched and nothing has to invert an accumulated stream. +""" + +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.modeling_lingbot_va import LingBotVAPolicy +from lerobot.utils.constants import ACTION, OBS_IMAGES + +USED_CHANNELS = [21, 22, 23, 24, 25, 26, 27, 29, 14, 15, 16, 17, 18, 19, 20, 28] + + +def _make_config(*, use_relative_actions: bool = True): + cfg = LingBotVAConfig( + device="cpu", + action_per_frame=16, + frame_chunk_size=2, + used_action_channel_ids=list(USED_CHANNELS), + use_relative_actions=use_relative_actions, + normalization_mapping={ + "VISUAL": NormalizationMode.IDENTITY, + "STATE": NormalizationMode.IDENTITY, + "ACTION": NormalizationMode.QUANTILES, + }, + ) + cfg.input_features = {f"{OBS_IMAGES}.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 128, 128))} + cfg.output_features = {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(16,))} + return cfg + + +def _policy(cfg) -> LingBotVAPolicy: + """A bare policy: the conversion reads only `config`, so the 5B transformer is irrelevant here.""" + policy = LingBotVAPolicy.__new__(LingBotVAPolicy) + policy.config = cfg + policy.dtype = torch.float32 + return policy + + +def _ramp_chunk(cfg, *, start: float = 0.0) -> torch.Tensor: + """A chunk of cumulative offsets growing from `start`, shaped [1, action_dim, F, apf, 1].""" + horizon = cfg.frame_chunk_size * cfg.action_per_frame + values = start + torch.arange(horizon, dtype=torch.float32) * 0.05 + chunk = torch.zeros(1, cfg.action_dim, cfg.frame_chunk_size, cfg.action_per_frame, 1) + for channel in USED_CHANNELS: + chunk[0, channel, :, :, 0] = values.reshape(cfg.frame_chunk_size, cfg.action_per_frame) + return chunk + + +def test_increments_are_the_first_difference_along_the_horizon() -> None: + """The horizon runs frame-major across F * action_per_frame, not per frame independently.""" + cfg = _make_config() + chunk = _ramp_chunk(cfg) + out = _policy(cfg)._to_action_increments(chunk) + + flat_in = chunk[0, USED_CHANNELS[0]].reshape(-1) + flat_out = out[0, USED_CHANNELS[0]].reshape(-1) + assert flat_out[0] == flat_in[0] # element 0 is passed through + torch.testing.assert_close(flat_out[1:], flat_in[1:] - flat_in[:-1], atol=1e-6, rtol=0) + # the frame boundary is a normal step, not a restart + boundary = cfg.action_per_frame + torch.testing.assert_close( + flat_out[boundary], flat_in[boundary] - flat_in[boundary - 1], atol=1e-6, rtol=0 + ) + + +def test_conditioning_is_independent_of_the_chunk_anchor() -> None: + """The property that removes the sawtooth: the same motion encodes identically in any chunk.""" + cfg = _make_config() + policy = _policy(cfg) + early = policy._to_action_increments(_ramp_chunk(cfg, start=0.0)) + late = policy._to_action_increments(_ramp_chunk(cfg, start=0.6)) # a later chunk, same motion + + # cumulative offsets differ everywhere; increments differ only at element 0, which carries the + # anchor offset by design (it is the command-vs-measured error, ~2 deg, not a chunk of motion). + assert not torch.allclose(_ramp_chunk(cfg, start=0.0), _ramp_chunk(cfg, start=0.6)) + torch.testing.assert_close(early[:, :, :, 1:], late[:, :, :, 1:], atol=1e-6, rtol=0) + + +def test_preprocess_action_state_applies_the_conversion() -> None: + cfg = _make_config() + chunk = _ramp_chunk(cfg) + out = _policy(cfg)._preprocess_action_state(chunk) + torch.testing.assert_close(out, _policy(cfg)._to_action_increments(chunk), atol=1e-6, rtol=0) + + +def test_absolute_action_runs_are_untouched() -> None: + """The one hard compatibility constraint: an absolute-action checkpoint must be unaffected. + + Its action stream is already globally consistent (the values are joint positions, not offsets + from a per-chunk anchor), so there is nothing to re-encode. + """ + cfg = _make_config(use_relative_actions=False) + chunk = _ramp_chunk(cfg) + out = _policy(cfg)._preprocess_action_state(chunk) + torch.testing.assert_close(out, chunk, atol=0, rtol=0) + + +def test_conversion_is_invertible_so_no_information_is_lost() -> None: + """cumsum recovers the original chunk -- the conditioning is a re-encoding, not a truncation.""" + cfg = _make_config() + chunk = _ramp_chunk(cfg, start=0.3) + increments = _policy(cfg)._to_action_increments(chunk) + b, c, f, n, _ = chunk.shape + recovered = torch.cumsum(increments.reshape(b, c, f * n), dim=-1).reshape(b, c, f, n, 1) + torch.testing.assert_close(recovered, chunk, atol=1e-5, rtol=0) + + +def test_unused_channels_stay_zero() -> None: + cfg = _make_config() + out = _policy(cfg)._to_action_increments(_ramp_chunk(cfg)) + unused = [c for c in range(cfg.action_dim) if c not in USED_CHANNELS] + assert torch.count_nonzero(out[0, unused]) == 0 + + +def test_predicted_actions_are_not_converted() -> None: + """Only the conditioning changes; the output path still emits cumulative offsets. + + Guards the reason this is safe to do inside the model: normalization is affine with a non-zero + constant, so cumsum does not commute with unnormalization. If the *prediction* were differenced, + the postprocessor would drift by (n-1)*(q99+q01)/2 per chunk -- 0 on a symmetric channel but + ~118 deg on right_joint_6. + """ + import inspect + + source = inspect.getsource(LingBotVAPolicy.predict_action_chunk) + assert "_to_action_increments" not in source + # the conversion is reachable only from the KV-cache path + assert "_to_action_increments" in inspect.getsource(LingBotVAPolicy._preprocess_action_state) + assert "_preprocess_action_state" in inspect.getsource(LingBotVAPolicy._compute_kv_cache) diff --git a/tests/policies/lingbot_va/test_configuration_lingbot_va.py b/tests/policies/lingbot_va/test_configuration_lingbot_va.py index 089d5ec02..98ec9c47e 100644 --- a/tests/policies/lingbot_va/test_configuration_lingbot_va.py +++ b/tests/policies/lingbot_va/test_configuration_lingbot_va.py @@ -44,10 +44,24 @@ def test_chunk_size_and_action_steps() -> None: assert cfg.chunk_size == 16 assert cfg.n_action_steps == 16 assert cfg.action_delta_indices == list(range(16)) - assert cfg.observation_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() diff --git a/tests/policies/lingbot_va/test_processor.py b/tests/policies/lingbot_va/test_processor.py index e79771600..0bce42732 100644 --- a/tests/policies/lingbot_va/test_processor.py +++ b/tests/policies/lingbot_va/test_processor.py @@ -18,7 +18,7 @@ from __future__ import annotations import torch -from lerobot.configs.types import FeatureType, PolicyFeature +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 @@ -31,8 +31,12 @@ from lerobot.utils.constants import ( ) -def _make_config() -> LingBotVAConfig: +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()