class ParallelRolloutSampler(SamplerBase, Serializable):
""" Class for sampling from multiple environments in parallel """
def __init__(self,
env,
policy,
num_workers: int,
*,
min_rollouts: int = None,
min_steps: int = None,
show_progress_bar: bool = True,
seed: int = None):
"""
Constructor
:param env: environment to sample from
:param policy: policy to act in the environment (can also be an exploration strategy)
:param num_workers: number of parallel samplers
:param min_rollouts: minimum number of complete rollouts to sample
:param min_steps: minimum total number of steps to sample
:param show_progress_bar: it `True`, display a progress bar using `tqdm`
:param seed: seed value for the random number generators, pass `None` for no seeding
"""
Serializable._init(self, locals())
super().__init__(min_rollouts=min_rollouts, min_steps=min_steps)
self.env = env
self.policy = policy
self.show_progress_bar = show_progress_bar
# Set method to spawn if using cuda
if self.policy.device == 'cuda':
mp.set_start_method('spawn', force=True)
# Create parallel pool. We use one thread per env because it's easier.
self.pool = SamplerPool(num_workers)
# Set all rngs' seeds
if seed is not None:
self.set_seed(seed)
# Distribute environments. We use pickle to make sure a copy is created for n_envs=1
self.pool.invoke_all(_ps_init, pickle.dumps(self.env),
pickle.dumps(self.policy))
def set_seed(self, seed):
"""
Set a deterministic seed on all workers.
:param seed: seed value for the random number generators
"""
self.pool.set_seed(seed)
def reinit(self, env=None, policy=None):
"""
Re-initialize the sampler.
:param env: the environment which the policy operates
:param policy: the policy used for sampling
"""
# Update env and policy if passed
if env is not None:
self.env = env
if policy is not None:
self.policy = policy
# Always broadcast to workers
self.pool.invoke_all(_ps_init, pickle.dumps(self.env),
pickle.dumps(self.policy))
def sample(self,
init_states: List[np.ndarray] = None,
domain_params: List[np.ndarray] = None,
eval: bool = False) -> List[StepSequence]:
"""
Do the sampling according to the previously given environment, policy, and number of steps/rollouts.
:param init_states: initial states forw `run_map()`, pass `None` (default) to sample from the environment's
initial state space
:param domain_params: domain parameters for `run_map()`, pass `None` (default) to not explicitly set them
:param eval: pass `False` if the rollout is executed during training, else `True`. Forwarded to `rollout()`.
:return: list of sampled rollouts
"""
# Update policy's state
self.pool.invoke_all(_ps_update_policy, self.policy.state_dict())
# Collect samples
with tqdm(leave=False,
file=sys.stdout,
desc='Sampling',
disable=(not self.show_progress_bar),
unit='steps'
if self.min_steps is not None else 'rollouts') as pb:
if self.min_steps is None:
if init_states is None and domain_params is None:
# Simply run min_rollouts times
func = partial(_ps_run_one, eval=eval)
arglist = range(self.min_rollouts)
elif init_states is not None and domain_params is None:
# Run every initial state so often that we at least get min_rollouts trajectories
func = partial(_ps_run_one_init_state, eval=eval)
rep_factor = ceil(self.min_rollouts / len(init_states))
arglist = rep_factor * init_states
elif init_states is not None and domain_params is not None:
# Run every combination of initial state and domain parameter so often that we at least get
# min_rollouts trajectories
func = partial(_ps_run_one_reset_kwargs, eval=eval)
allcombs = list(product(init_states, domain_params))
rep_factor = ceil(self.min_rollouts / len(allcombs))
arglist = rep_factor * allcombs
else:
raise NotImplementedError
# Only minimum number of rollouts given, thus use run_map
return self.pool.run_map(func, arglist, pb)
else:
# Minimum number of steps given, thus use run_collect (automatically handles min_runs=None)
if init_states is None:
return self.pool.run_collect(self.min_steps,
partial(_ps_sample_one,
eval=eval),
collect_progressbar=pb,
min_runs=self.min_rollouts)[0]
else:
raise NotImplementedError
policy_list.append(policy)
# Fix initial state (set to None if it should not be fixed)
init_state_list = [None]*args.num_ro_per_config
# Crate empty data frame
df = pd.DataFrame(columns=['policy', 'ret', 'len'])
# Evaluate all policies
for i, (env_sim, policy) in enumerate(zip(env_sim_list, policy_list)):
# Create a new sampler pool for every policy to synchronize the random seeds i.e. init states
pool = SamplerPool(args.num_workers)
# Seed the sampler
if args.seed is not None:
pool.set_seed(args.seed)
print_cbt(f"Set the random number generators' seed to {args.seed}.", 'w')
else:
print_cbt('No seed was set', 'y')
# Add the same wrappers as during training
env = wrap_like_other_env(env, env_sim)
# Sample rollouts
ros = eval_randomized_domain(pool, env, pert, policy, init_state_list) # internally calls DomainRandWrapperLive
# Compute results metrics
rets = [ro.undiscounted_return() for ro in ros]
lengths = [float(ro.length) for ro in ros] # int values are not numeric in pandas
df = df.append(pd.DataFrame(dict(policy=ex_labels[i], ret=rets, len=lengths)), ignore_index=True)
def evaluate_policy(args, ex_dir):
"""Helper function to evaluate the policy from an experiment in the associated environment."""
env, policy, _ = load_experiment(ex_dir, args)
# Create multi-dim evaluation grid
param_spec = dict()
param_spec_dim = None
if isinstance(inner_env(env), BallOnPlateSim):
param_spec["ball_radius"] = np.linspace(0.02, 0.08, num=2, endpoint=True)
param_spec["ball_rolling_friction_coefficient"] = np.linspace(0.0295, 0.9, num=2, endpoint=True)
elif isinstance(inner_env(env), QQubeSwingUpSim):
eval_num = 200
# Use nominal values for all other parameters.
for param, nominal_value in env.get_nominal_domain_param().items():
param_spec[param] = nominal_value
# param_spec["gravity_const"] = np.linspace(5.0, 15.0, num=eval_num, endpoint=True)
param_spec["damping_pend_pole"] = np.linspace(0.0, 0.0001, num=eval_num, endpoint=True)
param_spec["damping_rot_pole"] = np.linspace(0.0, 0.0006, num=eval_num, endpoint=True)
param_spec_dim = 2
elif isinstance(inner_env(env), QBallBalancerSim):
# param_spec["gravity_const"] = np.linspace(7.91, 11.91, num=11, endpoint=True)
# param_spec["ball_mass"] = np.linspace(0.003, 0.3, num=11, endpoint=True)
# param_spec["ball_radius"] = np.linspace(0.01, 0.1, num=11, endpoint=True)
param_spec["plate_length"] = np.linspace(0.275, 0.275, num=11, endpoint=True)
param_spec["arm_radius"] = np.linspace(0.0254, 0.0254, num=11, endpoint=True)
# param_spec["load_inertia"] = np.linspace(5.2822e-5*0.5, 5.2822e-5*1.5, num=11, endpoint=True)
# param_spec["motor_inertia"] = np.linspace(4.6063e-7*0.5, 4.6063e-7*1.5, num=11, endpoint=True)
# param_spec["gear_ratio"] = np.linspace(60, 80, num=11, endpoint=True)
# param_spec["gear_efficiency"] = np.linspace(0.6, 1.0, num=11, endpoint=True)
# param_spec["motor_efficiency"] = np.linspace(0.49, 0.89, num=11, endpoint=True)
# param_spec["motor_back_emf"] = np.linspace(0.006, 0.066, num=11, endpoint=True)
# param_spec["motor_resistance"] = np.linspace(2.6*0.5, 2.6*1.5, num=11, endpoint=True)
# param_spec["combined_damping"] = np.linspace(0.0, 0.05, num=11, endpoint=True)
# param_spec["friction_coeff"] = np.linspace(0, 0.015, num=11, endpoint=True)
# param_spec["voltage_thold_x_pos"] = np.linspace(0.0, 1.0, num=11, endpoint=True)
# param_spec["voltage_thold_x_neg"] = np.linspace(-1., 0.0, num=11, endpoint=True)
# param_spec["voltage_thold_y_pos"] = np.linspace(0.0, 1.0, num=11, endpoint=True)
# param_spec["voltage_thold_y_neg"] = np.linspace(-1.0, 0, num=11, endpoint=True)
# param_spec["offset_th_x"] = np.linspace(-5/180*np.pi, 5/180*np.pi, num=11, endpoint=True)
# param_spec["offset_th_y"] = np.linspace(-5/180*np.pi, 5/180*np.pi, num=11, endpoint=True)
else:
raise NotImplementedError
# Always add an action delay wrapper (with 0 delay by default)
if typed_env(env, ActDelayWrapper) is None:
env = ActDelayWrapper(env)
# param_spec['act_delay'] = np.linspace(0, 30, num=11, endpoint=True, dtype=int)
add_info = "-".join(param_spec.keys())
# Create multidimensional results grid and ensure right number of rollouts
param_list = param_grid(param_spec)
param_list *= args.num_rollouts_per_config
# Fix initial state (set to None if it should not be fixed)
init_state = np.array([0.0, 0.0, 0.0, 0.0])
# Create sampler
pool = SamplerPool(args.num_workers)
if args.seed is not None:
pool.set_seed(args.seed)
print_cbt(f"Set the random number generators' seed to {args.seed}.", "w")
else:
print_cbt("No seed was set", "y")
# Sample rollouts
ros = eval_domain_params(pool, env, policy, param_list, init_state)
# Compute metrics
lod = []
for ro in ros:
d = dict(**ro.rollout_info["domain_param"], ret=ro.undiscounted_return(), len=ro.length)
# Simply remove the observation noise from the domain parameters
try:
d.pop("obs_noise_mean")
d.pop("obs_noise_std")
except KeyError:
pass
lod.append(d)
df = pd.DataFrame(lod)
metrics = dict(
avg_len=df["len"].mean(),
avg_ret=df["ret"].mean(),
median_ret=df["ret"].median(),
min_ret=df["ret"].min(),
max_ret=df["ret"].max(),
std_ret=df["ret"].std(),
)
pprint(metrics, indent=4)
# Create subfolder and save
timestamp = datetime.datetime.now()
add_info = timestamp.strftime(pyrado.timestamp_format) + "--" + add_info
save_dir = osp.join(ex_dir, "eval_domain_grid", add_info)
os.makedirs(save_dir, exist_ok=True)
save_dicts_to_yaml(
{"ex_dir": str(ex_dir)},
{"varied_params": list(param_spec.keys())},
{"num_rpp": args.num_rollouts_per_config, "seed": args.seed},
{"metrics": dict_arraylike_to_float(metrics)},
save_dir=save_dir,
file_name="summary",
)
pyrado.save(df, f"df_sp_grid_{len(param_spec) if param_spec_dim is None else param_spec_dim}d.pkl", save_dir)
class ParameterExplorationSampler(Serializable):
""" Parallel sampler for parameter exploration """
def __init__(self,
env: Env,
policy: Policy,
num_workers: int,
num_rollouts_per_param: int,
seed: int = None):
"""
Constructor
:param env: environment to sample from
:param policy: policy used for sampling
:param num_workers: number of parallel samplers
:param num_rollouts_per_param: number of rollouts per policy parameter set (and init state if specified)
:param seed: seed value for the random number generators, pass `None` for no seeding
"""
if not isinstance(num_rollouts_per_param, int):
raise pyrado.TypeErr(given=num_rollouts_per_param,
expected_type=int)
if num_rollouts_per_param < 1:
raise pyrado.ValueErr(given=num_rollouts_per_param,
ge_constraint='1')
Serializable._init(self, locals())
# Check environment for domain randomization wrappers (stops after finding the outermost)
self._dr_wrapper = typed_env(env, DomainRandWrapper)
if self._dr_wrapper is not None:
assert isinstance(inner_env(env), SimEnv)
# Remove them all from the env chain since we sample the domain parameter later explicitly
env = remove_all_dr_wrappers(env)
self.env, self.policy = env, policy
self.num_rollouts_per_param = num_rollouts_per_param
# Create parallel pool. We use one thread per environment because it's easier.
self.pool = SamplerPool(num_workers)
# Set all rngs' seeds
if seed is not None:
self.pool.set_seed(seed)
# Distribute environments. We use pickle to make sure a copy is created for n_envs = 1
self.pool.invoke_all(_pes_init, pickle.dumps(self.env),
pickle.dumps(self.policy))
def _sample_domain_params(self) -> [list, dict]:
""" Sample domain parameters from the cached domain randomization wrapper. """
if self._dr_wrapper is None:
# No params
return [None] * self.num_rollouts_per_param
elif isinstance(self._dr_wrapper, DomainRandWrapperBuffer
) and self._dr_wrapper.buffer is not None:
# Use buffered param sets
idcs = np.random.randint(0,
len(self._dr_wrapper.buffer),
size=self.num_rollouts_per_param)
return [self._dr_wrapper.buffer[i] for i in idcs]
else:
# Sample new ones (same as in DomainRandWrapperBuffer.fill_buffer)
self._dr_wrapper.randomizer.randomize(self.num_rollouts_per_param)
return self._dr_wrapper.randomizer.get_params(-1,
format='list',
dtype='numpy')
def _sample_one_init_state(self, domain_param: dict) -> [np.ndarray, None]:
"""
Sample an init state for the given domain parameter set(s).
For some environments, the initial state space depends on the domain parameters, so we need to set them before
sampling it. We can just reset `self.env` here safely though, since it's not used for anything else.
:param domain_param: domain parameters to set
:return: initial state, `None` if no initial state space is defined
"""
self.env.reset(domain_param=domain_param)
ispace = attr_env_get(self.env, 'init_space')
if ispace is not None:
return ispace.sample_uniform()
else:
# No init space, no init state
return None
def sample(self, param_sets: to.Tensor) -> ParameterSamplingResult:
"""
Sample rollouts for a given set of parameters.
:param param_sets: sets of policy parameters
:return: data structure containing the policy parameter sets and the associated rollout data
"""
# Sample domain params for each rollout
domain_params = self._sample_domain_params()
if isinstance(domain_params, dict):
# There is only one domain parameter set (i.e. one init state)
init_states = [self._sample_one_init_state(domain_params)]
domain_params = [
domain_params
] # cast to list of dict to make iterable like the next case
elif isinstance(domain_params, list):
# There are more than one domain parameter set (i.e. multiple init states)
init_states = [
self._sample_one_init_state(dp) for dp in domain_params
]
else:
raise pyrado.TypeErr(given=domain_params,
expected_type=[list, dict])
# Explode parameter list for rollouts per param
all_params = [(p, *r) for p in param_sets
for r in zip(domain_params, init_states)]
# Sample rollouts in parallel
with tqdm(leave=False,
file=sys.stdout,
desc='Sampling',
unit='rollouts') as pb:
all_ros = self.pool.run_map(_pes_sample_one, all_params, pb)
# Group rollouts by parameters
ros_iter = iter(all_ros)
return ParameterSamplingResult([
ParameterSample(params=p,
rollouts=list(
itertools.islice(ros_iter,
self.num_rollouts_per_param)))
for p in param_sets
])
class ParallelSampler(SamplerBase, Serializable):
""" Class for sampling from multiple environments in parallel """
def __init__(self,
env,
policy,
num_envs: int,
*,
min_rollouts: int = None,
min_steps: int = None,
bernoulli_reset: bool = None,
seed: int = None):
"""
Constructor
:param env: environment to sample from
:param policy: policy to act in the environment (can also be an exploration strategy)
:param num_envs: number of parallel samplers
:param min_rollouts: minimum number of complete rollouts to sample.
:param min_steps: minimum total number of steps to sample.
:param bernoulli_reset: probability for resetting after the current time step
:param seed: Seed to use. Every subprocess is seeded with seed+thread_number
"""
Serializable._init(self, locals())
super().__init__(min_rollouts=min_rollouts, min_steps=min_steps)
self.env = env
self.policy = policy
self.bernoulli_reset = bernoulli_reset
# Set method to spawn if using cuda
if self.policy.device == 'cuda':
mp.set_start_method('spawn', force=True)
# Create parallel pool. We use one thread per env because it's easier.
self.pool = SamplerPool(num_envs)
if seed is not None:
self.pool.set_seed(seed)
# Distribute environments. We use pickle to make sure a copy is created for n_envs=1
self.pool.invoke_all(_ps_init, pickle.dumps(self.env),
pickle.dumps(self.policy),
pickle.dumps(self.bernoulli_reset))
def reinit(self, env=None, policy=None, bernoulli_reset: bool = None):
""" Re-initialize the sampler. """
# Update env and policy if passed
if env is not None:
self.env = env
if policy is not None:
self.policy = policy
if bernoulli_reset is not None:
self.bernoulli_reset = bernoulli_reset
# Always broadcast to workers
self.pool.invoke_all(_ps_init, pickle.dumps(self.env),
pickle.dumps(self.policy),
pickle.dumps(self.bernoulli_reset))
def sample(self) -> List[StepSequence]:
""" Do the sampling according to the previously given environment, policy, and number of steps/rollouts. """
# Update policy's state
self.pool.invoke_all(_ps_update_policy, self.policy.state_dict())
# Collect samples
with tqdm(leave=False,
file=sys.stdout,
desc='Sampling',
unit='steps'
if self.min_steps is not None else 'rollouts') as pb:
if self.min_steps is None:
# Only minimum number of rollouts given, thus use run_map
return self.pool.run_map(_ps_run_one, range(self.min_rollouts),
pb)
else:
# Minimum number of steps given, thus use run_collect (automatically handles min_runs=None)
return self.pool.run_collect(self.min_steps,
_ps_sample_one,
collect_progressbar=pb,
min_runs=self.min_rollouts)[0]
