[RLlib] MetricsLogger cleanup throughput logic.

rllib/core/learner/learner.py

@@ -1736,7 +1736,6 @@ def _log_steps_trained_metrics(self, batch: MultiAgentBatch):
             (ALL_MODULES, NUM_ENV_STEPS_TRAINED_LIFETIME),
             batch.env_steps(),
             reduce="sum",
-            with_throughput=True,
         )
 
     @Deprecated(

rllib/env/multi_agent_env_runner.py

@@ -792,7 +792,6 @@ def set_state(self, state: StateDict) -> None:
                 key=NUM_ENV_STEPS_SAMPLED_LIFETIME,
                 value=state[NUM_ENV_STEPS_SAMPLED_LIFETIME],
                 reduce="sum",
-                with_throughput=True,
             )
 
     @override(Checkpointable)
@@ -993,7 +992,6 @@ def _increase_sampled_metrics(self, num_steps, next_obs, episode):
             NUM_ENV_STEPS_SAMPLED_LIFETIME,
             num_steps,
             reduce="sum",
-            with_throughput=True,
         )
         # Completed episodes.
         if episode.is_done:

rllib/env/single_agent_env_runner.py

@@ -560,7 +560,6 @@ def set_state(self, state: StateDict) -> None:
                 key=NUM_ENV_STEPS_SAMPLED_LIFETIME,
                 value=state[NUM_ENV_STEPS_SAMPLED_LIFETIME],
                 reduce="sum",
-                with_throughput=True,
             )
 
     @override(Checkpointable)
@@ -799,7 +798,6 @@ def _increase_sampled_metrics(self, num_steps, num_episodes_completed):
             NUM_ENV_STEPS_SAMPLED_LIFETIME,
             num_steps,
             reduce="sum",
-            with_throughput=True,
         )
         self.metrics.log_value(
             (NUM_AGENT_STEPS_SAMPLED_LIFETIME, DEFAULT_AGENT_ID),

rllib/utils/metrics/metrics_logger.py

@@ -30,8 +30,8 @@ class MetricsLogger:
     - Reducing these collected values using a user specified reduction method (for
     example "min" or "mean") and other settings controlling the reduction and internal
     data, such as sliding windows or EMA coefficients.
-    - Resetting the logged values after a `reduce()` call in order to make space for
-    new values to be logged.
+    - Optionally clearing all logged values after a `reduce()` call to make space for
+    new data.
 
     .. testcode::
 
@@ -233,7 +233,6 @@ def log_value(
         window: Optional[Union[int, float]] = None,
         ema_coeff: Optional[float] = None,
         clear_on_reduce: bool = False,
-        with_throughput: bool = False,
     ) -> None:
         """Logs a new value under a (possibly nested) key to the logger.
 
@@ -329,13 +328,6 @@ def log_value(
                 `self.reduce()` is called. Setting this to True is useful for cases,
                 in which the internal values list would otherwise grow indefinitely,
                 for example if reduce is None and there is no `window` provided.
-            with_throughput: Whether to track a throughput estimate together with this
-                metric. This is only supported for `reduce=sum` and
-                `clear_on_reduce=False` metrics (aka. "lifetime counts"). The `Stats`
-                object under the logged key then keeps track of the time passed
-                between two consecutive calls to `reduce()` and update its throughput
-                estimate. The current throughput estimate of a key can be obtained
-                through: `MetricsLogger.peek([some key], throughput=True)`.
         """
         # No reduction (continue appending to list) AND no window.
         # -> We'll force-reset our values upon `reduce()`.
@@ -358,7 +350,6 @@ def log_value(
                             window=window,
                             ema_coeff=ema_coeff,
                             clear_on_reduce=clear_on_reduce,
-                            throughput=with_throughput,
                         )
                     ),
                 )
@@ -701,14 +692,9 @@ def log_time(
         window: Optional[Union[int, float]] = None,
         ema_coeff: Optional[float] = None,
         clear_on_reduce: bool = False,
-        key_for_throughput: Optional[Union[str, Tuple[str, ...]]] = None,
-        key_for_unit_count: Optional[Union[str, Tuple[str, ...]]] = None,
     ) -> Stats:
         """Measures and logs a time delta value under `key` when used with a with-block.
 
-        Additionally, measures and logs the throughput for the timed code, iff
-        `key_for_throughput` and `key_for_unit_count` are provided.
-
         .. testcode::
 
             import time
@@ -769,32 +755,13 @@ def log_time(
             clear_on_reduce = True
 
         if not self._key_in_stats(key):
-            measure_throughput = None
-            if key_for_unit_count is not None:
-                measure_throughput = True
-                key_for_throughput = key_for_throughput or (key + "_throughput_per_s")
-
             self._set_key(
                 key,
                 Stats(
                     reduce=reduce,
                     window=window,
                     ema_coeff=ema_coeff,
                     clear_on_reduce=clear_on_reduce,
-                    on_exit=(
-                        lambda time_delta_s, kt=key_for_throughput, ku=key_for_unit_count, r=reduce, w=window, e=ema_coeff, c=clear_on_reduce: (  # noqa
-                            self.log_value(
-                                kt,
-                                value=self.peek(ku) / time_delta_s,
-                                reduce=r,
-                                window=w,
-                                ema_coeff=e,
-                                clear_on_reduce=c,
-                            )
-                        )
-                    )
-                    if measure_throughput
-                    else None,
                 ),
             )
 
@@ -1005,7 +972,6 @@ def set_value(
         window: Optional[Union[int, float]] = None,
         ema_coeff: Optional[float] = None,
         clear_on_reduce: bool = False,
-        with_throughput: bool = False,
     ) -> None:
         """Overrides the logged values under `key` with `value`.
 
@@ -1042,13 +1008,6 @@ def set_value(
                 in which the internal values list would otherwise grow indefinitely,
                 for example if reduce is None and there is no `window` provided.
                 Note that this is only applied if `key` does not exist in `self` yet.
-            with_throughput: Whether to track a throughput estimate together with this
-                metric. This is only supported for `reduce=sum` and
-                `clear_on_reduce=False` metrics (aka. "lifetime counts"). The `Stats`
-                object under the logged key then keeps track of the time passed
-                between two consecutive calls to `reduce()` and update its throughput
-                estimate. The current throughput estimate of a key can be obtained
-                through: `MetricsLogger.peek([some key], throughput=True)`.
         """
         # Key already in self -> Erase internal values list with [`value`].
         if self._key_in_stats(key):
@@ -1064,7 +1023,6 @@ def set_value(
                 window=window,
                 ema_coeff=ema_coeff,
                 clear_on_reduce=clear_on_reduce,
-                with_throughput=with_throughput,
             )
 
     def reset(self) -> None:

rllib/utils/metrics/stats.py

@@ -137,7 +137,6 @@ def __init__(
         ema_coeff: Optional[float] = None,
         clear_on_reduce: bool = False,
         on_exit: Optional[Callable] = None,
-        throughput: Union[bool, float] = False,
     ):
         """Initializes a Stats instance.
 
@@ -175,13 +174,6 @@ def __init__(
                 to True is useful for cases, in which the internal values list would
                 otherwise grow indefinitely, for example if reduce is None and there
                 is no `window` provided.
-            with_throughput: Whether to track a throughput estimate together with this
-                Stats. This is only supported for `reduce=sum` and
-                `clear_on_reduce=False` metrics (aka. "lifetime counts"). The `Stats`
-                then keeps track of the time passed between two consecutive calls to
-                `reduce()` and update its throughput estimate. The current throughput
-                estimate of a key can be obtained through:
-                `Stats.peek([some key], throughput=True)`.
         """
         # Thus far, we only support mean, max, min, and sum.
         if reduce not in [None, "mean", "min", "max", "sum"]:
@@ -228,10 +220,14 @@ def __init__(
         # previous `reduce()` result in hist[1].
         self._hist = deque([0, 0, 0], maxlen=3)
 
-        self._throughput = throughput if throughput is not True else 0.0
-        if self._throughput is not False:
-            assert self._reduce_method == "sum"
-            assert self._window in [None, float("inf")]
+        self._throughput = 0.0
+        self._measure_throughput = False
+        if (
+            self._reduce_method == "sum"
+            and self._window in [None, float("inf")]
+            and not self._clear_on_reduce
+        ):
+            self._measure_throughput = True
             self._throughput_last_time = -1
 
     def push(self, value) -> None:
@@ -295,7 +291,7 @@ def peek(self, *, previous: Optional[int] = None, throughput: bool = False) -> A
             return self._hist[-abs(previous)]
         # Return the last measured throughput.
         elif throughput:
-            return self._throughput if self._throughput is not False else None
+            return self._throughput if self._measure_throughput else None
         return self._reduced_values()[0]
 
     def reduce(self) -> "Stats":
@@ -314,7 +310,7 @@ class for details on the reduction logic applied to the values list, based on
         reduced, values = self._reduced_values()
 
         # Keep track and update underlying throughput metric.
-        if self._throughput is not False:
+        if self._measure_throughput:
             # Take the delta between the new (upcoming) reduced value and the most
             # recently reduced value (one `reduce()` call ago).
             delta_sum = reduced - self._hist[-1]
@@ -356,7 +352,7 @@ def merge_on_time_axis(self, other: "Stats") -> None:
             self.values = self.values[-self._window :]
 
         # Adopt `other`'s current throughput estimate (it's the newer one).
-        if self._throughput is not False:
+        if self._measure_throughput:
             self._throughput = other._throughput
 
     def merge_in_parallel(self, *others: "Stats") -> None:
@@ -647,8 +643,8 @@ def similar_to(
             window=other._window,
             ema_coeff=other._ema_coeff,
             clear_on_reduce=other._clear_on_reduce,
-            throughput=other._throughput,
         )
+        stats._throughput = other._throughput
         stats._hist = other._hist
         return stats