Spaces:
Sleeping
Sleeping
| from collections import deque, OrderedDict | |
| from src.rlkit.core.eval_util import create_stats_ordered_dict | |
| from src.rlkit.samplers.rollout_functions import rollout, multitask_rollout, ensemble_rollout, ensemble_eval_rollout | |
| from src.rlkit.samplers.rollout_functions import ensemble_ucb_rollout | |
| from src.rlkit.samplers.data_collector.base import PathCollector | |
| class MdpPathCollector(PathCollector): | |
| def __init__( | |
| self, | |
| env, | |
| policy, | |
| noise_flag=0, | |
| max_num_epoch_paths_saved=None, | |
| render=False, | |
| render_kwargs=None, | |
| ): | |
| if render_kwargs is None: | |
| render_kwargs = {} | |
| self._env = env | |
| self._policy = policy | |
| self._max_num_epoch_paths_saved = max_num_epoch_paths_saved | |
| self._epoch_paths = deque(maxlen=self._max_num_epoch_paths_saved) | |
| self._render = render | |
| self._render_kwargs = render_kwargs | |
| self._noise_flag = noise_flag | |
| self._num_steps_total = 0 | |
| self._num_paths_total = 0 | |
| def collect_new_paths( | |
| self, | |
| max_path_length, | |
| num_steps, | |
| discard_incomplete_paths, | |
| ): | |
| paths = [] | |
| num_steps_collected = 0 | |
| while num_steps_collected < num_steps: | |
| max_path_length_this_loop = min( # Do not go over num_steps | |
| max_path_length, | |
| num_steps - num_steps_collected, | |
| ) | |
| path = rollout( | |
| self._env, | |
| self._policy, | |
| noise_flag=self._noise_flag, | |
| max_path_length=max_path_length_this_loop, | |
| ) | |
| path_len = len(path['actions']) | |
| if ( | |
| path_len != max_path_length | |
| and not path['terminals'][-1] | |
| and discard_incomplete_paths | |
| ): | |
| break | |
| num_steps_collected += path_len | |
| paths.append(path) | |
| self._num_paths_total += len(paths) | |
| self._num_steps_total += num_steps_collected | |
| self._epoch_paths.extend(paths) | |
| return paths | |
| def collect_normalized_new_paths( | |
| self, | |
| max_path_length, | |
| num_steps, | |
| discard_incomplete_paths, | |
| input_mean, | |
| input_std, | |
| ): | |
| paths = [] | |
| num_steps_collected = 0 | |
| while num_steps_collected < num_steps: | |
| max_path_length_this_loop = min( # Do not go over num_steps | |
| max_path_length, | |
| num_steps - num_steps_collected, | |
| ) | |
| path = normalized_rollout( | |
| self._env, | |
| self._policy, | |
| noise_flag=self._noise_flag, | |
| max_path_length=max_path_length_this_loop, | |
| input_mean=input_mean, | |
| input_std=input_std, | |
| ) | |
| path_len = len(path['actions']) | |
| if ( | |
| path_len != max_path_length | |
| and not path['terminals'][-1] | |
| and discard_incomplete_paths | |
| ): | |
| break | |
| num_steps_collected += path_len | |
| paths.append(path) | |
| self._num_paths_total += len(paths) | |
| self._num_steps_total += num_steps_collected | |
| self._epoch_paths.extend(paths) | |
| return paths | |
| def get_epoch_paths(self): | |
| return self._epoch_paths | |
| def end_epoch(self, epoch): | |
| self._epoch_paths = deque(maxlen=self._max_num_epoch_paths_saved) | |
| def get_diagnostics(self): | |
| path_lens = [len(path['actions']) for path in self._epoch_paths] | |
| stats = OrderedDict([ | |
| ('num steps total', self._num_steps_total), | |
| ('num paths total', self._num_paths_total), | |
| ]) | |
| stats.update(create_stats_ordered_dict( | |
| "path length", | |
| path_lens, | |
| always_show_all_stats=True, | |
| )) | |
| return stats | |
| def get_snapshot(self): | |
| return dict( | |
| env=self._env, | |
| policy=self._policy, | |
| ) | |
| class EnsembleMdpPathCollector(PathCollector): | |
| def __init__( | |
| self, | |
| env, | |
| policy, | |
| num_ensemble, | |
| noise_flag=0, | |
| ber_mean=0.5, | |
| eval_flag=False, | |
| max_num_epoch_paths_saved=None, | |
| render=False, | |
| render_kwargs=None, | |
| critic1=None, | |
| critic2=None, | |
| inference_type=0.0, | |
| feedback_type=1, | |
| ): | |
| if render_kwargs is None: | |
| render_kwargs = {} | |
| self._env = env | |
| self._policy = policy | |
| self._max_num_epoch_paths_saved = max_num_epoch_paths_saved | |
| self._epoch_paths = deque(maxlen=self._max_num_epoch_paths_saved) | |
| self._render = render | |
| self._render_kwargs = render_kwargs | |
| self.num_ensemble = num_ensemble | |
| self.eval_flag = eval_flag | |
| self.ber_mean = ber_mean | |
| self.critic1 = critic1 | |
| self.critic2 = critic2 | |
| self.inference_type = inference_type | |
| self.feedback_type = feedback_type | |
| self._noise_flag = noise_flag | |
| self._num_steps_total = 0 | |
| self._num_paths_total = 0 | |
| def collect_new_paths( | |
| self, | |
| max_path_length, | |
| num_steps, | |
| discard_incomplete_paths, | |
| ): | |
| paths = [] | |
| num_steps_collected = 0 | |
| while num_steps_collected < num_steps: | |
| max_path_length_this_loop = min( # Do not go over num_steps | |
| max_path_length, | |
| num_steps - num_steps_collected, | |
| ) | |
| if self.eval_flag: | |
| path = ensemble_eval_rollout( | |
| self._env, | |
| self._policy, | |
| self.num_ensemble, | |
| max_path_length=max_path_length_this_loop, | |
| ) | |
| else: | |
| if self.inference_type > 0: # UCB | |
| path = ensemble_ucb_rollout( | |
| self._env, | |
| self._policy, | |
| critic1=self.critic1, | |
| critic2=self.critic2, | |
| inference_type=self.inference_type, | |
| feedback_type=self.feedback_type, | |
| num_ensemble=self.num_ensemble, | |
| noise_flag=self._noise_flag, | |
| max_path_length=max_path_length_this_loop, | |
| ber_mean=self.ber_mean, | |
| ) | |
| else: | |
| path = ensemble_rollout( | |
| self._env, | |
| self._policy, | |
| self.num_ensemble, | |
| noise_flag=self._noise_flag, | |
| max_path_length=max_path_length_this_loop, | |
| ber_mean=self.ber_mean, | |
| ) | |
| path_len = len(path['actions']) | |
| if ( | |
| path_len != max_path_length | |
| and not path['terminals'][-1] | |
| and discard_incomplete_paths | |
| ): | |
| break | |
| num_steps_collected += path_len | |
| paths.append(path) | |
| self._num_paths_total += len(paths) | |
| self._num_steps_total += num_steps_collected | |
| self._epoch_paths.extend(paths) | |
| return paths | |
| def get_epoch_paths(self): | |
| return self._epoch_paths | |
| def end_epoch(self, epoch): | |
| self._epoch_paths = deque(maxlen=self._max_num_epoch_paths_saved) | |
| def get_diagnostics(self): | |
| path_lens = [len(path['actions']) for path in self._epoch_paths] | |
| stats = OrderedDict([ | |
| ('num steps total', self._num_steps_total), | |
| ('num paths total', self._num_paths_total), | |
| ]) | |
| stats.update(create_stats_ordered_dict( | |
| "path length", | |
| path_lens, | |
| always_show_all_stats=True, | |
| )) | |
| return stats | |
| def get_snapshot(self): | |
| return dict( | |
| env=self._env, | |
| policy=self._policy, | |
| ) | |
| class AsyncEnsembleMdpPathCollector(PathCollector): | |
| def __init__( | |
| self, | |
| env, | |
| policy, | |
| num_ensemble, | |
| noise_flag=0, | |
| ber_mean=0.5, | |
| eval_flag=False, | |
| max_num_epoch_paths_saved=None, | |
| render=False, | |
| render_kwargs=None, | |
| critic1=None, | |
| critic2=None, | |
| inference_type=0.0, | |
| feedback_type=1, | |
| ): | |
| if render_kwargs is None: | |
| render_kwargs = {} | |
| self._env = env | |
| self._policy = policy | |
| self._max_num_epoch_paths_saved = max_num_epoch_paths_saved | |
| self._epoch_paths = deque(maxlen=self._max_num_epoch_paths_saved) | |
| self._render = render | |
| self._render_kwargs = render_kwargs | |
| self.num_ensemble = num_ensemble | |
| self.eval_flag = eval_flag | |
| self.ber_mean = ber_mean | |
| self.critic1 = critic1 | |
| self.critic2 = critic2 | |
| self.inference_type = inference_type | |
| self.feedback_type = feedback_type | |
| self._noise_flag = noise_flag | |
| self._num_steps_total = 0 | |
| self._num_paths_total = 0 | |
| def collect_new_paths( | |
| self, | |
| max_path_length, | |
| num_steps, | |
| discard_incomplete_paths, | |
| ): | |
| paths = [] | |
| num_steps_collected = 0 | |
| while num_steps_collected < num_steps: | |
| max_path_length_this_loop = min( # Do not go over num_steps | |
| max_path_length, | |
| num_steps - num_steps_collected, | |
| ) | |
| if self.eval_flag: | |
| path = ensemble_eval_rollout( | |
| self._env, | |
| self._policy, | |
| self.num_ensemble, | |
| max_path_length=max_path_length_this_loop, | |
| ) | |
| else: | |
| if self.inference_type > 0: # UCB | |
| path = ensemble_ucb_rollout( | |
| self._env, | |
| self._policy, | |
| critic1=self.critic1, | |
| critic2=self.critic2, | |
| inference_type=self.inference_type, | |
| feedback_type=self.feedback_type, | |
| num_ensemble=self.num_ensemble, | |
| noise_flag=self._noise_flag, | |
| max_path_length=max_path_length_this_loop, | |
| ber_mean=self.ber_mean, | |
| ) | |
| else: | |
| path = ensemble_rollout( | |
| self._env, | |
| self._policy, | |
| self.num_ensemble, | |
| noise_flag=self._noise_flag, | |
| max_path_length=max_path_length_this_loop, | |
| ber_mean=self.ber_mean, | |
| ) | |
| path_len = len(path['actions']) | |
| if ( | |
| path_len != max_path_length | |
| and not path['terminals'][-1] | |
| and discard_incomplete_paths | |
| ): | |
| break | |
| num_steps_collected += path_len | |
| paths.append(path) | |
| self._num_paths_total += len(paths) | |
| self._num_steps_total += num_steps_collected | |
| self._epoch_paths.extend(paths) | |
| return paths | |