We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
variable initialization in local_tf_runner is slow and incorrect
The variable initialization code in function `initialize_tf_vars` of `garage/experiment/local_tf_runner.py` is written in a way that's both very inefficient and potentially incorrect. In particular, the list comprehension
````[
                    v for v in tf.global_variables()
                    if v.name.split(':')[0] in str(
                        self.sess.run(tf.report_uninitialized_variables()))
]
````
reruns the `tf.report_uninitialized_variables()` op every time the condition needs to be evaluated, which is very slow for larger graphs. It takes me several minutes to run for a convnet with 128x128x4 pixel inputs. Storing the result of this operation reduces runtime to mere seconds.

Further, the use of string operations can potentially lead to re-initializing some variables. If the name of some initialised variable A is a substring of an uninitialised variable B, the above list comprehension will include variable A in the result. A better alternative is to decode the bytestrings returned by `sess.run(tf.report_uninitialized_variables()` and construct a set.

The function code then becomes

````uninit_set = set(
            e.decode() for e in self.sess.run(tf.report_uninitialized_variables())
        )
        self.sess.run(
            tf.variables_initializer([
                v for v in tf.global_variables()
                if v.name.split(':')[0] in uninit_set
            ]))
````

If this idea seems reasonable, I am happy to make a pull request.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `garage/experiment/local_tf_runner.py`
Content:
```
1 """
2 The local runner for tensorflow algorithms.
3 
4 A runner setup context for algorithms during initialization and
5 pipelines data between sampler and algorithm during training.
6 """
7 import time
8 
9 import tensorflow as tf
10 
11 from garage.logger import logger
12 from garage.logger import snapshotter
13 from garage.logger import tabular
14 
15 # Note: Optional module should be imported ad hoc to break circular dependency.
16 
17 
18 class LocalRunner:
19     """This class implements a local runner for tensorflow algorithms.
20 
21     A local runner provides a default tensorflow session using python context.
22     This is useful for those experiment components (e.g. policy) that require a
23     tensorflow session during construction.
24 
25     Use Runner.setup(algo, env) to setup algorithm and environement for runner
26     and Runner.train() to start training.
27 
28     Examples:
29         with LocalRunner() as runner:
30             env = gym.make('CartPole-v1')
31             policy = CategoricalMLPPolicy(
32                 env_spec=env.spec,
33                 hidden_sizes=(32, 32))
34             algo = TRPO(
35                 env=env,
36                 policy=policy,
37                 baseline=baseline,
38                 max_path_length=100,
39                 discount=0.99,
40                 max_kl_step=0.01)
41             runner.setup(algo, env)
42             runner.train(n_epochs=100, batch_size=4000)
43 
44     """
45 
46     def __init__(self, sess=None, max_cpus=1):
47         """Create a new local runner.
48 
49         Args:
50             max_cpus: The maximum number of parallel sampler workers.
51             sess: An optional tensorflow session.
52                   A new session will be created immediately if not provided.
53 
54         Note:
55             The local runner will set up a joblib task pool of size max_cpus
56             possibly later used by BatchSampler. If BatchSampler is not used,
57             the processes in the pool will remain dormant.
58 
59             This setup is required to use tensorflow in a multiprocess
60             environment before a tensorflow session is created
61             because tensorflow is not fork-safe.
62 
63             See https://github.com/tensorflow/tensorflow/issues/2448.
64 
65         """
66         if max_cpus > 1:
67             from garage.sampler import singleton_pool
68             singleton_pool.initialize(max_cpus)
69         self.sess = sess or tf.Session()
70         self.has_setup = False
71         self.plot = False
72 
73     def __enter__(self):
74         """Set self.sess as the default session.
75 
76         Returns:
77             This local runner.
78 
79         """
80         if tf.get_default_session() is not self.sess:
81             self.sess.__enter__()
82         return self
83 
84     def __exit__(self, exc_type, exc_val, exc_tb):
85         """Leave session."""
86         if tf.get_default_session() is self.sess:
87             self.sess.__exit__(exc_type, exc_val, exc_tb)
88 
89     def setup(self, algo, env, sampler_cls=None, sampler_args=None):
90         """Set up runner for algorithm and environment.
91 
92         This method saves algo and env within runner and creates a sampler.
93 
94         Note:
95             After setup() is called all variables in session should have been
96             initialized. setup() respects existing values in session so
97             policy weights can be loaded before setup().
98 
99         Args:
100             algo: An algorithm instance.
101             env: An environement instance.
102             sampler_cls: A sampler class.
103             sampler_args: Arguments to be passed to sampler constructor.
104 
105         """
106         self.algo = algo
107         self.env = env
108         self.policy = self.algo.policy
109 
110         if sampler_args is None:
111             sampler_args = {}
112 
113         if sampler_cls is None:
114             from garage.tf.algos.batch_polopt import BatchPolopt
115             if isinstance(algo, BatchPolopt):
116                 if self.policy.vectorized:
117                     from garage.tf.samplers import OnPolicyVectorizedSampler
118                     sampler_cls = OnPolicyVectorizedSampler
119                 else:
120                     from garage.tf.samplers import BatchSampler
121                     sampler_cls = BatchSampler
122             else:
123                 from garage.tf.samplers import OffPolicyVectorizedSampler
124                 sampler_cls = OffPolicyVectorizedSampler
125 
126         self.sampler = sampler_cls(algo, env, **sampler_args)
127 
128         self.initialize_tf_vars()
129         logger.log(self.sess.graph)
130         self.has_setup = True
131 
132     def initialize_tf_vars(self):
133         """Initialize all uninitialized variables in session."""
134         with tf.name_scope("initialize_tf_vars"):
135             self.sess.run(
136                 tf.variables_initializer([
137                     v for v in tf.global_variables()
138                     if v.name.split(':')[0] in str(
139                         self.sess.run(tf.report_uninitialized_variables()))
140                 ]))
141 
142     def start_worker(self):
143         """Start Plotter and Sampler workers."""
144         self.sampler.start_worker()
145         if self.plot:
146             from garage.tf.plotter import Plotter
147             self.plotter = Plotter(self.env, self.policy)
148             self.plotter.start()
149 
150     def shutdown_worker(self):
151         """Shutdown Plotter and Sampler workers."""
152         self.sampler.shutdown_worker()
153         if self.plot:
154             self.plotter.close()
155 
156     def obtain_samples(self, itr, batch_size):
157         """Obtain one batch of samples.
158 
159         Args:
160             itr: Index of iteration (epoch).
161             batch_size: Number of steps in batch.
162                 This is a hint that the sampler may or may not respect.
163 
164         Returns:
165             One batch of samples.
166 
167         """
168         if self.n_epoch_cycles == 1:
169             logger.log('Obtaining samples...')
170         return self.sampler.obtain_samples(itr, batch_size)
171 
172     def save_snapshot(self, itr, paths=None):
173         """Save snapshot of current batch.
174 
175         Args:
176             itr: Index of iteration (epoch).
177             paths: Batch of samples after preprocessed.
178 
179         """
180         assert self.has_setup
181 
182         logger.log("Saving snapshot...")
183         params = self.algo.get_itr_snapshot(itr)
184         params['env'] = self.env
185         if paths:
186             params['paths'] = paths
187         snapshotter.save_itr_params(itr, params)
188         logger.log('Saved')
189 
190     def log_diagnostics(self, pause_for_plot=False):
191         """Log diagnostics.
192 
193         Args:
194             pause_for_plot: Pause for plot.
195 
196         """
197         logger.log('Time %.2f s' % (time.time() - self.start_time))
198         logger.log('EpochTime %.2f s' % (time.time() - self.itr_start_time))
199         logger.log(tabular)
200         if self.plot:
201             self.plotter.update_plot(self.policy, self.algo.max_path_length)
202             if pause_for_plot:
203                 input('Plotting evaluation run: Press Enter to " "continue...')
204 
205     def train(self,
206               n_epochs,
207               n_epoch_cycles=1,
208               batch_size=None,
209               plot=False,
210               store_paths=False,
211               pause_for_plot=False):
212         """Start training.
213 
214         Args:
215             n_epochs: Number of epochs.
216             n_epoch_cycles: Number of batches of samples in each epoch.
217                 This is only useful for off-policy algorithm.
218                 For on-policy algorithm this value should always be 1.
219             batch_size: Number of steps in batch.
220             plot: Visualize policy by doing rollout after each epoch.
221             store_paths: Save paths in snapshot.
222             pause_for_plot: Pause for plot.
223 
224         Returns:
225             The average return in last epoch cycle.
226 
227         """
228         assert self.has_setup, ('Use Runner.setup() to setup runner before '
229                                 'training.')
230         if batch_size is None:
231             from garage.tf.samplers import OffPolicyVectorizedSampler
232             if isinstance(self.sampler, OffPolicyVectorizedSampler):
233                 batch_size = self.algo.max_path_length
234             else:
235                 batch_size = 40 * self.algo.max_path_length
236 
237         self.n_epoch_cycles = n_epoch_cycles
238 
239         self.plot = plot
240         self.start_worker()
241         self.start_time = time.time()
242 
243         itr = 0
244         last_return = None
245         for epoch in range(n_epochs):
246             self.itr_start_time = time.time()
247             paths = None
248             with logger.prefix('epoch #%d | ' % epoch):
249                 for cycle in range(n_epoch_cycles):
250                     paths = self.obtain_samples(itr, batch_size)
251                     paths = self.sampler.process_samples(itr, paths)
252                     last_return = self.algo.train_once(itr, paths)
253                     itr += 1
254                 self.save_snapshot(epoch, paths if store_paths else None)
255                 self.log_diagnostics(pause_for_plot)
256                 logger.dump_all(itr)
257                 tabular.clear()
258 
259         self.shutdown_worker()
260         return last_return
261
```
--- 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/garage/experiment/local_tf_runner.py b/garage/experiment/local_tf_runner.py
--- a/garage/experiment/local_tf_runner.py
+++ b/garage/experiment/local_tf_runner.py
@@ -132,11 +132,14 @@
     def initialize_tf_vars(self):
         """Initialize all uninitialized variables in session."""
         with tf.name_scope("initialize_tf_vars"):
+            uninited_set = [
+                e.decode()
+                for e in self.sess.run(tf.report_uninitialized_variables())
+            ]
             self.sess.run(
                 tf.variables_initializer([
                     v for v in tf.global_variables()
-                    if v.name.split(':')[0] in str(
-                        self.sess.run(tf.report_uninitialized_variables()))
+                    if v.name.split(':')[0] in uninited_set
                 ]))
 
     def start_worker(self):

{"golden_diff": "diff --git a/garage/experiment/local_tf_runner.py b/garage/experiment/local_tf_runner.py\n--- a/garage/experiment/local_tf_runner.py\n+++ b/garage/experiment/local_tf_runner.py\n@@ -132,11 +132,14 @@\n     def initialize_tf_vars(self):\n         \"\"\"Initialize all uninitialized variables in session.\"\"\"\n         with tf.name_scope(\"initialize_tf_vars\"):\n+            uninited_set = [\n+                e.decode()\n+                for e in self.sess.run(tf.report_uninitialized_variables())\n+            ]\n             self.sess.run(\n                 tf.variables_initializer([\n                     v for v in tf.global_variables()\n-                    if v.name.split(':')[0] in str(\n-                        self.sess.run(tf.report_uninitialized_variables()))\n+                    if v.name.split(':')[0] in uninited_set\n                 ]))\n \n     def start_worker(self):\n", "issue": "variable initialization in local_tf_runner is slow and incorrect\nThe variable initialization code in function `initialize_tf_vars` of `garage/experiment/local_tf_runner.py` is written in a way that's both very inefficient and potentially incorrect. In particular, the list comprehension\r\n````[\r\n                    v for v in tf.global_variables()\r\n                    if v.name.split(':')[0] in str(\r\n                        self.sess.run(tf.report_uninitialized_variables()))\r\n]\r\n````\r\nreruns the `tf.report_uninitialized_variables()` op every time the condition needs to be evaluated, which is very slow for larger graphs. It takes me several minutes to run for a convnet with 128x128x4 pixel inputs. Storing the result of this operation reduces runtime to mere seconds.\r\n\r\nFurther, the use of string operations can potentially lead to re-initializing some variables. If the name of some initialised variable A is a substring of an uninitialised variable B, the above list comprehension will include variable A in the result. A better alternative is to decode the bytestrings returned by `sess.run(tf.report_uninitialized_variables()` and construct a set.\r\n\r\nThe function code then becomes\r\n\r\n````uninit_set = set(\r\n            e.decode() for e in self.sess.run(tf.report_uninitialized_variables())\r\n        )\r\n        self.sess.run(\r\n            tf.variables_initializer([\r\n                v for v in tf.global_variables()\r\n                if v.name.split(':')[0] in uninit_set\r\n            ]))\r\n````\r\n\r\nIf this idea seems reasonable, I am happy to make a pull request.\n", "before_files": [{"content": "\"\"\"\nThe local runner for tensorflow algorithms.\n\nA runner setup context for algorithms during initialization and\npipelines data between sampler and algorithm during training.\n\"\"\"\nimport time\n\nimport tensorflow as tf\n\nfrom garage.logger import logger\nfrom garage.logger import snapshotter\nfrom garage.logger import tabular\n\n# Note: Optional module should be imported ad hoc to break circular dependency.\n\n\nclass LocalRunner:\n    \"\"\"This class implements a local runner for tensorflow algorithms.\n\n    A local runner provides a default tensorflow session using python context.\n    This is useful for those experiment components (e.g. policy) that require a\n    tensorflow session during construction.\n\n    Use Runner.setup(algo, env) to setup algorithm and environement for runner\n    and Runner.train() to start training.\n\n    Examples:\n        with LocalRunner() as runner:\n            env = gym.make('CartPole-v1')\n            policy = CategoricalMLPPolicy(\n                env_spec=env.spec,\n                hidden_sizes=(32, 32))\n            algo = TRPO(\n                env=env,\n                policy=policy,\n                baseline=baseline,\n                max_path_length=100,\n                discount=0.99,\n                max_kl_step=0.01)\n            runner.setup(algo, env)\n            runner.train(n_epochs=100, batch_size=4000)\n\n    \"\"\"\n\n    def __init__(self, sess=None, max_cpus=1):\n        \"\"\"Create a new local runner.\n\n        Args:\n            max_cpus: The maximum number of parallel sampler workers.\n            sess: An optional tensorflow session.\n                  A new session will be created immediately if not provided.\n\n        Note:\n            The local runner will set up a joblib task pool of size max_cpus\n            possibly later used by BatchSampler. If BatchSampler is not used,\n            the processes in the pool will remain dormant.\n\n            This setup is required to use tensorflow in a multiprocess\n            environment before a tensorflow session is created\n            because tensorflow is not fork-safe.\n\n            See https://github.com/tensorflow/tensorflow/issues/2448.\n\n        \"\"\"\n        if max_cpus > 1:\n            from garage.sampler import singleton_pool\n            singleton_pool.initialize(max_cpus)\n        self.sess = sess or tf.Session()\n        self.has_setup = False\n        self.plot = False\n\n    def __enter__(self):\n        \"\"\"Set self.sess as the default session.\n\n        Returns:\n            This local runner.\n\n        \"\"\"\n        if tf.get_default_session() is not self.sess:\n            self.sess.__enter__()\n        return self\n\n    def __exit__(self, exc_type, exc_val, exc_tb):\n        \"\"\"Leave session.\"\"\"\n        if tf.get_default_session() is self.sess:\n            self.sess.__exit__(exc_type, exc_val, exc_tb)\n\n    def setup(self, algo, env, sampler_cls=None, sampler_args=None):\n        \"\"\"Set up runner for algorithm and environment.\n\n        This method saves algo and env within runner and creates a sampler.\n\n        Note:\n            After setup() is called all variables in session should have been\n            initialized. setup() respects existing values in session so\n            policy weights can be loaded before setup().\n\n        Args:\n            algo: An algorithm instance.\n            env: An environement instance.\n            sampler_cls: A sampler class.\n            sampler_args: Arguments to be passed to sampler constructor.\n\n        \"\"\"\n        self.algo = algo\n        self.env = env\n        self.policy = self.algo.policy\n\n        if sampler_args is None:\n            sampler_args = {}\n\n        if sampler_cls is None:\n            from garage.tf.algos.batch_polopt import BatchPolopt\n            if isinstance(algo, BatchPolopt):\n                if self.policy.vectorized:\n                    from garage.tf.samplers import OnPolicyVectorizedSampler\n                    sampler_cls = OnPolicyVectorizedSampler\n                else:\n                    from garage.tf.samplers import BatchSampler\n                    sampler_cls = BatchSampler\n            else:\n                from garage.tf.samplers import OffPolicyVectorizedSampler\n                sampler_cls = OffPolicyVectorizedSampler\n\n        self.sampler = sampler_cls(algo, env, **sampler_args)\n\n        self.initialize_tf_vars()\n        logger.log(self.sess.graph)\n        self.has_setup = True\n\n    def initialize_tf_vars(self):\n        \"\"\"Initialize all uninitialized variables in session.\"\"\"\n        with tf.name_scope(\"initialize_tf_vars\"):\n            self.sess.run(\n                tf.variables_initializer([\n                    v for v in tf.global_variables()\n                    if v.name.split(':')[0] in str(\n                        self.sess.run(tf.report_uninitialized_variables()))\n                ]))\n\n    def start_worker(self):\n        \"\"\"Start Plotter and Sampler workers.\"\"\"\n        self.sampler.start_worker()\n        if self.plot:\n            from garage.tf.plotter import Plotter\n            self.plotter = Plotter(self.env, self.policy)\n            self.plotter.start()\n\n    def shutdown_worker(self):\n        \"\"\"Shutdown Plotter and Sampler workers.\"\"\"\n        self.sampler.shutdown_worker()\n        if self.plot:\n            self.plotter.close()\n\n    def obtain_samples(self, itr, batch_size):\n        \"\"\"Obtain one batch of samples.\n\n        Args:\n            itr: Index of iteration (epoch).\n            batch_size: Number of steps in batch.\n                This is a hint that the sampler may or may not respect.\n\n        Returns:\n            One batch of samples.\n\n        \"\"\"\n        if self.n_epoch_cycles == 1:\n            logger.log('Obtaining samples...')\n        return self.sampler.obtain_samples(itr, batch_size)\n\n    def save_snapshot(self, itr, paths=None):\n        \"\"\"Save snapshot of current batch.\n\n        Args:\n            itr: Index of iteration (epoch).\n            paths: Batch of samples after preprocessed.\n\n        \"\"\"\n        assert self.has_setup\n\n        logger.log(\"Saving snapshot...\")\n        params = self.algo.get_itr_snapshot(itr)\n        params['env'] = self.env\n        if paths:\n            params['paths'] = paths\n        snapshotter.save_itr_params(itr, params)\n        logger.log('Saved')\n\n    def log_diagnostics(self, pause_for_plot=False):\n        \"\"\"Log diagnostics.\n\n        Args:\n            pause_for_plot: Pause for plot.\n\n        \"\"\"\n        logger.log('Time %.2f s' % (time.time() - self.start_time))\n        logger.log('EpochTime %.2f s' % (time.time() - self.itr_start_time))\n        logger.log(tabular)\n        if self.plot:\n            self.plotter.update_plot(self.policy, self.algo.max_path_length)\n            if pause_for_plot:\n                input('Plotting evaluation run: Press Enter to \" \"continue...')\n\n    def train(self,\n              n_epochs,\n              n_epoch_cycles=1,\n              batch_size=None,\n              plot=False,\n              store_paths=False,\n              pause_for_plot=False):\n        \"\"\"Start training.\n\n        Args:\n            n_epochs: Number of epochs.\n            n_epoch_cycles: Number of batches of samples in each epoch.\n                This is only useful for off-policy algorithm.\n                For on-policy algorithm this value should always be 1.\n            batch_size: Number of steps in batch.\n            plot: Visualize policy by doing rollout after each epoch.\n            store_paths: Save paths in snapshot.\n            pause_for_plot: Pause for plot.\n\n        Returns:\n            The average return in last epoch cycle.\n\n        \"\"\"\n        assert self.has_setup, ('Use Runner.setup() to setup runner before '\n                                'training.')\n        if batch_size is None:\n            from garage.tf.samplers import OffPolicyVectorizedSampler\n            if isinstance(self.sampler, OffPolicyVectorizedSampler):\n                batch_size = self.algo.max_path_length\n            else:\n                batch_size = 40 * self.algo.max_path_length\n\n        self.n_epoch_cycles = n_epoch_cycles\n\n        self.plot = plot\n        self.start_worker()\n        self.start_time = time.time()\n\n        itr = 0\n        last_return = None\n        for epoch in range(n_epochs):\n            self.itr_start_time = time.time()\n            paths = None\n            with logger.prefix('epoch #%d | ' % epoch):\n                for cycle in range(n_epoch_cycles):\n                    paths = self.obtain_samples(itr, batch_size)\n                    paths = self.sampler.process_samples(itr, paths)\n                    last_return = self.algo.train_once(itr, paths)\n                    itr += 1\n                self.save_snapshot(epoch, paths if store_paths else None)\n                self.log_diagnostics(pause_for_plot)\n                logger.dump_all(itr)\n                tabular.clear()\n\n        self.shutdown_worker()\n        return last_return\n", "path": "garage/experiment/local_tf_runner.py"}], "after_files": [{"content": "\"\"\"\nThe local runner for tensorflow algorithms.\n\nA runner setup context for algorithms during initialization and\npipelines data between sampler and algorithm during training.\n\"\"\"\nimport time\n\nimport tensorflow as tf\n\nfrom garage.logger import logger\nfrom garage.logger import snapshotter\nfrom garage.logger import tabular\n\n# Note: Optional module should be imported ad hoc to break circular dependency.\n\n\nclass LocalRunner:\n    \"\"\"This class implements a local runner for tensorflow algorithms.\n\n    A local runner provides a default tensorflow session using python context.\n    This is useful for those experiment components (e.g. policy) that require a\n    tensorflow session during construction.\n\n    Use Runner.setup(algo, env) to setup algorithm and environement for runner\n    and Runner.train() to start training.\n\n    Examples:\n        with LocalRunner() as runner:\n            env = gym.make('CartPole-v1')\n            policy = CategoricalMLPPolicy(\n                env_spec=env.spec,\n                hidden_sizes=(32, 32))\n            algo = TRPO(\n                env=env,\n                policy=policy,\n                baseline=baseline,\n                max_path_length=100,\n                discount=0.99,\n                max_kl_step=0.01)\n            runner.setup(algo, env)\n            runner.train(n_epochs=100, batch_size=4000)\n\n    \"\"\"\n\n    def __init__(self, sess=None, max_cpus=1):\n        \"\"\"Create a new local runner.\n\n        Args:\n            max_cpus: The maximum number of parallel sampler workers.\n            sess: An optional tensorflow session.\n                  A new session will be created immediately if not provided.\n\n        Note:\n            The local runner will set up a joblib task pool of size max_cpus\n            possibly later used by BatchSampler. If BatchSampler is not used,\n            the processes in the pool will remain dormant.\n\n            This setup is required to use tensorflow in a multiprocess\n            environment before a tensorflow session is created\n            because tensorflow is not fork-safe.\n\n            See https://github.com/tensorflow/tensorflow/issues/2448.\n\n        \"\"\"\n        if max_cpus > 1:\n            from garage.sampler import singleton_pool\n            singleton_pool.initialize(max_cpus)\n        self.sess = sess or tf.Session()\n        self.has_setup = False\n        self.plot = False\n\n    def __enter__(self):\n        \"\"\"Set self.sess as the default session.\n\n        Returns:\n            This local runner.\n\n        \"\"\"\n        if tf.get_default_session() is not self.sess:\n            self.sess.__enter__()\n        return self\n\n    def __exit__(self, exc_type, exc_val, exc_tb):\n        \"\"\"Leave session.\"\"\"\n        if tf.get_default_session() is self.sess:\n            self.sess.__exit__(exc_type, exc_val, exc_tb)\n\n    def setup(self, algo, env, sampler_cls=None, sampler_args=None):\n        \"\"\"Set up runner for algorithm and environment.\n\n        This method saves algo and env within runner and creates a sampler.\n\n        Note:\n            After setup() is called all variables in session should have been\n            initialized. setup() respects existing values in session so\n            policy weights can be loaded before setup().\n\n        Args:\n            algo: An algorithm instance.\n            env: An environement instance.\n            sampler_cls: A sampler class.\n            sampler_args: Arguments to be passed to sampler constructor.\n\n        \"\"\"\n        self.algo = algo\n        self.env = env\n        self.policy = self.algo.policy\n\n        if sampler_args is None:\n            sampler_args = {}\n\n        if sampler_cls is None:\n            from garage.tf.algos.batch_polopt import BatchPolopt\n            if isinstance(algo, BatchPolopt):\n                if self.policy.vectorized:\n                    from garage.tf.samplers import OnPolicyVectorizedSampler\n                    sampler_cls = OnPolicyVectorizedSampler\n                else:\n                    from garage.tf.samplers import BatchSampler\n                    sampler_cls = BatchSampler\n            else:\n                from garage.tf.samplers import OffPolicyVectorizedSampler\n                sampler_cls = OffPolicyVectorizedSampler\n\n        self.sampler = sampler_cls(algo, env, **sampler_args)\n\n        self.initialize_tf_vars()\n        logger.log(self.sess.graph)\n        self.has_setup = True\n\n    def initialize_tf_vars(self):\n        \"\"\"Initialize all uninitialized variables in session.\"\"\"\n        with tf.name_scope(\"initialize_tf_vars\"):\n            uninited_set = [\n                e.decode()\n                for e in self.sess.run(tf.report_uninitialized_variables())\n            ]\n            self.sess.run(\n                tf.variables_initializer([\n                    v for v in tf.global_variables()\n                    if v.name.split(':')[0] in uninited_set\n                ]))\n\n    def start_worker(self):\n        \"\"\"Start Plotter and Sampler workers.\"\"\"\n        self.sampler.start_worker()\n        if self.plot:\n            from garage.tf.plotter import Plotter\n            self.plotter = Plotter(self.env, self.policy)\n            self.plotter.start()\n\n    def shutdown_worker(self):\n        \"\"\"Shutdown Plotter and Sampler workers.\"\"\"\n        self.sampler.shutdown_worker()\n        if self.plot:\n            self.plotter.close()\n\n    def obtain_samples(self, itr, batch_size):\n        \"\"\"Obtain one batch of samples.\n\n        Args:\n            itr: Index of iteration (epoch).\n            batch_size: Number of steps in batch.\n                This is a hint that the sampler may or may not respect.\n\n        Returns:\n            One batch of samples.\n\n        \"\"\"\n        if self.n_epoch_cycles == 1:\n            logger.log('Obtaining samples...')\n        return self.sampler.obtain_samples(itr, batch_size)\n\n    def save_snapshot(self, itr, paths=None):\n        \"\"\"Save snapshot of current batch.\n\n        Args:\n            itr: Index of iteration (epoch).\n            paths: Batch of samples after preprocessed.\n\n        \"\"\"\n        assert self.has_setup\n\n        logger.log(\"Saving snapshot...\")\n        params = self.algo.get_itr_snapshot(itr)\n        params['env'] = self.env\n        if paths:\n            params['paths'] = paths\n        snapshotter.save_itr_params(itr, params)\n        logger.log('Saved')\n\n    def log_diagnostics(self, pause_for_plot=False):\n        \"\"\"Log diagnostics.\n\n        Args:\n            pause_for_plot: Pause for plot.\n\n        \"\"\"\n        logger.log('Time %.2f s' % (time.time() - self.start_time))\n        logger.log('EpochTime %.2f s' % (time.time() - self.itr_start_time))\n        logger.log(tabular)\n        if self.plot:\n            self.plotter.update_plot(self.policy, self.algo.max_path_length)\n            if pause_for_plot:\n                input('Plotting evaluation run: Press Enter to \" \"continue...')\n\n    def train(self,\n              n_epochs,\n              n_epoch_cycles=1,\n              batch_size=None,\n              plot=False,\n              store_paths=False,\n              pause_for_plot=False):\n        \"\"\"Start training.\n\n        Args:\n            n_epochs: Number of epochs.\n            n_epoch_cycles: Number of batches of samples in each epoch.\n                This is only useful for off-policy algorithm.\n                For on-policy algorithm this value should always be 1.\n            batch_size: Number of steps in batch.\n            plot: Visualize policy by doing rollout after each epoch.\n            store_paths: Save paths in snapshot.\n            pause_for_plot: Pause for plot.\n\n        Returns:\n            The average return in last epoch cycle.\n\n        \"\"\"\n        assert self.has_setup, ('Use Runner.setup() to setup runner before '\n                                'training.')\n        if batch_size is None:\n            from garage.tf.samplers import OffPolicyVectorizedSampler\n            if isinstance(self.sampler, OffPolicyVectorizedSampler):\n                batch_size = self.algo.max_path_length\n            else:\n                batch_size = 40 * self.algo.max_path_length\n\n        self.n_epoch_cycles = n_epoch_cycles\n\n        self.plot = plot\n        self.start_worker()\n        self.start_time = time.time()\n\n        itr = 0\n        last_return = None\n        for epoch in range(n_epochs):\n            self.itr_start_time = time.time()\n            paths = None\n            with logger.prefix('epoch #%d | ' % epoch):\n                for cycle in range(n_epoch_cycles):\n                    paths = self.obtain_samples(itr, batch_size)\n                    paths = self.sampler.process_samples(itr, paths)\n                    last_return = self.algo.train_once(itr, paths)\n                    itr += 1\n                self.save_snapshot(epoch, paths if store_paths else None)\n                self.log_diagnostics(pause_for_plot)\n                logger.dump_all(itr)\n                tabular.clear()\n\n        self.shutdown_worker()\n        return last_return\n", "path": "garage/experiment/local_tf_runner.py"}]}