gymcts 1.0.0__py3-none-any.whl

gymcts/colorful_console_utils.py

@@ -0,0 +1,142 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+CEND = "\33[0m"
+CBOLD = "\33[1m"
+CITALIC = "\33[3m"
+CURL = "\33[4m"
+CBLINK = "\33[5m"
+CBLINK2 = "\33[6m"
+CSELECTED = "\33[7m"
+
+CBLACK = "\33[30m"
+CRED = "\33[31m"
+CGREEN = "\33[32m"
+CYELLOW = "\33[33m"
+CBLUE = "\33[34m"
+CCYAN = '\33[96m'
+CMAGENTA = '\033[35m'
+CVIOLET = "\33[35m"
+CBEIGE = "\33[36m"
+CWHITE = "\33[37m"
+
+CBLACKBG = "\33[40m"
+CREDBG = "\33[41m"
+CGREENBG = "\33[42m"
+CYELLOWBG = "\33[43m"
+CBLUEBG = "\33[44m"
+CVIOLETBG = "\33[45m"
+CBEIGEBG = "\33[46m"
+CWHITEBG = "\33[47m"
+
+CGREY = "\33[90m"
+CRED2 = "\33[91m"
+CGREEN2 = "\33[92m"
+CYELLOW2 = "\33[93m"
+CBLUE2 = "\33[94m"
+CCYAN2 = "\033[36m"
+CVIOLET2 = "\33[95m"
+CBEIGE2 = "\33[96m"
+CWHITE2 = "\33[97m"
+
+CGREYBG = "\33[100m"
+CREDBG2 = "\33[101m"
+CGREENBG2 = "\33[102m"
+CYELLOWBG2 = "\33[103m"
+CBLUEBG2 = "\33[104m"
+CVIOLETBG2 = "\33[105m"
+CBEIGEBG2 = "\33[106m"
+CWHITEBG2 = "\33[107m"
+
+
+def rgb_color_sequence(r: int | float, g: int | float, b: int | float,
+                       *, format_type: str = 'foreground') -> str:
+    """
+    generates a color-codes, that change the color of text in console outputs.
+
+    rgb values must be numbers between 0 and 255 or 0.0 and 1.0.
+
+    :param r:               red value.
+    :param g:               green value
+    :param b:               blue value
+
+    :param format_type:     specifies weather the foreground-color or the background-color shall be adjusted.
+                            valid options: 'foreground','background'
+    :return:                a string that contains the color-codes.
+    """
+    # type: ignore # noqa: F401
+    if format_type == 'foreground':
+        f = '\033[38;2;{};{};{}m'.format  # font rgb format
+    elif format_type == 'background':
+        f = '\033[48;2;{};{};{}m'.format  # font background rgb format
+    else:
+        raise ValueError(f"format {format_type} is not defined. Use 'foreground' or 'background'.")
+    rgb = [r, g, b]
+
+    if isinstance(r, int) and isinstance(g, int) and isinstance(b, int):
+        if min(rgb) < 0 and max(rgb) > 255:
+            raise ValueError("rgb values must be numbers between 0 and 255 or 0.0 and 1.0")
+        return f(r, g, b)
+    if isinstance(r, float) and isinstance(g, float) and isinstance(b, float):
+        if min(rgb) < 0 and max(rgb) > 1:
+            raise ValueError("rgb values must be numbers between 0 and 255 or 0.0 and 1.0")
+        return f(*[int(n * 255) for n in [r, g, b]])
+
+
+def wrap_with_color_codes(s: object, /, r: int | float, g: int | float, b: int | float, **kwargs) \
+        -> str:
+    """
+    stringify an object and wrap it with console color codes. It adds the color control sequence in front and one
+    at the end that resolves the color again.
+
+    rgb values must be numbers between 0 and 255 or 0.0 and 1.0.
+
+    :param s: the object to stringify and wrap
+    :param r: red value.
+    :param g: green value.
+    :param b: blue value.
+    :param kwargs: additional argument for the 'DisjunctiveGraphJspVisualizer.rgb_color_sequence'-method.
+    :return:
+    """
+    return f"{rgb_color_sequence(r, g, b, **kwargs)}" \
+           f"{s}" \
+           f"{CEND}"
+
+
+
+def wrap_evenly_spaced_color(s: str, 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
+
+    c_map = plt.cm.get_cmap(c_map)  # select the desired cmap
+    arr = np.linspace(0, 1, n_classes + 1)  # create a list with numbers from 0 to 1 with n items
+
+    color_vals = c_map(arr[n_of_item])[:-1]
+    color_asni = rgb_color_sequence(*color_vals, format_type='foreground')
+
+    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:
+    if s is None or min_val is None or max_val is None or min_val >= max_val:
+        return s
+
+    if c_map is not None:
+        c_map = plt.cm.get_cmap(c_map)  # select the desired cmap
+    else:
+        from matplotlib.colors import LinearSegmentedColormap
+        colors = [
+            np.array([255 / 255, 100 / 255, 128 / 255, 1.0]),  # RGBA values
+            np.array([63 / 255, 197 / 255, 161 / 255, 1.0]),  # RGBA values
+        ]
+        c_map = LinearSegmentedColormap.from_list("custom_cmap", colors, N=256)
+
+    color_vals = c_map((value - min_val) / (max_val - min_val))[:-1]
+    color_asni = rgb_color_sequence(*color_vals, format_type='foreground')
+
+    return f"{color_asni}{s}{CEND}"
+
+
+if __name__ == '__main__':
+    res = wrap_with_color_scale("test", 1.0, 0, 1)
+    print(res)

gymcts/gymcts_agent.py

@@ -0,0 +1,261 @@
+import copy
+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.logger import log
+
+TSoloMCTSNode = TypeVar("TSoloMCTSNode", bound="SoloMCTSNode")
+
+
+class SoloMCTSAgent:
+    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!
+    clear_mcts_tree_after_step: bool
+
+    def __init__(self,
+                 env: SoloMCTSGymEnv,
+                 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
+                 ):
+        # 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.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 = SoloMCTSNode(
+            action=None,
+            parent=None,
+            env_reference=env,
+        )
+
+    def navigate_to_leaf(self, from_node: SoloMCTSNode) -> SoloMCTSNode:
+        log.debug(f"Navigate to leaf. from_node: {from_node}")
+        if from_node.terminal:
+            log.debug("Node is terminal. Returning from_node")
+            return from_node
+        if from_node.is_leaf():
+            log.debug("Node is leaf. Returning from_node")
+            return from_node
+
+        temp_node = from_node
+        # 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())
+        log.debug(f"Selected leaf node: {temp_node}")
+        return temp_node
+
+    def expand_node(self, node: SoloMCTSNode) -> None:
+        log.debug(f"expanding node: {node}")
+        # EXPANSION STRATEGY
+        # expand all children
+
+        child_dict = {}
+        for action in node.valid_actions:
+            # reconstruct state
+            # load state of leaf node
+            self._load_state(node)
+
+            obs, reward, terminal, truncated, _ = self.env.step(action)
+            child_dict[action] = SoloMCTSNode(
+                action=action,
+                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]:
+
+        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 {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 _load_state(self, node: SoloMCTSNode) -> None:
+        if isinstance(self.env, NaiveSoloMCTSGymEnvWrapper):
+            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]:
+
+        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
+
+        action = self.vanilla_mcts_search(
+            search_start_node=search_start_node,
+            num_simulations=num_simulations,
+        )
+        next_node = search_start_node.children[action]
+
+        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 vanilla_mcts_search(self, search_start_node: SoloMCTSNode = 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
+
+        for i in range(num_simulations):
+            log.debug(f"simulation {i}")
+            # navigate to leaf
+            leaf_node = self.navigate_to_leaf(from_node=search_start_node)
+
+            if leaf_node.visit_count > 0 and not leaf_node.terminal:
+                # expand leaf
+                self.expand_node(leaf_node)
+                leaf_node = leaf_node.get_random_child()
+
+            # load state of leaf node
+            self._load_state(leaf_node)
+
+            # rollout
+            episode_return = self.env.rollout()
+            # self.env.render()
+
+            self.backpropagation(node=leaf_node, episode_return=episode_return)
+
+        if self.render_tree_after_step:
+            self.show_mcts_tree()
+
+        return search_start_node.get_best_action()
+
+    def show_mcts_tree(self, start_node: SoloMCTSNode = 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 self._generate_mcts_tree(start_node=start_node, depth=tree_max_depth):
+            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:
+        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)
+            node.mean_value = node.mean_value + (episode_return - node.mean_value) / (node.visit_count + 1)
+            node.visit_count += 1
+            node.max_value = max(node.max_value, episode_return)
+            node.min_value = min(node.min_value, episode_return)
+            node = node.parent
+        # also update root node
+        # node.mean_value = ((node.mean_value * node.visit_count) + episode_return) / (node.visit_count + 1)
+        node.mean_value = node.mean_value + (episode_return - node.mean_value) / (node.visit_count + 1)
+        node.visit_count += 1
+        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_deterministic_wrapper.py

@@ -0,0 +1,107 @@
+import random
+import copy
+
+import numpy as np
+from typing import TypeVar, Any, SupportsFloat, Callable
+import gymnasium as gym
+from gymnasium.core import WrapperActType, WrapperObsType
+from gymnasium.wrappers import RecordEpisodeStatistics
+
+from gymcts.gymcts_gym_env import SoloMCTSGymEnv
+
+from gymcts.logger import log
+
+
+class DeterministicSoloMCTSGymEnvWrapper(SoloMCTSGymEnv, gym.Wrapper):
+    _terminal_flag: bool = False
+    _last_reward: SupportsFloat = 0
+    _step_tuple: tuple[WrapperObsType, SupportsFloat, bool, bool, dict[str, Any]] = None
+
+    _action_mask_fn: Callable[[gym.Env], np.ndarray] | None = None
+
+    def __init__(
+            self,
+            env,
+            action_mask_fn: str | Callable[[gym.Env], np.ndarray] | None = None,
+            buffer_length: int = 100,
+    ):
+        # wrap with RecordEpisodeStatistics if it is not already wrapped
+        env = RecordEpisodeStatistics(env, buffer_length=buffer_length)
+
+        gym.Wrapper.__init__(self, env)
+
+        self._wrapper_action_history = []
+
+        # assert that the action space is discrete
+        if not isinstance(env.action_space, gym.spaces.Discrete):
+            raise ValueError("Only discrete action spaces are supported.")
+
+        if action_mask_fn is not None:
+            # copy of stable baselines3 contrib implementation
+            if isinstance(action_mask_fn, str):
+                found_method = getattr(self.env, action_mask_fn)
+                if not callable(found_method):
+                    raise ValueError(f"Environment attribute {action_mask_fn} is not a method")
+
+                self._action_mask_fn = found_method
+            else:
+                self._action_mask_fn = action_mask_fn
+
+    def load_state(self, state: list[int]) -> None:
+        self.env.reset()
+        self._wrapper_action_history = []
+
+        for action in state:
+            self.env.step(action)
+            self._wrapper_action_history.append(action)
+
+    def is_terminal(self) -> bool:
+        if not len(self.get_valid_actions()):
+            return True
+        else:
+            return self._terminal_flag
+
+    def action_masks(self) -> np.ndarray | None:
+        return self._action_mask_fn(self.env) if self._action_mask_fn is not None else None
+
+    def get_valid_actions(self) -> list[int]:
+        if self._action_mask_fn is None:
+            action_space: gym.spaces.Discrete = self.env.action_space  # Type hinting
+            return list(range(action_space.n))
+        else:
+            return [i for i, mask in enumerate(self.action_masks()) if mask]
+
+    def rollout(self) -> float:
+        log.debug("performing rollout")
+        # random rollout
+        # perform random valid action util terminal
+        is_terminal_state = self.is_terminal()
+
+        if is_terminal_state:
+            _, _, _, _, info = self._step_tuple
+            episode_return = info["episode"]["r"]
+            return episode_return
+
+        while not is_terminal_state:
+            action = random.choice(self.get_valid_actions())
+            # print(f"Valid actions: {self.get_valid_actions()}, selected action: {action}")
+            _obs, _reward, is_terminal_state, _truncated, info = self.step(action)
+
+        episode_return = info["episode"]["r"]
+        log.debug(f"Rollout return: {episode_return}")
+        return episode_return
+
+    def get_state(self) -> list[int]:
+        return self._wrapper_action_history.copy()
+
+    def step(
+            self, action: WrapperActType
+    ) -> tuple[WrapperObsType, SupportsFloat, bool, bool, dict[str, Any]]:
+        step_tuple = self.env.step(action)
+        self._wrapper_action_history.append(action)
+        obs, reward, terminated, truncated, info = step_tuple
+
+        self._terminal_flag = terminated or truncated
+        self._step_tuple = step_tuple
+
+        return step_tuple

gymcts/gymcts_gym_env.py

@@ -0,0 +1,28 @@
+from typing import TypeVar, Any, SupportsFloat, Callable
+from abc import ABC, abstractmethod
+import gymnasium as gym
+
+TSoloMCTSNode = TypeVar("TSoloMCTSNode", bound="SoloMCTSNode")
+
+
+class SoloMCTSGymEnv(ABC, gym.Env):
+
+    @abstractmethod
+    def get_state(self) -> Any:
+        pass
+
+    @abstractmethod
+    def load_state(self, state: Any) -> None:
+        pass
+
+    @abstractmethod
+    def is_terminal(self) -> bool:
+        pass
+
+    @abstractmethod
+    def get_valid_actions(self) -> list[int]:
+        pass
+
+    @abstractmethod
+    def rollout(self) -> float:
+        pass

gymcts/gymcts_naive_wrapper.py

@@ -0,0 +1,114 @@
+import random
+import copy
+
+import numpy as np
+from typing import TypeVar, Any, SupportsFloat, Callable
+import gymnasium as gym
+from gymnasium.core import WrapperActType, WrapperObsType
+from gymnasium.wrappers import RecordEpisodeStatistics
+
+from gymcts.gymcts_gym_env import SoloMCTSGymEnv
+
+from gymcts.logger import log
+
+
+class NaiveSoloMCTSGymEnvWrapper(SoloMCTSGymEnv, gym.Wrapper):
+
+
+    _terminal_flag:bool = False
+    _last_reward: SupportsFloat = 0
+    _step_tuple: tuple[WrapperObsType, SupportsFloat, bool, bool, dict[str, Any]] = None
+
+    _action_mask_fn: Callable[[gym.Env], np.ndarray] | None = None
+
+    def is_terminal(self) -> bool:
+        return self._terminal_flag
+
+    def load_state(self, state: Any) -> 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').
+        'self' in a method cannot be replaced with another object (as far as i know). Therefore the copy is done by
+        MCTSaAgent here.
+        """
+        raise NotImplementedError(msg)
+
+    def __init__(self,
+                 env,
+                 action_mask_fn: str | Callable[[gym.Env], np.ndarray] | None = None,
+                 buffer_length: int = 100,
+                 record_video: bool = False,
+                 ):
+        # wrap with RecordEpisodeStatistics if it is not already wrapped
+        env = RecordEpisodeStatistics(env, buffer_length=buffer_length)
+
+        gym.Wrapper.__init__(self, env)
+        # super().__init__(env)
+
+        # assert that the action space is discrete
+        if not isinstance(env.action_space, gym.spaces.Discrete):
+            raise ValueError("Only discrete action spaces are supported.")
+
+        if action_mask_fn is not None:
+            # copy of stable baselines3 contrib implementation
+            if isinstance(action_mask_fn, str):
+                found_method = getattr(self.env, action_mask_fn)
+                if not callable(found_method):
+                    raise ValueError(f"Environment attribute {action_mask_fn} is not a method")
+
+                self._action_mask_fn = found_method
+            else:
+                self._action_mask_fn = action_mask_fn
+
+    def get_state(self) -> Any:
+        log.debug("getting state")
+        original_state = self
+        copied_state = copy.deepcopy(self)
+
+        log.debug(f"original state memory location: {hex(id(original_state))}")
+        log.debug(f"copied memory location: {hex(id(copied_state))}")
+
+        return copied_state
+
+    def action_masks(self) -> np.ndarray | None:
+        return self._action_mask_fn(self.env) if self._action_mask_fn is not None else None
+
+    def get_valid_actions(self) -> list[int]:
+        if self._action_mask_fn is None:
+            action_space: gym.spaces.Discrete = self.env.action_space  # Type hinting
+            return list(range(action_space.n))
+        else:
+            return [i for i, mask in enumerate(self.action_masks()) if mask]
+
+    def step(
+            self, action: WrapperActType
+    ) -> tuple[WrapperObsType, SupportsFloat, bool, bool, dict[str, Any]]:
+        step_tuple = self.env.step(action)
+
+        obs, reward, terminated, truncated, info = step_tuple
+        self._terminal_flag = terminated or truncated
+        self._step_tuple = step_tuple
+
+        return step_tuple
+
+
+    def rollout(self) -> float:
+        log.debug("performing rollout")
+        # random rollout
+        # perform random valid action util terminal
+        is_terminal_state = self.is_terminal()
+
+        if is_terminal_state:
+            _, _, _, _, info = self._step_tuple
+            episode_return = info["episode"]["r"]
+            return episode_return
+
+        while not is_terminal_state:
+            action = random.choice(self.get_valid_actions())
+            # print(f"Valid actions: {self.get_valid_actions()}, selected action: {action}")
+            _obs, _reward, is_terminal_state, _truncated, info = self.step(action)
+
+
+        episode_return = info["episode"]["r"]
+        log.debug(f"Rollout return: {episode_return}")
+        return episode_return