A library for multi-agent learning that aims to accelerate the research.
"""A research library for multi-agent learning."""


# Import all submodules.
from coaction import agents
from coaction import experiments
from coaction import games
from coaction import loggers
from coaction import utils


__all__ = ["agents", "experiments", "games", "loggers", "utils"]

coaction/__init__.py
yigityalin/coaction
"""This module contains the base and ready-to-use classes for agents."""

from coaction.agents.agent import Agent
from coaction.agents import fictitious_play
from coaction.agents.fictitious_play.async_fp import AsynchronousFictitiousPlay
from coaction.agents.fictitious_play.async_sfp import AsynchronousSmoothedFictitiousPlay
from coaction.agents.fictitious_play.sync_fp import SynchronousFictitiousPlay
from coaction.agents.fictitious_play.sync_sfp import SynchronousSmoothedFictitiousPlay
from coaction.agents.fictitious_play.model_free_fp import ModelFreeFictitiousPlay
from coaction.agents.fictitious_play.model_free_sfp import (
    ModelFreeSmoothedFictitiousPlay,
)
from coaction.agents.ind_q.ind_q import IndividualQLearning


__all__ = [
    "Agent",
    "fictitious_play",
    "AsynchronousFictitiousPlay",
    "AsynchronousSmoothedFictitiousPlay",
    "SynchronousFictitiousPlay",
    "SynchronousSmoothedFictitiousPlay",
    "ModelFreeFictitiousPlay",
    "ModelFreeSmoothedFictitiousPlay",
    "IndividualQLearning",
]

coaction/agents/__init__.py
yigityalin/coaction
"""Implements the base class for agents."""

from abc import ABC, abstractmethod
from collections.abc import Collection, Sequence
from copy import deepcopy
from typing import Any
import pickle

import numpy as np

from coaction.games.game import StateType, ActionType, RewardType


ParameterType = Any


class Agent(ABC):
    """Base class for all agents."""

    def __init__(
        self, name: str, seed: int, logged_params: Collection[str] = None, **kwargs
    ):
        """Initialize the agent.

        Args:
            name (str): The name of the agent.
            seed (int): The seed for the random number generator.
            logged_params (dict[str, Any]): The parameters to log. If None, default agent
                configuration is used. If the agent does not override `loggable_params`,
                this parameter is ignored and the default agent configuration is used.
                This is useful for agents that are implemented in coaction. For custom
                agents, it is suggested to implement the agent according to the
                __setattr__ method, i.e., all the params except the logged ones should
                be registered as buffers via 'register_buffer' or via adding '_' before
                the variable name.
            **kwargs: Keyword arguments.
        """
        _ = kwargs  # ignore unused kwargs
        self._name: str = name
        self._seed: int = seed
        self._rng = None
        self._buffer = {}
        self._logged_params = {}

        if self.loggable_params is NotImplemented:
            logged_params = None
        if logged_params:
            self._logged_params_names = set(logged_params) & self.loggable_params
        else:
            self._logged_params_names = None

    def __repr__(self):
        return f"{self.__class__.__name__}(name={self.name})"

    def __setattr__(self, key: str, value: Any):
        """Set an attribute.

        If the attribute name starts with an underscore, it is set as a normal attribute.
        Otherwise, it is registered as a parameter.

        In order to register a parameter invisible to loggers, use `register_buffer` instead.
        """
        logged_params = (
            self._logged_params_names if hasattr(self, "_logged_params_names") else None
        )
        if (logged_params is None and key.startswith("_")) or (
            logged_params is not None and key.lstrip("_") not in logged_params
        ):
            super().__setattr__(key, value)
        else:
            self.register_parameter(key, value)

    @property
    def name(self) -> str:
        """Return the name of the agent."""
        return self._name

    @property
    def seed(self) -> int:
        """Return the seed."""
        return self._seed

    @property
    def rng(self):
        """Return the random number generator."""
        if self._rng is None:
            self._rng = np.random.default_rng(self.seed)
        return self._rng

    @property
    def loggable_params(self) -> set[str]:
        """Return the agent parameters that can be logged."""
        return NotImplemented

    @property
    def params(self) -> dict[str, Any]:
        """Return the agent parameters."""
        return {**self._logged_params, **self._buffer}

    @property
    def logged_params(self) -> dict[str, Any]:
        """Return the agent parameters that can be logged."""
        return self._logged_params

    @property
    def buffer(self) -> dict[str, ParameterType]:
        """Return the agent parameters that cannot be logged."""
        return self._buffer

    @abstractmethod
    def reset(self):
        """Reset the agent's parameters."""
        self._rng = None

    @abstractmethod
    def act(self, state: StateType) -> ActionType:
        """Return the action to take given the current state.

        Args:
            state (StateType): The current state of the game.
        """

    @abstractmethod
    def update(
        self,
        state: StateType,
        actions: Sequence[ActionType],
        reward: RewardType,
        next_state: StateType,
        **kwargs,
    ):
        """Update the agent's parameters.

        Args:
            state (StateType): The current state of the game.
            actions (Sequence[ActionType]): The actions taken by the agents.
            reward (RewardType): The reward received by the agents.
            next_state (StateType): The next state of the game.
            **kwargs: Keyword arguments.
        """

    def clone(self, seed: int = None):
        """Return a deep copy of the agent."""
        agent = deepcopy(self)
        agent._seed = seed  # pylint: disable=protected-access
        return agent

    def save(self, path: str):
        """Save the agent to a file.

        Args:
            path (str): The path to the file.
        """
        with open(path, "wb") as file:
            pickle.dump(self, file)

    @staticmethod
    def load(path: str):
        """Load an agent from a file.

        Args:
            path (str): The path to the file.
        """
        with open(path, "rb") as file:
            return pickle.load(file)

    def register_buffer(self, name: str, value: ParameterType):
        """Register a parameter which will not be seen by loggers but will be saved.

        This method is useful for registering parameters that are not tracked by loggers.

        Args:
            name (str): The name of the parameter.
            value (Any): The value of the parameter.
        """
        if name in self._logged_params:
            raise ValueError(
                f"Parameter {name} is already registered as a logged parameter"
            )
        self._buffer[name] = value
        super().__setattr__(name, value)

    def register_parameter(self, name: str, value: ParameterType):
        """Register a parameter visible to loggers.

        This method is useful for registering parameters that are tracked by loggers.

        Args:
            name (str): The name of the parameter.
            value (Any): The value of the parameter.
        """
        if name in self._buffer:
            self.register_buffer(name, value)
        else:
            self._logged_params[name] = value
            super().__setattr__(name, value)


class TwoPlayerAgent(Agent):
    """Base class for agents in two-player games.

    This class is a convenience class for agents in two-player games.
    It does not add any functionality over the base `Agent` class.
    """


class MultiPlayerAgent(Agent):
    """Base class for agents in multi-player games.

    This class is a convenience class for agents in multi-player games.
    It does not add any functionality over the base `Agent` class.
    """

coaction/agents/agent.py
yigityalin/coaction
"""Fictitious play agents."""

from coaction.agents.fictitious_play.async_fp import AsynchronousFictitiousPlay
from coaction.agents.fictitious_play.async_sfp import AsynchronousSmoothedFictitiousPlay
from coaction.agents.fictitious_play.sync_fp import SynchronousFictitiousPlay
from coaction.agents.fictitious_play.sync_sfp import SynchronousSmoothedFictitiousPlay
from coaction.agents.fictitious_play.model_free_fp import ModelFreeFictitiousPlay
from coaction.agents.fictitious_play.model_free_sfp import (
    ModelFreeSmoothedFictitiousPlay,
)


__all__ = [
    "AsynchronousFictitiousPlay",
    "AsynchronousSmoothedFictitiousPlay",
    "SynchronousFictitiousPlay",
    "SynchronousSmoothedFictitiousPlay",
    "ModelFreeFictitiousPlay",
    "ModelFreeSmoothedFictitiousPlay",
]

coaction/agents/fictitious_play/__init__.py
yigityalin/coaction
"""Implementation of the Fictitious Play algorithm."""

from collections.abc import Callable, Collection
from typing import Final, Sequence

import numpy as np
import numpy.typing as npt

from coaction.agents.agent import TwoPlayerAgent
from coaction.agents.fictitious_play import utils as fp_utils
from coaction.games.game import ActionType, RewardType, StateType


class AsynchronousFictitiousPlay(TwoPlayerAgent):
    """Implementation of the Fictitious Play algorithm.

    Parameters
    ----------
    name : str
        The name of the agent.
    seed : int
        The seed for the random number generator.
    reward_matrix : npt.NDArray[RewardType]
        The reward matrix.
    transition_matrix : npt.NDArray[np.float_]
        The transition matrix.
    alpha : Callable[[int], float]
        Step size for beliefs.
    beta : Callable[[int], float]
        Step size for Q function.
    gamma : float
        Discount factor
    """

    def __init__(
        self,
        name: str,
        seed: int,
        reward_matrix: npt.NDArray[RewardType],
        transition_matrix: npt.NDArray[np.float_],
        alpha: Callable[[int], float],
        beta: Callable[[int], float],
        gamma: float,
        initial_Q: None | int | float | np.ndarray = None,
        initial_pi: None | int | float | np.ndarray = None,
        logged_params: Collection[str] = None,
        **kwargs,
    ):
        """Initialize the agent.

        Args:
            name (str): The name of the agent.
            seed (int): The seed for the random number generator.
            reward_matrix (npt.NDArray[RewardType]): The reward matrix.
            transition_matrix (npt.NDArray[np.float_]): The transition matrix.
            alpha (Callable[[int], float]): Step size for beliefs.
            beta (Callable[[int], float]): Step size for Q function.
            gamma (float): Discount factor
            initial_Q (None | int | float | np.ndarray): The initial Q matrix.
            initial_pi (None | int | float | np.ndarray): The initial pi matrix.
            logged_params (Collection[str]): The parameters to log.
        """
        super().__init__(name, seed, logged_params, **kwargs)
        self._R = reward_matrix  # pylint: disable=invalid-name
        self._T = transition_matrix  # pylint: disable=invalid-name

        self._alpha = alpha
        self._beta = beta
        self._gamma = gamma
        self._initial_Q = (  # pylint: disable=invalid-name
            reward_matrix if initial_Q is None else initial_Q
        )
        self._initial_pi = initial_pi

        self._n_states = reward_matrix.shape[0]
        self._n_actions = reward_matrix.shape[1]
        self._n_opponent_actions = reward_matrix.shape[2]

        self._counts = np.zeros(self._n_states, dtype=np.int_)
        self._Q = fp_utils.get_initial_Q(  # pylint: disable=invalid-name
            self._initial_Q, self._n_states, self._n_actions, self._n_opponent_actions
        )
        self._pi = fp_utils.get_initial_pi(  # pylint: disable=invalid-name
            self._initial_pi, self._n_states, self._n_opponent_actions, self.rng
        )
        self.v = np.zeros(  # pylint: disable=invalid-name
            self._n_states, dtype=np.float_
        )

        self._utils = fp_utils.FictitiousPlayUtils(self._Q, self._pi, self.rng)

        self._eye: Final = np.eye(reward_matrix.shape[-1])

    @property
    def loggable_params(self) -> set[str]:
        return fp_utils.LOGGABLE_PARAMS

    def reset(self):
        super().reset()
        self._counts = np.zeros(self._n_states, dtype=np.int_)
        self._Q = fp_utils.get_initial_Q(
            self._initial_Q, self._n_states, self._n_actions, self._n_opponent_actions
        )
        self._pi = fp_utils.get_initial_pi(
            self._initial_pi, self._n_states, self._n_opponent_actions, self.rng
        )
        self.v = np.zeros(  # pylint: disable=invalid-name
            self._n_states, dtype=np.float_
        )
        self._utils.reset(self._Q, self._pi)

    def act(self, state: StateType) -> ActionType:
        return self._utils.best_response(state)

    def update(
        self,
        state: StateType,
        actions: Sequence[ActionType],
        reward: RewardType,
        next_state: StateType,
        **kwargs,
    ):
        self._pi[state] += fp_utils.belief_update(
            self._pi[state],
            self._alpha(self._counts[state]),
            actions[1],  # type: ignore
        )
        self._Q[state] += fp_utils.model_based_sync_q_update(
            self._Q[state],
            self.v,
            self._beta(self._counts[state]),
            self._gamma,
            self._R[state],
            self._T[state],
        )
        self.v = self._utils.value_function(self._Q, self._pi)
        self._counts[state] += 1

coaction/agents/fictitious_play/async_fp.py
yigityalin/coaction
"""Implementation of the Fictitious Play algorithm."""

from collections.abc import Callable, Collection
from typing import Final, Sequence

import numpy as np
import numpy.typing as npt

from coaction.agents.agent import TwoPlayerAgent
from coaction.agents.fictitious_play import utils as fp_utils
from coaction.games.game import ActionType, RewardType, StateType


class AsynchronousSmoothedFictitiousPlay(TwoPlayerAgent):
    """Implementation of the Fictitious Play algorithm.

    Parameters
    ----------
    name : str
        The name of the agent.
    seed : int
        The seed for the random number generator.
    reward_matrix : npt.NDArray[RewardType]
        The reward matrix.
    transition_matrix : npt.NDArray[np.float_]
        The transition matrix.
    alpha : Callable[[int], float]
        Step size for beliefs.
    beta : Callable[[int], float]
        Step size for Q function.
    gamma : float
        Discount factor
    tau : float
        Temperature parameter
    """

    def __init__(
        self,
        name: str,
        seed: int,
        reward_matrix: npt.NDArray[RewardType],
        transition_matrix: npt.NDArray[np.float_],
        alpha: Callable[[int], float],
        beta: Callable[[int], float],
        gamma: float,
        tau: float,
        initial_Q: None | int | float | np.ndarray = None,
        initial_pi: None | int | float | np.ndarray = None,
        logged_params: Collection[str] = None,
        **kwargs,
    ):
        """Initialize the agent.

        Args:
            name (str): The name of the agent.
            seed (int): The seed for the random number generator.
            reward_matrix (npt.NDArray[RewardType]): The reward matrix.
            transition_matrix (npt.NDArray[np.float_]): The transition matrix.
            alpha (Callable[[int], float]): Step size for beliefs.
            beta (Callable[[int], float]): Step size for Q function.
            gamma (float): Discount factor
            tau (float): Temperature parameter
            initial_Q (None | int | float | np.ndarray): The initial Q matrix.
            initial_pi (None | int | float | np.ndarray): The initial pi matrix.
            logged_params (Collection[str]): The parameters to log.
        """
        super().__init__(name, seed, logged_params, **kwargs)
        self._R = reward_matrix  # pylint: disable=invalid-name
        self._T = transition_matrix  # pylint: disable=invalid-name

        self._alpha = alpha
        self._beta = beta
        self._gamma = gamma
        self._tau = tau
        self._initial_Q = (  # pylint: disable=invalid-name
            reward_matrix if initial_Q is None else initial_Q
        )
        self._initial_pi = initial_pi

        self._n_states = reward_matrix.shape[0]
        self._n_actions = reward_matrix.shape[1]
        self._n_opponent_actions = reward_matrix.shape[2]

        self._counts = np.zeros(self._n_states, dtype=np.int_)
        self._Q = fp_utils.get_initial_Q(  # pylint: disable=invalid-name
            self._initial_Q, self._n_states, self._n_actions, self._n_opponent_actions
        )
        self._pi = fp_utils.get_initial_pi(  # pylint: disable=invalid-name
            self._initial_pi, self._n_states, self._n_opponent_actions, self.rng
        )
        self.v = np.zeros(  # pylint: disable=invalid-name
            self._n_states, dtype=np.float_
        )

        self._utils = fp_utils.SmoothedFictitiousPlayUtils(
            self._Q, self._pi, self._tau, self.rng
        )
        self._eye: Final = np.eye(reward_matrix.shape[-1])

    @property
    def loggable_params(self) -> set[str]:
        return fp_utils.LOGGABLE_PARAMS

    def reset(self):
        super().reset()
        self._counts = np.zeros(self._n_states, dtype=np.int_)
        self._Q = fp_utils.get_initial_Q(
            self._initial_Q, self._n_states, self._n_actions, self._n_opponent_actions
        )
        self._pi = fp_utils.get_initial_pi(
            self._initial_pi, self._n_states, self._n_opponent_actions, self.rng
        )
        self.v = np.zeros(  # pylint: disable=invalid-name
            self._n_states, dtype=np.float_
        )
        self._utils.reset(self._Q, self._pi)

    def act(self, state: StateType) -> ActionType:
        return self._utils.smoothed_best_response(state)

    def update(
        self,
        state: StateType,
        actions: Sequence[ActionType],
        reward: RewardType,
        next_state: StateType,
        **kwargs,
    ):
        self._pi[state] += fp_utils.belief_update(
            self._pi[state],
            self._alpha(self._counts[state]),
            actions[1],  # type: ignore
        )
        self._Q[state] += fp_utils.model_based_sync_q_update(
            self._Q[state],
            self.v,
            self._beta(self._counts[state]),
            self._gamma,
            self._R[state],
            self._T[state],
        )
        self.v = self._utils.smoothed_value_function(self._Q, self._pi)
        self._counts[state] += 1

coaction/agents/fictitious_play/async_sfp.py
yigityalin/coaction
"""Implementation of the Fictitious Play algorithm."""

from collections.abc import Callable, Collection
from typing import Final, Sequence

import numpy as np
import numpy.typing as npt

from coaction.agents.agent import TwoPlayerAgent
from coaction.agents.fictitious_play import utils as fp_utils
from coaction.games.game import ActionType, RewardType, StateType


class ModelFreeFictitiousPlay(TwoPlayerAgent):
    """Implementation of the Fictitious Play algorithm.

    Parameters
    ----------
    name : str
        The name of the agent.
    seed : int
        The seed for the random number generator.
    alpha : Callable[[int], float]
        Step size for beliefs.
    beta : Callable[[int], float]
        Step size for Q function.
    gamma : float
        Discount factor
    tau : float
        Temperature parameter
    """

    def __init__(
        self,
        name: str,
        seed: int,
        reward_matrix: npt.NDArray[RewardType],
        alpha: Callable[[int], float],
        beta: Callable[[int], float],
        gamma: float,
        initial_Q: None | int | float | np.ndarray = None,
        initial_pi: None | int | float | np.ndarray = None,
        logged_params: Collection[str] = None,
        **kwargs,
    ):
        """Initialize the agent.

        Args:
            name (str): The name of the agent.
            seed (int): The seed for the random number generator.
            reward_matrix (npt.NDArray[RewardType]): The reward matrix.
            transition_matrix (npt.NDArray[np.float_]): The transition matrix.
            alpha (Callable[[int], float]): Step size for beliefs.
            beta (Callable[[int], float]): Step size for Q function.
            gamma (float): Discount factor
            initial_Q (None | int | float | np.ndarray): The initial Q matrix.
            initial_pi (None | int | float | np.ndarray): The initial pi matrix.
            logged_params (Collection[str]): The parameters to log.
        """
        super().__init__(name, seed, logged_params, **kwargs)

        self._alpha = alpha
        self._beta = beta
        self._gamma = gamma
        self._initial_Q = initial_Q  # pylint: disable=invalid-name
        self._initial_pi = initial_pi

        self._n_states = reward_matrix.shape[0]
        self._n_actions = reward_matrix.shape[1]
        self._n_opponent_actions = reward_matrix.shape[2]

        self._counts = np.zeros_like(reward_matrix)
        self._Q = fp_utils.get_initial_Q(  # pylint: disable=invalid-name
            self._initial_Q, self._n_states, self._n_actions, self._n_opponent_actions
        )
        self._pi = fp_utils.get_initial_pi(  # pylint: disable=invalid-name
            self._initial_pi, self._n_states, self._n_opponent_actions, self.rng
        )
        self.v = np.zeros(  # pylint: disable=invalid-name
            self._n_states, dtype=np.float_
        )
        self._utils = fp_utils.FictitiousPlayUtils(self._Q, self._pi, self.rng)

        self._eye: Final = np.eye(reward_matrix.shape[-1])

    @property
    def loggable_params(self) -> set[str]:
        return fp_utils.LOGGABLE_PARAMS

    def reset(self):
        super().reset()
        self._counts = np.zeros_like(self._counts)
        self._Q = fp_utils.get_initial_Q(
            self._initial_Q, self._n_states, self._n_actions, self._n_opponent_actions
        )
        self._pi = fp_utils.get_initial_pi(
            self._initial_pi, self._n_states, self._n_opponent_actions, self.rng
        )
        self.v = np.zeros(  # pylint: disable=invalid-name
            self._n_states, dtype=np.float_
        )
        self._utils.reset(self._Q, self._pi)

    def act(self, state: StateType) -> ActionType:
        return self._utils.best_response(state)

    def update(
        self,
        state: StateType,
        actions: Sequence[ActionType],
        reward: RewardType,
        next_state: StateType,
        **kwargs,
    ):
        self._pi[state] += fp_utils.belief_update(
            self._pi[state], self._alpha(self._counts[state].sum()), actions[1]  # type: ignore
        )
        self._Q[state, *actions] += fp_utils.model_free_q_update(
            self._Q[state, *actions],
            self.v[state],
            self._beta(self._counts[state, *actions]),
            self._gamma,
            reward,
        )
        self.v = self._utils.value_function(self._Q, self._pi)
        self._counts[state, *actions] += 1

coaction/agents/fictitious_play/model_free_fp.py
yigityalin/coaction
"""Implementation of the Fictitious Play algorithm."""

from collections.abc import Callable, Collection
from typing import Final, Sequence

import numpy as np
import numpy.typing as npt

from coaction.agents.agent import TwoPlayerAgent
from coaction.agents.fictitious_play import utils as fp_utils
from coaction.games.game import ActionType, RewardType, StateType


class ModelFreeSmoothedFictitiousPlay(TwoPlayerAgent):
    """Implementation of the Fictitious Play algorithm.

    Parameters
    ----------
    name : str
        The name of the agent.
    seed : int
        The seed for the random number generator.
    alpha : Callable[[int], float]
        Step size for beliefs.
    beta : Callable[[int], float]
        Step size for Q function.
    gamma : float
        Discount factor
    tau : float
        Temperature parameter
    """

    def __init__(
        self,
        name: str,
        seed: int,
        reward_matrix: npt.NDArray[RewardType],
        alpha: Callable[[int], float],
        beta: Callable[[int], float],
        gamma: float,
        tau: float,
        initial_Q: None | int | float | np.ndarray = None,
        initial_pi: None | int | float | np.ndarray = None,
        logged_params: Collection[str] = None,
        **kwargs,
    ):
        """Initialize the agent.

        Args:
            name (str): The name of the agent.
            seed (int): The seed for the random number generator.
            reward_matrix (npt.NDArray[RewardType]): The reward matrix.
            alpha (Callable[[int], float]): Step size for beliefs.
            beta (Callable[[int], float]): Step size for Q function.
            gamma (float): Discount factor
            tau (float): Temperature parameter
            initial_Q (None | int | float | np.ndarray): The initial Q matrix.
            initial_pi (None | int | float | np.ndarray): The initial pi matrix.
            logged_params (Collection[str]): The parameters to log.
        """
        super().__init__(name, seed, logged_params, **kwargs)

        self._alpha = alpha
        self._beta = beta
        self._gamma = gamma
        self._tau = tau
        self._initial_Q = initial_Q  # pylint: disable=invalid-name
        self._initial_pi = initial_pi

        self._n_states = reward_matrix.shape[0]
        self._n_actions = reward_matrix.shape[1]
        self._n_opponent_actions = reward_matrix.shape[2]

        self._counts = np.zeros_like(reward_matrix)
        self._Q = fp_utils.get_initial_Q(  # pylint: disable=invalid-name
            self._initial_Q, self._n_states, self._n_actions, self._n_opponent_actions
        )
        self._pi = fp_utils.get_initial_pi(  # pylint: disable=invalid-name
            self._initial_pi, self._n_states, self._n_opponent_actions, self.rng
        )
        self.v = np.zeros(  # pylint: disable=invalid-name
            self._n_states, dtype=np.float_
        )

        self._utils = fp_utils.SmoothedFictitiousPlayUtils(
            self._Q, self._pi, self._tau, self.rng
        )

        self._t = 0
        self._eye: Final = np.eye(reward_matrix.shape[-1])

    @property
    def loggable_params(self) -> set[str]:
        return fp_utils.LOGGABLE_PARAMS

    def reset(self):
        super().reset()
        self._counts = np.zeros_like(self._counts, dtype=np.int_)
        self._Q = fp_utils.get_initial_Q(
            self._initial_Q, self._n_states, self._n_actions, self._n_opponent_actions
        )
        self._pi = fp_utils.get_initial_pi(
            self._initial_pi, self._n_states, self._n_opponent_actions, self.rng
        )
        self.v = np.zeros(  # pylint: disable=invalid-name
            self._n_states, dtype=np.float_
        )
        self._utils.reset(self._Q, self._pi)

    def act(self, state: StateType) -> ActionType:
        return self._utils.smoothed_best_response(state)

    def update(
        self,
        state: StateType,
        actions: Sequence[ActionType],
        reward: RewardType,
        next_state: StateType,
        **kwargs,
    ):
        self._pi[state] += fp_utils.belief_update(
            self._pi[state], self._alpha(self._counts[state].sum()), actions[1]  # type: ignore
        )
        self._Q[state, *actions] += fp_utils.model_free_q_update(
            self._Q[state, *actions],
            self.v[state],
            self._beta(self._counts[state, *actions]),
            self._gamma,
            reward,
        )
        self.v = self._utils.smoothed_value_function(self._Q, self._pi)
        self._counts[state, *actions] += 1

coaction/agents/fictitious_play/model_free_sfp.py
yigityalin/coaction
"""Implementation of the Fictitious Play algorithm."""

from collections.abc import Callable, Collection
from typing import Final, Sequence

import numpy as np
import numpy.typing as npt

from coaction.agents.agent import TwoPlayerAgent
from coaction.agents.fictitious_play import utils as fp_utils
from coaction.games.game import ActionType, RewardType, StateType


class SynchronousFictitiousPlay(TwoPlayerAgent):
    """Implementation of the Fictitious Play algorithm.

    Parameters
    ----------
    name : str
        The name of the agent.
    seed : int
        The seed for the random number generator.
    reward_matrix : npt.NDArray[RewardType]
        The reward matrix.
    transition_matrix : npt.NDArray[np.float_]
        The transition matrix.
    alpha : Callable[[int], float]
        Step size for beliefs.
    beta : Callable[[int], float]
        Step size for Q function.
    gamma : float
        Discount factor
    """

    def __init__(
        self,
        name: str,
        seed: int,
        reward_matrix: npt.NDArray[RewardType],
        transition_matrix: npt.NDArray[np.float_],
        alpha: Callable[[int], float],
        beta: Callable[[int], float],
        gamma: float,
        initial_Q: None | int | float | np.ndarray = None,
        initial_pi: None | int | float | np.ndarray = None,
        logged_params: Collection[str] = None,
        **kwargs,
    ):
        """Initialize the agent.

        Args:
            name (str): The name of the agent.
            seed (int): The seed for the random number generator.
            reward_matrix (npt.NDArray[RewardType]): The reward matrix.
            transition_matrix (npt.NDArray[np.float_]): The transition matrix.
            alpha (Callable[[int], float]): Step size for beliefs.
            beta (Callable[[int], float]): Step size for Q function.
            gamma (float): Discount factor
            initial_Q (None | int | float | np.ndarray): The initial Q matrix.
            initial_pi (None | int | float | np.ndarray): The initial pi matrix.
            logged_params (bool): Whether to log the parameters.
        """
        super().__init__(name, seed, logged_params, **kwargs)
        self._R = reward_matrix  # pylint: disable=invalid-name
        self._T = transition_matrix  # pylint: disable=invalid-name

        self._alpha = alpha
        self._beta = beta
        self._gamma = gamma
        self._initial_Q = (  # pylint: disable=invalid-name
            reward_matrix if initial_Q is None else initial_Q
        )
        self._initial_pi = initial_pi

        self._n_states = reward_matrix.shape[0]
        self._n_actions = reward_matrix.shape[1]
        self._n_opponent_actions = reward_matrix.shape[2]

        self._counts = np.zeros(self._n_states, dtype=np.int_)
        self._Q = fp_utils.get_initial_Q(  # pylint: disable=invalid-name
            self._initial_Q, self._n_states, self._n_actions, self._n_opponent_actions
        )
        self._pi = fp_utils.get_initial_pi(  # pylint: disable=invalid-name
            self._initial_pi, self._n_states, self._n_opponent_actions, self.rng
        )
        self.v = np.zeros(  # pylint: disable=invalid-name
            self._n_states, dtype=np.float_
        )

        self._utils = fp_utils.FictitiousPlayUtils(self._Q, self._pi, self.rng)

        self._t = 0
        self._eye: Final = np.eye(reward_matrix.shape[-1])

    @property
    def loggable_params(self) -> set[str]:
        return fp_utils.LOGGABLE_PARAMS

    def reset(self):
        super().reset()
        self._counts = np.zeros(self._n_states, dtype=np.int_)
        self._Q = fp_utils.get_initial_Q(
            self._initial_Q, self._n_states, self._n_actions, self._n_opponent_actions
        )
        self._pi = fp_utils.get_initial_pi(
            self._initial_pi, self._n_states, self._n_opponent_actions, self.rng
        )
        self.v = np.zeros(  # pylint: disable=invalid-name
            self._n_states, dtype=np.float_
        )
        self._t = 0
        self._utils.reset(self._Q, self._pi)

    def act(self, state: StateType) -> ActionType:
        return self._utils.best_response(state)

    def update(
        self,
        state: StateType,
        actions: Sequence[ActionType],
        reward: RewardType,
        next_state: StateType,
        **kwargs,
    ):
        self._pi[state] += fp_utils.belief_update(
            self._pi[state],
            self._alpha(self._counts[state]),
            actions[1],  # type: ignore
        )
        self._Q += fp_utils.model_based_sync_q_update(
            self._Q,
            self.v,
            self._beta(self._t),
            self._gamma,
            self._R,
            self._T,
        )
        self.v = self._utils.value_function(self._Q, self._pi)
        self._counts[state] += 1
        self._t += 1

coaction/agents/fictitious_play/sync_fp.py
yigityalin/coaction
"""Implementation of the Fictitious Play algorithm."""

from collections.abc import Callable, Collection
from typing import Final, Sequence

import numpy as np
import numpy.typing as npt

from coaction.agents.agent import TwoPlayerAgent
from coaction.agents.fictitious_play import utils as fp_utils
from coaction.games.game import ActionType, RewardType, StateType


class SynchronousSmoothedFictitiousPlay(TwoPlayerAgent):
    """Implementation of the Fictitious Play algorithm.

    Parameters
    ----------
    name : str
        The name of the agent.
    seed : int
        The seed for the random number generator.
    reward_matrix : npt.NDArray[RewardType]
        The reward matrix.
    transition_matrix : npt.NDArray[np.float_]
        The transition matrix.
    alpha : Callable[[int], float]
        Step size for beliefs.
    beta : Callable[[int], float]
        Step size for Q function.
    gamma : float
        Discount factor
    tau : float
        Temperature parameter
    """

    def __init__(
        self,
        name: str,
        seed: int,
        reward_matrix: npt.NDArray[RewardType],
        transition_matrix: npt.NDArray[np.float_],
        alpha: Callable[[int], float],
        beta: Callable[[int], float],
        gamma: float,
        tau: float,
        initial_Q: None | int | float | np.ndarray = None,
        initial_pi: None | int | float | np.ndarray = None,
        logged_params: Collection[str] = None,
        **kwargs,
    ):
        """Initialize the agent.

        Args:
            name (str): The name of the agent.
            seed (int): The seed for the random number generator.
            reward_matrix (npt.NDArray[RewardType]): The reward matrix.
            transition_matrix (npt.NDArray[np.float_]): The transition matrix.
            alpha (Callable[[int], float]): Step size for beliefs.
            beta (Callable[[int], float]): Step size for Q function.
            gamma (float): Discount factor
            tau (float): Temperature parameter
            initial_Q (None | int | float | np.ndarray): The initial Q matrix.
            initial_pi (None | int | float | np.ndarray): The initial pi matrix.
            logged_params (bool): Whether to log the parameters.
        """
        super().__init__(name, seed, logged_params, **kwargs)
        self._R = reward_matrix  # pylint: disable=invalid-name
        self._T = transition_matrix  # pylint: disable=invalid-name

        self._alpha = alpha
        self._beta = beta
        self._gamma = gamma
        self._tau = tau
        self._initial_Q = (  # pylint: disable=invalid-name
            reward_matrix if initial_Q is None else initial_Q
        )
        self._initial_pi = initial_pi

        self._n_states = reward_matrix.shape[0]
        self._n_actions = reward_matrix.shape[1]
        self._n_opponent_actions = reward_matrix.shape[2]

        self._counts = np.zeros(self._n_states, dtype=np.int_)
        self._Q = fp_utils.get_initial_Q(  # pylint: disable=invalid-name
            self._initial_Q, self._n_states, self._n_actions, self._n_opponent_actions
        )
        self._pi = fp_utils.get_initial_pi(  # pylint: disable=invalid-name
            self._initial_pi, self._n_states, self._n_opponent_actions, self.rng
        )
        self.v = np.zeros(  # pylint: disable=invalid-name
            self._n_states, dtype=np.float_
        )

        self._utils = fp_utils.SmoothedFictitiousPlayUtils(
            self._Q, self._pi, self._tau, self.rng
        )

        self._t = 0
        self._eye: Final = np.eye(reward_matrix.shape[-1])

    @property
    def loggable_params(self) -> set[str]:
        return fp_utils.LOGGABLE_PARAMS

    def reset(self):
        super().reset()
        self._counts = np.zeros(self._n_states, dtype=np.int_)
        self._Q = fp_utils.get_initial_Q(
            self._initial_Q, self._n_states, self._n_actions, self._n_opponent_actions
        )
        self._pi = fp_utils.get_initial_pi(
            self._initial_pi, self._n_states, self._n_opponent_actions, self.rng
        )
        self.v = np.zeros(  # pylint: disable=invalid-name
            self._n_states, dtype=np.float_
        )
        self._t = 0
        self._utils.reset(self._Q, self._pi)

    def act(self, state: StateType) -> ActionType:
        return self._utils.smoothed_best_response(state)

    def update(
        self,
        state: StateType,
        actions: Sequence[ActionType],
        reward: RewardType,
        next_state: StateType,
        **kwargs,
    ):
        self._pi[state] += fp_utils.belief_update(
            self._pi[state],
            self._alpha(self._counts[state]),
            actions[1],  # type: ignore
        )
        self._Q += fp_utils.model_based_sync_q_update(
            self._Q,
            self.v,
            self._beta(self._t),
            self._gamma,
            self._R,
            self._T,
        )
        self.v = self._utils.smoothed_value_function(self._Q, self._pi)
        self._counts[state] += 1
        self._t += 1

coaction/agents/fictitious_play/sync_sfp.py
yigityalin/coaction
"""Utilities for fictitious play agents."""

from typing import Final

import numpy as np
import numpy.typing as npt

from coaction.games.game import ActionType, RewardType, StateType
from coaction.utils.math import softmax


LOGGABLE_PARAMS: Final[set] = {"Q", "pi", "counts", "v"}


def get_initial_Q(  # pylint: disable=invalid-name
    _initial_Q: None | int | float | np.ndarray,  # pylint: disable=invalid-name
    _n_states: int,
    _n_actions: int,
    _n_opponent_actions: int,
) -> npt.NDArray[RewardType]:
    """Return the initial Q matrix.

    Args:
        _initial_q (None | int | float | np.ndarray): The initial Q matrix.
        _n_states (int): The number of states.
        _n_actions (int): The number of actions.
        _n_opponent_actions (int): The number of opponent actions.
    """
    if _initial_Q is None:
        initial_Q = np.zeros(  # pylint: disable=invalid-name
            (_n_states, _n_actions, _n_opponent_actions)
        )
    elif isinstance(_initial_Q, (int, float)):
        initial_Q = (  # pylint: disable=invalid-name
            np.ones((_n_states, _n_actions, _n_opponent_actions)) * _initial_Q
        )
    elif isinstance(_initial_Q, np.ndarray):
        if _initial_Q.shape != (
            _n_states,
            _n_actions,
            _n_opponent_actions,
        ):
            raise ValueError(
                f"initial_q must have shape {(_n_states, _n_actions, _n_opponent_actions)}, "
                f"but has shape {_initial_Q.shape}."
            )
        initial_Q = _initial_Q.copy()  # pylint: disable=invalid-name
    else:
        raise TypeError(
            f"initial_q must be None, int, float, or np.ndarray, "
            f"but is {type(_initial_Q)}."
        )
    return initial_Q


def get_initial_pi(
    _initial_pi: None | np.ndarray,
    _n_states: int,
    _n_opponent_actions: int,
    rng: np.random.Generator,
) -> npt.NDArray[np.float_]:
    """Return the initial belief matrix.

    Args:
        _initial_pi (None | np.ndarray): The initial belief matrix.
        _n_states (int): The number of states.
        _n_opponent_actions (int): The number of opponent actions.
        rng (np.random.Generator): The random number generator.
    """
    if _initial_pi is None:
        initial_pi = rng.dirichlet(np.ones(_n_opponent_actions), size=_n_states)
    elif isinstance(_initial_pi, np.ndarray):
        if _initial_pi.shape != (_n_states, _n_opponent_actions):
            raise ValueError(
                f"initial_pi must have shape {(_n_states, _n_opponent_actions)}, "
                f"but has shape {_initial_pi.shape}."
            )
        initial_pi = _initial_pi.copy()
    else:
        raise TypeError(
            f"initial_pi must be None or np.ndarray, but is {type(_initial_pi)}."
        )
    return initial_pi


def q_function(
    Q: npt.NDArray[RewardType], pi: npt.NDArray[np.float_]
):  # pylint: disable=invalid-name
    """Return the Q matrix.

    Args:
        Q (npt.NDArray[RewardType]): The Q matrix.
        pi (npt.NDArray[np.float_]): The belief matrix.
    """
    return np.squeeze(Q @ pi[..., np.newaxis], axis=-1)


def belief_update(
    pi: npt.NDArray[np.float_], step_size: float, action: ActionType
):  # pylint: disable=invalid-name
    """Return the belief matrix update.

    Args:
        pi (npt.NDArray[np.float_]): The belief vector for a state.
        step_size (float): The step size.
        action (int): The action taken.
    """
    return step_size * (np.eye(pi.shape[-1])[action] - pi)


def model_free_q_update(
    Q: RewardType,  # pylint: disable=invalid-name
    v: RewardType,  # pylint: disable=invalid-name
    step_size: float,
    gamma: float,
    reward: RewardType,
):
    """Return the Q matrix update.

    Args:
        Q (RewardType): The Q matrix.
        v (RewardType): The value function.
        step_size (float): The step size.
        gamma (float): The discount factor.
        reward (RewardType): The reward received.
    """
    return step_size * (reward + gamma * v - Q)


def model_based_async_q_update(
    Q: npt.NDArray[RewardType],  # pylint: disable=invalid-name
    v: npt.NDArray[RewardType],  # pylint: disable=invalid-name
    step_size: float,
    gamma: float,
    state: StateType,
    reward_matrix: npt.NDArray[RewardType],
    transition_matrix: npt.NDArray[np.float_],
):
    """Return the Q matrix update.

    Args:
        Q (npt.NDArray[RewardType]): The Q matrix.
        v (npt.NDArray[np.float_]): The value function.
        step_size (float): The step size.
        gamma (float): The discount factor.
        state (StateType): The state.
        reward_matrix (npt.NDArray[RewardType]): The reward matrix.
        transition_matrix (npt.NDArray[np.float_]): The transition matrix.
    """
    return step_size * (
        reward_matrix[state] + gamma * transition_matrix[state] @ v - Q[state]
    )


def model_based_sync_q_update(
    Q: npt.NDArray[RewardType],  # pylint: disable=invalid-name
    v: npt.NDArray[RewardType],  # pylint: disable=invalid-name
    step_size: float,
    gamma: float,
    reward_matrix: npt.NDArray[RewardType],
    transition_matrix: npt.NDArray[np.float_],
):
    """Return the Q matrix update.

    Args:
        Q (npt.NDArray[RewardType]): The Q matrix.
        v (npt.NDArray[np.float_]): The value function.
        step_size (float): The step size.
        gamma (float): The discount factor.
        reward_matrix (npt.NDArray[RewardType]): The reward matrix.
        transition_matrix (npt.NDArray[np.float_]): The transition matrix.
    """
    return step_size * (reward_matrix + gamma * transition_matrix @ v - Q)


class FictitiousPlayUtils:
    """Utilities for fictitious play agents."""

    def __init__(
        self,
        Q: npt.NDArray,  # pylint: disable=invalid-name
        pi: npt.NDArray,  # pylint: disable=invalid-name
        rng: np.random.Generator,
    ):
        """Initialize the utilities.

        Args:
            Q (npt.NDArray[RewardType]): The Q matrix.
            pi (npt.NDArray[np.float_]): The belief matrix.
            tau (float): The temperature parameter.
            rng (np.random.Generator): The random number generator.
        """
        self._rng = rng

        self.q = q_function(Q, pi)  # pylint: disable=invalid-name

    def reset(self, Q: npt.NDArray, pi: npt.NDArray):  # pylint: disable=invalid-name
        """Reset the utilities.

        Args:
            Q (npt.NDArray[RewardType]): The Q matrix.
            pi (npt.NDArray[np.float_]): The belief matrix.
        """
        self.q = q_function(Q, pi)

    def value_function(
        self, Q: npt.NDArray[RewardType], pi: npt.NDArray[np.float_]
    ):  # pylint: disable=invalid-name
        """Return the value function.

        Args:
            state (StateType): The state.
            Q (npt.NDArray[RewardType]): The Q matrix.
            pi (npt.NDArray[np.float_]): The belief matrix.
        """
        self.q = q_function(Q, pi)  # pylint: disable=invalid-name
        return self.q.max(axis=-1)

    def best_response(
        self,
        state: StateType,
        *,
        Q: npt.NDArray[RewardType] | None = None,  # pylint: disable=invalid-name
        pi: npt.NDArray[np.float_] | None = None,  # pylint: disable=invalid-name
    ):
        """Return the best response.

        Args:
            Q (npt.NDArray[RewardType]): The Q matrix.
            pi (npt.NDArray[np.float_]): The belief matrix.
        """
        if Q is not None and pi is not None:
            self.q = q_function(Q, pi)
        q = self.q[state]  # pylint: disable=invalid-name
        if np.isclose(q.min(), q.max()):
            return self._rng.choice(len(q))
        return self.q[state].argmax()


class SmoothedFictitiousPlayUtils:
    """Utilities for smoothed fictitious play agents."""

    def __init__(
        self,
        Q: npt.NDArray,  # pylint: disable=invalid-name
        pi: npt.NDArray,  # pylint: disable=invalid-name
        tau: float,
        rng: np.random.Generator,
    ):
        """Initialize the utilities.

        Args:
            Q (npt.NDArray[RewardType]): The Q matrix.
            pi (npt.NDArray[np.float_]): The belief matrix.
            tau (float): The temperature parameter.
            rng (np.random.Generator): The random number generator.
        """
        self._tau = tau
        self._rng = rng

        self.q = q_function(Q, pi)  # pylint: disable=invalid-name
        self.mu = softmax(self.q, self._tau)  # pylint: disable=invalid-name

    def reset(self, Q: npt.NDArray, pi: npt.NDArray):  # pylint: disable=invalid-name
        """Reset the utilities.

        Args:
            Q (npt.NDArray[RewardType]): The Q matrix.
            pi (npt.NDArray[np.float_]): The belief matrix.
        """
        self.q = q_function(Q, pi)
        self.mu = softmax(self.q, self._tau)

    def smoothed_value_function(
        self,
        Q: npt.NDArray[RewardType],  # pylint: disable=invalid-name
        pi: npt.NDArray[np.float_],  # pylint: disable=invalid-name
    ) -> npt.NDArray:
        """Return the smoothed value function.

        Args:
            Q (npt.NDArray[RewardType]): The Q matrix.
            pi (npt.NDArray[np.float_]): The belief matrix.
        """
        self.q = q_function(Q, pi)
        self.mu = softmax(self.q, self._tau)
        return np.sum(self.mu * self.q, axis=-1)

    def smoothed_best_response(
        self,
        state: StateType,
        *,
        Q: npt.NDArray[RewardType] | None = None,  # pylint: disable=invalid-name
        pi: npt.NDArray[np.float_] | None = None,  # pylint: disable=invalid-name
    ) -> ActionType:
        """Return the smoothed best response.

        Args:
            state (StateType): The state.
            Q (npt.NDArray[RewardType]): The Q matrix.
            pi (npt.NDArray[np.float_]): The belief matrix.
        """
        if Q is not None and pi is not None:
            self.q = q_function(Q, pi)
            self.mu = softmax(self.q[state], self._tau)
        mu = self.mu[state]  # pylint: disable=invalid-name
        return self._rng.choice(len(mu), p=mu)  # type: ignore

coaction/agents/fictitious_play/utils.py
yigityalin/coaction
"""Implementation of Individual Q-Learning."""

from coaction.agents.ind_q.ind_q import IndividualQLearning


__all__ = ["IndividualQLearning"]

coaction/agents/ind_q/__init__.py
yigityalin/coaction
"""Implementation of individual Q-learning agent."""

from collections.abc import Callable
from typing import Sequence

import numpy as np
import numpy.typing as npt

from coaction.agents.agent import TwoPlayerAgent
from coaction.agents.ind_q import utils as ind_q_utils
from coaction.games.game import ActionType, RewardType, StateType
from coaction.utils.math import softmax


class IndividualQLearning(TwoPlayerAgent):
    """Implementation of individual Q-learning agent.

    Parameters
    ----------
    name : str
        The name of the agent.
    seed : int
        The seed for the random number generator.
    alpha : Callable[[int], float]
        Step size for q update.
    beta : Callable[[int], float]
        Step size for v update.
    gamma : float
        Discount factor
    tau : float
        Temperature parameter
    max_step_size : float
        Maximum step size for q update
    """

    def __init__(
        self,
        name: str,
        seed: int,
        reward_matrix: npt.NDArray[RewardType],
        alpha: Callable[[int], float],
        beta: Callable[[int], float],
        gamma: float,
        tau: float,
        max_step_size: float,
        initial_q: None | int | float | np.ndarray = None,
        **kwargs,
    ):
        """Initialize the agent.

        Args:
            name (str): The name of the agent.
            seed (int): The seed for the random number generator.
            reward_matrix (npt.NDArray[RewardType]): The reward matrix.
            alpha (Callable[[int], float]): Step size for q update.
            beta (Callable[[int], float]): Step size for v update.
            gamma (float): Discount factor
            tau (float): Temperature parameter
            max_step_size (float): Maximum step size for q update
            initial_q (None | int | float | np.ndarray): The initial Q matrix.
        """
        super().__init__(name, seed, **kwargs)
        self._alpha = alpha
        self._beta = beta
        self._gamma = gamma
        self._tau = tau
        self._max_step_size = max_step_size
        self._initial_q = initial_q

        self._n_states = reward_matrix.shape[0]
        self._n_actions = reward_matrix.shape[1]

        self._counts = np.zeros(self._n_states)
        self._q = ind_q_utils.get_initial_q(
            self._initial_q, self._n_states, self._n_actions
        )
        self.v = np.zeros(self._n_states)  # pylint: disable=invalid-name

        self._mu = None

    @property
    def loggable_params(self) -> set[str]:
        return ind_q_utils.LOGGABLE_PARAMS

    def reset(self):
        super().reset()
        self._counts = np.zeros_like(self._counts)
        self._q = ind_q_utils.get_initial_q(
            self._initial_q, self._n_states, self._n_actions
        )
        self.v = np.zeros_like(self.v)

    def act(self, state: StateType) -> ActionType:
        self._mu = softmax(self._q[state], self._tau)
        return self.rng.choice(self._n_actions, p=self._mu)

    def update(
        self,
        state: StateType,
        actions: Sequence[ActionType],
        reward: RewardType,
        next_state: StateType,
        **kwargs,
    ):
        action = actions[0]
        self._q[state, action] += min(
            self._alpha(self._counts[state]) / self._mu[action], self._max_step_size
        ) * (reward + self._gamma * self.v[next_state] - self._q[state, action])
        self.v[state] += self._beta(self._counts[state]) * (
            np.dot(self._q[state], self._mu) - self.v[state]
        )
        self._counts[state] += 1

coaction/agents/ind_q/ind_q.py
yigityalin/coaction
"""Utility functions for independent Q-learning agents."""

from typing import Final

import numpy as np
import numpy.typing as npt

from coaction.games.game import RewardType


LOGGABLE_PARAMS: Final[set] = {"q", "v", "counts"}


def get_initial_q(
    _initial_q: None | int | float | np.ndarray,
    _n_states: int,
    _n_actions: int,
) -> npt.NDArray[RewardType]:
    """Return the initial Q matrix.

    Args:
        _initial_q (None | int | float | np.ndarray): The initial Q matrix.
        _n_states (int): The number of states.
        _n_actions (int): The number of actions.
        _n_opponent_actions (int): The number of opponent actions.
    """
    if _initial_q is None:
        initial_q = np.zeros((_n_states, _n_actions))
    elif isinstance(_initial_q, (int, float)):
        initial_q = np.ones((_n_states, _n_actions)) * _initial_q
    elif isinstance(_initial_q, np.ndarray):
        if _initial_q.shape != (_n_states, _n_actions):
            raise ValueError(
                f"initial_q must have shape {(_n_states, _n_actions)}, "
                f"but has shape {_initial_q.shape}."
            )
        initial_q = _initial_q
    else:
        raise TypeError(
            f"initial_q must be None, int, float, or np.ndarray, "
            f"but is {type(_initial_q)}."
        )
    return initial_q

coaction/agents/ind_q/utils.py
yigityalin/coaction
"""Create a new project config template.

An example usage is given in create_project.sh.
"""

from collections.abc import Collection
from typing import Final
import argparse
import inspect

from coaction import agents, loggers, games
from coaction.loggers import AgentLogger, GameLogger, ProgressLogger
from coaction.utils.paths import ProjectPaths


_EXCLUDED_PARAMETERS: Final[frozenset[str]] = frozenset(
    [
        "self",
        "paths",
        "config",
        "queue",
        "project_dir",
        "project_name",
        "experiment_name",
        "agent",
        "agent_names",
        "game",
        "reward_matrix",
        "transition_matrix",
    ]
)


class ClassNotFoundError(Exception):
    """Raised when a class is not found."""


def _parse_args():
    """Parse command line arguments.

    An example usage is given in create_project.sh.
    """
    parser = argparse.ArgumentParser(
        description="Create a new project config template."
    )
    parser.add_argument(
        "--parent_dir",
        help="Path to the directory in which project config will be initialized.",
    )
    parser.add_argument("--project_name", help="Name of the project.")
    parser.add_argument(
        "--experiment_names", nargs="+", help="Names of the experiments."
    )
    parser.add_argument(
        "--agent_types", nargs="+", action="append", help="Types of the agents."
    )
    parser.add_argument(
        "--game_types",
        nargs="+",
        help="Types of the games.",
    )
    return parser.parse_args()


def _check_args(args: argparse.Namespace):
    """Validate the command line arguments.

    Args:
        args (argparse.Namespace): The command line arguments.
    """
    path = ProjectPaths.from_parent(
        args.parent_dir, args.project_name
    ).get_project_dir()
    if path.exists():
        raise ValueError(f"Project {args.project_name} already exists.")
    if len(args.experiment_names) != len(args.agent_types):
        raise ValueError(
            f"Number of experiments ({len(args.experiment_names)}) "
            f"does not match number of agent types ({len(args.agent_types)})."
        )
    if len(args.experiment_names) != len(args.game_types):
        raise ValueError(
            f"Number of experiments ({len(args.experiment_names)}) "
            f"does not match number of game types ({len(args.game_types)})."
        )


def _count_endline_characters(template: str) -> int:
    """Count the number of endline characters at the end of a string.

    Args:
        template (str): The string.
    """
    return len(template) - len(template.rstrip("\n"))


def _get_class_from_string(cls_name: str) -> type:
    """Return the class from a string.

    Args:
        cls_name (str): The name of the class.
    """
    if cls_name in agents.__all__:
        return getattr(agents, cls_name)
    if cls_name in loggers.__all__:
        return getattr(loggers, cls_name)
    if cls_name in games.__all__:
        return getattr(games, cls_name)
    raise ClassNotFoundError(f"Class {cls_name} not found.")


def _get_signature(cls: type, exclude: Collection[str] = _EXCLUDED_PARAMETERS):
    """Return the signature of a class.

    Args:
        cls (type): The class.
    """
    signature = inspect.signature(cls)
    return {k: v for k, v in signature.parameters.items() if k not in exclude}


def _get_param_template(
    param_name: str, cls: type, exclude: Collection[str] = _EXCLUDED_PARAMETERS
):
    """Return the template for a parameter.

    Args:
        param_name (str): The name of the parameter.
        cls_name (str): The name of the class.
    """
    signature = _get_signature(cls, exclude)
    template = param_name + " = {   # type: " + cls.__name__ + "\n"
    for k, val in signature.items():
        template += f"    '{k}': {val.default if val.default is not val.empty else ''},"
        type_ = (
            val.annotation
            if val.annotation is not val.empty
            else "typing.Any (additional keyword parameters)"
        )
        template += f"  # type: {type_} \n"
    return template + "}\n\n"


def _create_dir(args: argparse.Namespace):
    """Create the project directory.

    Args:
        args (argparse.Namespace): The command line arguments.
    """
    path = ProjectPaths.from_parent(
        args.parent_dir, args.project_name
    ).get_project_dir()
    path.mkdir(parents=True)


def _rm_dir(args: argparse.Namespace):
    """Remove the project directory.

    Args:
        args (argparse.Namespace): The command line arguments.
    """
    path = ProjectPaths.from_parent(
        args.parent_dir, args.project_name
    ).get_project_dir()
    path.rmdir()


def _create_project_config(args: argparse.Namespace):
    """Create the project config file.

    Args:
        args (argparse.Namespace): The command line arguments.
    """
    template_path = ProjectPaths.get_project_config_template_path()
    with open(template_path, "r", encoding="utf-8") as file:
        file_content = file.read()

    config_path = ProjectPaths.from_parent(
        args.parent_dir, args.project_name
    ).get_project_config_path()
    with open(config_path, "w", encoding="utf-8") as file:
        file.write(file_content)


def _create_files(args: argparse.Namespace):
    """Create the project config files.

    Args:
        args (argparse.Namespace): The command line arguments.
    """
    ProjectPaths.from_parent(
        args.parent_dir, args.project_name
    ).get_project_config_dir().mkdir(parents=True)
    for experiment_name, agent_types, game_type in zip(
        args.experiment_names,
        args.agent_types,
        args.game_types,
    ):
        file_content = _create_file_content(agent_types, game_type)
        path = ProjectPaths.from_parent(
            args.parent_dir, args.project_name, experiment_name
        ).get_experiment_config_path()
        with open(path, "w", encoding="utf-8") as file:
            file.write(file_content)


def _create_file_content(agent_types: list[str], game_type: str):
    """Create the content of an experiment config file.

    Args:
        agent_types (list[type]): The types of the agents.
        game_type (type): The type of the game.
    """
    imports_template = _get_imports_template()
    agents_template = _get_agents_template(agent_types)
    game_template = _get_game_template(game_type)
    loggers_template = _get_loggers_template()
    episodes_and_stages_template = _get_episodes_and_stages_template()
    multiprocessing_template = _get_multiprocessing_template()
    template = (
        imports_template
        + multiprocessing_template
        + episodes_and_stages_template
        + agents_template
        + game_template
        + loggers_template
    )
    # Count the number of endline characters at the end of the template
    # and ensure that there is only one.
    endline_characters = _count_endline_characters(template)
    if endline_characters > 1:
        template = template[: -endline_characters + 1]
    elif endline_characters == 0:
        template += "\n"
    return template


def _load_agents_template(
    agent_index: int, agent_name: str, load_file: bool = True
) -> str:
    """Return the template for the agents.

    Args:
        agent_types (list[str]): The types of the agents.
    """
    if load_file:
        path = ProjectPaths.get_agent_template_path()
        with open(path, "r", encoding="utf-8") as file:
            template = file.read()
        template = template.format_map({"agent_name": agent_name})
    else:
        template = ""
    template += (
        "agent_{agent_index}_kwargs = ".format_map({"agent_index": agent_index}) + "{\n"
    )
    template += (
        "     # TODO: specify any keyword arguments for agent {agent_index} here\n"
    )
    template += "}\n\n"
    return template


def _get_imports_template() -> str:
    template = "from coaction import agents, loggers, games, utils\n"
    return template + "\n\n"


def _get_agents_template(agent_types: list[str]) -> str:
    """Return the template for the agents.

    Args:
        agent_types (list[str]): The types of the agents.
    """
    agent_type_names = []
    agents_kwargs_template = "\n# Agent kwargs\n"
    for agent_index, agent_type in enumerate(agent_types):
        try:
            agent_cls = _get_class_from_string(agent_type)
            if agent_cls is agents.Agent:
                raise ValueError(
                    "Agent cannot be used as an agent type. "
                    "Please specify a subclass of Agent or a name for your custom type."
                )
            agents_kwargs_template += _get_param_template(
                f"agent_{agent_index}_kwargs", agent_cls
            )
            agent_type_names.append(f"agents.{agent_cls.__name__}")
        except ClassNotFoundError:
            load_file = agent_type not in agent_type_names
            agents_kwargs_template += _load_agents_template(
                agent_index, agent_type, load_file
            )
            agent_type_names.append(agent_type)

    agent_types_template = f'agent_types = [{", ".join(agent_type_names)}] \n'

    agent_kwargs_names = [
        f"agent_{agent_index}_kwargs" for agent_index in range(len(agent_types))
    ]
    agents_kwargs_list_template = (
        f'agent_kwargs = [{", ".join(agent_kwargs_names)}] \n\n'
    )
    return agents_kwargs_template + agent_types_template + agents_kwargs_list_template


def _get_game_template(game_type: str) -> str:
    """Return the template for the game.

    Args:
        game_type (type): The type of the game.
    """
    game_cls = _get_class_from_string(game_type)
    game_template = "# Game kwargs\n"
    game_template += f"game_type = games.{game_cls.__name__}\n"
    game_template += _get_param_template("game_kwargs", game_cls, exclude={"self"})
    return game_template


def _get_loggers_template() -> str:
    """Return the template for the logger.

    Args:
        logger_type (type): The type of the logger.
    """
    logger_template = "# Logger kwargs\n"
    logger_template += _get_param_template("agent_logger_kwargs", AgentLogger)
    logger_template += _get_param_template("game_logger_kwargs", GameLogger)
    logger_template += _get_param_template("progress_logger_kwargs", ProgressLogger)
    return logger_template


def _get_episodes_and_stages_template() -> str:
    """Return the template for the total episodes and stages.

    Args:
        total_episodes (int): The total number of episodes.
        total_stages (int): The total number of stages.
    """
    return "total_episodes = \ntotal_stages = \n\n"


def _get_multiprocessing_template() -> str:
    """Return the template for multiprocessing."""
    return "num_parallel_episodes = None  # 'None' means use all available cores\n\n"


def main(args: argparse.Namespace):
    """Run the coaction command line tool.

    Args:
        args (argparse.Namespace): The command line arguments.
    """
    _check_args(args)
    _create_dir(args)
    try:
        _create_files(args)
        _create_project_config(args)
    except ClassNotFoundError as exc:
        _rm_dir(args)
        raise exc


if __name__ == "__main__":
    main(_parse_args())

coaction/create.py
yigityalin/coaction
"""Configuration for experiments."""

from pathlib import Path
from typing import Any
import dataclasses
import shutil

from coaction.utils.modules import load_module
from coaction.utils.paths import ProjectPaths


@dataclasses.dataclass
class ExperimentConfig:
    """Configuration for an experiment.

    An experiment contains multiple agents and a game.

    Attributes:
        name (str): The name of the experiment.
        agent_types (list[type]): The types of the agents.
        game_type (type): The type of the game.
        game_kwargs (dict[str, Any]): The keyword arguments for the game.
        agent_kwargs (list[dict[str, Any]]): The keyword arguments for the agents.
        agent_logger_kwargs (dict[str, Any]): The keyword arguments for the
            agent loggers.
        game_logger_kwargs (dict[str, Any]): The keyword arguments for the game
            logger.
        progress_logger_kwargs (dict[str, Any]): The keyword arguments for the
            progress logger.
        total_episodes (int): The total number of episodes.
        total_stages (int): The total number of stages.
        num_parallel_episodes (int): The number of episodes to run in parallel.
    """

    name: str
    agent_types: list[type]
    game_type: type
    game_kwargs: dict[str, Any]
    agent_kwargs: list[dict[str, Any]]
    agent_logger_kwargs: dict[str, Any]
    game_logger_kwargs: dict[str, Any]
    progress_logger_kwargs: dict[str, Any]
    total_episodes: int
    total_stages: int
    num_parallel_episodes: int

    @classmethod
    def from_py_file(cls, path: Path | str):
        """Return an experiment config from a Python file."""
        module = load_module(path)
        fields = [field for field in dataclasses.fields(cls) if field.name != "name"]
        kwargs = {field.name: getattr(module, field.name) for field in fields}
        kwargs["name"] = getattr(module, "name", None) or module.__name__
        config = cls(**kwargs)  # type: ignore
        config._validate()
        return config

    def _validate(self):
        agent_names = [kwargs["name"] for kwargs in self.agent_kwargs]  # type: ignore
        if len(agent_names) != len(set(agent_names)):
            raise ValueError("Agent names must be unique.")

    def to_dict(self):
        """Return a dictionary representation of the config."""
        return dataclasses.asdict(self)


@dataclasses.dataclass
class GlobalConfig:
    """Global configuration for a project.

    Attributes:
        num_parallel_experiments (int): The number of experiments to run in
            parallel.
        order (list[str]): The order in which to run the experiments.
    """

    num_parallel_experiments: int
    num_parallel_episodes: int
    order: list[str]

    @classmethod
    def from_py_file(cls, path: Path | str):
        """Return a global config from a Python file."""
        module = load_module(path)
        fields = [field for field in dataclasses.fields(cls)]
        kwargs = {field.name: getattr(module, field.name) for field in fields}
        return cls(**kwargs)

    def to_dict(self):
        """Return a dictionary representation of the config."""
        return dataclasses.asdict(self)


@dataclasses.dataclass
class ProjectConfig:
    """Configuration for a project.

    A project contains multiple experiments.

    Attributes:
        project_dir (Path | str): The path to the project directory.
        name (str): The name of the project.
        experiments (list[ExperimentConfig]): The experiments in the project.
    """

    project_dir: Path | str
    name: str
    global_config: GlobalConfig
    experiments: list[ExperimentConfig]
    paths: ProjectPaths

    @classmethod
    def from_dir(cls, path: Path | str):
        """Return a project config from a directory."""
        paths = ProjectPaths(path, increment_run=True)
        project_dir = paths.get_project_dir()
        if not project_dir.is_dir() or not paths.get_project_config_dir().is_dir():
            paths.cleanup()
            raise ValueError(f"Path {path} is not a directory.")
        global_config = GlobalConfig.from_py_file(paths.get_project_config_path())
        experiments = sorted(
            [
                ExperimentConfig.from_py_file(experiment_path)
                for experiment_path in paths.get_experiment_config_paths()
            ],
            key=ProjectConfig._get_sort_key(global_config),  # type: ignore
        )
        config = cls(
            project_dir=project_dir,
            name=project_dir.stem,
            global_config=global_config,
            experiments=experiments,
            paths=paths,
        )
        config.copy_config_files()
        return config

    @staticmethod
    def _get_sort_key(config: GlobalConfig):
        """Return the sort key for an experiment."""
        if config.order is None:
            return lambda cfg: cfg.name
        return lambda cfg: config.order.index(cfg.name)

    def copy_config_files(self):
        """Copy the config to the project directory."""
        shutil.copytree(
            self.paths.get_project_config_dir(),
            self.paths.get_project_run_config_dir(),
            ignore=shutil.ignore_patterns("__pycache__"),
        )

    def to_dict(self):
        """Return a dictionary representation of the config."""
        return dataclasses.asdict(self)

coaction/experiments/config.py
yigityalin/coaction
"""Implementation of episodes representing a single game between agents."""

from typing import Any
import multiprocessing as mp

from coaction.agents.agent import Agent
from coaction.experiments.multiprocessing import DummySemaphore
from coaction.games.game import MarkovGame
from coaction.loggers.agent import AgentLogger
from coaction.loggers.game import GameLogger
from coaction.loggers.progress import ProgressLogger


class Episode(mp.Process):
    """A single game between agents."""

    def __init__(
        self,
        game: MarkovGame,
        agents: list[Agent],
        agent_logger: AgentLogger,
        game_logger: GameLogger,
        progress_logger: ProgressLogger,
        episode: int,
        total_stages: int,
        semaphore: DummySemaphore
        | Any,  # TODO: Use multiprocessing.Semaphore instead of Any.
        global_semaphore: DummySemaphore
        | Any,  # TODO: Use multiprocessing.Semaphore instead of Any.
    ) -> None:
        """Initialize the episode."""
        super().__init__()
        self.game: MarkovGame = game
        self.agents: list[Agent] = agents
        self.agent_logger: AgentLogger = agent_logger
        self.game_logger: GameLogger = game_logger
        self.progress_logger: ProgressLogger = progress_logger
        self.episode: int = episode
        self.total_stages: int = total_stages
        self.semaphore = semaphore
        self.global_semaphore = global_semaphore

    def run(self) -> None:
        """Run the episode."""
        # Acquire the semaphore to limit the number of parallel episodes.
        self.global_semaphore.acquire()
        self.semaphore.acquire()

        game = self.game
        agents = self.agents
        agent_logger = self.agent_logger
        game_logger = self.game_logger
        progress_logger = self.progress_logger

        agent_logger.on_episode_begin(self.episode, agents)
        game_logger.on_episode_begin(self.episode, game)
        progress_logger.on_episode_begin(self.episode)

        state = game.reset()
        for agent in self.agents:
            agent.reset()

        for stage in range(1, 1 + self.total_stages):
            agent_logger.on_stage_begin(stage, agents)
            game_logger.on_stage_begin(stage, game)
            progress_logger.on_stage_begin(stage, agents=agents, episode=self.episode)

            actions = [agent.act(game.state) for agent in self.agents]
            next_state, rewards = game.step(actions)

            for i, agent in enumerate(agents):
                actions_view = actions[i:] + actions[:i]
                agent.update(
                    state=state,
                    actions=actions_view,
                    reward=rewards[i],
                    next_state=next_state,
                )

            agent_logger.on_stage_end(stage, agents)
            game_logger.on_stage_end(stage, game, state, actions, rewards)
            progress_logger.on_stage_end(stage, agents=agents, episode=self.episode)

            state = next_state

        agent_logger.on_episode_end(self.episode, agents)
        game_logger.on_episode_end(self.episode, game)
        progress_logger.on_episode_end(self.episode)

        # Release the semaphore to allow another episode to run.
        self.semaphore.release()
        self.global_semaphore.release()

coaction/experiments/episode.py
yigityalin/coaction
"""Experiment class."""

from copy import deepcopy
from typing import Any, NamedTuple
import inspect
import multiprocessing as mp

import numpy as np

from coaction.agents.agent import Agent
from coaction.experiments.config import ExperimentConfig, ProjectConfig
from coaction.experiments.episode import Episode
from coaction.experiments.multiprocessing import DummySemaphore
from coaction.games.game import MarkovGame
from coaction.loggers.agent import AgentLogger
from coaction.loggers.game import GameLogger
from coaction.loggers.progress import ProgressLogger


class _AgentRequirements(NamedTuple):
    """Requirements for an agent."""

    reward_matrix: bool
    transition_matrix: bool


def _inspect_agent(agent_type: type) -> _AgentRequirements:
    """Return the requirements for an agent."""
    signature = inspect.signature(agent_type.__init__)
    reward_matrix = "reward_matrix" in signature.parameters
    transition_matrix = "transition_matrix" in signature.parameters
    return _AgentRequirements(reward_matrix, transition_matrix)


class Experiment(mp.Process):
    """A single experiment."""

    def __init__(
        self,
        project_config: ProjectConfig,
        config: ExperimentConfig,
        semaphore: DummySemaphore
        | Any,  # TODO: Use multiprocessing.Semaphore instead of Any.
        global_semaphore: DummySemaphore
        | Any,  # TODO: Use multiprocessing.Semaphore instead of Any.
    ) -> None:
        """Initialize the experiment."""
        super().__init__()
        self.project_config: ProjectConfig = project_config
        self.config: ExperimentConfig = config
        self.paths = project_config.paths.with_experiment_name(config.name)
        self.semaphore = semaphore
        self.global_semaphore = global_semaphore
        self._agent_seed_sequences: dict[int, np.random.SeedSequence] = {}

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(name={self.config.name})"

    def _construct_experiment(
        self,
    ) -> tuple[list[Agent], MarkovGame, AgentLogger, GameLogger, ProgressLogger]:
        """Construct the experiment."""
        game = self._construct_game()
        agents = self._construct_agents(game)
        agent_logger = self._construct_agent_logger()
        game_logger = self._construct_game_logger(agents)
        progress_logger = self._construct_progress_logger()
        return agents, game, agent_logger, game_logger, progress_logger

    def _construct_game(self) -> MarkovGame:
        """Construct the game."""
        return self.config.game_type(**self.config.game_kwargs).as_markov_game()

    def _construct_agents(self, game: MarkovGame) -> list[Agent]:
        """Construct the agents."""
        requirements = [
            _inspect_agent(agent_type) for agent_type in self.config.agent_types
        ]
        agent_kwargs = deepcopy(self.config.agent_kwargs)
        for agent_idx, requirement in enumerate(requirements):
            reward_matrix, transition_matrix = game.view(agent_idx)
            if requirement.reward_matrix:
                agent_kwargs[agent_idx]["reward_matrix"] = reward_matrix
            if requirement.transition_matrix:
                agent_kwargs[agent_idx]["transition_matrix"] = transition_matrix
        agents = [
            agent_type(**agent_kwargs)
            for agent_type, agent_kwargs in zip(self.config.agent_types, agent_kwargs)
        ]
        return agents

    def _construct_agent_logger(self) -> AgentLogger:
        """Construct the agent loggers."""
        agent_logger = AgentLogger(self.paths, **self.config.agent_logger_kwargs)
        return agent_logger

    def _construct_game_logger(self, agents: list[Agent]) -> GameLogger:
        """Construct the game logger."""
        agent_names = [agent.name for agent in agents]
        game_logger = GameLogger(
            agent_names,
            self.paths,
            **self.config.game_logger_kwargs,  # type: ignore
        )
        return game_logger

    def _construct_progress_logger(self) -> ProgressLogger:
        """Construct the progress logger."""
        progress_logger = ProgressLogger(
            self.config, self.paths, **self.config.progress_logger_kwargs
        )
        return progress_logger

    def _update_and_get_seed(self, agent_idx: int) -> int:
        if agent_idx not in self._agent_seed_sequences:
            self._agent_seed_sequences[agent_idx] = np.random.SeedSequence(
                self.config.agent_kwargs[agent_idx]["seed"]
            )
        else:
            self._agent_seed_sequences[agent_idx] = self._agent_seed_sequences[
                agent_idx
            ].spawn(1)[0]
        return self._agent_seed_sequences[agent_idx].generate_state(1)[0]

    def run(self):
        """Run the experiment."""
        # Acquire the semaphore to limit the number of parallel experiments.
        self.semaphore.acquire()

        (
            agents,
            game,
            agent_logger,
            game_logger,
            progress_logger,
        ) = self._construct_experiment()
        agent_logger.on_experiment_begin(agents)
        game_logger.on_experiment_begin(game)
        progress_logger.on_experiment_begin()

        if self.config.num_parallel_episodes is not None:
            semaphore = mp.Semaphore(self.config.num_parallel_episodes)
        else:
            semaphore = DummySemaphore()

        episodes: list[Episode] = []
        for episode in range(self.config.total_episodes):
            episode = Episode(
                game=game.clone(),
                agents=[
                    agent.clone(
                        seed=self._update_and_get_seed(agent_idx)
                    )
                    for agent_idx, agent in enumerate(agents)
                ],
                agent_logger=agent_logger.clone(),
                game_logger=game_logger.clone(),
                progress_logger=progress_logger.clone(),
                episode=episode,
                total_stages=self.config.total_stages,
                semaphore=semaphore,
                global_semaphore=self.global_semaphore,
            )
            episodes.append(episode)
            episode.start()

        for episode in episodes:
            episode.join()

        agent_logger.on_experiment_end(agents)
        game_logger.on_experiment_end(game)
        progress_logger.on_experiment_end()

        # Release the semaphore to allow another experiment to run.
        self.semaphore.release()

coaction/experiments/experiment.py
yigityalin/coaction
"""Multiprocessing utilities."""


class DummySemaphore:
    """Dummy semaphore class for cases the number of processes are not limited."""

    def __init__(self):
        """Initialize the dummy semaphore."""

    def __repr__(self) -> str:
        """Return a string representation of the dummy semaphore."""
        return f"{self.__class__.__name__}()"

    def acquire(self):
        """Dummy acquire method."""

    def release(self):
        """Dummy release method."""

coaction/experiments/multiprocessing.py
yigityalin/coaction
"""A project that contains multiple experiments."""

import multiprocessing as mp
import time

from coaction.experiments.config import ProjectConfig
from coaction.experiments.experiment import Experiment
from coaction.experiments.multiprocessing import DummySemaphore
from coaction.utils.time import get_strftime


class Project:
    """A project that contains multiple experiments."""

    def __init__(self, config: ProjectConfig):
        self.config = config
        config.paths.get_project_run_log_dir().mkdir(parents=True, exist_ok=True)
        self._log_file = config.paths.get_project_run_progress_log_path().open(
            "w+", encoding="utf-8"
        )

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(name={self.config.name})"

    def run(self):
        """Run the project."""
        self.log_text(f"Project {self.config.name} started.")

        start = time.perf_counter()
        if self.config.global_config.num_parallel_experiments is not None:
            semaphore = mp.Semaphore(self.config.global_config.num_parallel_experiments)
        else:
            semaphore = DummySemaphore()

        if self.config.global_config.num_parallel_episodes is not None:
            global_semaphore = mp.Semaphore(
                self.config.global_config.num_parallel_episodes
            )
        else:
            global_semaphore = DummySemaphore()

        experiments: list[Experiment] = []
        for experiment_config in self.config.experiments:
            experiment = Experiment(
                self.config, experiment_config, semaphore, global_semaphore
            )
            experiments.append(experiment)
            experiment.start()

        for experiment in experiments:
            experiment.join()
        end = time.perf_counter()
        self.log_text(
            f"Project {self.config.name} finished in {end - start:0.4f} seconds."
        )
        self.cleanup()

    def cleanup(self):
        """Clean up the project."""
        self._log_file.close()

    def log_text(self, text):
        """Log text to the log file."""
        self._log_file.write(f"[{get_strftime()}] - {text}\n")
        self._log_file.flush()

coaction/experiments/project.py
yigityalin/coaction
"""Implementations of games.

This module contains implementations of games. Each game is implemented as a
class that inherits from the abstract base class :class:`Game`. The
:class:`Game` class defines the interface that all games must implement.

For the sake of simplicity, we mainly focus on two-player matrix/Markov games. These games
are implemented as subclasses of :class:`TwoPlayerMatrixGame` and :class:`TwoPlayerMarkovGame`.
"""

from coaction.games.game import (
    Game,
    MarkovGame,
    MatrixGame,
    StateType,
    ActionType,
    RewardType,
)


__all__ = [
    "Game",
    "MarkovGame",
    "MatrixGame",
    "StateType",
    "ActionType",
    "RewardType",
]

coaction/games/__init__.py
yigityalin/coaction
"""Core classes for games.

This module contains the base classes for games. Each game is implemented as a
class that inherits from the abstract base class :class:`Game`. The
:class:`Game` class defines the interface that all games must implement.

For the sake of simplicity, we mainly focus on matrix/Markov games. These games
are implemented as subclasses of :class:`MatrixGame` and :class:`MarkovGame`.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Sequence
from copy import deepcopy
from pathlib import Path
import json

import numpy as np
import numpy.typing as npt


StateType = int
ActionType = int
RewardType = np.float_


class Game(ABC):
    """Base class for games.

    Parameters
    ----------
    name : str
        The name of the game.
    seed : int
        The seed for the random number generator.
    metadata : dict
        Metadata for the game.
    """

    def __init__(self, name: str, seed: int, **metadata):
        """Initialize the game.

        Args:
            name (str): The name of the game.
            seed (int): The seed for the random number generator.
            metadata (dict): Metadata for the game.
        """
        self.name: str = name
        self.seed: int = seed
        self.metadata: dict = metadata
        self._rng: np.random.Generator | None = None

    def __repr__(self):
        return f"{self.__class__.__name__}({self.__dict__})"

    @property
    def rng(self) -> np.random.Generator:
        """Return the random number generator."""
        if self._rng is None:
            self._rng = np.random.default_rng(self.seed)
        return self._rng

    @property
    @abstractmethod
    def n_agents(self) -> int:
        """Return the number of agents."""

    def clone(self):
        """Return a copy of the game."""
        return deepcopy(self)

    @abstractmethod
    def reset(self):
        """Reset the game."""

    @abstractmethod
    def step(
        self, actions: Sequence[ActionType]
    ) -> npt.NDArray[RewardType] | tuple[StateType, npt.NDArray[RewardType]]:
        """Execute a step in the game."""

    @abstractmethod
    def view(
        self, agent_id: int, only_viewer: bool = True
    ) -> npt.NDArray | tuple[npt.NDArray, npt.NDArray]:
        """Return a view of the game for the given agent."""

    @abstractmethod
    def to_json(self) -> dict:
        """Return a JSON representation of the game."""

    @classmethod
    @abstractmethod
    def from_json(cls, path: Path):
        """Load the game from a JSON representation."""

    @abstractmethod
    def as_markov_game(self) -> MarkovGame:
        """Return a Markov game with the given transition matrix."""


class MatrixGame(Game):
    """Base class for matrix games."""

    def __init__(self, name: str, reward_matrix: npt.NDArray, seed: int, **metadata):
        super().__init__(name, seed, **metadata)
        self._reward_matrix: npt.NDArray[np.float_] = reward_matrix

    @property
    def n_agents(self) -> int:
        """Return the number of agents."""
        return self.reward_matrix.shape[0]

    @property
    def reward_matrix(self) -> npt.NDArray[RewardType]:
        """Return the reward matrix."""
        return self._reward_matrix

    def reset(self):
        """Reset the game."""
        self._rng = None

    def step(self, actions: Sequence[ActionType]) -> npt.NDArray[RewardType]:
        """Execute a step in the game."""
        return self.reward_matrix[:, *actions]

    def view(
        self,
        agent_id: int,
        only_viewer: bool = True,
    ) -> npt.NDArray:
        """Return a view of the game for the given agent."""
        axes = np.asarray(
            list(range(agent_id, self.n_agents)) + list(range(0, agent_id)), dtype=int
        )
        reward_matrix = self.reward_matrix[axes].transpose(0, *(axes + 1))
        if only_viewer:
            return reward_matrix[0]
        return reward_matrix

    def as_markov_game(self) -> MarkovGame:
        """Return a Markov game with the given transition matrix."""
        reward_matrix = np.expand_dims(self.reward_matrix, axis=1)
        transition_matrix = np.ones((1, *self.reward_matrix.shape[1:], 1))
        return MarkovGame(
            self.name,
            reward_matrix,
            transition_matrix,
            self.seed,
            **self.metadata,
        )

    def to_json(self) -> dict:
        """Return a JSON representation of the game."""
        return {
            "name": self.name,
            "seed": self.seed,
            "metadata": self.metadata,
            "reward_matrix": self.reward_matrix.tolist(),
        }

    @classmethod
    def from_json(cls, path: Path):
        """Load the game from a JSON representation."""
        with open(path, "r", encoding="utf-8") as file:
            data = json.load(file)
        return cls(
            data["name"],
            np.asarray(data["reward_matrix"]),
            data["seed"],
            **data["metadata"],
        )


class MarkovGame(MatrixGame):
    """Base class for Markov games."""

    def __init__(
        self,
        name: str,
        reward_matrix: npt.NDArray,
        transition_matrix: npt.NDArray[np.float_],
        seed: int,
        **metadata,
    ):
        if reward_matrix.shape[1:] != transition_matrix.shape[:-1]:
            raise ValueError("The reward and transition matrices are not compatible.")
        super().__init__(name, reward_matrix, seed, **metadata)
        self._transition_matrix: npt.NDArray[np.float_] = transition_matrix

    @property
    def state(self) -> StateType:
        """Return the current state."""
        return self._state

    @state.setter
    def state(self, state: StateType):
        """Set the current state."""
        self._state = state

    @property
    def n_states(self) -> int:
        """Return the number of states."""
        return self.transition_matrix.shape[0]

    @property
    def transition_matrix(self) -> npt.NDArray:
        """Return the transition matrix."""
        return self._transition_matrix

    def reset(self):
        """Reset the game."""
        super().reset()
        self.state = self.rng.integers(self.n_states)
        return self.state

    def step(
        self, actions: Sequence[ActionType]
    ) -> npt.NDArray[RewardType] | tuple[StateType, npt.NDArray[RewardType]]:
        """Execute a step in the game."""
        reward = self.reward_matrix[:, self.state, *actions]
        self.state = self.rng.choice(
            self.n_states, p=self.transition_matrix[self.state, *actions]
        )
        return self.state, reward

    def view(
        self, agent_id: int, only_viewer: bool = True
    ) -> tuple[npt.NDArray, npt.NDArray]:
        """Return a view of the game for the given agent."""
        axes = np.asarray(
            list(range(agent_id, self.n_agents)) + list(range(0, agent_id)), dtype=int
        )
        reward_matrix = self.reward_matrix[axes].transpose(0, 1, *(axes + 2))
        transition_matrix = self.transition_matrix.transpose(0, *(axes + 1), -1)
        if only_viewer:
            return reward_matrix[0], transition_matrix
        return reward_matrix, transition_matrix

    def to_json(self) -> dict:
        """Return a JSON representation of the game."""
        return {
            "name": self.name,
            "seed": self.seed,
            "metadata": self.metadata,
            "reward_matrix": self.reward_matrix.tolist(),
            "transition_matrix": self.transition_matrix.tolist(),
        }

    @classmethod
    def from_json(cls, path: Path):
        """Load the game from a JSON representation."""
        with open(path, "r", encoding="utf-8") as file:
            data = json.load(file)
        return cls(
            data["name"],
            np.asarray(data["reward_matrix"]),
            np.asarray(data["transition_matrix"]),
            data["seed"],
            **data["metadata"],
        )

    def as_markov_game(self) -> MarkovGame:
        """Return a Markov game with the given transition matrix."""
        return self.clone()

coaction/games/game.py
yigityalin/coaction
"""Implementations of loggers."""

from coaction.loggers.logger import Logger
from coaction.loggers.loader import LogLoader
from coaction.loggers.agent import AgentLogger
from coaction.loggers.game import GameLogger
from coaction.loggers.progress import ProgressLogger


__all__ = ["Logger", "LogLoader", "AgentLogger", "GameLogger", "ProgressLogger"]

coaction/loggers/__init__.py
yigityalin/coaction
"""Implementation of the default agent logger."""

from collections import defaultdict
from copy import deepcopy
from typing import Any

import numpy as np

from coaction.agents.agent import Agent
from coaction.utils.io import save_object
from coaction.utils.paths import ProjectPaths


# TODO: Add np.memmap support
class AgentLogger:
    """Default agent logger.

    Logs the agent's logged parameters at the end of steps and writes the logs to a file.
    """

    def __init__(
        self,
        paths: ProjectPaths,
        log_each: int = 1,
        save_in_chunks: int | None = None,
    ):
        """Initialize the agent logger.

        Args:
            project_dir (Path | str): The path to the project directory.
            project_name (str): The name of the project.
            experiment_name (str): The name of the experiment.
            log_each (int, optional): The number of stages between each log. Defaults to 1.
            save_in_chunks (int | None, optional): The number of episodes to save in each
                chunk. Defaults to None. If None, all episodes are saved in a single file.
                If not None, the episodes are saved in chunks of size `save_in_chunks`.
        """
        self.paths = paths
        self._log_each: int = log_each
        self._logged_params: list[defaultdict[str, list[Any]]]
        self._save_in_chunks: int | None = save_in_chunks
        if self._save_in_chunks is not None:
            if self._save_in_chunks < 1:
                raise ValueError(
                    f"save_in_chunks must be greater than or equal to 1, got {self._save_in_chunks}"
                )
            self._current_chunk: int = 0
        else:
            self._current_chunk: int | None = None
        self._episode = 0

    def on_experiment_begin(self, agents: list[Agent]):
        """Called when an experiment begins."""
        for agent in agents:
            save_object(
                agent.buffer,
                self.paths.get_agent_log_path(agent.name, "buffer"),
            )

    def on_experiment_end(self, agents: list[Agent]):
        """Called when an experiment ends."""

    def on_episode_begin(
        self, episode: int, agents: list[Agent]
    ):  # pylint: disable=unused-argument
        """Called when an episode begins."""
        self._reset(agents)
        self._episode = episode

    def on_episode_end(
        self, episode: int, agents: list[Agent]
    ):  # pylint: disable=unused-argument
        """Called when an episode ends."""
        self._save_chunk(agents)

    def on_stage_begin(self, stage: int, agents: list[Agent]):
        """Called when a stage begins."""

    def on_stage_end(self, stage: int, agents: list[Agent]):
        """Called when a stage ends."""
        if stage % self._log_each == 0:
            for agent, params in zip(agents, self._logged_params):
                for key, value in agent.logged_params.items():
                    params[key].append(deepcopy(value))
            if self._save_in_chunks is not None:
                h_length = int(np.ceil(stage / self._log_each))
                if h_length % self._save_in_chunks == 0:
                    self._save_chunk(agents)

    def _save_chunk(self, agents: list[Agent]):
        self._save_agent_params(self._episode, agents)
        if self._save_in_chunks is not None:
            self._current_chunk += 1
        for params in self._logged_params:
            for key in params.keys():
                params[key].clear()

    def _reset(self, agents: list[Agent]):
        """Reset the logger."""
        self._logged_params = [defaultdict(list) for _ in agents]
        if self._current_chunk is not None:
            self._current_chunk = 0

    def _save_agent_params(self, episode: int, agents: list[Agent]):
        """Save the agent's parameters."""
        for agent, params in zip(agents, self._logged_params):
            for key, value in params.items():
                if not len(value) == 0:
                    chunk = (
                        self._current_chunk if self._current_chunk is not None else 0
                    )
                    save_object(
                        value,
                        self.paths.get_agent_episode_log_path(
                            episode,
                            agent.name,
                            key.lstrip("_"),
                            chunk,
                        ),
                        allow_file_exists=True,
                    )

    def clone(self):
        """Clone the logger."""
        return deepcopy(self)

coaction/loggers/agent.py
yigityalin/coaction
"""Implementations of loggers for the game module."""

from copy import deepcopy
from typing import Sequence

import numpy as np
import numpy.typing as npt

from coaction.games.game import Game, StateType, ActionType, RewardType
from coaction.utils.io import save_object
from coaction.utils.paths import ProjectPaths


# TODO: Add np.memmap support
class GameLogger:
    """Base class for all game loggers."""

    def __init__(
        self,
        agent_names: list[str],
        paths: ProjectPaths,
        log_each: int = 1,
        save_state_history: bool = False,
        save_action_history: bool = False,
        save_reward_history: bool = True,
        save_in_chunks: int | None = None,
    ):
        """Initialize the game logger.

        Args:
            game (mal.Game): The game to log.
            agent_names (list[str]): The names of the agents in the game.
            paths (ProjectPaths): The paths to the project directories.
            log_each (int, optional): The number of stages between each log. Defaults to 1.
            save_state_history (bool, optional): Whether to save the state history.
                Defaults to False.
            save_action_history (bool, optional): Whether to save the action history.
                Defaults to False.
            save_reward_history (bool, optional): Whether to save the reward history.
                Defaults to True.
            save_in_chunks (int | None, optional): The number of episodes to save in each
                chunk. Defaults to None. If None, all episodes are saved in a single file.
                If not None, the episodes are saved in chunks of size `save_in_chunks`.
        """
        self.paths = paths
        self._agent_names = agent_names
        self._log_each = log_each
        self._save_reward_history = save_reward_history
        self._save_state_history = save_state_history
        self._save_action_history = save_action_history

        self._save_in_chunks: int | None = save_in_chunks
        if self._save_in_chunks is not None:
            if self._save_in_chunks < 1:
                raise ValueError(
                    f"save_in_chunks must be greater than or equal to 1, got {self._save_in_chunks}"
                )
            self._current_chunk: int = 0
        else:
            self._current_chunk: int | None = None
        self._episode = 0

        self._state_history: list[StateType]
        self._action_history: dict[str, list[ActionType]]
        self._reward_history: dict[str, list[RewardType]]

    def on_experiment_begin(self, game: Game):
        """Called when an experiment begins."""
        save_object(
            game.to_json(),
            self.paths.get_game_log_path("game"),
        )

    def on_experiment_end(self, game: Game):
        """Called when an experiment ends."""

    def on_episode_begin(
        self, episode: int, game: Game
    ):  # pylint: disable=unused-argument
        """Called when an episode begins."""
        self._reset()
        self._episode = episode

    def on_episode_end(
        self, episode: int, game: Game
    ):  # pylint: disable=unused-argument
        """Called when an episode ends."""
        self._save_histories()

    def on_stage_begin(self, stage: int, game: Game):
        """Called when a stage begins."""

    def on_stage_end(
        self,
        stage: int,
        game: Game,
        state: StateType,
        actions: Sequence[ActionType],
        rewards: npt.NDArray[RewardType],
    ):  # pylint: disable=unused-argument
        """Called when a stage ends."""
        if stage % self._log_each == 0:
            if self._save_state_history:
                self._state_history.append(state)
            for agent_name, action, reward in zip(self._agent_names, actions, rewards):
                if self._save_action_history:
                    self._action_history[agent_name].append(action)
                if self._save_reward_history:
                    self._reward_history[agent_name].append(reward)
            if self._save_in_chunks is not None:
                h_length = int(np.ceil(stage / self._log_each))
                if h_length % self._save_in_chunks == 0:
                    self._save_histories()

    def _reset(self):
        """Reset the logger."""
        if self._save_state_history:
            self._state_history = []
        if self._save_action_history:
            self._action_history = {name: [] for name in self._agent_names}
        if self._save_reward_history:
            self._reward_history = {name: [] for name in self._agent_names}
        if self._current_chunk is not None:
            self._current_chunk = 0

    def _save_histories(self):
        """Save the histories."""
        if self._save_state_history:
            save_object(
                self._state_history,
                self.paths.get_game_episode_log_path(self._episode, "states"),
                allow_file_exists=True,
            )
            self._state_history.clear()
        if self._save_action_history:
            self._save_history(self._action_history, self._episode, "actions")
            for key in self._action_history.keys():
                self._action_history[key].clear()
        if self._save_reward_history:
            self._save_history(self._reward_history, self._episode, "rewards")
            for key in self._reward_history.keys():
                self._reward_history[key].clear()
        if self._save_in_chunks is not None:
            self._current_chunk += 1

    def _save_history(self, history: dict[str, list], episode: int, log_name: str):
        """Save a history."""
        for agent_name, hist in history.items():
            if not len(hist) == 0:
                chunk = self._current_chunk if self._current_chunk is not None else 0
                save_object(
                    hist,
                    self.paths.get_agent_episode_log_path(
                        episode, agent_name, log_name, chunk
                    ),
                    allow_file_exists=True,
                )

    def clone(self):
        """Clone the logger."""
        return deepcopy(self)

coaction/loggers/game.py
yigityalin/coaction
"""Implements the LogLoader class, which loads logs of a project run."""

from pathlib import Path

import numpy as np

from coaction.experiments.config import ExperimentConfig
from coaction.utils.io import find_and_load_object
from coaction.utils.modules import load_module
from coaction.utils.paths import ProjectPaths


# TODO: Add np.memmap support
class LogLoader:
    """Load the logs from a project."""

    def __init__(self, project_path: Path | str) -> None:
        """Initialize the log loader.

        Args:
            project_path (Path | str): The path to the project.
        """
        self.project_path = Path(project_path).resolve()
        self.paths = ProjectPaths(self.project_path)
        self._loader = find_and_load_object

    def get_project_paths(self, run: int, experiment_name: str):
        """Get the paths to the project.

        Args:
            run (int): The run number.
            experiment_name (str): The name of the experiment.
        """
        paths = self.paths
        if run is not None:
            paths = paths.with_run(run)
        if experiment_name is not None:
            paths = paths.with_experiment_name(experiment_name)
        return paths

    def load_agent_log(
        self, run: int, experiment_name: str, agent_name: str, log_name: str
    ):
        """Load an agent's log.

        Args:
            experiment_name (str): The name of the experiment.
            agent_name (str): The name of the agent.
            log_name (str): The name of the log.
        """
        return self._loader(
            self.paths.with_run(run)
            .with_experiment_name(experiment_name)
            .get_agent_log_path(agent_name, log_name)
        )

    def load_game_log(self, run: int, experiment_name: str, log_name: str):
        """Load a game's log.

        Args:
            experiment_name (str): The name of the experiment.
            log_name (str): The name of the log.
        """
        return self._loader(
            self.paths.with_run(run)
            .with_experiment_name(experiment_name)
            .get_game_log_path(log_name)
        )

    def load_episode_log(
        self,
        run: int,
        experiment_name: str,
        episode: int,
        agent_name: str,
        log_name: str,
        chunk: int | None = None,
    ):
        """Load an episode's log.

        Args:
            experiment_name (str): The name of the experiment.
            episode (int): The episode number.
            agent_name (str): The name of the agent.
            log_name (str): The name of the log.
            chunk (int | None, optional): The chunk number. Defaults to None. If None,
                the entire log is loaded. If not None, the chunk is loaded.
        """
        if chunk is not None:
            return self._loader(
                self.paths.with_run(run)
                .with_experiment_name(experiment_name)
                .get_agent_episode_log_path(episode, agent_name, log_name, chunk)
            )
        log_files = (
            self.paths.with_run(run)
            .with_experiment_name(experiment_name)
            .get_agent_episode_log_paths(episode, agent_name, log_name)
        )
        logs = [self._loader(log_file) for log_file in log_files]
        if len(logs) == 0:
            raise ValueError(
                f"Could not find log for episode {episode} of agent {agent_name} in experiment {experiment_name}"
            )
        if len(logs) == 1:
            return logs[0]
        return np.concatenate(logs)

    def load_experiment_logs(
        self, run: int, experiment_name: str, agent_name: str, log_name: str
    ):
        """Load an experiment's log.

        Args:
            experiment_name (str): The name of the experiment.
            log_name (str): The name of the log.
        """
        config = self.load_experiment_config(run, experiment_name)
        logs = [
            self.load_episode_log(run, experiment_name, episode, agent_name, log_name)
            for episode in range(config.total_episodes)
        ]
        return np.array(logs)

    def load_experiment_config(
        self, run: int, experiment_name: str
    ) -> ExperimentConfig:
        """Load the experiment config.

        Args:
            run (int): The run number.
            experiment_name (str): The name of the experiment.
        """
        config_path = (
            self.paths.with_run(run)
            .with_experiment_name(experiment_name)
            .get_project_run_config_path()
        )
        return load_module(config_path)

coaction/loggers/loader.py
yigityalin/coaction
"""Implementation of the base logger class."""

from typing import Protocol


class Logger(Protocol):
    """Base class for all loggers."""

    def on_experiment_begin(self, *args, **kwargs):
        """Called when an experiment begins."""

    def on_experiment_end(self, *args, **kwargs):
        """Called when an experiment ends."""

    def on_episode_begin(self, episode: int, *args, **kwargs):
        """Called when an episode begins."""

    def on_episode_end(self, episode: int, *args, **kwargs):
        """Called when an episode ends."""

    def on_stage_begin(self, stage: int, *args, **kwargs):
        """Called when a stage begins."""

    def on_stage_end(self, stage: int, *args, **kwargs):
        """Called when a stage ends."""

coaction/loggers/logger.py
yigityalin/coaction
"""Implementation of Progress logger"""

from copy import deepcopy
from io import TextIOWrapper

from coaction.experiments.config import ExperimentConfig
from coaction.loggers.logger import Logger
from coaction.utils.paths import ProjectPaths
from coaction.utils.progress import Progress


# TODO: implement logging to a queue
class ProgressLogger(Logger):
    """Logs progress of an experiment."""

    def __init__(
        self,
        config: ExperimentConfig,
        paths: ProjectPaths,
        log_each: int,
    ):
        """Initialize the progress logger.

        Args:
            config (ExperimentConfig): The configuration for the experiment.
            paths (ProjectPaths): The paths to the project directories.
            log_each (int): The number of episodes to log.
            queue (mp.Queue): The queue to log to.
        """
        super().__init__()
        self._config = config
        self._paths = paths
        self._log_each = log_each

        self._exp_pbar = Progress(
            initial=0, total=config.total_episodes, prefix=f"Experiment {config.name}"
        )
        self._exp_file: TextIOWrapper

        self._episode_pbars: list[Progress] = [  # type: ignore
            None for _ in range(config.total_episodes)
        ]
        self._episode_files: list[TextIOWrapper] = [  # type: ignore
            None for _ in range(config.total_episodes)
        ]

    def on_experiment_begin(self, *args, **kwargs):
        """Called when an experiment begins."""
        self._paths.get_experiment_log_dir().mkdir(parents=True, exist_ok=True)
        self._exp_file = self._paths.get_experiment_progress_log_path().open(
            "w+", encoding="utf-8"
        )
        self._exp_pbar.start()
        self._log_exp_text(f"Experiment {self._config.name} started.")
        self._display_exp(flush=True)

    def on_experiment_end(self, *args, **kwargs):
        """Called when an experiment ends."""
        self._log_exp_text(f"Experiment {self._config.name} ended.")
        self._exp_pbar.update(self._config.total_episodes - self._exp_pbar.current)
        self._display_exp(flush=True)
        self._exp_file.close()

    def on_episode_begin(self, episode: int, *args, **kwargs):
        """Called when an episode begins."""
        self._paths.get_episode_log_dir(episode).mkdir(parents=True, exist_ok=True)
        self._episode_files[episode] = self._paths.get_episode_progress_log_path(
            episode
        ).open("w", encoding="utf-8")
        self._episode_pbars[episode] = Progress(
            initial=0,
            total=self._config.total_stages,
            prefix=f"Episode {episode}",
        )
        self._episode_pbars[episode].start()
        self._log_episode_text(episode, f"Episode {episode} started.")
        self._display_episode(episode, flush=True)

    def on_episode_end(self, episode: int, *args, **kwargs):
        """Called when an episode ends."""
        self._log_episode_text(episode, f"Episode {episode} ended.")
        self._episode_pbars[episode].update(
            self._config.total_stages - self._episode_pbars[episode].current
        )
        self._display_episode(episode, flush=True)
        self._episode_files[episode].close()

    def on_stage_end(self, stage: int, *args, **kwargs):
        """Called when a stage ends."""
        if stage % self._log_each == 0:
            episode = kwargs["episode"]
            self._episode_pbars[episode].update(self._log_each)
            self._display_episode(episode, flush=True)

    def _log_exp_text(self, text: str):
        """Log text to experiment file."""
        self._exp_pbar.add_text(text)

    def _log_episode_text(self, episode: int, text: str):
        """Log text to episode file."""
        self._episode_pbars[episode].add_text(text)

    def _display_exp(self, flush: bool = True):
        display = self._exp_pbar.display()
        self._exp_file.seek(0)
        self._exp_file.write(display)
        if flush:
            self._exp_file.flush()

    def _display_episode(self, episode: int, flush: bool = True):
        display = self._episode_pbars[episode].display()
        self._episode_files[episode].seek(0)
        self._episode_files[episode].write(display)
        if flush:
            self._episode_files[episode].flush()

    def __deepcopy__(self, memo):
        """Return a deepcopy of the logger."""
        logger = ProgressLogger(
            config=self._config,
            paths=self._paths,
            log_each=self._log_each,
        )
        logger._exp_pbar = self._exp_pbar
        logger._exp_file = self._exp_file
        return logger

    def clone(self):
        """Return a clone of the logger."""
        return deepcopy(self)

coaction/loggers/progress.py
yigityalin/coaction
"""Main entry point for the coaction command line tool.

This module is the main entry point for the coaction command line tool.
It parses the command line arguments and runs the project."""

import argparse

from coaction.experiments.config import ProjectConfig
from coaction.experiments.project import Project


def _parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument("--project", help="Path to the project config directory.")
    return parser.parse_args()


def main(args: argparse.Namespace):
    """Run the mal command line tool.

    Args:
        args (argparse.Namespace): The command line arguments.
    """

    config = ProjectConfig.from_dir(args.project)
    project = Project(config)
    project.run()


if __name__ == "__main__":
    main(_parse_args())

coaction/run.py
yigityalin/coaction
"""Utilities for the coaction project."""

from coaction.utils import io
from coaction.utils import games
from coaction.utils import math
from coaction.utils import modules
from coaction.utils import paths
from coaction.utils import progress
from coaction.utils import time


__all__ = ["io", "games", "math", "modules", "paths", "progress", "time"]

coaction/utils/__init__.py
yigityalin/coaction
"""Utility functions for games."""

from typing import Sequence

import numpy as np
import numpy.typing as npt
import scipy.optimize as opt

from coaction.games.game import RewardType


def continuation_payoff(
    value_function: npt.NDArray[RewardType],
    reward_matrix: npt.NDArray[RewardType],
    transition_matrix: npt.NDArray[RewardType],
    gamma: float,
) -> np.ndarray:
    """Return the continuation payoff.

    Parameters
    ----------
    value_function : np.ndarray
        The value function.
    reward_matrix : np.ndarray
        The reward matrix.
    transition_matrix : np.ndarray
        The transition matrix.
    gamma : float
        The discount factor.

    Returns
    -------
    np.ndarray
        The continuation payoff.
    """
    return reward_matrix + gamma * transition_matrix @ value_function


def solve_minimax(arr: npt.NDArray[RewardType]):
    """Solve the linear program for the minimax theorem.

    Parameters
    ----------
    arr : npt.NDArray[RewardType]
        The payoff matrix.
    """
    n_rows, n_cols = arr.shape

    # coefficient vector
    c = np.zeros(1 + n_rows)  # pylint: disable=invalid-name
    c[-1] = -1

    # upper bound matrix
    A_ub = np.hstack((-arr.T, np.ones((n_cols, 1))))  # pylint: disable=invalid-name

    # upper bound vector
    b_ub = np.zeros((n_cols, 1))  # pylint: disable=invalid-name

    # equality constraint matrix
    A_eq = np.ones((1, n_rows + 1))  # pylint: disable=invalid-name
    A_eq[0, -1] = 0

    # equality constraint vector
    b_eq = 1

    # bounds
    bounds = [(0, None)] * n_rows + [(None, None)]

    # solve the linear program
    result = opt.linprog(c, A_ub, b_ub, A_eq, b_eq, bounds)
    return result.x[-1]


def shapley_value(
    T: npt.NDArray[RewardType],  # pylint: disable=invalid-name
    R: npt.NDArray[RewardType],  # pylint: disable=invalid-name
    gamma: float,
    n_iterations=100,
):
    """Return the Shapley value of a Markov game.

    Parameters
    ----------
    T : npt.NDArray[RewardType]
        The transition matrix.
    R : npt.NDArray[RewardType]
        The reward matrix.
    gamma : float
        The discount factor.
    n_iterations : int, optional
        The number of iterations to run, by default 100
    """
    n_states = R.shape[0]
    Q = np.zeros_like(R)  # pylint: disable=invalid-name
    v = np.zeros(n_states)  # pylint: disable=invalid-name

    for _ in range(n_iterations):
        for state in range(n_states):
            v[state] = solve_minimax(Q[state])
        Q = continuation_payoff(v, R, T, gamma)  # pylint: disable=invalid-name
    return v


def generate_zsmg(
    n_states, n_actions: Sequence[int], min_reward=0.0, max_reward=1.0, seed=None
) -> tuple[npt.NDArray[RewardType], npt.NDArray[RewardType]]:
    """Generate a two-player zero-sum Markov game.

    Parameters
    ----------
    n_states : int
        The number of states.
    n_actions : Sequence[int]
        The number of actions for each player.
    min_reward : float, optional
        The minimum reward for agent 0, by default 0.0
    max_reward : float, optional
        The maximum reward for agent 0, by default 1.0
    seed : int, optional
        The seed for the random number generator, by default None

    Returns
    -------
    tuple[npt.NDArray[RewardType], npt.NDArray[RewardType]]
        The reward matrix and the transition matrix.

    """
    rng = np.random.default_rng(seed=seed)
    R = rng.uniform(  # pylint: disable=invalid-name
        min_reward, max_reward, size=(n_states, *n_actions)
    )
    T = rng.dirichlet(  # pylint: disable=invalid-name
        np.ones(n_states), size=(n_states, *n_actions)
    )
    return np.stack((R, -R)), T

coaction/utils/games.py
yigityalin/coaction
"""Utilities for IO operations."""

from collections.abc import Sequence
from functools import singledispatch
from pathlib import Path
import json
import pickle

import numpy as np


def _check_save_path(path: Path, allow_file_exists: bool = False) -> Path:
    """Check if a save path is valid.

    This function will create the parent directory if it does not exist.

    Parameters
    ----------
    path : Path
        The path to the file.
    allow_file_exists : bool
        Whether to allow overwriting an existing file.
    """
    path = Path(path).resolve()
    if not allow_file_exists and path.exists():
        raise FileExistsError(f"Path {path} already exists.")
    if not path.parent.exists():
        path.parent.mkdir(parents=True)
    return path


def _check_load_path(path: Path) -> Path:
    """Check if a load path is valid.

    Parameters
    ----------
    path : Path
        The path to the file.
    """
    path = Path(path).resolve()
    if not path.exists():
        raise FileNotFoundError(f"Path {path} does not exist.")
    if not path.is_file():
        raise FileNotFoundError(f"Path {path} is not a file.")
    if not path.suffix:
        raise ValueError(f"Path {path} has no suffix.")
    return path


@singledispatch
def save_object(obj, path: Path, allow_file_exists: bool = False):
    """Save an object to a file.

    Parameters
    ----------
    obj : Any
        The object to save.
    path : Path
        The path to the file.
    allow_file_exists : bool
        Whether to allow overwriting an existing file.
    """
    if not path.suffix:
        path = path.with_suffix(".pkl")
    path = _check_save_path(path, allow_file_exists=allow_file_exists)
    with open(path, "wb") as file:
        pickle.dump(obj, file)


@save_object.register
def _(obj: dict, path: Path, allow_file_exists: bool = False):
    """Save a dictionary to a file.

    This function will save the dictionary as a JSON file.

    Parameters
    ----------
    obj : dict
        The dictionary to save.
    path : Path
        The path to the file.
    allow_file_exists : bool
        Whether to allow overwriting an existing file.
    """
    if not path.suffix:
        path = path.with_suffix(".json")
    path = _check_save_path(path, allow_file_exists=allow_file_exists)
    with open(path, "w", encoding="utf-8") as file:
        json.dump(obj, file, indent=4)


@save_object.register
def _(obj: Sequence | np.ndarray, path: Path, allow_file_exists: bool = False):
    """Save a sequence to a file.

    This function will save the sequence as a NumPy file.

    Parameters
    ----------
    obj : Sequence | np.ndarray
        The sequence to save.
    path : Path
        The path to the file.
    allow_file_exists : bool
        Whether to allow overwriting an existing file.
    """
    if not path.suffix:
        path = path.with_suffix(".npy")
    path = _check_save_path(path, allow_file_exists=allow_file_exists)
    np.save(path, obj)


def find_and_load_object(path: Path):
    """Find and load an object from a file.

    This function will try to load an object from a file with the following suffixes:
    - .json
    - .npy
    - .pkl

    Parameters
    ----------
    path : Path
        The path to the file.
    """
    for suffix in [".json", ".npy", ".pkl"]:
        path = path.with_suffix(suffix)
        if path.exists():
            return load_object(path)
    raise FileNotFoundError(f"Path {path} does not exist.")


def load_object(path: Path):
    """Load an object from a file.

    Parameters
    ----------
    path : Path
        The path to the file.
    """
    path = _check_load_path(path)
    match path.suffix:
        case ".json":
            return load_json(path)
        case ".npy":
            return load_npy(path)
        case ".pkl":
            return load_pkl(path)
        case _:
            raise ValueError(f"Path {path} has an invalid suffix.")


def load_json(path: Path) -> dict:
    """Load a JSON file.

    Parameters
    ----------
    path : Path
        The path to the file.
    """
    with open(path, "r", encoding="utf-8") as file:
        return json.load(file)


def load_npy(path: Path) -> np.ndarray:
    """Load a NumPy file.

    Parameters
    ----------
    path : Path
        The path to the file.
    """
    return np.load(path)


def load_pkl(path: Path):
    """Load a pickle file.

    Parameters
    ----------
    path : Path
        The path to the file.
    """
    with open(path, "rb") as file:
        return pickle.load(file)

coaction/utils/io.py
yigityalin/coaction
"""Math utilities."""

import numpy as np
import numpy.typing as npt


def softmax(arr: npt.NDArray[np.float_], tau: float):
    """Compute the softmax of an array.

    Parameters
    ----------
    arr : npt.NDArray[float]
        Array to compute the softmax of.
    tau : float
        Temperature parameter.

    Returns
    -------
    npt.NDArray[float]
        Softmax of `arr`.
    """
    arr = arr - np.max(arr, axis=-1, keepdims=True)
    exp = np.exp(arr / tau)
    return exp / np.sum(exp, axis=-1, keepdims=True)

coaction/utils/math.py
yigityalin/coaction
"""Utilities for modules."""

from pathlib import Path
import importlib.util


def load_module(module_path: Path | str):
    """Load a module from a file."""
    module_path = Path(module_path).resolve()
    spec = importlib.util.spec_from_file_location(module_path.stem, module_path)
    if spec is None or spec.loader is None:
        raise ValueError(f"Could not load module from {module_path}")
    config = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(config)
    return config

coaction/utils/modules.py
yigityalin/coaction
"""Path utilities.

This module contains utilities for paths. The :class:`ProjectPaths` class contains
methods for getting paths to various directories and files.
"""

from __future__ import annotations

from pathlib import Path
from typing import final


class ProjectPaths:
    """Path utilities for the coaction project."""

    def __init__(
        self,
        project_dir: Path | str,
        experiment_name: str | None = None,
        increment_run: bool = False,
    ) -> None:
        """Initialize the logger paths.

        Args:
            project_dir (Path | str): The path to the project directory.
            project_name (str): The name of the project.
            experiment_name (str | None): The name of the experiment.
        """
        self.project_dir: Path = Path(project_dir).resolve()
        self.project_name: str = self.project_dir.stem
        self.experiment_name: str | None = experiment_name
        self.run: int = self._get_run()
        if increment_run:
            self._get_project_run_log_dir(self.run).mkdir(parents=True, exist_ok=True)
        elif self.run != 0:
            self.run -= 1

    def __repr__(self) -> str:
        """Return the string representation of the logger paths."""
        return (
            f"{self.__class__.__name__}(project_name={self.project_name}, "
            f"experiment_name={self.experiment_name})"
        )

    def _get_run(self) -> int:
        run = 0
        while self._get_project_run_log_dir(run).exists():
            run += 1
        return run

    def _get_project_run_log_dir(self, run: int) -> Path:
        """Return the path to a project run log directory.

        Args:
            run (int): The run number.
        """
        return self.get_project_logs_dir() / f"run_{run}"

    def cleanup(self) -> None:
        """Remove the project run log directory."""
        self.get_project_run_log_dir().rmdir()
        project_logs_dir = self.get_project_logs_dir()
        if project_logs_dir.exists() and not list(project_logs_dir.iterdir()):
            project_logs_dir.rmdir()
        project_dir = self.get_project_dir()
        if project_dir.exists() and not list(project_dir.iterdir()):
            project_dir.rmdir()

    @classmethod
    def from_parent(
        cls,
        parent_dir: Path | str,
        project_name: str,
        experiment_name: str | None = None,
    ) -> "ProjectPaths":
        """Return the logger paths from a parent directory.

        Args:
            parent_dir (Path | str): The path to the parent directory.
            project_name (str): The name of the project.
            experiment_name (str | None): The name of the experiment.
        """
        project_dir = Path(parent_dir).resolve() / project_name
        return cls(project_dir, experiment_name)

    @final
    @staticmethod
    def get_templates_dir() -> Path:
        """Return the path to the templates directory.

        The templates directory is the directory where all the templates are stored.
        """
        return Path(__file__).parent.parent / "_templates"

    @final
    @staticmethod
    def get_agent_template_path() -> Path:
        """Return the path to the agent template file.

        The agent template file is the file that contains the template for an agent.
        """
        return ProjectPaths.get_templates_dir() / "_agent.txt"

    @final
    @staticmethod
    def get_project_config_template_path() -> Path:
        """Return the path to the project config template file.

        The project config template file is the file containing the template for a project config.
        """
        return ProjectPaths.get_templates_dir() / "_project_config.txt"

    @final
    def get_project_dir(self) -> Path:
        """Return the path to a project directory.

        The project directory is the directory where all the data for a single project are stored.
        """
        return self.project_dir

    @final
    def get_project_config_dir(self) -> Path:
        """Return the path to a project config directory.

        The project config directory is the directory
        where all the configs for a single project are stored.
        """
        return self.get_project_dir() / "configs"

    @final
    def get_project_config_path(self) -> Path:
        """Return the path to a project config file.

        The project config file is the file that contains the configuration for a project.
        """
        return self.get_project_config_dir() / f"{self.project_name}.py"

    @final
    def get_experiment_config_path(self) -> Path:
        """Return the path to an experiment config file.

        The experiment config file is the file that contains the configuration for an experiment.
        """
        return self.get_project_config_dir() / f"{self.experiment_name}.py"

    @final
    def get_experiment_config_paths(self) -> list[Path]:
        """Return the paths to all the experiment config files.

        The experiment config file is the file that contains the configuration for an experiment.
        """
        return [
            path
            for path in self.get_project_config_dir().iterdir()
            if path.is_file()
            and not path.stem.startswith("_")
            and path.suffix == ".py"
            and path.stem != self.project_name
        ]

    @final
    def get_experiment_run_config_paths(self) -> list[Path]:
        """Return the paths to all the experiment config files.

        The experiment config file is the file that contains the configuration for an experiment.
        """
        return [
            path
            for path in self.get_project_run_config_dir().iterdir()
            if path.is_file()
            and not path.stem.startswith("_")
            and path.suffix == ".py"
            and path.stem != self.project_name
        ]

    @final
    def get_experiment_names(self) -> list[str]:
        """Return the names of all the experiments."""
        return [path.stem for path in self.get_experiment_config_paths()]

    @final
    def get_run_experiment_names(self) -> list[str]:
        """Return the names of all the experiments."""
        return [path.stem for path in self.get_experiment_run_config_paths()]

    @final
    def get_project_logs_dir(self) -> Path:
        """Return the path to a project logs directory.

        The project log directory is the directory
        where all the logs for all the project runs are stored.
        """
        return self.get_project_dir() / "logs"

    @final
    def get_project_run_log_dir(self) -> Path:
        """Return the path to a project log directory.

        The project log directory is the directory
        where all the logs for a single project run is stored.
        """
        return self._get_project_run_log_dir(self.run)

    @final
    def get_project_run_progress_log_path(self) -> Path:
        """Return the path to a project's progress log file."""
        return self.get_project_run_log_dir() / "progress.log"

    @final
    def get_project_run_config_dir(self) -> Path:
        """Return the path to a project run config directory."""
        return self.get_project_run_log_dir() / "configs"

    @final
    def get_project_run_config_path(self) -> Path:
        """Return the path to a project run config file."""
        return self.get_project_run_config_dir() / f"{self.experiment_name}.py"

    @final
    def get_experiment_log_dirs(self) -> list[Path]:
        """Return the paths to all the experiment log directories.

        The experiment log directory is the directory where all the logs for a single experiment.
        """
        return sorted(self.get_project_run_log_dir().iterdir())

    @final
    def get_experiment_log_dir(self) -> Path:
        """Return the path to an experiment directory.

        The experiment directory is the directory where all the logs for a single experiment.
        """
        if self.experiment_name is None:
            raise ValueError("Experiment name must be set.")
        return self.get_project_run_log_dir() / self.experiment_name

    @final
    def get_experiment_progress_log_path(self) -> Path:
        """Return the path to an experiment's progress log file.

        The progress log file is the file that contains the progress of an experiment.
        """
        return self.get_experiment_log_dir() / "progress.log"

    @final
    def get_agents_log_dir(self) -> Path:
        """Return the path to an episode's agents directory.

        The agents directory is the directory where all the logs for agents are stored.
        """
        return self.get_experiment_log_dir() / "agents"

    @final
    def get_agent_log_dir(self, agent_name: str) -> Path:
        """Return the path to an agent directory.

        The agent directory is the directory where all the logs for a single agent are stored.

        Args:
            agent_name (str): The name of the agent.
        """
        return self.get_agents_log_dir() / agent_name

    @final
    def get_agent_log_path(self, agent_name: str, log_name: str) -> Path:
        """Return the path to an agent's log file.

        Args:
            agent_name (str): The name of the agent.
            log_name (str): The name of the log file.
        """
        return self.get_agent_log_dir(agent_name) / log_name

    @final
    def get_game_log_dir(self) -> Path:
        """Return the path to an episode log directory.

        The game directory is the directory where all the data for games are stored.
        """
        return self.get_experiment_log_dir() / "games"

    @final
    def get_game_log_path(self, log_name: str) -> Path:
        """Return the path to an episode's game's log file.

        Args:
            log_name (str): The name of the log file.
        """
        return self.get_game_log_dir() / log_name

    @final
    def get_episodes_log_dir(self) -> Path:
        """Return the path to an episode log directory.

        The episode directory is the directory where all the data for episodes are stored.
        """
        return self.get_experiment_log_dir() / "episodes"

    @final
    def get_episode_log_dir(self, episode: int) -> Path:
        """Return the path to an episode directory.

        The episode directory is the directory where all the logs for a single episode are stored.

        Args:
            episode (int): The episode number.
        """
        return self.get_episodes_log_dir() / str(episode)

    @final
    def get_episode_progress_log_path(self, episode: int) -> Path:
        """Return the path to an experiment's config log file.

        The config log file is the file that contains the configuration for an experiment.
        """
        return self.get_episode_log_dir(episode) / "progress.log"

    @final
    def get_agents_episode_log_dir(self, episode: int) -> Path:
        """Return the path to an episode's agents directory.

        The agents directory is the directory where all the logs for agents are stored.

        Args:
            episode (int): The episode number.
        """
        return self.get_episode_log_dir(episode) / "agents"

    @final
    def get_agent_episode_log_dir(self, episode: int, agent_name: str) -> Path:
        """Return the path to an episode's agent's directory.

        The agent directory is the directory where all the logs for a single agent are stored.

        Args:
            episode (int): The episode number.
            agent_name (str): The name of the agent.
        """
        return self.get_agents_episode_log_dir(episode) / agent_name

    @final
    def get_agent_episode_log_path(
        self, episode: int, agent_name: str, log_name: str, chunk: int
    ) -> Path:
        """Return the path to an episode's agent's log file.

        Args:
            episode (int): The episode number.
            agent_name (str): The name of the agent.
            log_name (str): The name of the log file.
            chunk (int): The chunk number
        """
        stem = f"{log_name}_chunk_{chunk}"
        return self.get_agent_episode_log_dir(episode, agent_name) / stem

    @final
    def get_agent_episode_log_paths(self, episode: int, agent_name: str, log_name):
        """Return the paths to an episode's agent's log files.

        Args:
            episode (int): The episode number.
            agent_name (str): The name of the agent.
            log_name (str): The name of the log file.
        """
        directory = self.get_agent_episode_log_dir(episode, agent_name)
        return sorted(
            directory.glob(f"{log_name}_chunk_*"),
            key=lambda x: int(x.stem.split("_")[-1]),
        )

    @final
    def get_game_episode_log_dir(self, episode: int) -> Path:
        """Return the path to an episode's games directory.

        The games directory is the directory where all the logs for games are stored.

        Args:
            episode (int): The episode number.
        """
        return self.get_episode_log_dir(episode) / "game"

    @final
    def get_game_episode_log_path(
        self, episode: int, log_name: str, chunk: int | None = None
    ) -> Path:
        """Return the path to an episode's game's log file.

        Args:
            episode (int): The episode number.
            log_name (str): The name of the log file.
            chunk (int | None, optional): The chunk number. Defaults to None.
        """
        stem = f"{log_name}_chunk_{chunk or 0}"
        return self.get_game_episode_log_dir(episode) / stem

    @final
    def with_run(self, run: int) -> ProjectPaths:
        """Return a new ProjectPaths object with a new run number.

        Args:
            experiment_name (str): The name of the experiment.
        """
        paths = ProjectPaths(self.project_dir, self.experiment_name)
        paths.run = run
        return paths

    @final
    def with_experiment_name(self, experiment_name: str) -> ProjectPaths:
        """Return a new ProjectPaths object with a new experiment name.

        Args:
            experiment_name (str): The name of the experiment.
        """
        paths = ProjectPaths(self.project_dir, experiment_name)
        paths.run = self.run
        return paths

coaction/utils/paths.py
yigityalin/coaction
"""Progress bar for logging"""

from io import TextIOWrapper
from typing import Final
import datetime
import time

from coaction.utils.time import get_strftime


_PB_PROGRESS_SYMBOL: Final = "█"
_PB_REMAINING_SYMBOL: Final = "."


def time_to_str(time_delta: float) -> str:
    """Return a string representation of a time delta."""
    return str(datetime.timedelta(seconds=time_delta))


class Progress:
    """A progress bar for logging"""

    def __init__(
        self,
        initial: int,
        total: int,
        prefix: str,
        text_only: bool = False,
        show_bar: bool = True,
        size: int = 60,
    ):
        self._initial = initial
        self._total = total
        self._prefix = prefix
        self._text_only = text_only
        self._show_bar = show_bar
        self._size = size

        self.current: int
        self._start_time: float
        self._current_time: float
        self._file: TextIOWrapper
        self._desc: str = ""
        self._text: str = ""

    @property
    def desc(self):
        """Return the description."""
        return self._desc

    @desc.setter
    def desc(self, value):
        self._desc = value

    @property
    def text(self):
        """Return the text."""
        return self._text

    @text.setter
    def text(self, value):
        self._text = value

    def start(self):
        """Start the progress bar."""
        self.current = self._initial
        self._text = ""
        self._start_time = time.perf_counter()
        self._current_time = self._start_time

    def update(self, n: int):  # pylint: disable=invalid-name
        """Update the progress bar.

        Args:
            n (int): The number of steps to update.
        """
        self.current = min(self.current + n, self._total)
        self._current_time = time.perf_counter()

    def add_text(self, text: str):
        """Add text to the progress bar."""
        self._text += f"[{get_strftime()}] - {text}\n"

    def display(self):
        """Return the progress bar as a string."""
        if self._text_only:
            return f"{self._text}\n"
        percentage = (self.current * 100) // self._total
        fill = (self.current * self._size) // self._total
        remaining = self._size - fill
        timedelta = self._current_time - self._start_time
        time_per_step = timedelta / self.current if self.current else float("NaN")
        time_remaining = time_per_step * (self._total - self.current)
        timedelta_str = time_to_str(timedelta)
        time_remaining_str = (
            time_to_str(time_remaining) if self.current else float("NaN")
        )
        prefix = f"{self._prefix}: {percentage}%"
        postfix = f"{self.current}/{self._total} [{timedelta_str}<{time_remaining_str}]"
        if not self._show_bar:
            return f"{prefix} ({postfix}) {self.desc}\n\n{self._text}"
        pb = f"|{_PB_PROGRESS_SYMBOL * fill}{_PB_REMAINING_SYMBOL * remaining}|"  # pylint: disable=invalid-name
        return f"{prefix} {pb} {postfix} {self.desc}\n\n{self._text}"

coaction/utils/progress.py
yigityalin/coaction
"""Time utils."""

import time


def get_strftime():
    """Return a string representation of the current time."""
    return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())

coaction/utils/time.py
yigityalin/coaction