colosseum.agent.agents.infinite_horizon.actor_critic

  1from typing import Dict, Any, TYPE_CHECKING
  2
  3import gin
  4import numpy as np
  5import sonnet as snt
  6import tensorflow as tf
  7from bsuite.baselines.tf.actor_critic import PolicyValueNet, ActorCritic
  8from ray import tune
  9
 10from colosseum.utils.non_tabular.bsuite import NonTabularBsuiteAgentWrapper
 11
 12if TYPE_CHECKING:
 13    from colosseum.agent.agents.base import BaseAgent
 14    from colosseum.utils.acme.specs import MDPSpec
 15
 16
 17@gin.configurable
 18class ActorCriticContinuous(NonTabularBsuiteAgentWrapper):
 19    """
 20    The wrapper for the `ActorCritic` agent from `bsuite`.
 21    """
 22
 23    @staticmethod
 24    def produce_gin_file_from_parameters(parameters: Dict[str, Any], index: int = 0):
 25        string = ""
 26        for k, v in parameters.items():
 27            string += f"prms_{index}/ActorCriticContinuous.{k} = {v}\n"
 28        return string[:-1]
 29
 30    @staticmethod
 31    def is_episodic() -> bool:
 32        return False
 33
 34    @staticmethod
 35    def get_hyperparameters_search_spaces() -> Dict[str, tune.sample.Domain]:
 36        return {
 37            "network_width": tune.choice([64, 128, 256]),
 38            "network_depth": tune.choice([2, 4]),
 39            "max_sequence_length": tune.choice([16, 32, 64, 128]),
 40            "td_lambda": tune.choice([0.7, 0.8, 0.9]),
 41        }
 42
 43    @staticmethod
 44    def get_agent_instance_from_parameters(
 45        seed: int,
 46        optimization_horizon: int,
 47        mdp_specs: "MDPSpec",
 48        parameters: Dict[str, Any],
 49    ) -> "BaseAgent":
 50
 51        return ActorCriticContinuous(
 52            seed,
 53            mdp_specs,
 54            optimization_horizon,
 55            parameters["network_width"],
 56            parameters["network_depth"],
 57            parameters["max_sequence_length"],
 58            parameters["td_lambda"],
 59        )
 60
 61    @property
 62    def current_optimal_stochastic_policy(self) -> np.ndarray:
 63        logits = tf.stop_gradient(
 64            self._agent._network(
 65                tf.convert_to_tensor(self.emission_map.all_observations)
 66            )[0].logits
 67        ).numpy()
 68
 69        return (logits >= logits.max(-1, keepdims=True)).astype(np.float32)
 70
 71    def __init__(
 72        self,
 73        seed: int,
 74        mdp_specs: "MDPSpec",
 75        optimization_horizon: int,
 76        # MDP model parameters
 77        network_width: int,
 78        network_depth: int,
 79        max_sequence_length: int,
 80        td_lambda: float,
 81    ):
 82        r"""
 83        Parameters
 84        ----------
 85        seed : int
 86            The random seed.
 87        mdp_specs : MDPSpec
 88            The full specification of the MDP.
 89        optimization_horizon : int
 90            The total number of interactions that the agent is expected to have with the MDP.
 91        network_width : int
 92            The width of the neural networks of the agent.
 93        network_depth : int
 94            The depth of the neural networks of the agent.
 95        max_sequence_length : int
 96            The maximum sequence length for training the agent.
 97        td_lambda : float
 98            The TD(:math:`\lambda`) parameter for training the agent.
 99        """
100
101        tf.random.set_seed(seed)
102        np.random.seed(seed)
103
104        network = PolicyValueNet(
105            hidden_sizes=[network_width] * network_depth,
106            action_spec=mdp_specs.actions,
107        )
108        agent = ActorCritic(
109            obs_spec=mdp_specs.observations,
110            action_spec=mdp_specs.actions,
111            network=network,
112            optimizer=snt.optimizers.Adam(learning_rate=3e-3),
113            max_sequence_length=max_sequence_length,
114            td_lambda=td_lambda,
115            discount=0.99,
116            seed=seed,
117        )
118        super(ActorCriticContinuous, self).__init__(seed, agent, mdp_specs)