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

gymcts/gymcts_distributed_agent.py

@@ -0,0 +1,299 @@
+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,
+                 clear_mcts_tree_after_step: bool = False,
+                 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.clear_mcts_tree_after_step = clear_mcts_tree_after_step
+
+        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
+                if not self.clear_mcts_tree_after_step:
+                    self.backpropagation(search_start_node, ready_node.mean_value, ready_node.visit_count)
+                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
+            )
+
+        if self.clear_mcts_tree_after_step:
+            # 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
+            next_node.reset()
+
+        self.search_root_node = next_node
+
+        return action, next_node
+
+    def backpropagation(self, node: GymctsNode, average_episode_return: float, num_episodes: int) -> None:
+        log.debug(f"performing backpropagation from leaf node: {node}")
+        while not node.is_root():
+            node.mean_value = (node.mean_value * node.visit_count + average_episode_return * num_episodes) / (
+                        node.visit_count + num_episodes)
+            node.visit_count += num_episodes
+            node.max_value = max(node.max_value, average_episode_return)
+            node.min_value = min(node.min_value, average_episode_return)
+            node = node.parent
+        # also update root node
+        node.mean_value = (node.mean_value * node.visit_count + average_episode_return * num_episodes) / (
+                    node.visit_count + num_episodes)
+        node.visit_count += num_episodes
+        node.max_value = max(node.max_value, average_episode_return)
+        node.min_value = min(node.min_value, average_episode_return)
+
+    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=10,
+        exclude_unvisited_nodes_from_render=True,
+        num_parallel=1,
+    )
+    import time
+
+    start_time = time.perf_counter()
+    actions = agent1.solve()
+    end_time = time.perf_counter()
+
+    agent1.show_mcts_tree_from_root()
+
+    print(f"solution time pro action: {end_time - start_time}/{len(actions)}")

gymcts/gymcts_env_abc.py

@@ -0,0 +1,61 @@
+from typing import TypeVar, Any, SupportsFloat, Callable
+from abc import ABC, abstractmethod
+import gymnasium as gym
+
+TSoloMCTSNode = TypeVar("TSoloMCTSNode", bound="SoloMCTSNode")
+
+
+class GymctsABC(ABC, gym.Env):
+
+    @abstractmethod
+    def get_state(self) -> Any:
+        """
+        Returns the current state of the environment. The state can be any datatype in principle, that allows to restore
+        the environment to the same state. The state is used to restore the environment unsing the load_state method.
+
+        It's recommended to use a numpy array if possible, as it is easy to serialize and deserialize.
+
+        :return: the current state of the environment
+        """
+        pass
+
+    @abstractmethod
+    def load_state(self, state: Any) -> None:
+        """
+        Loads the state of the environment. The state can be any datatype in principle, that allows to restore the
+        environment to the same state. The state is used to restore the environment unsing the load_state method.
+
+        :param state: the state to load
+        :return: None
+        """
+        pass
+
+    @abstractmethod
+    def is_terminal(self) -> bool:
+        """
+        Returns True if the environment is in a terminal state, False otherwise.
+        :return:
+        """
+        pass
+
+    @abstractmethod
+    def get_valid_actions(self) -> list[int]:
+        """
+        Returns a list of valid actions for the current state of the environment.
+        This used to obtain potential actions/subsequent sates for the MCTS tree.
+        :return: the list of valid actions
+        """
+        pass
+
+    @abstractmethod
+    def rollout(self) -> float:
+        """
+        Performs a rollout from the current state of the environment and returns the return (sum of rewards) of the rollout.
+
+        Please make sure the return value is in the interval [-1, 1].
+        Otherwise, the MCTS algorithm will not work as expected (due to a male-fitted exploration coefficient;
+        exploration and exploitation are not well-balanced then).
+
+        :return: the return of the rollout
+        """
+        pass

gymcts/gymcts_node.py

@@ -4,32 +4,41 @@ 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
+    best_action_weight: float = 0.05 # weight for the best action
+    ubc_c = 0.707 # exploration coefficient
 
 
-    # attributes
-    visit_count: int = 0
-    mean_value: float = 0
-    max_value: float = -float("inf")
-    min_value: float = +float("inf")
-    terminal: bool = False
-    state: Any
 
+    # attributes
+    #
+    # Note these attributes are not static. Their defined here to give developers a hint what fields are available
+    # in the class. They are not static because they are not shared between instances of the class in scope of
+    # this library.
+    visit_count: int = 0 # number of times the node has been visited
+    mean_value: float = 0 # mean value of the node
+    max_value: float = -float("inf") # maximum value of the node
+    min_value: float = +float("inf") # minimum value of the node
+    terminal: bool = False # whether the node is terminal or not
+    state: Any = None # state of the node
 
     def __str__(self, colored=False, action_space_n=None) -> str:
+        """
+        Returns a string representation of the node. The string representation is used for visualisation purposes.
+        It is used for example in the mcts tree visualisation functionality.
+
+        :param colored: true if the string representation should be colored, false otherwise. (ture is used by the mcts tree visualisation)
+        :param action_space_n: the number of actions in the action space. This is used for coloring the action in the string representation.
+        :return: a potentially colored string representation of the node.
+        """
         if not colored:
 
             if not self.is_root():
@@ -39,11 +48,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,52 +75,78 @@ 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[TGymctsNode, None, None]:
+        """
+        Traverse the tree and yield all nodes in the tree.
 
-    def traverse_nodes(self) -> Generator[TSoloMCTSNode, None, None]:
+        :return: a generator that yields all nodes in the tree.
+        """
         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:
+        """
+        Returns the root node of the tree. The root node is the node that has no parent.
+
+        :return: the root node of the tree.
+        """
         if self.is_root():
             return self
         return self.parent.get_root()
 
     def max_tree_depth(self):
+        """
+        Returns the maximum depth of the tree. The depth of a node is the number of edges from
+        the node to the root node.
+
+        :return: the maximum depth of the tree.
+        """
         if self.is_leaf():
             return 0
         return 1 + max(child.max_tree_depth() for child in self.children.values())
 
     def n_children_recursively(self):
+        """
+        Returns the number of children of the node recursively. The number of children of a node is the number of
+        children of the node plus the number of children of all children of the node.
+
+        :return: the number of children of the node recursively.
+        """
         if self.is_leaf():
             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,
                  ):
+        """
+        Initializes the node. The node is initialized with the state of the environment and the action that was taken to
+        reach the node. The node is also initialized with the parent node and the environment reference.
+
+        :param action: the action that was taken to reach the node. If the node is a root node, this parameter is None.
+        :param parent: the parent node of the node. If the node is a root node, this parameter is None.
+        :param env_reference: a reference to the environment. The environment is used to get the state of the node and the valid actions.
+        """
 
         # 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 +166,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 +176,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,40 +185,71 @@ 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:
+        """
+        Returns true if the node is a root node. A root node is a node that has no parent.
+
+        :return: true if the node is a root node, false otherwise.
+        """
         return self.parent is None
 
     def is_leaf(self) -> bool:
+        """
+        Returns true if the node is a leaf node. A leaf node is a node that has no children. A leaf node is a node that has no children.
+
+        :return: true if the node is a leaf node, false otherwise.
+        """
         return self.children is None or len(self.children) == 0
 
-    def get_random_child(self) -> TSoloMCTSNode:
+    def get_random_child(self) -> TGymctsNode:
+        """
+        Returns a random child of the node. A random child is a child that is selected randomly from the list of children.
+        :return:
+        """
         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:
+        """
+        Returns the best action of the node. The best action is the action that has the highest score.
+        The score is calculated using the get_score() method. The best action is the action that has the highest score.
+        The best action is the action that has the highest score.
+
+        :return: the best action of the node.
+        """
         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?
+        """
+        Returns the score of the node. The score is calculated using the mean value and the maximum value of the node.
+        The score is calculated using the formula: score = (1 - a) * mean_value + a * max_value
+        where a is the best action weight.
+
+        :return: the score of the node.
+        """
         # 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
 
     def get_max_value(self) -> float:
+        """
+        Returns the maximum value of the node. The maximum value is the maximum value of the node.
+
+        :return: the maximum value of the node.
+        """
         return self.max_value
 
     def ucb_score(self):
@@ -207,7 +270,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))