fix(rewards): harden distributional value target projection

- Clamp HL-Gauss targets
- Handle terminal mask shapes safely
- Align the default smoothing ratio with Stop Regressing

src/lerobot/rewards/distributional_value_function/common.py

@@ -40,6 +40,7 @@ class DistributionalValueMixin:
         return (cdf[:, 1:] - cdf[:, :-1]) / normalization
 
     def dirac_delta_target(self, target_value: Tensor) -> Tensor:
+        """Two-hot scalar projection (legacy configuration name: ``dirac_delta``)."""
         target_value = target_value.reshape(-1).clamp(
             self.config.value_support_min, self.config.value_support_max
         )
@@ -73,6 +74,9 @@ class DistributionalValueMixin:
             raise ValueError(f"Unknown target method: {self.config.target_method}")
         if not self.config.use_one_hot_terminal:
             return base
+        is_terminal = is_terminal.reshape(-1)
+        if is_terminal.numel() != base.shape[0]:
+            raise ValueError(f"Expected {base.shape[0]} terminal flags, got {is_terminal.numel()}")
         nearest = torch.argmin(
             torch.abs(
                 self.value_head.bin_centers[None]
@@ -81,7 +85,7 @@ class DistributionalValueMixin:
             dim=-1,
         )
         terminal = F.one_hot(nearest, num_classes=self.config.num_value_bins).to(base.dtype)
-        return torch.where(is_terminal.reshape(-1, 1).bool(), terminal, base)
+        return torch.where(is_terminal[:, None].bool(), terminal, base)
 
     def _distributional_forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Any]]:
         readout = self._get_value_readout(batch)

src/lerobot/rewards/distributional_value_function/configuration_distributional_value_function.py

@@ -67,7 +67,9 @@ class DistributionalVFConfig(RewardModelConfig):
     num_value_bins: int = 201
     value_support_min: float = -1.0
     value_support_max: float = 0.0
-    hl_gauss_sigma_ratio: float = 5.0
+    # Stop Regressing (Farebrother et al., 2024) default: spreads most
+    # probability mass across approximately six neighboring bins.
+    hl_gauss_sigma_ratio: float = 0.75
 
     # Target distribution method: "dirac_delta" (paper-faithful C51) or "hl_gauss" (soft)
     target_method: str = "dirac_delta"

src/lerobot/rewards/distributional_value_function/modeling_distributional_value_function.py

@@ -325,9 +325,9 @@ class DistributionalVFRewardModel(PreTrainedRewardModel):
         Returns:
             [batch_size, num_value_bins] target probability distribution.
         """
-        if target_value.ndim == 2:
+        target_value = target_value.reshape(-1).clamp(
-            target_value = target_value.squeeze(-1)
+            self.config.value_support_min, self.config.value_support_max
-
+        )
         target_value = target_value.to(dtype=self.value_head.bin_centers.dtype)
 
         # Bin edges: half a bin-width outside the first/last center
@@ -360,11 +360,11 @@ class DistributionalVFRewardModel(PreTrainedRewardModel):
         return bin_probabilities
 
     def dirac_delta_target(self, target_value: Tensor) -> Tensor:
-        """Dirac delta (C51) projection: split probability between two nearest bins.
+        """Two-hot scalar projection between adjacent bins.
 
-        Standard distributional RL projection from Bellemare et al. 2017.
+        ``dirac_delta`` is retained as the public configuration name for
-        "A Distributional Perspective on Reinforcement Learning"
+        checkpoint compatibility. This is the scalar two-hot projection, not
-        arXiv:1707.06887
+        the full distributional Bellman projection used by C51.
 
         Args:
             target_value: [batch_size] or [batch_size, 1] target values.
@@ -372,10 +372,9 @@ class DistributionalVFRewardModel(PreTrainedRewardModel):
         Returns:
             [batch_size, num_value_bins] target probability distribution.
         """
-        if target_value.ndim == 2:
+        target_value = target_value.reshape(-1).clamp(
-            target_value = target_value.squeeze(-1)
+            self.config.value_support_min, self.config.value_support_max
-
+        )
-        target_value = target_value.clamp(self.config.value_support_min, self.config.value_support_max)
         target_value = target_value.to(dtype=self.value_head.bin_centers.dtype)
 
         bin_width = self.value_head.bin_centers[1] - self.value_head.bin_centers[0]
@@ -391,7 +390,12 @@ class DistributionalVFRewardModel(PreTrainedRewardModel):
         weight_lower = torch.where(same_bin, torch.ones_like(weight_lower), weight_lower)
 
         batch_size = target_value.shape[0]
-        target_distribution = torch.zeros(batch_size, self.config.num_value_bins, device=target_value.device)
+        target_distribution = torch.zeros(
+            batch_size,
+            self.config.num_value_bins,
+            device=target_value.device,
+            dtype=target_value.dtype,
+        )
         batch_indices = torch.arange(batch_size, device=target_value.device)
         target_distribution[batch_indices, lower_bin_idx] += weight_lower
         target_distribution[batch_indices, upper_bin_idx] += weight_upper
@@ -399,7 +403,7 @@ class DistributionalVFRewardModel(PreTrainedRewardModel):
         return target_distribution
 
     def one_hot_target(self, target_value: Tensor) -> Tensor:
-        """One-hot target for terminal states (exact return, no smoothing).
+        """One-hot target at the nearest support bin for terminal states.
 
         Args:
             target_value: [batch_size] or [batch_size, 1] target values.
@@ -407,8 +411,9 @@ class DistributionalVFRewardModel(PreTrainedRewardModel):
         Returns:
             [batch_size, num_value_bins] one-hot distribution at the nearest bin.
         """
-        if target_value.ndim == 2:
+        target_value = target_value.reshape(-1).clamp(
-            target_value = target_value.squeeze(-1)
+            self.config.value_support_min, self.config.value_support_max
+        )
         target_value = target_value.to(dtype=self.value_head.bin_centers.dtype)
         nearest_bin_idx = torch.argmin(
             torch.abs(self.value_head.bin_centers.unsqueeze(0) - target_value.unsqueeze(-1)), dim=-1
@@ -446,6 +451,11 @@ class DistributionalVFRewardModel(PreTrainedRewardModel):
         if not use_one_hot_terminal:
             return base_distribution
 
+        is_terminal = is_terminal.reshape(-1)
+        if is_terminal.numel() != base_distribution.shape[0]:
+            raise ValueError(
+                f"Expected {base_distribution.shape[0]} terminal flags, got {is_terminal.numel()}"
+            )
         terminal_distribution = self.one_hot_target(target_value)
         return torch.where(is_terminal[:, None].bool(), terminal_distribution, base_distribution)
 

tests/rewards/test_distributional_value_function.py

@@ -120,6 +120,7 @@ def test_config_defaults_match_pi06_gemma3_layout():
     assert config.image_resolution == (448, 448)
     assert config.num_image_tokens == 256
     assert config.target_method == "dirac_delta"
+    assert config.hl_gauss_sigma_ratio == 0.75
 
 
 # ------------------------------------------------------------------
@@ -170,6 +171,15 @@ def test_hl_gauss_handles_2d_input():
     torch.testing.assert_close(dist.sum(dim=-1), torch.ones(2), atol=1e-5, rtol=0)
 
 
+@skip_if_package_missing("transformers")
+def test_hl_gauss_clamps_values_outside_support():
+    model = _make_model()
+    outside = model.hl_gauss_target(torch.tensor([-1.5, 0.5]))
+    boundaries = model.hl_gauss_target(torch.tensor([-1.0, 0.0]))
+
+    torch.testing.assert_close(outside, boundaries)
+
+
 @skip_if_package_missing("transformers")
 def test_dirac_delta_sums_to_one():
     """Dirac delta target distribution sums to 1 for each sample."""
@@ -258,6 +268,34 @@ def test_terminal_gets_one_hot():
     assert (dist[2] > 0).sum() > 2
 
 
+@skip_if_package_missing("transformers")
+def test_terminal_column_vector_preserves_batch_shape():
+    model = _make_model()
+    targets = torch.tensor([[-0.5], [-0.3]])
+    is_terminal = torch.tensor([[False], [True]])
+
+    dist = model.compute_target_distribution(
+        targets, is_terminal, method="hl_gauss", use_one_hot_terminal=True
+    )
+
+    assert dist.shape == (2, NUM_BINS)
+    assert (dist[0] > 0).sum() > 2
+    assert (dist[1] > 0).sum() == 1
+
+
+@skip_if_package_missing("transformers")
+def test_terminal_count_must_match_batch_size():
+    model = _make_model()
+
+    with pytest.raises(ValueError, match="Expected 2 terminal flags, got 1"):
+        model.compute_target_distribution(
+            torch.tensor([-0.5, -0.3]),
+            torch.tensor([True]),
+            method="hl_gauss",
+            use_one_hot_terminal=True,
+        )
+
+
 @skip_if_package_missing("transformers")
 def test_no_terminal_override_when_disabled():
     """When use_one_hot_terminal=False, terminal states use the base method."""

tests/rewards/test_temporal_siglip_value_function.py

@@ -133,3 +133,9 @@ def test_temporal_model_forward(monkeypatch):
         eval_loss, eval_metrics = model(batch)
     assert torch.isfinite(eval_loss)
     assert -1.0 <= eval_metrics["predicted_value_mean"] <= 0.0
+
+    targets = model.compute_target_distribution(
+        torch.tensor([[-0.5], [-0.3]]),
+        torch.tensor([[False], [True]]),
+    )
+    assert targets.shape == (2, model.config.num_value_bins)