Fix 2.8 issue per sample grad

advanced_source/pendulum.py

@@ -100,7 +100,7 @@
 from tensordict.nn import TensorDictModule
 from torch import nn
 
-from torchrl.data import BoundedTensorSpec, CompositeSpec, UnboundedContinuousTensorSpec
+from torchrl.data import Bounded, Composite, Unbounded
 from torchrl.envs import (
     CatTensors,
     EnvBase,
@@ -403,14 +403,14 @@ def _reset(self, tensordict):
 
 def _make_spec(self, td_params):
     # Under the hood, this will populate self.output_spec["observation"]
-    self.observation_spec = CompositeSpec(
-        th=BoundedTensorSpec(
+    self.observation_spec = Composite(
+        th=Bounded(
             low=-torch.pi,
             high=torch.pi,
             shape=(),
             dtype=torch.float32,
         ),
-        thdot=BoundedTensorSpec(
+        thdot=Bounded(
             low=-td_params["params", "max_speed"],
             high=td_params["params", "max_speed"],
             shape=(),
@@ -426,24 +426,26 @@ def _make_spec(self, td_params):
     self.state_spec = self.observation_spec.clone()
     # action-spec will be automatically wrapped in input_spec when
     # `self.action_spec = spec` will be called supported
-    self.action_spec = BoundedTensorSpec(
+    self.action_spec = Bounded(
         low=-td_params["params", "max_torque"],
         high=td_params["params", "max_torque"],
         shape=(1,),
         dtype=torch.float32,
     )
-    self.reward_spec = UnboundedContinuousTensorSpec(shape=(*td_params.shape, 1))
+    self.reward_spec = Unbounded(shape=(*td_params.shape, 1))
 
 
 def make_composite_from_td(td):
     # custom function to convert a ``tensordict`` in a similar spec structure
     # of unbounded values.
-    composite = CompositeSpec(
+    composite = Composite(
         {
-            key: make_composite_from_td(tensor)
-            if isinstance(tensor, TensorDictBase)
-            else UnboundedContinuousTensorSpec(
-                dtype=tensor.dtype, device=tensor.device, shape=tensor.shape
+            key: (
+                make_composite_from_td(tensor)
+                if isinstance(tensor, TensorDictBase)
+                else Unbounded(
+                    dtype=tensor.dtype, device=tensor.device, shape=tensor.shape
+                )
             )
             for key, tensor in td.items()
         },
@@ -687,7 +689,7 @@ def _reset(
     # is of type ``Composite``
     @_apply_to_composite
     def transform_observation_spec(self, observation_spec):
-        return BoundedTensorSpec(
+        return Bounded(
             low=-1,
             high=1,
             shape=observation_spec.shape,
@@ -711,7 +713,7 @@ def _reset(
     # is of type ``Composite``
     @_apply_to_composite
     def transform_observation_spec(self, observation_spec):
-        return BoundedTensorSpec(
+        return Bounded(
             low=-1,
             high=1,
             shape=observation_spec.shape,

intermediate_source/per_sample_grads.py

@@ -168,7 +168,36 @@ def compute_loss(params, buffers, sample, target):
 # we can double check that the results using ``grad`` and ``vmap`` match the
 # results of hand processing each one individually:
 
-for per_sample_grad, ft_per_sample_grad in zip(per_sample_grads, ft_per_sample_grads.values()):
+for i, (per_sample_grad, ft_per_sample_grad) in enumerate(
+    zip(per_sample_grads, ft_per_sample_grads.values())
+):
+    is_close = torch.allclose(per_sample_grad, ft_per_sample_grad, atol=3e-3, rtol=1e-5)
+    if not is_close:
+        # Calculate and print the maximum absolute difference
+        abs_diff = (per_sample_grad - ft_per_sample_grad).abs()
+        max_diff = abs_diff.max().item()
+        mean_diff = abs_diff.mean().item()
+        print(f"Gradient {i} mismatch:")
+        print(f"  Max absolute difference: {max_diff}")
+        print(f"  Mean absolute difference: {mean_diff}")
+        print(f"  Shape of tensors: {per_sample_grad.shape}")
+        # Print a sample of values from both tensors where the difference is largest
+        max_idx = abs_diff.argmax().item()
+        flat_idx = max_idx
+        if len(abs_diff.shape) > 1:
+            # Convert flat index to multi-dimensional index
+            indices = []
+            temp_shape = abs_diff.shape
+            for dim in reversed(temp_shape):
+                indices.insert(0, flat_idx % dim)
+                flat_idx //= dim
+            print(f"  Max difference at index: {indices}")
+            print(f"  Manual gradient value: {per_sample_grad[tuple(indices)].item()}")
+            print(
+                f"  Functional gradient value: {ft_per_sample_grad[tuple(indices)].item()}"
+            )
+
+    # Keep the original assertion
     assert torch.allclose(per_sample_grad, ft_per_sample_grad, atol=3e-3, rtol=1e-5)
 
 ######################################################################