colosseum.experiment.agent_mdp_interaction

  1import io
  2from time import time
  3from typing import TYPE_CHECKING, Dict, Iterable, List, Tuple, Union, Set
  4
  5import numpy as np
  6import pandas as pd
  7import seaborn as sns
  8import toolz
  9import tqdm
 10from matplotlib import pyplot as plt
 11from tqdm import trange
 12from wrapt_timeout_decorator import timeout
 13
 14from colosseum import config
 15from colosseum.config import process_debug_output
 16from colosseum.emission_maps import Tabular
 17from colosseum.experiment.indicators import (
 18    get_episodic_regrets_and_average_reward_at_time_zero,
 19)
 20from colosseum.mdp.utils.markov_chain import get_average_reward
 21from colosseum.utils.acme import InMemoryLogger
 22from colosseum.utils.acme.base_logger import Logger
 23from colosseum.utils.formatter import clear_agent_mdp_class_name
 24
 25if TYPE_CHECKING:
 26    from colosseum.mdp import ContinuousMDP, EpisodicMDP, BaseMDP
 27    from colosseum.agent.agents.base import BaseAgent
 28
 29sns.set_theme()
 30
 31
 32class MDPLoop:
 33    """
 34    The `MDPLoop` is the object in charge of the agent/MDP interactions and the computation of the performance indicators.
 35    It also provides limited plotting functionalities.
 36    """
 37
 38    @staticmethod
 39    def get_indicators() -> List[str]:
 40        """
 41        Returns
 42        -------
 43        List[str]
 44            The code names for the indicators that are computed by the MDPLoop.
 45        """
 46        return [
 47            "cumulative_expected_reward",
 48            "cumulative_regret",
 49            "cumulative_reward",
 50            "normalized_cumulative_expected_reward",
 51            "normalized_cumulative_regret",
 52            "normalized_cumulative_reward",
 53            "steps_per_second",
 54        ]
 55
 56    @staticmethod
 57    def get_baseline_indicators() -> List[str]:
 58        """
 59        Returns
 60        -------
 61        List[str]
 62            The code names for the baseline indicators that are computed by the MDPLoop.
 63        """
 64        return [
 65            "random_cumulative_regret",
 66            "random_cumulative_expected_reward",
 67            "random_normalized_cumulative_regret",
 68            "random_normalized_cumulative_expected_reward",
 69            "optimal_cumulative_expected_reward",
 70            "optimal_normalized_cumulative_expected_reward",
 71            "worst_cumulative_regret",
 72            "worst_cumulative_expected_reward",
 73            "worst_normalized_cumulative_regret",
 74            "worst_normalized_cumulative_expected_reward",
 75        ]
 76
 77    @staticmethod
 78    def get_baselines() -> Set[str]:
 79        """
 80        Returns
 81        -------
 82        Set[str]
 83            The baselines available for comparison.
 84        """
 85        return set(b[: b.find("_")] for b in MDPLoop.get_baseline_indicators())
 86
 87    @staticmethod
 88    def get_baselines_color_dict() -> Dict[str, str]:
 89        """
 90        Returns
 91        -------
 92        Dict[str, str]
 93            The color associated by default to the baselines.
 94        """
 95        return dict(random="black", worst="crimson", optimal="gold")
 96
 97    @staticmethod
 98    def get_baselines_style_dict():
 99        """
100        Returns
101        -------
102        Dict[str, str]
103            The line style associated by default to the baselines.
104        """
105        return dict(random=(0, (6, 12)), worst=(9, (6, 12)), optimal=(0, (6, 12)))
106
107    def __init__(
108        self,
109        mdp: Union["BaseMDP", "EpisodicMDP", "ContinuousMDP"],
110        agent: "BaseAgent",
111        logger: Logger = None,
112        n_log_intervals_to_check_for_agent_optimality: int = 10,
113        enforce_time_constraint: bool = True,
114    ) -> object:
115        """
116        Parameters
117        ----------
118        mdp: Union["EpisodicMDP", "ContinuousMDP"]
119            The MDP.
120        agent : BaseAgent
121            The agent.
122        logger : Logger
123            The logger where the results of the interaction between the agent and the MDP are stored. By default, the
124            `InMemoryLogger` is used.
125        n_log_intervals_to_check_for_agent_optimality : int
126            The length of the interval between check is the policy has reached optimality. By default, the check happens
127            every ten interactions.
128        enforce_time_constraint : bool
129            If True, the computational time constraint given in the `run` function is enforced through multithreading.
130            By default, it is enforced.
131        """
132
133        if logger is None:
134            logger = InMemoryLogger()
135
136        self.logger = logger
137        self._enforce_time_constraint = enforce_time_constraint
138        self._mdp = mdp
139        self._agent = agent
140        self._episodic = self._mdp.is_episodic()
141        self._n_steps_to_check_for_agent_optimality = (
142            n_log_intervals_to_check_for_agent_optimality
143        )
144        assert self._episodic == agent.is_episodic()
145        assert self._agent.is_emission_map_accepted(
146            Tabular if self._mdp.emission_map is None else self._mdp.emission_map
147        )
148        self.actions_sequence = []
149
150    @property
151    def remaining_time(self) -> float:
152        """
153        Returns
154        -------
155        float
156            The remaining computational time for training the agent.
157        """
158        return self._max_time - (time() - self._mdp_loop_timer)
159
160    def _limit_update_time(self, t, f):
161        try:
162            if self.remaining_time < 0.5:
163                raise TimeoutError()
164            timeout(self.remaining_time)(f)()
165        except TimeoutError or SystemError:
166            if config._DEBUG_LEVEL > 0:
167                print("Time exceeded with function ", f)
168            self._limit_exceeded(t)
169
170    def _limit_exceeded(self, t):
171        self._is_training = False
172        self._last_training_step = t
173        if config._DEBUG_LEVEL > 0:
174            do = f"Stopped training at {time() - self._mdp_loop_timer:.2f}"
175            process_debug_output(do)
176        if self._verbose:
177            self._verbose_postfix["is_training"] = f"No, time exhausted at {t}"
178
179    def run(
180        self,
181        T: int,
182        log_every: int = -1,
183        max_time: float = np.inf,
184    ) -> Tuple[int, Dict[str, float]]:
185        """
186        runs the agent/MDP interactions.
187
188        Parameters
189        ----------
190        T : int
191            The number of total interactions between the agent and the MDP.
192        log_every : int
193            The number of time steps after which performance indicators are calculated. By default, it does not calculate
194            them at any time except at the last one.
195        max_time : float
196            The maximum number of seconds the interactions can take. If it is surpassed then the loop is interrupted.
197            By default, the maximum given time is infinite.
198
199        Returns
200        ----------
201        int
202            The time step at which the training has been interrupted due to the time constraint. If the constraint has
203            been respected it returns -1.
204        Dict[str, float]
205            The performance indicators computed at the end of the interactions.
206        """
207
208        if max_time == np.inf:
209            enforce_time_constraint = False
210        else:
211            enforce_time_constraint = self._enforce_time_constraint
212
213        assert (
214            type(log_every) == int
215        ), f"The log_every variable should be an integer, received value: {log_every}."
216        log_every = -1 if log_every == 0 else log_every
217
218        # Reset the visitation count of the MDP
219        self._mdp.reset_visitation_counts()
220
221        self._reset_run_variables()
222        self._max_time = max_time
223
224        ts = self._mdp.reset()
225        first_before_new_episode_timer = time()
226        if enforce_time_constraint and self.remaining_time < np.inf:
227            self._limit_update_time(0, self._agent.before_start_interacting)
228        else:
229            self._agent.before_start_interacting()
230        if config._DEBUG_LEVEL > 0:
231            if self._is_training:
232                do = f"before_start_interacting completed in {time() - first_before_new_episode_timer:.2f}."
233            else:
234                do = "before_start_interacting exceeded the time limit."
235            process_debug_output(do)
236
237        self._set_loop(T)
238        for t in self._loop:
239            if self._is_training and self.remaining_time < 0.5:
240                self._limit_exceeded(t)
241
242            # MDP step
243            h = self._mdp.h
244            action = self._agent.select_action(ts, h)
245            new_ts = self._mdp.step(action)
246            self.actions_sequence.append(new_ts.reward)
247
248            # Single step agent update
249            if self._is_training:
250                if enforce_time_constraint and self.remaining_time < np.inf:
251                    self._limit_update_time(
252                        t,
253                        lambda: self._agent.step_update(ts, action, new_ts, h),
254                    )
255                else:
256                    self._agent.step_update(ts, action, new_ts, h)
257
258            # End of episode agent update
259            if self._is_training and self._agent.is_episode_end(ts, action, new_ts, h):
260                if enforce_time_constraint and self.remaining_time < np.inf:
261                    self._limit_update_time(t, self._agent.episode_end_update)
262                else:
263                    self._agent.episode_end_update()
264
265            if t > 0 and log_every > 0 and t % log_every == 0:
266                # Log the performance of the agent
267                self._update_performance_logs(t)
268                self._n_steps_since_last_log = 0
269
270                # User defined custom log
271                self._agent.agent_logs()
272
273                # Verbose loggings
274                self._update_user_loggings(t)
275
276                # Storing the latest regrets
277                self._latest_expected_regrets.append(self._normalized_regret)
278                if (
279                    len(self._latest_expected_regrets)
280                    > self._n_steps_to_check_for_agent_optimality
281                ):
282                    self._latest_expected_regrets.pop(0)
283
284                # Stop training if the agent has confidently reached the optimal policy
285                if self._is_training and t > 0.2 * T and self._is_policy_optimal():
286                    if type(self._loop) == tqdm.std.tqdm:
287                        self._verbose_postfix["is_training"] = f"No, optimal at {t}"
288                    self._is_training = False
289
290            self._n_steps_since_last_log += 1
291            self._cumulative_reward += new_ts.reward
292            ts = new_ts
293
294            # Resetting episodic MDPs
295            if self._mdp.is_episodic() and new_ts.last():
296                assert self._mdp.necessary_reset or t == T - 2
297                ts = self._mdp.reset()
298                self._n_episodes += 1
299
300        self._update_performance_logs(t)
301        self.logger.close()
302        return self._last_training_step, self._last_logs
303
304    def _reset_run_variables(self):
305        self._cumulative_reward = 0.0
306        self._cumulative_regret = 0.0
307        self._normalized_cumulative_regret = 0.0
308        self._random_cumulative_expected_reward = 0.0
309        self._random_cumulative_regret = 0.0
310        self._normalized_random_cumulative_regret = 0.0
311        self._cumulative_expected_reward_agent = 0.0
312
313        self._verbose = False
314        self._verbose_postfix = dict(is_training="True")
315        self._is_training = True
316        self._n_steps_since_last_log = 0
317        self._last_training_step = -1
318        self._n_episodes = 0
319        self._last_logs = None
320        self._past_logs = None
321        self._cached_episodic_regrets = None
322        self._cached_continuous_regrets = None
323        self._latest_expected_regrets = []
324
325        # Cache the regret for the random agent
326        if self._episodic:
327
328            # Random agent regret
329            self._episodic_regret_random_agent = (
330                self._mdp.episodic_optimal_average_reward
331                - self._mdp.episodic_random_average_reward
332            )
333            self._episodic_normalized_regret_random_agent = (
334                self._episodic_regret_random_agent
335                / (
336                    self._mdp.episodic_optimal_average_reward
337                    - self._mdp.episodic_worst_average_reward
338                )
339            )
340
341            # Worst agent regret
342            self._episodic_regret_worst_agent = (
343                self._mdp.episodic_optimal_average_reward
344                - self._mdp.episodic_worst_average_reward
345            )
346            self._episodic_normalized_regret_worst_agent = (
347                self._episodic_regret_worst_agent
348                / (
349                    self._mdp.episodic_optimal_average_reward
350                    - self._mdp.episodic_worst_average_reward
351                )
352            )
353
354            # Reward normalized
355            self._cumulative_reward_normalizer = lambda t, cr: (
356                cr - t * self._mdp.episodic_worst_average_reward
357            ) / (
358                self._mdp.episodic_optimal_average_reward
359                - self._mdp.episodic_worst_average_reward
360            )
361        else:
362
363            # Random agent regret
364            self._regret_random_agent = (
365                self._mdp.optimal_average_reward - self._mdp.random_average_reward
366            )
367            self._normalized_regret_random_agent = self._regret_random_agent / (
368                self._mdp.optimal_average_reward - self._mdp.worst_average_reward
369            )
370
371            # Worst agent regret
372            self._regret_worst_agent = (
373                self._mdp.optimal_average_reward - self._mdp.worst_average_reward
374            )
375            self._normalized_regret_worst_agent = self._regret_worst_agent / (
376                self._mdp.optimal_average_reward - self._mdp.worst_average_reward
377            )
378
379            assert (
380                self._mdp.optimal_average_reward - self._mdp.worst_average_reward
381                > 0.0002
382            ), type(self._mdp).__name__ + str(self._mdp.parameters)
383
384            self._cumulative_reward_normalizer = lambda t, cr: (
385                cr - t * self._mdp.worst_average_reward
386            ) / (self._mdp.optimal_average_reward - self._mdp.worst_average_reward)
387
388        self.logger.reset()
389        self._mdp_loop_timer = time()
390        self._verbose_time = time()
391
392    def _update_performance_logs(self, t: int):
393        self._compute_performance_indicators(t + 1)
394
395        self._last_logs = dict(
396            steps=t,
397            cumulative_regret=self._cumulative_regret,
398            cumulative_reward=self._cumulative_reward,
399            cumulative_expected_reward=self._cumulative_expected_reward_agent,
400            normalized_cumulative_regret=self._normalized_cumulative_regret,
401            normalized_cumulative_reward=self._cumulative_reward_normalizer(
402                t, self._cumulative_reward
403            ),
404            normalized_cumulative_expected_reward=self._cumulative_reward_normalizer(
405                t, self._cumulative_expected_reward_agent
406            ),
407            random_cumulative_regret=self._cumulative_regret_random_agent,
408            random_cumulative_expected_reward=self._cumulative_reward_random_agent,
409            random_normalized_cumulative_regret=self._normalized_cumulative_regret_random_agent,
410            random_normalized_cumulative_expected_reward=self._cumulative_reward_normalizer(
411                t, self._cumulative_reward_random_agent
412            ),
413            worst_cumulative_regret=self._cumulative_regret_worst_agent,
414            worst_cumulative_expected_reward=self._cumulative_reward_worst_agent,
415            worst_normalized_cumulative_regret=self._normalized_cumulative_regret_worst_agent,
416            worst_normalized_cumulative_expected_reward=self._cumulative_reward_normalizer(
417                t, self._cumulative_reward_worst_agent
418            ),
419            optimal_cumulative_expected_reward=self._cumulative_reward_optimal_agent,
420            optimal_normalized_cumulative_expected_reward=self._cumulative_reward_normalizer(
421                t, self._cumulative_reward_optimal_agent
422            ),
423            steps_per_second=t / (time() - self._mdp_loop_timer),
424        )
425
426        # Communicate the indicators to the logger with a maximum of five digits
427        a = toolz.valmap(lambda x: np.round(x, 5), self._last_logs)
428        self.logger.write(a)
429
430    def _compute_regrets(self):
431        if self._episodic:
432            return self._compute_episodic_regret()
433        return self._compute_continuous_regret()
434
435    def _compute_performance_indicators(self, t: int):
436        self._compute_regrets()
437
438        if self._episodic:
439            # Randon agent (regret)
440            self._cumulative_regret_random_agent = (
441                self._episodic_regret_random_agent * t
442            )
443            self._normalized_cumulative_regret_random_agent = (
444                self._episodic_normalized_regret_random_agent * t
445            )
446
447            # Worst agent (regret)
448            self._cumulative_regret_worst_agent = self._episodic_regret_worst_agent * t
449            self._normalized_cumulative_regret_worst_agent = (
450                self._episodic_normalized_regret_worst_agent * t
451            )
452
453            # Random agent (reward)
454            self._cumulative_reward_random_agent = (
455                self._mdp.episodic_random_average_reward * t
456            )
457
458            # Worst agent (reward)
459            self._cumulative_reward_worst_agent = (
460                self._mdp.episodic_worst_average_reward * t
461            )
462
463            # Optimal agent (reward)
464            self._cumulative_reward_optimal_agent = (
465                self._mdp.episodic_optimal_average_reward * t
466            )
467
468        else:
469            # Randon agent (regret)
470            self._cumulative_regret_random_agent = self._regret_random_agent * t
471            self._normalized_cumulative_regret_random_agent = (
472                self._normalized_regret_random_agent * t
473            )
474
475            # Worst agent (regret)
476            self._cumulative_regret_worst_agent = self._regret_worst_agent * t
477            self._normalized_cumulative_regret_worst_agent = (
478                self._normalized_regret_worst_agent * t
479            )
480
481            # Random agent (reward)
482            self._cumulative_reward_random_agent = self._mdp.random_average_reward * t
483
484            # Worst agent (reward)
485            self._cumulative_reward_worst_agent = self._mdp.worst_average_reward * t
486
487            # Optimal agent (reward)
488            self._cumulative_reward_optimal_agent = self._mdp.optimal_average_reward * t
489
490        # Avoid numerical errors that lead to negative rewards
491        assert (
492            self._regret >= 0.0
493        ), f"{self._regret} on {type(self._mdp).__name__} {self._mdp.parameters} for policy {self._agent.current_optimal_stochastic_policy}"
494        assert self._normalized_regret >= 0.0, self._normalized_regret
495
496        self._cumulative_regret += self._regret * self._n_steps_since_last_log
497        self._normalized_cumulative_regret += (
498            self._normalized_regret * self._n_steps_since_last_log
499        )
500        self._cumulative_expected_reward_agent += (
501            self._agent_average_reward * self._n_steps_since_last_log
502        )
503
504    @property
505    def _agent_average_reward(self):
506        if self._episodic:
507            return self._episodic_agent_average_reward / self._mdp.H
508        return self._agent_continuous_average_reward
509
510    def _compute_continuous_regret(self):
511        if not self._is_training:
512            if self._cached_continuous_regrets is None:
513                self._cached_continuous_regrets = self._get_continuous_regrets()
514            self._regret, self._normalized_regret = self._cached_continuous_regrets
515        else:
516            self._regret, self._normalized_regret = self._get_continuous_regrets()
517
518    def _get_continuous_regrets(self):
519        self._agent_continuous_average_reward = get_average_reward(
520            self._mdp.T,
521            self._mdp.R,
522            self._agent.current_optimal_stochastic_policy,
523            [(self._mdp.node_to_index[self._mdp.cur_node], 1.0)],
524        )
525
526        r = self._mdp.optimal_average_reward - self._agent_continuous_average_reward
527        if np.isclose(r, 0.0, atol=1e-3):
528            r = 0.0
529        if r < 0:
530            r = 0
531        nr = r / (self._mdp.optimal_average_reward - self._mdp.worst_average_reward)
532        return r, nr
533
534    def _compute_episodic_regret(self):
535        if not self._is_training:
536            # If the agent is not training, the policy will not change we can cache and reuse the regret for each given
537            # starting state.
538            if self._cached_episodic_regrets is None:
539                Rs, epi_agent_ar = get_episodic_regrets_and_average_reward_at_time_zero(
540                    self._mdp.H,
541                    self._mdp.T,
542                    self._mdp.R,
543                    self._agent.current_optimal_stochastic_policy,
544                    self._mdp.starting_state_distribution,
545                    self._mdp.optimal_value_functions[1],
546                )
547                self._episodic_agent_average_reward = epi_agent_ar
548                self._cached_episodic_regrets = {
549                    n: (
550                        Rs[self._mdp.node_to_index[n]] / self._mdp.H,  # expected regret
551                        Rs[self._mdp.node_to_index[n]]  # normalized expected regret
552                        / self._mdp.get_minimal_regret_for_starting_node(n),
553                    )
554                    for n in self._mdp.starting_nodes
555                }
556            self._regret, self._normalized_regret = self._cached_episodic_regrets[
557                self._mdp.last_starting_node
558            ]
559        else:
560            Rs, epi_agent_ar = get_episodic_regrets_and_average_reward_at_time_zero(
561                self._mdp.H,
562                self._mdp.T,
563                self._mdp.R,
564                self._agent.current_optimal_stochastic_policy,
565                self._mdp.starting_state_distribution,
566                self._mdp.optimal_value_functions[1],
567            )
568            self._episodic_agent_average_reward = epi_agent_ar
569            self._regret = (
570                Rs[self._mdp.node_to_index[self._mdp.last_starting_node]] / self._mdp.H
571            )
572            self._normalized_regret = (
573                self._regret
574                / self._mdp.get_minimal_regret_for_starting_node(
575                    self._mdp.last_starting_node
576                )
577                * self._mdp.H
578            )
579
580    def _is_policy_optimal(self) -> bool:
581        if (
582            len(self._latest_expected_regrets)
583            == self._n_steps_to_check_for_agent_optimality
584            and np.isclose(
585                0,
586                self._latest_expected_regrets,
587                atol=1e-4 if self._mdp.is_episodic() else 1e-5,
588            ).all()
589        ):
590            # After we get an empirical suggestions that the policy may be optimal, we check if the expected regret is
591            # zero as well
592            self._compute_regrets()
593            return np.isclose(self._normalized_regret, 0).all()
594        return False
595
596    def _set_loop(self, T: int) -> Iterable:
597        """
598        creates a loop lasting for T steps taking into account the verbosity level.
599        """
600        if config.VERBOSE_LEVEL != 0:
601            desc = f"Experiment loop {type(self._agent).__name__}@{type(self._mdp).__name__}"
602            if type(config.VERBOSE_LEVEL) == str:
603                self.s = io.StringIO()  # we need this reference
604                self._loop = trange(T, desc=desc, file=self.s, mininterval=5)
605            else:
606                self._loop = trange(T, desc=desc, mininterval=5)
607            self._verbose = True
608        else:
609            self._loop = range(T)
610
611    def _update_user_loggings(self, t: int):
612        if self._verbose:  # and time() - self._verbose_time > 5:
613            self._verbose_postfix["Instantaneous normalized regret"] = np.round(
614                self._normalized_regret / t, 8
615            )
616            self._loop.set_postfix(self._verbose_postfix, refresh=False)
617
618    def plot(
619        self,
620        indicator: str = "cumulative_regret",
621        ax=None,
622        baselines=("random", "worst", "optimal"),
623        label=None,
624    ):
625        """
626        plots the values of the indicator obtained by the agent during the interactions along with the baseline values.
627
628        Parameters
629        ----------
630        indicator : str
631            The code name of the performance indicator that will be shown in the plot. Check `MDPLoop.get_indicators()`
632            to get a list of the available indicators. By default, the 'cumulative_regret' is shown.
633        ax : plt.Axes
634            The ax object where the plot will be put. By default, a new axis is created.
635        baselines : List[str]
636            The baselines to be included in the plot. Check `MDPLoop.get_baselines()` to get a list of the available
637            baselines. By default, all baselines are shown.
638        label : str
639            The label to be given to the agent. By default, a cleaned version of the agent class name is used.
640        """
641
642        show = ax is None
643        if ax is None:
644            fig, ax = plt.subplots()
645
646        assert indicator in self.get_indicators(), (
647            f"{indicator} is not an indicator. The indicators available are: "
648            + ",".join(self.get_indicators())
649            + "."
650        )
651
652        df_e = pd.DataFrame(self.logger.data)
653        time_steps = [0] + df_e.loc[:, "steps"].tolist()
654        ax.plot(
655            time_steps[1:] if indicator == "steps_per_second" else time_steps,
656            ([] if indicator == "steps_per_second" else [0])
657            + df_e.loc[:, indicator].tolist(),
658            label=clear_agent_mdp_class_name(type(self._agent).__name__)
659            if label is None
660            else label,
661        )
662        ax.set_ylabel(indicator.replace("_", " ").capitalize())
663
664        for b in baselines:
665            indicator = indicator.replace(
666                "cumulative_reward", "cumulative_expected_reward"
667            )
668            if b + "_" + indicator in self.get_baseline_indicators():
669                ax.plot(
670                    time_steps,
671                    [0] + df_e.loc[:, b + "_" + indicator].tolist(),
672                    label=b.capitalize(),
673                    # alpha=0.9,
674                    linestyle=(0, (5, 10)),
675                    color="darkolivegreen"
676                    if "optimal" in b
677                    else ("darkred" if "worst" in b else "darkslategray"),
678                    linewidth=2,
679                )
680
681        ax.set_xlabel("time step")
682        ax.legend()
683        if show:
684            plt.tight_layout()
685            plt.show()