wehub-resource-sync
88 lines
3.5 KiB
Python
88 lines
3.5 KiB
Python
"""Example of plugging a custom EMA Stats class into the MetricsLogger via
|
|
``AlgorithmConfig.reporting(custom_stats_cls_lookup=...)``.
|
|
|
|
RLlib aggregates metrics through Stats objects (see
|
|
:py:class:`~ray.rllib.utils.metrics.stats.base.StatsBase`). A ``reduce=...``
|
|
keyword (e.g. ``"ema"``, ``"mean"``) maps to a Stats class via
|
|
:py:data:`~ray.rllib.utils.metrics.metrics_logger.DEFAULT_STATS_CLS_LOOKUP`.
|
|
By passing your own dictionary to
|
|
``AlgorithmConfig.reporting(custom_stats_cls_lookup=...)`` you can either
|
|
add a new key (referenced via ``log_value(..., reduce="<your_key>")``) or
|
|
replace an existing key to change RLlib's default reduction behaviour for
|
|
all metrics that internally use that reduction.
|
|
|
|
This example replaces the default ``"ema"`` key with ``ClearingEmaStats``,
|
|
a subclass of :py:class:`~ray.rllib.utils.metrics.stats.ema.EmaStats` that
|
|
resets its running value to ``NaN`` after every ``reduce()`` call. The
|
|
default ``EmaStats`` carries the running value forward across reduces so a
|
|
metric smoothly tracks across training iterations; ``ClearingEmaStats``
|
|
instead restarts each iteration's EMA from scratch. This reproduces the
|
|
short-lived behaviour from Ray 2.53/2.54 and is useful when you want each
|
|
reported value to reflect only the data observed within a single
|
|
iteration. Because RLlib's internal timers (and other ``reduce="ema"``
|
|
metrics) go through the lookup, swapping the key is enough -- no per-metric
|
|
plumbing required.
|
|
|
|
``VerifyClearingEmaCallback`` is attached purely as a sanity
|
|
check: every ``on_train_result`` it asserts that RLlib's built-in
|
|
training-iteration timer is now an instance of ``ClearingEmaStats`` and
|
|
that its internal ``_value`` was reset to ``NaN`` by the reduce that just
|
|
happened.
|
|
|
|
How to run this script
|
|
----------------------
|
|
``python clearing_ema.py --stop-iters=3``
|
|
|
|
For debugging:
|
|
``python clearing_ema.py --no-tune --num-env-runners=0 --stop-iters=3``
|
|
"""
|
|
import math
|
|
|
|
import numpy as np
|
|
|
|
from ray.rllib.callbacks.callbacks import RLlibCallback
|
|
from ray.rllib.examples.utils import (
|
|
add_rllib_example_script_args,
|
|
run_rllib_example_script_experiment,
|
|
)
|
|
from ray.rllib.utils.metrics import TIMERS, TRAINING_ITERATION_TIMER
|
|
from ray.rllib.utils.metrics.metrics_logger import DEFAULT_STATS_CLS_LOOKUP
|
|
from ray.rllib.utils.metrics.stats.ema import EmaStats
|
|
from ray.tune.registry import get_trainable_cls
|
|
|
|
|
|
class ClearingEmaStats(EmaStats):
|
|
def reduce(self, compile: bool = True):
|
|
result = super().reduce(compile=compile)
|
|
self._value = np.nan
|
|
return result
|
|
|
|
|
|
class VerifyClearingEmaCallback(RLlibCallback):
|
|
def on_train_result(self, *, algorithm, metrics_logger, result, **kwargs) -> None:
|
|
timer_stats = metrics_logger.stats[TIMERS][TRAINING_ITERATION_TIMER]
|
|
# Only check that our value was cleared
|
|
assert math.isnan(timer_stats._value)
|
|
|
|
|
|
parser = add_rllib_example_script_args(default_reward=50.0, default_iters=5)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = parser.parse_args()
|
|
|
|
# Override the default "ema" key in the lookup. Every metric RLlib logs
|
|
# with `reduce="ema"` (e.g. all of its built-in timers) will now use
|
|
# `ClearingEmaStats`.
|
|
custom_stats_lookup = {**DEFAULT_STATS_CLS_LOOKUP, "ema": ClearingEmaStats}
|
|
|
|
base_config = (
|
|
get_trainable_cls(args.algo)
|
|
.get_default_config()
|
|
.environment("CartPole-v1")
|
|
.reporting(custom_stats_cls_lookup=custom_stats_lookup)
|
|
.callbacks(VerifyClearingEmaCallback)
|
|
)
|
|
|
|
run_rllib_example_script_experiment(base_config, args)
|