Repository: ray-project/ray

Issue:
[rllib] ModuleNotFoundError: No module named 'tree'
*Ray version and other system information (Python version, TensorFlow version, OS):*
Python3.6 
pip install -U ray-0.9.0.dev0-cp36-cp36m-manylinux1_x86_64.whl
On 2020-03-18
### Reproduction (REQUIRED)
$ rllib train -f atari-ddppo.yaml

Traceback (most recent call last):
  File "/home/simon/anaconda3/bin/rllib", line 5, in <module>
    from ray.rllib.scripts import cli
  File "/home/simon/anaconda3/lib/python3.6/site-packages/ray/rllib/__init__.py", line 9, in <module>
    from ray.rllib.evaluation.policy_graph import PolicyGraph
  File "/home/simon/anaconda3/lib/python3.6/site-packages/ray/rllib/evaluation/__init__.py", line 2, in <module>
    from ray.rllib.evaluation.rollout_worker import RolloutWorker
  File "/home/simon/anaconda3/lib/python3.6/site-packages/ray/rllib/evaluation/rollout_worker.py", line 19, in <module>
    from ray.rllib.evaluation.sampler import AsyncSampler, SyncSampler
  File "/home/simon/anaconda3/lib/python3.6/site-packages/ray/rllib/evaluation/sampler.py", line 11, in <module>
    from ray.rllib.evaluation.sample_batch_builder import \
  File "/home/simon/anaconda3/lib/python3.6/site-packages/ray/rllib/evaluation/sample_batch_builder.py", line 6, in <module>
    from ray.rllib.policy.sample_batch import SampleBatch, MultiAgentBatch
  File "/home/simon/anaconda3/lib/python3.6/site-packages/ray/rllib/policy/__init__.py", line 2, in <module>
    from ray.rllib.policy.torch_policy import TorchPolicy
  File "/home/simon/anaconda3/lib/python3.6/site-packages/ray/rllib/policy/torch_policy.py", line 10, in <module>
    from ray.rllib.utils.torch_ops import convert_to_non_torch_type
  File "/home/simon/anaconda3/lib/python3.6/site-packages/ray/rllib/utils/torch_ops.py", line 1, in <module>
    import tree
ModuleNotFoundError: No module named 'tree'




Discussion:

Resolving patch:
diff --git a/rllib/policy/eager_tf_policy.py b/rllib/policy/eager_tf_policy.py
--- a/rllib/policy/eager_tf_policy.py
+++ b/rllib/policy/eager_tf_policy.py
@@ -5,7 +5,6 @@
 import logging
 import functools
 import numpy as np
-import tree
 
 from ray.util.debug import log_once
 from ray.rllib.evaluation.episode import _flatten_action
@@ -24,12 +23,12 @@
 def _convert_to_tf(x):
     if isinstance(x, SampleBatch):
         x = {k: v for k, v in x.items() if k != SampleBatch.INFOS}
-        return tree.map_structure(_convert_to_tf, x)
+        return tf.nest.map_structure(_convert_to_tf, x)
     if isinstance(x, Policy):
         return x
 
     if x is not None:
-        x = tree.map_structure(
+        x = tf.nest.map_structure(
             lambda f: tf.convert_to_tensor(f) if f is not None else None, x)
     return x
 
@@ -38,7 +37,7 @@ def _convert_to_numpy(x):
     if x is None:
         return None
     try:
-        return tree.map_structure(lambda component: component.numpy(), x)
+        return tf.nest.map_structure(lambda component: component.numpy(), x)
     except AttributeError:
         raise TypeError(
             ("Object of type {} has no method to convert to numpy.").format(
@@ -66,7 +65,7 @@ def convert_eager_outputs(func):
     def _func(*args, **kwargs):
         out = func(*args, **kwargs)
         if tf.executing_eagerly():
-            out = tree.map_structure(_convert_to_numpy, out)
+            out = tf.nest.map_structure(_convert_to_numpy, out)
         return out
 
     return _func
@@ -551,7 +550,7 @@ def _initialize_loss_with_dummy_batch(self):
                 SampleBatch.NEXT_OBS: np.array(
                     [self.observation_space.sample()]),
                 SampleBatch.DONES: np.array([False], dtype=np.bool),
-                SampleBatch.ACTIONS: tree.map_structure(
+                SampleBatch.ACTIONS: tf.nest.map_structure(
                     lambda c: np.array([c]), self.action_space.sample()),
                 SampleBatch.REWARDS: np.array([0], dtype=np.float32),
             }
@@ -568,7 +567,8 @@ def _initialize_loss_with_dummy_batch(self):
                 dummy_batch["seq_lens"] = np.array([1], dtype=np.int32)
 
             # Convert everything to tensors.
-            dummy_batch = tree.map_structure(tf.convert_to_tensor, dummy_batch)
+            dummy_batch = tf.nest.map_structure(tf.convert_to_tensor,
+                                                dummy_batch)
 
             # for IMPALA which expects a certain sample batch size.
             def tile_to(tensor, n):
@@ -576,7 +576,7 @@ def tile_to(tensor, n):
                                [n] + [1 for _ in tensor.shape.as_list()[1:]])
 
             if get_batch_divisibility_req:
-                dummy_batch = tree.map_structure(
+                dummy_batch = tf.nest.map_structure(
                     lambda c: tile_to(c, get_batch_divisibility_req(self)),
                     dummy_batch)
 
@@ -595,7 +595,7 @@ def tile_to(tensor, n):
             # overwrite any tensor state from that call)
             self.model.from_batch(dummy_batch)
 
-            postprocessed_batch = tree.map_structure(
+            postprocessed_batch = tf.nest.map_structure(
                 lambda c: tf.convert_to_tensor(c), postprocessed_batch.data)
 
             loss_fn(self, self.model, self.dist_class, postprocessed_batch)
diff --git a/rllib/policy/tf_policy.py b/rllib/policy/tf_policy.py
--- a/rllib/policy/tf_policy.py
+++ b/rllib/policy/tf_policy.py
@@ -1,7 +1,6 @@
 import errno
 import logging
 import os
-import tree
 
 import numpy as np
 import ray
@@ -212,10 +211,8 @@ def _initialize_loss(self, loss, loss_inputs):
             self._loss = loss
 
         self._optimizer = self.optimizer()
-        self._grads_and_vars = [
-            (g, v) for (g, v) in self.gradients(self._optimizer, self._loss)
-            if g is not None
-        ]
+        self._grads_and_vars = [(g, v) for (g, v) in self.gradients(
+            self._optimizer, self._loss) if g is not None]
         self._grads = [g for (g, v) in self._grads_and_vars]
 
         # TODO(sven/ekl): Deprecate support for v1 models.
@@ -493,7 +490,7 @@ def _build_signature_def(self):
 
         # build output signatures
         output_signature = self._extra_output_signature_def()
-        for i, a in enumerate(tree.flatten(self._sampled_action)):
+        for i, a in enumerate(tf.nest.flatten(self._sampled_action)):
             output_signature["actions_{}".format(i)] = \
                 tf.saved_model.utils.build_tensor_info(a)
 
diff --git a/rllib/utils/torch_ops.py b/rllib/utils/torch_ops.py
--- a/rllib/utils/torch_ops.py
+++ b/rllib/utils/torch_ops.py
@@ -1,9 +1,17 @@
-import tree
+import logging
 
 from ray.rllib.utils.framework import try_import_torch
 
 torch, _ = try_import_torch()
 
+logger = logging.getLogger(__name__)
+
+try:
+    import tree
+except (ImportError, ModuleNotFoundError) as e:
+    logger.warning("`dm-tree` is not installed! Run `pip install dm-tree`.")
+    raise e
+
 
 def sequence_mask(lengths, maxlen, dtype=None):
     """
@@ -34,6 +42,7 @@ def convert_to_non_torch_type(stats):
         dict: A new dict with the same structure as stats_dict, but with all
             values converted to non-torch Tensor types.
     """
+
     # The mapping function used to numpyize torch Tensors.
     def mapping(item):
         if isinstance(item, torch.Tensor):


Test patch:
diff --git a/rllib/policy/tests/test_compute_log_likelihoods.py b/rllib/policy/tests/test_compute_log_likelihoods.py
--- a/rllib/policy/tests/test_compute_log_likelihoods.py
+++ b/rllib/policy/tests/test_compute_log_likelihoods.py
@@ -1,5 +1,6 @@
 import numpy as np
 from scipy.stats import norm
+from tensorflow.python.eager.context import eager_mode
 import unittest
 
 import ray.rllib.agents.dqn as dqn
@@ -43,6 +44,14 @@ def do_test_log_likelihood(run,
         config["eager"] = fw == "eager"
         config["use_pytorch"] = fw == "torch"
 
+        eager_ctx = None
+        if fw == "eager":
+            eager_ctx = eager_mode()
+            eager_ctx.__enter__()
+            assert tf.executing_eagerly()
+        elif fw == "tf":
+            assert not tf.executing_eagerly()
+
         trainer = run(config=config, env=env)
         policy = trainer.get_policy()
         vars = policy.get_weights()
@@ -104,6 +113,9 @@ def do_test_log_likelihood(run,
                     prev_reward_batch=np.array([prev_r]))
                 check(np.exp(logp), expected_prob, atol=0.2)
 
+        if eager_ctx:
+            eager_ctx.__exit__(None, None, None)
+
 
 class TestComputeLogLikelihood(unittest.TestCase):
     def test_dqn(self):