From b37a1620761475f2c534a3b39f4a9bd62973514c Mon Sep 17 00:00:00 2001
From: Sven Mika <sven@anyscale.io>
Date: Wed, 3 Jun 2020 08:14:29 +0200
Subject: [PATCH] [RLlib] Make envs specifiable in configs by their class path.
 (#8750)

---
 rllib/BUILD                                             | 2 +-
 rllib/agents/trainer.py                                 | 7 +++++++
 rllib/utils/from_config.py                              | 2 ++
 rllib/utils/tests/test_framework_agnostic_components.py | 8 ++++++++
 4 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/rllib/BUILD b/rllib/BUILD
index 11d91e101..29aeead47 100644
--- a/rllib/BUILD
+++ b/rllib/BUILD
@@ -1120,7 +1120,7 @@ py_test(
 py_test(
     name = "tests/test_avail_actions_qmix",
     tags = ["tests_dir", "tests_dir_A"],
-    size = "small",
+    size = "medium",
     srcs = ["tests/test_avail_actions_qmix.py"]
 )
 
diff --git a/rllib/agents/trainer.py b/rllib/agents/trainer.py
index 5b4221d68..64bb4d93a 100644
--- a/rllib/agents/trainer.py
+++ b/rllib/agents/trainer.py
@@ -20,6 +20,7 @@ from ray.rllib.utils import FilterManager, deep_update, merge_dicts
 from ray.rllib.utils.framework import check_framework, try_import_tf
 from ray.rllib.utils.annotations import override, PublicAPI, DeveloperAPI
 from ray.rllib.utils.deprecation import DEPRECATED_VALUE, deprecation_warning
+from ray.rllib.utils.from_config import from_config
 from ray.tune.registry import ENV_CREATOR, register_env, _global_registry
 from ray.tune.trainable import Trainable
 from ray.tune.trial import ExportFormat
@@ -535,8 +536,14 @@ class Trainer(Trainable):
         env = self._env_id
         if env:
             config["env"] = env
+            # An already registered env.
             if _global_registry.contains(ENV_CREATOR, env):
                 self.env_creator = _global_registry.get(ENV_CREATOR, env)
+            # A class specifier.
+            elif "." in env:
+                self.env_creator = \
+                    lambda env_config: from_config(env, env_config)
+            # Try gym.
             else:
                 import gym  # soft dependency
                 self.env_creator = lambda env_config: gym.make(env)
diff --git a/rllib/utils/from_config.py b/rllib/utils/from_config.py
index dc401284d..af8bcd30a 100644
--- a/rllib/utils/from_config.py
+++ b/rllib/utils/from_config.py
@@ -73,6 +73,8 @@ def from_config(cls, config=None, **kwargs):
         pass
     if isinstance(config, dict):
         type_ = config.pop("type", None)
+        if type_ is None and isinstance(cls, str):
+            type_ = cls
         ctor_kwargs = config
         # Give kwargs priority over things defined in config dict.
         # This way, one can pass a generic `spec` and then override single
diff --git a/rllib/utils/tests/test_framework_agnostic_components.py b/rllib/utils/tests/test_framework_agnostic_components.py
index 79d970793..400c16c40 100644
--- a/rllib/utils/tests/test_framework_agnostic_components.py
+++ b/rllib/utils/tests/test_framework_agnostic_components.py
@@ -165,6 +165,14 @@ class TestFrameWorkAgnosticComponents(unittest.TestCase):
                 value = sess.run(value)
             check(value, np.array([-6.6]))  # prop_b == -1.5
 
+    def test_unregistered_envs(self):
+        """Tests, whether an Env can be specified simply by its absolute class.
+        """
+        env_cls = "ray.rllib.examples.env.stateless_cartpole.StatelessCartPole"
+        env = from_config(env_cls, {"config": 42.0})
+        state = env.reset()
+        self.assertTrue(state.shape == (2, ))
+
 
 if __name__ == "__main__":
     import pytest