We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Time to initialize a policy grows linearly with the number of agents
<!--
General questions should be asked on the mailing list ray-dev@googlegroups.com.
Questions about how to use Ray should be asked on
[StackOverflow](https://stackoverflow.com/questions/tagged/ray).

Before submitting an issue, please fill out the following form.
-->

### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux Ubuntu 18.04
- **Ray installed from (source or binary)**: Binary
- **Ray version**: 0.7.4
- **Python version**: 3.7.4
- **Exact command to reproduce**: N/A

<!--
You can obtain the Ray version with

python -c "import ray; print(ray.__version__)"
-->

### Describe the problem
<!-- Describe the problem clearly here. -->
I noticed that in multi agent settings, the time to initialize a policy per agent increases as more agents are initialized. In the sample output I provided below, you can see that the time to initialize a single DynamicTFPolicy grows from 4.6 seconds to 15.3 seconds from the first agent to the tenth agent created. Line 291 of `rllib/policy/dynamic_tf_policy.py` is
```python
self._sess.run(tf.global_variables_initializer())
```
which I believe will run one time for each agent initialized. If I'm not mistaken, this means that every variable in the computation graph is being initialized each time that we initialize a DynamicTFPolicy. If initializing a DynamicTFPolicy adds new variables to the computation graph (as I believe it does), this would explain why the time to initialize a DynamicTFPolicy grows over time: We are initializing every variable in the computation graph, and the computation graph is growing. My question is, why does line 291 run a global variables initializer? Is there a reason for this that I can't see inside this method? How hard would it be to modify this to only initialize variables in the individual policy that we care to initialize?

I'm asking this because as detailed in #5753, I'm trying to modify rllib to allow initialization and removal of policies during training. The overhead incurred by this initialization quickly slows the training script down enough to be useless. Also, if anyone knows what the resource bottleneck is for policy initialization, that would be very helpful to know for when we're picking new hardware. Does it need a ton of cores to run in parallel, or more memory, or a bigger GPU or more GPUs or something? Thanks.


### Source code / logs
<!-- Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. -->

```
(pytorch) root@e3a955e42cae:~/bees/bees# python trainer.py settings/settings.json
2019-10-23 12:38:04,168 WARNING worker.py:1426 -- WARNING: Not updating worker name since `setproctitle` is not installed. Install this with `pip install setproctitle` (or ray[debug]) to enable monitoring of worker processes.
2019-10-23 12:38:04,169 INFO resource_spec.py:205 -- Starting Ray with 3.52 GiB memory available for workers and up to 1.78 GiB for objects. You can adjust these settings with ray.init(memory=<bytes>, object_store_memory=<bytes>).
2019-10-23 12:38:05,354 INFO trainer.py:344 -- Tip: set 'eager': true or the --eager flag to enable TensorFlow eager execution
2019-10-23 12:38:05,754 WARNING ppo.py:149 -- Using the simple minibatch optimizer. This will significantly reduce performance, consider simple_optimizer=False.
DTFP: 4.604122s
DTFP: 4.856234s
DTFP: 5.630484s
DTFP: 6.850456s
DTFP: 7.856700s
DTFP: 9.624164s
DTFP: 10.894944s
DTFP: 12.129192s
DTFP: 14.210247s
DTFP: 15.342738s
```

Line 130 in `tf_policy_template.py`  (modified to print debug output above)
```python
            t = time.time()

            DynamicTFPolicy.__init__(
                self,
                obs_space,
                action_space,
                config,
                loss_fn,
                stats_fn=stats_fn,
                grad_stats_fn=grad_stats_fn,
                before_loss_init=before_loss_init_wrapper,
                make_model=make_model,
                action_sampler_fn=action_sampler_fn,
                existing_model=existing_model,
                existing_inputs=existing_inputs,
                get_batch_divisibility_req=get_batch_divisibility_req,
                obs_include_prev_action_reward=obs_include_prev_action_reward)

            print("DTFP: %fs" % (time.time() - t))
```

Snippet of trainer script used.
```python
# pylint: disable=invalid-name
if __name__ == "__main__":
    ray.init()

    # Get ``settings`` file for now.
    settings_file = sys.argv[1]
    with open(settings_file, "r") as f:
        settings = json.load(f)

    env_config = settings["env"]
    time_steps = env_config["time_steps"]

    space_env = create_env(settings)
    env = create_env(settings)

    # Register environment
    register_env("world", lambda _: env)

    # Build environment instance to get ``obs_space``.
    obs_space = space_env.observation_space
    act_space = space_env.action_space

    # You can also have multiple policies per trainer, but here we just
    # show one each for PPO and DQN.
    policies: Dict[str, Tuple[Any, gym.Space, gym.Space, Dict[Any, Any]]] = {
        "0": (PPOTFPolicy, obs_space, act_space, {}),
        "1": (PPOTFPolicy, obs_space, act_space, {}),
        "2": (PPOTFPolicy, obs_space, act_space, {}),
        "3": (PPOTFPolicy, obs_space, act_space, {}),
        "4": (PPOTFPolicy, obs_space, act_space, {}),
        "5": (PPOTFPolicy, obs_space, act_space, {}),
        "6": (PPOTFPolicy, obs_space, act_space, {}),
        "7": (PPOTFPolicy, obs_space, act_space, {}),
        "8": (PPOTFPolicy, obs_space, act_space, {}),
        "9": (PPOTFPolicy, obs_space, act_space, {}),
    }

    def policy_mapping_fn(agent_id: int) -> str:
        """ Returns the given agent's policy identifier. """
        return str(agent_id)

    ppo_trainer = PPOTrainer(
        env="bee_world",
        config={
            "multiagent": {
                "policies": policies,
                "policy_mapping_fn": policy_mapping_fn,
                "policies_to_train": ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9"],
            },
            "simple_optimizer": True,
            # Disable filters, otherwise we would need to synchronize those
            # as well to the DQN agent.
            "observation_filter": "NoFilter",
            "num_workers": 2,
            "num_gpus": 1,
            "train_batch_size": 2,
            "sample_batch_size": 1,
            "sgd_minibatch_size": 2,
        },
    )
```




--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `python/ray/experimental/tf_utils.py`
Content:
```
1 from collections import deque, OrderedDict
2 import numpy as np
3 
4 from ray.rllib.utils import try_import_tf
5 
6 tf = try_import_tf()
7 
8 
9 def unflatten(vector, shapes):
10     i = 0
11     arrays = []
12     for shape in shapes:
13         size = np.prod(shape, dtype=np.int)
14         array = vector[i:(i + size)].reshape(shape)
15         arrays.append(array)
16         i += size
17     assert len(vector) == i, "Passed weight does not have the correct shape."
18     return arrays
19 
20 
21 class TensorFlowVariables:
22     """A class used to set and get weights for Tensorflow networks.
23 
24     Attributes:
25         sess (tf.Session): The tensorflow session used to run assignment.
26         variables (Dict[str, tf.Variable]): Extracted variables from the loss
27             or additional variables that are passed in.
28         placeholders (Dict[str, tf.placeholders]): Placeholders for weights.
29         assignment_nodes (Dict[str, tf.Tensor]): Nodes that assign weights.
30     """
31 
32     def __init__(self, output, sess=None, input_variables=None):
33         """Creates TensorFlowVariables containing extracted variables.
34 
35         The variables are extracted by performing a BFS search on the
36         dependency graph with loss as the root node. After the tree is
37         traversed and those variables are collected, we append input_variables
38         to the collected variables. For each variable in the list, the
39         variable has a placeholder and assignment operation created for it.
40 
41         Args:
42             output (tf.Operation, List[tf.Operation]): The tensorflow
43                 operation to extract all variables from.
44             sess (tf.Session): Session used for running the get and set
45                 methods.
46             input_variables (List[tf.Variables]): Variables to include in the
47                 list.
48         """
49         self.sess = sess
50         if not isinstance(output, (list, tuple)):
51             output = [output]
52         queue = deque(output)
53         variable_names = []
54         explored_inputs = set(output)
55 
56         # We do a BFS on the dependency graph of the input function to find
57         # the variables.
58         while len(queue) != 0:
59             tf_obj = queue.popleft()
60             if tf_obj is None:
61                 continue
62             # The object put into the queue is not necessarily an operation,
63             # so we want the op attribute to get the operation underlying the
64             # object. Only operations contain the inputs that we can explore.
65             if hasattr(tf_obj, "op"):
66                 tf_obj = tf_obj.op
67             for input_op in tf_obj.inputs:
68                 if input_op not in explored_inputs:
69                     queue.append(input_op)
70                     explored_inputs.add(input_op)
71             # Tensorflow control inputs can be circular, so we keep track of
72             # explored operations.
73             for control in tf_obj.control_inputs:
74                 if control not in explored_inputs:
75                     queue.append(control)
76                     explored_inputs.add(control)
77             if ("Variable" in tf_obj.node_def.op
78                     or "VarHandle" in tf_obj.node_def.op):
79                 variable_names.append(tf_obj.node_def.name)
80         self.variables = OrderedDict()
81         variable_list = [
82             v for v in tf.global_variables()
83             if v.op.node_def.name in variable_names
84         ]
85         if input_variables is not None:
86             variable_list += input_variables
87         for v in variable_list:
88             self.variables[v.op.node_def.name] = v
89 
90         self.placeholders = {}
91         self.assignment_nodes = {}
92 
93         # Create new placeholders to put in custom weights.
94         for k, var in self.variables.items():
95             self.placeholders[k] = tf.placeholder(
96                 var.value().dtype,
97                 var.get_shape().as_list(),
98                 name="Placeholder_" + k)
99             self.assignment_nodes[k] = var.assign(self.placeholders[k])
100 
101     def set_session(self, sess):
102         """Sets the current session used by the class.
103 
104         Args:
105             sess (tf.Session): Session to set the attribute with.
106         """
107         self.sess = sess
108 
109     def get_flat_size(self):
110         """Returns the total length of all of the flattened variables.
111 
112         Returns:
113             The length of all flattened variables concatenated.
114         """
115         return sum(
116             np.prod(v.get_shape().as_list()) for v in self.variables.values())
117 
118     def _check_sess(self):
119         """Checks if the session is set, and if not throw an error message."""
120         assert self.sess is not None, ("The session is not set. Set the "
121                                        "session either by passing it into the "
122                                        "TensorFlowVariables constructor or by "
123                                        "calling set_session(sess).")
124 
125     def get_flat(self):
126         """Gets the weights and returns them as a flat array.
127 
128         Returns:
129             1D Array containing the flattened weights.
130         """
131         self._check_sess()
132         return np.concatenate([
133             v.eval(session=self.sess).flatten()
134             for v in self.variables.values()
135         ])
136 
137     def set_flat(self, new_weights):
138         """Sets the weights to new_weights, converting from a flat array.
139 
140         Note:
141             You can only set all weights in the network using this function,
142             i.e., the length of the array must match get_flat_size.
143 
144         Args:
145             new_weights (np.ndarray): Flat array containing weights.
146         """
147         self._check_sess()
148         shapes = [v.get_shape().as_list() for v in self.variables.values()]
149         arrays = unflatten(new_weights, shapes)
150         placeholders = [
151             self.placeholders[k] for k, v in self.variables.items()
152         ]
153         self.sess.run(
154             list(self.assignment_nodes.values()),
155             feed_dict=dict(zip(placeholders, arrays)))
156 
157     def get_weights(self):
158         """Returns a dictionary containing the weights of the network.
159 
160         Returns:
161             Dictionary mapping variable names to their weights.
162         """
163         self._check_sess()
164         return {
165             k: v.eval(session=self.sess)
166             for k, v in self.variables.items()
167         }
168 
169     def set_weights(self, new_weights):
170         """Sets the weights to new_weights.
171 
172         Note:
173             Can set subsets of variables as well, by only passing in the
174             variables you want to be set.
175 
176         Args:
177             new_weights (Dict): Dictionary mapping variable names to their
178                 weights.
179         """
180         self._check_sess()
181         assign_list = [
182             self.assignment_nodes[name] for name in new_weights.keys()
183             if name in self.assignment_nodes
184         ]
185         assert assign_list, ("No variables in the input matched those in the "
186                              "network. Possible cause: Two networks were "
187                              "defined in the same TensorFlow graph. To fix "
188                              "this, place each network definition in its own "
189                              "tf.Graph.")
190         self.sess.run(
191             assign_list,
192             feed_dict={
193                 self.placeholders[name]: value
194                 for (name, value) in new_weights.items()
195                 if name in self.placeholders
196             })
197
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/python/ray/experimental/tf_utils.py b/python/ray/experimental/tf_utils.py
--- a/python/ray/experimental/tf_utils.py
+++ b/python/ray/experimental/tf_utils.py
@@ -161,10 +161,7 @@
             Dictionary mapping variable names to their weights.
         """
         self._check_sess()
-        return {
-            k: v.eval(session=self.sess)
-            for k, v in self.variables.items()
-        }
+        return self.sess.run(self.variables)
 
     def set_weights(self, new_weights):
         """Sets the weights to new_weights.

{"golden_diff": "diff --git a/python/ray/experimental/tf_utils.py b/python/ray/experimental/tf_utils.py\n--- a/python/ray/experimental/tf_utils.py\n+++ b/python/ray/experimental/tf_utils.py\n@@ -161,10 +161,7 @@\n             Dictionary mapping variable names to their weights.\n         \"\"\"\n         self._check_sess()\n-        return {\n-            k: v.eval(session=self.sess)\n-            for k, v in self.variables.items()\n-        }\n+        return self.sess.run(self.variables)\n \n     def set_weights(self, new_weights):\n         \"\"\"Sets the weights to new_weights.\n", "issue": "Time to initialize a policy grows linearly with the number of agents\n<!--\r\nGeneral questions should be asked on the mailing list ray-dev@googlegroups.com.\r\nQuestions about how to use Ray should be asked on\r\n[StackOverflow](https://stackoverflow.com/questions/tagged/ray).\r\n\r\nBefore submitting an issue, please fill out the following form.\r\n-->\r\n\r\n### System information\r\n- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux Ubuntu 18.04\r\n- **Ray installed from (source or binary)**: Binary\r\n- **Ray version**: 0.7.4\r\n- **Python version**: 3.7.4\r\n- **Exact command to reproduce**: N/A\r\n\r\n<!--\r\nYou can obtain the Ray version with\r\n\r\npython -c \"import ray; print(ray.__version__)\"\r\n-->\r\n\r\n### Describe the problem\r\n<!-- Describe the problem clearly here. -->\r\nI noticed that in multi agent settings, the time to initialize a policy per agent increases as more agents are initialized. In the sample output I provided below, you can see that the time to initialize a single DynamicTFPolicy grows from 4.6 seconds to 15.3 seconds from the first agent to the tenth agent created. Line 291 of `rllib/policy/dynamic_tf_policy.py` is\r\n```python\r\nself._sess.run(tf.global_variables_initializer())\r\n```\r\nwhich I believe will run one time for each agent initialized. If I'm not mistaken, this means that every variable in the computation graph is being initialized each time that we initialize a DynamicTFPolicy. If initializing a DynamicTFPolicy adds new variables to the computation graph (as I believe it does), this would explain why the time to initialize a DynamicTFPolicy grows over time: We are initializing every variable in the computation graph, and the computation graph is growing. My question is, why does line 291 run a global variables initializer? Is there a reason for this that I can't see inside this method? How hard would it be to modify this to only initialize variables in the individual policy that we care to initialize?\r\n\r\nI'm asking this because as detailed in #5753, I'm trying to modify rllib to allow initialization and removal of policies during training. The overhead incurred by this initialization quickly slows the training script down enough to be useless. Also, if anyone knows what the resource bottleneck is for policy initialization, that would be very helpful to know for when we're picking new hardware. Does it need a ton of cores to run in parallel, or more memory, or a bigger GPU or more GPUs or something? Thanks.\r\n\r\n\r\n### Source code / logs\r\n<!-- Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. -->\r\n\r\n```\r\n(pytorch) root@e3a955e42cae:~/bees/bees# python trainer.py settings/settings.json\r\n2019-10-23 12:38:04,168 WARNING worker.py:1426 -- WARNING: Not updating worker name since `setproctitle` is not installed. Install this with `pip install setproctitle` (or ray[debug]) to enable monitoring of worker processes.\r\n2019-10-23 12:38:04,169 INFO resource_spec.py:205 -- Starting Ray with 3.52 GiB memory available for workers and up to 1.78 GiB for objects. You can adjust these settings with ray.init(memory=<bytes>, object_store_memory=<bytes>).\r\n2019-10-23 12:38:05,354 INFO trainer.py:344 -- Tip: set 'eager': true or the --eager flag to enable TensorFlow eager execution\r\n2019-10-23 12:38:05,754 WARNING ppo.py:149 -- Using the simple minibatch optimizer. This will significantly reduce performance, consider simple_optimizer=False.\r\nDTFP: 4.604122s\r\nDTFP: 4.856234s\r\nDTFP: 5.630484s\r\nDTFP: 6.850456s\r\nDTFP: 7.856700s\r\nDTFP: 9.624164s\r\nDTFP: 10.894944s\r\nDTFP: 12.129192s\r\nDTFP: 14.210247s\r\nDTFP: 15.342738s\r\n```\r\n\r\nLine 130 in `tf_policy_template.py`  (modified to print debug output above)\r\n```python\r\n            t = time.time()\r\n\r\n            DynamicTFPolicy.__init__(\r\n                self,\r\n                obs_space,\r\n                action_space,\r\n                config,\r\n                loss_fn,\r\n                stats_fn=stats_fn,\r\n                grad_stats_fn=grad_stats_fn,\r\n                before_loss_init=before_loss_init_wrapper,\r\n                make_model=make_model,\r\n                action_sampler_fn=action_sampler_fn,\r\n                existing_model=existing_model,\r\n                existing_inputs=existing_inputs,\r\n                get_batch_divisibility_req=get_batch_divisibility_req,\r\n                obs_include_prev_action_reward=obs_include_prev_action_reward)\r\n\r\n            print(\"DTFP: %fs\" % (time.time() - t))\r\n```\r\n\r\nSnippet of trainer script used.\r\n```python\r\n# pylint: disable=invalid-name\r\nif __name__ == \"__main__\":\r\n    ray.init()\r\n\r\n    # Get ``settings`` file for now.\r\n    settings_file = sys.argv[1]\r\n    with open(settings_file, \"r\") as f:\r\n        settings = json.load(f)\r\n\r\n    env_config = settings[\"env\"]\r\n    time_steps = env_config[\"time_steps\"]\r\n\r\n    space_env = create_env(settings)\r\n    env = create_env(settings)\r\n\r\n    # Register environment\r\n    register_env(\"world\", lambda _: env)\r\n\r\n    # Build environment instance to get ``obs_space``.\r\n    obs_space = space_env.observation_space\r\n    act_space = space_env.action_space\r\n\r\n    # You can also have multiple policies per trainer, but here we just\r\n    # show one each for PPO and DQN.\r\n    policies: Dict[str, Tuple[Any, gym.Space, gym.Space, Dict[Any, Any]]] = {\r\n        \"0\": (PPOTFPolicy, obs_space, act_space, {}),\r\n        \"1\": (PPOTFPolicy, obs_space, act_space, {}),\r\n        \"2\": (PPOTFPolicy, obs_space, act_space, {}),\r\n        \"3\": (PPOTFPolicy, obs_space, act_space, {}),\r\n        \"4\": (PPOTFPolicy, obs_space, act_space, {}),\r\n        \"5\": (PPOTFPolicy, obs_space, act_space, {}),\r\n        \"6\": (PPOTFPolicy, obs_space, act_space, {}),\r\n        \"7\": (PPOTFPolicy, obs_space, act_space, {}),\r\n        \"8\": (PPOTFPolicy, obs_space, act_space, {}),\r\n        \"9\": (PPOTFPolicy, obs_space, act_space, {}),\r\n    }\r\n\r\n    def policy_mapping_fn(agent_id: int) -> str:\r\n        \"\"\" Returns the given agent's policy identifier. \"\"\"\r\n        return str(agent_id)\r\n\r\n    ppo_trainer = PPOTrainer(\r\n        env=\"bee_world\",\r\n        config={\r\n            \"multiagent\": {\r\n                \"policies\": policies,\r\n                \"policy_mapping_fn\": policy_mapping_fn,\r\n                \"policies_to_train\": [\"0\", \"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\"],\r\n            },\r\n            \"simple_optimizer\": True,\r\n            # Disable filters, otherwise we would need to synchronize those\r\n            # as well to the DQN agent.\r\n            \"observation_filter\": \"NoFilter\",\r\n            \"num_workers\": 2,\r\n            \"num_gpus\": 1,\r\n            \"train_batch_size\": 2,\r\n            \"sample_batch_size\": 1,\r\n            \"sgd_minibatch_size\": 2,\r\n        },\r\n    )\r\n```\r\n\r\n\r\n\n", "before_files": [{"content": "from collections import deque, OrderedDict\nimport numpy as np\n\nfrom ray.rllib.utils import try_import_tf\n\ntf = try_import_tf()\n\n\ndef unflatten(vector, shapes):\n    i = 0\n    arrays = []\n    for shape in shapes:\n        size = np.prod(shape, dtype=np.int)\n        array = vector[i:(i + size)].reshape(shape)\n        arrays.append(array)\n        i += size\n    assert len(vector) == i, \"Passed weight does not have the correct shape.\"\n    return arrays\n\n\nclass TensorFlowVariables:\n    \"\"\"A class used to set and get weights for Tensorflow networks.\n\n    Attributes:\n        sess (tf.Session): The tensorflow session used to run assignment.\n        variables (Dict[str, tf.Variable]): Extracted variables from the loss\n            or additional variables that are passed in.\n        placeholders (Dict[str, tf.placeholders]): Placeholders for weights.\n        assignment_nodes (Dict[str, tf.Tensor]): Nodes that assign weights.\n    \"\"\"\n\n    def __init__(self, output, sess=None, input_variables=None):\n        \"\"\"Creates TensorFlowVariables containing extracted variables.\n\n        The variables are extracted by performing a BFS search on the\n        dependency graph with loss as the root node. After the tree is\n        traversed and those variables are collected, we append input_variables\n        to the collected variables. For each variable in the list, the\n        variable has a placeholder and assignment operation created for it.\n\n        Args:\n            output (tf.Operation, List[tf.Operation]): The tensorflow\n                operation to extract all variables from.\n            sess (tf.Session): Session used for running the get and set\n                methods.\n            input_variables (List[tf.Variables]): Variables to include in the\n                list.\n        \"\"\"\n        self.sess = sess\n        if not isinstance(output, (list, tuple)):\n            output = [output]\n        queue = deque(output)\n        variable_names = []\n        explored_inputs = set(output)\n\n        # We do a BFS on the dependency graph of the input function to find\n        # the variables.\n        while len(queue) != 0:\n            tf_obj = queue.popleft()\n            if tf_obj is None:\n                continue\n            # The object put into the queue is not necessarily an operation,\n            # so we want the op attribute to get the operation underlying the\n            # object. Only operations contain the inputs that we can explore.\n            if hasattr(tf_obj, \"op\"):\n                tf_obj = tf_obj.op\n            for input_op in tf_obj.inputs:\n                if input_op not in explored_inputs:\n                    queue.append(input_op)\n                    explored_inputs.add(input_op)\n            # Tensorflow control inputs can be circular, so we keep track of\n            # explored operations.\n            for control in tf_obj.control_inputs:\n                if control not in explored_inputs:\n                    queue.append(control)\n                    explored_inputs.add(control)\n            if (\"Variable\" in tf_obj.node_def.op\n                    or \"VarHandle\" in tf_obj.node_def.op):\n                variable_names.append(tf_obj.node_def.name)\n        self.variables = OrderedDict()\n        variable_list = [\n            v for v in tf.global_variables()\n            if v.op.node_def.name in variable_names\n        ]\n        if input_variables is not None:\n            variable_list += input_variables\n        for v in variable_list:\n            self.variables[v.op.node_def.name] = v\n\n        self.placeholders = {}\n        self.assignment_nodes = {}\n\n        # Create new placeholders to put in custom weights.\n        for k, var in self.variables.items():\n            self.placeholders[k] = tf.placeholder(\n                var.value().dtype,\n                var.get_shape().as_list(),\n                name=\"Placeholder_\" + k)\n            self.assignment_nodes[k] = var.assign(self.placeholders[k])\n\n    def set_session(self, sess):\n        \"\"\"Sets the current session used by the class.\n\n        Args:\n            sess (tf.Session): Session to set the attribute with.\n        \"\"\"\n        self.sess = sess\n\n    def get_flat_size(self):\n        \"\"\"Returns the total length of all of the flattened variables.\n\n        Returns:\n            The length of all flattened variables concatenated.\n        \"\"\"\n        return sum(\n            np.prod(v.get_shape().as_list()) for v in self.variables.values())\n\n    def _check_sess(self):\n        \"\"\"Checks if the session is set, and if not throw an error message.\"\"\"\n        assert self.sess is not None, (\"The session is not set. Set the \"\n                                       \"session either by passing it into the \"\n                                       \"TensorFlowVariables constructor or by \"\n                                       \"calling set_session(sess).\")\n\n    def get_flat(self):\n        \"\"\"Gets the weights and returns them as a flat array.\n\n        Returns:\n            1D Array containing the flattened weights.\n        \"\"\"\n        self._check_sess()\n        return np.concatenate([\n            v.eval(session=self.sess).flatten()\n            for v in self.variables.values()\n        ])\n\n    def set_flat(self, new_weights):\n        \"\"\"Sets the weights to new_weights, converting from a flat array.\n\n        Note:\n            You can only set all weights in the network using this function,\n            i.e., the length of the array must match get_flat_size.\n\n        Args:\n            new_weights (np.ndarray): Flat array containing weights.\n        \"\"\"\n        self._check_sess()\n        shapes = [v.get_shape().as_list() for v in self.variables.values()]\n        arrays = unflatten(new_weights, shapes)\n        placeholders = [\n            self.placeholders[k] for k, v in self.variables.items()\n        ]\n        self.sess.run(\n            list(self.assignment_nodes.values()),\n            feed_dict=dict(zip(placeholders, arrays)))\n\n    def get_weights(self):\n        \"\"\"Returns a dictionary containing the weights of the network.\n\n        Returns:\n            Dictionary mapping variable names to their weights.\n        \"\"\"\n        self._check_sess()\n        return {\n            k: v.eval(session=self.sess)\n            for k, v in self.variables.items()\n        }\n\n    def set_weights(self, new_weights):\n        \"\"\"Sets the weights to new_weights.\n\n        Note:\n            Can set subsets of variables as well, by only passing in the\n            variables you want to be set.\n\n        Args:\n            new_weights (Dict): Dictionary mapping variable names to their\n                weights.\n        \"\"\"\n        self._check_sess()\n        assign_list = [\n            self.assignment_nodes[name] for name in new_weights.keys()\n            if name in self.assignment_nodes\n        ]\n        assert assign_list, (\"No variables in the input matched those in the \"\n                             \"network. Possible cause: Two networks were \"\n                             \"defined in the same TensorFlow graph. To fix \"\n                             \"this, place each network definition in its own \"\n                             \"tf.Graph.\")\n        self.sess.run(\n            assign_list,\n            feed_dict={\n                self.placeholders[name]: value\n                for (name, value) in new_weights.items()\n                if name in self.placeholders\n            })\n", "path": "python/ray/experimental/tf_utils.py"}], "after_files": [{"content": "from collections import deque, OrderedDict\nimport numpy as np\n\nfrom ray.rllib.utils import try_import_tf\n\ntf = try_import_tf()\n\n\ndef unflatten(vector, shapes):\n    i = 0\n    arrays = []\n    for shape in shapes:\n        size = np.prod(shape, dtype=np.int)\n        array = vector[i:(i + size)].reshape(shape)\n        arrays.append(array)\n        i += size\n    assert len(vector) == i, \"Passed weight does not have the correct shape.\"\n    return arrays\n\n\nclass TensorFlowVariables:\n    \"\"\"A class used to set and get weights for Tensorflow networks.\n\n    Attributes:\n        sess (tf.Session): The tensorflow session used to run assignment.\n        variables (Dict[str, tf.Variable]): Extracted variables from the loss\n            or additional variables that are passed in.\n        placeholders (Dict[str, tf.placeholders]): Placeholders for weights.\n        assignment_nodes (Dict[str, tf.Tensor]): Nodes that assign weights.\n    \"\"\"\n\n    def __init__(self, output, sess=None, input_variables=None):\n        \"\"\"Creates TensorFlowVariables containing extracted variables.\n\n        The variables are extracted by performing a BFS search on the\n        dependency graph with loss as the root node. After the tree is\n        traversed and those variables are collected, we append input_variables\n        to the collected variables. For each variable in the list, the\n        variable has a placeholder and assignment operation created for it.\n\n        Args:\n            output (tf.Operation, List[tf.Operation]): The tensorflow\n                operation to extract all variables from.\n            sess (tf.Session): Session used for running the get and set\n                methods.\n            input_variables (List[tf.Variables]): Variables to include in the\n                list.\n        \"\"\"\n        self.sess = sess\n        if not isinstance(output, (list, tuple)):\n            output = [output]\n        queue = deque(output)\n        variable_names = []\n        explored_inputs = set(output)\n\n        # We do a BFS on the dependency graph of the input function to find\n        # the variables.\n        while len(queue) != 0:\n            tf_obj = queue.popleft()\n            if tf_obj is None:\n                continue\n            # The object put into the queue is not necessarily an operation,\n            # so we want the op attribute to get the operation underlying the\n            # object. Only operations contain the inputs that we can explore.\n            if hasattr(tf_obj, \"op\"):\n                tf_obj = tf_obj.op\n            for input_op in tf_obj.inputs:\n                if input_op not in explored_inputs:\n                    queue.append(input_op)\n                    explored_inputs.add(input_op)\n            # Tensorflow control inputs can be circular, so we keep track of\n            # explored operations.\n            for control in tf_obj.control_inputs:\n                if control not in explored_inputs:\n                    queue.append(control)\n                    explored_inputs.add(control)\n            if (\"Variable\" in tf_obj.node_def.op\n                    or \"VarHandle\" in tf_obj.node_def.op):\n                variable_names.append(tf_obj.node_def.name)\n        self.variables = OrderedDict()\n        variable_list = [\n            v for v in tf.global_variables()\n            if v.op.node_def.name in variable_names\n        ]\n        if input_variables is not None:\n            variable_list += input_variables\n        for v in variable_list:\n            self.variables[v.op.node_def.name] = v\n\n        self.placeholders = {}\n        self.assignment_nodes = {}\n\n        # Create new placeholders to put in custom weights.\n        for k, var in self.variables.items():\n            self.placeholders[k] = tf.placeholder(\n                var.value().dtype,\n                var.get_shape().as_list(),\n                name=\"Placeholder_\" + k)\n            self.assignment_nodes[k] = var.assign(self.placeholders[k])\n\n    def set_session(self, sess):\n        \"\"\"Sets the current session used by the class.\n\n        Args:\n            sess (tf.Session): Session to set the attribute with.\n        \"\"\"\n        self.sess = sess\n\n    def get_flat_size(self):\n        \"\"\"Returns the total length of all of the flattened variables.\n\n        Returns:\n            The length of all flattened variables concatenated.\n        \"\"\"\n        return sum(\n            np.prod(v.get_shape().as_list()) for v in self.variables.values())\n\n    def _check_sess(self):\n        \"\"\"Checks if the session is set, and if not throw an error message.\"\"\"\n        assert self.sess is not None, (\"The session is not set. Set the \"\n                                       \"session either by passing it into the \"\n                                       \"TensorFlowVariables constructor or by \"\n                                       \"calling set_session(sess).\")\n\n    def get_flat(self):\n        \"\"\"Gets the weights and returns them as a flat array.\n\n        Returns:\n            1D Array containing the flattened weights.\n        \"\"\"\n        self._check_sess()\n        return np.concatenate([\n            v.eval(session=self.sess).flatten()\n            for v in self.variables.values()\n        ])\n\n    def set_flat(self, new_weights):\n        \"\"\"Sets the weights to new_weights, converting from a flat array.\n\n        Note:\n            You can only set all weights in the network using this function,\n            i.e., the length of the array must match get_flat_size.\n\n        Args:\n            new_weights (np.ndarray): Flat array containing weights.\n        \"\"\"\n        self._check_sess()\n        shapes = [v.get_shape().as_list() for v in self.variables.values()]\n        arrays = unflatten(new_weights, shapes)\n        placeholders = [\n            self.placeholders[k] for k, v in self.variables.items()\n        ]\n        self.sess.run(\n            list(self.assignment_nodes.values()),\n            feed_dict=dict(zip(placeholders, arrays)))\n\n    def get_weights(self):\n        \"\"\"Returns a dictionary containing the weights of the network.\n\n        Returns:\n            Dictionary mapping variable names to their weights.\n        \"\"\"\n        self._check_sess()\n        return self.sess.run(self.variables)\n\n    def set_weights(self, new_weights):\n        \"\"\"Sets the weights to new_weights.\n\n        Note:\n            Can set subsets of variables as well, by only passing in the\n            variables you want to be set.\n\n        Args:\n            new_weights (Dict): Dictionary mapping variable names to their\n                weights.\n        \"\"\"\n        self._check_sess()\n        assign_list = [\n            self.assignment_nodes[name] for name in new_weights.keys()\n            if name in self.assignment_nodes\n        ]\n        assert assign_list, (\"No variables in the input matched those in the \"\n                             \"network. Possible cause: Two networks were \"\n                             \"defined in the same TensorFlow graph. To fix \"\n                             \"this, place each network definition in its own \"\n                             \"tf.Graph.\")\n        self.sess.run(\n            assign_list,\n            feed_dict={\n                self.placeholders[name]: value\n                for (name, value) in new_weights.items()\n                if name in self.placeholders\n            })\n", "path": "python/ray/experimental/tf_utils.py"}]}