Repository: ray-project/ray

Issue:
[rllib] Evaluation doesn’t work properly for PyTorch.
This is not a contribution.

Ray version: 0.8.1
Python version: 3.6.8
Pytorch version: 1.4
OS: Ubuntu 18.04 Docker

Training with Pytorch with evaluation enabled doesn’t work properly as it ends up calling tensorflow code. Here's a script that reproduces the problem:

```python
import gym
from gym.spaces import Box
from ray import tune

class DummyEnv(gym.Env):
    def __init__(self, config):
        self.action_space = Box(0.0, 1.0, shape=(1,))
        self.observation_space = Box(0.0, 1.0, shape=(1, ))

    def reset(self):
        self.num_steps = 0
        return [0.0]

    def step(self, action):
        self.num_steps += 1
        return [0.0], 1.0, True if self.num_steps == 30 else False, {}

tune.run(
    "PPO",
    config={
        "env": DummyEnv,
        "use_pytorch": True,
        "num_workers": 1,
        "evaluation_interval": 1,
        "evaluation_config": {"seed": 1},
    }
)
```
This will result in the following exception: 
```
ray.exceptions.RayTaskError(AssertionError): ray::PPO.train() (pid=33990)
  File "python/ray/_raylet.pyx", line 452, in ray._raylet.execute_task
  File "python/ray/_raylet.pyx", line 430, in ray._raylet.execute_task.function_executor
  File "/usr/local/lib/python3.6/dist-packages/ray/rllib/agents/trainer.py", line 513, in train
    evaluation_metrics = self._evaluate()
  File "/usr/local/lib/python3.6/dist-packages/ray/rllib/agents/trainer.py", line 685, in _evaluate
    lambda w: w.restore(ray.get(weights)))
  File "/usr/local/lib/python3.6/dist-packages/ray/rllib/evaluation/worker_set.py", line 110, in foreach_worker
    local_result = [func(self.local_worker())]
  File "/usr/local/lib/python3.6/dist-packages/ray/rllib/agents/trainer.py", line 685, in <lambda>
    lambda w: w.restore(ray.get(weights)))
  File "/usr/local/lib/python3.6/dist-packages/ray/rllib/evaluation/rollout_worker.py", line 764, in restore
    self.policy_map[pid].set_state(state)
  File "/usr/local/lib/python3.6/dist-packages/ray/rllib/policy/policy.py", line 320, in set_state
    self.set_weights(state)
  File "/usr/local/lib/python3.6/dist-packages/ray/rllib/policy/tf_policy.py", line 290, in set_weights
    return self._variables.set_weights(weights)
  File "/usr/local/lib/python3.6/dist-packages/ray/experimental/tf_utils.py", line 185, in set_weights
    assert assign_list, ("No variables in the input matched those in the "
AssertionError: No variables in the input matched those in the network. Possible cause: Two networks were defined in the same TensorFlow graph. To fix this, place each network definition in its own tf.Graph.
```


Discussion:
This is a bug (seems to have been there for a while). Fixing this now ...

Resolving patch:
diff --git a/rllib/agents/trainer_template.py b/rllib/agents/trainer_template.py
--- a/rllib/agents/trainer_template.py
+++ b/rllib/agents/trainer_template.py
@@ -38,40 +38,43 @@ def build_trainer(name,
     Arguments:
         name (str): name of the trainer (e.g., "PPO")
         default_policy (cls): the default Policy class to use
-        default_config (dict): The default config dict of the algorithm,
-            otherwise uses the Trainer default config.
-        validate_config (func): optional callback that checks a given config
-            for correctness. It may mutate the config as needed.
-        get_initial_state (func): optional function that returns the initial
-            state dict given the trainer instance as an argument. The state
-            dict must be serializable so that it can be checkpointed, and will
-            be available as the `trainer.state` variable.
-        get_policy_class (func): optional callback that takes a config and
-            returns the policy class to override the default with
-        before_init (func): optional function to run at the start of trainer
-            init that takes the trainer instance as argument
-        make_workers (func): override the method that creates rollout workers.
-            This takes in (trainer, env_creator, policy, config) as args.
-        make_policy_optimizer (func): optional function that returns a
-            PolicyOptimizer instance given (WorkerSet, config)
-        after_init (func): optional function to run at the end of trainer init
-            that takes the trainer instance as argument
-        before_train_step (func): optional callback to run before each train()
-            call. It takes the trainer instance as an argument.
-        after_optimizer_step (func): optional callback to run after each
-            step() call to the policy optimizer. It takes the trainer instance
-            and the policy gradient fetches as arguments.
-        after_train_result (func): optional callback to run at the end of each
-            train() call. It takes the trainer instance and result dict as
-            arguments, and may mutate the result dict as needed.
-        collect_metrics_fn (func): override the method used to collect metrics.
-            It takes the trainer instance as argumnt.
-        before_evaluate_fn (func): callback to run before evaluation. This
-            takes the trainer instance as argument.
-        mixins (list): list of any class mixins for the returned trainer class.
-            These mixins will be applied in order and will have higher
-            precedence than the Trainer class
-        training_pipeline (func): Experimental support for custom
+        default_config (Optional[dict]): The default config dict of the
+            algorithm. If None, uses the Trainer default config.
+        validate_config (Optional[callable]): Optional callback that checks a
+            given config for correctness. It may mutate the config as needed.
+        get_initial_state (Optional[callable]): Optional callable that returns
+            the initial state dict given the trainer instance as an argument.
+            The state dict must be serializable so that it can be checkpointed,
+            and will be available as the `trainer.state` variable.
+        get_policy_class (Optional[callable]): Optional callable that takes a
+            Trainer config and returns the policy class to override the default
+            with.
+        before_init (Optional[callable]): Optional callable to run at the start
+            of trainer init that takes the trainer instance as argument.
+        make_workers (Optional[callable]): Override the default method that
+            creates rollout workers. This takes in (trainer, env_creator,
+            policy, config) as args.
+        make_policy_optimizer (Optional[callable]): Optional callable that
+            returns a PolicyOptimizer instance given (WorkerSet, config).
+        after_init (Optional[callable]): Optional callable to run at the end of
+            trainer init that takes the trainer instance as argument.
+        before_train_step (Optional[callable]): Optional callable to run before
+            each train() call. It takes the trainer instance as an argument.
+        after_optimizer_step (Optional[callable]): Optional callable to run
+            after each step() call to the policy optimizer. It takes the
+            trainer instance and the policy gradient fetches as arguments.
+        after_train_result (Optional[callable]): Optional callable to run at
+            the end of each train() call. It takes the trainer instance and
+            result dict as arguments, and may mutate the result dict as needed.
+        collect_metrics_fn (Optional[callable]): Optional callable to override
+            the default method used to collect metrics. Takes the trainer
+            instance as argumnt.
+        before_evaluate_fn (Optional[callable]): Optional callable to run
+            before evaluation. Takes the trainer instance as argument.
+        mixins (Optional[List[class]]): Optional list of mixin class(es) for
+            the returned trainer class. These mixins will be applied in order
+            and will have higher precedence than the Trainer class.
+        training_pipeline (Optional[callable]): Experimental support for custom
             training pipelines. This overrides `make_policy_optimizer`.
 
     Returns:
@@ -92,20 +95,26 @@ def __init__(self, config=None, env=None, logger_creator=None):
         def _init(self, config, env_creator):
             if validate_config:
                 validate_config(config)
+
             if get_initial_state:
                 self.state = get_initial_state(self)
             else:
                 self.state = {}
-            if get_policy_class is None:
-                policy = default_policy
-            else:
-                policy = get_policy_class(config)
+
+            # Override default policy if `get_policy_class` is provided.
+            if get_policy_class is not None:
+                self._policy = get_policy_class(config)
+
             if before_init:
                 before_init(self)
+
+            # Creating all workers (excluding evaluation workers).
             if make_workers:
-                self.workers = make_workers(self, env_creator, policy, config)
+                self.workers = make_workers(self, env_creator, self._policy,
+                                            config)
             else:
-                self.workers = self._make_workers(env_creator, policy, config,
+                self.workers = self._make_workers(env_creator, self._policy,
+                                                  config,
                                                   self.config["num_workers"])
             self.train_pipeline = None
             self.optimizer = None


Test patch: