[RLlib] MetricsLogger: Fix `get/set_state` to handle tensors in `self.values`. #53514

sven1977 · 2025-06-03T11:13:23Z

MetricsLogger: Fix get/set_state to handle tensors in self.values.

Stats.get_state should never return tensors (converts to numpy now).
Stats keeps track of potential tensor insertions, triggering the new self._may_have_tensors flag to be True.
State.from_state now looks at the self._may_have_tensors flag and does NOT populate the self.values field, if it's True.

Why are these changes needed?

Related issue number

Closes #53467

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

Signed-off-by: sven1977 <[email protected]>

Copilot

Pull Request Overview

This PR updates Stats.get_state and Stats.from_state to correctly handle tensor values in self.values, ensuring state serialization avoids returning raw tensors.

Introduces convert_to_numpy in get_state to serialize any tensors to NumPy.
Adds a _could_be_tensor flag to track potential tensor insertions and skips repopulating values in from_state when tensors were present.
Refactors check_value to detect zero-dimensional tensors and enforce scalar reduction behavior.

Comments suppressed due to low confidence (3)

rllib/utils/metrics/stats.py:187

[nitpick] The private flag _could_be_tensor may be clearer as _may_have_tensors or _has_tensors to reflect its boolean purpose more directly.

self._could_be_tensor = False

rllib/utils/metrics/stats.py:639

There are no explicit unit tests covering tensor serialization in get_state and recovery in from_state. Adding tests for pushing Torch/TF tensors and ensuring correct NumPy output will prevent regressions.

def get_state(self) -> Dict[str, Any]:

rllib/utils/metrics/stats.py:642

convert_to_numpy may not handle deque inputs uniformly. It could be safer to convert self.values to a list first, e.g., convert_to_numpy(list(self.values)).

"values": convert_to_numpy(self.values),

rllib/utils/metrics/stats.py

Copilot · 2025-06-03T11:15:20Z

rllib/utils/metrics/stats.py

@@ -350,17 +358,20 @@ class for details on the reduction logic applied to the values list, based on
        """
        len_before_reduce = len(self)
        if self._has_new_values:
-            # Only calculate and update history if there were new values pushed since last reduce
+            # Only calculate and update history if there were new values pushed since'


There's a stray apostrophe after since. Removing it will clean up the comment.

Suggested change

# Only calculate and update history if there were new values pushed since'

# Only calculate and update history if there were new values pushed since

Copilot · 2025-06-03T11:15:20Z

rllib/utils/metrics/stats.py

        }
        if self._throughput_stats is not None:
            state["throughput_stats"] = self._throughput_stats.get_state()
        return state

    @staticmethod
-    def from_state(state: Dict[str, Any], throughputs=False) -> "Stats":
+    def from_state(state: Dict[str, Any]) -> "Stats":


The throughputs parameter was removed from from_state, which may break existing callers relying on that signature. Consider retaining a backward-compatible overload or documenting the change.

ArturNiederfahrenhorst · 2025-06-03T11:52:06Z

rllib/utils/metrics/stats.py

+        # whether we are on a supported device).
+        values = state["values"]
+        if "_could_be_tensor" in state and state["_could_be_tensor"]:
+            values = []


If we do this, we should state it as a known limitation in metrics-logger.rst.

Alternatively, can we noch check whether we are on a supported device and keep track of GPU tensors?

That's true. This is a limitation for now. Albeit a very small one since - normally - tensor logging is only done for loss metrics, like loss, entropy, etc.. and these are very very often ephemeral values where it's not a problem at all to just start fresh after a checkpoint loading (window=1 anyways?).

Yeah, but then we would have to store the original device as well, which quickly becomes super messy (when you transfer a checkpoint from one cluster type (GPUs) to another (no GPUs?)).

What if window!=1? It's a user-facing class.

Rodger! Yeah there is some ugly complexity here.

ArturNiederfahrenhorst · 2025-06-03T11:53:08Z

rllib/utils/metrics/stats.py

@@ -183,29 +184,33 @@ def __init__(
        self.values: Union[List, deque.Deque] = None
        self._set_values(force_list(init_values))

+        self._could_be_tensor = False


From reading the code self._could_be_tensor can be renamed to self._is_tensor?

ArturNiederfahrenhorst

Left some comments but - no blockers.

Signed-off-by: sven1977 <[email protected]>

…ics_logger_get_set_state_handling_tensors

Signed-off-by: sven1977 <[email protected]>

….values`. (#53514) Signed-off-by: elliot-barn <[email protected]>

….values`. (ray-project#53514)

….values`. (#53514) Signed-off-by: elliot-barn <[email protected]>

wip

25af3c8

Signed-off-by: sven1977 <[email protected]>

Copilot AI review requested due to automatic review settings June 3, 2025 11:13

sven1977 requested a review from a team as a code owner June 3, 2025 11:13

Copilot AI reviewed Jun 3, 2025

View reviewed changes

ArturNiederfahrenhorst reviewed Jun 3, 2025

View reviewed changes

ArturNiederfahrenhorst approved these changes Jun 3, 2025

View reviewed changes

wip

40b6490

Signed-off-by: sven1977 <[email protected]>

sven1977 enabled auto-merge (squash) June 3, 2025 12:20

github-actions bot disabled auto-merge June 3, 2025 12:20

github-actions bot added the go add ONLY when ready to merge, run all tests label Jun 3, 2025

sven1977 enabled auto-merge (squash) June 3, 2025 12:23

wip

caa40e8

Signed-off-by: sven1977 <[email protected]>

github-actions bot disabled auto-merge June 3, 2025 13:25

wip

86f795e

Signed-off-by: sven1977 <[email protected]>

sven1977 enabled auto-merge (squash) June 5, 2025 14:17

wip

c19d560

Signed-off-by: sven1977 <[email protected]>

github-actions bot disabled auto-merge June 13, 2025 08:08

wip

4288532

Signed-off-by: sven1977 <[email protected]>

sven1977 enabled auto-merge (squash) June 13, 2025 11:04

sven1977 added 2 commits June 17, 2025 14:31

Merge branch 'master' of https://github.com/ray-project/ray into metr…

14d0315

…ics_logger_get_set_state_handling_tensors

wip

cf23225

Signed-off-by: sven1977 <[email protected]>

github-actions bot disabled auto-merge June 17, 2025 12:40

sven1977 merged commit 8ab0868 into ray-project:master Jun 17, 2025
5 checks passed

elliot-barn pushed a commit that referenced this pull request Jun 18, 2025

[RLlib] MetricsLogger: Fix get/set_state to handle tensors in `self…

3775dfd

….values`. (#53514) Signed-off-by: elliot-barn <[email protected]>

minerharry pushed a commit to minerharry/ray that referenced this pull request Jun 27, 2025

[RLlib] MetricsLogger: Fix get/set_state to handle tensors in `self…

15abbf9

….values`. (ray-project#53514)

elliot-barn pushed a commit that referenced this pull request Jul 2, 2025

[RLlib] MetricsLogger: Fix get/set_state to handle tensors in `self…

d4c0dd0

….values`. (#53514) Signed-off-by: elliot-barn <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[RLlib] MetricsLogger: Fix `get/set_state` to handle tensors in `self.values`. #53514

[RLlib] MetricsLogger: Fix `get/set_state` to handle tensors in `self.values`. #53514

Uh oh!

sven1977 commented Jun 3, 2025 •

edited by ArturNiederfahrenhorst

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Copilot AI Jun 3, 2025

Uh oh!

Copilot AI Jun 3, 2025

Uh oh!

ArturNiederfahrenhorst Jun 3, 2025

Uh oh!

sven1977 Jun 3, 2025 •

edited

Loading

Uh oh!

ArturNiederfahrenhorst Jun 5, 2025

Uh oh!

ArturNiederfahrenhorst Jun 3, 2025

Uh oh!

ArturNiederfahrenhorst left a comment

Uh oh!

Uh oh!

Uh oh!

	# Only calculate and update history if there were new values pushed since'
	# Only calculate and update history if there were new values pushed since

[RLlib] MetricsLogger: Fix get/set_state to handle tensors in self.values. #53514

[RLlib] MetricsLogger: Fix get/set_state to handle tensors in self.values. #53514

Uh oh!

Conversation

sven1977 commented Jun 3, 2025 • edited by ArturNiederfahrenhorst Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why are these changes needed?

Related issue number

Checks

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Uh oh!

Uh oh!

Copilot AI Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

ArturNiederfahrenhorst Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

sven1977 Jun 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ArturNiederfahrenhorst Jun 5, 2025

Choose a reason for hiding this comment

Uh oh!

ArturNiederfahrenhorst Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

ArturNiederfahrenhorst left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

[RLlib] MetricsLogger: Fix `get/set_state` to handle tensors in `self.values`. #53514

[RLlib] MetricsLogger: Fix `get/set_state` to handle tensors in `self.values`. #53514

sven1977 commented Jun 3, 2025 •

edited by ArturNiederfahrenhorst

Loading

sven1977 Jun 3, 2025 •

edited

Loading