We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
RaySampler does not use all workers
In release 2019.10, [RaySampler](https://github.com/rlworkgroup/garage/blob/release-2019.10/src/garage/sampler/ray_sampler.py) sees only one of the workers as idle for all iterations after the first iteration. This can be seen with [examples/tf/trpo_swimmer_ray_sampler.py](https://github.com/rlworkgroup/garage/blob/release-2019.10/examples/tf/trpo_swimmer_ray_sampler.py). 

To me it looks like `self._idle_worker_ids` is not correctly updated at the end of `obtain_samples()` [here](https://github.com/rlworkgroup/garage/blob/d6c993526c23c289ebc94635a78186d9bf197f32/src/garage/sampler/ray_sampler.py#L130). Only those workers which have returned with a result are put back in the idle pool. This is a problem because results from some workers are not processed at the end and these workers don't go back into the idle pool. 

A quick fix (?) is to reset the list of idle workers ids to all workers ids at the start of `obtain samples()`. This does result in all the workers being used for all iterations. Since the agent parameters are updated in the next iteration, the samples from the previously not ready workers must be valid. However, I am not totally sure. 

@avnishn @ryanjulian Appreciate it if the fix can be back-ported to 2019.10!

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/garage/sampler/ray_sampler.py`
Content:
```
1 """This is an implementation of an on policy batch sampler.
2 
3 Uses a data parallel design.
4 Included is a sampler that deploys sampler workers.
5 
6 The sampler workers must implement some type of set agent parameters
7 function, and a rollout function
8 """
9 from collections import defaultdict
10 import pickle
11 
12 import numpy as np
13 import psutil
14 import ray
15 
16 from garage.experiment import deterministic
17 from garage.misc.prog_bar_counter import ProgBarCounter
18 from garage.sampler.base import BaseSampler
19 
20 
21 class RaySampler(BaseSampler):
22     """Collects Policy Rollouts in a data parallel fashion.
23 
24     Args:
25         algo (garage.np.algo.RLAlgorithm): A garage algo object
26         env (gym.Env): A gym/akro env object
27         seed (int): Random seed.
28         should_render (bool): If the sampler render the trajectories.
29         num_processors (int): Number of processors to be used.
30         sampler_worker_cls (garage.sampler.ray_sampler.SamplerWorker):
31             If none, uses the default SamplerWorker class
32 
33     """
34 
35     def __init__(self,
36                  algo,
37                  env,
38                  seed,
39                  should_render=False,
40                  num_processors=None,
41                  sampler_worker_cls=None):
42         super().__init__(algo, env)
43         self._sampler_worker = ray.remote(SamplerWorker if sampler_worker_cls
44                                           is None else sampler_worker_cls)
45         self._seed = seed
46         deterministic.set_seed(self._seed)
47         self._max_path_length = self.algo.max_path_length
48         self._should_render = should_render
49         if not ray.is_initialized():
50             ray.init(log_to_driver=False)
51         self._num_workers = (num_processors if num_processors else
52                              psutil.cpu_count(logical=False))
53         self._all_workers = defaultdict(None)
54         self._idle_worker_ids = list(range(self._num_workers))
55         self._active_worker_ids = []
56 
57     def start_worker(self):
58         """Initialize a new ray worker."""
59         # Pickle the environment once, instead of once per worker.
60         env_pkl = pickle.dumps(self.env)
61         # We need to pickle the agent so that we can e.g. set up the TF Session
62         # in the worker *before* unpicling it.
63         agent_pkl = pickle.dumps(self.algo.policy)
64         for worker_id in range(self._num_workers):
65             self._all_workers[worker_id] = self._sampler_worker.remote(
66                 worker_id, env_pkl, agent_pkl, self._seed,
67                 self._max_path_length, self._should_render)
68 
69     # pylint: disable=arguments-differ
70     def obtain_samples(self, itr, num_samples):
71         """Sample the policy for new trajectories.
72 
73         Args:
74             itr (int): Iteration number.
75             num_samples (int): Number of steps the the sampler should collect.
76 
77         Returns:
78             list[dict]: Sample paths, each path with key
79                 * observations: (numpy.ndarray)
80                 * actions: (numpy.ndarray)
81                 * rewards: (numpy.ndarray)
82                 * agent_infos: (dict)
83                 * env_infos: (dict)
84 
85         """
86         _active_workers = []
87         self._active_worker_ids = []
88         pbar = ProgBarCounter(num_samples)
89         completed_samples = 0
90         traj = []
91         updating_workers = []
92 
93         # update the policy params of each worker before sampling
94         # for the current iteration
95         curr_policy_params = self.algo.policy.get_param_values()
96         params_id = ray.put(curr_policy_params)
97         while self._idle_worker_ids:
98             worker_id = self._idle_worker_ids.pop()
99             worker = self._all_workers[worker_id]
100             updating_workers.append(worker.set_agent.remote(params_id))
101 
102         while completed_samples < num_samples:
103             # if there are workers still being updated, check
104             # which ones are still updating and take the workers that
105             # are done updating, and start collecting trajectories on
106             # those workers.
107             if updating_workers:
108                 updated, updating_workers = ray.wait(updating_workers,
109                                                      num_returns=1,
110                                                      timeout=0.1)
111                 upd = [ray.get(up) for up in updated]
112                 self._idle_worker_ids.extend(upd)
113 
114             # if there are idle workers, use them to collect trajectories
115             # mark the newly busy workers as active
116             while self._idle_worker_ids:
117                 idle_worker_id = self._idle_worker_ids.pop()
118                 self._active_worker_ids.append(idle_worker_id)
119                 worker = self._all_workers[idle_worker_id]
120                 _active_workers.append(worker.rollout.remote())
121 
122             # check which workers are done/not done collecting a sample
123             # if any are done, send them to process the collected trajectory
124             # if they are not, keep checking if they are done
125             ready, not_ready = ray.wait(_active_workers,
126                                         num_returns=1,
127                                         timeout=0.001)
128             _active_workers = not_ready
129             for result in ready:
130                 trajectory, num_returned_samples = self._process_trajectory(
131                     result)
132                 completed_samples += num_returned_samples
133                 pbar.inc(num_returned_samples)
134                 traj.append(trajectory)
135         pbar.stop()
136         return traj
137 
138     def shutdown_worker(self):
139         """Shuts down the worker."""
140         ray.shutdown()
141 
142     def _process_trajectory(self, result):
143         """Collect trajectory from ray object store.
144 
145         Converts that trajectory to garage friendly format.
146 
147         Args:
148             result (obj): Ray object id of ready to be collected trajectory.
149 
150         Returns:
151             dict: One trajectory, with keys
152                 * observations: (numpy.ndarray)
153                 * actions: (numpy.ndarray)
154                 * rewards: (numpy.ndarray)
155                 * agent_infos: (dict)
156                 * env_infos: (dict)
157             int: Number of returned samples in the trajectory
158 
159         """
160         trajectory = ray.get(result)
161         ready_worker_id = trajectory[0]
162         self._active_worker_ids.remove(ready_worker_id)
163         self._idle_worker_ids.append(ready_worker_id)
164         trajectory = dict(observations=np.asarray(trajectory[1]),
165                           actions=np.asarray(trajectory[2]),
166                           rewards=np.asarray(trajectory[3]),
167                           agent_infos=trajectory[4],
168                           env_infos=trajectory[5])
169         num_returned_samples = len(trajectory['observations'])
170         return trajectory, num_returned_samples
171 
172 
173 class SamplerWorker:
174     """Constructs a single sampler worker.
175 
176     The worker can have its parameters updated, and sampler its policy for
177     trajectories or rollouts.
178 
179     Args:
180         worker_id (int): the id of the sampler_worker
181         env_pkl (bytes): A pickled gym or akro env object
182         agent_pkl (bytes): A pickled agent
183         seed (int): Random seed.
184         max_path_length (int): max trajectory length
185         should_render (bool): if true, renders trajectories after
186             sampling them
187 
188     """
189 
190     def __init__(self,
191                  worker_id,
192                  env_pkl,
193                  agent_pkl,
194                  seed,
195                  max_path_length,
196                  should_render=False):
197         self.worker_id = worker_id
198         self._env = pickle.loads(env_pkl)
199         self.agent = pickle.loads(agent_pkl)
200         self._seed = seed + self.worker_id
201         deterministic.set_seed(self._seed)
202         self._max_path_length = max_path_length
203         self._should_render = should_render
204         self.agent_updates = 0
205 
206     def set_agent(self, flattened_params):
207         """Set the agent params.
208 
209         Args:
210             flattened_params (list[np.ndarray]): model parameters
211 
212         Returns:
213             int: Worker id of this sampler worker.
214 
215         """
216         self.agent.set_param_values(flattened_params)
217         self.agent_updates += 1
218         return self.worker_id
219 
220     def rollout(self):
221         """Sample a single rollout from the agent/policy.
222 
223         The following value for the following keys will be a 2D array,
224         with the first dimension corresponding to the time dimension.
225 
226         - observations
227         - actions
228         - rewards
229         - next_observations
230         - terminals
231         The next two elements will be lists of dictionaries, with
232         the index into the list being the index into the time
233         - agent_infos
234         - env_infos
235 
236         Returns:
237             int: ID of this work
238             numpy.ndarray: observations
239             numpy.ndarray: actions
240             numpy.ndarray: rewards
241             dict[list]: agent info
242             dict[list]: environment info
243 
244         """
245         observations = []
246         actions = []
247         rewards = []
248         agent_infos = defaultdict(list)
249         env_infos = defaultdict(list)
250         o = self._env.reset()
251         self.agent.reset()
252         next_o = None
253         path_length = 0
254         while path_length < self._max_path_length:
255             a, agent_info = self.agent.get_action(o)
256             next_o, r, d, env_info = self._env.step(a)
257             observations.append(o)
258             rewards.append(r)
259             actions.append(a)
260             for k, v in agent_info.items():
261                 agent_infos[k].append(v)
262             for k, v in env_info.items():
263                 env_infos[k].append(v)
264             path_length += 1
265             if d:
266                 break
267             o = next_o
268         for k, v in agent_infos.items():
269             agent_infos[k] = np.asarray(v)
270         for k, v in env_infos.items():
271             env_infos[k] = np.asarray(v)
272         return self.worker_id,\
273             np.array(observations),\
274             np.array(actions),\
275             np.array(rewards),\
276             dict(agent_infos),\
277             dict(env_infos)
278
```
--- 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/src/garage/sampler/ray_sampler.py b/src/garage/sampler/ray_sampler.py
--- a/src/garage/sampler/ray_sampler.py
+++ b/src/garage/sampler/ray_sampler.py
@@ -89,7 +89,7 @@
         completed_samples = 0
         traj = []
         updating_workers = []
-
+        samples_to_be_collected = 0
         # update the policy params of each worker before sampling
         # for the current iteration
         curr_policy_params = self.algo.policy.get_param_values()
@@ -98,7 +98,6 @@
             worker_id = self._idle_worker_ids.pop()
             worker = self._all_workers[worker_id]
             updating_workers.append(worker.set_agent.remote(params_id))
-
         while completed_samples < num_samples:
             # if there are workers still being updated, check
             # which ones are still updating and take the workers that
@@ -113,12 +112,20 @@
 
             # if there are idle workers, use them to collect trajectories
             # mark the newly busy workers as active
-            while self._idle_worker_ids:
+            workers_to_use = int(
+                np.clip(
+                    np.ceil(
+                        (num_samples - completed_samples -
+                         samples_to_be_collected) / self._max_path_length) -
+                    len(self._active_worker_ids), 0, len(self._all_workers)))
+            workers_started = 0
+            while self._idle_worker_ids and workers_started < workers_to_use:
                 idle_worker_id = self._idle_worker_ids.pop()
+                workers_started += 1
                 self._active_worker_ids.append(idle_worker_id)
+                samples_to_be_collected += self._max_path_length
                 worker = self._all_workers[idle_worker_id]
                 _active_workers.append(worker.rollout.remote())
-
             # check which workers are done/not done collecting a sample
             # if any are done, send them to process the collected trajectory
             # if they are not, keep checking if they are done
@@ -129,9 +136,12 @@
             for result in ready:
                 trajectory, num_returned_samples = self._process_trajectory(
                     result)
+                samples_to_be_collected -= self._max_path_length
                 completed_samples += num_returned_samples
                 pbar.inc(num_returned_samples)
                 traj.append(trajectory)
+        self._idle_worker_ids = list(range(self._num_workers))
+        self._active_worker_ids = []
         pbar.stop()
         return traj

{"golden_diff": "diff --git a/src/garage/sampler/ray_sampler.py b/src/garage/sampler/ray_sampler.py\n--- a/src/garage/sampler/ray_sampler.py\n+++ b/src/garage/sampler/ray_sampler.py\n@@ -89,7 +89,7 @@\n         completed_samples = 0\n         traj = []\n         updating_workers = []\n-\n+        samples_to_be_collected = 0\n         # update the policy params of each worker before sampling\n         # for the current iteration\n         curr_policy_params = self.algo.policy.get_param_values()\n@@ -98,7 +98,6 @@\n             worker_id = self._idle_worker_ids.pop()\n             worker = self._all_workers[worker_id]\n             updating_workers.append(worker.set_agent.remote(params_id))\n-\n         while completed_samples < num_samples:\n             # if there are workers still being updated, check\n             # which ones are still updating and take the workers that\n@@ -113,12 +112,20 @@\n \n             # if there are idle workers, use them to collect trajectories\n             # mark the newly busy workers as active\n-            while self._idle_worker_ids:\n+            workers_to_use = int(\n+                np.clip(\n+                    np.ceil(\n+                        (num_samples - completed_samples -\n+                         samples_to_be_collected) / self._max_path_length) -\n+                    len(self._active_worker_ids), 0, len(self._all_workers)))\n+            workers_started = 0\n+            while self._idle_worker_ids and workers_started < workers_to_use:\n                 idle_worker_id = self._idle_worker_ids.pop()\n+                workers_started += 1\n                 self._active_worker_ids.append(idle_worker_id)\n+                samples_to_be_collected += self._max_path_length\n                 worker = self._all_workers[idle_worker_id]\n                 _active_workers.append(worker.rollout.remote())\n-\n             # check which workers are done/not done collecting a sample\n             # if any are done, send them to process the collected trajectory\n             # if they are not, keep checking if they are done\n@@ -129,9 +136,12 @@\n             for result in ready:\n                 trajectory, num_returned_samples = self._process_trajectory(\n                     result)\n+                samples_to_be_collected -= self._max_path_length\n                 completed_samples += num_returned_samples\n                 pbar.inc(num_returned_samples)\n                 traj.append(trajectory)\n+        self._idle_worker_ids = list(range(self._num_workers))\n+        self._active_worker_ids = []\n         pbar.stop()\n         return traj\n", "issue": "RaySampler does not use all workers\nIn release 2019.10, [RaySampler](https://github.com/rlworkgroup/garage/blob/release-2019.10/src/garage/sampler/ray_sampler.py) sees only one of the workers as idle for all iterations after the first iteration. This can be seen with [examples/tf/trpo_swimmer_ray_sampler.py](https://github.com/rlworkgroup/garage/blob/release-2019.10/examples/tf/trpo_swimmer_ray_sampler.py). \r\n\r\nTo me it looks like `self._idle_worker_ids` is not correctly updated at the end of `obtain_samples()` [here](https://github.com/rlworkgroup/garage/blob/d6c993526c23c289ebc94635a78186d9bf197f32/src/garage/sampler/ray_sampler.py#L130). Only those workers which have returned with a result are put back in the idle pool. This is a problem because results from some workers are not processed at the end and these workers don't go back into the idle pool. \r\n\r\nA quick fix (?) is to reset the list of idle workers ids to all workers ids at the start of `obtain samples()`. This does result in all the workers being used for all iterations. Since the agent parameters are updated in the next iteration, the samples from the previously not ready workers must be valid. However, I am not totally sure. \r\n\r\n@avnishn @ryanjulian Appreciate it if the fix can be back-ported to 2019.10!\n", "before_files": [{"content": "\"\"\"This is an implementation of an on policy batch sampler.\n\nUses a data parallel design.\nIncluded is a sampler that deploys sampler workers.\n\nThe sampler workers must implement some type of set agent parameters\nfunction, and a rollout function\n\"\"\"\nfrom collections import defaultdict\nimport pickle\n\nimport numpy as np\nimport psutil\nimport ray\n\nfrom garage.experiment import deterministic\nfrom garage.misc.prog_bar_counter import ProgBarCounter\nfrom garage.sampler.base import BaseSampler\n\n\nclass RaySampler(BaseSampler):\n    \"\"\"Collects Policy Rollouts in a data parallel fashion.\n\n    Args:\n        algo (garage.np.algo.RLAlgorithm): A garage algo object\n        env (gym.Env): A gym/akro env object\n        seed (int): Random seed.\n        should_render (bool): If the sampler render the trajectories.\n        num_processors (int): Number of processors to be used.\n        sampler_worker_cls (garage.sampler.ray_sampler.SamplerWorker):\n            If none, uses the default SamplerWorker class\n\n    \"\"\"\n\n    def __init__(self,\n                 algo,\n                 env,\n                 seed,\n                 should_render=False,\n                 num_processors=None,\n                 sampler_worker_cls=None):\n        super().__init__(algo, env)\n        self._sampler_worker = ray.remote(SamplerWorker if sampler_worker_cls\n                                          is None else sampler_worker_cls)\n        self._seed = seed\n        deterministic.set_seed(self._seed)\n        self._max_path_length = self.algo.max_path_length\n        self._should_render = should_render\n        if not ray.is_initialized():\n            ray.init(log_to_driver=False)\n        self._num_workers = (num_processors if num_processors else\n                             psutil.cpu_count(logical=False))\n        self._all_workers = defaultdict(None)\n        self._idle_worker_ids = list(range(self._num_workers))\n        self._active_worker_ids = []\n\n    def start_worker(self):\n        \"\"\"Initialize a new ray worker.\"\"\"\n        # Pickle the environment once, instead of once per worker.\n        env_pkl = pickle.dumps(self.env)\n        # We need to pickle the agent so that we can e.g. set up the TF Session\n        # in the worker *before* unpicling it.\n        agent_pkl = pickle.dumps(self.algo.policy)\n        for worker_id in range(self._num_workers):\n            self._all_workers[worker_id] = self._sampler_worker.remote(\n                worker_id, env_pkl, agent_pkl, self._seed,\n                self._max_path_length, self._should_render)\n\n    # pylint: disable=arguments-differ\n    def obtain_samples(self, itr, num_samples):\n        \"\"\"Sample the policy for new trajectories.\n\n        Args:\n            itr (int): Iteration number.\n            num_samples (int): Number of steps the the sampler should collect.\n\n        Returns:\n            list[dict]: Sample paths, each path with key\n                * observations: (numpy.ndarray)\n                * actions: (numpy.ndarray)\n                * rewards: (numpy.ndarray)\n                * agent_infos: (dict)\n                * env_infos: (dict)\n\n        \"\"\"\n        _active_workers = []\n        self._active_worker_ids = []\n        pbar = ProgBarCounter(num_samples)\n        completed_samples = 0\n        traj = []\n        updating_workers = []\n\n        # update the policy params of each worker before sampling\n        # for the current iteration\n        curr_policy_params = self.algo.policy.get_param_values()\n        params_id = ray.put(curr_policy_params)\n        while self._idle_worker_ids:\n            worker_id = self._idle_worker_ids.pop()\n            worker = self._all_workers[worker_id]\n            updating_workers.append(worker.set_agent.remote(params_id))\n\n        while completed_samples < num_samples:\n            # if there are workers still being updated, check\n            # which ones are still updating and take the workers that\n            # are done updating, and start collecting trajectories on\n            # those workers.\n            if updating_workers:\n                updated, updating_workers = ray.wait(updating_workers,\n                                                     num_returns=1,\n                                                     timeout=0.1)\n                upd = [ray.get(up) for up in updated]\n                self._idle_worker_ids.extend(upd)\n\n            # if there are idle workers, use them to collect trajectories\n            # mark the newly busy workers as active\n            while self._idle_worker_ids:\n                idle_worker_id = self._idle_worker_ids.pop()\n                self._active_worker_ids.append(idle_worker_id)\n                worker = self._all_workers[idle_worker_id]\n                _active_workers.append(worker.rollout.remote())\n\n            # check which workers are done/not done collecting a sample\n            # if any are done, send them to process the collected trajectory\n            # if they are not, keep checking if they are done\n            ready, not_ready = ray.wait(_active_workers,\n                                        num_returns=1,\n                                        timeout=0.001)\n            _active_workers = not_ready\n            for result in ready:\n                trajectory, num_returned_samples = self._process_trajectory(\n                    result)\n                completed_samples += num_returned_samples\n                pbar.inc(num_returned_samples)\n                traj.append(trajectory)\n        pbar.stop()\n        return traj\n\n    def shutdown_worker(self):\n        \"\"\"Shuts down the worker.\"\"\"\n        ray.shutdown()\n\n    def _process_trajectory(self, result):\n        \"\"\"Collect trajectory from ray object store.\n\n        Converts that trajectory to garage friendly format.\n\n        Args:\n            result (obj): Ray object id of ready to be collected trajectory.\n\n        Returns:\n            dict: One trajectory, with keys\n                * observations: (numpy.ndarray)\n                * actions: (numpy.ndarray)\n                * rewards: (numpy.ndarray)\n                * agent_infos: (dict)\n                * env_infos: (dict)\n            int: Number of returned samples in the trajectory\n\n        \"\"\"\n        trajectory = ray.get(result)\n        ready_worker_id = trajectory[0]\n        self._active_worker_ids.remove(ready_worker_id)\n        self._idle_worker_ids.append(ready_worker_id)\n        trajectory = dict(observations=np.asarray(trajectory[1]),\n                          actions=np.asarray(trajectory[2]),\n                          rewards=np.asarray(trajectory[3]),\n                          agent_infos=trajectory[4],\n                          env_infos=trajectory[5])\n        num_returned_samples = len(trajectory['observations'])\n        return trajectory, num_returned_samples\n\n\nclass SamplerWorker:\n    \"\"\"Constructs a single sampler worker.\n\n    The worker can have its parameters updated, and sampler its policy for\n    trajectories or rollouts.\n\n    Args:\n        worker_id (int): the id of the sampler_worker\n        env_pkl (bytes): A pickled gym or akro env object\n        agent_pkl (bytes): A pickled agent\n        seed (int): Random seed.\n        max_path_length (int): max trajectory length\n        should_render (bool): if true, renders trajectories after\n            sampling them\n\n    \"\"\"\n\n    def __init__(self,\n                 worker_id,\n                 env_pkl,\n                 agent_pkl,\n                 seed,\n                 max_path_length,\n                 should_render=False):\n        self.worker_id = worker_id\n        self._env = pickle.loads(env_pkl)\n        self.agent = pickle.loads(agent_pkl)\n        self._seed = seed + self.worker_id\n        deterministic.set_seed(self._seed)\n        self._max_path_length = max_path_length\n        self._should_render = should_render\n        self.agent_updates = 0\n\n    def set_agent(self, flattened_params):\n        \"\"\"Set the agent params.\n\n        Args:\n            flattened_params (list[np.ndarray]): model parameters\n\n        Returns:\n            int: Worker id of this sampler worker.\n\n        \"\"\"\n        self.agent.set_param_values(flattened_params)\n        self.agent_updates += 1\n        return self.worker_id\n\n    def rollout(self):\n        \"\"\"Sample a single rollout from the agent/policy.\n\n        The following value for the following keys will be a 2D array,\n        with the first dimension corresponding to the time dimension.\n\n        - observations\n        - actions\n        - rewards\n        - next_observations\n        - terminals\n        The next two elements will be lists of dictionaries, with\n        the index into the list being the index into the time\n        - agent_infos\n        - env_infos\n\n        Returns:\n            int: ID of this work\n            numpy.ndarray: observations\n            numpy.ndarray: actions\n            numpy.ndarray: rewards\n            dict[list]: agent info\n            dict[list]: environment info\n\n        \"\"\"\n        observations = []\n        actions = []\n        rewards = []\n        agent_infos = defaultdict(list)\n        env_infos = defaultdict(list)\n        o = self._env.reset()\n        self.agent.reset()\n        next_o = None\n        path_length = 0\n        while path_length < self._max_path_length:\n            a, agent_info = self.agent.get_action(o)\n            next_o, r, d, env_info = self._env.step(a)\n            observations.append(o)\n            rewards.append(r)\n            actions.append(a)\n            for k, v in agent_info.items():\n                agent_infos[k].append(v)\n            for k, v in env_info.items():\n                env_infos[k].append(v)\n            path_length += 1\n            if d:\n                break\n            o = next_o\n        for k, v in agent_infos.items():\n            agent_infos[k] = np.asarray(v)\n        for k, v in env_infos.items():\n            env_infos[k] = np.asarray(v)\n        return self.worker_id,\\\n            np.array(observations),\\\n            np.array(actions),\\\n            np.array(rewards),\\\n            dict(agent_infos),\\\n            dict(env_infos)\n", "path": "src/garage/sampler/ray_sampler.py"}], "after_files": [{"content": "\"\"\"This is an implementation of an on policy batch sampler.\n\nUses a data parallel design.\nIncluded is a sampler that deploys sampler workers.\n\nThe sampler workers must implement some type of set agent parameters\nfunction, and a rollout function\n\"\"\"\nfrom collections import defaultdict\nimport pickle\n\nimport numpy as np\nimport psutil\nimport ray\n\nfrom garage.experiment import deterministic\nfrom garage.misc.prog_bar_counter import ProgBarCounter\nfrom garage.sampler.base import BaseSampler\n\n\nclass RaySampler(BaseSampler):\n    \"\"\"Collects Policy Rollouts in a data parallel fashion.\n\n    Args:\n        algo (garage.np.algo.RLAlgorithm): A garage algo object\n        env (gym.Env): A gym/akro env object\n        seed (int): Random seed.\n        should_render (bool): If the sampler render the trajectories.\n        num_processors (int): Number of processors to be used.\n        sampler_worker_cls (garage.sampler.ray_sampler.SamplerWorker):\n            If none, uses the default SamplerWorker class\n\n    \"\"\"\n\n    def __init__(self,\n                 algo,\n                 env,\n                 seed,\n                 should_render=False,\n                 num_processors=None,\n                 sampler_worker_cls=None):\n        super().__init__(algo, env)\n        self._sampler_worker = ray.remote(SamplerWorker if sampler_worker_cls\n                                          is None else sampler_worker_cls)\n        self._seed = seed\n        deterministic.set_seed(self._seed)\n        self._max_path_length = self.algo.max_path_length\n        self._should_render = should_render\n        if not ray.is_initialized():\n            ray.init(log_to_driver=False)\n        self._num_workers = (num_processors if num_processors else\n                             psutil.cpu_count(logical=False))\n        self._all_workers = defaultdict(None)\n        self._idle_worker_ids = list(range(self._num_workers))\n        self._active_worker_ids = []\n\n    def start_worker(self):\n        \"\"\"Initialize a new ray worker.\"\"\"\n        # Pickle the environment once, instead of once per worker.\n        env_pkl = pickle.dumps(self.env)\n        # We need to pickle the agent so that we can e.g. set up the TF Session\n        # in the worker *before* unpicling it.\n        agent_pkl = pickle.dumps(self.algo.policy)\n        for worker_id in range(self._num_workers):\n            self._all_workers[worker_id] = self._sampler_worker.remote(\n                worker_id, env_pkl, agent_pkl, self._seed,\n                self._max_path_length, self._should_render)\n\n    # pylint: disable=arguments-differ\n    def obtain_samples(self, itr, num_samples):\n        \"\"\"Sample the policy for new trajectories.\n\n        Args:\n            itr (int): Iteration number.\n            num_samples (int): Number of steps the the sampler should collect.\n\n        Returns:\n            list[dict]: Sample paths, each path with key\n                * observations: (numpy.ndarray)\n                * actions: (numpy.ndarray)\n                * rewards: (numpy.ndarray)\n                * agent_infos: (dict)\n                * env_infos: (dict)\n\n        \"\"\"\n        _active_workers = []\n        self._active_worker_ids = []\n        pbar = ProgBarCounter(num_samples)\n        completed_samples = 0\n        traj = []\n        updating_workers = []\n        samples_to_be_collected = 0\n        # update the policy params of each worker before sampling\n        # for the current iteration\n        curr_policy_params = self.algo.policy.get_param_values()\n        params_id = ray.put(curr_policy_params)\n        while self._idle_worker_ids:\n            worker_id = self._idle_worker_ids.pop()\n            worker = self._all_workers[worker_id]\n            updating_workers.append(worker.set_agent.remote(params_id))\n        while completed_samples < num_samples:\n            # if there are workers still being updated, check\n            # which ones are still updating and take the workers that\n            # are done updating, and start collecting trajectories on\n            # those workers.\n            if updating_workers:\n                updated, updating_workers = ray.wait(updating_workers,\n                                                     num_returns=1,\n                                                     timeout=0.1)\n                upd = [ray.get(up) for up in updated]\n                self._idle_worker_ids.extend(upd)\n\n            # if there are idle workers, use them to collect trajectories\n            # mark the newly busy workers as active\n            workers_to_use = int(\n                np.clip(\n                    np.ceil(\n                        (num_samples - completed_samples -\n                         samples_to_be_collected) / self._max_path_length) -\n                    len(self._active_worker_ids), 0, len(self._all_workers)))\n            workers_started = 0\n            while self._idle_worker_ids and workers_started < workers_to_use:\n                idle_worker_id = self._idle_worker_ids.pop()\n                workers_started += 1\n                self._active_worker_ids.append(idle_worker_id)\n                samples_to_be_collected += self._max_path_length\n                worker = self._all_workers[idle_worker_id]\n                _active_workers.append(worker.rollout.remote())\n            # check which workers are done/not done collecting a sample\n            # if any are done, send them to process the collected trajectory\n            # if they are not, keep checking if they are done\n            ready, not_ready = ray.wait(_active_workers,\n                                        num_returns=1,\n                                        timeout=0.001)\n            _active_workers = not_ready\n            for result in ready:\n                trajectory, num_returned_samples = self._process_trajectory(\n                    result)\n                samples_to_be_collected -= self._max_path_length\n                completed_samples += num_returned_samples\n                pbar.inc(num_returned_samples)\n                traj.append(trajectory)\n        self._idle_worker_ids = list(range(self._num_workers))\n        self._active_worker_ids = []\n        pbar.stop()\n        return traj\n\n    def shutdown_worker(self):\n        \"\"\"Shuts down the worker.\"\"\"\n        ray.shutdown()\n\n    def _process_trajectory(self, result):\n        \"\"\"Collect trajectory from ray object store.\n\n        Converts that trajectory to garage friendly format.\n\n        Args:\n            result (obj): Ray object id of ready to be collected trajectory.\n\n        Returns:\n            dict: One trajectory, with keys\n                * observations: (numpy.ndarray)\n                * actions: (numpy.ndarray)\n                * rewards: (numpy.ndarray)\n                * agent_infos: (dict)\n                * env_infos: (dict)\n            int: Number of returned samples in the trajectory\n\n        \"\"\"\n        trajectory = ray.get(result)\n        ready_worker_id = trajectory[0]\n        self._active_worker_ids.remove(ready_worker_id)\n        self._idle_worker_ids.append(ready_worker_id)\n        trajectory = dict(observations=np.asarray(trajectory[1]),\n                          actions=np.asarray(trajectory[2]),\n                          rewards=np.asarray(trajectory[3]),\n                          agent_infos=trajectory[4],\n                          env_infos=trajectory[5])\n        num_returned_samples = len(trajectory['observations'])\n        return trajectory, num_returned_samples\n\n\nclass SamplerWorker:\n    \"\"\"Constructs a single sampler worker.\n\n    The worker can have its parameters updated, and sampler its policy for\n    trajectories or rollouts.\n\n    Args:\n        worker_id (int): the id of the sampler_worker\n        env_pkl (bytes): A pickled gym or akro env object\n        agent_pkl (bytes): A pickled agent\n        seed (int): Random seed.\n        max_path_length (int): max trajectory length\n        should_render (bool): if true, renders trajectories after\n            sampling them\n\n    \"\"\"\n\n    def __init__(self,\n                 worker_id,\n                 env_pkl,\n                 agent_pkl,\n                 seed,\n                 max_path_length,\n                 should_render=False):\n        self.worker_id = worker_id\n        self._env = pickle.loads(env_pkl)\n        self.agent = pickle.loads(agent_pkl)\n        self._seed = seed + self.worker_id\n        deterministic.set_seed(self._seed)\n        self._max_path_length = max_path_length\n        self._should_render = should_render\n        self.agent_updates = 0\n\n    def set_agent(self, flattened_params):\n        \"\"\"Set the agent params.\n\n        Args:\n            flattened_params (list[np.ndarray]): model parameters\n\n        Returns:\n            int: Worker id of this sampler worker.\n\n        \"\"\"\n        self.agent.set_param_values(flattened_params)\n        self.agent_updates += 1\n        return self.worker_id\n\n    def rollout(self):\n        \"\"\"Sample a single rollout from the agent/policy.\n\n        The following value for the following keys will be a 2D array,\n        with the first dimension corresponding to the time dimension.\n\n        - observations\n        - actions\n        - rewards\n        - next_observations\n        - terminals\n        The next two elements will be lists of dictionaries, with\n        the index into the list being the index into the time\n        - agent_infos\n        - env_infos\n\n        Returns:\n            int: ID of this work\n            numpy.ndarray: observations\n            numpy.ndarray: actions\n            numpy.ndarray: rewards\n            dict[list]: agent info\n            dict[list]: environment info\n\n        \"\"\"\n        observations = []\n        actions = []\n        rewards = []\n        agent_infos = defaultdict(list)\n        env_infos = defaultdict(list)\n        o = self._env.reset()\n        self.agent.reset()\n        next_o = None\n        path_length = 0\n        while path_length < self._max_path_length:\n            a, agent_info = self.agent.get_action(o)\n            next_o, r, d, env_info = self._env.step(a)\n            observations.append(o)\n            rewards.append(r)\n            actions.append(a)\n            for k, v in agent_info.items():\n                agent_infos[k].append(v)\n            for k, v in env_info.items():\n                env_infos[k].append(v)\n            path_length += 1\n            if d:\n                break\n            o = next_o\n        for k, v in agent_infos.items():\n            agent_infos[k] = np.asarray(v)\n        for k, v in env_infos.items():\n            env_infos[k] = np.asarray(v)\n        return self.worker_id,\\\n            np.array(observations),\\\n            np.array(actions),\\\n            np.array(rewards),\\\n            dict(agent_infos),\\\n            dict(env_infos)\n", "path": "src/garage/sampler/ray_sampler.py"}]}