Skip to content

[RLlib] Fix device check in Learner.#53706

Merged
sven1977 merged 4 commits intoray-project:masterfrom
simonsays1980:fix-device-check-in-learner
Jun 13, 2025
Merged

[RLlib] Fix device check in Learner.#53706
sven1977 merged 4 commits intoray-project:masterfrom
simonsays1980:fix-device-check-in-learner

Conversation

@simonsays1980
Copy link
Copy Markdown
Contributor

@simonsays1980 simonsays1980 commented Jun 10, 2025

Why are these changes needed?

In the Learner is is checked, if the training batch is already on the correct device. For this it is checked, if

  • the "obs" column of the batch is a numpy array or
  • in case that it is already a tensor (other cases should not be possible at this stage of the pipeline) if the tensor is already on the correct device.

This can throw an exception, if the observations are no an array-type (e.g. gym.spaces.Dict, gym.spaces.Tuple). This PR fixes this by switching the column of the batch to be checked to a simpler one, namely the "rewards" (which are always of simple vector format).

Related issue number

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: simonsays1980 <simon.zehnder@gmail.com>
@simonsays1980 simonsays1980 marked this pull request as ready for review June 10, 2025 17:41
Copilot AI review requested due to automatic review settings June 10, 2025 17:41
@simonsays1980 simonsays1980 requested a review from a team as a code owner June 10, 2025 17:41
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR fixes the device check in the Learner by switching the validation from the "obs" column to the "rewards" column, ensuring that the correct device property is checked for training data.

  • Updated the import to include the Columns constant.
  • Replaced "obs" with Columns.REWARDS in the device check condition.

Signed-off-by: simonsays1980 <simon.zehnder@gmail.com>
Signed-off-by: simonsays1980 <simon.zehnder@gmail.com>
@simonsays1980 simonsays1980 requested a review from sven1977 June 11, 2025 16:49
@sven1977 sven1977 enabled auto-merge (squash) June 13, 2025 11:50
@github-actions github-actions bot added the go add ONLY when ready to merge, run all tests label Jun 13, 2025
Copy link
Copy Markdown
Contributor

@sven1977 sven1977 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! Thanks for the fix @simonsays1980 !

@sven1977 sven1977 merged commit c7436b7 into ray-project:master Jun 13, 2025
7 checks passed
elliot-barn pushed a commit that referenced this pull request Jun 18, 2025
Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>
elliot-barn pushed a commit that referenced this pull request Jul 2, 2025
Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

go add ONLY when ready to merge, run all tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants