From 64613d8f03d91549cf7c847d2a1e8374ab5f8922 Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Sat, 25 Jul 2026 17:17:35 +0200 Subject: [PATCH] fix(rewards): harden distributional value target projection - Clamp HL-Gauss targets - Handle terminal mask shapes safely - Align the default smoothing ratio with Stop Regressing --- .../distributional_value_function/common.py | 6 ++- ...iguration_distributional_value_function.py | 4 +- .../modeling_distributional_value_function.py | 40 ++++++++++++------- .../test_distributional_value_function.py | 38 ++++++++++++++++++ .../test_temporal_siglip_value_function.py | 6 +++ 5 files changed, 77 insertions(+), 17 deletions(-) diff --git a/src/lerobot/rewards/distributional_value_function/common.py b/src/lerobot/rewards/distributional_value_function/common.py index 0ffaedede..a7bda4cad 100644 --- a/src/lerobot/rewards/distributional_value_function/common.py +++ b/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) diff --git a/src/lerobot/rewards/distributional_value_function/configuration_distributional_value_function.py b/src/lerobot/rewards/distributional_value_function/configuration_distributional_value_function.py index 812275b53..96f2576a9 100644 --- a/src/lerobot/rewards/distributional_value_function/configuration_distributional_value_function.py +++ b/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" diff --git a/src/lerobot/rewards/distributional_value_function/modeling_distributional_value_function.py b/src/lerobot/rewards/distributional_value_function/modeling_distributional_value_function.py index 53427b95d..d60867ea7 100644 --- a/src/lerobot/rewards/distributional_value_function/modeling_distributional_value_function.py +++ b/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.squeeze(-1) - + target_value = target_value.reshape(-1).clamp( + 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. - "A Distributional Perspective on Reinforcement Learning" - arXiv:1707.06887 + ``dirac_delta`` is retained as the public configuration name for + checkpoint compatibility. This is the scalar two-hot projection, not + 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.squeeze(-1) - - target_value = target_value.clamp(self.config.value_support_min, self.config.value_support_max) + target_value = target_value.reshape(-1).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.squeeze(-1) + target_value = target_value.reshape(-1).clamp( + 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) diff --git a/tests/rewards/test_distributional_value_function.py b/tests/rewards/test_distributional_value_function.py index 77a383ecf..769b982c1 100644 --- a/tests/rewards/test_distributional_value_function.py +++ b/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.""" diff --git a/tests/rewards/test_temporal_siglip_value_function.py b/tests/rewards/test_temporal_siglip_value_function.py index beaec5f7a..e434464ff 100644 --- a/tests/rewards/test_temporal_siglip_value_function.py +++ b/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)