From 876a1ba5bd0c5570dae7de7e3334bbe85e561d17 Mon Sep 17 00:00:00 2001
From: Sven Mika <sven@anyscale.io>
Date: Fri, 6 Mar 2020 21:45:30 +0100
Subject: [PATCH] [RLlib] Issue 7421: can't convert cuda tensor to numpy in
 torch ppo. (#7445)

---
 rllib/policy/torch_policy_template.py |  5 +++--
 rllib/utils/torch_ops.py              | 23 ++++++++++++++---------
 2 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/rllib/policy/torch_policy_template.py b/rllib/policy/torch_policy_template.py
index 01d963ea7..df67c4f13 100644
--- a/rllib/policy/torch_policy_template.py
+++ b/rllib/policy/torch_policy_template.py
@@ -105,8 +105,9 @@ def build_torch_policy(name,
             # Do all post-processing always with no_grad().
             # Not using this here will introduce a memory leak (issue #6962).
             with torch.no_grad():
-                return postprocess_fn(self, sample_batch, other_agent_batches,
-                                      episode)
+                return postprocess_fn(
+                    self, convert_to_non_torch_type(sample_batch),
+                    convert_to_non_torch_type(other_agent_batches), episode)
 
         @override(TorchPolicy)
         def extra_grad_process(self):
diff --git a/rllib/utils/torch_ops.py b/rllib/utils/torch_ops.py
index c762b0670..c53d53e54 100644
--- a/rllib/utils/torch_ops.py
+++ b/rllib/utils/torch_ops.py
@@ -1,3 +1,5 @@
+import tree
+
 from ray.rllib.utils.framework import try_import_torch
 
 torch, _ = try_import_torch()
@@ -20,21 +22,24 @@ def sequence_mask(lengths, maxlen, dtype=None):
     return mask
 
 
-def convert_to_non_torch_type(stats_dict):
+def convert_to_non_torch_type(stats):
     """Converts values in stats_dict to non-Tensor numpy or python types.
 
     Args:
-        stats_dict (dict): A flat key, value dict, the values of which will be
-            converted and returned as a new dict.
+        stats (any): Any (possibly nested) struct, the values in which will be
+            converted and returned as a new struct with all torch tensors
+            being converted to numpy types.
 
     Returns:
         dict: A new dict with the same structure as stats_dict, but with all
             values converted to non-torch Tensor types.
     """
-    ret = {}
-    for k, v in stats_dict.items():
-        if isinstance(v, torch.Tensor):
-            ret[k] = v.cpu().item() if len(v.size()) == 0 else v.cpu().numpy()
+    # The mapping function used to numpyize torch Tensors.
+    def mapping(item):
+        if isinstance(item, torch.Tensor):
+            return item.cpu().item() if len(item.size()) == 0 else \
+                item.cpu().numpy()
         else:
-            ret[k] = v
-    return ret
+            return item
+
+    return tree.map_structure(mapping, stats)