class BCValidationEvaluator(Evaluator): """Evaluates a given policy on validation data. Expects that the first two items returned in the dataset tuple are the observation_dict and action_dict. :param data_loader: The data used for evaluation. :param loss: Loss function to be used. :param model_selection: Model selection interface that will be notified of the recorded rewards. """ def __init__(self, loss: BCLoss, model_selection: Optional[ModelSelectionBase], data_loader: DataLoader, logging_prefix: Optional[str] = "eval"): self.loss = loss self.data_loader = data_loader self.model_selection = model_selection self.env = None if logging_prefix: self.eval_stats = LogStatsAggregator( LogStatsLevel.EPOCH, get_stats_logger(logging_prefix)) else: self.eval_stats = LogStatsAggregator(LogStatsLevel.EPOCH) self.eval_events = self.eval_stats.create_event_topic(ImitationEvents) @override(Evaluator) def evaluate(self, policy: TorchPolicy) -> None: """Evaluate given policy (results are stored in stat logs) and dump the model if the reward improved. :param policy: Policy to evaluate. """ policy.eval() with torch.no_grad(): total_loss = [] for iteration, data in enumerate(self.data_loader, 0): observations, actions, actor_ids = data[0], data[1], data[-1] actor_ids = debatch_actor_ids(actor_ids) # Convert only actions to torch, since observations are converted in # policy.compute_substep_policy_output method convert_to_torch(actions, device=policy.device, cast=None, in_place=True) total_loss.append( self.loss.calculate_loss(policy=policy, observations=observations, actions=actions, events=self.eval_events, actor_ids=actor_ids).item()) if self.model_selection: self.model_selection.update(-np.mean(total_loss).item())
class ESTrainer(Trainer): """Trainer class for OpenAI Evolution Strategies. :param algorithm_config: Algorithm parameters. :param torch_policy: Multi-step policy encapsulating the policy networks :param shared_noise: The noise table, with the same content for every worker and the master. :param normalization_stats: Normalization statistics as calculated by the NormalizeObservationWrapper. """ def __init__( self, algorithm_config: ESAlgorithmConfig, torch_policy: TorchPolicy, shared_noise: SharedNoiseTable, normalization_stats: Optional[Dict[str, Tuple[np.ndarray, np.ndarray]]] ) -> None: super().__init__(algorithm_config) # --- training setup --- self.model_selection: Optional[ModelSelectionBase] = None self.policy: Union[Policy, TorchModel] = torch_policy self.shared_noise = shared_noise self.normalization_stats = normalization_stats # setup the optimizer, now that the policy is available self.optimizer = Factory(Optimizer).instantiate( algorithm_config.optimizer) self.optimizer.setup(self.policy) # prepare statistics collection self.eval_stats = LogStatsAggregator(LogStatsLevel.EPOCH, get_stats_logger("eval")) self.train_stats = LogStatsAggregator(LogStatsLevel.EPOCH, get_stats_logger("train")) # injection of ES-specific events self.es_events = self.train_stats.create_event_topic(ESEvents) @override(Trainer) def train(self, distributed_rollouts: ESDistributedRollouts, n_epochs: Optional[int] = None, model_selection: Optional[ModelSelectionBase] = None) -> None: """ Run the ES training loop. :param distributed_rollouts: The distribution interface for experience collection. :param n_epochs: Number of epochs to train. :param model_selection: Optional model selection class, receives model evaluation results. """ n_epochs = self.algorithm_config.n_epochs if n_epochs is None else n_epochs self.model_selection = model_selection for epoch in itertools.count(): # check if we reached the max number of epochs if n_epochs and epoch == n_epochs: break print('********** Iteration {} **********'.format(epoch)) step_start_time = time.time() # do the actual update step (disable autograd, as we calculate the gradient from the rollout returns) with torch.no_grad(): self._update(distributed_rollouts) step_end_time = time.time() # log the step duration self.es_events.real_time(step_end_time - step_start_time) # update the epoch count increment_log_step() def load_state_dict(self, state_dict: Dict) -> None: """Set the model and optimizer state. :param state_dict: The state dict. """ self.policy.load_state_dict(state_dict) @override(Trainer) def state_dict(self): """implementation of :class:`~maze.train.trainers.common.trainer.Trainer` """ return self.policy.state_dict() @override(Trainer) def load_state(self, file_path: Union[str, BinaryIO]) -> None: """implementation of :class:`~maze.train.trainers.common.trainer.Trainer` """ state_dict = torch.load(file_path, map_location=torch.device(self.policy.device)) self.load_state_dict(state_dict) def _update(self, distributed_rollouts: ESDistributedRollouts): # Pop off results for the current task n_train_episodes, n_timesteps_popped = 0, 0 # aggregate all collected training rollouts for this episode epoch_results = ESRolloutResult(is_eval=False) # obtain a generator from the distribution interface rollouts_generator = distributed_rollouts.generate_rollouts( policy=self.policy, max_steps=self.algorithm_config.max_steps, noise_stddev=self.algorithm_config.noise_stddev, normalization_stats=self.normalization_stats) # collect eval and training rollouts for result in rollouts_generator: if result.is_eval: # This was an eval job for e in result.episode_stats: self.eval_stats.receive(e) continue # we received training experience from perturbed policy networks epoch_results.noise_indices.extend(result.noise_indices) epoch_results.episode_stats.extend(result.episode_stats) # update the training statistics for e in result.episode_stats: self.train_stats.receive(e) n_train_episodes += 1 n_timesteps_popped += e[(BaseEnvEvents.reward, "count", None)] # continue until we collected enough episodes and timesteps if (n_train_episodes >= self.algorithm_config.n_rollouts_per_update and n_timesteps_popped >= self.algorithm_config.n_timesteps_per_update): break # notify the model selection of the evaluation results eval_stats = self.eval_stats.reduce() if self.model_selection and len(eval_stats): reward = eval_stats[(BaseEnvEvents.reward, "mean", None)] self.model_selection.update(reward) # prepare returns, reshape the positive/negative antithetic estimation as (rollouts, 2) returns_n2 = np.array([ e[(BaseEnvEvents.reward, "sum", None)] for e in epoch_results.episode_stats ]).reshape(-1, 2) # improve robustness: weight by rank, not by reward proc_returns_n2 = self._compute_centered_ranks(returns_n2) # compute the gradient g = self._batched_weighted_sum( proc_returns_n2[:, 0] - proc_returns_n2[:, 1], (self.shared_noise.get(idx, self.policy.num_params) for idx in epoch_results.noise_indices), batch_size=500) g /= n_train_episodes / 2.0 # apply the weight update theta = get_flat_parameters(self.policy) update_ratio = self.optimizer.update(-g + self.algorithm_config.l2_penalty * theta.numpy()) # statistics logging self.es_events.update_ratio(update_ratio) for i in self.policy.state_dict().keys(): self.es_events.policy_grad_norm(policy_id=i, value=np.square(g).sum()**0.5) self.es_events.policy_norm(policy_id=i, value=np.square(theta).sum()**0.5) @classmethod def _iter_groups(cls, items: Iterable, group_size: int) -> Generator[Tuple, None, None]: assert group_size >= 1 group = [] for x in items: group.append(x) if len(group) == group_size: yield tuple(group) del group[:] if group: yield tuple(group) @classmethod def _batched_weighted_sum(cls, weights: Iterable[float], vectors: Iterable[np.ndarray], batch_size: int) -> np.ndarray: """calculate a weighted sum of the given vectors, in steps of at most `batch_size` vectors""" # start with float, at the first operation numpy broadcasting takes care of the correct shape total: Union[np.array, float] = 0. for batch_weights, batch_vectors in zip( cls._iter_groups(weights, batch_size), cls._iter_groups(vectors, batch_size)): assert len(batch_weights) == len(batch_vectors) <= batch_size total += np.dot(np.asarray(batch_weights, dtype=np.float32), np.asarray(batch_vectors, dtype=np.float32)) return total @classmethod def _compute_ranks(cls, x: np.ndarray) -> np.ndarray: """ Returns ranks in [0, len(x)) Note: This is different from scipy.stats.rankdata, which returns ranks in [1, len(x)]. """ assert x.ndim == 1 ranks = np.empty(len(x), dtype=int) ranks[x.argsort()] = np.arange(len(x)) return ranks @classmethod def _compute_centered_ranks(cls, x): y = cls._compute_ranks(x.ravel()).reshape(x.shape).astype(np.float32) y /= (x.size - 1) y -= .5 return y