[Feature] MCTSForest

ghstack-source-id: 761863077202685046b093070361970c88081e36 Pull Request resolved: #2307
pytorch · Jul 23, 2024 · d1bc98a · d1bc98a
1 parent b2536e3
commit d1bc98a
Show file tree

Hide file tree

Showing 4 changed files with 457 additions and 0 deletions.
diff --git a/test/test_storage_map.py b/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()

diff --git a/torchrl/data/__init__.py b/torchrl/data/__init__.py
@@ -6,6 +6,9 @@
 from .map import (
     BinaryToDecimal,
     HashToInt,
+    MCTSChildren,
+    MCTSForest,
+    MCTSNode,
     QueryModule,
     RandomProjectionHash,
     SipHash,

diff --git a/torchrl/data/map/__init__.py b/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