Repository: ray-project/ray

Issue:
[RLlib] Cannot rollout ES after training.
Using ray 0.8.1. 

Having trained an agent with ES:
```
rllib train --run=ES --env=LunarLander-v2 --checkpoint-freq=10
```

I can't play it. (It works with PPO).

Executing this:
```
rllib rollout checkpoint_50/checkpoint-50 --run ES --env LunarLander-v2 --steps 600
```
results in:
```
2020-02-12 05:41:30,922 INFO trainable.py:423 -- Current state after restoring: {'_iteration': 50, '_timesteps_total': 16979666, '_time_total': 2578.4048051834106, '_episodes_total': None}
Traceback (most recent call last):
  File "/home/andriy/miniconda3/envs/myproj/bin/rllib", line 8, in <module>
    sys.exit(cli())
  File "/home/andriy/miniconda3/envs/myproj/lib/python3.7/site-packages/ray/rllib/scripts.py", line 36, in cli
    rollout.run(options, rollout_parser)
  File "/home/andriy/miniconda3/envs/myproj/lib/python3.7/site-packages/ray/rllib/rollout.py", line 265, in run
    args.no_render, args.monitor)
  File "/home/andriy/miniconda3/envs/myproj/lib/python3.7/site-packages/ray/rllib/rollout.py", line 364, in rollout
    prev_action=prev_actions[agent_id],
  File "/home/andriy/miniconda3/envs/myproj/lib/python3.7/site-packages/ray/rllib/rollout.py", line 272, in __missing__
    self[key] = value = self.default_factory(key)
  File "/home/andriy/miniconda3/envs/myproj/lib/python3.7/site-packages/ray/rllib/rollout.py", line 340, in <lambda>
    lambda agent_id: action_init[mapping_cache[agent_id]])
NameError: free variable 'action_init' referenced before assignment in enclosing scope
```



Discussion:
I'll take a look.

Resolving patch:
diff --git a/rllib/agents/ars/ars.py b/rllib/agents/ars/ars.py
--- a/rllib/agents/ars/ars.py
+++ b/rllib/agents/ars/ars.py
@@ -165,8 +165,7 @@ def _init(self, config, env_creator):
         # PyTorch check.
         if config["use_pytorch"]:
             raise ValueError(
-                "ARS does not support PyTorch yet! Use tf instead."
-            )
+                "ARS does not support PyTorch yet! Use tf instead.")
 
         env = env_creator(config["env_config"])
         from ray.rllib import models
@@ -301,7 +300,7 @@ def _stop(self):
             w.__ray_terminate__.remote()
 
     @override(Trainer)
-    def compute_action(self, observation):
+    def compute_action(self, observation, *args, **kwargs):
         return self.policy.compute(observation, update=True)[0]
 
     def _collect_results(self, theta_id, min_episodes):
diff --git a/rllib/agents/es/es.py b/rllib/agents/es/es.py
--- a/rllib/agents/es/es.py
+++ b/rllib/agents/es/es.py
@@ -171,8 +171,7 @@ def _init(self, config, env_creator):
         # PyTorch check.
         if config["use_pytorch"]:
             raise ValueError(
-                "ES does not support PyTorch yet! Use tf instead."
-            )
+                "ES does not support PyTorch yet! Use tf instead.")
 
         policy_params = {"action_noise_std": 0.01}
 
@@ -292,7 +291,7 @@ def _train(self):
         return result
 
     @override(Trainer)
-    def compute_action(self, observation):
+    def compute_action(self, observation, *args, **kwargs):
         return self.policy.compute(observation, update=False)[0]
 
     @override(Trainer)
diff --git a/rllib/rollout.py b/rllib/rollout.py
--- a/rllib/rollout.py
+++ b/rllib/rollout.py
@@ -15,6 +15,7 @@
 from ray.rllib.env import MultiAgentEnv
 from ray.rllib.env.base_env import _DUMMY_AGENT_ID
 from ray.rllib.evaluation.episode import _flatten_action
+from ray.rllib.evaluation.worker_set import WorkerSet
 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
@@ -339,7 +340,7 @@ def rollout(agent,
     if saver is None:
         saver = RolloutSaver()
 
-    if hasattr(agent, "workers"):
+    if hasattr(agent, "workers") and isinstance(agent.workers, WorkerSet):
         env = agent.workers.local_worker().env
         multiagent = isinstance(env, MultiAgentEnv)
         if agent.workers.local_worker().multiagent:
@@ -349,15 +350,22 @@ def rollout(agent,
         policy_map = agent.workers.local_worker().policy_map
         state_init = {p: m.get_initial_state() for p, m in policy_map.items()}
         use_lstm = {p: len(s) > 0 for p, s in state_init.items()}
-        action_init = {
-            p: _flatten_action(m.action_space.sample())
-            for p, m in policy_map.items()
-        }
     else:
         env = gym.make(env_name)
         multiagent = False
+        try:
+            policy_map = {DEFAULT_POLICY_ID: agent.policy}
+        except AttributeError:
+            raise AttributeError(
+                "Agent ({}) does not have a `policy` property! This is needed "
+                "for performing (trained) agent rollouts.".format(agent))
         use_lstm = {DEFAULT_POLICY_ID: False}
 
+    action_init = {
+        p: _flatten_action(m.action_space.sample())
+        for p, m in policy_map.items()
+    }
+
     # If monitoring has been requested, manually wrap our environment with a
     # gym monitor, which is set to record every episode.
     if video_dir:


Test patch: