colosseum.agent.utils

  1from typing import Type, TYPE_CHECKING, Any, Dict, List
  2
  3import numpy as np
  4
  5from colosseum.utils.miscellanea import rounding_nested_structure
  6
  7if TYPE_CHECKING:
  8    from colosseum.agent.agents.base import BaseAgent
  9
 10
 11def sample_agent_hyperparameters(
 12    agent_class: Type["BaseAgent"], seed: int
 13) -> Dict[str, Any]:
 14    """
 15    samples parameters from the agent class sample spaces.
 16
 17    Parameters
 18    ----------
 19    agent_class : Type["BaseAgent"]
 20        The agent class for which we are sampling from.
 21    seed : int
 22        The random seed.
 23
 24    Returns
 25    -------
 26    Dict[str, Any]
 27        The parameters sampled from the agent hyperparameter spaces.
 28    """
 29    np.random.seed(seed)
 30    search_spaces = agent_class.get_hyperparameters_search_spaces()
 31    return rounding_nested_structure({k: v.sample() for k, v in search_spaces.items()})
 32
 33
 34def sample_n_agent_hyperparameters(
 35    n: int, agent_class: Type["BaseAgent"], seed: int
 36) -> List[Dict[str, Any]]:
 37    """
 38    samples n parameters from the agent class sample spaces.
 39
 40    Parameters
 41    ----------
 42    n : int
 43        The number of samples.
 44    agent_class : Type["BaseAgent"]
 45        The agent class for which we are sampling from.
 46    seed : int
 47        The random seed.
 48
 49    Returns
 50    -------
 51    List[Dict[str, Any]]
 52        The list of n parameters sampled from the agent hyperparameter spaces.
 53    """
 54    return [sample_agent_hyperparameters(agent_class, seed + i) for i in range(n)]
 55
 56
 57def sample_agent_gin_configs(
 58    agent_class: Type["BaseAgent"], n: int = 1, seed: int = 42
 59) -> List[str]:
 60    """
 61    samples gin configurations from the agent class sample spaces.
 62
 63    Parameters
 64    ----------
 65    agent_class : Type["BaseAgent"]
 66        The agent class for which we are sampling from.
 67    n : int
 68        The number of samples. By default, it is set to one.
 69    seed : int
 70        The random seed. By default, it is set to :math:`42`.
 71
 72    Returns
 73    -------
 74    List[str]
 75        The list containing the sampled gin configs.
 76    """
 77    return [
 78        agent_class.produce_gin_file_from_parameters(params, i)
 79        for i, params in enumerate(sample_n_agent_hyperparameters(n, agent_class, seed))
 80    ]
 81
 82
 83def sample_agent_gin_configs_file(
 84    agent_class: Type["BaseAgent"], n: int = 1, seed: int = 42
 85) -> str:
 86    """
 87    samples gin configurations from the agent class sample spaces and store them in a string that can be used to create
 88    a gin config file.
 89
 90    Parameters
 91    ----------
 92    agent_class : Type["BaseAgent"]
 93        The agent class for which we are sampling from.
 94    n : int
 95        The number of samples. By default, it is set to one.
 96    seed : int
 97        The random seed. By default, it is set to :math:`42`.
 98
 99    Returns
100    -------
101    str
102        The gin configuration file.
103    """
104    return "\n".join(sample_agent_gin_configs(agent_class, n, seed))