text: #######################################################################
  # Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com)    #
  # Permission given to modify the code as long as you keep this        #
  # declaration at the top                                              #
  #######################################################################
  
  from .BaseAgent import *
  from ..network import *
  
  
  class QuantileRegressionDQNActor(BaseActor):
      def __init__(self, config):
          BaseActor.__init__(self, config)
          self.config = config
          self.start()
  
      def _transition(self):
          if self._state is None:
              self._state = self._task.reset()
          config = self.config
          with config.lock:
              q_values = self._network(config.state_normalizer(self._state)).mean(-1)
          q_values = to_np(q_values).flatten()
          if self._total_steps < config.exploration_steps \
                  or np.random.rand() < config.random_action_prob():
              action = np.random.randint(0, len(q_values))
          else:
              action = np.argmax(q_values)
          next_state, reward, done, info = self._task.step([action])
          entry = [self._state[0], action, reward[0], next_state[0], int(done[0]), info]
          self._total_steps += 1
          self._state = next_state
          return entry
  
  
  class QuantileRegressionDQNAgent(BaseAgent):
      def __init__(self, config):
          BaseAgent.__init__(self, config)
          self.config = config
          config.lock = mp.Lock()
  
          self.replay = config.replay_fn()
          self.actor = QuantileRegressionDQNActor(config)
  
          self.network = config.network_fn()
          self.network.share_memory()
          self.target_network = config.network_fn()
          self.target_network.load_state_dict(self.network.state_dict())
          self.optimizer = config.optimizer_fn(self.network.parameters())
  
          self.actor.set_network(self.network)
  
          self.episode_reward = 0
          self.episode_rewards = []
  
          self.total_steps = 0
          self.batch_indices = range_tensor(self.replay.batch_size)
  
          self.quantile_weight = 1.0 / self.config.num_quantiles
          self.cumulative_density = tensor(
              (2 * np.arange(self.config.num_quantiles) + 1) / (2.0 * self.config.num_quantiles)).view(1, -1)
  
      def close(self):
          close_obj(self.replay)
          close_obj(self.actor)
  
      def eval_step(self, state):
          self.config.state_normalizer.set_read_only()
          state = self.config.state_normalizer(state)
          q = self.network(state).mean(-1)
          action = np.argmax(to_np(q).flatten())
          self.config.state_normalizer.unset_read_only()
          return action
  
      def step(self):
          config = self.config
          transitions = self.actor.step()
          experiences = []
          for state, action, reward, next_state, done, _ in transitions:
              self.episode_reward += reward
              self.total_steps += 1
              reward = config.reward_normalizer(reward)
              if done:
                  self.episode_rewards.append(self.episode_reward)
                  self.episode_reward = 0
              experiences.append([state, action, reward, next_state, done])
          self.replay.feed_batch(experiences)
  
          if self.total_steps > self.config.exploration_steps:
              experiences = self.replay.sample()
              states, actions, rewards, next_states, terminals = experiences
              states = self.config.state_normalizer(states)
              next_states = self.config.state_normalizer(next_states)
  
              quantiles_next = self.target_network(next_states).detach()
              a_next = torch.argmax(quantiles_next.sum(-1), dim=-1)
              quantiles_next = quantiles_next[self.batch_indices, a_next, :]
  
              rewards = tensor(rewards).unsqueeze(-1)
              terminals = tensor(terminals).unsqueeze(-1)
              quantiles_next = rewards + self.config.discount * (1 - terminals) * quantiles_next
  
              quantiles = self.network(states)
              actions = tensor(actions).long()
              quantiles = quantiles[self.batch_indices, actions, :]
  
              quantiles_next = quantiles_next.t().unsqueeze(-1)
              diff = quantiles_next - quantiles
              loss = huber(diff) * (self.cumulative_density - (diff.detach() < 0).float()).abs()
  
              self.optimizer.zero_grad()
              loss.mean(0).mean(1).sum().backward()
              nn.utils.clip_grad_norm_(self.network.parameters(), self.config.gradient_clip)
              with config.lock:
                  self.optimizer.step()
  
          if self.total_steps / self.config.sgd_update_frequency % \
                  self.config.target_network_update_freq == 0:
              self.target_network.load_state_dict(self.network.state_dict())
  