From a17935e7631b3f0d2e8c60d34b121a4ebf594c8a Mon Sep 17 00:00:00 2001 From: Sonica-B Date: Sun, 25 Oct 2026 21:13:07 -0700 Subject: [PATCH 1/2] [BugFix] Accept a single bound in RewardClipping RewardClipping declares clamp_min and clamp_max as optional (None by default, and RewardClippingConfig exposes them as optional fields), and _apply_transform has branches for a one-sided clip. The constructor, however, passes both values to torch.as_tensor, so RewardClipping(clamp_max=1.0) or RewardClipping(clamp_min=0.0) raises "RuntimeError: Could not infer dtype of NoneType" before the transform is built. Keep a missing bound as None, bound the reward spec on the open side by the extreme value of the reward dtype (as ClipTransform does), print None in __repr__, and raise a TypeError when neither bound is given. The docstring now names the real arguments (clamp_min/clamp_max, not clip_min/clip_max), documents in_keys/out_keys and has an example. --- test/transforms/test_reward_transforms.py | 29 ++++++++++ torchrl/envs/transforms/_reward.py | 66 +++++++++++++++++------ 2 files changed, 79 insertions(+), 16 deletions(-) diff --git a/test/transforms/test_reward_transforms.py b/test/transforms/test_reward_transforms.py index 1c000c29a70..925b666d41e 100644 --- a/test/transforms/test_reward_transforms.py +++ b/test/transforms/test_reward_transforms.py @@ -274,6 +274,35 @@ def test_transform_no_env(self, reward_key): assert (td[reward_key] <= 0.1).all() assert (td[reward_key] >= -0.1).all() + @pytest.mark.parametrize( + "clamp_min,clamp_max,expected", + [ + (-0.1, None, [-0.1, -0.05, 0.05, 1.0]), + (None, 0.1, [-1.0, -0.05, 0.05, 0.1]), + ], + ) + def test_one_sided_bound(self, clamp_min, clamp_max, expected): + reward = torch.tensor([[-1.0], [-0.05], [0.05], [1.0]]) + t = RewardClipping(clamp_min=clamp_min, clamp_max=clamp_max) + out = t(TensorDict(reward=reward, batch_size=[4]))["reward"] + torch.testing.assert_close(out, torch.tensor(expected).unsqueeze(-1)) + + env = TransformedEnv( + ContinuousActionVecMockEnv(), + RewardClipping(clamp_min=clamp_min, clamp_max=clamp_max), + ) + check_env_specs(env) + space = env.reward_spec.space + dtype_max = torch.finfo(env.reward_spec.dtype).max + low = -dtype_max if clamp_min is None else clamp_min + high = dtype_max if clamp_max is None else clamp_max + torch.testing.assert_close(space.low, torch.full_like(space.low, low)) + torch.testing.assert_close(space.high, torch.full_like(space.high, high)) + + def test_no_bound_raises(self): + with pytest.raises(TypeError, match="At least one of"): + RewardClipping() + def test_transform_compose(self): t = Compose(RewardClipping(-0.1, 0.1)) td = TensorDict({"reward": torch.randn(10)}, []) diff --git a/torchrl/envs/transforms/_reward.py b/torchrl/envs/transforms/_reward.py index 94009d79e6c..fe673cfc082 100644 --- a/torchrl/envs/transforms/_reward.py +++ b/torchrl/envs/transforms/_reward.py @@ -267,15 +267,41 @@ class RewardClipping(Transform): """Clips the reward between `clamp_min` and `clamp_max`. Args: - clip_min (scalar): minimum value of the resulting reward. - clip_max (scalar): maximum value of the resulting reward. + clamp_min (float or torch.Tensor, optional): minimum value of the + resulting reward. If ``None``, the reward is not bounded from below. + Defaults to ``None``. + clamp_max (float or torch.Tensor, optional): maximum value of the + resulting reward. If ``None``, the reward is not bounded from above. + Defaults to ``None``. + in_keys (sequence of NestedKey, optional): the reward entries to clip. + Defaults to ``["reward"]``. + out_keys (sequence of NestedKey, optional): the entries where the + clipped rewards are written. Defaults to ``in_keys``. + + At least one of ``clamp_min`` and ``clamp_max`` must be provided. When only + one bound is given, the transformed reward spec is bounded by that value on + one side and by the most extreme value of the reward dtype on the other. + + Examples: + >>> import torch + >>> from tensordict import TensorDict + >>> from torchrl.envs import GymEnv, RewardClipping, TransformedEnv + >>> t = RewardClipping(clamp_max=1.0) + >>> td = TensorDict(reward=torch.tensor([[-3.0], [0.5], [2.0]]), batch_size=[3]) + >>> t(td)["reward"].squeeze(-1) + tensor([-3.0000, 0.5000, 1.0000]) + >>> env = TransformedEnv(GymEnv("Pendulum-v1"), RewardClipping(clamp_min=-1.0)) + >>> env.reward_spec.space.low + tensor([-1.]) + >>> rollout = env.rollout(10) + >>> assert (rollout["next", "reward"] >= -1.0).all() """ def __init__( self, - clamp_min: float | None = None, - clamp_max: float | None = None, + clamp_min: float | Tensor | None = None, + clamp_max: float | Tensor | None = None, in_keys: Sequence[NestedKey] | None = None, out_keys: Sequence[NestedKey] | None = None, ): @@ -284,14 +310,16 @@ def __init__( if out_keys is None: out_keys = copy(in_keys) super().__init__(in_keys=in_keys, out_keys=out_keys) - clamp_min_tensor = ( - clamp_min if isinstance(clamp_min, Tensor) else torch.as_tensor(clamp_min) - ) - clamp_max_tensor = ( - clamp_max if isinstance(clamp_max, Tensor) else torch.as_tensor(clamp_max) - ) - self.register_buffer("clamp_min", clamp_min_tensor) - self.register_buffer("clamp_max", clamp_max_tensor) + if clamp_min is None and clamp_max is None: + raise TypeError( + "At least one of `clamp_min` and `clamp_max` must be provided." + ) + if clamp_min is not None and not isinstance(clamp_min, Tensor): + clamp_min = torch.as_tensor(clamp_min) + if clamp_max is not None and not isinstance(clamp_max, Tensor): + clamp_max = torch.as_tensor(clamp_max) + self.register_buffer("clamp_min", clamp_min) + self.register_buffer("clamp_max", clamp_max) def _apply_transform(self, reward: torch.Tensor) -> torch.Tensor: if self.clamp_max is not None and self.clamp_min is not None: @@ -305,9 +333,14 @@ def _apply_transform(self, reward: torch.Tensor) -> torch.Tensor: @_apply_to_composite def transform_reward_spec(self, reward_spec: TensorSpec) -> TensorSpec: if isinstance(reward_spec, Unbounded): + dtype_info = ( + torch.finfo(reward_spec.dtype) + if reward_spec.dtype.is_floating_point + else torch.iinfo(reward_spec.dtype) + ) return Bounded( - self.clamp_min, - self.clamp_max, + self.clamp_min if self.clamp_min is not None else dtype_info.min, + self.clamp_max if self.clamp_max is not None else dtype_info.max, shape=reward_spec.shape, device=reward_spec.device, dtype=reward_spec.dtype, @@ -320,10 +353,11 @@ def transform_reward_spec(self, reward_spec: TensorSpec) -> TensorSpec: ) def __repr__(self) -> str: + clamp_min = None if self.clamp_min is None else f"{float(self.clamp_min):4.4f}" + clamp_max = None if self.clamp_max is None else f"{float(self.clamp_max):4.4f}" return ( f"{self.__class__.__name__}(" - f"clamp_min={float(self.clamp_min):4.4f}, clamp_max" - f"={float(self.clamp_max):4.4f}, keys={self.in_keys})" + f"clamp_min={clamp_min}, clamp_max={clamp_max}, keys={self.in_keys})" ) From 31b4786ac8502e28d057a48b52d78fecf2d75f79 Mon Sep 17 00:00:00 2001 From: Sonica-B Date: Tue, 27 Oct 2026 14:14:10 -0700 Subject: [PATCH 2/2] [BugFix] Raise ValueError and use 4 significant figures in RewardClipping repr Address review on #4622: RewardClipping() with neither clamp_min nor clamp_max now raises ValueError instead of TypeError, and __repr__ formats the bounds with :.4g instead of :4.4f, so a bound such as 1e-5 is no longer printed as 0.0000. Add a test for the repr. --- test/transforms/test_reward_transforms.py | 10 +++++++++- torchrl/envs/transforms/_reward.py | 6 +++--- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/test/transforms/test_reward_transforms.py b/test/transforms/test_reward_transforms.py index 925b666d41e..af943ee3ce8 100644 --- a/test/transforms/test_reward_transforms.py +++ b/test/transforms/test_reward_transforms.py @@ -300,9 +300,17 @@ def test_one_sided_bound(self, clamp_min, clamp_max, expected): torch.testing.assert_close(space.high, torch.full_like(space.high, high)) def test_no_bound_raises(self): - with pytest.raises(TypeError, match="At least one of"): + with pytest.raises(ValueError, match="At least one of"): RewardClipping() + def test_repr(self): + assert repr(RewardClipping(-0.1, 0.1)) == ( + "RewardClipping(clamp_min=-0.1, clamp_max=0.1, keys=['reward'])" + ) + assert repr(RewardClipping(clamp_max=1e-5)) == ( + "RewardClipping(clamp_min=None, clamp_max=1e-05, keys=['reward'])" + ) + def test_transform_compose(self): t = Compose(RewardClipping(-0.1, 0.1)) td = TensorDict({"reward": torch.randn(10)}, []) diff --git a/torchrl/envs/transforms/_reward.py b/torchrl/envs/transforms/_reward.py index fe673cfc082..f4fea1eb298 100644 --- a/torchrl/envs/transforms/_reward.py +++ b/torchrl/envs/transforms/_reward.py @@ -311,7 +311,7 @@ def __init__( out_keys = copy(in_keys) super().__init__(in_keys=in_keys, out_keys=out_keys) if clamp_min is None and clamp_max is None: - raise TypeError( + raise ValueError( "At least one of `clamp_min` and `clamp_max` must be provided." ) if clamp_min is not None and not isinstance(clamp_min, Tensor): @@ -353,8 +353,8 @@ def transform_reward_spec(self, reward_spec: TensorSpec) -> TensorSpec: ) def __repr__(self) -> str: - clamp_min = None if self.clamp_min is None else f"{float(self.clamp_min):4.4f}" - clamp_max = None if self.clamp_max is None else f"{float(self.clamp_max):4.4f}" + clamp_min = None if self.clamp_min is None else f"{float(self.clamp_min):.4g}" + clamp_max = None if self.clamp_max is None else f"{float(self.clamp_max):.4g}" return ( f"{self.__class__.__name__}(" f"clamp_min={clamp_min}, clamp_max={clamp_max}, keys={self.in_keys})"