[Feature] MCTSForest

ghstack-source-id: efee967055c8dab87bbef442451ce20f21730b13 Pull Request resolved: #2307
pytorch · Jul 22, 2024 · dda9371 · dda9371
1 parent 8bbb86f
commit dda9371
Show file tree

Hide file tree

Showing 4 changed files with 453 additions and 0 deletions.
diff --git a/test/test_storage_map.py b/test/test_storage_map.py
@@ -237,6 +237,8 @@ def test_map_rollout(self):
         assert contains[: rollout.shape[-1]].all()
         assert not contains[rollout.shape[-1] :].any()
 
+class TestMCTSForest:
+    ...
 
 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