From daf0bef2858441e3d2da953d9a76400b6ce7a77d Mon Sep 17 00:00:00 2001
From: Sven Mika <sven@anyscale.io>
Date: Thu, 21 Jan 2021 16:30:26 +0100
Subject: [PATCH] [RLlib] Dreamer: Fix broken import and add compilation test
 case. (#13553)

---
 rllib/BUILD                                | 23 ++++++++----
 rllib/agents/dreamer/dreamer.py            |  2 ++
 rllib/agents/dreamer/dreamer_model.py      |  2 +-
 rllib/agents/dreamer/tests/test_dreamer.py | 41 ++++++++++++++++++++++
 rllib/env/wrappers/dm_control_wrapper.py   |  2 +-
 5 files changed, 61 insertions(+), 9 deletions(-)
 create mode 100644 rllib/agents/dreamer/tests/test_dreamer.py

diff --git a/rllib/BUILD b/rllib/BUILD
index daa623dff..f8f1cbd3c 100644
--- a/rllib/BUILD
+++ b/rllib/BUILD
@@ -436,13 +436,13 @@ py_test(
     srcs = ["agents/a3c/tests/test_a3c.py"]
 )
 
-## APEXTrainer (DQN)
-#py_test(
-#    name = "test_apex_dqn",
-#    tags = ["agents_dir"],
-#    size = "large",
-#    srcs = ["agents/dqn/tests/test_apex_dqn.py"]
-#)
+# APEXTrainer (DQN)
+py_test(
+    name = "test_apex_dqn",
+    tags = ["agents_dir"],
+    size = "medium",
+    srcs = ["agents/dqn/tests/test_apex_dqn.py"]
+)
 
 # APEXDDPGTrainer
 py_test(
@@ -482,6 +482,15 @@ py_test(
     srcs = ["agents/dqn/tests/test_simple_q.py"]
 )
 
+# TODO: enable once we have a MuJoCo-independent test case.
+## Dreamer
+#py_test(
+#    name = "test_dreamer",
+#    tags = ["agents_dir"],
+#    size = "small",
+#    srcs = ["agents/dreamer/tests/test_dreamer.py"]
+#)
+
 # ES
 py_test(
     name = "test_es",
diff --git a/rllib/agents/dreamer/dreamer.py b/rllib/agents/dreamer/dreamer.py
index 94774d9fe..21646d618 100644
--- a/rllib/agents/dreamer/dreamer.py
+++ b/rllib/agents/dreamer/dreamer.py
@@ -31,6 +31,8 @@ DEFAULT_CONFIG = with_common_config({
     "discount": 0.99,
     # Lambda
     "lambda": 0.95,
+    # Clipping is done inherently via policy tanh.
+    "clip_actions": False,
     # Training iterations per data collection from real env
     "dreamer_train_iters": 100,
     # Horizon for Enviornment (1000 for Mujoco/DMC)
diff --git a/rllib/agents/dreamer/dreamer_model.py b/rllib/agents/dreamer/dreamer_model.py
index 5483f664f..f2db417e5 100644
--- a/rllib/agents/dreamer/dreamer_model.py
+++ b/rllib/agents/dreamer/dreamer_model.py
@@ -1,6 +1,6 @@
 import numpy as np
 from typing import Any, List, Tuple
-from ray.rllib.models.torch.modules.reshape import Reshape
+from ray.rllib.models.torch.misc import Reshape
 from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
 from ray.rllib.utils.framework import try_import_torch
 from ray.rllib.utils.framework import TensorType
diff --git a/rllib/agents/dreamer/tests/test_dreamer.py b/rllib/agents/dreamer/tests/test_dreamer.py
new file mode 100644
index 000000000..2b318866c
--- /dev/null
+++ b/rllib/agents/dreamer/tests/test_dreamer.py
@@ -0,0 +1,41 @@
+import unittest
+
+import ray
+from ray import tune
+import ray.rllib.agents.dreamer as dreamer
+from ray.rllib.examples.env.dm_control_suite import hopper_hop
+from ray.rllib.utils.test_utils import check_compute_single_action, \
+    framework_iterator
+
+
+class TestDreamer(unittest.TestCase):
+    """Sanity tests for DreamerTrainer."""
+
+    def setUp(self):
+        ray.init()
+
+    def tearDown(self):
+        ray.shutdown()
+
+    def test_dreamer_compilation(self):
+        """Test whether an DreamerTrainer can be built with all frameworks."""
+        config = dreamer.DEFAULT_CONFIG.copy()
+        tune.register_env("dm_control_hopper_hop", lambda _: hopper_hop())
+
+        num_iterations = 1
+
+        # Test against all frameworks.
+        for _ in framework_iterator(config, frameworks="torch"):
+            for env in ["dm_control_hopper_hop"]:
+                trainer = dreamer.DREAMERTrainer(config=config, env=env)
+                for i in range(num_iterations):
+                    results = trainer.train()
+                    print(results)
+                check_compute_single_action(trainer)
+                trainer.stop()
+
+
+if __name__ == "__main__":
+    import pytest
+    import sys
+    sys.exit(pytest.main(["-v", __file__]))
diff --git a/rllib/env/wrappers/dm_control_wrapper.py b/rllib/env/wrappers/dm_control_wrapper.py
index 6734e2a3a..3286aae28 100644
--- a/rllib/env/wrappers/dm_control_wrapper.py
+++ b/rllib/env/wrappers/dm_control_wrapper.py
@@ -31,7 +31,7 @@ except ImportError:
     specs = None
 try:
     from dm_control import suite
-except ImportError:
+except (ImportError, OSError):
     suite = None
 import numpy as np