lingbot-va: increment-encoded action KV memory under relative actions

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>

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]:

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):

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)

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()

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()