[rllib] [rfc] add contrib module and guideline for merging (#3565)

This adds guidelines for merging code into `rllib/contrib` vs `rllib/agents`. Also, clean up the agent import code to make registration easier.

python/ray/rllib/agents/agent.py

@@ -10,7 +10,6 @@ import pickle
 import six
 import tempfile
 import tensorflow as tf
-import traceback
 from types import FunctionType
 
 import ray
@@ -542,69 +541,3 @@ def _register_if_needed(env_object):
         name = env_object.__name__
         register_env(name, lambda config: env_object(config))
         return name
-
-
-def get_agent_class(alg):
-    """Returns the class of a known agent given its name."""
-
-    try:
-        return _get_agent_class(alg)
-    except ImportError:
-        from ray.rllib.agents.mock import _agent_import_failed
-        return _agent_import_failed(traceback.format_exc())
-
-
-def _get_agent_class(alg):
-    if alg == "DDPG":
-        from ray.rllib.agents import ddpg
-        return ddpg.DDPGAgent
-    elif alg == "APEX_DDPG":
-        from ray.rllib.agents import ddpg
-        return ddpg.ApexDDPGAgent
-    elif alg == "PPO":
-        from ray.rllib.agents import ppo
-        return ppo.PPOAgent
-    elif alg == "ES":
-        from ray.rllib.agents import es
-        return es.ESAgent
-    elif alg == "ARS":
-        from ray.rllib.agents import ars
-        return ars.ARSAgent
-    elif alg == "DQN":
-        from ray.rllib.agents import dqn
-        return dqn.DQNAgent
-    elif alg == "APEX":
-        from ray.rllib.agents import dqn
-        return dqn.ApexAgent
-    elif alg == "A3C":
-        from ray.rllib.agents import a3c
-        return a3c.A3CAgent
-    elif alg == "A2C":
-        from ray.rllib.agents import a3c
-        return a3c.A2CAgent
-    elif alg == "PG":
-        from ray.rllib.agents import pg
-        return pg.PGAgent
-    elif alg == "IMPALA":
-        from ray.rllib.agents import impala
-        return impala.ImpalaAgent
-    elif alg == "QMIX":
-        from ray.rllib.agents import qmix
-        return qmix.QMixAgent
-    elif alg == "APEX_QMIX":
-        from ray.rllib.agents import qmix
-        return qmix.ApexQMixAgent
-    elif alg == "script":
-        from ray.tune import script_runner
-        return script_runner.ScriptRunner
-    elif alg == "__fake":
-        from ray.rllib.agents.mock import _MockAgent
-        return _MockAgent
-    elif alg == "__sigmoid_fake_data":
-        from ray.rllib.agents.mock import _SigmoidFakeData
-        return _SigmoidFakeData
-    elif alg == "__parameter_tuning":
-        from ray.rllib.agents.mock import _ParameterTuningAgent
-        return _ParameterTuningAgent
-    else:
-        raise Exception(("Unknown algorithm {}.").format(alg))

python/ray/rllib/agents/registry.py

@@ -0,0 +1,122 @@
+"""Registry of algorithm names for `rllib train --run=<alg_name>`"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import traceback
+
+from ray.rllib.contrib.registry import CONTRIBUTED_ALGORITHMS
+
+
+def _import_qmix():
+    from ray.rllib.agents import qmix
+    return qmix.QMixAgent
+
+
+def _import_apex_qmix():
+    from ray.rllib.agents import qmix
+    return qmix.ApexQMixAgent
+
+
+def _import_ddpg():
+    from ray.rllib.agents import ddpg
+    return ddpg.DDPGAgent
+
+
+def _import_apex_ddpg():
+    from ray.rllib.agents import ddpg
+    return ddpg.ApexDDPGAgent
+
+
+def _import_ppo():
+    from ray.rllib.agents import ppo
+    return ppo.PPOAgent
+
+
+def _import_es():
+    from ray.rllib.agents import es
+    return es.ESAgent
+
+
+def _import_ars():
+    from ray.rllib.agents import ars
+    return ars.ARSAgent
+
+
+def _import_dqn():
+    from ray.rllib.agents import dqn
+    return dqn.DQNAgent
+
+
+def _import_apex():
+    from ray.rllib.agents import dqn
+    return dqn.ApexAgent
+
+
+def _import_a3c():
+    from ray.rllib.agents import a3c
+    return a3c.A3CAgent
+
+
+def _import_a2c():
+    from ray.rllib.agents import a3c
+    return a3c.A2CAgent
+
+
+def _import_pg():
+    from ray.rllib.agents import pg
+    return pg.PGAgent
+
+
+def _import_impala():
+    from ray.rllib.agents import impala
+    return impala.ImpalaAgent
+
+
+ALGORITHMS = {
+    "DDPG": _import_ddpg,
+    "APEX_DDPG": _import_apex_ddpg,
+    "PPO": _import_ppo,
+    "ES": _import_es,
+    "ARS": _import_ars,
+    "DQN": _import_dqn,
+    "APEX": _import_apex,
+    "A3C": _import_a3c,
+    "A2C": _import_a2c,
+    "PG": _import_pg,
+    "IMPALA": _import_impala,
+    "QMIX": _import_qmix,
+    "APEX_QMIX": _import_apex_qmix,
+}
+
+
+def get_agent_class(alg):
+    """Returns the class of a known agent given its name."""
+
+    try:
+        return _get_agent_class(alg)
+    except ImportError:
+        from ray.rllib.agents.mock import _agent_import_failed
+        return _agent_import_failed(traceback.format_exc())
+
+
+def _get_agent_class(alg):
+    if alg in ALGORITHMS:
+        return ALGORITHMS[alg]()
+    elif alg in CONTRIBUTED_ALGORITHMS:
+        return CONTRIBUTED_ALGORITHMS[alg]()
+    elif alg == "script":
+        from ray.tune import script_runner
+        return script_runner.ScriptRunner
+    elif alg == "__fake":
+        from ray.rllib.agents.mock import _MockAgent
+        return _MockAgent
+    elif alg == "__sigmoid_fake_data":
+        from ray.rllib.agents.mock import _SigmoidFakeData
+        return _SigmoidFakeData
+    elif alg == "__parameter_tuning":
+        from ray.rllib.agents.mock import _ParameterTuningAgent
+        return _ParameterTuningAgent
+    else:
+        raise Exception(("Unknown algorithm {}.").format(alg))