from typing import Optional, Union, Any, List from easydict import EasyDict from ding.utils import deep_merge_dicts, SequenceType from collections import namedtuple import numpy as np import torch class LevelSampler(): """ Overview: Policy class of Prioritized Level Replay algorithm. https://arxiv.org/pdf/2010.03934.pdf PLR is a method for improving generalization and sample-efficiency of \ deep RL agents on procedurally-generated environments by adaptively updating \ a sampling distribution over the training levels based on a score of the learning \ potential of replaying each level. """ config = dict( strategy='policy_entropy', replay_schedule='fixed', score_transform='rank', temperature=1.0, eps=0.05, rho=0.2, nu=0.5, alpha=1.0, staleness_coef=0, staleness_transform='power', staleness_temperature=1.0, ) def __init__( self, seeds: Optional[List[int]], obs_space: Union[int, SequenceType], action_space: int, num_actors: int, cfg: EasyDict, ): self.cfg = EasyDict(deep_merge_dicts(self.config, cfg)) self.cfg.update(cfg) self.obs_space = obs_space self.action_space = action_space self.strategy = self.cfg.strategy self.replay_schedule = self.cfg.replay_schedule self.score_transform = self.cfg.score_transform self.temperature = self.cfg.temperature # Eps means the level replay epsilon for eps-greedy sampling self.eps = self.cfg.eps # Rho means the minimum size of replay set relative to total number of levels before sampling replays self.rho = self.cfg.rho # Nu means the probability of sampling a new level instead of a replay level self.nu = self.cfg.nu # Alpha means the level score EWA smoothing factor self.alpha = self.cfg.alpha self.staleness_coef = self.cfg.staleness_coef self.staleness_transform = self.cfg.staleness_transform self.staleness_temperature = self.cfg.staleness_temperature # Track seeds and scores as in np arrays backed by shared memory self.seeds = np.array(seeds, dtype=np.int64) self.seed2index = {seed: i for i, seed in enumerate(seeds)} self.unseen_seed_weights = np.ones(len(seeds)) self.seed_scores = np.zeros(len(seeds)) self.partial_seed_scores = np.zeros((num_actors, len(seeds)), dtype=np.float32) self.partial_seed_steps = np.zeros((num_actors, len(seeds)), dtype=np.int64) self.seed_staleness = np.zeros(len(seeds)) self.next_seed_index = 0 # Only used for sequential strategy def update_with_rollouts(self, train_data: dict, num_actors: int): total_steps = train_data['reward'].shape[0] if self.strategy == 'random': return if self.strategy == 'policy_entropy': score_function = self._entropy elif self.strategy == 'least_confidence': score_function = self._least_confidence elif self.strategy == 'min_margin': score_function = self._min_margin elif self.strategy == 'gae': score_function = self._gae elif self.strategy == 'value_l1': score_function = self._value_l1 elif self.strategy == 'one_step_td_error': score_function = self._one_step_td_error else: raise ValueError('Not supported strategy: {}'.format(self.strategy)) self._update_with_rollouts(train_data, num_actors, total_steps, score_function) for actor_index in range(self.partial_seed_scores.shape[0]): for seed_idx in range(self.partial_seed_scores.shape[1]): if self.partial_seed_scores[actor_index][seed_idx] != 0: self.update_seed_score(actor_index, seed_idx, 0, 0) self.partial_seed_scores.fill(0) self.partial_seed_steps.fill(0) def update_seed_score(self, actor_index: int, seed_idx: int, score: float, num_steps: int): score = self._partial_update_seed_score(actor_index, seed_idx, score, num_steps, done=True) self.unseen_seed_weights[seed_idx] = 0. # No longer unseen old_score = self.seed_scores[seed_idx] self.seed_scores[seed_idx] = (1 - self.alpha) * old_score + self.alpha * score def _partial_update_seed_score( self, actor_index: int, seed_idx: int, score: float, num_steps: int, done: bool = False ): partial_score = self.partial_seed_scores[actor_index][seed_idx] partial_num_steps = self.partial_seed_steps[actor_index][seed_idx] running_num_steps = partial_num_steps + num_steps merged_score = partial_score + (score - partial_score) * num_steps / float(running_num_steps) if done: self.partial_seed_scores[actor_index][seed_idx] = 0. # zero partial score, partial num_steps self.partial_seed_steps[actor_index][seed_idx] = 0 else: self.partial_seed_scores[actor_index][seed_idx] = merged_score self.partial_seed_steps[actor_index][seed_idx] = running_num_steps return merged_score def _entropy(self, **kwargs): episode_logits = kwargs['episode_logits'] num_actions = self.action_space max_entropy = -(1. / num_actions) * np.log(1. / num_actions) * num_actions return (-torch.exp(episode_logits) * episode_logits).sum(-1).mean().item() / max_entropy def _least_confidence(self, **kwargs): episode_logits = kwargs['episode_logits'] return (1 - torch.exp(episode_logits.max(-1, keepdim=True)[0])).mean().item() def _min_margin(self, **kwargs): episode_logits = kwargs['episode_logits'] top2_confidence = torch.exp(episode_logits.topk(2, dim=-1)[0]) return 1 - (top2_confidence[:, 0] - top2_confidence[:, 1]).mean().item() def _gae(self, **kwargs): advantages = kwargs['adv'] return advantages.mean().item() def _value_l1(self, **kwargs): advantages = kwargs['adv'] # If the absolute value of ADV is large, it means that the level can significantly change # the policy and can be used to learn more return advantages.abs().mean().item() def _one_step_td_error(self, **kwargs): rewards = kwargs['rewards'] value = kwargs['value'] max_t = len(rewards) td_errors = (rewards[:-1] + value[:max_t - 1] - value[1:max_t]).abs() return td_errors.abs().mean().item() def _update_with_rollouts(self, train_data: dict, num_actors: int, all_total_steps: int, score_function): level_seeds = train_data['seed'].reshape(num_actors, int(all_total_steps / num_actors)).transpose(0, 1) policy_logits = train_data['logit'].reshape(num_actors, int(all_total_steps / num_actors), -1).transpose(0, 1) done = train_data['done'].reshape(num_actors, int(all_total_steps / num_actors)).transpose(0, 1) total_steps, num_actors = policy_logits.shape[:2] num_decisions = len(policy_logits) for actor_index in range(num_actors): done_steps = done[:, actor_index].nonzero()[:total_steps, 0] start_t = 0 for t in done_steps: if not start_t < total_steps: break if t == 0: # if t is 0, then this done step caused a full update of previous seed last cycle continue seed_t = level_seeds[start_t, actor_index].item() seed_t = int(seed_t) seed_idx_t = self.seed2index[seed_t] score_function_kwargs = {} episode_logits = policy_logits[start_t:t, actor_index] score_function_kwargs['episode_logits'] = torch.log_softmax(episode_logits, -1) if self.strategy in ['gae', 'value_l1', 'one_step_td_error']: rewards = train_data['reward'].reshape(num_actors, int(all_total_steps / num_actors)).transpose(0, 1) adv = train_data['adv'].reshape(num_actors, int(all_total_steps / num_actors)).transpose(0, 1) value = train_data['value'].reshape(num_actors, int(all_total_steps / num_actors)).transpose(0, 1) score_function_kwargs['adv'] = adv[start_t:t, actor_index] score_function_kwargs['rewards'] = rewards[start_t:t, actor_index] score_function_kwargs['value'] = value[start_t:t, actor_index] score = score_function(**score_function_kwargs) num_steps = len(episode_logits) self.update_seed_score(actor_index, seed_idx_t, score, num_steps) start_t = t.item() if start_t < total_steps: seed_t = level_seeds[start_t, actor_index].item() seed_idx_t = self.seed2index[seed_t] score_function_kwargs = {} episode_logits = policy_logits[start_t:, actor_index] score_function_kwargs['episode_logits'] = torch.log_softmax(episode_logits, -1) if self.strategy in ['gae', 'value_l1', 'one_step_td_error']: rewards = train_data['reward'].reshape(num_actors, int(all_total_steps / num_actors)).transpose(0, 1) adv = train_data['adv'].reshape(num_actors, int(all_total_steps / num_actors)).transpose(0, 1) value = train_data['value'].reshape(num_actors, int(all_total_steps / num_actors)).transpose(0, 1) score_function_kwargs['adv'] = adv[start_t:, actor_index] score_function_kwargs['rewards'] = rewards[start_t:, actor_index] score_function_kwargs['value'] = value[start_t:, actor_index] score = score_function(**score_function_kwargs) num_steps = len(episode_logits) self._partial_update_seed_score(actor_index, seed_idx_t, score, num_steps) def _update_staleness(self, selected_idx: int): if self.staleness_coef > 0: self.seed_staleness += 1 self.seed_staleness[selected_idx] = 0 def _sample_replay_level(self): sample_weights = self._sample_weights() if np.isclose(np.sum(sample_weights), 0): sample_weights = np.ones_like(sample_weights, dtype=np.float32) / len(sample_weights) seed_idx = np.random.choice(range(len(self.seeds)), 1, p=sample_weights)[0] seed = self.seeds[seed_idx] self._update_staleness(seed_idx) return int(seed) def _sample_unseen_level(self): sample_weights = self.unseen_seed_weights / self.unseen_seed_weights.sum() seed_idx = np.random.choice(range(len(self.seeds)), 1, p=sample_weights)[0] seed = self.seeds[seed_idx] self._update_staleness(seed_idx) return int(seed) def sample(self, strategy: Optional[str] = None): if not strategy: strategy = self.strategy if strategy == 'random': seed_idx = np.random.choice(range(len(self.seeds))) seed = self.seeds[seed_idx] return int(seed) elif strategy == 'sequential': seed_idx = self.next_seed_index self.next_seed_index = (self.next_seed_index + 1) % len(self.seeds) seed = self.seeds[seed_idx] return int(seed) num_unseen = (self.unseen_seed_weights > 0).sum() proportion_seen = (len(self.seeds) - num_unseen) / len(self.seeds) if self.replay_schedule == 'fixed': if proportion_seen >= self.rho: # Sample replay level with fixed prob = 1 - nu OR if all levels seen if np.random.rand() > self.nu or not proportion_seen < 1.0: return self._sample_replay_level() # Otherwise, sample a new level return self._sample_unseen_level() else: # Default to proportionate schedule if proportion_seen >= self.rho and np.random.rand() < proportion_seen: return self._sample_replay_level() else: return self._sample_unseen_level() def _sample_weights(self): weights = self._score_transform(self.score_transform, self.temperature, self.seed_scores) weights = weights * (1 - self.unseen_seed_weights) # zero out unseen levels z = np.sum(weights) if z > 0: weights /= z staleness_weights = 0 if self.staleness_coef > 0: staleness_weights = self._score_transform( self.staleness_transform, self.staleness_temperature, self.seed_staleness ) staleness_weights = staleness_weights * (1 - self.unseen_seed_weights) z = np.sum(staleness_weights) if z > 0: staleness_weights /= z weights = (1 - self.staleness_coef) * weights + self.staleness_coef * staleness_weights return weights def _score_transform(self, transform: Optional[str], temperature: float, scores: Optional[List[float]]): if transform == 'rank': temp = np.flip(scores.argsort()) ranks = np.empty_like(temp) ranks[temp] = np.arange(len(temp)) + 1 weights = 1 / ranks ** (1. / temperature) elif transform == 'power': eps = 0 if self.staleness_coef > 0 else 1e-3 weights = (np.array(scores) + eps) ** (1. / temperature) return weights