text: # coding=utf-8
  # Copyright 2020 The Google Research Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  # Lint as: python2, python3
  """DeepQNetwork models for molecule generation."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import functools
  
  import numpy as np
  
  from rdkit import Chem
  from rdkit import DataStructs
  from rdkit.Chem import AllChem
  
  from six.moves import range
  from six.moves import zip
  import tensorflow.compat.v1 as tf
  from tensorflow.contrib import layers as contrib_layers
  from tensorflow.contrib import training as contrib_training
  
  
  class DeepQNetwork(object):
    """Deep Q Network.
  
    This class implements the network as used in the Nature
    (2015) paper.
    Human-level control through deep reinforcement learning
    https://www.nature.com/articles/nature14236
    https://storage.googleapis.com/deepmind-data/assets/papers/DeepMindNature14236Paper.pdf
    """
  
    def __init__(self,
                 input_shape,
                 q_fn,
                 learning_rate=0.001,
                 learning_rate_decay_steps=10000,
                 learning_rate_decay_rate=0.8,
                 optimizer='Adam',
                 grad_clipping=None,
                 gamma=1.0,
                 epsilon=0.2,
                 double_q=True,
                 num_bootstrap_heads=10,
                 scope='dqn',
                 reuse=None):
      """Creates the model function.
  
      Args:
        input_shape: Tuple. The shape of input.
        q_fn: A function, whose input is the observation features, and the
          output is the Q value of the observation.
        learning_rate: Float. The learning rate of the optimizer.
        learning_rate_decay_steps: Integer. The number of steps between each
          learning rate decay.
        learning_rate_decay_rate: Float. The rate of learning rate decay.
        optimizer: String. Which optimizer to use.
        grad_clipping: Boolean. Whether to clip gradient.
        gamma: Float. Discount factor.
        epsilon: Float. The probability of choosing a random action.
        double_q: Boolean. Whether to use double q learning.
        num_bootstrap_heads: Integer. The number of bootstrap heads to use.
        scope: String or VariableScope. Variable Scope.
        reuse: Boolean or None. Whether or not the variable should be reused.
      """
      self.input_shape = input_shape
      self.q_fn = q_fn
      self.learning_rate = learning_rate
      self.learning_rate_decay_steps = learning_rate_decay_steps
      self.learning_rate_decay_rate = learning_rate_decay_rate
      self.optimizer = optimizer
      self.grad_clipping = grad_clipping
      self.gamma = gamma
      self.num_bootstrap_heads = num_bootstrap_heads
      self.double_q = double_q
      self.scope = scope
      self.reuse = reuse
      self.epsilon = epsilon
  
    def build(self):
      """Builds the computational graph and training operations."""
      self._build_graph()
      self._build_training_ops()
      self._build_summary_ops()
  
    def _build_single_q_network(self, observations, head, state_t, state_tp1,
                                done_mask, reward_t, error_weight):
      """Builds the computational graph for a single Q network.
  
      Briefly, this part is calculating the following two quantities:
      1. q_value = q_fn(observations)
      2. td_error = q_fn(state_t) - reward_t - gamma * q_fn(state_tp1)
      The optimization target is to minimize the td_error.
  
      Args:
        observations: shape = [batch_size, hparams.fingerprint_length].
          The input of the Q function.
        head: shape = [1].
          The index of the head chosen for decision in bootstrap DQN.
        state_t: shape = [batch_size, hparams.fingerprint_length].
          The state at time step t.
        state_tp1: a list of tensors, with total number of batch_size,
          each has shape = [num_actions, hparams.fingerprint_length].
          Note that the num_actions can be different for each tensor.
          The state at time step t+1, tp1 is short for t plus 1.
        done_mask: shape = [batch_size, 1]
          Whether state_tp1 is the terminal state.
        reward_t: shape = [batch_size, 1]
          the reward at time step t.
        error_weight: shape = [batch_size, 1]
          weight for the loss.
  
      Returns:
        q_values: Tensor of [batch_size, 1]. The q values for the observations.
        td_error: Tensor of [batch_size, 1]. The TD error.
        weighted_error: Tensor of [batch_size, 1]. The TD error weighted by
          error_weight.
        q_fn_vars: List of tf.Variables. The variables of q_fn when computing
          the q_values of state_t
        q_fn_vars: List of tf.Variables. The variables of q_fn when computing
          the q_values of state_tp1
  
      """
      with tf.variable_scope('q_fn'):
        # q_value have shape [batch_size, 1].
        q_values = tf.gather(self.q_fn(observations), head, axis=-1)
  
      # calculating q_fn(state_t)
      # The Q network shares parameters with the action graph.
      with tf.variable_scope('q_fn', reuse=True):
        q_t = self.q_fn(state_t, reuse=True)
      q_fn_vars = tf.trainable_variables(scope=tf.get_variable_scope().name +
                                         '/q_fn')
  
      # calculating q_fn(state_tp1)
      with tf.variable_scope('q_tp1', reuse=tf.AUTO_REUSE):
        q_tp1 = [self.q_fn(s_tp1, reuse=tf.AUTO_REUSE) for s_tp1 in state_tp1]
      q_tp1_vars = tf.trainable_variables(scope=tf.get_variable_scope().name +
                                          '/q_tp1')
  
      if self.double_q:
        with tf.variable_scope('q_fn', reuse=True):
          q_tp1_online = [self.q_fn(s_tp1, reuse=True) for s_tp1 in state_tp1]
        if self.num_bootstrap_heads:
          num_heads = self.num_bootstrap_heads
        else:
          num_heads = 1
        # determine the action to choose based on online Q estimator.
        q_tp1_online_idx = [
            tf.stack(
                [tf.argmax(q, axis=0),
                 tf.range(num_heads, dtype=tf.int64)],
                axis=1) for q in q_tp1_online
        ]
        # use the index from max online q_values to compute the value
        # function
        v_tp1 = tf.stack(
            [tf.gather_nd(q, idx) for q, idx in zip(q_tp1, q_tp1_online_idx)],
            axis=0)
      else:
        v_tp1 = tf.stack([tf.reduce_max(q) for q in q_tp1], axis=0)
  
      # if s_{t+1} is the terminal state, we do not evaluate the Q value of
      # the state.
      q_tp1_masked = (1.0 - done_mask) * v_tp1
  
      q_t_target = reward_t + self.gamma * q_tp1_masked
  
      # stop gradient from flowing to the computating graph which computes
      # the Q value of s_{t+1}.
      # td_error has shape [batch_size, 1]
      td_error = q_t - tf.stop_gradient(q_t_target)
  
      # If use bootstrap, each head is trained with a different subset of the
      # training sample. Like the idea of dropout.
      if self.num_bootstrap_heads:
        head_mask = tf.keras.backend.random_binomial(
            shape=(1, self.num_bootstrap_heads), p=0.6)
        td_error = tf.reduce_mean(td_error * head_mask, axis=1)
      # The loss comes from a traditional trick in convex optimization:
      # http://web.stanford.edu/~boyd/cvxbook/.
      # See Chapter 6 pp. 298
      # It will makes the optimization robust.
      # Specifically, the loss will use l1 instead of l2 loss when the td error
      # gets larger than 1.0. The l2 loss has the disadvantage that it has
      # the tendency to be dominated by outliers. In terms of estimation theory,
      # the asymptotic relative efficiency of the l1 loss estimator is better
      # for heavy-tailed distributions.
      errors = tf.where(
          tf.abs(td_error) < 1.0, tf.square(td_error) * 0.5,
          1.0 * (tf.abs(td_error) - 0.5))
      weighted_error = tf.reduce_mean(error_weight * errors)
      return q_values, td_error, weighted_error, q_fn_vars, q_tp1_vars
  
    def _build_input_placeholder(self):
      """Creates the input placeholders.
  
      Input placeholders created:
        observations: shape = [batch_size, hparams.fingerprint_length].
          The input of the Q function.
        head: shape = [1].
          The index of the head chosen for decision.
        state_t: shape = [batch_size, hparams.fingerprint_length].
          The state at time step t.
        state_tp1: a list of tensors,
          each has shape = [num_actions, hparams.fingerprint_length].
          Note that the num_actions can be different for each tensor.
          The state at time step t+1.
        done_mask: shape = [batch_size, 1]
          Whether state_tp1 is the terminal state.
  
        error_weight: shape = [batch_size, 1]
          weight for the loss.
      """
      batch_size, fingerprint_length = self.input_shape
  
      with tf.variable_scope(self.scope, reuse=self.reuse):
        # Build the action graph to choose an action.
        # The observations, which are the inputs of the Q function.
        self.observations = tf.placeholder(
            tf.float32, [None, fingerprint_length], name='observations')
        # head is the index of the head we want to choose for decison.
        # See https://arxiv.org/abs/1703.07608
        self.head = tf.placeholder(tf.int32, [], name='head')
  
        # When sample from memory, the batch_size can be fixed, as it is
        # possible to sample any number of samples from memory.
        # state_t is the state at time step t
        self.state_t = tf.placeholder(
            tf.float32, self.input_shape, name='state_t')
        # state_tp1 is the state at time step t + 1, tp1 is short for t plus 1.
        self.state_tp1 = [
            tf.placeholder(
                tf.float32, [None, fingerprint_length], name='state_tp1_%i' % i)
            for i in range(batch_size)
        ]
        # done_mask is a {0, 1} tensor indicating whether state_tp1 is the
        # terminal state.
        self.done_mask = tf.placeholder(
            tf.float32, (batch_size, 1), name='done_mask')
  
        self.error_weight = tf.placeholder(
            tf.float32, (batch_size, 1), name='error_weight')
  
    def _build_graph(self):
      """Builds the computational graph.
  
      Input placeholders created:
        reward_t: shape = [batch_size, 1]
          the reward at time step t.
  
      Instance attributes created:
        q_values: the q values of the observations.
        q_fn_vars: the variables in q function.
        q_tp1_vars: the variables in q_tp1 function.
        td_error: the td_error.
        weighted_error: the weighted td error.
        action: the action to choose next step.
      """
      batch_size, _ = self.input_shape
      with tf.variable_scope(self.scope, reuse=self.reuse):
        self._build_input_placeholder()
        self.reward_t = tf.placeholder(
            tf.float32, (batch_size, 1), name='reward_t')
        # The Q network shares parameters with the action graph.
        # tenors start with q or v have shape [batch_size, 1] when not using
        # bootstrap. When using bootstrap, the shapes are
        # [batch_size, num_bootstrap_heads]
        (self.q_values, self.td_error, self.weighted_error,
         self.q_fn_vars, self.q_tp1_vars) = self._build_single_q_network(
             self.observations, self.head, self.state_t, self.state_tp1,
             self.done_mask, self.reward_t, self.error_weight)
        self.action = tf.argmax(self.q_values)
  
    def _build_training_ops(self):
      """Creates the training operations.
  
      Instance attributes created:
        optimization_op: the operation of optimize the loss.
        update_op: the operation to update the q network.
      """
      with tf.variable_scope(self.scope, reuse=self.reuse):
  
        self.optimization_op = contrib_layers.optimize_loss(
            loss=self.weighted_error,
            global_step=tf.train.get_or_create_global_step(),
            learning_rate=self.learning_rate,
            optimizer=self.optimizer,
            clip_gradients=self.grad_clipping,
            learning_rate_decay_fn=functools.partial(
                tf.train.exponential_decay,
                decay_steps=self.learning_rate_decay_steps,
                decay_rate=self.learning_rate_decay_rate),
            variables=self.q_fn_vars)
  
        self.update_op = []
        for var, target in zip(
            sorted(self.q_fn_vars, key=lambda v: v.name),
            sorted(self.q_tp1_vars, key=lambda v: v.name)):
          self.update_op.append(target.assign(var))
        self.update_op = tf.group(*self.update_op)
  
    def _build_summary_ops(self):
      """Creates the summary operations.
  
      Input placeholders created:
        smiles: the smiles string.
        reward: the reward.
  
      Instance attributes created:
        error_summary: the operation to log the summary of error.
        episode_summary: the operation to log the smiles string and reward.
      """
      with tf.variable_scope(self.scope, reuse=self.reuse):
        with tf.name_scope('summaries'):
          # The td_error here is the difference between q_t and q_t_target.
          # Without abs(), the summary of td_error is actually underestimated.
          self.error_summary = tf.summary.scalar(
              'td_error', tf.reduce_mean(tf.abs(self.td_error)))
          self.smiles = tf.placeholder(tf.string, [], 'summary_smiles')
          self.reward = tf.placeholder(tf.float32, [], 'summary_reward')
          smiles_summary = tf.summary.text('SMILES', self.smiles)
          reward_summary = tf.summary.scalar('reward', self.reward)
          self.episode_summary = tf.summary.merge(
              [smiles_summary, reward_summary])
  
    def log_result(self, smiles, reward):
      """Summarizes the SMILES string and reward at the end of an episode.
  
      Args:
        smiles: String. The SMILES string.
        reward: Float. The reward.
  
      Returns:
        the summary protobuf
      """
      return tf.get_default_session().run(
          self.episode_summary,
          feed_dict={
              self.smiles: smiles,
              self.reward: reward
          })
  
    def _run_action_op(self, observations, head):
      """Function that runs the op calculating an action given the observations.
  
      Args:
        observations: np.array. shape = [num_actions, fingerprint_length].
          Observations that can be feed into the Q network.
        head: Integer. The output index to use.
  
      Returns:
        Integer. which action to be performed.
      """
      return np.asscalar(tf.get_default_session().run(
          self.action,
          feed_dict={
              self.observations: observations,
              self.head: head
          }))
  
    def get_action(self,
                   observations,
                   stochastic=True,
                   head=0,
                   update_epsilon=None):
      """Function that chooses an action given the observations.
  
      Args:
        observations: np.array. shape = [num_actions, fingerprint_length].
          Observations that can be feed into the Q network.
        stochastic: Boolean. If set to False all the actions are always
          deterministic (default True).
        head: Integer. The output index to use.
        update_epsilon: Float or None. update epsilon a new value, if None
          no update happens (default: no update).
  
      Returns:
        Integer. which action to be performed.
      """
      if update_epsilon is not None:
        self.epsilon = update_epsilon
  
      if stochastic and np.random.uniform() < self.epsilon:
        return np.random.randint(0, observations.shape[0])
      else:
        return self._run_action_op(observations, head)
  
    def train(self, states, rewards, next_states, done, weight, summary=True):
      """Function that takes a transition (s,a,r,s') and optimizes Bellman error.
  
      Args:
        states: object, a batch of observations.
        rewards: np.array, immediate reward attained after executing those actions
          dtype must be float32 and shape must be (batch_size,).
        next_states: object, observations that followed states.
        done: np.array, 1 if obs_t was the last observation in the episode and 0
          otherwise obs_tp1 gets ignored, but must be of the valid shape. dtype
          must be float32 and shape must be (batch_size,).
        weight: np.array, importance sampling weights for every element of the
          batch. dtype must be float32 and shape must be (batch_size,).
        summary: Boolean, whether to get summary.
  
      Returns:
        td_error: np.array. a list of differences between Q(s,a) and the
          target in Bellman's equation.
          dtype is float32 and shape is (batch_size,).
      """
      if summary:
        ops = [self.td_error, self.error_summary, self.optimization_op]
      else:
        ops = [self.td_error, self.optimization_op]
      feed_dict = {
          self.state_t: states,
          self.reward_t: rewards,
          self.done_mask: done,
          self.error_weight: weight
      }
      for i, next_state in enumerate(next_states):
        feed_dict[self.state_tp1[i]] = next_state
      return tf.get_default_session().run(ops, feed_dict=feed_dict)
  
  
  class MultiObjectiveDeepQNetwork(DeepQNetwork):
    """Multi Objective Deep Q Network.
  
    The idea is described in
    Multiobjective Reinforcement Learning: A Comprehensive Overview
    https://ieeexplore.ieee.org/document/6918520/
  
    Briefly, the difference between this Multi Objective Deep Q Network and
    a naive Deep Q Network is that this one uses one Q network for approximating
    each of the objectives. And a weighted sum of those Q values are used for
    decision making.
    The loss is the summation of the losses of each Q network.
    """
  
    def __init__(self, objective_weight, **kwargs):
      """Creates the model function.
  
      Args:
        objective_weight: np.array with shape [num_objectives, 1]. The weight
          vector for the objectives.
        **kwargs: arguments for the DeepQNetworks class.
  
      """
      # Normalize the sum to 1.
      self.objective_weight = objective_weight / np.sum(objective_weight)
      self.num_objectives = objective_weight.shape[0]
      super(MultiObjectiveDeepQNetwork, self).__init__(**kwargs)
  
    def _build_graph(self):
      """Builds the computational graph.
  
      Input placeholders created:
        observations: shape = [batch_size, hparams.fingerprint_length].
          The input of the Q function.
        head: shape = [1].
          The index of the head chosen for decision.
        objective_weight: shape = [num_objectives, 1].
          objective_weight is the weight to scalarize the objective vector:
          reward = sum (objective_weight_i * objective_i)
        state_t: shape = [batch_size, hparams.fingerprint_length].
          The state at time step t.
        state_tp1: a list of tensors,
          each has shape = [num_actions, hparams.fingerprint_length].
          Note that the num_actions can be different for each tensor.
          The state at time step t+1.
        done_mask: shape = [batch_size, 1]
          Whether state_tp1 is the terminal state.
        reward_t: shape = [batch_size, num_objectives]
          the reward at time step t.
        error weight: shape = [batch_size, 1]
          weight for the loss.
  
      Instance attributes created:
        q_values: List of Tensors of [batch_size, 1]. The q values for the
          observations.
        td_error: List of Tensor of [batch_size, 1]. The TD error.
          weighted_error: List of Tensor of [batch_size, 1]. The TD error weighted
          by importance sampling weight.
        q_fn_vars: List of tf.Variables. The variables of q_fn when computing
          the q_values of state_t
        q_fn_vars: List of tf.Variables. The variables of q_fn when computing
          the q_values of state_tp1
  
      """
      batch_size, _ = self.input_shape
      with tf.variable_scope(self.scope, reuse=self.reuse):
        self._build_input_placeholder()
        self.reward_t = tf.placeholder(
            tf.float32, (batch_size, self.num_objectives), name='reward_t')
        # objective_weight is the weight to scalarize the objective vector:
        # reward = sum (objective_weight_i * objective_i)
        self.objective_weight_input = tf.placeholder(
            tf.float32, [self.num_objectives, 1], name='objective_weight')
  
        # split reward for each q network
        rewards_list = tf.split(self.reward_t, self.num_objectives, axis=1)
        q_values_list = []
        self.td_error = []
        self.weighted_error = 0
        self.q_fn_vars = []
        self.q_tp1_vars = []
  
        # build a Q network for each objective
        for obj_idx in range(self.num_objectives):
          with tf.variable_scope('objective_%i' % obj_idx):
            (q_values, td_error, weighted_error,
             q_fn_vars, q_tp1_vars) = self._build_single_q_network(
                 self.observations, self.head, self.state_t, self.state_tp1,
                 self.done_mask, rewards_list[obj_idx], self.error_weight)
            q_values_list.append(tf.expand_dims(q_values, 1))
            # td error is for summary only.
            # weighted error is the optimization goal.
            self.td_error.append(td_error)
            self.weighted_error += weighted_error / self.num_objectives
            self.q_fn_vars += q_fn_vars
            self.q_tp1_vars += q_tp1_vars
        q_values = tf.concat(q_values_list, axis=1)
        # action is the one that leads to the maximum weighted reward.
        self.action = tf.argmax(
            tf.matmul(q_values, self.objective_weight_input), axis=0)
  
    def _build_summary_ops(self):
      """Creates the summary operations.
  
      Input placeholders created:
        smiles: the smiles string.
        rewards: the rewards.
        weighted_reward: the weighted sum of the rewards.
  
      Instance attributes created:
        error_summary: the operation to log the summary of error.
        episode_summary: the operation to log the smiles string and reward.
      """
      with tf.variable_scope(self.scope, reuse=self.reuse):
        with tf.name_scope('summaries'):
          # The td_error here is the difference between q_t and q_t_target.
          # Without abs(), the summary of td_error is actually underestimated.
          error_summaries = [
              tf.summary.scalar('td_error_%i' % i,
                                tf.reduce_mean(tf.abs(self.td_error[i])))
              for i in range(self.num_objectives)
          ]
          self.error_summary = tf.summary.merge(error_summaries)
          self.smiles = tf.placeholder(tf.string, [], 'summary_smiles')
          self.rewards = [
              tf.placeholder(tf.float32, [], 'summary_reward_obj_%i' % i)
              for i in range(self.num_objectives)
          ]
          # Weighted sum of the rewards.
          self.weighted_reward = tf.placeholder(tf.float32, [],
                                                'summary_reward_sum')
          smiles_summary = tf.summary.text('SMILES', self.smiles)
          reward_summaries = [
              tf.summary.scalar('reward_obj_%i' % i, self.rewards[i])
              for i in range(self.num_objectives)
          ]
          reward_summaries.append(
              tf.summary.scalar('sum_reward', self.rewards[-1]))
  
          self.episode_summary = tf.summary.merge([smiles_summary] +
                                                  reward_summaries)
  
    def log_result(self, smiles, reward):
      """Summarizes the SMILES string and reward at the end of an episode.
  
      Args:
        smiles: String. The SMILES string.
        reward: List of Float. The rewards for each objective.
  
      Returns:
        the summary protobuf.
      """
      feed_dict = {
          self.smiles: smiles,
      }
      for i, reward_value in enumerate(reward):
        feed_dict[self.rewards[i]] = reward_value
      # calculated the weighted sum of the rewards.
      feed_dict[self.weighted_reward] = np.asscalar(
          np.array([reward]).dot(self.objective_weight))
      return tf.get_default_session().run(
          self.episode_summary, feed_dict=feed_dict)
  
    def _run_action_op(self, observations, head):
      """Function that runs the op calculating an action given the observations.
  
      Args:
        observations: np.array. shape = [num_actions, fingerprint_length].
          Observations that can be feed into the Q network.
        head: Integer. The output index to use.
  
      Returns:
        Integer. which action to be performed.
      """
      return np.asscalar(tf.get_default_session().run(
          self.action,
          feed_dict={
              self.observations: observations,
              self.objective_weight_input: self.objective_weight,
              self.head: head
          }))
  
  
  def multi_layer_model(inputs, hparams, reuse=None):
    """Multi-layer model for q learning.
  
    Args:
      inputs: Tensor. The input.
      hparams: tf.HParameters. The hyper-parameters.
      reuse: Boolean. Whether the parameters should be reused.
  
    Returns:
      Tensor. shape = [batch_size, hparams.num_bootstrap_heads]. The output.
    """
    output = inputs
    for i, units in enumerate(hparams.dense_layers):
      output = tf.layers.dense(output, units, name='dense_%i' % i, reuse=reuse)
      output = getattr(tf.nn, hparams.activation)(output)
      if hparams.batch_norm:
        output = tf.layers.batch_normalization(
            output, fused=True, name='bn_%i' % i, reuse=reuse)
    if hparams.num_bootstrap_heads:
      output_dim = hparams.num_bootstrap_heads
    else:
      output_dim = 1
    output = tf.layers.dense(output, output_dim, name='final', reuse=reuse)
    return output
  
  
  def get_hparams(**kwargs):
    """Get the hyperparameters for the model from a json object.
  
    Args:
      **kwargs: Dict of parameter overrides.
    Possible keyword arguments:
      atom_types: Dict. The possible atom types in the molecule.
      max_steps_per_episode: Integer. The maximum number of steps for one episode.
      allow_removal: Boolean. Whether to allow removal of a bond.
      allow_no_modification: Boolean. If true, the valid action set will include
        doing nothing to the current molecule, i.e., the current molecule itself
        will be added to the action set.
      replay_buffer_size: Integer. The size of the replay buffer.
      learning_rate: Float. Learning rate.
      learning_rate_decay_steps: Integer. The number of steps between each
        learning rate decay.
      learning_rate_decay_rate: Float. The rate of learning rate decay.
      num_episodes: Integer. Number of episodes to run.
      batch_size: Integer. The batch size.
      learning_frequency: Integer. The number of steps between each training
        operation.
      update_frequency: Integer. The number of steps between each update of the
        target Q network
      grad_clipping: Integer. maximum value of the gradient norm.
      gamma: Float. The discount factor for the reward.
      double_q: Boolean. Whether to used double Q learning.
        See https://arxiv.org/abs/1509.06461 for detail.
      bootstrap: Integer. The number of bootstrap heads. See
        https://arxiv.org/abs/1703.07608 for detail.
      prioritized: Boolean. Whether to use prioritized replay. See
        https://arxiv.org/abs/1511.05952 for detail.
      prioritized_alpha: Float. The parameter alpha in the prioritized replay.
      prioritized_beta: Float. The parameter beta in the prioritized replay.
      prioritized_epsilon: Float. The parameter epsilon in the prioritized replay.
      fingerprint_radius: Integer. The radius of the Morgan fingerprint.
      fingerprint_length: Integer. The length of the Morgan fingerprint.
      dense_layers: List of integers. The hidden units in the dense layers.
      activation: String. The activation function to use.
      optimizer: String. The optimizer to use.
      batch_norm: Boolean. Whether to use batch normalization.
      save_frequency: Integer. The number of episodes between each saving.
  
    Returns:
      A HParams object containing all the hyperparameters.
    """
    hparams = contrib_training.HParams(
        atom_types=['C', 'O', 'N'],
        max_steps_per_episode=40,
        allow_removal=True,
        allow_no_modification=True,
        allow_bonds_between_rings=False,
        allowed_ring_sizes=[3, 4, 5, 6],
        replay_buffer_size=1000000,
        learning_rate=1e-4,
        learning_rate_decay_steps=10000,
        learning_rate_decay_rate=0.8,
        num_episodes=5000,
        batch_size=64,
        learning_frequency=4,
        update_frequency=20,
        grad_clipping=10.0,
        gamma=0.9,
        double_q=True,
        num_bootstrap_heads=12,
        prioritized=False,
        prioritized_alpha=0.6,
        prioritized_beta=0.4,
        prioritized_epsilon=1e-6,
        fingerprint_radius=3,
        fingerprint_length=2048,
        dense_layers=[1024, 512, 128, 32],
        activation='relu',
        optimizer='Adam',
        batch_norm=False,
        save_frequency=1000,
        max_num_checkpoints=100,
        discount_factor=0.7)
    return hparams.override_from_dict(kwargs)
  
  
  def get_fingerprint(smiles, hparams):
    """Get Morgan Fingerprint of a specific SMILES string.
  
    Args:
      smiles: String. The SMILES string of the molecule.
      hparams: tf.contrib.training.HParams. Hyper parameters.
  
    Returns:
      np.array. shape = [hparams.fingerprint_length]. The Morgan fingerprint.
    """
    if smiles is None:
      return np.zeros((hparams.fingerprint_length,))
    molecule = Chem.MolFromSmiles(smiles)
    if molecule is None:
      return np.zeros((hparams.fingerprint_length,))
    fingerprint = AllChem.GetMorganFingerprintAsBitVect(
        molecule, hparams.fingerprint_radius, hparams.fingerprint_length)
    arr = np.zeros((1,))
    # ConvertToNumpyArray takes ~ 0.19 ms, while
    # np.asarray takes ~ 4.69 ms
    DataStructs.ConvertToNumpyArray(fingerprint, arr)
    return arr
  
  
  def get_fingerprint_with_steps_left(smiles, steps_left, hparams):
    """Get Morgan Fingerprint of a SMILES string with number of steps left.
  
    If fixing the max num of steps can be taken in a MDP, the MDP is then
    a time-heterogeneous one. Therefore a time dependent policy is needed
    for optimal performance.
  
    Args:
      smiles: String. The SMILES string of the molecule.
      steps_left: Integer. The number of steps left in the environment.
      hparams: tf.contrib.training.HParams. Hyper parameters.
  
    Returns:
      np.array. shape = [hparams.fingerprint_length + 1]. The fingerprint.
    """
    fingerprint = get_fingerprint(smiles, hparams)
    return np.append(fingerprint, steps_left)
  