gymcts 1.0.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

gymcts/colorful_console_utils.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -103,8 +105,7 @@ def wrap_with_color_codes(s: object, /, r: int | float, g: int | float, b: int |
            f"{CEND}"
 
 
-
-def wrap_evenly_spaced_color(s: str, n_of_item:int, n_classes:int, c_map="rainbow") -> str:
+def wrap_evenly_spaced_color(s: Any, n_of_item: int, n_classes: int, c_map="rainbow") -> str:
     if s is None or n_of_item is None or n_classes is None:
         return s
 
@@ -117,7 +118,7 @@ def wrap_evenly_spaced_color(s: str, n_of_item:int, n_classes:int, c_map="rainbo
     return f"{color_asni}{s}{CEND}"
 
 
-def wrap_with_color_scale(s: str, value: float, min_val:float, max_val:float, c_map=None) -> str:
+def wrap_with_color_scale(s: str, value: float, min_val: float, max_val: float, c_map=None) -> str:
     if s is None or min_val is None or max_val is None or min_val >= max_val:
         return s
 

gymcts/{gymcts_deterministic_wrapper.py → gymcts_action_history_wrapper.py}

@@ -7,12 +7,12 @@ import gymnasium as gym
 from gymnasium.core import WrapperActType, WrapperObsType
 from gymnasium.wrappers import RecordEpisodeStatistics
 
-from gymcts.gymcts_gym_env import SoloMCTSGymEnv
+from gymcts.gymcts_env_abc import GymctsABC
 
 from gymcts.logger import log
 
 
-class DeterministicSoloMCTSGymEnvWrapper(SoloMCTSGymEnv, gym.Wrapper):
+class ActionHistoryMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
     _terminal_flag: bool = False
     _last_reward: SupportsFloat = 0
     _step_tuple: tuple[WrapperObsType, SupportsFloat, bool, bool, dict[str, Any]] = None

gymcts/gymcts_agent.py

@@ -3,27 +3,28 @@ import gymnasium as gym
 
 from typing import TypeVar, Any, SupportsFloat, Callable
 
-from gymcts.gymcts_gym_env import SoloMCTSGymEnv
-from gymcts.gymcts_naive_wrapper import NaiveSoloMCTSGymEnvWrapper
-from gymcts.gymcts_node import SoloMCTSNode
+from gymcts.gymcts_env_abc import GymctsABC
+from gymcts.gymcts_deepcopy_wrapper import DeepCopyMCTSGymEnvWrapper
+from gymcts.gymcts_node import GymctsNode
+from gymcts.gymcts_tree_plotter import _generate_mcts_tree
 
 from gymcts.logger import log
 
 TSoloMCTSNode = TypeVar("TSoloMCTSNode", bound="SoloMCTSNode")
 
 
-class SoloMCTSAgent:
+class GymctsAgent:
     render_tree_after_step: bool = False
     render_tree_max_depth: int = 2
     exclude_unvisited_nodes_from_render: bool = False
     number_of_simulations_per_step: int = 25
 
-    env: SoloMCTSGymEnv
-    search_root_node: SoloMCTSNode  # NOTE: this is not the same as the root of the tree!
+    env: GymctsABC
+    search_root_node: GymctsNode  # NOTE: this is not the same as the root of the tree!
     clear_mcts_tree_after_step: bool
 
     def __init__(self,
-                 env: SoloMCTSGymEnv,
+                 env: GymctsABC,
                  clear_mcts_tree_after_step: bool = True,
                  render_tree_after_step: bool = False,
                  render_tree_max_depth: int = 2,
@@ -43,13 +44,13 @@ class SoloMCTSAgent:
         self.env = env
         self.clear_mcts_tree_after_step = clear_mcts_tree_after_step
 
-        self.search_root_node = SoloMCTSNode(
+        self.search_root_node = GymctsNode(
             action=None,
             parent=None,
             env_reference=env,
         )
 
-    def navigate_to_leaf(self, from_node: SoloMCTSNode) -> SoloMCTSNode:
+    def navigate_to_leaf(self, from_node: GymctsNode) -> GymctsNode:
         log.debug(f"Navigate to leaf. from_node: {from_node}")
         if from_node.terminal:
             log.debug("Node is terminal. Returning from_node")
@@ -66,7 +67,7 @@ class SoloMCTSAgent:
         log.debug(f"Selected leaf node: {temp_node}")
         return temp_node
 
-    def expand_node(self, node: SoloMCTSNode) -> None:
+    def expand_node(self, node: GymctsNode) -> None:
         log.debug(f"expanding node: {node}")
         # EXPANSION STRATEGY
         # expand all children
@@ -78,7 +79,7 @@ class SoloMCTSAgent:
             self._load_state(node)
 
             obs, reward, terminal, truncated, _ = self.env.step(action)
-            child_dict[action] = SoloMCTSNode(
+            child_dict[action] = GymctsNode(
                 action=action,
                 parent=node,
                 env_reference=self.env,
@@ -110,14 +111,14 @@ class SoloMCTSAgent:
         # restore state of current node
         return action_list
 
-    def _load_state(self, node: SoloMCTSNode) -> None:
-        if isinstance(self.env, NaiveSoloMCTSGymEnvWrapper):
+    def _load_state(self, node: GymctsNode) -> None:
+        if isinstance(self.env, DeepCopyMCTSGymEnvWrapper):
             self.env = copy.deepcopy(node.state)
         else:
             self.env.load_state(node.state)
 
-    def perform_mcts_step(self, search_start_node: SoloMCTSNode = None, num_simulations: int = None,
-                          render_tree_after_step: bool = None) -> tuple[int, SoloMCTSNode]:
+    def perform_mcts_step(self, search_start_node: GymctsNode = None, num_simulations: int = None,
+                          render_tree_after_step: bool = None) -> tuple[int, GymctsNode]:
 
         if render_tree_after_step is None:
             render_tree_after_step = self.render_tree_after_step
@@ -149,7 +150,7 @@ class SoloMCTSAgent:
 
         return action, next_node
 
-    def vanilla_mcts_search(self, search_start_node: SoloMCTSNode = None, num_simulations=10) -> int:
+    def vanilla_mcts_search(self, search_start_node: GymctsNode = None, num_simulations=10) -> int:
         log.debug(f"performing one MCTS search step with {num_simulations} simulations")
         if search_start_node is None:
             search_start_node = self.search_root_node
@@ -178,7 +179,7 @@ class SoloMCTSAgent:
 
         return search_start_node.get_best_action()
 
-    def show_mcts_tree(self, start_node: SoloMCTSNode = None, tree_max_depth: int = None) -> None:
+    def show_mcts_tree(self, start_node: GymctsNode = None, tree_max_depth: int = None) -> None:
 
         if start_node is None:
             start_node = self.search_root_node
@@ -187,13 +188,17 @@ class SoloMCTSAgent:
             tree_max_depth = self.render_tree_max_depth
 
         print(start_node.__str__(colored=True, action_space_n=self.env.action_space.n))
-        for line in self._generate_mcts_tree(start_node=start_node, depth=tree_max_depth):
+        for line in _generate_mcts_tree(
+                start_node=start_node,
+                depth=tree_max_depth,
+                action_space_n=self.env.action_space.n,
+        ):
             print(line)
 
     def show_mcts_tree_from_root(self, tree_max_depth: int = None) -> None:
         self.show_mcts_tree(start_node=self.search_root_node.get_root(), tree_max_depth=tree_max_depth)
 
-    def backpropagation(self, node: SoloMCTSNode, episode_return: float) -> None:
+    def backpropagation(self, node: GymctsNode, episode_return: float) -> None:
         log.debug(f"performing backpropagation from leaf node: {node}")
         while not node.is_root():
             # node.mean_value = ((node.mean_value * node.visit_count) + episode_return) / (node.visit_count + 1)
@@ -209,53 +214,4 @@ class SoloMCTSAgent:
         node.max_value = max(node.max_value, episode_return)
         node.min_value = min(node.min_value, episode_return)
 
-    def _generate_mcts_tree(self, start_node: SoloMCTSNode = None, prefix: str = None, depth: int = None) -> list[str]:
 
-        if prefix is None:
-            prefix = ""
-        import gymcts.colorful_console_utils as ccu
-
-        if start_node is None:
-            start_node = self.search_root_node
-
-        # prefix components:
-        space = '    '
-        branch = '│   '
-        # pointers:
-        tee = '├── '
-        last = '└── '
-
-        contents = start_node.children.values() if start_node.children is not None else []
-        if self.exclude_unvisited_nodes_from_render:
-            contents = [node for node in contents if node.visit_count > 0]
-        # contents each get pointers that are ├── with a final └── :
-        # pointers = [tee] * (len(contents) - 1) + [last]
-        pointers = [tee for _ in range(len(contents) - 1)] + [last]
-
-        for pointer, current_node in zip(pointers, contents):
-            n_item = current_node.parent.action if current_node.parent is not None else 0
-            n_classes = self.env.action_space.n
-
-            pointer = ccu.wrap_evenly_spaced_color(
-                s=pointer,
-                n_of_item=n_item,
-                n_classes=n_classes,
-            )
-
-            yield prefix + pointer + f"{current_node.__str__(colored=True, action_space_n=n_classes)}"
-            if current_node.children and len(current_node.children):  # extend the prefix and recurse:
-                # extension = branch if pointer == tee else space
-                extension = branch if tee in pointer else space
-                # i.e. space because last, └── , above so no more |
-                extension = ccu.wrap_evenly_spaced_color(
-                    s=extension,
-                    n_of_item=n_item,
-                    n_classes=n_classes,
-                )
-                if depth is not None and depth <= 0:
-                    continue
-                yield from self._generate_mcts_tree(
-                    current_node,
-                    prefix=prefix + extension,
-                    depth=depth - 1 if depth is not None else None
-                )

gymcts/{gymcts_naive_wrapper.py → gymcts_deepcopy_wrapper.py}

@@ -7,12 +7,12 @@ import gymnasium as gym
 from gymnasium.core import WrapperActType, WrapperObsType
 from gymnasium.wrappers import RecordEpisodeStatistics
 
-from gymcts.gymcts_gym_env import SoloMCTSGymEnv
+from gymcts.gymcts_env_abc import GymctsABC
 
 from gymcts.logger import log
 
 
-class NaiveSoloMCTSGymEnvWrapper(SoloMCTSGymEnv, gym.Wrapper):
+class DeepCopyMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
 
 
     _terminal_flag:bool = False

gymcts/gymcts_distributed_agent.py

@@ -0,0 +1,281 @@
+import copy
+import gymnasium as gym
+
+from typing import TypeVar, Any, SupportsFloat, Callable
+
+from ray.types import ObjectRef
+
+from gymcts.gymcts_agent import GymctsAgent
+from gymcts.gymcts_env_abc import GymctsABC
+from gymcts.gymcts_deepcopy_wrapper import DeepCopyMCTSGymEnvWrapper
+from gymcts.gymcts_node import GymctsNode
+from gymcts.gymcts_tree_plotter import _generate_mcts_tree
+
+from gymcts.logger import log
+
+import ray
+import copy
+
+TSoloMCTSNode = TypeVar("TSoloMCTSNode", bound="SoloMCTSNode")
+
+
+@ray.remote
+def mcts_lookahead(
+        gymcts_start_node: GymctsNode,
+        env: GymctsABC,
+        num_simulations: int) -> GymctsNode:
+    agent = GymctsAgent(
+        env=env,
+        clear_mcts_tree_after_step=False,
+        number_of_simulations_per_step=num_simulations,
+    )
+    agent.search_root_node = gymcts_start_node
+
+    agent.vanilla_mcts_search(
+        search_start_node=gymcts_start_node,
+        num_simulations=num_simulations,
+    )
+    return agent.search_root_node
+
+
+def merge_nodes(gymcts_node1, gymcts_node2, perform_state_equality_check=False):
+    log.debug(f"merging {gymcts_node1} and {gymcts_node2}")
+    # maybe add some state equality check here
+    if perform_state_equality_check:
+        if gymcts_node1.state != gymcts_node2.state:
+            raise ValueError("States are different")
+
+    if gymcts_node1 is None:
+        log.debug(f"first node is None, returning second node ({gymcts_node2})")
+        return gymcts_node2
+    if gymcts_node2 is None:
+        log.debug(f"second node is None, returning first node ({gymcts_node1})")
+        return gymcts_node1
+    if gymcts_node1 is None and gymcts_node2 is None:
+        log.error("Both nodes are None")
+        raise ValueError("Both nodes are None")
+
+    if gymcts_node1.is_leaf() and not gymcts_node2.is_leaf():
+        log.debug(f"first node is leaf, second node is not leaf")
+        gymcts_node2.parent = gymcts_node1.parent
+        log.debug(f"returning first node: {gymcts_node2}")
+        return gymcts_node2
+
+    if gymcts_node2.is_leaf() and not gymcts_node1.is_leaf():
+        log.debug(f"second node is leaf, first node is not leaf")
+        log.debug(f"returning first node: {gymcts_node1}")
+        return gymcts_node1
+
+    if gymcts_node1.is_leaf() and gymcts_node2.is_leaf():
+        log.debug(f"both nodes are leafs, returning first node")
+        log.debug(f"returning first node: {gymcts_node1}")
+        return gymcts_node1
+
+    # check if gymcts_node1 and gymcts_node2 have the same children
+    if gymcts_node1.children.keys() != gymcts_node2.children.keys():
+        log.error("Nodes have different children")
+        raise ValueError("Nodes have different children")
+
+    for (action1, child1), (action2, child2) in zip(gymcts_node1.children.items(), gymcts_node2.children.items()):
+        if action1 != action2:
+            log.error("Actions are different")
+            raise ValueError("Actions are different")
+        log.debug(f"merging children with action {action1} for node {gymcts_node1}")
+        gymcts_node1.children[action1] = merge_nodes(
+            child1,
+            child2,
+            perform_state_equality_check=perform_state_equality_check
+        )
+
+    visit_count = gymcts_node1.visit_count + gymcts_node2.visit_count
+    mean_value = (
+                         gymcts_node1.mean_value * gymcts_node1.visit_count + gymcts_node2.mean_value * gymcts_node2.visit_count) / visit_count
+    max_value = max(gymcts_node1.max_value, gymcts_node2.max_value)
+    min_value = min(gymcts_node1.min_value, gymcts_node2.min_value)
+
+    gymcts_node1.visit_count = visit_count
+    gymcts_node1.mean_value = mean_value
+    gymcts_node1.max_value = max_value
+    gymcts_node1.min_value = min_value
+    log.debug(f"merged node: {gymcts_node1}")
+    log.debug(f"returning node: {gymcts_node1}")
+    return gymcts_node1
+
+
+class DistributedGymctsAgent:
+    render_tree_after_step: bool = False
+    render_tree_max_depth: int = 2
+    exclude_unvisited_nodes_from_render: bool = False
+    number_of_simulations_per_step: int = 25
+    num_parallel: int = 4
+
+    env: GymctsABC
+    search_root_node: GymctsNode  # NOTE: this is not the same as the root of the tree!
+    clear_mcts_tree_after_step: bool
+
+    def __init__(self,
+                 env: GymctsABC,
+                 render_tree_after_step: bool = False,
+                 render_tree_max_depth: int = 2,
+                 num_parallel: int = 4,
+                 number_of_simulations_per_step: int = 25,
+                 exclude_unvisited_nodes_from_render: bool = False
+                 ):
+        # check if action space of env is discrete
+        if not isinstance(env.action_space, gym.spaces.Discrete):
+            raise ValueError("Action space must be discrete.")
+
+        self.num_parallel = num_parallel
+
+        self.render_tree_after_step = render_tree_after_step
+        self.exclude_unvisited_nodes_from_render = exclude_unvisited_nodes_from_render
+        self.render_tree_max_depth = render_tree_max_depth
+
+        self.number_of_simulations_per_step = number_of_simulations_per_step
+
+        self.env = env
+
+        self.search_root_node = GymctsNode(
+            action=None,
+            parent=None,
+            env_reference=env,
+        )
+
+    def solve(self, num_simulations_per_step: int = None, render_tree_after_step: bool = None) -> list[int]:
+
+        if num_simulations_per_step is None:
+            num_simulations_per_step = self.number_of_simulations_per_step
+        if render_tree_after_step is None:
+            render_tree_after_step = self.render_tree_after_step
+
+        log.debug(f"Solving from root node: {self.search_root_node}")
+
+        current_node = self.search_root_node
+
+        action_list = []
+
+        while not current_node.terminal:
+            next_action, current_node = self.perform_mcts_step(num_simulations=num_simulations_per_step,
+                                                               render_tree_after_step=render_tree_after_step)
+            log.info(
+                f"selected action {next_action} after {self.num_parallel} x {num_simulations_per_step} simulations.")
+            action_list.append(next_action)
+            log.info(f"current action list: {action_list}")
+
+
+        log.info(f"Final action list: {action_list}")
+        # restore state of current node
+        return action_list
+
+    def perform_mcts_step(self, search_start_node: GymctsNode = None, num_simulations: int = None,
+                          render_tree_after_step: bool = None, num_parallel: int = None) -> tuple[int, GymctsNode]:
+
+        if render_tree_after_step is None:
+            render_tree_after_step = self.render_tree_after_step
+
+        if render_tree_after_step is None:
+            render_tree_after_step = self.render_tree_after_step
+
+        if num_simulations is None:
+            num_simulations = self.number_of_simulations_per_step
+
+        if search_start_node is None:
+            search_start_node = self.search_root_node
+
+        if num_parallel is None:
+            num_parallel = self.num_parallel
+
+        # action = self.vanilla_mcts_search(
+        #    search_start_node=search_start_node,
+        #   num_simulations=num_simulations,
+        # )
+        # next_node = search_start_node.children[action]
+
+        mcts_interation_futures = [
+            mcts_lookahead.remote(
+                copy.deepcopy(search_start_node),
+                copy.deepcopy(self.env),
+                num_simulations=num_simulations
+            )
+            for _ in range(num_parallel)
+        ]
+
+        while mcts_interation_futures:
+            ready_gymcts_nodes, mcts_interation_futures = ray.wait(mcts_interation_futures)
+            for ready_node_ref in ready_gymcts_nodes:
+                ready_node = ray.get(ready_node_ref)
+
+                # merge the tree
+                search_start_node = merge_nodes(search_start_node, ready_node)
+
+        action = search_start_node.get_best_action()
+        next_node = search_start_node.children[action]
+
+        if self.render_tree_after_step:
+            self.show_mcts_tree(
+                start_node=search_start_node,
+                tree_max_depth=self.render_tree_max_depth
+            )
+
+
+        # to clear memory we need to remove all nodes except the current node
+        # this is done by setting the root node to the current node
+        # and setting the parent of the current node to None
+        # we also need to reset the children of the current node
+        # this is done by calling the reset method
+        #
+        # in a distributed setting we need we delete all previous nodes
+        # this is because backpropagation merging trees is already computationally expensive
+        # and backpropagating the whole tree would be even more expensive
+        next_node.reset()
+
+        self.search_root_node = next_node
+
+        return action, next_node
+
+    def show_mcts_tree(self, start_node: GymctsNode = None, tree_max_depth: int = None) -> None:
+
+        if start_node is None:
+            start_node = self.search_root_node
+
+        if tree_max_depth is None:
+            tree_max_depth = self.render_tree_max_depth
+
+        print(start_node.__str__(colored=True, action_space_n=self.env.action_space.n))
+        for line in _generate_mcts_tree(
+                start_node=start_node,
+                depth=tree_max_depth,
+                action_space_n=self.env.action_space.n
+        ):
+            print(line)
+
+    def show_mcts_tree_from_root(self, tree_max_depth: int = None) -> None:
+        self.show_mcts_tree(start_node=self.search_root_node.get_root(), tree_max_depth=tree_max_depth)
+
+
+if __name__ == '__main__':
+    ray.init()
+
+    log.setLevel(20)  # 10=DEBUG, 20=INFO, 30=WARNING, 40=ERROR, 50=CR
+    env = gym.make('FrozenLake-v1', desc=None, map_name="4x4", is_slippery=False)
+    env.reset()
+
+    # 1. wrap the environment with the naive wrapper or a custom gymcts wrapper
+    # env1 = ActionHistoryMCTSGymEnvWrapper(env1)
+    env = DeepCopyMCTSGymEnvWrapper(env)
+
+    # 2. create the agent
+    agent1 = DistributedGymctsAgent(
+        env=env,
+        render_tree_after_step=True,
+        number_of_simulations_per_step=1000,
+        exclude_unvisited_nodes_from_render=True,
+        num_parallel=1,
+    )
+    import time
+
+    start_time = time.perf_counter()
+    actions = agent1.solve()
+    end_time = time.perf_counter()
+
+    print(f"solution time pro action: {end_time - start_time}/{len(actions)}")

gymcts/{gymcts_gym_env.py → gymcts_env_abc.py}

@@ -5,7 +5,7 @@ import gymnasium as gym
 TSoloMCTSNode = TypeVar("TSoloMCTSNode", bound="SoloMCTSNode")
 
 
-class SoloMCTSGymEnv(ABC, gym.Env):
+class GymctsABC(ABC, gym.Env):
 
     @abstractmethod
     def get_state(self) -> Any:

gymcts/gymcts_node.py

@@ -4,21 +4,17 @@ import math
 
 from typing import TypeVar, Any, SupportsFloat, Callable, Generator
 
-from gymcts.gymcts_gym_env import SoloMCTSGymEnv
+from gymcts.gymcts_env_abc import GymctsABC
 
 from gymcts.logger import log
 
-TSoloMCTSNode = TypeVar("TSoloMCTSNode", bound="SoloMCTSNode")
+TGymctsNode = TypeVar("TGymctsNode", bound="GymctsNode")
 
 
-
-
-class SoloMCTSNode:
-
+class GymctsNode:
     # static properties
     best_action_weight: float = 0.05
-    ubc_c  = 0.707
-
+    ubc_c = 0.707
 
     # attributes
     visit_count: int = 0
@@ -28,7 +24,6 @@ class SoloMCTSNode:
     terminal: bool = False
     state: Any
 
-
     def __str__(self, colored=False, action_space_n=None) -> str:
         if not colored:
 
@@ -39,11 +34,9 @@ class SoloMCTSNode:
 
         import gymcts.colorful_console_utils as ccu
 
-
         if self.is_root():
             return f"({ccu.CYELLOW}N{ccu.CEND}={self.visit_count}, {ccu.CYELLOW}Q_v{ccu.CEND}={self.mean_value:.2f}, {ccu.CYELLOW}best{ccu.CEND}={self.max_value:.2f})"
 
-
         if action_space_n is None:
             raise ValueError("action_space_n must be provided if colored is True")
 
@@ -68,25 +61,23 @@ class SoloMCTSNode:
             if isinstance(value, int):
                 return f"{color}{value}{e}"
 
-
         root_node = self.get_root()
         mean_val = f"{self.mean_value:.2f}"
 
         return ((f"("
-                f"{p}a{e}={ccu.wrap_evenly_spaced_color(s=self.action, n_of_item=self.action, n_classes=action_space_n)}, "
-                f"{p}N{e}={colorful_value(self.visit_count)}, "
-                f"{p}Q_v{e}={ccu.wrap_with_color_scale(s=mean_val, value=self.mean_value, min_val=root_node.min_value, max_val=root_node.max_value)}, "
-                f"{p}best{e}={colorful_value(self.max_value)}") +
+                 f"{p}a{e}={ccu.wrap_evenly_spaced_color(s=self.action, n_of_item=self.action, n_classes=action_space_n)}, "
+                 f"{p}N{e}={colorful_value(self.visit_count)}, "
+                 f"{p}Q_v{e}={ccu.wrap_with_color_scale(s=mean_val, value=self.mean_value, min_val=root_node.min_value, max_val=root_node.max_value)}, "
+                 f"{p}best{e}={colorful_value(self.max_value)}") +
                 (f", {p}ubc{e}={colorful_value(self.ucb_score())})" if not self.is_root() else ")"))
 
-
-    def traverse_nodes(self) -> Generator[TSoloMCTSNode, None, None]:
+    def traverse_nodes(self) -> Generator[TGymctsNode, None, None]:
         yield self
         if self.children:
             for child in self.children.values():
                 yield from child.traverse_nodes()
 
-    def get_root(self) -> TSoloMCTSNode:
+    def get_root(self) -> TGymctsNode:
         if self.is_root():
             return self
         return self.parent.get_root()
@@ -101,19 +92,17 @@ class SoloMCTSNode:
             return 0
         return len(self.children) + sum(child.n_children_recursively() for child in self.children.values())
 
-
-
     def __init__(self,
                  action: int | None,
-                 parent: TSoloMCTSNode | None,
-                 env_reference: SoloMCTSGymEnv,
+                 parent: TGymctsNode | None,
+                 env_reference: GymctsABC,
                  ):
 
         # field depending on whether the node is a root node or not
         self.action: int | None
 
-        self.env_reference: SoloMCTSGymEnv
-        self.parent: SoloMCTSNode | None
+        self.env_reference: GymctsABC
+        self.parent: GymctsNode | None
         self.uuid = uuid.uuid4()
 
         if parent is None:
@@ -133,7 +122,7 @@ class SoloMCTSNode:
 
         from copy import copy
         self.state = env_reference.get_state()
-        #log.debug(f"saving state of node '{str(self)}' to memory location: {hex(id(self.state))}")
+        # log.debug(f"saving state of node '{str(self)}' to memory location: {hex(id(self.state))}")
         self.visit_count: int = 0
 
         self.mean_value: float = 0
@@ -143,8 +132,7 @@ class SoloMCTSNode:
         # safe valid action instead of calling the environment
         # this reduces the compute but increases the memory usage
         self.valid_actions: list[int] = env_reference.get_valid_actions()
-        self.children: dict[int, SoloMCTSNode] | None = None  # may be expanded later
-
+        self.children: dict[int, GymctsNode] | None = None  # may be expanded later
 
     def reset(self) -> None:
         self.parent = None
@@ -153,35 +141,33 @@ class SoloMCTSNode:
         self.mean_value: float = 0
         self.max_value: float = -float("inf")
         self.min_value: float = +float("inf")
-        self.children: dict[int, SoloMCTSNode] | None = None  # may be expanded later
+        self.children: dict[int, GymctsNode] | None = None  # may be expanded later
 
         # just setting the children of the parent node to None should be enough to trigger garbage collection
         # however, we also set the parent to None to make sure that the parent is not referenced anymore
         if self.parent:
             self.parent.reset()
 
-
-
     def is_root(self) -> bool:
         return self.parent is None
 
     def is_leaf(self) -> bool:
         return self.children is None or len(self.children) == 0
 
-    def get_random_child(self) -> TSoloMCTSNode:
+    def get_random_child(self) -> TGymctsNode:
         if self.is_leaf():
-            raise ValueError("cannot get random child of leaf node") #todo: maybe return self instead?
+            raise ValueError("cannot get random child of leaf node")  # todo: maybe return self instead?
 
         return list(self.children.values())[random.randint(0, len(self.children) - 1)]
 
     def get_best_action(self) -> int:
         return max(self.children.values(), key=lambda child: child.get_score()).action
 
-    def get_score(self) -> float: # todo: make it an attribute?
+    def get_score(self) -> float:  # todo: make it an attribute?
         # return self.mean_value
-        assert 0 <= SoloMCTSNode.best_action_weight <= 1
-        a = SoloMCTSNode.best_action_weight
-        return (1-a) * self.mean_value + a * self.max_value
+        assert 0 <= GymctsNode.best_action_weight <= 1
+        a = GymctsNode.best_action_weight
+        return (1 - a) * self.mean_value + a * self.max_value
 
     def get_mean_value(self) -> float:
         return self.mean_value
@@ -207,7 +193,7 @@ class SoloMCTSNode:
         if self.is_root():
             raise ValueError("ucb_score can only be called on non-root nodes")
         # c = 0.707 # todo: make it an attribute?
-        c = SoloMCTSNode.ubc_c
+        c = GymctsNode.ubc_c
         if self.visit_count == 0:
             return float("inf")
-        return self.mean_value + c * math.sqrt(math.log(self.parent.visit_count) / (self.visit_count))
+        return self.mean_value + c * math.sqrt(math.log(self.parent.visit_count) / (self.visit_count))

gymcts/gymcts_tree_plotter.py

@@ -0,0 +1,75 @@
+from gymcts.gymcts_node import GymctsNode
+
+from gymcts.logger import log
+
+
+def _generate_mcts_tree(
+        start_node: GymctsNode = None,
+        prefix: str = None,
+        depth: int = None,
+        exclude_unvisited_nodes_from_render: bool = True,
+        action_space_n: int = None
+) -> list[str]:
+    if prefix is None:
+        prefix = ""
+    import gymcts.colorful_console_utils as ccu
+
+    if start_node is None:
+        raise ValueError("start_node must not be None")
+
+    if action_space_n is None:
+        log.warning("action_space_n is None, defaulting to 100")
+        action_space_n = 100
+
+    # prefix components:
+    space = '    '
+    branch = '│   '
+    # pointers:
+    tee = '├── '
+    last = '└── '
+
+    contents = start_node.children.values() if start_node.children is not None else []
+    if exclude_unvisited_nodes_from_render:
+        contents = [node for node in contents if node.visit_count > 0]
+    # contents each get pointers that are ├── with a final └── :
+    # pointers = [tee] * (len(contents) - 1) + [last]
+    pointers = [tee for _ in range(len(contents) - 1)] + [last]
+
+    for pointer, current_node in zip(pointers, contents):
+        n_item = current_node.parent.action if current_node.parent is not None else 0
+        n_classes = action_space_n
+
+        pointer = ccu.wrap_evenly_spaced_color(
+            s=pointer,
+            n_of_item=n_item,
+            n_classes=n_classes,
+        )
+
+        yield prefix + pointer + f"{current_node.__str__(colored=True, action_space_n=n_classes)}"
+        if current_node.children and len(current_node.children):  # extend the prefix and recurse:
+            # extension = branch if pointer == tee else space
+            extension = branch if tee in pointer else space
+            # i.e. space because last, └── , above so no more |
+            extension = ccu.wrap_evenly_spaced_color(
+                s=extension,
+                n_of_item=n_item,
+                n_classes=n_classes,
+            )
+            if depth is not None and depth <= 0:
+                continue
+            yield from _generate_mcts_tree(
+                current_node,
+                prefix=prefix + extension,
+                action_space_n=action_space_n,
+                depth=depth - 1 if depth is not None else None
+            )
+
+
+def show_mcts_tree(
+        start_node: GymctsNode = None,
+        tree_max_depth: int = None,
+        action_space_n: int = None
+) -> None:
+    print(start_node.__str__(colored=True, action_space_n=action_space_n))
+    for line in _generate_mcts_tree(start_node=start_node, depth=tree_max_depth):
+        print(line)

{gymcts-1.0.0.dist-info → gymcts-1.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: gymcts
-Version: 1.0.0
+Version: 1.2.0
 Summary: A minimalistic implementation of the Monte Carlo Tree Search algorithm for planning problems fomulated as gymnaisum reinforcement learning environments.
 Author: Alexander Nasuta
 Author-email: Alexander Nasuta <alexander.nasuta@wzl-iqs.rwth-aachen.de>
@@ -25,7 +25,7 @@ License: MIT License
         LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
         OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
         SOFTWARE.
-Project-URL: Homepage, https://github.com/Alexander-Nasuta/pypitemplate
+Project-URL: Homepage, https://github.com/Alexander-Nasuta/gymcts
 Platform: unix
 Platform: linux
 Platform: osx
@@ -34,7 +34,7 @@ Platform: win32
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
-Requires-Python: >=3.9
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: rich
@@ -63,6 +63,9 @@ Requires-Dist: furo; extra == "dev"
 Requires-Dist: twine; extra == "dev"
 Requires-Dist: sphinx-copybutton; extra == "dev"
 Requires-Dist: nbsphinx; extra == "dev"
+Requires-Dist: jupytext; extra == "dev"
+Requires-Dist: jupyter; extra == "dev"
+Dynamic: license-file
 
 # Graph Matrix Job Shop Env
 
@@ -118,8 +121,8 @@ The NaiveSoloMCTSGymEnvWrapper can be used with non-deterministic environments,
 ```python
 import gymnasium as gym
 
-from gymcts.gymcts_agent import SoloMCTSAgent
-from gymcts.gymcts_naive_wrapper import NaiveSoloMCTSGymEnvWrapper
+from gymcts.gymcts_agent import GymctsAgent
+from gymcts.gymcts_deepcopy_wrapper import DeepCopyMCTSGymEnvWrapper
 
 from gymcts.logger import log
 
@@ -133,10 +136,10 @@ if __name__ == '__main__':
     env.reset()
 
     # 1. wrap the environment with the naive wrapper or a custom gymcts wrapper
-    env = NaiveSoloMCTSGymEnvWrapper(env)
+    env = DeepCopyMCTSGymEnvWrapper(env)
 
     # 2. create the agent
-    agent = SoloMCTSAgent(
+    agent = GymctsAgent(
         env=env,
         clear_mcts_tree_after_step=False,
         render_tree_after_step=True,
@@ -170,13 +173,13 @@ if __name__ == '__main__':
 A minimal example of how to use the package with the FrozenLake environment and the DeterministicSoloMCTSGymEnvWrapper is provided in the following code snippet below.
 The DeterministicSoloMCTSGymEnvWrapper can be used with deterministic environments, such as the FrozenLake environment without slippery ice.
 
-The DeterministicSoloMCTSGymEnvWrapper saves the action sequence that lead to the current state in the MCTS node. 
+The DeterministicSoloMCTSGymEnvWrapper saves the action sequence that lead to the current state in the MCTS node.
 
 ```python
 import gymnasium as gym
 
-from gymcts.gymcts_agent import SoloMCTSAgent
-from gymcts.gymcts_deterministic_wrapper import DeterministicSoloMCTSGymEnvWrapper
+from gymcts.gymcts_agent import GymctsAgent
+from gymcts.gymcts_action_history_wrapper import ActionHistoryMCTSGymEnvWrapper
 
 from gymcts.logger import log
 
@@ -190,10 +193,10 @@ if __name__ == '__main__':
     env.reset()
 
     # 1. wrap the environment with the wrapper
-    env = DeterministicSoloMCTSGymEnvWrapper(env)
+    env = ActionHistoryMCTSGymEnvWrapper(env)
 
     # 2. create the agent
-    agent = SoloMCTSAgent(
+    agent = GymctsAgent(
         env=env,
         clear_mcts_tree_after_step=False,
         render_tree_after_step=True,
@@ -232,8 +235,8 @@ To create a video of the solution of the FrozenLake environment, you can use the
 ```python  
 import gymnasium as gym
 
-from gymcts.gymcts_agent import SoloMCTSAgent
-from gymcts.gymcts_naive_wrapper import NaiveSoloMCTSGymEnvWrapper
+from gymcts.gymcts_agent import GymctsAgent
+from gymcts.gymcts_deepcopy_wrapper import DeepCopyMCTSGymEnvWrapper
 
 from gymcts.logger import log
 
@@ -249,10 +252,10 @@ if __name__ == '__main__':
     env.reset()
 
     # 1. wrap the environment with the naive wrapper or a custom gymcts wrapper
-    env = NaiveSoloMCTSGymEnvWrapper(env)
+    env = DeepCopyMCTSGymEnvWrapper(env)
 
     # 2. create the agent
-    agent = SoloMCTSAgent(
+    agent = GymctsAgent(
         env=env,
         clear_mcts_tree_after_step=False,
         render_tree_after_step=True,
@@ -413,13 +416,12 @@ The color gradient is based on the minimum and maximum values of the respective
 The visualisation is rendered in the terminal and can be limited to a certain depth of the tree.
 The default depth is 2.
 
-
 ```python
 import gymnasium as gym
 
-from gymcts.gymcts_agent import SoloMCTSAgent
-from gymcts.gymcts_deterministic_wrapper import DeterministicSoloMCTSGymEnvWrapper
-from gymcts.gymcts_naive_wrapper import NaiveSoloMCTSGymEnvWrapper
+from gymcts.gymcts_agent import GymctsAgent
+from gymcts.gymcts_action_history_wrapper import ActionHistoryMCTSGymEnvWrapper
+from gymcts.gymcts_deepcopy_wrapper import DeepCopyMCTSGymEnvWrapper
 
 from gymcts.logger import log
 
@@ -433,10 +435,10 @@ if __name__ == '__main__':
     env.reset()
 
     # wrap the environment with the naive wrapper or a custom gymcts wrapper
-    env = DeterministicSoloMCTSGymEnvWrapper(env)
+    env = ActionHistoryMCTSGymEnvWrapper(env)
 
     # create the agent
-    agent = SoloMCTSAgent(
+    agent = GymctsAgent(
         env=env,
         clear_mcts_tree_after_step=False,
         render_tree_after_step=False,

gymcts-1.2.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+gymcts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+gymcts/colorful_console_utils.py,sha256=OhULcXHKbEA4uJDAEYCTcW6wUv0LsHX_XSYzZ_Szsv4,4553
+gymcts/gymcts_action_history_wrapper.py,sha256=AjvBBwd1t9-nTYP09aMdlScAkFNXf5vOagejpjWYOPo,3810
+gymcts/gymcts_agent.py,sha256=O2y98jKFjR5TzqVV7DO1jlcYDyzAgd_H2RF4-w4NP0g,8499
+gymcts/gymcts_deepcopy_wrapper.py,sha256=OleQTnvxv3gLEo8-2asyeo-CpZ4HEbgyFGS5DTCD7NM,4167
+gymcts/gymcts_distributed_agent.py,sha256=M7dyBfC8u3M99PJFoXKgIc_CPTyHGppmktkH-y9ci4U,10448
+gymcts/gymcts_env_abc.py,sha256=7nCRiiClmmVLX-d_Q1dxeztmuvmAtmWZwjT81zrG1_w,575
+gymcts/gymcts_node.py,sha256=PT_YZFwt1zjuvd8i9Wb5LEkHAqmJOFyPDp3GFD05lqM,7138
+gymcts/gymcts_tree_plotter.py,sha256=eg207wHcDepwWODXzmDYQn1Aai29Cs4jFS1HNvAhlXs,2651
+gymcts/logger.py,sha256=nAkUa4djiuCR7hF0EUsplhqFHCp76QcOX1cV3lIPzOI,937
+gymcts-1.2.0.dist-info/licenses/LICENSE,sha256=UGe75WojDiw_77SEnK2aysEDlElRlkWie7U7NaAFx00,1072
+gymcts-1.2.0.dist-info/METADATA,sha256=zhEIFo0rOnv5hCv6ukImkq-9nshO4EfXMbHlhNlYhyA,23640
+gymcts-1.2.0.dist-info/WHEEL,sha256=DK49LOLCYiurdXXOXwGJm6U4DkHkg4lcxjhqwRa0CP4,91
+gymcts-1.2.0.dist-info/top_level.txt,sha256=E8MoLsPimUPD0H1Y6lum4TVe-lhSDAyBAXGrkYIT52w,7
+gymcts-1.2.0.dist-info/RECORD,,

{gymcts-1.0.0.dist-info → gymcts-1.2.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.8.2)
+Generator: setuptools (78.0.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

gymcts-1.0.0.dist-info/RECORD

@@ -1,13 +0,0 @@
-gymcts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-gymcts/colorful_console_utils.py,sha256=bbZzRFzimhsIhbT-nmz6v62WJLxFDgzFvqI_pmIsckE,4526
-gymcts/gymcts_agent.py,sha256=TJXJH77T95EP3ZNtzWqlGw9iFF1R-nsItp7UA1ZlXUs,10537
-gymcts/gymcts_deterministic_wrapper.py,sha256=PILGPaQnyG2u_2u48MEE3aeJCtdgjjO55ZFDxeIVeH0,3824
-gymcts/gymcts_gym_env.py,sha256=R1Z1fhoywdXmPt_FYgrarIh0YFQvCifayAWnCcEiJKE,580
-gymcts/gymcts_naive_wrapper.py,sha256=qeQ7rzBz7BFv2yCJj3GmdFt5UlTx5VHMw5ImZUl9H5k,4178
-gymcts/gymcts_node.py,sha256=jxdtuC1iqeRtEA-Qfvq-mOuM8vdDl43iWe5hqItG90w,7185
-gymcts/logger.py,sha256=nAkUa4djiuCR7hF0EUsplhqFHCp76QcOX1cV3lIPzOI,937
-gymcts-1.0.0.dist-info/LICENSE,sha256=UGe75WojDiw_77SEnK2aysEDlElRlkWie7U7NaAFx00,1072
-gymcts-1.0.0.dist-info/METADATA,sha256=sAXJQreADqEOviVL8nT8fmrx7hP-qM7C_-SC5FNw-94,23572
-gymcts-1.0.0.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
-gymcts-1.0.0.dist-info/top_level.txt,sha256=E8MoLsPimUPD0H1Y6lum4TVe-lhSDAyBAXGrkYIT52w,7
-gymcts-1.0.0.dist-info/RECORD,,