fix temporal siglip value function

src/lerobot/rewards/temporal_siglip_value_function/modeling_temporal_siglip_value_function.py

@@ -142,8 +142,9 @@ class TemporalSiglipVFRewardModel(DistributionalValueMixin, PreTrainedRewardMode
             src_key_padding_mask=~frame_valid,
             is_causal=True,
         )
-        last_valid = frame_valid.long().sum(-1).sub(1).clamp_min(0)
-        return hidden[torch.arange(batch_size, device=hidden.device), last_valid]
+        # The history window is ordered oldest→current and the current frame is
+        # always the final, non-padding element.
+        return hidden[:, -1]
 
     def _fit_state_dim(self, state: Tensor) -> Tensor:
         if state.shape[-1] > self.config.state_dim:

src/lerobot/rewards/temporal_siglip_value_function/processor_temporal_siglip_value_function.py

@@ -52,20 +52,24 @@ class TemporalSiglipImageProcessorStep(ProcessorStep):
                 )
                 continue
 
-            image = observation[key].float()
+            image = observation[key]
             if image.ndim == 4:
                 image = image[:, None]
             if image.ndim != 5 or image.shape[2] != 3:
                 raise ValueError(f"Expected {key} as [B,T,3,H,W], got {tuple(image.shape)}")
             batch_size, history = image.shape[:2]
+            image = image.float() / 127.5 - 1.0 if image.dtype == torch.uint8 else image.float() * 2.0 - 1.0
             image = image.flatten(0, 1).permute(0, 2, 3, 1)
-            image = image * 2.0 - 1.0
             if image.shape[1:3] != self.image_resolution:
                 image = resize_with_pad_torch(image, *self.image_resolution)
             observation[key] = image.permute(0, 3, 1, 2).unflatten(0, (batch_size, history))
-            observation[key + IMAGE_MASK_SUFFIX] = torch.ones(
-                batch_size, history, dtype=torch.bool, device=image.device
-            )
+            padding_key = f"{key}_is_pad"
+            if padding_key in observation:
+                observation[key + IMAGE_MASK_SUFFIX] = ~observation[padding_key].bool()
+            else:
+                observation[key + IMAGE_MASK_SUFFIX] = torch.ones(
+                    batch_size, history, dtype=torch.bool, device=image.device
+                )
         transition[TransitionKey.OBSERVATION] = observation
         return transition
 

tests/rewards/test_temporal_siglip_value_function.py

@@ -53,12 +53,15 @@ def test_temporal_image_processor():
     )
     transition = {
         TransitionKey.OBSERVATION: {
-            CAMERAS[0]: torch.full((1, 2, 3, 20, 16), 0.5),
+            CAMERAS[0]: torch.full((1, 2, 3, 20, 16), 128, dtype=torch.uint8),
+            CAMERAS[0] + "_is_pad": torch.tensor([[True, False]]),
         }
     }
     observation = step(transition)[TransitionKey.OBSERVATION]
     assert observation[CAMERAS[0]].shape == (1, 2, 3, 32, 32)
     assert observation[CAMERAS[0] + ".mask"].shape == (1, 2)
+    assert observation[CAMERAS[0] + ".mask"].tolist() == [[False, True]]
+    assert -1.0 <= observation[CAMERAS[0]].min() <= observation[CAMERAS[0]].max() <= 1.0
 
 
 def test_temporal_model_forward(monkeypatch):
@@ -89,7 +92,7 @@ def test_temporal_model_forward(monkeypatch):
     model = modeling.TemporalSiglipVFRewardModel(_config())
     batch = {
         **{key: torch.rand(1, 2, 3, 16, 16) for key in CAMERAS},
-        **{key + ".mask": torch.ones(1, 2, dtype=torch.bool) for key in CAMERAS},
+        **{key + ".mask": torch.tensor([[False, True]]) for key in CAMERAS},
         OBS_STATE: torch.rand(1, 2, 4),
         OBS_LANGUAGE_TOKENS: torch.ones(1, 4, dtype=torch.long),
         OBS_LANGUAGE_ATTENTION_MASK: torch.ones(1, 4, dtype=torch.bool),