-
Notifications
You must be signed in to change notification settings - Fork 5.6k
/
__init__.py
55 lines (43 loc) · 1.6 KB
/
__init__.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
import logging
from ray._private.usage import usage_lib
# Note: do not introduce unnecessary library dependencies here, e.g. gym.
# This file is imported from the tune module in order to register RLlib agents.
from ray.rllib.env.base_env import BaseEnv
from ray.rllib.env.external_env import ExternalEnv
from ray.rllib.env.multi_agent_env import MultiAgentEnv
from ray.rllib.env.vector_env import VectorEnv
from ray.rllib.evaluation.rollout_worker import RolloutWorker
from ray.rllib.policy.policy import Policy
from ray.rllib.policy.sample_batch import SampleBatch
from ray.rllib.policy.tf_policy import TFPolicy
from ray.rllib.policy.torch_policy import TorchPolicy
from ray.tune.registry import register_trainable
def _setup_logger():
logger = logging.getLogger("ray.rllib")
handler = logging.StreamHandler()
handler.setFormatter(
logging.Formatter(
"%(asctime)s\t%(levelname)s %(filename)s:%(lineno)s -- %(message)s"
)
)
logger.addHandler(handler)
logger.propagate = False
def _register_all():
from ray.rllib.algorithms.registry import ALGORITHMS, _get_algorithm_class
for key, get_trainable_class_and_config in ALGORITHMS.items():
register_trainable(key, get_trainable_class_and_config()[0])
for key in ["__fake", "__sigmoid_fake_data", "__parameter_tuning"]:
register_trainable(key, _get_algorithm_class(key))
_setup_logger()
usage_lib.record_library_usage("rllib")
__all__ = [
"Policy",
"TFPolicy",
"TorchPolicy",
"RolloutWorker",
"SampleBatch",
"BaseEnv",
"MultiAgentEnv",
"VectorEnv",
"ExternalEnv",
]