Repository: ray-project/ray

Issue:
[rllib] ray/rllib/rollout.py does not find registered trainables
### What is the problem?
When using rollout.py with custom trainable (registered via `tune.register_trainable`) the script does not find the trainable.

This seems to be caused because rollout.py uses `ray.rllib.agents.registry.get_agent_class` instead of `ray.tune.registry.get_trainable_cls`.

*Ray version and other system information (Python version, TensorFlow version, OS):*
python=3.7.3, ray[rllib,debug]==ray-0.9.0.dev0, tensorflow==1.15.0, Ubuntu 18.04 LTS

### Reproduction
```
#generate checkpoint
rllib train --run DQN --env CartPole-v0 --stop '{"timesteps_total": 5000}' --checkpoint-freq 1
python rollout.py PATH_TO_CHECKPOINT --run OtherDQN --episodes 10 --out rollout.pkl

2020-03-26 16:28:25,858	INFO resource_spec.py:212 -- Starting Ray with 11.62 GiB memory available for workers and up to 5.83 GiB for objects. You can adjust these settings with ray.init(memory=<bytes>, object_store_memory=<bytes>).
2020-03-26 16:28:26,332	INFO services.py:1123 -- View the Ray dashboard at localhost:8265
Traceback (most recent call last):
  File "rollout.py", line 475, in <module>
    run(args, parser)
  File "rollout.py", line 285, in run
    cls = get_agent_class(args.run)
  File "/home/carl/miniconda3/envs/rollout_test/lib/python3.7/site-packages/ray/rllib/agents/registry.py", line 130, in get_agent_class
    return _get_agent_class(alg)
  File "/home/carl/miniconda3/envs/rollout_test/lib/python3.7/site-packages/ray/rllib/agents/registry.py", line 154, in _get_agent_class
    raise Exception(("Unknown algorithm {}.").format(alg))
Exception: Unknown algorithm OtherDQN.
```
The provided rollout.py adds the following lines at the start:
```
from ray.rllib.agents.dqn import DQNTrainer
OtherDQNTrainer = DQNTrainer.with_updates(
    name="OtherDQN")
ray.tune.register_trainable("OtherDQN", OtherDQNTrainer)
```
If `ray.rllib.agents.registry.get_agent_class` is replaces with `ray.tune.registry.get_trainable_cls` it works (with both "OtherDQN" and "DQN")

[rollout.zip](https://github.com/ray-project/ray/files/4388161/rollout.zip)

- [x] I have verified my script runs in a clean environment and reproduces the issue.
- [x] I have verified the issue also occurs with the [latest wheels](https://ray.readthedocs.io/en/latest/installation.html).



Discussion:

Resolving patch:
diff --git a/rllib/rollout.py b/rllib/rollout.py
--- a/rllib/rollout.py
+++ b/rllib/rollout.py
@@ -11,7 +11,6 @@
 
 import gym
 import ray
-from ray.rllib.agents.registry import get_agent_class
 from ray.rllib.env import MultiAgentEnv
 from ray.rllib.env.base_env import _DUMMY_AGENT_ID
 from ray.rllib.evaluation.episode import _flatten_action
@@ -19,6 +18,7 @@
 from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID
 from ray.rllib.utils.deprecation import deprecation_warning
 from ray.tune.utils import merge_dicts
+from ray.tune.registry import get_trainable_cls
 
 EXAMPLE_USAGE = """
 Example Usage via RLlib CLI:
@@ -274,7 +274,7 @@ def run(args, parser):
     ray.init()
 
     # Create the Trainer from config.
-    cls = get_agent_class(args.run)
+    cls = get_trainable_cls(args.run)
     agent = cls(env=args.env, config=config)
     # Load state from checkpoint.
     agent.restore(args.checkpoint)


Test patch: