colosseum.agent.agents.episodic.posterior_sampling

  1from typing import Any, Callable, Dict, Union
  2
  3import dm_env
  4import gin
  5import numpy as np
  6from ray import tune
  7
  8from colosseum.agent.actors import QValuesActor
  9from colosseum.agent.agents.base import BaseAgent
 10from colosseum.agent.mdp_models.bayesian_model import BayesianMDPModel
 11from colosseum.agent.mdp_models.bayesian_models import RewardsConjugateModel
 12from colosseum.agent.mdp_models.bayesian_models import TransitionsConjugateModel
 13from colosseum.dynamic_programming import episodic_value_iteration
 14from colosseum.dynamic_programming.utils import get_policy_from_q_values
 15from colosseum.emission_maps import EmissionMap
 16from colosseum.utils.acme.specs import MDPSpec
 17
 18
 19@gin.configurable
 20class PSRLEpisodic(BaseAgent):
 21    """
 22    The posterior sampling for reinforcement learning algorithm.
 23
 24    Osband, Ian, Daniel Russo, and Benjamin Van Roy. "(More) efficient reinforcement learning via posterior sampling."
 25    Advances in Neural Information Processing Systems 26 (2013).
 26    """
 27
 28    def step_update(
 29        self,
 30        ts_t: dm_env.TimeStep,
 31        a_t: "ACTION_TYPE",
 32        ts_tp1: dm_env.TimeStep,
 33        time: int,
 34    ):
 35        super(PSRLEpisodic, self).step_update(ts_t, a_t, ts_tp1, time)
 36
 37    @staticmethod
 38    def is_emission_map_accepted(emission_map: "EmissionMap") -> bool:
 39        return emission_map.is_tabular
 40
 41    @staticmethod
 42    def produce_gin_file_from_parameters(parameters: Dict[str, Any], index: int = 0):
 43        return (
 44            "from colosseum.agent.mdp_models import bayesian_models\n"
 45            f"prms_{index}/PSRLEpisodic.reward_prior_model = %bayesian_models.RewardsConjugateModel.N_NIG\n"
 46            f"prms_{index}/PSRLEpisodic.transitions_prior_model = %bayesian_models.TransitionsConjugateModel.M_DIR\n"
 47            f"prms_{index}/PSRLEpisodic.rewards_prior_prms = [{parameters['prior_mean']}, 1, 1, 1]\n"
 48            f"prms_{index}/PSRLEpisodic.transitions_prior_prms = [{parameters['transition_prior']}]"
 49        )
 50
 51    @staticmethod
 52    def is_episodic() -> bool:
 53        return True
 54
 55    @staticmethod
 56    def get_hyperparameters_search_spaces() -> Dict[str, tune.sample.Domain]:
 57        return dict(
 58            prior_mean=tune.uniform(0.001, 2.0), transition_prior=tune.uniform(0.001, 2)
 59        )
 60
 61    @staticmethod
 62    def get_agent_instance_from_parameters(
 63        seed: int,
 64        optimization_horizon: int,
 65        mdp_specs: MDPSpec,
 66        parameters: Dict[str, Any],
 67    ) -> "BaseAgent":
 68        return PSRLEpisodic(
 69            mdp_specs=mdp_specs,
 70            seed=seed,
 71            optimization_horizon=optimization_horizon,
 72            reward_prior_model=RewardsConjugateModel.N_NIG,
 73            transitions_prior_model=TransitionsConjugateModel.M_DIR,
 74            rewards_prior_prms=[parameters["prior_mean"], 1, 1, 1],
 75            transitions_prior_prms=[parameters["transition_prior"]],
 76        )
 77
 78    @property
 79    def current_optimal_stochastic_policy(self) -> np.ndarray:
 80        T_map, R_map = self._mdp_model.get_map_estimate()
 81        Q, _ = episodic_value_iteration(self._time_horizon, T_map, R_map)
 82        return get_policy_from_q_values(Q, True)
 83
 84    def __init__(
 85        self,
 86        seed: int,
 87        mdp_specs: MDPSpec,
 88        optimization_horizon: int,
 89        # MDP model parameters
 90        reward_prior_model: RewardsConjugateModel = None,
 91        transitions_prior_model: TransitionsConjugateModel = None,
 92        rewards_prior_prms=None,
 93        transitions_prior_prms=None,
 94        # Actor parameters
 95        epsilon_greedy: Union[float, Callable] = None,
 96        boltzmann_temperature: Union[float, Callable] = None,
 97    ):
 98        """
 99
100        Parameters
101        ----------
102        seed : int
103            The random seed.
104        mdp_specs : MDPSpec
105            The full specification of the MDP.
106        optimization_horizon : int
107            The total number of interactions that the agent is expected to have with the MDP.
108        reward_prior_model : RewardsConjugateModel, optional
109            The reward priors.
110        transitions_prior_model : TransitionsConjugateModel, optional
111            The transitions priors.
112        rewards_prior_prms : Any
113            The reward prior parameters.
114        transitions_prior_prms : Any
115            The transitions prior parameters.
116        epsilon_greedy : Union[float, Callable], optional
117            The probability of selecting an action at random. It can be provided as a float or as a function of the
118            total number of interactions. By default, the probability is set to zero.
119        boltzmann_temperature : Union[float, Callable], optional
120            The parameter that controls the Boltzmann exploration. It can be provided as a float or as a function of
121            the total number of interactions. By default, Boltzmann exploration is disabled.
122        """
123
124        mdp_model = BayesianMDPModel(
125            seed,
126            mdp_specs,
127            reward_prior_model=reward_prior_model,
128            transitions_prior_model=transitions_prior_model,
129            rewards_prior_prms=rewards_prior_prms,
130            transitions_prior_prms=transitions_prior_prms,
131        )
132        actor = QValuesActor(seed, mdp_specs, epsilon_greedy, boltzmann_temperature)
133
134        super(PSRLEpisodic, self).__init__(
135            seed,
136            mdp_specs,
137            mdp_model,
138            actor,
139            optimization_horizon,
140        )
141
142    def episode_end_update(self):
143        Q, _ = episodic_value_iteration(self._time_horizon, *self._mdp_model.sample())
144        self._actor.set_q_values(Q)
145
146    def before_start_interacting(self):
147        self.episode_end_update()