shared/stats.py

import numpy as np

from dataclasses import dataclass
from torch.utils.tensorboard.writer import SummaryWriter
from typing import Dict, List, Optional, Sequence, TypeVar


@dataclass
class Episode:
 score: float = 0
 length: int = 0


StatisticSelf = TypeVar("StatisticSelf", bound="Statistic")


@dataclass
class Statistic:
 values: np.ndarray
 round_digits: int = 2

 @property
 def mean(self) -> float:
 return np.mean(self.values).item()

 @property
 def std(self) -> float:
 return np.std(self.values).item()

 @property
 def min(self) -> float:
 return np.min(self.values).item()

 @property
 def max(self) -> float:
 return np.max(self.values).item()

 def sum(self) -> float:
 return np.sum(self.values).item()

 def __len__(self) -> int:
 return len(self.values)

 def _diff(self: StatisticSelf, o: StatisticSelf) -> float:
 return (self.mean - self.std) - (o.mean - o.std)

 def __gt__(self: StatisticSelf, o: StatisticSelf) -> bool:
 return self._diff(o) > 0

 def __ge__(self: StatisticSelf, o: StatisticSelf) -> bool:
 return self._diff(o) >= 0

 def __repr__(self) -> str:
 mean = round(self.mean, self.round_digits)
 std = round(self.std, self.round_digits)
 if self.round_digits == 0:
 mean = int(mean)
 std = int(std)
 return f"{mean} +/- {std}"

 def to_dict(self) -> Dict[str, float]:
 return {
 "mean": self.mean,
 "std": self.std,
 "min": self.min,
 "max": self.max,
 }


EpisodesStatsSelf = TypeVar("EpisodesStatsSelf", bound="EpisodesStats")


class EpisodesStats:
 episodes: Sequence[Episode]
 simple: bool
 score: Statistic
 length: Statistic

 def __init__(self, episodes: Sequence[Episode], simple: bool = False) -> None:
 self.episodes = episodes
 self.simple = simple
 self.score = Statistic(np.array([e.score for e in episodes]))
 self.length = Statistic(np.array([e.length for e in episodes]), round_digits=0)

 def __gt__(self: EpisodesStatsSelf, o: EpisodesStatsSelf) -> bool:
 return self.score > o.score

 def __ge__(self: EpisodesStatsSelf, o: EpisodesStatsSelf) -> bool:
 return self.score >= o.score

 def __repr__(self) -> str:
 return (
 f"Score: {self.score} ({round(self.score.mean - self.score.std, 2)}) | "
 f"Length: {self.length}"
 )

 def __len__(self) -> int:
 return len(self.episodes)

 def _asdict(self) -> dict:
 return {
 "n_episodes": len(self.episodes),
 "score": self.score.to_dict(),
 "length": self.length.to_dict(),
 }

 def write_to_tensorboard(
 self, tb_writer: SummaryWriter, main_tag: str, global_step: Optional[int] = None
 ) -> None:
 stats = {"mean": self.score.mean}
 if not self.simple:
 stats.update(
 {
 "min": self.score.min,
 "max": self.score.max,
 "result": self.score.mean - self.score.std,
 "n_episodes": len(self.episodes),
 "length": self.length.mean,
 }
 )
 tb_writer.add_scalars(
 main_tag,
 stats,
 global_step=global_step,
 )


class EpisodeAccumulator:
 def __init__(self, num_envs: int):
 self._episodes = []
 self.current_episodes = [Episode() for _ in range(num_envs)]

 @property
 def episodes(self) -> List[Episode]:
 return self._episodes

 def step(self, reward: np.ndarray, done: np.ndarray) -> None:
 for idx, current in enumerate(self.current_episodes):
 current.score += reward[idx]
 current.length += 1
 if done[idx]:
 self._episodes.append(current)
 self.current_episodes[idx] = Episode()
 self.on_done(idx, current)

 def __len__(self) -> int:
 return len(self.episodes)

 def on_done(self, ep_idx: int, episode: Episode) -> None:
 pass

 def stats(self) -> EpisodesStats:
 return EpisodesStats(self.episodes)