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]