[Feature] TensorDictMap Query module

ghstack-source-id: 9d70a3ea62402ee0619822c87dc9c05d8219101b
Pull Request resolved: #2305

test/test_storage_map.py

@@ -9,7 +9,8 @@
 
 import torch
 
-from torchrl.data.map import BinaryToDecimal, RandomProjectionHash, SipHash
+from tensordict import TensorDict
+from torchrl.data.map import BinaryToDecimal, QueryModule, RandomProjectionHash, SipHash
 
 _has_gym = importlib.util.find_spec("gymnasium", None) or importlib.util.find_spec(
     "gym", None
@@ -51,6 +52,69 @@ def test_randomprojection_hash(self, n_components, scale):
         assert y.unique().numel() == y.numel()
 
 
+class TestQuery:
+    def test_query_construct(self):
+        query_module = QueryModule(
+            in_keys=[(("key1",),), (("another",), "key2")],
+            index_key=("some", ("_index",)),
+            hash_module=SipHash(),
+            clone=False,
+        )
+        assert not query_module.clone
+        assert query_module.in_keys == ["key1", ("another", "key2")]
+        assert query_module.index_key == ("some", "_index")
+        assert isinstance(query_module.hash_module, dict)
+        assert isinstance(
+            query_module.aggregator,
+            type(query_module.hash_module[query_module.in_keys[0]]),
+        )
+        query_module = QueryModule(
+            in_keys=[(("key1",),), (("another",), "key2")],
+            index_key=("some", ("_index",)),
+            hash_module=SipHash(),
+            clone=False,
+            aggregator=SipHash(),
+        )
+        # assert not isinstance(query_module.aggregator is not query_module.hash_module[0]
+        assert isinstance(query_module.aggregator, SipHash)
+        query_module = QueryModule(
+            in_keys=[(("key1",),), (("another",), "key2")],
+            index_key=("some", ("_index",)),
+            hash_module=[SipHash(), SipHash()],
+            clone=False,
+        )
+        # assert query_module.aggregator is not query_module.hash_module[0]
+        assert isinstance(query_module.aggregator, SipHash)
+
+    @pytest.mark.parametrize("index_key", ["index", ("another", "index")])
+    @pytest.mark.parametrize("clone", [True, False])
+    def test_query(self, clone, index_key):
+        query_module = QueryModule(
+            in_keys=["key1", "key2"],
+            index_key=index_key,
+            hash_module=SipHash(),
+            clone=clone,
+        )
+
+        query = TensorDict(
+            {
+                "key1": torch.Tensor([[1], [1], [1], [2]]),
+                "key2": torch.Tensor([[3], [3], [2], [3]]),
+            },
+            batch_size=(4,),
+        )
+        res = query_module(query)
+        if clone:
+            assert res is not query
+        else:
+            assert res is query
+        assert index_key in res
+
+        assert res[index_key][0] == res[index_key][1]
+        for i in range(1, 3):
+            assert res[index_key][i].item() != res[index_key][i + 1].item()
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

torchrl/data/__init__.py

@@ -3,7 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from .map import BinaryToDecimal, RandomProjectionHash, SipHash
+from .map import BinaryToDecimal, HashToInt, QueryModule, RandomProjectionHash, SipHash
 from .postprocs import MultiStep
 from .replay_buffers import (
     Flat2TED,

torchrl/data/map/__init__.py

@@ -4,3 +4,4 @@
 # LICENSE file in the root directory of this source tree.
 
 from .hash import BinaryToDecimal, RandomProjectionHash, SipHash
+from .query import HashToInt, QueryModule

torchrl/data/map/query.py

@@ -0,0 +1,197 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import Any, Callable, Dict, List, Mapping, TypeVar
+
+import torch
+import torch.nn as nn
+from tensordict import NestedKey, TensorDictBase
+from tensordict.nn.common import TensorDictModuleBase
+from torchrl._utils import logger as torchrl_logger
+from torchrl.data.map import SipHash
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class HashToInt(nn.Module):
+    """Converts a hash value to an integer that can be used for indexing a contiguous storage."""
+
+    def __init__(self):
+        super().__init__()
+        self._index_to_index = {}
+
+    def __call__(self, key: torch.Tensor, extend: bool = False) -> torch.Tensor:
+        result = []
+        if extend:
+            for _item in key.tolist():
+                result.append(
+                    self._index_to_index.setdefault(_item, len(self._index_to_index))
+                )
+        else:
+            for _item in key.tolist():
+                result.append(
+                    self._index_to_index.get(_item, len(self._index_to_index))
+                )
+        return torch.tensor(result, device=key.device, dtype=key.dtype)
+
+    def state_dict(self) -> Dict[str, torch.Tensor]:
+        values = torch.tensor(self._index_to_index.values())
+        keys = torch.tensor(self._index_to_index.keys())
+        return {"keys": keys, "values": values}
+
+    def load_state_dict(
+        self, state_dict: Mapping[str, Any], strict: bool = True, assign: bool = False
+    ):
+        keys = state_dict["keys"]
+        values = state_dict["values"]
+        self._index_to_index = {
+            key: val for key, val in zip(keys.tolist(), values.tolist())
+        }
+
+
+class QueryModule(TensorDictModuleBase):
+    """A Module to generate compatible indices for storage.
+
+    A module that queries a storage and return required index of that storage.
+    Currently, it only outputs integer indices (torch.int64).
+
+    Args:
+        in_keys (list of NestedKeys): keys of the input tensordict that
+            will be used to generate the hash value.
+        index_key (NestedKey): the output key where the index value will be written.
+            Defaults to ``"_index"``.
+
+    Keyword Args:
+        hash_key (NestedKey): the output key where the hash value will be written.
+            Defaults to ``"_hash"``.
+        hash_module (Callable[[Any], int] or a list of these, optional): a hash
+            module similar to :class:`~tensordict.nn.SipHash` (default).
+            If a list of callables is provided, its length must equate the number of in_keys.
+        hash_to_int (Callable[[int], int], optional): a stateful function that
+            maps a hash value to a non-negative integer corresponding to an index in a
+            storage. Defaults to :class:`~torchrl.data.map.HashToInt`.
+        aggregator (Callable[[int], int], optional): a hash function to group multiple hashes
+            together. This argument should only be passed when there is more than one ``in_keys``.
+            If a single ``hash_module`` is provided but no aggregator is passed, it will take
+            the value of the hash_module. If no ``hash_module`` or a list of ``hash_modules`` is
+            provided but no aggregator is passed, it will default to ``SipHash``.
+       clone (bool, optional): if ``True``, a shallow clone of the input TensorDict will be
+            returned. This can be used to retrieve the integer index within the storage,
+            corresponding to a given input tensordict.
+            Defaults to ``False``.
+    d
+        Examples:
+            >>> query_module = QueryModule(
+            ...     in_keys=["key1", "key2"],
+            ...     index_key="index",
+            ...     hash_module=SipHash(),
+            ... )
+            >>> query = TensorDict(
+            ...     {
+            ...         "key1": torch.Tensor([[1], [1], [1], [2]]),
+            ...         "key2": torch.Tensor([[3], [3], [2], [3]]),
+            ...         "other": torch.randn(4),
+            ...     },
+            ...     batch_size=(4,),
+            ... )
+            >>> res = query_module(query)
+            >>> # The first two pairs of key1 and key2 match
+            >>> assert res["index"][0] == res["index"][1]
+            >>> # The last three pairs of key1 and key2 have at least one mismatching value
+            >>> assert res["index"][1] != res["index"][2]
+            >>> assert res["index"][2] != res["index"][3]
+    """
+
+    def __init__(
+        self,
+        in_keys: List[NestedKey],
+        index_key: NestedKey = "_index",
+        hash_key: NestedKey = "_hash",
+        *,
+        hash_module: Callable[[Any], int] | List[Callable[[Any], int]] | None = None,
+        hash_to_int: Callable[[int], int] | None = None,
+        aggregator: Callable[[Any], int] = None,
+        clone: bool = False,
+    ):
+        if len(in_keys) == 0:
+            raise ValueError("`in_keys` cannot be empty.")
+        in_keys = in_keys if isinstance(in_keys, List) else [in_keys]
+
+        super().__init__()
+        in_keys = self.in_keys = in_keys
+        self.out_keys = [index_key, hash_key]
+        index_key = self.out_keys[0]
+        self.hash_key = self.out_keys[1]
+
+        if aggregator is not None and len(self.in_keys) == 1:
+            torchrl_logger.warn(
+                "An aggregator was provided but there is only one in-key to be read. "
+                "This module will be ignored."
+            )
+        elif aggregator is None:
+            if hash_module is not None and not isinstance(hash_module, list):
+                aggregator = hash_module
+            else:
+                aggregator = SipHash()
+        if hash_module is None:
+            hash_module = [SipHash() for _ in range(len(self.in_keys))]
+        elif not isinstance(hash_module, list):
+            try:
+                hash_module = [
+                    deepcopy(hash_module) if len(self.in_keys) > 1 else hash_module
+                    for _ in range(len(self.in_keys))
+                ]
+            except Exception as err:
+                raise RuntimeError(
+                    "failed to deepcopy the hash module. Please provide a list of hash modules instead."
+                ) from err
+        elif len(hash_module) != len(self.in_keys):
+            raise ValueError(
+                "The number of hash_modules must match the number of in_keys. "
+                f"Got {len(hash_module)} hash modules but {len(in_keys)} in_keys."
+            )
+        if hash_to_int is None:
+            hash_to_int = HashToInt()
+
+        self.aggregator = aggregator
+        self.hash_module = dict(zip(self.in_keys, hash_module))
+        self.hash_to_int = hash_to_int
+
+        self.index_key = index_key
+        self.clone = clone
+
+    def forward(
+        self,
+        tensordict: TensorDictBase,
+        extend: bool = True,
+        write_hash: bool = True,
+    ) -> TensorDictBase:
+        hash_values = []
+
+        for k in self.in_keys:
+            hash_values.append(self.hash_module[k](tensordict.get(k)))
+        if len(self.in_keys) > 1:
+            hash_values = torch.stack(
+                hash_values,
+                dim=-1,
+            )
+            hash_values = self.aggregator(hash_values)
+        else:
+            hash_values = hash_values[0]
+
+        td_hash_value = self.hash_to_int(hash_values, extend=extend)
+
+        if self.clone:
+            output = tensordict.copy()
+        else:
+            output = tensordict
+
+        output.set(self.index_key, td_hash_value)
+        if write_hash:
+            output.set(self.hash_key, hash_values)
+        return output