[Feature] Max pool Transform (#841)

Author: albertbou92

docs/source/reference/envs.rst

@@ -229,6 +229,7 @@ in the environment. The keys to be included in this inverse transform are passed
     SqueezeTransform
     StepCounter
     TensorDictPrimer
+    TimeMaxPool
     ToTensorImage
     UnsqueezeTransform
     VecNorm

test/test_transforms.py

@@ -56,6 +56,7 @@
     SqueezeTransform,
     StepCounter,
     TensorDictPrimer,
+    TimeMaxPool,
     ToTensorImage,
     TransformedEnv,
     UnsqueezeTransform,
@@ -106,7 +107,6 @@ def _test_vecnorm_subproc_auto(
 
     @pytest.mark.parametrize("nprc", [2, 5])
     def test_vecnorm_parallel_auto(self, nprc):
-
         queues = []
         prcs = []
         if _has_gym:
@@ -864,6 +864,34 @@ def test_sum_reward(self, keys, device):
         assert "some_extra_observation" in transformed_observation_spec2.keys()
         assert "episode_reward" in transformed_observation_spec2.keys()
 
+    @pytest.mark.parametrize("T", [2, 4])
+    @pytest.mark.parametrize("seq_len", [8])
+    @pytest.mark.parametrize("device", get_available_devices())
+    def test_time_max_pool(self, T, seq_len, device):
+        batch = 1
+        nodes = 4
+        keys = ["observation"]
+        time_max_pool = TimeMaxPool(keys, T=T)
+
+        tensor_list = []
+        for _ in range(seq_len):
+            tensor_list.append(torch.rand(batch, nodes).to(device))
+        max_vals, _ = torch.max(torch.stack(tensor_list[-T:]), dim=0)
+
+        print(f"max vals: {max_vals}")
+
+        for i in range(seq_len):
+            env_td = TensorDict(
+                {
+                    "observation": tensor_list[i],
+                },
+                device=device,
+                batch_size=[batch],
+            )
+            transformed_td = time_max_pool(env_td)
+
+        assert (max_vals == transformed_td["observation"]).all()
+
     @pytest.mark.parametrize("batch", [[], [1], [3, 2]])
     @pytest.mark.parametrize(
         "keys",
@@ -1667,7 +1695,6 @@ def test_append(self):
         assert obs_spec.shape[-1] == 4 * env.base_env.observation_spec[key].shape[-1]
 
     def test_insert(self):
-
         env = ContinuousActionVecMockEnv()
         obs_spec = env.observation_spec
         (key,) = itertools.islice(obs_spec.keys(), 1)

torchrl/envs/__init__.py

@@ -34,6 +34,7 @@
     SqueezeTransform,
     StepCounter,
     TensorDictPrimer,
+    TimeMaxPool,
     ToTensorImage,
     Transform,
     TransformedEnv,

torchrl/envs/transforms/__init__.py

@@ -30,6 +30,7 @@
     SqueezeTransform,
     StepCounter,
     TensorDictPrimer,
+    TimeMaxPool,
     ToTensorImage,
     Transform,
     TransformedEnv,

torchrl/envs/transforms/transforms.py

@@ -2682,7 +2682,6 @@ def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec
 
         episode_specs = {}
         if isinstance(reward_spec, CompositeSpec):
-
             # If reward_spec is a CompositeSpec, all in_keys should be keys of reward_spec
             if not all(k in reward_spec.keys() for k in self.in_keys):
                 raise KeyError("Not all in_keys are present in ´reward_spec´")
@@ -2697,7 +2696,6 @@ def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec
                 episode_specs.update({out_key: episode_spec})
 
         else:
-
             # If reward_spec is not a CompositeSpec, the only in_key should be ´reward´
             if not set(self.in_keys) == {"reward"}:
                 raise KeyError(
@@ -2882,3 +2880,106 @@ def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec
                 if key in self.selected_keys
             }
         )
+
+
+class TimeMaxPool(Transform):
+    """Take the maximum value in each position over the last T observations.
+
+    This transform take the maximum value in each position for all in_keys tensors over the last T time steps.
+
+    Args:
+        in_keys (sequence of str, optional): input keys on which the max pool will be applied. Defaults to "observation" if left empty.
+        out_keys (sequence of str, optional): output keys where the output will be written. Defaults to `in_keys` if left empty.
+        T (int, optional): Number of time steps over which to apply max pooling.
+    """
+
+    inplace = False
+    invertible = False
+
+    def __init__(
+        self,
+        in_keys: Optional[Sequence[str]] = None,
+        out_keys: Optional[Sequence[str]] = None,
+        T: int = 1,
+    ):
+        if in_keys is None:
+            in_keys = ["observation"]
+        super().__init__(in_keys=in_keys, out_keys=out_keys)
+        if T < 1:
+            raise ValueError(
+                "TimeMaxPoolTranform T parameter should have a value greater or equal to one."
+            )
+        if len(self.in_keys) != len(self.out_keys):
+            raise ValueError(
+                "TimeMaxPoolTranform in_keys and out_keys don't have the same number of elements"
+            )
+        self.buffer_size = T
+        for in_key in self.in_keys:
+            buffer_name = f"_maxpool_buffer_{in_key}"
+            setattr(
+                self,
+                buffer_name,
+                torch.nn.parameter.UninitializedBuffer(
+                    device=torch.device("cpu"), dtype=torch.get_default_dtype()
+                ),
+            )
+
+    def reset(self, tensordict: TensorDictBase) -> TensorDictBase:
+        """Resets _buffers."""
+        # Non-batched environments
+        if len(tensordict.batch_size) < 1 or tensordict.batch_size[0] == 1:
+            for in_key in self.in_keys:
+                buffer_name = f"_maxpool_buffer_{in_key}"
+                buffer = getattr(self, buffer_name)
+                if isinstance(buffer, torch.nn.parameter.UninitializedBuffer):
+                    continue
+                buffer.fill_(0.0)
+
+        # Batched environments
+        else:
+            _reset = tensordict.get(
+                "_reset",
+                torch.ones(
+                    tensordict.batch_size,
+                    dtype=torch.bool,
+                    device=tensordict.device,
+                ),
+            )
+            for in_key in self.in_keys:
+                buffer_name = f"_maxpool_buffer_{in_key}"
+                buffer = getattr(self, buffer_name)
+                if isinstance(buffer, torch.nn.parameter.UninitializedBuffer):
+                    continue
+                buffer[:, _reset] = 0.0
+
+        return tensordict
+
+    def _make_missing_buffer(self, data, buffer_name):
+        buffer = getattr(self, buffer_name)
+        buffer.materialize((self.buffer_size,) + data.shape)
+        buffer = buffer.to(data.dtype).to(data.device).zero_()
+        setattr(self, buffer_name, buffer)
+
+    def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
+        """Update the episode tensordict with max pooled keys."""
+        for in_key, out_key in zip(self.in_keys, self.out_keys):
+            # Lazy init of buffers
+            buffer_name = f"_maxpool_buffer_{in_key}"
+            buffer = getattr(self, buffer_name)
+            if isinstance(buffer, torch.nn.parameter.UninitializedBuffer):
+                data = tensordict[in_key]
+                self._make_missing_buffer(data, buffer_name)
+            # shift obs 1 position to the right
+            buffer.copy_(torch.roll(buffer, shifts=1, dims=0))
+            # add new obs
+            buffer[0].copy_(tensordict[in_key])
+            # apply max pooling
+            pooled_tensor, _ = buffer.max(dim=0)
+            # add to tensordict
+            tensordict.set(out_key, pooled_tensor)
+
+        return tensordict
+
+    @_apply_to_composite
+    def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
+        return observation_spec