colosseum.agent.agents.episodic.boot_dqn

  1from typing import Any, Dict, Callable, TYPE_CHECKING
  2
  3import gin
  4import numpy as np
  5import sonnet as snt
  6import tensorflow as tf
  7from bsuite.baselines.tf.boot_dqn import BootstrappedDqn, make_ensemble
  8from ray import tune
  9
 10from colosseum.dynamic_programming.utils import get_policy_from_q_values
 11from colosseum.utils.non_tabular.bsuite import NonTabularBsuiteAgentWrapper
 12
 13if TYPE_CHECKING:
 14    from colosseum.agent.agents.base import BaseAgent
 15    from colosseum.utils.acme.specs import MDPSpec
 16
 17
 18@gin.configurable
 19class BootDQNEpisodic(NonTabularBsuiteAgentWrapper):
 20    """
 21    The wrapper for the `BootDQN` agent from `bsuite`.
 22    """
 23
 24    @staticmethod
 25    def produce_gin_file_from_parameters(parameters: Dict[str, Any], index: int = 0):
 26        string = ""
 27        for k, v in parameters.items():
 28            string += f"prms_{index}/BootDQNEpisodic.{k} = {v}\n"
 29        return string[:-1]
 30
 31    @staticmethod
 32    def is_episodic() -> bool:
 33        return True
 34
 35    @staticmethod
 36    def get_hyperparameters_search_spaces() -> Dict[str, tune.sample.Domain]:
 37        return {
 38            "network_width": tune.choice([64, 128, 256]),
 39            "network_depth": tune.choice([2, 4]),
 40            "batch_size": tune.choice([32, 64, 128]),
 41            "sgd_period": tune.choice([1, 4, 8]),
 42            "target_update_period": tune.choice([4, 16, 32]),
 43            "mask_prob": tune.choice([0.8, 0.9, 1.0]),
 44            "noise_scale": tune.choice([0.0, 0.05, 0.1]),
 45            "n_ensemble": tune.choice([8, 16, 20]),
 46        }
 47
 48    @staticmethod
 49    def get_agent_instance_from_parameters(
 50        seed: int,
 51        optimization_horizon: int,
 52        mdp_specs: "MDPSpec",
 53        parameters: Dict[str, Any],
 54    ) -> "BaseAgent":
 55        return BootDQNEpisodic(
 56            seed,
 57            mdp_specs,
 58            optimization_horizon,
 59            parameters["network_width"],
 60            parameters["network_depth"],
 61            parameters["batch_size"],
 62            parameters["sgd_period"],
 63            parameters["target_update_period"],
 64            parameters["mask_prob"],
 65            parameters["noise_scale"],
 66            parameters["n_ensemble"],
 67        )
 68
 69    @property
 70    def current_optimal_stochastic_policy(self) -> np.ndarray:
 71        H, S, d = self.emission_map.all_observations.shape
 72        qvals = (
 73            tf.stop_gradient(
 74                self._agent._forward[self._agent._active_head](
 75                    self.emission_map.all_observations.reshape(-1, d)
 76                )
 77            )
 78            .numpy()
 79            .reshape(H, S, self._mdp_spec.actions.num_values)
 80        )
 81        return get_policy_from_q_values(qvals, True)
 82
 83    def __init__(
 84        self,
 85        seed: int,
 86        mdp_specs: "MDPSpec",
 87        optimization_horizon: int,
 88        # MDP model parameters
 89        network_width: int,
 90        network_depth: int,
 91        batch_size: int,
 92        sgd_period: int,
 93        target_update_period: int,
 94        # Actor parameters
 95        mask_prob: float,
 96        noise_scale: float,
 97        n_ensemble: int,
 98        learning_rate: float = 1e-3,
 99        replay_capacity: int = 10000,
100        epsilon_fn: Callable[[int], float] = lambda t: 0,
101    ):
102        r"""
103        Parameters
104        ----------
105        seed : int
106            The random seed.
107        mdp_specs : MDPSpec
108            The full specification of the MDP.
109        optimization_horizon : int
110            The total number of interactions that the agent is expected to have with the MDP.
111        network_width : int
112            The width of the neural networks of the agent.
113        network_depth : int
114            The depth of the neural networks of the agent.
115        batch_size : int
116            The batch size for training the agent.
117        sgd_period : int
118            The stochastic gradient descent update period.
119        target_update_period : int
120            The interval length between updating the target network.
121        mask_prob : float
122            The masking probability for the bootstrapping procedure.
123        noise_scale : float
124            The scale of the Gaussian noise_class added to the value estimates.
125        n_ensemble : int
126            The number of ensembles.
127        learning_rate : float
128            The learning rate of the optimizer. By default, it is set to 1e-3.
129        replay_capacity : int
130            The maximum capacity of the replay buffer. By default, it is set to 10 000.
131        epsilon_fn : Callable[[int], float]]
132            The :math:`\epsilon` greedy probability as a function of the time. By default, it is set to zero.
133        """
134
135        tf.random.set_seed(seed)
136        np.random.seed(seed)
137
138        ensemble = make_ensemble(
139            mdp_specs.actions.num_values, n_ensemble, network_depth, network_width
140        )
141        optimizer = snt.optimizers.Adam(learning_rate=learning_rate)
142
143        agent = BootstrappedDqn(
144            obs_spec=mdp_specs.observations,
145            action_spec=mdp_specs.actions,
146            ensemble=ensemble,
147            batch_size=batch_size,
148            discount=1,
149            replay_capacity=replay_capacity,
150            min_replay_size=batch_size,
151            sgd_period=sgd_period,
152            target_update_period=target_update_period,
153            optimizer=optimizer,
154            mask_prob=mask_prob,
155            noise_scale=noise_scale,
156            seed=seed,
157            epsilon_fn=epsilon_fn,
158        )
159        super(BootDQNEpisodic, self).__init__(seed, agent, mdp_specs)