wehub-resource-sync
30 lines
988 B
Python
30 lines
988 B
Python
import numpy as np
|
|
|
|
from ray.rllib.utils.framework import try_import_torch
|
|
from ray.rllib.utils.metrics.stats.series import SeriesStats
|
|
from ray.util.annotations import DeveloperAPI
|
|
|
|
torch, _ = try_import_torch()
|
|
|
|
|
|
@DeveloperAPI
|
|
class MaxStats(SeriesStats):
|
|
"""A Stats object that tracks the max of a series of singular values (not vectors)."""
|
|
|
|
stats_cls_identifier = "max"
|
|
|
|
def _np_reduce_fn(self, values):
|
|
return np.nanmax(values)
|
|
|
|
def _torch_reduce_fn(self, values):
|
|
"""Reduce function for torch tensors (stays on GPU)."""
|
|
# torch.nanmax not available, use workaround
|
|
clean_values = values[~torch.isnan(values)]
|
|
if len(clean_values) == 0:
|
|
return torch.tensor(float("nan"), device=values.device)
|
|
# Cast to float32 to avoid errors from Long tensors
|
|
return torch.max(clean_values.float())
|
|
|
|
def __repr__(self) -> str:
|
|
return f"MaxStats({self.peek()}; window={self._window}; len={len(self)})"
|