diff --git a/rllib/utils/torch_ops.py b/rllib/utils/torch_ops.py index 53c26dc4d..ce6c86a16 100644 --- a/rllib/utils/torch_ops.py +++ b/rllib/utils/torch_ops.py @@ -63,11 +63,20 @@ def convert_to_torch_tensor(x, device=None): return RepeatedValues( tree.map_structure(mapping, item.values), item.lengths, item.max_len) - # Non-writable numpy-arrays will cause PyTorch warning. - if isinstance(item, np.ndarray) and item.flags.writeable is False: - with warnings.catch_warnings(): - warnings.simplefilter("ignore") + # Numpy arrays. + if isinstance(item, np.ndarray): + # np.object_ type (e.g. info dicts in train batch): leave as-is. + if item.dtype == np.object_: + return item + # Non-writable numpy-arrays will cause PyTorch warning. + elif item.flags.writeable is False: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + tensor = torch.from_numpy(item) + # Already numpy: Wrap as torch tensor. + else: tensor = torch.from_numpy(item) + # Everything else: Convert to numpy, then wrap as torch tensor. else: tensor = torch.from_numpy(np.asarray(item)) # Floatify all float64 tensors.