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

gymcts/colorful_console_utils.py

@@ -106,6 +106,18 @@ def wrap_with_color_codes(s: object, /, r: int | float, g: int | float, b: int |
 
 
 def wrap_evenly_spaced_color(s: Any, n_of_item: int, n_classes: int, c_map="rainbow") -> str:
+    """
+    Wraps a string with a color scale (a matplotlib c_map) based on the n_of_item and n_classes.
+    This function is used to color code the available actions in the MCTS tree visualisation.
+    The children of the MCTS tree are colored based on their action for a clearer visualisation.
+
+    :param s: the string (or object) to be wrapped. objects are converted to string (using the __str__ function).
+    :param n_of_item: the index of the item to be colored. In a mcts tree, this is the (parent-)action of the node.
+    :param n_classes: the number of classes (or items) to be colored. In a mcts tree, this is the number of available actions.
+    :param c_map: the colormap to be used (default is 'rainbow').
+                  The colormap can be any matplotlib colormap, e.g. 'viridis', 'plasma', 'inferno', 'magma', 'cividis'.
+    :return: a string that contains the color-codes (prefix and suffix) and the string s in between.
+    """
     if s is None or n_of_item is None or n_classes is None:
         return s
 
@@ -119,6 +131,16 @@ def wrap_evenly_spaced_color(s: Any, n_of_item: int, n_classes: int, c_map="rain
 
 
 def wrap_with_color_scale(s: str, value: float, min_val: float, max_val: float, c_map=None) -> str:
+    """
+    Wraps a string with a color scale (a matplotlib c_map) based on the value, min_val, and max_val.
+
+    :param s: the string to be wrapped
+    :param value: the value to be mapped to a color
+    :param min_val: the minimum value of the scale
+    :param max_val: the maximum value of the scale
+    :param c_map: the colormap to be used (default is 'rainbow')
+    :return:
+    """
     if s is None or min_val is None or max_val is None or min_val >= max_val:
         return s
 

gymcts/gymcts_action_history_wrapper.py

@@ -1,8 +1,7 @@
 import random
-import copy
 
 import numpy as np
-from typing import TypeVar, Any, SupportsFloat, Callable
+from typing import Any, SupportsFloat, Callable
 import gymnasium as gym
 from gymnasium.core import WrapperActType, WrapperObsType
 from gymnasium.wrappers import RecordEpisodeStatistics
@@ -13,6 +12,21 @@ from gymcts.logger import log
 
 
 class ActionHistoryMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
+    """
+    A wrapper for gym environments that implements the GymctsABC interface.
+    It uses the action history as state representation.
+    Please note that this is not the most efficient way to implement the state representation.
+    It is supposed to be used to see if your use-case works well with the MCTS algorithm.
+    If it does, you can consider implementing all GymctsABC methods in a more efficient way.
+    The action history is a list of actions taken in the environment.
+    The state is represented as a list of actions taken in the environment.
+    The state is used to restore the environment using the load_state method.
+
+    It is supposed to be used to see if your use-case works well with the MCTS algorithm.
+    If it does, you can consider implementing all GymctsABC methods in a more efficient way.
+    """
+
+    # helper attributes for the wrapper
     _terminal_flag: bool = False
     _last_reward: SupportsFloat = 0
     _step_tuple: tuple[WrapperObsType, SupportsFloat, bool, bool, dict[str, Any]] = None
@@ -25,6 +39,17 @@ class ActionHistoryMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
             action_mask_fn: str | Callable[[gym.Env], np.ndarray] | None = None,
             buffer_length: int = 100,
     ):
+        """
+        A wrapper for gym environments that implements the GymctsABC interface.
+        It uses the action history as state representation.
+        Please note that this is not the most efficient way to implement the state representation.
+        It is supposed to be used to see if your use-case works well with the MCTS algorithm.
+        If it does, you can consider implementing all GymctsABC methods in a more efficient way.
+
+        :param env: the environment to wrap
+        :param action_mask_fn: a function that takes the environment as input and returns a mask of valid actions
+        :param buffer_length: the length of the buffer for recording episodes for determining their rollout returns
+        """
         # wrap with RecordEpisodeStatistics if it is not already wrapped
         env = RecordEpisodeStatistics(env, buffer_length=buffer_length)
 
@@ -48,6 +73,17 @@ class ActionHistoryMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
                 self._action_mask_fn = action_mask_fn
 
     def load_state(self, state: list[int]) -> None:
+        """
+        Loads the state of the environment. The state is a list of actions taken in the environment.
+
+        The environment is reset and all actions in the state are performed in order to restore the environment to the
+        same state.
+
+        This works only for deterministic environments!
+
+        :param state: the state to load
+        :return: None
+        """
         self.env.reset()
         self._wrapper_action_history = []
 
@@ -56,15 +92,30 @@ class ActionHistoryMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
             self._wrapper_action_history.append(action)
 
     def is_terminal(self) -> bool:
+        """
+        Returns True if the environment is in a terminal state, False otherwise.
+
+        :return:
+        """
         if not len(self.get_valid_actions()):
             return True
         else:
             return self._terminal_flag
 
     def action_masks(self) -> np.ndarray | None:
+        """
+        Returns the action masks for the environment. If the action_mask_fn is not set, it returns None.
+
+        :return:
+        """
         return self._action_mask_fn(self.env) if self._action_mask_fn is not None else None
 
     def get_valid_actions(self) -> list[int]:
+        """
+        Returns a list of valid actions for the current state of the environment.
+
+        :return: a list of valid actions
+        """
         if self._action_mask_fn is None:
             action_space: gym.spaces.Discrete = self.env.action_space  # Type hinting
             return list(range(action_space.n))
@@ -72,6 +123,12 @@ class ActionHistoryMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
             return [i for i, mask in enumerate(self.action_masks()) if mask]
 
     def rollout(self) -> float:
+        """
+        Performs a random rollout from the current state of the environment and returns the return (sum of rewards)
+        of the rollout.
+
+        :return: the return of the rollout
+        """
         log.debug("performing rollout")
         # random rollout
         # perform random valid action util terminal
@@ -92,11 +149,24 @@ class ActionHistoryMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
         return episode_return
 
     def get_state(self) -> list[int]:
+        """
+        Returns the current state of the environment. The state is a list of actions taken in the environment,
+        namely all action that have been taken in the environment so far (since the last reset).
+
+        :return: a list of actions taken in the environment
+        """
+
         return self._wrapper_action_history.copy()
 
     def step(
             self, action: WrapperActType
     ) -> tuple[WrapperObsType, SupportsFloat, bool, bool, dict[str, Any]]:
+        """
+        Performs a step in the environment. It adds the action to the action history and updates the terminal flag.
+
+        :param action: action to perform in the environment
+        :return: the step tuple of the environment (obs, reward, terminated, truncated, info)
+        """
         step_tuple = self.env.step(action)
         self._wrapper_action_history.append(action)
         obs, reward, terminated, truncated, info = step_tuple

gymcts/gymcts_agent.py

@@ -1,7 +1,8 @@
 import copy
+import random
 import gymnasium as gym
 
-from typing import TypeVar, Any, SupportsFloat, Callable
+from typing import TypeVar, Any, SupportsFloat, Callable, Literal
 
 from gymcts.gymcts_env_abc import GymctsABC
 from gymcts.gymcts_deepcopy_wrapper import DeepCopyMCTSGymEnvWrapper
@@ -10,7 +11,9 @@ from gymcts.gymcts_tree_plotter import _generate_mcts_tree
 
 from gymcts.logger import log
 
-TSoloMCTSNode = TypeVar("TSoloMCTSNode", bound="SoloMCTSNode")
+
+
+
 
 
 class GymctsAgent:
@@ -23,17 +26,50 @@ class GymctsAgent:
     search_root_node: GymctsNode  # NOTE: this is not the same as the root of the tree!
     clear_mcts_tree_after_step: bool
 
+
+    # (num_simulations: int, step_idx: int) -> int
+    @staticmethod
+    def calc_number_of_simulations_per_step(num_simulations: int, step_idx: int) -> int:
+        """
+        A function that returns a constant number of simulations per step.
+
+        :param num_simulations: The number of simulations to return.
+        :param step_idx: The current step index (not used in this function).
+        :return: A callable that takes an environment as input and returns the constant number of simulations.
+        """
+        return num_simulations
+
     def __init__(self,
                  env: GymctsABC,
                  clear_mcts_tree_after_step: bool = True,
                  render_tree_after_step: bool = False,
                  render_tree_max_depth: int = 2,
                  number_of_simulations_per_step: int = 25,
-                 exclude_unvisited_nodes_from_render: bool = False
+                 exclude_unvisited_nodes_from_render: bool = False,
+                 calc_number_of_simulations_per_step: Callable[[int,int], int] = None,
+                 score_variate: Literal["UCT_v0", "UCT_v1", "UCT_v2",] = "UCT_v0",
+                 best_action_weight=None,
                  ):
         # check if action space of env is discrete
         if not isinstance(env.action_space, gym.spaces.Discrete):
             raise ValueError("Action space must be discrete.")
+        if calc_number_of_simulations_per_step is not None:
+            # check if the provided function is callable
+            if not callable(calc_number_of_simulations_per_step):
+                raise ValueError("calc_number_of_simulations_per_step must be a callable accepting two arguments: num_simulations and step_idx.")
+            # assign the provided function to the attribute
+            # it needs to be staticmethod to be used as a class attribute
+            print("Using provided calc_number_of_simulations_per_step function.")
+            self.calc_number_of_simulations_per_step = staticmethod(calc_number_of_simulations_per_step)
+        if score_variate not in ["UCT_v0", "UCT_v1", "UCT_v2"]:
+            raise ValueError("score_variate must be one of ['UCT_v0', 'UCT_v1', 'UCT_v2'].")
+        GymctsNode.score_variate = score_variate
+
+        if best_action_weight is not None:
+            if best_action_weight < 0 or best_action_weight > 1:
+                raise ValueError("best_action_weight must be in range [0, 1].")
+            GymctsNode.best_action_weight = best_action_weight
+
 
         self.render_tree_after_step = render_tree_after_step
         self.exclude_unvisited_nodes_from_render = exclude_unvisited_nodes_from_render
@@ -63,7 +99,10 @@ class GymctsAgent:
         # NAVIGATION STRATEGY
         # select child with highest UCB score
         while not temp_node.is_leaf():
-            temp_node = max(temp_node.children.values(), key=lambda child: child.ucb_score())
+            children = list(temp_node.children.values())
+            max_ucb_score = max(child.tree_policy_score() for child in children)
+            best_children = [child for child in children if child.tree_policy_score() == max_ucb_score]
+            temp_node = random.choice(best_children)
         log.debug(f"Selected leaf node: {temp_node}")
         return temp_node
 
@@ -84,7 +123,6 @@ class GymctsAgent:
                 parent=node,
                 env_reference=self.env,
             )
-
         node.children = child_dict
 
     def solve(self, num_simulations_per_step: int = None, render_tree_after_step: bool = None) -> list[int]:
@@ -100,13 +138,20 @@ class GymctsAgent:
 
         action_list = []
 
+        idx = 0
         while not current_node.terminal:
-            next_action, current_node = self.perform_mcts_step(num_simulations=num_simulations_per_step,
+            num_sims = self.calc_number_of_simulations_per_step(num_simulations_per_step, idx)
+
+            log.info(f"Performing MCTS step {idx} with {num_sims} simulations.")
+
+            next_action, current_node = self.perform_mcts_step(num_simulations=num_sims,
                                                                render_tree_after_step=render_tree_after_step)
-            log.info(f"selected action {next_action} after {num_simulations_per_step} simulations.")
+            log.info(f"selected action {next_action} after {num_sims} simulations.")
             action_list.append(next_action)
             log.info(f"current action list: {action_list}")
 
+            idx += 1
+
         log.info(f"Final action list: {action_list}")
         # restore state of current node
         return action_list
@@ -145,6 +190,8 @@ class GymctsAgent:
             # we also need to reset the children of the current node
             # this is done by calling the reset method
             next_node.reset()
+        else:
+            next_node.remove_parent()
 
         self.search_root_node = next_node
 

gymcts/gymcts_deepcopy_wrapper.py

@@ -13,8 +13,15 @@ from gymcts.logger import log
 
 
 class DeepCopyMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
-
-
+    """
+    A wrapper for gym environments that implements the GymctsABC interface.
+    It uses deepcopys as state representation.
+    Please note that this is not the most efficient way to implement the state representation.
+    It is supposed to be used to see if your use-case works well with the MCTS algorithm.
+    If it does, you can consider implementing all GymctsABC methods in a more efficient way.
+    """
+
+    # helper attributes for the wrapper
     _terminal_flag:bool = False
     _last_reward: SupportsFloat = 0
     _step_tuple: tuple[WrapperObsType, SupportsFloat, bool, bool, dict[str, Any]] = None
@@ -22,9 +29,21 @@ class DeepCopyMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
     _action_mask_fn: Callable[[gym.Env], np.ndarray] | None = None
 
     def is_terminal(self) -> bool:
+        """
+        Returns True if the environment is in a terminal state, False otherwise.
+
+        :return: True if the environment is in a terminal state, False otherwise.
+        """
         return self._terminal_flag
 
     def load_state(self, state: Any) -> None:
+        """
+        The load_state method is not implemented. The state is loaded by replacing the env with the 'state' (the copy
+        provided my 'get_state'). 'self' in a method cannot be replaced with another object (as far as i know).
+
+        :param state: a deepcopy of the environment
+        :return: None
+        """
         msg = """
         The NaiveSoloMCTSGymEnvWrapper uses deepcopies of the entire env as the state.
         The loading of the state is done by replacing the env with the 'state' (the copy provided my 'get_state').
@@ -39,6 +58,16 @@ class DeepCopyMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
                  buffer_length: int = 100,
                  record_video: bool = False,
                  ):
+        """
+        The constructor of the wrapper. It wraps the environment with RecordEpisodeStatistics and checks if the action
+        space is discrete. It also checks if the action_mask_fn is a string or a callable. If it is a string, it tries to
+        find the method in the environment. If it is a callable, it assigns it to the _action_mask_fn attribute.
+
+        :param env: the environment to wrap
+        :param action_mask_fn:
+        :param buffer_length:
+        :param record_video:
+        """
         # wrap with RecordEpisodeStatistics if it is not already wrapped
         env = RecordEpisodeStatistics(env, buffer_length=buffer_length)
 
@@ -61,6 +90,10 @@ class DeepCopyMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
                 self._action_mask_fn = action_mask_fn
 
     def get_state(self) -> Any:
+        """
+        Returns the current state of the environment as a deepcopy of the environment.
+        :return: a deepcopy of the environment
+        """
         log.debug("getting state")
         original_state = self
         copied_state = copy.deepcopy(self)
@@ -71,9 +104,19 @@ class DeepCopyMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
         return copied_state
 
     def action_masks(self) -> np.ndarray | None:
+        """
+        Returns the action masks for the environment. If the action_mask_fn is not set, it returns None.
+        :return: the action masks for the environment
+        """
         return self._action_mask_fn(self.env) if self._action_mask_fn is not None else None
 
     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
+        """
         if self._action_mask_fn is None:
             action_space: gym.spaces.Discrete = self.env.action_space  # Type hinting
             return list(range(action_space.n))
@@ -83,6 +126,14 @@ class DeepCopyMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
     def step(
             self, action: WrapperActType
     ) -> tuple[WrapperObsType, SupportsFloat, bool, bool, dict[str, Any]]:
+        """
+        Performs a step in the environment.
+        This method is used to update the wrapper with the new state and the new action, to realize the terminal state
+        functionality.
+
+        :param action: action to perform in the environment
+        :return: the step tuple of the environment (obs, reward, terminated, truncated, info)
+        """
         step_tuple = self.env.step(action)
 
         obs, reward, terminated, truncated, info = step_tuple
@@ -93,6 +144,12 @@ class DeepCopyMCTSGymEnvWrapper(GymctsABC, gym.Wrapper):
 
 
     def rollout(self) -> float:
+        """
+        Performs a rollout from the current state of the environment and returns the return (sum of rewards) of the
+        rollout.
+
+        :return: the return of the rollout
+        """
         log.debug("performing rollout")
         # random rollout
         # perform random valid action util terminal

gymcts/gymcts_distributed_agent.py

@@ -118,6 +118,7 @@ class DistributedGymctsAgent:
                  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
                  ):
@@ -134,6 +135,7 @@ class DistributedGymctsAgent:
         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,
@@ -206,6 +208,8 @@ class DistributedGymctsAgent:
                 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()
@@ -217,22 +221,34 @@ class DistributedGymctsAgent:
                 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()
+        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:
@@ -268,7 +284,7 @@ if __name__ == '__main__':
     agent1 = DistributedGymctsAgent(
         env=env,
         render_tree_after_step=True,
-        number_of_simulations_per_step=1000,
+        number_of_simulations_per_step=10,
         exclude_unvisited_nodes_from_render=True,
         num_parallel=1,
     )
@@ -278,4 +294,6 @@ if __name__ == '__main__':
     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

@@ -1,28 +1,71 @@
 from typing import TypeVar, Any, SupportsFloat, Callable
 from abc import ABC, abstractmethod
 import gymnasium as gym
-
-TSoloMCTSNode = TypeVar("TSoloMCTSNode", bound="SoloMCTSNode")
+import numpy as np
 
 
 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 action_masks(self) -> np.ndarray | None:
+        """
+        Returns a numpy array of action masks for the environment. The array should have the same length as the number
+        of actions in the action space. If an action is valid, the corresponding mask value should be 1, otherwise 0.
+        If no action mask is available, it should return None.
+
+        :return: a numpy array of action masks or None
+        """
         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