We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
LocalRunner closes tf sessions it does not own
`garage.experiment.local_tf_runner.LocalRunner` will close its tf session in `__exit__`. However, the session may not have been created by LocalRunner and may be managed by some other entity. Closing externally-owned sessions can lead to unexpected behavior or failures.

Example:
```
with tf.Session() as sess:
  with LocalRunner(sess=sess) as runner:
    # Run an experiment
```
Code will fail with error:
`AttributeError: 'NoneType' object has no attribute '__exit__'`

--- 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             uninited_set = [
136                 e.decode()
137                 for e in self.sess.run(tf.report_uninitialized_variables())
138             ]
139             self.sess.run(
140                 tf.variables_initializer([
141                     v for v in tf.global_variables()
142                     if v.name.split(':')[0] in uninited_set
143                 ]))
144 
145     def start_worker(self):
146         """Start Plotter and Sampler workers."""
147         self.sampler.start_worker()
148         if self.plot:
149             from garage.tf.plotter import Plotter
150             self.plotter = Plotter(self.env, self.policy)
151             self.plotter.start()
152 
153     def shutdown_worker(self):
154         """Shutdown Plotter and Sampler workers."""
155         self.sampler.shutdown_worker()
156         if self.plot:
157             self.plotter.close()
158 
159     def obtain_samples(self, itr, batch_size):
160         """Obtain one batch of samples.
161 
162         Args:
163             itr: Index of iteration (epoch).
164             batch_size: Number of steps in batch.
165                 This is a hint that the sampler may or may not respect.
166 
167         Returns:
168             One batch of samples.
169 
170         """
171         if self.n_epoch_cycles == 1:
172             logger.log('Obtaining samples...')
173         return self.sampler.obtain_samples(itr, batch_size)
174 
175     def save_snapshot(self, itr, paths=None):
176         """Save snapshot of current batch.
177 
178         Args:
179             itr: Index of iteration (epoch).
180             paths: Batch of samples after preprocessed.
181 
182         """
183         assert self.has_setup
184 
185         logger.log("Saving snapshot...")
186         params = self.algo.get_itr_snapshot(itr)
187         params['env'] = self.env
188         if paths:
189             params['paths'] = paths
190         snapshotter.save_itr_params(itr, params)
191         logger.log('Saved')
192 
193     def log_diagnostics(self, pause_for_plot=False):
194         """Log diagnostics.
195 
196         Args:
197             pause_for_plot: Pause for plot.
198 
199         """
200         logger.log('Time %.2f s' % (time.time() - self.start_time))
201         logger.log('EpochTime %.2f s' % (time.time() - self.itr_start_time))
202         logger.log(tabular)
203         if self.plot:
204             self.plotter.update_plot(self.policy, self.algo.max_path_length)
205             if pause_for_plot:
206                 input('Plotting evaluation run: Press Enter to " "continue...')
207 
208     def train(self,
209               n_epochs,
210               n_epoch_cycles=1,
211               batch_size=None,
212               plot=False,
213               store_paths=False,
214               pause_for_plot=False):
215         """Start training.
216 
217         Args:
218             n_epochs: Number of epochs.
219             n_epoch_cycles: Number of batches of samples in each epoch.
220                 This is only useful for off-policy algorithm.
221                 For on-policy algorithm this value should always be 1.
222             batch_size: Number of steps in batch.
223             plot: Visualize policy by doing rollout after each epoch.
224             store_paths: Save paths in snapshot.
225             pause_for_plot: Pause for plot.
226 
227         Returns:
228             The average return in last epoch cycle.
229 
230         """
231         assert self.has_setup, ('Use Runner.setup() to setup runner before '
232                                 'training.')
233         if batch_size is None:
234             from garage.tf.samplers import OffPolicyVectorizedSampler
235             if isinstance(self.sampler, OffPolicyVectorizedSampler):
236                 batch_size = self.algo.max_path_length
237             else:
238                 batch_size = 40 * self.algo.max_path_length
239 
240         self.n_epoch_cycles = n_epoch_cycles
241 
242         self.plot = plot
243         self.start_worker()
244         self.start_time = time.time()
245 
246         itr = 0
247         last_return = None
248         for epoch in range(n_epochs):
249             self.itr_start_time = time.time()
250             paths = None
251             with logger.prefix('epoch #%d | ' % epoch):
252                 for cycle in range(n_epoch_cycles):
253                     paths = self.obtain_samples(itr, batch_size)
254                     paths = self.sampler.process_samples(itr, paths)
255                     last_return = self.algo.train_once(itr, paths)
256                     itr += 1
257                 self.save_snapshot(epoch, paths if store_paths else None)
258                 self.log_diagnostics(pause_for_plot)
259                 logger.dump_all(itr)
260                 tabular.clear()
261 
262         self.shutdown_worker()
263         return last_return
264
```
--- 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
@@ -67,6 +67,7 @@
             from garage.sampler import singleton_pool
             singleton_pool.initialize(max_cpus)
         self.sess = sess or tf.Session()
+        self.sess_entered = False
         self.has_setup = False
         self.plot = False
 
@@ -79,12 +80,14 @@
         """
         if tf.get_default_session() is not self.sess:
             self.sess.__enter__()
+            self.sess_entered = True
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Leave session."""
-        if tf.get_default_session() is self.sess:
+        if tf.get_default_session() is self.sess and self.sess_entered:
             self.sess.__exit__(exc_type, exc_val, exc_tb)
+            self.sess_entered = False
 
     def setup(self, algo, env, sampler_cls=None, sampler_args=None):
         """Set up runner for algorithm and environment.

{"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@@ -67,6 +67,7 @@\n             from garage.sampler import singleton_pool\n             singleton_pool.initialize(max_cpus)\n         self.sess = sess or tf.Session()\n+        self.sess_entered = False\n         self.has_setup = False\n         self.plot = False\n \n@@ -79,12 +80,14 @@\n         \"\"\"\n         if tf.get_default_session() is not self.sess:\n             self.sess.__enter__()\n+            self.sess_entered = True\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+        if tf.get_default_session() is self.sess and self.sess_entered:\n             self.sess.__exit__(exc_type, exc_val, exc_tb)\n+            self.sess_entered = False\n \n     def setup(self, algo, env, sampler_cls=None, sampler_args=None):\n         \"\"\"Set up runner for algorithm and environment.\n", "issue": "LocalRunner closes tf sessions it does not own\n`garage.experiment.local_tf_runner.LocalRunner` will close its tf session in `__exit__`. However, the session may not have been created by LocalRunner and may be managed by some other entity. Closing externally-owned sessions can lead to unexpected behavior or failures.\r\n\r\nExample:\r\n```\r\nwith tf.Session() as sess:\r\n  with LocalRunner(sess=sess) as runner:\r\n    # Run an experiment\r\n```\r\nCode will fail with error:\r\n`AttributeError: 'NoneType' object has no attribute '__exit__'`\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            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"}], "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.sess_entered = False\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            self.sess_entered = True\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 and self.sess_entered:\n            self.sess.__exit__(exc_type, exc_val, exc_tb)\n            self.sess_entered = False\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"}]}