diff --git a/src/lerobot/rewards/temporal_siglip_value_function/modeling_temporal_siglip_value_function.py b/src/lerobot/rewards/temporal_siglip_value_function/modeling_temporal_siglip_value_function.py index 9c6cfa13c..a605658e7 100644 --- a/src/lerobot/rewards/temporal_siglip_value_function/modeling_temporal_siglip_value_function.py +++ b/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: diff --git a/src/lerobot/rewards/temporal_siglip_value_function/processor_temporal_siglip_value_function.py b/src/lerobot/rewards/temporal_siglip_value_function/processor_temporal_siglip_value_function.py index 263e3ba43..81520691e 100644 --- a/src/lerobot/rewards/temporal_siglip_value_function/processor_temporal_siglip_value_function.py +++ b/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 diff --git a/tests/rewards/test_temporal_siglip_value_function.py b/tests/rewards/test_temporal_siglip_value_function.py index a0e40c589..faf88758b 100644 --- a/tests/rewards/test_temporal_siglip_value_function.py +++ b/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),