colosseum.agent.actors.Q_values_actor

 1from typing import TYPE_CHECKING, Callable, Union
 2
 3import dm_env
 4import numpy as np
 5
 6from colosseum.agent.actors import BaseActor
 7from colosseum.utils.acme.specs import MDPSpec
 8
 9if TYPE_CHECKING:
10    from colosseum.mdp import ACTION_TYPE
11
12
13class QValuesActor(BaseActor):
14    """
15    The `QValuesActor` is an actor component that selects actions that maximises a given q-values estimate. It also
16    supports the epsilon-greedy and the Boltzmann action selection strategies to boost exploration.
17    """
18
19    def __init__(
20        self,
21        seed: int,
22        mdp_specs: MDPSpec,
23        epsilon_greedy: Union[float, Callable[[int], float]] = None,
24        boltzmann_temperature: Union[float, Callable[[int], float]] = None,
25    ):
26        """
27        Parameters
28        ----------
29        seed : int
30            The random seed.
31        mdp_specs : MDPSpec
32            The full specification of the MDP.
33        epsilon_greedy : Union[float, Callable], optional
34            The probability of selecting an action at random. It can be provided as a float or as a function of the
35            total number of interactions. By default, the probability is set to zero.
36        boltzmann_temperature : Union[float, Callable], optional
37            The parameter that controls the Boltzmann exploration. It can be provided as a float or as a function of
38            the total number of interactions. By default, Boltzmann exploration is disabled.
39        """
40        super(QValuesActor, self).__init__(seed, mdp_specs)
41
42        if epsilon_greedy is not None:
43            if type(epsilon_greedy) == float:
44                epsilon_greedy = lambda t: epsilon_greedy
45        if boltzmann_temperature is not None:
46            if type(boltzmann_temperature) == float:
47                boltzmann_temperature = lambda t: boltzmann_temperature
48
49        self._epsilon_greedy = epsilon_greedy
50        self._boltzmann_temperature = boltzmann_temperature
51        self._total_interactions = 0
52        self._q_values = None
53        self._n_states = self._mdp_spec.observations.num_values
54        self._n_actions = self._mdp_spec.actions.num_values
55
56    def set_q_values(self, Q: np.ndarray):
57        """
58        updates the q-values estimates of the component with the one given in input.
59        Parameters
60        ----------
61        Q : np.ndarray
62            The q-values estimates.
63        """
64        self._q_values = Q
65        self._episodic = Q.ndim == 3
66
67    def select_action(self, ts: dm_env.TimeStep, time: int) -> "ACTION_TYPE":
68        assert self._q_values is not None, "The q values have not been initialized."
69
70        self._total_interactions += 1
71
72        # Epsilon greedy policy
73        if self._epsilon_greedy is not None:
74            if self._rng_fast.random() < self._epsilon_greedy(self._total_interactions):
75                return self._rng_fast.randint(0, self._n_actions - 1)
76
77        # Retrieve the q-values
78        q = self._q_values[(time, ts.observation) if self._episodic else ts.observation]
79
80        # Boltzmann exploration
81        if self._boltzmann_temperature is not None:
82            q = np.exp(self._boltzmann_temperature(self._total_interactions) * q)
83            return self._rng.choice(
84                range(self._n_actions), replace=False, p=q / q.sum()
85            )
86
87        # Greedy selection
88        return self._rng.choice(np.where(q == q.max())[0])