-
Notifications
You must be signed in to change notification settings - Fork 714
/
parallel_py_environment.py
392 lines (329 loc) · 12.4 KB
/
parallel_py_environment.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
# coding=utf-8
# Copyright 2018 The TF-Agents Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Runs multiple environments in parallel processes and steps them in batch."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import atexit
import multiprocessing
import sys
import traceback
from absl import logging
import gin
import numpy as np
import tensorflow as tf
from tf_agents.environments import py_environment
from tf_agents.utils import nest_utils
@gin.configurable
class ParallelPyEnvironment(py_environment.PyEnvironment):
"""Batch together environments and simulate them in external processes.
The environments are created in external processes by calling the provided
callables. This can be an environment class, or a function creating the
environment and potentially wrapping it. The returned environment should not
access global variables.
"""
def __init__(self, env_constructors, start_serially=True, blocking=False,
flatten=False):
"""Batch together environments and simulate them in external processes.
The environments can be different but must use the same action and
observation specs.
Args:
env_constructors: List of callables that create environments.
start_serially: Whether to start environments serially or in parallel.
blocking: Whether to step environments one after another.
flatten: Boolean, whether to use flatten action and time_steps during
communication to reduce overhead.
Raises:
ValueError: If the action or observation specs don't match.
"""
super(ParallelPyEnvironment, self).__init__()
self._envs = [ProcessPyEnvironment(ctor, flatten=flatten)
for ctor in env_constructors]
self._num_envs = len(env_constructors)
self._blocking = blocking
self._start_serially = start_serially
self.start()
self._action_spec = self._envs[0].action_spec()
self._observation_spec = self._envs[0].observation_spec()
self._time_step_spec = self._envs[0].time_step_spec()
if any(env.action_spec() != self._action_spec for env in self._envs):
raise ValueError('All environments must have the same action spec.')
if any(env.time_step_spec() != self._time_step_spec for env in self._envs):
raise ValueError('All environments must have the same time_step_spec.')
self._flatten = flatten
def start(self):
logging.info('Spawning all processes.')
for env in self._envs:
env.start(wait_to_start=self._start_serially)
if not self._start_serially:
logging.info('Waiting for all processes to start.')
for env in self._envs:
env.wait_start()
logging.info('All processes started.')
@property
def batched(self):
return True
@property
def batch_size(self):
return self._num_envs
def observation_spec(self):
return self._observation_spec
def action_spec(self):
return self._action_spec
def time_step_spec(self):
return self._time_step_spec
def _reset(self):
"""Reset all environments and combine the resulting observation.
Returns:
Time step with batch dimension.
"""
time_steps = [env.reset(self._blocking) for env in self._envs]
if not self._blocking:
time_steps = [promise() for promise in time_steps]
return self._stack_time_steps(time_steps)
def _step(self, actions):
"""Forward a batch of actions to the wrapped environments.
Args:
actions: Batched action, possibly nested, to apply to the environment.
Raises:
ValueError: Invalid actions.
Returns:
Batch of observations, rewards, and done flags.
"""
time_steps = [
env.step(action, self._blocking)
for env, action in zip(self._envs, self._unstack_actions(actions))]
# When blocking is False we get promises that need to be called.
if not self._blocking:
time_steps = [promise() for promise in time_steps]
return self._stack_time_steps(time_steps)
def close(self):
"""Close all external process."""
logging.info('Closing all processes.')
for env in self._envs:
env.close()
logging.info('All processes closed.')
def _stack_time_steps(self, time_steps):
"""Given a list of TimeStep, combine to one with a batch dimension."""
if self._flatten:
return nest_utils.fast_map_structure_flatten(
lambda *arrays: np.stack(arrays), self._time_step_spec, *time_steps)
else:
return nest_utils.fast_map_structure(
lambda *arrays: np.stack(arrays), *time_steps)
def _unstack_actions(self, batched_actions):
"""Returns a list of actions from potentially nested batch of actions."""
flattened_actions = tf.nest.flatten(batched_actions)
if self._flatten:
unstacked_actions = zip(*flattened_actions)
else:
unstacked_actions = [
tf.nest.pack_sequence_as(batched_actions, actions)
for actions in zip(*flattened_actions)
]
return unstacked_actions
def seed(self, seeds):
"""Seeds the parallel environments."""
if len(seeds) != len(self._envs):
raise ValueError(
'Number of seeds should match the number of parallel_envs.')
promises = [env.call('seed', seed) for seed, env in zip(seeds, self._envs)]
# Block until all envs are seeded.
return [promise() for promise in promises]
class ProcessPyEnvironment(object):
"""Step a single env in a separate process for lock free paralellism."""
# Message types for communication via the pipe.
_READY = 1
_ACCESS = 2
_CALL = 3
_RESULT = 4
_EXCEPTION = 5
_CLOSE = 6
def __init__(self, env_constructor, flatten=False):
"""Step environment in a separate process for lock free paralellism.
The environment is created in an external process by calling the provided
callable. This can be an environment class, or a function creating the
environment and potentially wrapping it. The returned environment should
not access global variables.
Args:
env_constructor: Callable that creates and returns a Python environment.
flatten: Boolean, whether to assume flattened actions and time_steps
during communication to avoid overhead.
Attributes:
observation_spec: The cached observation spec of the environment.
action_spec: The cached action spec of the environment.
time_step_spec: The cached time step spec of the environment.
"""
self._env_constructor = env_constructor
self._flatten = flatten
self._observation_spec = None
self._action_spec = None
self._time_step_spec = None
def start(self, wait_to_start=True):
"""Start the process.
Args:
wait_to_start: Whether the call should wait for an env initialization.
"""
self._conn, conn = multiprocessing.Pipe()
self._process = multiprocessing.Process(
target=self._worker,
args=(conn, self._env_constructor, self._flatten))
atexit.register(self.close)
self._process.start()
if wait_to_start:
self.wait_start()
def wait_start(self):
"""Wait for the started process to finish initialization."""
result = self._conn.recv()
if isinstance(result, Exception):
self._conn.close()
self._process.join(5)
raise result
assert result == self._READY, result
def observation_spec(self):
if not self._observation_spec:
self._observation_spec = self.call('observation_spec')()
return self._observation_spec
def action_spec(self):
if not self._action_spec:
self._action_spec = self.call('action_spec')()
return self._action_spec
def time_step_spec(self):
if not self._time_step_spec:
self._time_step_spec = self.call('time_step_spec')()
return self._time_step_spec
def __getattr__(self, name):
"""Request an attribute from the environment.
Note that this involves communication with the external process, so it can
be slow.
Args:
name: Attribute to access.
Returns:
Value of the attribute.
"""
self._conn.send((self._ACCESS, name))
return self._receive()
def call(self, name, *args, **kwargs):
"""Asynchronously call a method of the external environment.
Args:
name: Name of the method to call.
*args: Positional arguments to forward to the method.
**kwargs: Keyword arguments to forward to the method.
Returns:
Promise object that blocks and provides the return value when called.
"""
payload = name, args, kwargs
self._conn.send((self._CALL, payload))
return self._receive
def close(self):
"""Send a close message to the external process and join it."""
try:
self._conn.send((self._CLOSE, None))
self._conn.close()
except IOError:
# The connection was already closed.
pass
self._process.join(5)
def step(self, action, blocking=True):
"""Step the environment.
Args:
action: The action to apply to the environment.
blocking: Whether to wait for the result.
Returns:
time step when blocking, otherwise callable that returns the time step.
"""
promise = self.call('step', action)
if blocking:
return promise()
else:
return promise
def reset(self, blocking=True):
"""Reset the environment.
Args:
blocking: Whether to wait for the result.
Returns:
New observation when blocking, otherwise callable that returns the new
observation.
"""
promise = self.call('reset')
if blocking:
return promise()
else:
return promise
def _receive(self):
"""Wait for a message from the worker process and return its payload.
Raises:
Exception: An exception was raised inside the worker process.
KeyError: The reveived message is of an unknown type.
Returns:
Payload object of the message.
"""
message, payload = self._conn.recv()
# Re-raise exceptions in the main process.
if message == self._EXCEPTION:
stacktrace = payload
raise Exception(stacktrace)
if message == self._RESULT:
return payload
self.close()
raise KeyError('Received message of unexpected type {}'.format(message))
def _worker(self, conn, env_constructor, flatten=False):
"""The process waits for actions and sends back environment results.
Args:
conn: Connection for communication to the main process.
env_constructor: env_constructor for the OpenAI Gym environment.
flatten: Boolean, whether to assume flattened actions and time_steps
during communication to avoid overhead.
Raises:
KeyError: When receiving a message of unknown type.
"""
try:
env = env_constructor()
action_spec = env.action_spec()
conn.send(self._READY) # Ready.
while True:
try:
# Only block for short times to have keyboard exceptions be raised.
if not conn.poll(0.1):
continue
message, payload = conn.recv()
except (EOFError, KeyboardInterrupt):
break
if message == self._ACCESS:
name = payload
result = getattr(env, name)
conn.send((self._RESULT, result))
continue
if message == self._CALL:
name, args, kwargs = payload
if flatten and name == 'step':
args = [tf.nest.pack_sequence_as(action_spec, args[0])]
result = getattr(env, name)(*args, **kwargs)
if flatten and name in ['step', 'reset']:
result = tf.nest.flatten(result)
conn.send((self._RESULT, result))
continue
if message == self._CLOSE:
assert payload is None
break
raise KeyError('Received message of unknown type {}'.format(message))
except Exception: # pylint: disable=broad-except
etype, evalue, tb = sys.exc_info()
stacktrace = ''.join(traceback.format_exception(etype, evalue, tb))
message = 'Error in environment process: {}'.format(stacktrace)
logging.error(message)
conn.send((self._EXCEPTION, stacktrace))
finally:
conn.close()