def declare_memory(self): if self.fix_buffer: self.memory = ExperienceReplayMemory( self.experience_replay_size ) if not self.priority_replay else PrioritizedReplayMemory( self.experience_replay_size, self.priority_alpha, self.priority_beta_start, self.priority_beta_frames) else: self.memory = MutExperienceReplayMemory( self.experience_replay_size ) if not self.priority_replay else MutPrioritizedReplayMemory( self.priority_alpha, self.priority_beta_start, self.priority_beta_frames)
class Model(DQN_Agent): def __init__(self, static_policy=False, env=None, config=None, log_dir='/tmp/gym'): self.atoms = config.ATOMS self.v_max = config.V_MAX self.v_min = config.V_MIN self.supports = torch.linspace(self.v_min, self.v_max, self.atoms).view( 1, 1, self.atoms).to(config.device) self.delta = (self.v_max - self.v_min) / (self.atoms - 1) super(Model, self).__init__(static_policy, env, config, log_dir=log_dir) self.nsteps = max(self.nsteps, 3) def declare_networks(self): self.model = CategoricalDuelingDQN(self.env.observation_space.shape, self.env.action_space.n, noisy=True, sigma_init=self.sigma_init, atoms=self.atoms) self.target_model = CategoricalDuelingDQN( self.env.observation_space.shape, self.env.action_space.n, noisy=True, sigma_init=self.sigma_init, atoms=self.atoms) def declare_memory(self): self.memory = PrioritizedReplayMemory(self.experience_replay_size, self.priority_alpha, self.priority_beta_start, self.priority_beta_frames) def projection_distribution(self, batch_vars): batch_state, batch_action, batch_reward, non_final_next_states, non_final_mask, empty_next_state_values, indices, weights = batch_vars with torch.no_grad(): max_next_dist = torch.zeros((self.batch_size, 1, self.atoms), device=self.device, dtype=torch.float) + 1. / self.atoms if not empty_next_state_values: max_next_action = self.get_max_next_state_action( non_final_next_states) self.target_model.sample_noise() max_next_dist[non_final_mask] = self.target_model( non_final_next_states).gather(1, max_next_action) max_next_dist = max_next_dist.squeeze() Tz = batch_reward.view( -1, 1) + (self.gamma**self.nsteps) * self.supports.view( 1, -1) * non_final_mask.to(torch.float).view(-1, 1) Tz = Tz.clamp(self.v_min, self.v_max) b = (Tz - self.v_min) / self.delta l = b.floor().to(torch.int64) u = b.ceil().to(torch.int64) l[(u > 0) * (l == u)] -= 1 u[(l < (self.atoms - 1)) * (l == u)] += 1 offset = torch.linspace(0, (self.batch_size - 1) * self.atoms, self.batch_size).unsqueeze(dim=1).expand( self.batch_size, self.atoms).to(batch_action) m = batch_state.new_zeros(self.batch_size, self.atoms) m.view(-1).index_add_( 0, (l + offset).view(-1), (max_next_dist * (u.float() - b)).view(-1)) # m_l = m_l + p(s_t+n, a*)(u - b) m.view(-1).index_add_( 0, (u + offset).view(-1), (max_next_dist * (b - l.float())).view(-1)) # m_u = m_u + p(s_t+n, a*)(b - l) return m def compute_loss(self, batch_vars): batch_state, batch_action, batch_reward, non_final_next_states, non_final_mask, empty_next_state_values, indices, weights = batch_vars batch_action = batch_action.unsqueeze(dim=-1).expand( -1, -1, self.atoms) batch_reward = batch_reward.view(-1, 1, 1) #estimate self.model.sample_noise() current_dist = self.model(batch_state).gather(1, batch_action).squeeze() target_prob = self.projection_distribution(batch_vars) loss = -(target_prob * current_dist.log()).sum(-1) self.memory.update_priorities( indices, loss.detach().squeeze().abs().cpu().numpy().tolist()) loss = loss * weights loss = loss.mean() return loss def get_action(self, s, eps): with torch.no_grad(): X = torch.tensor([s], device=self.device, dtype=torch.float) self.model.sample_noise() a = self.model(X) * self.supports a = a.sum(dim=2).max(1)[1].view(1, 1) return a.item() def get_max_next_state_action(self, next_states): next_dist = self.model(next_states) * self.supports return next_dist.sum(dim=2).max(1)[1].view(next_states.size(0), 1, 1).expand( -1, -1, self.atoms)
def declare_memory(self): self.memory = PrioritizedReplayMemory(self.experience_replay_size, self.priority_alpha, self.priority_beta_start, self.priority_beta_frames)
class Model(DQN_Agent): def __init__(self, static_policy=False, env=None, config=None): self.num_quantiles = config.QUANTILES self.cumulative_density = torch.tensor( (2 * np.arange(self.num_quantiles) + 1) / (2.0 * self.num_quantiles), device=config.device, dtype=torch.float) self.quantile_weight = 1.0 / self.num_quantiles super(Model, self).__init__(static_policy, env, config) self.nsteps = max(self.nsteps, 3) def declare_networks(self): self.model = DuelingQRDQN(self.env.observation_space.shape, self.env.action_space.n, noisy=True, sigma_init=self.sigma_init, quantiles=self.num_quantiles) self.target_model = DuelingQRDQN(self.env.observation_space.shape, self.env.action_space.n, noisy=True, sigma_init=self.sigma_init, quantiles=self.num_quantiles) def declare_memory(self): self.memory = PrioritizedReplayMemory(self.experience_replay_size, self.priority_alpha, self.priority_beta_start, self.priority_beta_frames) def next_distribution(self, batch_vars): batch_state, batch_action, batch_reward, non_final_next_states, non_final_mask, empty_next_state_values, indices, weights = batch_vars with torch.no_grad(): quantiles_next = torch.zeros((self.batch_size, self.num_quantiles), device=self.device, dtype=torch.float) if not empty_next_state_values: self.target_model.sample_noise() max_next_action = self.get_max_next_state_action( non_final_next_states) quantiles_next[non_final_mask] = self.target_model( non_final_next_states).gather( 1, max_next_action).squeeze(dim=1) quantiles_next = batch_reward + (self.gamma * quantiles_next) return quantiles_next def compute_loss(self, batch_vars): batch_state, batch_action, batch_reward, non_final_next_states, non_final_mask, empty_next_state_values, indices, weights = batch_vars batch_action = batch_action.unsqueeze(dim=-1).expand( -1, -1, self.num_quantiles) self.model.sample_noise() quantiles = self.model(batch_state) quantiles = quantiles.gather(1, batch_action).squeeze(1) quantiles_next = self.next_distribution(batch_vars) diff = quantiles_next.t().unsqueeze(-1) - quantiles.unsqueeze(0) loss = self.huber(diff) * torch.abs( self.cumulative_density.view(1, -1) - (diff < 0).to(torch.float)) loss = loss.transpose(0, 1) self.memory.update_priorities( indices, loss.detach().mean(1).sum(-1).abs().cpu().numpy().tolist()) loss = loss * weights.view(self.batch_size, 1, 1) loss = loss.mean(1).sum(-1).mean() return loss def get_action(self, s, eps): with torch.no_grad(): X = torch.tensor([s], device=self.device, dtype=torch.float) self.model.sample_noise() a = (self.model(X) * self.quantile_weight).sum(dim=2).max(dim=1)[1] return a.item() def get_max_next_state_action(self, next_states): next_dist = self.model(next_states) * self.quantile_weight return next_dist.sum(dim=2).max(1)[1].view( next_states.size(0), 1, 1).expand(-1, -1, self.num_quantiles)