treequest 0.1.2__tar.gz → 0.2.0__tar.gz

{treequest-0.1.2 → treequest-0.2.0}/.gitignore

@@ -8,3 +8,6 @@ wheels/
 
 # Virtual environments
 .venv
+
+# test images
+*.png

{treequest-0.1.2 → treequest-0.2.0}/PKG-INFO

@@ -1,11 +1,13 @@
 Metadata-Version: 2.4
 Name: treequest
-Version: 0.1.2
+Version: 0.2.0
 Summary: A Tree Search Library with Flexible API for LLM Inference-Time Scaling
 Author-email: Kou Misaki <kou.misaki@sakana.ai>, Takuya Akiba <takiba@sakana.ai>
 License-File: LICENSE
 Requires-Python: >=3.11
+Requires-Dist: joblib>=1.5.2
 Requires-Dist: scipy>=1.0.0
+Requires-Dist: ulid-py>=1.1.0
 Provides-Extra: abmcts-m
 Requires-Dist: jax>=0.5.2; extra == 'abmcts-m'
 Requires-Dist: numpy>=1.26.4; extra == 'abmcts-m'
@@ -62,6 +64,44 @@ best_state, best_node_score = tq.top_k(search_tree, algo, k=1)[0]
 print(f"Best state: {best_state}, Score: {best_node_score}")
 ```
 
+Alternatively, you can use an ask–tell interface with batched AB-MCTS sampling steps:
+```python
+import random
+import treequest as tq
+
+State = str
+
+def generate(parent_state: State | None) -> tuple[State, float]:
+    ...
+
+generate_fns = {"Action A": generate}
+actions = list(generate_fns.keys())
+
+# We use batch_size=5 here
+batch_size = 5
+
+# It runs AB-MCTS sampling step with 5 processes in parallel
+algo = tq.ABMCTSM(max_process_workers=batch_size)
+search_tree = algo.init_tree()
+
+total_budget = 50
+num_steps = total_budget // batch_size
+for _ in range(num_steps):
+    # ask_batch returns a list of `Trial` object, which has action, parent_state and trial_id attrs
+    search_tree, trials = algo.ask_batch(search_tree, batch_size, actions)
+
+    for trial in trials:
+        result = generate_fns[trial.action](trial.parent_state)
+        # Call tell method with trial_id to update search_tree
+        search_tree = algo.tell(search_tree, trial.trial_id, result)
+
+best_state, best_node_score = tq.top_k(search_tree, algo, k=1)[0]
+```
+
+In particular for AB-MCTS-M, each `step` call can be slow. If you encounter slow execution, prefer `ask_batch` over `step`.
+Please note that using a large `batch_size` can skew the search-tree shape (i.e., the tree may become too wide), so it is best to avoid overly large `batch_size`, see [PROFILING.md](./docs/PROFILING.md) for example trees.
+We recommend `batch_size<=5` as a starting point.
+
 ## Features
 - Easy-to-use API with customizable node generation and node scoring logic.
 - **AB-MCTS-A** and **AB-MCTS-M**, as well as **Multi-LLM AB-MCTS** support (See [our paper](https://arxiv.org/abs/2503.04412) for algorithm details).
@@ -155,6 +195,17 @@ for _ in range(20):
 ```
 The variation is not limited to LLM types; you can use different prompts, actions, scoring logic, etc. in `generate_fns`.
 
+### Batch Semantics and Concurrency
+- Algorithms are stateless objects; the evolving tree/search state is returned from `init_tree`, `step`, `ask`, and `tell`.
+- `ask_batch(state, batch_size, actions)` returns exactly `batch_size` Trial objects to expand next.
+  - Non-queue algorithms (e.g., `ABMCTSM`, `ABMCTSA`, `MultiArmedBanditUCB`) return exactly `batch_size` Trials.
+  - Queue-based algorithms (e.g., `StandardMCTS`, `BestFirstSearchAlgo`, `TreeOfThoughtsBFS`) precompute a set of parent/action pairs and duplicate them if needed to fill `batch_size`.
+- `tell(state, trial_id, (new_state, score))` reflects the result for the corresponding Trial.
+  - Order-independent: you can call `tell` in any order; reflection is tied to `trial_id`.
+  - Idempotent: calling `tell` twice on the same `trial_id` does not add extra nodes.
+  - For queue-based algorithms, over-told Trials beyond possible number of childs from a parent node (e.g., `(# actions)*samples_per_action` for StandardMCTS) become `INVALID` and are not reflected.
+- Scores are expected to be normalized to the `[0, 1]` range.
+
 ## Algorithms
 
 ### ABMCTS-A: ABMCTS with Node Aggregation

{treequest-0.1.2 → treequest-0.2.0}/README.md

@@ -44,6 +44,44 @@ best_state, best_node_score = tq.top_k(search_tree, algo, k=1)[0]
 print(f"Best state: {best_state}, Score: {best_node_score}")
 ```
 
+Alternatively, you can use an ask–tell interface with batched AB-MCTS sampling steps:
+```python
+import random
+import treequest as tq
+
+State = str
+
+def generate(parent_state: State | None) -> tuple[State, float]:
+    ...
+
+generate_fns = {"Action A": generate}
+actions = list(generate_fns.keys())
+
+# We use batch_size=5 here
+batch_size = 5
+
+# It runs AB-MCTS sampling step with 5 processes in parallel
+algo = tq.ABMCTSM(max_process_workers=batch_size)
+search_tree = algo.init_tree()
+
+total_budget = 50
+num_steps = total_budget // batch_size
+for _ in range(num_steps):
+    # ask_batch returns a list of `Trial` object, which has action, parent_state and trial_id attrs
+    search_tree, trials = algo.ask_batch(search_tree, batch_size, actions)
+
+    for trial in trials:
+        result = generate_fns[trial.action](trial.parent_state)
+        # Call tell method with trial_id to update search_tree
+        search_tree = algo.tell(search_tree, trial.trial_id, result)
+
+best_state, best_node_score = tq.top_k(search_tree, algo, k=1)[0]
+```
+
+In particular for AB-MCTS-M, each `step` call can be slow. If you encounter slow execution, prefer `ask_batch` over `step`.
+Please note that using a large `batch_size` can skew the search-tree shape (i.e., the tree may become too wide), so it is best to avoid overly large `batch_size`, see [PROFILING.md](./docs/PROFILING.md) for example trees.
+We recommend `batch_size<=5` as a starting point.
+
 ## Features
 - Easy-to-use API with customizable node generation and node scoring logic.
 - **AB-MCTS-A** and **AB-MCTS-M**, as well as **Multi-LLM AB-MCTS** support (See [our paper](https://arxiv.org/abs/2503.04412) for algorithm details).
@@ -137,6 +175,17 @@ for _ in range(20):
 ```
 The variation is not limited to LLM types; you can use different prompts, actions, scoring logic, etc. in `generate_fns`.
 
+### Batch Semantics and Concurrency
+- Algorithms are stateless objects; the evolving tree/search state is returned from `init_tree`, `step`, `ask`, and `tell`.
+- `ask_batch(state, batch_size, actions)` returns exactly `batch_size` Trial objects to expand next.
+  - Non-queue algorithms (e.g., `ABMCTSM`, `ABMCTSA`, `MultiArmedBanditUCB`) return exactly `batch_size` Trials.
+  - Queue-based algorithms (e.g., `StandardMCTS`, `BestFirstSearchAlgo`, `TreeOfThoughtsBFS`) precompute a set of parent/action pairs and duplicate them if needed to fill `batch_size`.
+- `tell(state, trial_id, (new_state, score))` reflects the result for the corresponding Trial.
+  - Order-independent: you can call `tell` in any order; reflection is tied to `trial_id`.
+  - Idempotent: calling `tell` twice on the same `trial_id` does not add extra nodes.
+  - For queue-based algorithms, over-told Trials beyond possible number of childs from a parent node (e.g., `(# actions)*samples_per_action` for StandardMCTS) become `INVALID` and are not reflected.
+- Scores are expected to be normalized to the `[0, 1]` range.
+
 ## Algorithms
 
 ### ABMCTS-A: ABMCTS with Node Aggregation

treequest-0.2.0/docs/PROFILING.md

@@ -0,0 +1,98 @@
+# Performance Investigation of `ask_batch` method
+Here, we report the profiling result of `ask_batch` with different `batch_size`. As we increase `batch_size`, AB-MCTS runs are accelerated, while the search tree shapes get more wide.
+
+I profiled TreeQuest on a 32-core x86-64 machine running Ubuntu 20.04. As for AB-MCTS-M, it runs 4 parallel MCMC chains per process, so I tested `batch_size` from 1 to 10.
+
+## Profiling Commands
+
+### AB-MCTS-M
+```bash
+hyperfine \
+  'uv run tests/profiling/ab_mcts_m.py -b 1' \
+  'uv run tests/profiling/ab_mcts_m.py -b 2' \
+  'uv run tests/profiling/ab_mcts_m.py -b 5' \
+  'uv run tests/profiling/ab_mcts_m.py -b 10' \
+  -w 0 -r 3 --export-markdown tests/profiling/benchmark_m.md
+```
+
+### AB-MCTS-A
+```bash
+hyperfine \
+  'uv run tests/profiling/ab_mcts_a.py -b 1' \
+  'uv run tests/profiling/ab_mcts_a.py -b 2' \
+  'uv run tests/profiling/ab_mcts_a.py -b 5' \
+  'uv run tests/profiling/ab_mcts_a.py -b 10' \
+  -w 0 -r 3 --export-markdown tests/profiling/benchmark_a.md
+```
+
+## Analysis for various `batch_size`
+### AB-MCTS-M
+#### Time Profiling
+For a search tree with (# non-root nodes) = 50,
+| Command | Mean [s] | Min [s] | Max [s] | Relative |
+|:---|---:|---:|---:|---:|
+| `uv run tests/profiling/ab_mcts_m.py -b 1` | 372.156 ± 37.808 | 342.916 | 414.852 | 6.98 ± 0.72 |
+| `uv run tests/profiling/ab_mcts_m.py -b 2` | 202.324 ± 8.006 | 194.178 | 210.182 | 3.80 ± 0.16 |
+| `uv run tests/profiling/ab_mcts_m.py -b 5` | 91.418 ± 2.391 | 90.026 | 94.179 | 1.71 ± 0.05 |
+| `uv run tests/profiling/ab_mcts_m.py -b 10` | 53.307 ± 0.747 | 52.468 | 53.902 | 1.00 |
+
+> ![NOTE]
+> For more resource-constrained systems, we get more mild boost. I ran profiling for M3 MacBook Air with 4 Pcores and 4 Ecores:
+> | Command | Mean [s] | Min [s] | Max [s] | Relative |
+> |:---|---:|---:|---:|---:|
+> | `uv run tests/profiling/ab_mcts_m.py -b 1` | 147.564 ± 4.527 | 144.297 | 152.731 | 1.90 ± 0.08 |
+> | `uv run tests/profiling/ab_mcts_m.py -b 2` | 116.557 ± 3.015 | 113.391 | 119.395 | 1.50 ± 0.06 |
+> | `uv run tests/profiling/ab_mcts_m.py -b 5` | 90.628 ± 1.760 | 89.299 | 92.624 | 1.17 ± 0.04 |
+> | `uv run tests/profiling/ab_mcts_m.py -b 10` | 77.773 ± 2.050 | 75.819 | 79.908 | 1.00 |
+
+
+We note that different batch size leads to different tree shape; The larger batch size generates wider search tree:
+
+#### `batch_size=1`
+![batch_size=1](../images/batch/abmcts_m_batch_size_1.jpg)
+
+#### `batch_size=2`
+![batch_size=2](../images/batch/abmcts_m_batch_size_2.jpg)
+
+#### `batch_size=5`
+![batch_size=5](../images/batch/abmcts_m_batch_size_5.jpg)
+
+#### `batch_size=10`
+![batch_size=10](../images/batch/abmcts_m_batch_size_10.jpg)
+
+## Analysis for various `batch_size`
+### AB-MCTS-A
+#### Time Profiling
+AB-MCTS-A is faster than AB-MCTS-M, so I tested (# non-root nodes) = 1000:
+| Command | Mean [s] | Min [s] | Max [s] | Relative |
+|:---|---:|---:|---:|---:|
+| `uv run tests/profiling/ab_mcts_a.py -b 1` | 26.116 ± 0.656 | 25.677 | 26.871 | 1.00 |
+| `uv run tests/profiling/ab_mcts_a.py -b 2` | 26.191 ± 0.709 | 25.599 | 26.976 | 1.00 ± 0.04 |
+| `uv run tests/profiling/ab_mcts_a.py -b 5` | 26.404 ± 0.441 | 26.061 | 26.902 | 1.01 ± 0.03 |
+| `uv run tests/profiling/ab_mcts_a.py -b 10` | 26.399 ± 0.364 | 25.998 | 26.709 | 1.01 ± 0.03 |
+
+
+> [!NOTE] 
+> I didn't parallelize AB-MCTS-A, so the command running time is the same for different batch sizes. 
+> Intially I also implemented parallel run for AB-MCTS-A as well, but it didn't lead to much speedup:
+> | Command | Mean [s] | Min [s] | Max [s] | Relative |
+> |:---|---:|---:|---:|---:|
+> | `uv run tests/profiling/ab_mcts_a.py -b 1` | 49.973 ± 2.909 | 47.645 | 53.235 | 2.79 ± 0.16 |
+> | `uv run tests/profiling/ab_mcts_a.py -b 2` | 30.971 ± 0.546 | 30.528 | 31.581 | 1.73 ± 0.03 |
+> | `uv run tests/profiling/ab_mcts_a.py -b 5` | 21.256 ± 0.706 | 20.732 | 22.058 | 1.18 ± 0.04 |
+> | `uv run tests/profiling/ab_mcts_a.py -b 10` | 17.941 ± 0.110 | 17.867 | 18.068 | 1.00 |
+
+
+We note that different batch size leads to different tree shape; The larger batch size generates wider search tree:
+
+#### `batch_size=1`
+![batch_size=1](../images/batch/abmcts_m_batch_size_1.jpg)
+
+#### `batch_size=2`
+![batch_size=2](../images/batch/abmcts_m_batch_size_2.jpg)
+
+#### `batch_size=5`
+![batch_size=5](../images/batch/abmcts_m_batch_size_5.jpg)
+
+#### `batch_size=10`
+![batch_size=10](../images/batch/abmcts_m_batch_size_10.jpg)

{treequest-0.1.2 → treequest-0.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "treequest"
-version = "0.1.2"
+version = "0.2.0"
 description = "A Tree Search Library with Flexible API for LLM Inference-Time Scaling"
 readme = "README.md"
 authors = [
@@ -9,7 +9,9 @@ authors = [
 ]
 requires-python = ">=3.11"
 dependencies = [
+    "joblib>=1.5.2",
     "scipy>=1.0.0",
+    "ulid-py>=1.1.0",
 ]
 
 [project.optional-dependencies]

{treequest-0.1.2 → treequest-0.2.0}/src/treequest/__init__.py

@@ -1,8 +1,6 @@
 from .algos.ab_mcts_m._ab_mcts_m_imports import _import as _ab_mcts_m_import
 
-if _ab_mcts_m_import.is_successful():
-    from .algos.ab_mcts_m.algo import ABMCTSM
-else:
+if not _ab_mcts_m_import.is_successful():
     # Create a placeholder that raises an informative error when accessed
     class _ABMCTSMPlaceholder:
         def __getattr__(self, name):  # type: ignore
@@ -14,6 +12,8 @@ else:
             raise ImportError("ABMCTSM import failed.")
 
     ABMCTSM = _ABMCTSMPlaceholder()  # type: ignore
+else:
+    from .algos.ab_mcts_m.algo import ABMCTSM  # type: ignore[assignment]
 
 from .algos.ab_mcts_a.algo import ABMCTSA
 from .algos.base import Algorithm

{treequest-0.1.2 → treequest-0.2.0}/src/treequest/algos/ab_mcts_a/algo.py

@@ -8,6 +8,7 @@ from typing import Dict, Generic, List, Literal, Optional, Tuple, TypeVar, Union
 from treequest.algos.ab_mcts_a.prob_state import NodeProbState, PriorConfig
 from treequest.algos.base import Algorithm
 from treequest.algos.tree import Node, Tree
+from treequest.trial import Trial, TrialId, TrialStore
 from treequest.types import GenerateFnType, StateScoreType
 
 # Type variable for state
@@ -20,7 +21,7 @@ logger = getLogger(__name__)
 class ABMCTSAStateManager(Generic[StateT]):
     """Manager for ABMCTSAState instances associated with expand_idx values."""
 
-    states: Dict[int, NodeProbState] = field(default_factory=dict)
+    states: Dict[int, NodeProbState] = field(default_factory=dict[int, NodeProbState])
     default_prior_config: Optional["PriorConfig"] = None
     default_reward_average_priors: Optional[Union[float, Dict[str, float]]] = None
     default_model_selection_strategy: str = "multiarm_bandit_thompson"
@@ -32,13 +33,13 @@ class ABMCTSAStateManager(Generic[StateT]):
     def __len__(self) -> int:
         return len(self.states)
 
-    def get(self, node: Node) -> Optional[NodeProbState]:
+    def get(self, node: Node[StateT]) -> Optional[NodeProbState]:
         """Get thompson state for the given expand_idx if it exists."""
         return self.states.get(node.expand_idx)
 
     def create(
         self,
-        node: Node,
+        node: Node[StateT],
         actions: List[str],
         prior_config: Optional["PriorConfig"] = None,
         reward_average_priors: Optional[Union[float, Dict[str, float]]] = None,
@@ -65,7 +66,7 @@ class ABMCTSAStateManager(Generic[StateT]):
 
     def get_or_create(
         self,
-        node: Node,
+        node: Node[StateT],
         actions: List[str],
         prior_config: Optional["PriorConfig"] = None,
         reward_average_priors: Optional[Union[float, Dict[str, float]]] = None,
@@ -96,10 +97,13 @@ class ABMCTSAAlgoState(Generic[StateT]):
     """State for ABMCTSA algorithm."""
 
     tree: Tree[StateT]
-    thompson_states: ABMCTSAStateManager = field(default_factory=ABMCTSAStateManager)
+    thompson_states: ABMCTSAStateManager[StateT] = field(
+        default_factory=ABMCTSAStateManager[StateT]
+    )
     all_rewards_store: Dict[str, List[float]] = field(
         default_factory=build_default_dict_of_list
     )
+    trial_store: TrialStore[StateT] = field(default_factory=TrialStore[StateT])
 
 
 class ABMCTSA(Algorithm[StateT, ABMCTSAAlgoState[StateT]]):
@@ -155,14 +159,14 @@ class ABMCTSA(Algorithm[StateT, ABMCTSAAlgoState[StateT]]):
             )
         self.model_selection_strategy = model_selection_strategy
 
-    def init_tree(self) -> ABMCTSAAlgoState:
+    def init_tree(self) -> ABMCTSAAlgoState[StateT]:
         """
         Initialize the algorithm state with an empty tree.
 
         Returns:
             Initial algorithm state
         """
-        tree: Tree = Tree.with_root_node()
+        tree: Tree[StateT] = Tree.with_root_node()
         state_manager = ABMCTSAStateManager[StateT](
             default_prior_config=self.prior_config,
             default_reward_average_priors=self.reward_average_priors,
@@ -172,55 +176,56 @@ class ABMCTSA(Algorithm[StateT, ABMCTSAAlgoState[StateT]]):
 
     def step(
         self,
-        state: ABMCTSAAlgoState,
+        state: ABMCTSAAlgoState[StateT],
         generate_fn: Mapping[str, GenerateFnType[StateT]],
         inplace: bool = False,
-    ) -> ABMCTSAAlgoState:
+    ) -> ABMCTSAAlgoState[StateT]:
         """
-        Perform one step of the Thompson Sampling MCTS algorithm.
-
-        Args:
-            state: Current algorithm state
-            generate_fn: Mapping of action names to generation functions
-
-        Returns:
-            Updated algorithm state
+        Perform one step of the AB-MCTS-A algorithm and generate a new node.
         """
         if not inplace:
             state = copy.deepcopy(state)
 
-        # initialize all_rewards_store
-        if len(state.all_rewards_store) == 0:
-            for action in generate_fn:
-                state.all_rewards_store[action] = []
+        actions = list(generate_fn.keys())
+        state, trial = self.ask(state, actions)
 
+        action = trial.action
+        node = state.tree.get_node(trial.node_to_expand)
+        result = generate_fn[action](node.state)
+
+        self.tell(state, trial.trial_id, result)
+        return state
+
+    def _get_expand_node_and_action(
+        self,
+        state: ABMCTSAAlgoState[StateT],
+        actions: list[str],
+    ) -> tuple[Node[StateT], str]:
         # If the tree is empty (only root), expand the root
         if not state.tree.root.children:
-            self._expand_node(state, state.tree.root, generate_fn)
-            return state
+            return state.tree.root, self._get_generation_action(
+                state, state.tree.root, actions
+            )
 
         # Run one simulation step
         node = state.tree.root
 
         # Selection phase: traverse tree until we reach a leaf node or need to create a new node
         while node.children:
-            node, action_used = self._select_child(state, node, generate_fn)
-
-            # If action is not None, it means we've generated a new node
-            if action_used is not None:
-                return state
+            node, action = self._select_child(state, node, actions)
 
-        # Expansion phase: expand leaf node
-        self._expand_node(state, node, generate_fn)
-
-        return state
+            # If action is not None, we will generate a new node from `node``
+            if action is not None:
+                return node, action
+        action = self._get_generation_action(state, node, actions)
+        return node, action
 
     def _select_child(
         self,
-        state: ABMCTSAAlgoState,
-        node: Node,
-        generate_fn: Mapping[str, GenerateFnType[StateT]],
-    ) -> Tuple[Node, Optional[str]]:
+        state: ABMCTSAAlgoState[StateT],
+        node: Node[StateT],
+        actions: list[str],
+    ) -> Tuple[Node[StateT], Optional[str]]:
         """
         Select a child node using Thompson Sampling.
 
@@ -235,7 +240,7 @@ class ABMCTSA(Algorithm[StateT, ABMCTSAAlgoState[StateT]]):
         # Get or create thompson state for this node
         thompson_state = state.thompson_states.get_or_create(
             node,
-            list(generate_fn.keys()),
+            actions,
         )
 
         # Ask for next node or action using Thompson Sampling
@@ -243,8 +248,7 @@ class ABMCTSA(Algorithm[StateT, ABMCTSAAlgoState[StateT]]):
 
         # If string returned, we need to generate a new node with that action
         if isinstance(selection, str):
-            new_node = self._generate_new_child(state, node, generate_fn, selection)
-            return new_node, selection
+            return node, selection
         else:
             # Otherwise, we return the existing child with that index
             if selection >= len(node.children):
@@ -253,27 +257,13 @@ class ABMCTSA(Algorithm[StateT, ABMCTSAAlgoState[StateT]]):
                 )
             return node.children[selection], None
 
-    def _expand_node(
-        self,
-        state: ABMCTSAAlgoState,
-        node: Node,
-        generate_fn: Mapping[str, GenerateFnType[StateT]],
-    ) -> Tuple[Node, str]:
-        """
-        Expand a leaf node by generating a new child.
-
-        Args:
-            state: Current algorithm state
-            node: Node to expand
-            generate_fn: Mapping of action names to generation functions
-
-        Returns:
-            Tuple of (new node, action used)
-        """
+    def _get_generation_action(
+        self, state: ABMCTSAAlgoState[StateT], node: Node[StateT], actions: list[str]
+    ) -> str:
         # Create thompson state for this node if it doesn't exist
         thompson_state = state.thompson_states.get_or_create(
             node,
-            list(generate_fn.keys()),
+            actions,
         )
 
         # Get action to use for generating child
@@ -285,53 +275,14 @@ class ABMCTSA(Algorithm[StateT, ABMCTSAAlgoState[StateT]]):
                 f"Something went wrong in ABMCTSA algorithm: selection should always be str when the expansion is from the leaf node, whle got {selection}"
             )
 
-        new_node = self._generate_new_child(state, node, generate_fn, selection)
-        return new_node, selection
+        return selection
 
-    def _generate_new_child(
+    def _backpropagate(
         self,
-        state: ABMCTSAAlgoState,
-        node: Node,
-        generate_fn: Mapping[str, GenerateFnType[StateT]],
+        state: ABMCTSAAlgoState[StateT],
+        node: Node[StateT],
+        score: float,
         action: str,
-    ) -> Node:
-        """
-        Generate a new child node using the specified action.
-
-        Args:
-            state: Current algorithm state
-            node: Parent node
-            generate_fn: Mapping of action names to generation functions
-            action: Name of action to use for generation
-
-        Returns:
-            Newly created node
-        """
-        # Generate new state and score using the selected action
-        node_state = None if node.is_root() else node.state
-        new_state, new_score = generate_fn[action](node_state)
-
-        # Add new node to the tree
-        new_node = state.tree.add_node((new_state, new_score), node)
-
-        # Update Thompson state with the new node
-        thompson_state = state.thompson_states.get(node)
-        if thompson_state:
-            thompson_state.register_new_child_node(
-                action, new_node, self.model_selection_strategy
-            )
-        else:
-            raise RuntimeError(
-                f"Internal Error in ABMCTSA: thompson_state should not be None for node {node}"
-            )
-
-        # Backpropagate the score through the parents
-        self._backpropagate(state, new_node, new_score, action)
-
-        return new_node
-
-    def _backpropagate(
-        self, state: ABMCTSAAlgoState, node: Node, score: float, action: str
     ) -> None:
         """
         Update Thompson Sampling statistics for all nodes in the path from node to root.
@@ -369,7 +320,7 @@ class ABMCTSA(Algorithm[StateT, ABMCTSAAlgoState[StateT]]):
             current = current.parent
 
     def get_state_score_pairs(
-        self, state: ABMCTSAAlgoState
+        self, state: ABMCTSAAlgoState[StateT]
     ) -> List[StateScoreType[StateT]]:
         """
         Get all the state-score pairs from the tree.
@@ -381,3 +332,57 @@ class ABMCTSA(Algorithm[StateT, ABMCTSAAlgoState[StateT]]):
             List of (state, score) pairs
         """
         return state.tree.get_state_score_pairs()
+
+    def ask_batch(
+        self, state: ABMCTSAAlgoState[StateT], batch_size: int, actions: list[str]
+    ) -> tuple[ABMCTSAAlgoState[StateT], list[Trial[StateT]]]:
+        """
+        ABMCTSA is lightweight, so we don't parallelize it.
+        TODO: If we need to optimize it, ProcessPoolExecutor seems suffice
+        """
+        # initialize all_rewards_store
+        if len(state.all_rewards_store) == 0:
+            for a in actions:
+                state.all_rewards_store[a] = []
+
+        trials: list[Trial[StateT]] = []
+        for _ in range(batch_size):
+            node, action = self._get_expand_node_and_action(state, actions)
+            trials.append(state.trial_store.create_trial(node, action))
+
+        return state, trials
+
+    def tell(
+        self,
+        state: ABMCTSAAlgoState[StateT],
+        trial_id: TrialId,
+        result: tuple[StateT, float],
+    ) -> ABMCTSAAlgoState[StateT]:
+        _new_state, new_score = result
+
+        finished_trial = state.trial_store.get_finished_trial(trial_id, new_score)
+        if (
+            finished_trial is None
+        ):  # Trial is no longer valid, so we do not reflect the result to state
+            return state
+
+        parent_node = state.tree.get_node(finished_trial.node_to_expand)
+        action = finished_trial.action
+        # Add new node to the tree
+        new_node = state.tree.add_node(result, parent_node)
+
+        # Update Thompson state with the new node
+        thompson_state = state.thompson_states.get(parent_node)
+        if thompson_state:
+            thompson_state.register_new_child_node(
+                action, new_node, self.model_selection_strategy
+            )
+        else:
+            raise RuntimeError(
+                f"Internal Error in ABMCTSA: thompson_state should not be None for node {parent_node}"
+            )
+
+        # Backpropagate the score through the parents
+        self._backpropagate(state, new_node, new_score, action)
+
+        return state