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 | from collections import defaultdict
from collections.abc import Callable
import multiprocessing as mp
from typing import Optional, Union
import numpy as np
import torch
from lit_sb3.core.evaluation import evaluate_policy
from lit_sb3.core.running_mean_std import RunningMeanStd
from lit_sb3.policies import BasePolicy
from lit_sb3.vec_env import VecEnv, unwrap_vec_normalize
from lit_sb3.vec_env.base_vec_env import CloudpickleWrapper
def _worker(
remote: mp.connection.Connection, # ty:ignore[possibly-missing-submodule]
parent_remote: mp.connection.Connection, # ty:ignore[possibly-missing-submodule]
worker_env_wrapper: CloudpickleWrapper,
train_policy_wrapper: CloudpickleWrapper,
n_eval_episodes: int = 1,
) -> None:
"""Function that will be run in each process.
It is in charge of creating environments, evaluating candidates
and communicating with the main process.
Args:
remote: Pipe to communicate with the parent process.
parent_remote:
worker_env_wrapper: Callable used to create the environment inside the process.
train_policy_wrapper: Callable used to create the policy inside the process.
n_eval_episodes: Number of evaluation episodes per candidate.
parent_remote: The parent remote parameter.
"""
parent_remote.close()
vec_env: VecEnv = worker_env_wrapper.var()
train_policy = train_policy_wrapper.var
vec_normalize = unwrap_vec_normalize(vec_env)
obs_rms = vec_normalize.obs_rms if vec_normalize is not None else None
while True:
try:
cmd, data = remote.recv()
if cmd == "eval":
results = []
# Evaluate each candidate and save results
for weights_idx, candidate_weights in data:
train_policy.load_from_vector(candidate_weights.cpu())
episode_rewards, episode_lengths = evaluate_policy(
train_policy,
vec_env,
n_eval_episodes=n_eval_episodes,
return_episode_rewards=True,
warn=False,
)
results.append((weights_idx, (episode_rewards, episode_lengths)))
remote.send(results)
elif cmd == "seed":
# Note: the seed will only be effective at the next reset
remote.send(vec_env.seed(seed=data))
elif cmd == "set_options":
# Note: the options will only be effective at the next reset
remote.send(vec_env.set_options(data))
elif cmd == "get_obs_rms":
remote.send(obs_rms)
elif cmd == "sync_obs_rms":
assert vec_normalize is not None, "Tried to call `sync_obs_rms` when not using VecNormalize"
vec_normalize.obs_rms = data
obs_rms = data
elif cmd == "close":
vec_env.close()
remote.close()
break
else:
raise NotImplementedError(f"`{cmd}` is not implemented in the worker")
except EOFError:
break
class AsyncEval:
"""Helper class to do asynchronous evaluation of different policies with multiple processes.
It is useful when implementing population based methods like Evolution Strategies (ES),
Cross Entropy Method (CEM) or Augmented Random Search (ARS).
.. warning::
Only 'forkserver' and 'spawn' start methods are thread-safe,
which is important to avoid race conditions.
However, compared to
'fork' they incur a small start-up cost and have restrictions on
global variables. With those methods, users must wrap the code in an
``if __name__ == "__main__":`` block.
For more information, see the multiprocessing documentation.
Args:
envs_fn: Vectorized environments to run in subprocesses (callable)
train_policy: The policy object that will load the different candidate
weights.
start_method: method used to start the subprocesses.
Must be one of the methods returned by ``multiprocessing.get_all_start_methods()``.
Defaults to 'forkserver' on available platforms, and 'spawn' otherwise.
n_eval_episodes: The number of episodes to test each agent
"""
def __init__(
self,
envs_fn: list[Callable[[], VecEnv]],
train_policy: BasePolicy,
start_method: str | None = None,
n_eval_episodes: int = 1,
) -> None:
self.waiting = False
self.closed = False
n_envs = len(envs_fn)
if start_method is None:
# Fork is not a thread safe method (see issue #217)
# but is more user friendly (does not require to wrap the code in
# a `if __name__ == "__main__":`)
forkserver_available = "forkserver" in mp.get_all_start_methods()
start_method = "forkserver" if forkserver_available else "spawn"
ctx = mp.get_context(start_method)
self.remotes, self.work_remotes = zip(*[ctx.Pipe() for _ in range(n_envs)], strict=False)
self.processes = []
for work_remote, remote, worker_env in zip(self.work_remotes, self.remotes, envs_fn, strict=False):
args = (
work_remote,
remote,
CloudpickleWrapper(worker_env),
CloudpickleWrapper(train_policy),
n_eval_episodes,
)
# daemon=True: if the main process crashes, we should not cause things to hang
process = ctx.Process(target=_worker, args=args, daemon=True) # type: ignore[attr-defined]
process.start()
self.processes.append(process)
work_remote.close()
def send_jobs(self, candidate_weights: torch.Tensor, pop_size: int) -> None:
"""Send jobs to the workers to evaluate new candidates.
Args:
candidate_weights: The weights to be evaluated.
:pop_size: The number of candidate (size of the population)
pop_size: The pop size parameter.
"""
jobs_per_worker = defaultdict(list)
for weights_idx in range(pop_size):
jobs_per_worker[weights_idx % len(self.remotes)].append((weights_idx, candidate_weights[weights_idx]))
for remote_idx, remote in enumerate(self.remotes):
remote.send(("eval", jobs_per_worker[remote_idx]))
self.waiting = True
def seed(self, seed: int | None = None) -> list[None | int]:
"""Seed the environments.
Args:
seed: The seed for the pseudo-random generators.
"""
if seed is None:
# Do nothing
return []
for idx, remote in enumerate(self.remotes):
remote.send(("seed", seed + idx))
return [remote.recv() for remote in self.remotes]
def set_options(self, options: list[dict] | dict | None = None) -> list[None | int]:
"""Set environment options for all environments.
If a dict is passed instead of a list, the same options will be used for all environments.
WARNING: Those options will only be passed to the environment at the next reset.
Args:
options: A dictionary of environment options to pass to each environment at the next reset.
"""
for remote in self.remotes:
remote.send(("set_options", options))
return [remote.recv() for remote in self.remotes]
def get_results(self) -> list[tuple[int, tuple[np.ndarray, np.ndarray]]]:
"""Retreive episode rewards and lengths from each worker.
for all candidates (there might be multiple candidates per worker).
Returns:
A list of tuples containing each candidate index and its
result (episodic reward and episode length)
"""
results = [remote.recv() for remote in self.remotes]
flat_results = [result for worker_results in results for result in worker_results]
self.waiting = False
return flat_results
def get_obs_rms(self) -> list[RunningMeanStd]:
"""Retrieve the observation filters (observation running mean std).
of each process, they will be combined in the main process.
Synchronisation is done afterward using ``sync_obs_rms()``.
Returns:
A list of ``RunningMeanStd`` objects (one per process)
"""
for remote in self.remotes:
remote.send(("get_obs_rms", None))
return [remote.recv() for remote in self.remotes]
def sync_obs_rms(self, obs_rms: RunningMeanStd) -> None:
"""Synchronise (and update) the observation filters.
(observation running mean std).
Args:
obs_rms: The updated ``RunningMeanStd`` to be used
by workers for normalizing observations.
"""
for remote in self.remotes:
remote.send(("sync_obs_rms", obs_rms))
def close(self) -> None:
"""Close the processes."""
if self.closed:
return
if self.waiting:
for remote in self.remotes:
remote.recv()
for remote in self.remotes:
remote.send(("close", None))
for process in self.processes:
process.join()
self.closed = True
|