def get_metrics(
self,
action: Union[Actions, Any],
reward: Union[Rewards, Any],
done: Union[bool, Sequence[bool]],
) -> Optional[EpisodeMetrics]:
# TODO: Add some metric about the entropy of the policy's distribution?
rewards = reward.y if isinstance(reward, Rewards) else reward
actions = action.y_pred if isinstance(action, Actions) else action
dones: Sequence[bool]
if not self.is_vectorized:
rewards = [rewards]
actions = [actions]
assert isinstance(done, bool)
dones = [done]
else:
assert isinstance(done, (np.ndarray, Tensor))
dones = done
metrics: List[EpisodeMetrics] = []
for env_index, (env_is_done, reward) in enumerate(zip(dones, rewards)):
if env_is_done:
metrics.append(
EpisodeMetrics(
n_samples=1,
# The average reward per episode.
mean_episode_reward=self.
_current_episode_reward[env_index],
# The average length of each episode.
mean_episode_length=self.
_current_episode_steps[env_index],
))
self._current_episode_reward[env_index] = 0
self._current_episode_steps[env_index] = 0
if not metrics:
return None
metric = sum(metrics, Metrics())
if wandb.run:
log_dict = metric.to_log_dict()
if self.wandb_prefix:
log_dict = add_prefix(log_dict,
prefix=self.wandb_prefix,
sep="/")
log_dict["steps"] = self._steps
log_dict["episode"] = self._episodes
wandb.log(log_dict)
return metric
def get_loss(self, x: Tensor, h_x: Tensor, y_pred: Tensor=None, y: Tensor=None) -> Loss:
loss_info = Loss(self.name)
batch_size = x.shape[0]
assert self.alphas is not None, "set the `self.alphas` attribute in the base class."
assert self.function_args is not None, "set the `self.function_args` attribute in the base class."
# Get the loss for each transformation argument.
for fn_arg, alpha in zip(self.function_args, self.alphas):
loss_i = self.get_loss_for_arg(x=x, h_x=h_x, fn_arg=fn_arg, alpha=alpha)
loss_info += loss_i
# print(f"{self.name}_{fn_arg}", loss_i.metrics)
# Fuse all the sub-metrics into a total metric.
# For instance, all the "rotate_0", "rotate_90", "rotate_180", etc.
metrics = loss_info.metrics
total_metrics = sum(loss_info.metrics.values(), Metrics())
# we actually add up all the metrics to get the "overall" metric.
metrics.clear()
metrics[self.name] = total_metrics
return loss_info
def average_metrics(self) -> MetricType:
return sum(self.average_metrics_per_task, Metrics())
def average_final_performance(self) -> MetricType:
return sum(self.final_performance_metrics, Metrics())
def average_online_performance(self) -> MetricType:
return sum(self.online_performance_metrics, Metrics())
def online_performance_metrics(self) -> List[MetricType]:
return [
sum(online_performance_dict.values(), Metrics())
for online_performance_dict in self.online_performance
]
def average_metrics(self) -> MetricType:
""" Returns the average 'Metrics' object for this task. """
return sum(self.metrics, Metrics())
def __post_init__(self):
if self.metrics and isinstance(self.metrics[0], dict):
self.metrics = [
Metrics.from_dict(metrics, drop_extra_fields=False)
for metrics in self.metrics
]
def online_performance_metrics(self) -> MetricsType:
return sum(self.online_performance.values(), Metrics())