[RLlib] Fix convert to torch tensor (ray-project#31023)

Signed-off-by: Artur Niederfahrenhorst <[email protected]>

rllib/BUILD

@@ -2109,6 +2109,13 @@ py_test(
  srcs = ["utils/exploration/tests/test_random_encoder.py"]
 )
 
+py_test(
+ name = "utils/tests/test_torch_utils",
+ tags = ["team:rllib", "policutils"],
+ size = "small",
+ srcs = ["utils/tests/test_torch_utils.py"]
+)
+
 # Schedules
 py_test(
  name = "test_schedules",

rllib/utils/tests/test_torch_utils.py

@@ -0,0 +1,45 @@
+import unittest
+
+import numpy as np
+import torch.cuda
+
+import ray
+from ray.rllib.utils.torch_utils import convert_to_torch_tensor
+
+
+class TestTorchUtils(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls) -> None:
+ ray.init()
+
+ @classmethod
+ def tearDownClass(cls) -> None:
+ ray.shutdown()
+
+ def test_convert_to_torch_tensor(self):
+ # Tests whether convert_to_torch_tensor works as expected
+
+ # Test single array
+ array = np.array([1, 2, 3])
+ tensor = torch.from_numpy(array)
+ self.assertTrue(all(convert_to_torch_tensor(array) == tensor))
+
+ # Test torch tensor
+ self.assertTrue(convert_to_torch_tensor(tensor) is tensor)
+
+ # Test conversion to 32-bit float
+ tensor_2 = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64)
+ self.assertTrue(convert_to_torch_tensor(tensor_2).dtype is torch.float32)
+
+ # Test nested structure with objects tested above
+ converted = convert_to_torch_tensor({"a": (array, tensor), "b": tensor_2})
+ self.assertTrue(all(convert_to_torch_tensor(converted["a"][0]) == tensor))
+ self.assertTrue(convert_to_torch_tensor(converted["a"][1]) is tensor)
+ self.assertTrue(convert_to_torch_tensor(converted["b"]).dtype is torch.float32)
+
+
+if __name__ == "__main__":
+ import pytest
+ import sys
+
+ sys.exit(pytest.main(["-v", __file__]))

rllib/utils/torch_utils.py

@@ -149,7 +149,7 @@ def mapping(item):
  elif isinstance(item, np.ndarray):
  # Object type (e.g. info dicts in train batch): leave as-is.
  # str type (e.g. agent_id in train batch): leave as-is.
- if item.dtype == object or isinstance(item.dtype, type(np.dtype("str_"))):
+ if item.dtype == object or item.dtype.type is np.str_:
  return item
  # Non-writable numpy-arrays will cause PyTorch warning.
  elif item.flags.writeable is False: