[Feature] MCTSForest

ghstack-source-id: dd7d939
Pull Request resolved: #2307

Author: Vincent Moens

test/test_storage_map.py

@@ -237,6 +237,12 @@ def test_map_rollout(self):
         assert contains[: rollout.shape[-1]].all()
         assert not contains[rollout.shape[-1] :].any()
 
+class TestMCTSForest:
+    def test_forest_build(self):
+        ...
+    def test_forest_extend_and_get(self):
+        ...
+
 
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()

torchrl/data/__init__.py

@@ -6,6 +6,9 @@
 from .map import (
     BinaryToDecimal,
     HashToInt,
+    MCTSChildren,
+    MCTSForest,
+    MCTSNode,
     QueryModule,
     RandomProjectionHash,
     SipHash,

torchrl/data/map/__init__.py

@@ -6,3 +6,4 @@
 from .hash import BinaryToDecimal, RandomProjectionHash, SipHash
 from .query import HashToInt, QueryModule
 from .tdstorage import TensorDictMap, TensorMap
+from .tree import MCTSChildren, MCTSForest, MCTSNode