docstr: Fintuning model using reinforce algorithm
  Args:
      existing_mol: molecule to be optimized. Practically, we provide 64 mols per call and the function may take less then 64 mols
      temperature: generation temperature
      batch_size: batch_size for collecting data
      max_size_rl: maximal num of atoms allowed for generation
  
  Returns:
docstr_tok: ["Fintuning","model","using","reinforce","algorithm","Args",":","existing_mol",":","molecule","to","be","optimized",".","Practically",",","we","provide","64","mols","per","call","and","the","function","may","take","less","then","64","mols","temperature",":","generation","temperature","batch_size",":","batch_size","for","collecting","data","max_size_rl",":","maximal","num","of","atoms","allowed","for","generation","Returns",":"]
code: def reinforce_forward_constrained_optim(
      self,
      mol_xs,
      mol_adjs,
      mol_sizes,
      raw_smiles,
      bfs_perm_origin,
      in_baseline=None,
      cur_iter=None,
  ):
      """
      Fintuning model using reinforce algorithm
      Args:
          existing_mol: molecule to be optimized. Practically, we provide 64 mols per call and the function may take less then 64 mols
          temperature: generation temperature
          batch_size: batch_size for collecting data
          max_size_rl: maximal num of atoms allowed for generation
  
      Returns:
  
      """
      assert cur_iter is not None
      atom_list, temperature, batch_size, max_size_rl = (
          self.conf_rl["atom_list"],
          self.conf_rl["temperature"],
          self.conf_rl["batch_size"],
          self.conf_rl["max_size_rl"],
      )
      batch_size = min(
          batch_size, mol_sizes.size(0)
      )  # last batch of one iter may be less batch_size
      if (
          cur_iter % self.conf_rl["update_iters"] == 0
      ):  # uodate the demenstration net every 4 iter.
          print("copying to old model at iter {}".format(cur_iter))
          self.flow_core_old.load_state_dict(
              self.flow_core.state_dict()
          )
          self.node_base_log_probs_old = nn.Parameter(
              self.node_base_log_probs.detach().clone(),
              requires_grad=False,
          )
          self.edge_base_log_probs_old = nn.Parameter(
              self.edge_base_log_probs.detach().clone(),
              requires_grad=False,
          )
  
      # assert cur_baseline is not None
      num2bond = {
          0: Chem.rdchem.BondType.SINGLE,
          1: Chem.rdchem.BondType.DOUBLE,
          2: Chem.rdchem.BondType.TRIPLE,
      }
      num2atom = {i: atom_list[i] for i in range(len(atom_list))}
  
      node_inputs = {}
      node_inputs["node_features"] = []
      node_inputs["adj_features"] = []
      node_inputs["node_features_cont"] = []
      node_inputs["rewards"] = []
      node_inputs["baseline_index"] = []
      node_inputs["node_cnt"] = []
  
      adj_inputs = {}
      adj_inputs["node_features"] = []
      adj_inputs["adj_features"] = []
      adj_inputs["edge_features_cont"] = []
      adj_inputs["index"] = []
      adj_inputs["rewards"] = []
      adj_inputs["baseline_index"] = []
      adj_inputs["edge_cnt"] = []
  
      reward_baseline = torch.zeros([max_size_rl + 5, 2]).cuda()
  
      max_action_size = batch_size * (
          int(
              max_size_rl
              + (self.edge_unroll - 1) * self.edge_unroll / 2
              + (max_size_rl - self.edge_unroll) * self.edge_unroll
          )
      )
  
      batch_length = 0
      total_node_step = 0
      total_edge_step = 0
  
      per_mol_reward = []
      per_mol_property_score = []
  
      ### gather training data from generation
      self.eval()  #! very important. Because we use batch normalization, training mode will result in unrealistic molecules
  
      with torch.no_grad():
          while (
              total_node_step + total_edge_step < max_action_size
              and batch_length < batch_size
          ):
              traj_node_inputs = {}
              traj_node_inputs["node_features"] = []
              traj_node_inputs["adj_features"] = []
              traj_node_inputs["node_features_cont"] = []
              traj_node_inputs["rewards"] = []
              traj_node_inputs["baseline_index"] = []
              traj_node_inputs["node_cnt"] = []
              traj_adj_inputs = {}
              traj_adj_inputs["node_features"] = []
              traj_adj_inputs["adj_features"] = []
              traj_adj_inputs["edge_features_cont"] = []
              traj_adj_inputs["index"] = []
              traj_adj_inputs["rewards"] = []
              traj_adj_inputs["baseline_index"] = []
              traj_adj_inputs["edge_cnt"] = []
  
              step_cnt = 1.0
  
              flag_reconstruct_from_node_adj = True
  
              rand_num = np.random.rand()
              if rand_num <= 0.5:
                  cur_modify_size = np.random.randint(
                      low=0, high=self.conf_rl["modify_size"]
                  )
              else:
                  cur_modify_size = 0
  
              keep_size = int(mol_sizes[batch_length]) - cur_modify_size
  
              org_bfs_perm_origin = bfs_perm_origin[batch_length]
              org_node_features = mol_xs[
                  batch_length : batch_length + 1
              ]  # (1, 38, 9)
              org_adj_features = mol_adjs[
                  batch_length : batch_length + 1
              ]  # (1, 4, 38, 38)
  
              cur_node_features = mol_xs[
                  batch_length : batch_length + 1
              ].clone()  # (1, 38, 9)
              cur_adj_features = mol_adjs[
                  batch_length : batch_length + 1
              ].clone()  # (1, 4, 38, 38)
  
              cur_node_features[:, keep_size:, :] = 0.0
              cur_adj_features[:, :, keep_size:, :] = 0.0
              cur_adj_features[:, :, :, keep_size:] = 0.0
  
              min_action_node = 1
  
              rw_mol = Chem.RWMol()  # editable mol
              mol = None
  
              for i in range(mol_sizes[batch_length]):
                  node_id = torch.argmax(
                      org_node_features[0, i]
                  ).item()  # (9, )
                  if i < keep_size:
                      rw_mol.AddAtom(Chem.Atom(num2atom[node_id]))
                  else:
                      pass
  
                  if i < self.edge_unroll:
                      edge_total = i  # edge to sample for current node
                      start = 0
                  else:
                      edge_total = self.edge_unroll
                      start = i - self.edge_unroll
  
                  for j in range(edge_total):
                      edge_id = torch.argmax(
                          org_adj_features[0, :, i, j + start]
                      ).item()  # (4, )
                      if edge_id == 3:  # no bond to add
                          continue
                      if i < keep_size:
                          rw_mol.AddBond(
                              i, j + start, num2bond[edge_id]
                          )
                      else:
                          pass
  
              mol = rw_mol.GetMol()
              s_raw = raw_smiles[batch_length]
              org_mol_true_raw = Chem.MolFromSmiles(s_raw)
  
              if (
                  check_chemical_validity(mol) is False
              ):  # do not use subgraph, use original graph
                  rw_mol = Chem.RWMol(org_mol_true_raw)
                  mol = rw_mol.GetMol()
                  flag_reconstruct_from_node_adj = (
                      False  # bfs is not reflected in current mol
                  )
                  keep_size = int(mol_sizes[batch_length])
                  cur_node_features = org_node_features.clone()
                  cur_adj_features = org_adj_features.clone()
                  cur_node_features[:, keep_size:, :] = 0.0
                  cur_adj_features[:, :, keep_size:, :] = 0.0
                  cur_adj_features[:, :, :, keep_size:] = 0.0
  
              assert (
                  check_chemical_validity(org_mol_true_raw) is True
              ), "s_raw is %s" % (s_raw)
              assert check_chemical_validity(mol) is True
  
              assert mol.GetNumAtoms() == keep_size
              assert (
                  org_mol_true_raw.GetNumAtoms()
                  == mol_sizes[batch_length]
              )
  
              if self.dp:
                  cur_node_features = cur_node_features.cuda()
                  cur_adj_features = cur_adj_features.cuda()
  
              is_continue = True
              if keep_size <= self.edge_unroll:
                  edge_idx = int(keep_size * (keep_size - 1) / 2)
              else:
                  edge_idx = int(
                      self.edge_unroll * (self.edge_unroll - 1) / 2
                      + (keep_size - self.edge_unroll)
                      * self.edge_unroll
                  )
  
              step_num_data_edge = 0
              added_num = 0
              node_features_each_iter_backup = cur_node_features.clone()
              adj_features_each_iter_backup = cur_adj_features.clone()
              for i in range(keep_size, max_size_rl):
                  if not is_continue:
                      break
  
                  step_num_data_edge = 0  # generating new node and its edges. Not sure if this will add into the final mol.
  
                  if i < self.edge_unroll:
                      edge_total = i  # edge to sample for current node
                      start = 0
                  else:
                      edge_total = self.edge_unroll
                      start = i - self.edge_unroll
  
                  # first generate node
                  ## reverse flow
                  prior_node_dist = (
                      torch.distributions.OneHotCategorical(
                          logits=self.node_base_log_probs_old[i]
                          * temperature[0]
                      )
                  )
                  latent_node = prior_node_dist.sample().view(1, -1)
  
                  if self.dp:
                      latent_node = self.flow_core_old.module.reverse(
                          cur_node_features,
                          cur_adj_features,
                          latent_node,
                          mode=0,
                      ).view(
                          -1
                      )  # (9, )
                  else:
                      latent_node = self.flow_core_old.reverse(
                          cur_node_features,
                          cur_adj_features,
                          latent_node,
                          mode=0,
                      ).view(
                          -1
                      )  # (9, )
  
                  feature_id = torch.argmax(latent_node).item()
                  total_node_step += 1
                  node_feature_cont = torch.zeros(
                      [1, self.node_dim]
                  ).cuda()
                  node_feature_cont[0, feature_id] = 1.0
  
                  # update traj inputs for node_id
                  traj_node_inputs["node_features"].append(
                      cur_node_features.clone()
                  )  # (1, max_size_rl, self.node_dim)
                  traj_node_inputs["adj_features"].append(
                      cur_adj_features.clone()
                  )  # (1, self.bond_dim, max_size_rl, max_size_rl)
                  traj_node_inputs["node_features_cont"].append(
                      node_feature_cont
                  )  # (1, self.node_dim)
                  traj_node_inputs["rewards"].append(
                      torch.full(
                          size=(1, 1), fill_value=step_cnt
                      ).cuda()
                  )  # (1, 1)
                  traj_node_inputs["baseline_index"].append(
                      torch.full(size=(1,), fill_value=step_cnt)
                      .long()
                      .cuda()
                  )  # (1, 1)
                  traj_node_inputs["node_cnt"].append(
                      torch.full(size=(1,), fill_value=float(i))
                      .long()
                      .cuda()
                  )
  
                  cur_node_features[0, i, feature_id] = 1.0
                  cur_adj_features[0, :, i, i] = 1.0
                  rw_mol.AddAtom(Chem.Atom(num2atom[feature_id]))
  
                  # then generate edges
                  if i == 0:
                      is_connect = True
                  else:
                      is_connect = False
                  # cur_mol_size = mol.GetNumAtoms
                  for j in range(edge_total):
                      valid = False
                      resample_edge = 0
                      edge_dis = self.edge_base_log_probs_old[
                          edge_idx
                      ].clone()
                      invalid_bond_type_set = set()
                      while not valid:
                          if (
                              len(invalid_bond_type_set) < 3
                              and resample_edge <= 50
                          ):  # haven't sampled all possible bond type or is not stuck in the loop
                              prior_edge_dist = (
                                  torch.distributions.OneHotCategorical(
                                      logits=edge_dis / temperature[1]
                                  )
                              )
                              latent_edge = (
                                  prior_edge_dist.sample().view(1, -1)
                              )
                              latent_id = torch.argmax(
                                  latent_edge, dim=1
                              )
  
                              if self.dp:
                                  latent_edge = (
                                      self.flow_core_old.module.reverse(
                                          cur_node_features,
                                          cur_adj_features,
                                          latent_edge,
                                          mode=1,
                                          edge_index=torch.Tensor(
                                              [[j + start, i]]
                                          )
                                          .long()
                                          .cuda(),
                                      ).view(-1)
                                  )  # (4, )
                              else:
                                  latent_edge = (
                                      self.flow_core_old.reverse(
                                          cur_node_features,
                                          cur_adj_features,
                                          latent_edge,
                                          mode=1,
                                          edge_index=torch.Tensor(
                                              [[j + start, i]]
                                          )
                                          .long()
                                          .cuda(),
                                      ).view(-1)
                                  )  # (4, )
                              edge_discrete_id = torch.argmax(
                                  latent_edge
                              ).item()
                          else:
                              assert (
                                  resample_edge > 50
                                  or len(invalid_bond_type_set) == 3
                              )
                              edge_discrete_id = 3  # we have no choice but to choose not to add edge between (i, j+start)
  
                          total_edge_step += 1
                          edge_feature_cont = torch.zeros(
                              [1, self.bond_dim]
                          ).cuda()
                          edge_feature_cont[0, edge_discrete_id] = 1.0
  
                          # update traj inputs for edge_id
                          traj_adj_inputs["node_features"].append(
                              cur_node_features.clone()
                          )  # 1, max_size_rl, self.node_dim
                          traj_adj_inputs["adj_features"].append(
                              cur_adj_features.clone()
                          )  # 1, self.bond_dim, max_size_rl, max_size_rl
                          traj_adj_inputs["edge_features_cont"].append(
                              edge_feature_cont
                          )  # 1, self.bond_dim
                          traj_adj_inputs["index"].append(
                              torch.Tensor([[j + start, i]])
                              .long()
                              .cuda()
                              .view(1, -1)
                          )  # (1, 2)
                          traj_adj_inputs["edge_cnt"].append(
                              torch.full(
                                  size=(1,), fill_value=float(edge_idx)
                              )
                              .long()
                              .cuda()
                          )
                          step_num_data_edge += 1  # add one edge data, not sure if this should be added to the final train data
  
                          cur_adj_features[
                              0, edge_discrete_id, i, j + start
                          ] = 1.0
                          cur_adj_features[
                              0, edge_discrete_id, j + start, i
                          ] = 1.0
                          if edge_discrete_id == 3:  # virtual edge
                              valid = (
                                  True  # virtual edge is alway valid
                              )
                          else:  # single/double/triple bond
                              if flag_reconstruct_from_node_adj:
                                  rw_mol.AddBond(
                                      i,
                                      j + start,
                                      num2bond[edge_discrete_id],
                                  )
                              else:
                                  # current rw_mol is copy from original_mol, which does not reflect the bfs order
                                  # so we use bfs_perm_origin to map j+start to the true node_id in rw_mol
                                  rw_mol.AddBond(
                                      i,
                                      int(
                                          org_bfs_perm_origin[
                                              j + start
                                          ].item()
                                      ),
                                      num2bond[edge_discrete_id],
                                  )
                              valid = check_valency(rw_mol)
                              if valid:
                                  is_connect = True
                              else:  # backtrack
                                  edge_dis[latent_id] = float("-inf")
                                  if flag_reconstruct_from_node_adj:
                                      rw_mol.RemoveBond(i, j + start)
                                  else:
                                      rw_mol.RemoveBond(
                                          i,
                                          int(
                                              org_bfs_perm_origin[
                                                  j + start
                                              ].item()
                                          ),
                                      )
                                  cur_adj_features[
                                      0, edge_discrete_id, i, j + start
                                  ] = 0.0
                                  cur_adj_features[
                                      0, edge_discrete_id, j + start, i
                                  ] = 0.0
                                  resample_edge += 1
  
                                  invalid_bond_type_set.add(
                                      edge_discrete_id
                                  )
  
                          if valid:
                              traj_adj_inputs["rewards"].append(
                                  torch.full(
                                      size=(1, 1), fill_value=step_cnt
                                  ).cuda()
                              )  # (1, 1)
                              traj_adj_inputs["baseline_index"].append(
                                  torch.full(
                                      size=(1,), fill_value=step_cnt
                                  )
                                  .long()
                                  .cuda()
                              )  # (1)
                          else:
                              if self.conf_rl["penalty"]:
                                  traj_adj_inputs["rewards"].append(
                                      torch.full(
                                          size=(1, 1), fill_value=-1.0
                                      ).cuda()
                                  )  # (1, 1) invalid edge penalty
                                  traj_adj_inputs[
                                      "baseline_index"
                                  ].append(
                                      torch.full(
                                          size=(1,), fill_value=step_cnt
                                      )
                                      .long()
                                      .cuda()
                                  )  # (1,)
                              else:
                                  traj_adj_inputs["node_features"].pop(
                                      -1
                                  )
                                  traj_adj_inputs["adj_features"].pop(
                                      -1
                                  )
                                  traj_adj_inputs[
                                      "edge_features_cont"
                                  ].pop(-1)
                                  traj_adj_inputs["index"].pop(-1)
                                  traj_adj_inputs["edge_cnt"].pop(-1)
                                  step_num_data_edge -= 1  # if we do not penalize invalid edge, pop train data, decrease counter by 1
  
                      edge_idx += 1
  
                  if (
                      is_connect
                  ):  # new generated node has at least one bond with previous node, do not stop generation, backup mol from rw_mol to mol
                      is_continue = True
                      mol = rw_mol.GetMol()
                      node_features_each_iter_backup = (
                          cur_node_features.clone()
                      )
                      adj_features_each_iter_backup = (
                          cur_adj_features.clone()
                      )
                      added_num += 1
                      # current_i_continue = False
  
                  else:
                      #! we may need to satisfy min action here
                      if added_num >= min_action_node:
                          # if we have already add 'min_action_node', ignore and let continue be false.
                          # the data added in last iter will be pop afterwards
                          is_continue = False
                      else:
                          # added_num is less than min_action_node
                          # first pop the appended train data
                          is_continue = False  # first set as False
                          rw_mol = Chem.RWMol(
                              mol
                          )  # important, recover the mol
                          cur_node_features = (
                              node_features_each_iter_backup.clone()
                          )  # recover the backuped features
                          cur_adj_features = (
                              adj_features_each_iter_backup.clone()
                          )
                          cur_node_features_tmp = (
                              cur_node_features.clone()
                          )
                          cur_adj_features_tmp = (
                              cur_adj_features.clone()
                          )
                          cur_mol_size = rw_mol.GetNumAtoms()
  
                          traj_node_inputs["node_features"].pop(-1)
                          traj_node_inputs["adj_features"].pop(-1)
                          traj_node_inputs["node_features_cont"].pop(-1)
                          traj_node_inputs["rewards"].pop(-1)
                          traj_node_inputs["baseline_index"].pop(-1)
                          traj_node_inputs["node_cnt"].pop(-1)
  
                          ## pop adj
                          for _ in range(step_num_data_edge):
                              traj_adj_inputs["node_features"].pop(-1)
                              traj_adj_inputs["adj_features"].pop(-1)
                              traj_adj_inputs["edge_features_cont"].pop(
                                  -1
                              )
                              traj_adj_inputs["index"].pop(-1)
                              traj_adj_inputs["rewards"].pop(-1)
                              traj_adj_inputs["baseline_index"].pop(-1)
                              traj_adj_inputs["edge_cnt"].pop(-1)
  
                          mol_demon_edit = Chem.RWMol(rw_mol)
                          last_id2 = mol_demon_edit.AddAtom(
                              Chem.Atom(6)
                          )
                          traj_node_inputs["node_features"].append(
                              cur_node_features_tmp.clone()
                          )
                          traj_node_inputs["adj_features"].append(
                              cur_adj_features_tmp.clone()
                          )
                          node_feature_cont = torch.zeros(
                              [1, self.node_dim]
                          ).cuda()
                          node_feature_cont[0, feature_id] = 1.0
                          traj_node_inputs["node_features_cont"].append(
                              node_feature_cont
                          )
                          traj_node_inputs["rewards"].append(
                              torch.full(
                                  size=(1, 1), fill_value=step_cnt
                              ).cuda()
                          )  # (1, 1)
                          traj_node_inputs["baseline_index"].append(
                              torch.full(size=(1,), fill_value=step_cnt)
                              .long()
                              .cuda()
                          )  # (1, 1)
                          traj_node_inputs["node_cnt"].append(
                              torch.full(size=(1,), fill_value=float(i))
                              .long()
                              .cuda()
                          )
  
                          cur_node_features_tmp[0, last_id2, 0] = 1.0
                          cur_adj_features_tmp[
                              0, :, last_id2, last_id2
                          ] = 1.0
  
                          flag_success = False
                          if cur_mol_size <= self.edge_unroll:
                              edge_total = cur_mol_size
                              start = 0
                              edge_idx = int(
                                  cur_mol_size * (cur_mol_size - 1) / 2
                              )
                          else:
                              edge_total = self.edge_unroll
                              start = cur_mol_size - self.edge_unroll
                              edge_idx = int(
                                  self.edge_unroll
                                  * (self.edge_unroll - 1)
                                  / 2
                                  + (cur_mol_size - self.edge_unroll)
                                  * self.edge_unroll
                              )
  
                          if cur_mol_size > keep_size:
                              mol_demon_edit.AddBond(
                                  cur_mol_size - 1,
                                  last_id2,
                                  Chem.rdchem.BondType.SINGLE,
                              )
                              valid = check_valency(mol_demon_edit)
                              if valid:
                                  flag_success = True
                                  for j in range(edge_total):
                                      edge_feature_cont = torch.zeros(
                                          [1, self.bond_dim]
                                      ).cuda()
                                      if j == edge_total - 1:
                                          edge_feature_cont[0, 0] = 1.0
                                      else:
                                          edge_feature_cont[0, 3] = 1.0
  
                                      traj_adj_inputs[
                                          "node_features"
                                      ].append(
                                          cur_node_features_tmp.clone()
                                      )
                                      traj_adj_inputs[
                                          "adj_features"
                                      ].append(
                                          cur_adj_features_tmp.clone()
                                      )
                                      traj_adj_inputs[
                                          "edge_features_cont"
                                      ].append(edge_feature_cont)
                                      traj_adj_inputs["index"].append(
                                          torch.Tensor(
                                              [[j + start, last_id2]]
                                          )
                                          .long()
                                          .cuda()
                                          .view(1, -1)
                                      )
                                      traj_adj_inputs[
                                          "edge_cnt"
                                      ].append(
                                          torch.full(
                                              size=(1,),
                                              fill_value=float(
                                                  edge_idx
                                              ),
                                          )
                                          .long()
                                          .cuda()
                                      )
                                      traj_adj_inputs["rewards"].append(
                                          torch.full(
                                              size=(1, 1),
                                              fill_value=step_cnt,
                                          ).cuda()
                                      )  # (1, 1)
                                      traj_adj_inputs[
                                          "baseline_index"
                                      ].append(
                                          torch.full(
                                              size=(1,),
                                              fill_value=step_cnt,
                                          )
                                          .long()
                                          .cuda()
                                      )  # (1)
  
                                      if j == edge_total - 1:
                                          cur_adj_features_tmp[
                                              0, 0, last_id2, j + start
                                          ] = 1.0
                                          cur_adj_features_tmp[
                                              0, 0, j + start, last_id2
                                          ] = 1.0
                                      else:
                                          cur_adj_features_tmp[
                                              0, 3, last_id2, j + start
                                          ] = 1.0
                                          cur_adj_features_tmp[
                                              0, 3, j + start, last_id2
                                          ] = 1.0
  
                                      edge_idx += 1
                              else:
                                  mol_demon_edit.RemoveBond(
                                      cur_mol_size - 1, last_id2
                                  )
  
                                  count = 0
                                  while True:
                                      if count > 100:
                                          break
                                      if last_id2 > 12:
                                          k = np.random.randint(1, 13)
                                      else:
                                          k = np.random.randint(
                                              1, last_id2 + 1
                                          )
                                      if flag_reconstruct_from_node_adj:
                                          mol_demon_edit.AddBond(
                                              int(last_id2 - k),
                                              int(last_id2),
                                              Chem.rdchem.BondType.SINGLE,
                                          )
                                      else:
                                          mol_demon_edit.AddBond(
                                              int(
                                                  org_bfs_perm_origin[
                                                      last_id2 - k
                                                  ]
                                              ),
                                              int(last_id2),
                                              Chem.rdchem.BondType.SINGLE,
                                          )
  
                                      valid = check_valency(
                                          mol_demon_edit
                                      )
                                      if valid:
                                          flag_success = True
                                          for j in range(edge_total):
                                              edge_feature_cont = (
                                                  torch.zeros(
                                                      [1, self.bond_dim]
                                                  ).cuda()
                                              )
                                              if (
                                                  j + start
                                                  == last_id2 - k
                                              ):
                                                  edge_feature_cont[
                                                      0, 0
                                                  ] = 1.0
                                              else:
                                                  edge_feature_cont[
                                                      0, 3
                                                  ] = 1.0
  
                                              traj_adj_inputs[
                                                  "node_features"
                                              ].append(
                                                  cur_node_features_tmp.clone()
                                              )
                                              traj_adj_inputs[
                                                  "adj_features"
                                              ].append(
                                                  cur_adj_features_tmp.clone()
                                              )
                                              traj_adj_inputs[
                                                  "edge_features_cont"
                                              ].append(
                                                  edge_feature_cont
                                              )
                                              traj_adj_inputs[
                                                  "index"
                                              ].append(
                                                  torch.Tensor(
                                                      [
                                                          [
                                                              j + start,
                                                              last_id2
                                                              - k,
                                                          ]
                                                      ]
                                                  )
                                                  .long()
                                                  .cuda()
                                                  .view(1, -1)
                                              )
                                              traj_adj_inputs[
                                                  "edge_cnt"
                                              ].append(
                                                  torch.full(
                                                      size=(1,),
                                                      fill_value=float(
                                                          edge_idx
                                                      ),
                                                  )
                                                  .long()
                                                  .cuda()
                                              )
                                              traj_adj_inputs[
                                                  "rewards"
                                              ].append(
                                                  torch.full(
                                                      size=(1, 1),
                                                      fill_value=step_cnt,
                                                  ).cuda()
                                              )  # (1, 1)
                                              traj_adj_inputs[
                                                  "baseline_index"
                                              ].append(
                                                  torch.full(
                                                      size=(1,),
                                                      fill_value=step_cnt,
                                                  )
                                                  .long()
                                                  .cuda()
                                              )  # (1)
  
                                              if (
                                                  j + start
                                                  == last_id2 - k
                                              ):
                                                  cur_adj_features_tmp[
                                                      0,
                                                      0,
                                                      last_id2,
                                                      last_id2 - k,
                                                  ] = 1.0
                                                  cur_adj_features_tmp[
                                                      0,
                                                      0,
                                                      last_id2 - k,
                                                      last_id2,
                                                  ] = 1.0
                                              else:
                                                  cur_adj_features_tmp[
                                                      0,
                                                      3,
                                                      last_id2,
                                                      last_id2 - k,
                                                  ] = 1.0
                                                  cur_adj_features_tmp[
                                                      0,
                                                      3,
                                                      last_id2 - k,
                                                      last_id2,
                                                  ] = 1.0
  
                                              edge_idx += 1
                                          break
                                      else:
                                          if flag_reconstruct_from_node_adj:
                                              mol_demon_edit.RemoveBond(
                                                  int(last_id2 - k),
                                                  int(last_id2),
                                              )
                                          else:
                                              mol_demon_edit.RemoveBond(
                                                  int(
                                                      org_bfs_perm_origin[
                                                          last_id2 - k
                                                      ]
                                                  ),
                                                  int(last_id2),
                                              )
                                          count += 1
                          else:
                              count = 0
                              while True:
                                  if count > 100:
                                      break
                                  if keep_size > 12:
                                      k = np.random.randint(1, 13)
                                  else:
                                      k = np.random.randint(
                                          1, keep_size + 1
                                      )
                                  if flag_reconstruct_from_node_adj:
                                      mol_demon_edit.AddBond(
                                          int(keep_size - k),
                                          int(keep_size),
                                          Chem.rdchem.BondType.SINGLE,
                                      )
                                  else:
                                      mol_demon_edit.AddBond(
                                          int(
                                              org_bfs_perm_origin[
                                                  keep_size - k
                                              ]
                                          ),
                                          int(keep_size),
                                          Chem.rdchem.BondType.SINGLE,
                                      )
                                  cur_adj_features_tmp[
                                      0, 0, keep_size - k, keep_size
                                  ] = 1.0
                                  cur_adj_features_tmp[
                                      0, 0, keep_size, keep_size - k
                                  ] = 1.0
  
                                  valid = check_valency(mol_demon_edit)
                                  if valid:
                                      flag_success = True
                                      for j in range(edge_total):
                                          edge_feature_cont = (
                                              torch.zeros(
                                                  [1, self.bond_dim]
                                              ).cuda()
                                          )
                                          if j + start == last_id2 - k:
                                              edge_feature_cont[
                                                  0, 0
                                              ] = 1.0
                                          else:
                                              edge_feature_cont[
                                                  0, 3
                                              ] = 1.0
  
                                          traj_adj_inputs[
                                              "node_features"
                                          ].append(
                                              cur_node_features_tmp.clone()
                                          )
                                          traj_adj_inputs[
                                              "adj_features"
                                          ].append(
                                              cur_adj_features_tmp.clone()
                                          )
                                          traj_adj_inputs[
                                              "edge_features_cont"
                                          ].append(edge_feature_cont)
                                          traj_adj_inputs[
                                              "index"
                                          ].append(
                                              torch.Tensor(
                                                  [
                                                      [
                                                          j + start,
                                                          last_id2 - k,
                                                      ]
                                                  ]
                                              )
                                              .long()
                                              .cuda()
                                              .view(1, -1)
                                          )
                                          traj_adj_inputs[
                                              "edge_cnt"
                                          ].append(
                                              torch.full(
                                                  size=(1,),
                                                  fill_value=float(
                                                      edge_idx
                                                  ),
                                              )
                                              .long()
                                              .cuda()
                                          )
                                          traj_adj_inputs[
                                              "rewards"
                                          ].append(
                                              torch.full(
                                                  size=(1, 1),
                                                  fill_value=step_cnt,
                                              ).cuda()
                                          )  # (1, 1)
                                          traj_adj_inputs[
                                              "baseline_index"
                                          ].append(
                                              torch.full(
                                                  size=(1,),
                                                  fill_value=step_cnt,
                                              )
                                              .long()
                                              .cuda()
                                          )  # (1)
  
                                          if j + start == last_id2 - k:
                                              cur_adj_features_tmp[
                                                  0,
                                                  0,
                                                  last_id2,
                                                  last_id2 - k,
                                              ] = 1.0
                                              cur_adj_features_tmp[
                                                  0,
                                                  0,
                                                  last_id2 - k,
                                                  last_id2,
                                              ] = 1.0
                                          else:
                                              cur_adj_features_tmp[
                                                  0,
                                                  3,
                                                  last_id2,
                                                  last_id2 - k,
                                              ] = 1.0
                                              cur_adj_features_tmp[
                                                  0,
                                                  3,
                                                  last_id2 - k,
                                                  last_id2,
                                              ] = 1.0
  
                                          edge_idx += 1
                                      break
                                  else:
                                      if flag_reconstruct_from_node_adj:
                                          mol_demon_edit.RemoveBond(
                                              int(keep_size - k),
                                              int(keep_size),
                                          )
                                      else:
                                          mol_demon_edit.RemoveBond(
                                              int(
                                                  org_bfs_perm_origin[
                                                      keep_size - k
                                                  ]
                                              ),
                                              int(keep_size),
                                          )
                                      count += 1
  
                          if flag_success and check_chemical_validity(
                              mol_demon_edit
                          ):  # successfully take one min action.
                              rw_mol = Chem.RWMol(mol_demon_edit)
                              cur_node_features = (
                                  cur_node_features_tmp.clone()
                              )
                              cur_adj_features = (
                                  cur_adj_features_tmp.clone()
                              )
                              is_continue = True
                              mol = rw_mol.GetMol()
                              node_features_each_iter_backup = (
                                  cur_node_features.clone()
                              )  # update node backup since new node is valid
                              adj_features_each_iter_backup = (
                                  cur_adj_features.clone()
                              )  #
                              added_num += 1
                          else:
                              traj_node_inputs["node_features"].pop(-1)
                              traj_node_inputs["adj_features"].pop(-1)
                              traj_node_inputs[
                                  "node_features_cont"
                              ].pop(-1)
                              traj_node_inputs["rewards"].pop(-1)
                              traj_node_inputs["baseline_index"].pop(-1)
                              traj_node_inputs["node_cnt"].pop(-1)
                              step_cnt += 1
                              continue
                  step_cnt += 1
  
              batch_length += 1
  
              # TODO: check the last iter of generation
              # (Thinking)
              # The last node was not added. So after we generate the second to last
              # node and all its edges, the rest adjacent matrix and node features should all be zero
              # But current implementation append
              num_atoms = mol.GetNumAtoms()
              assert num_atoms <= max_size_rl
  
              if num_atoms < max_size_rl:
                  #! this implementation is buggy. we only mask the last node feature cont
                  #! But we ignore the non-zero node features in generating edges
                  #! this pattern will make model not to generated any edges between
                  #! the new-generated isolated node and exsiting subgraph.
                  #! this may be the biggest bug in Reinforce algorithm!!!!!
                  #! since the final iter/(step) has largest reward....!!!!!!!
                  #! work around1: add a counter and mask out all node feautres in generating edges of last iter.
                  #! work around2: do not append any data if the isolated node is not connected to subgraph.
                  # currently use work around2
  
                  # pop all the reinforce train-data add by at the generating the last isolated node and its edge
                  ## pop node
                  try:
                      traj_node_inputs["node_features"].pop(-1)
                      traj_node_inputs["adj_features"].pop(-1)
                      traj_node_inputs["node_features_cont"].pop(-1)
                      traj_node_inputs["rewards"].pop(-1)
                      traj_node_inputs["baseline_index"].pop(-1)
                      traj_node_inputs["node_cnt"].pop(-1)
  
                      ## pop adj
                      for pop_cnt in range(step_num_data_edge):
                          traj_adj_inputs["node_features"].pop(-1)
                          traj_adj_inputs["adj_features"].pop(-1)
                          traj_adj_inputs["edge_features_cont"].pop(-1)
                          traj_adj_inputs["index"].pop(-1)
                          traj_adj_inputs["rewards"].pop(-1)
                          traj_adj_inputs["baseline_index"].pop(-1)
                          traj_adj_inputs["edge_cnt"].pop(-1)
                  except:
                      print(
                          "pop from empty list, take min action fail."
                      )
  
              reward_valid = 2
              reward_property = 0
              reward_length = 0
  
              assert mol is not None, "mol is None..."
              final_valid = check_chemical_validity(mol)
              s_tmp = Chem.MolToSmiles(mol, isomericSmiles=True)
              assert final_valid is True, (
                  "warning: use valency check during generation but the final molecule is invalid!!!, \
                                  raw is %s, cur is %s"
                  % (s_raw, s_tmp)
              )
  
              if not final_valid:
                  reward_valid -= 5
              else:
                  final_mol = convert_radical_electrons_to_hydrogens(
                      mol
                  )
                  s = Chem.MolToSmiles(final_mol, isomericSmiles=True)
  
                  final_mol = Chem.MolFromSmiles(s)
                  # mol filters with negative rewards
                  if not steric_strain_filter(
                      final_mol
                  ):  # passes 3D conversion, no excessive strain
                      reward_valid -= (
                          1  # TODO: check the magnitude of this reward.
                      )
                  if not zinc_molecule_filter(
                      final_mol
                  ):  # does not contain any problematic functional groups
                      reward_valid -= 1
  
                  try:
                      score = calculate_min_plogp(final_mol)
                      score_raw = calculate_min_plogp(org_mol_true_raw)
                      if self.conf_rl["reward_type"] == "exp":
                          reward_property += (
                              np.exp(
                                  score
                                  / self.conf_rl["exp_temperature"]
                              )
                              - self.conf_rl["exp_bias"]
                          )
                      elif self.conf_rl["reward_type"] == "linear":
                          reward_property += (
                              score * self.conf_rl["linear_coeff"]
                          )
                      elif self.conf_rl["reward_type"] == "imp":
                          reward_property += score - score_raw
                  except:
                      print("generated mol does not pass qed/plogp")
  
              reward_final_total = (
                  reward_valid + reward_property + reward_length
              )
              per_mol_reward.append(reward_final_total)
              per_mol_property_score.append(reward_property)
              reward_decay = self.conf_rl["reward_decay"]
  
              if added_num > 0:
                  node_inputs["node_features"].append(
                      torch.cat(
                          traj_node_inputs["node_features"], dim=0
                      )
                  )  # append tensor of shape (max_size_rl, max_size_rl, self.node_dim)
                  node_inputs["adj_features"].append(
                      torch.cat(traj_node_inputs["adj_features"], dim=0)
                  )  # append tensor of shape (max_size_rl, bond_dim, max_size_rl, max_size_rl)
                  node_inputs["node_features_cont"].append(
                      torch.cat(
                          traj_node_inputs["node_features_cont"], dim=0
                      )
                  )  # append tensor of shape (max_size_rl, 9)
  
                  traj_node_inputs_baseline_index = torch.cat(
                      traj_node_inputs["baseline_index"], dim=0
                  )  # (max_size_rl)
                  traj_node_inputs_rewards = torch.cat(
                      traj_node_inputs["rewards"], dim=0
                  )  # tensor of shape (max_size_rl, 1)
                  traj_node_inputs_rewards[
                      traj_node_inputs_rewards > 0
                  ] = reward_final_total * torch.pow(
                      reward_decay,
                      step_cnt
                      - 1.0
                      - traj_node_inputs_rewards[
                          traj_node_inputs_rewards > 0
                      ],
                  )
                  node_inputs["rewards"].append(
                      traj_node_inputs_rewards
                  )  # append tensor of shape (max_size_rl, 1)
                  node_inputs["baseline_index"].append(
                      traj_node_inputs_baseline_index
                  )
                  node_inputs["node_cnt"].append(
                      torch.cat(traj_node_inputs["node_cnt"], dim=0)
                  )
  
                  for ss in range(traj_node_inputs_rewards.size(0)):
                      reward_baseline[
                          traj_node_inputs_baseline_index[ss]
                      ][0] += 1.0
                      reward_baseline[
                          traj_node_inputs_baseline_index[ss]
                      ][1] += traj_node_inputs_rewards[ss][0]
  
                  adj_inputs["node_features"].append(
                      torch.cat(traj_adj_inputs["node_features"], dim=0)
                  )  # (step, max_size_rl, self.node_dim)
                  adj_inputs["adj_features"].append(
                      torch.cat(traj_adj_inputs["adj_features"], dim=0)
                  )  # (step, bond_dim, max_size_rl, max_size_rl)
                  adj_inputs["edge_features_cont"].append(
                      torch.cat(
                          traj_adj_inputs["edge_features_cont"], dim=0
                      )
                  )  # (step, 4)
                  adj_inputs["index"].append(
                      torch.cat(traj_adj_inputs["index"], dim=0)
                  )  # (step, 2)
                  adj_inputs["edge_cnt"].append(
                      torch.cat(traj_adj_inputs["edge_cnt"], dim=0)
                  )
  
                  traj_adj_inputs_baseline_index = torch.cat(
                      traj_adj_inputs["baseline_index"], dim=0
                  )  # (step)
                  traj_adj_inputs_rewards = torch.cat(
                      traj_adj_inputs["rewards"], dim=0
                  )
                  traj_adj_inputs_rewards[
                      traj_adj_inputs_rewards > 0
                  ] = reward_final_total * torch.pow(
                      reward_decay,
                      step_cnt
                      - 1.0
                      - traj_adj_inputs_rewards[
                          traj_adj_inputs_rewards > 0
                      ],
                  )
                  adj_inputs["rewards"].append(traj_adj_inputs_rewards)
                  adj_inputs["baseline_index"].append(
                      traj_adj_inputs_baseline_index
                  )
  
                  for ss in range(traj_adj_inputs_rewards.size(0)):
                      reward_baseline[
                          traj_adj_inputs_baseline_index[ss]
                      ][0] += 1.0
                      reward_baseline[
                          traj_adj_inputs_baseline_index[ss]
                      ][1] += traj_adj_inputs_rewards[ss][0]
  
      self.flow_core.train()
      for module in self.modules():
          if isinstance(module, torch.nn.modules.BatchNorm1d):
              module.eval()
  
      for i in range(reward_baseline.size(0)):
          if reward_baseline[i, 0] == 0:
              reward_baseline[i, 0] += 1.0
  
      reward_baseline_per_step = (
          reward_baseline[:, 1] / reward_baseline[:, 0]
      )  # (max_size_rl, )
  
      if in_baseline is not None:
          assert in_baseline.size() == reward_baseline_per_step.size()
          reward_baseline_per_step = (
              reward_baseline_per_step
              * (1.0 - self.conf_rl["moving_coeff"])
              + in_baseline * self.conf_rl["moving_coeff"]
          )
  
      node_inputs_node_features = torch.cat(
          node_inputs["node_features"], dim=0
      )  # (total_size, max_size_rl, 9)
      node_inputs_adj_features = torch.cat(
          node_inputs["adj_features"], dim=0
      )  # (total_size, 4, max_size_rl, max_size_rl)
      node_inputs_node_features_cont = torch.cat(
          node_inputs["node_features_cont"], dim=0
      )  # (total_size, 9)
      node_inputs_rewards = torch.cat(
          node_inputs["rewards"], dim=0
      ).view(
          -1
      )  # (total_size,)
      node_inputs_baseline_index = torch.cat(
          node_inputs["baseline_index"], dim=0
      ).long()  # (total_size,)
      node_inputs_baseline = torch.index_select(
          reward_baseline_per_step,
          dim=0,
          index=node_inputs_baseline_index,
      )  # (total_size, )
      node_inputs_node_cnts = torch.cat(node_inputs["node_cnt"], dim=0)
  
      adj_inputs_node_features = torch.cat(
          adj_inputs["node_features"], dim=0
      )  # (total_size, max_size_rl, 9)
      adj_inputs_adj_features = torch.cat(
          adj_inputs["adj_features"], dim=0
      )  # (total_size, 4, max_size_rl, max_size_rl)
      adj_inputs_edge_features_cont = torch.cat(
          adj_inputs["edge_features_cont"], dim=0
      )  # (total_size, 4)
      adj_inputs_index = torch.cat(
          adj_inputs["index"], dim=0
      )  # (total_size, 2)
      adj_inputs_rewards = torch.cat(adj_inputs["rewards"], dim=0).view(
          -1
      )  # (total_size,)
      adj_inputs_baseline_index = torch.cat(
          adj_inputs["baseline_index"], dim=0
      ).long()  # (total_size,)
      adj_inputs_baseline = torch.index_select(
          reward_baseline_per_step,
          dim=0,
          index=adj_inputs_baseline_index,
      )  # (total_size, )
      adj_inputs_edge_cnts = torch.cat(adj_inputs["edge_cnt"], dim=0)
  
      if self.dp:
          node_function = self.flow_core.module.forward_rl_node
          edge_function = self.flow_core.module.forward_rl_edge
  
          node_function_old = self.flow_core_old.module.forward_rl_node
          edge_function_old = self.flow_core_old.module.forward_rl_edge
      else:
          node_function = self.flow_core.forward_rl_node
          edge_function = self.flow_core.forward_rl_edge
  
          node_function_old = self.flow_core_old.forward_rl_node
          edge_function_old = self.flow_core_old.forward_rl_edge
  
      z_node, _ = node_function(
          node_inputs_node_features,
          node_inputs_adj_features,
          node_inputs_node_features_cont,
      )  # (total_step, 9), (total_step, )
  
      z_edge, _ = edge_function(
          adj_inputs_node_features,
          adj_inputs_adj_features,
          adj_inputs_edge_features_cont,
          adj_inputs_index,
      )  # (total_step, 4), (total_step, )
  
      with torch.no_grad():
          z_node_old, _ = node_function_old(
              node_inputs_node_features,
              node_inputs_adj_features,
              node_inputs_node_features_cont,
          )  # (total_step, 9), (total_step, )
  
          z_edge_old, _ = edge_function_old(
              adj_inputs_node_features,
              adj_inputs_adj_features,
              adj_inputs_edge_features_cont,
              adj_inputs_index,
          )  # (total_step, 4), (total_step, )
  
      node_total_length = z_node.size(0) * float(self.node_dim)
      edge_total_length = z_edge.size(0) * float(self.bond_dim)
  
      node_base_log_probs_sm = torch.nn.functional.log_softmax(
          self.node_base_log_probs, dim=-1
      )
      node_base_log_probs_sm_select = torch.index_select(
          node_base_log_probs_sm, dim=0, index=node_inputs_node_cnts
      )
      ll_node = torch.sum(
          z_node * node_base_log_probs_sm_select, dim=(-1, -2)
      )
      edge_base_log_probs_sm = torch.nn.functional.log_softmax(
          self.edge_base_log_probs, dim=-1
      )
      edge_base_log_probs_sm_select = torch.index_select(
          edge_base_log_probs_sm, dim=0, index=adj_inputs_edge_cnts
      )
      ll_edge = torch.sum(
          z_edge * edge_base_log_probs_sm_select, dim=(-1, -2)
      )
  
      node_base_log_probs_sm_old = torch.nn.functional.log_softmax(
          self.node_base_log_probs_old, dim=-1
      )
      node_base_log_probs_sm_old_select = torch.index_select(
          node_base_log_probs_sm_old, dim=0, index=node_inputs_node_cnts
      )
      ll_node_old = torch.sum(
          z_node_old * node_base_log_probs_sm_old_select, dim=(-1, -2)
      )
      edge_base_log_probs_sm_old = torch.nn.functional.log_softmax(
          self.edge_base_log_probs_old, dim=-1
      )
      edge_base_log_probs_sm_old_select = torch.index_select(
          edge_base_log_probs_sm_old, dim=0, index=adj_inputs_edge_cnts
      )
      ll_edge_old = torch.sum(
          z_edge_old * edge_base_log_probs_sm_old_select, dim=(-1, -2)
      )
  
      ratio_node = torch.exp(
          (ll_node - ll_node_old.detach()).clamp(max=10.0, min=-10.0)
      )
      ratio_edge = torch.exp(
          (ll_edge - ll_edge_old.detach()).clamp(max=10.0, min=-10.0)
      )
  
      if torch.isinf(ratio_node).any():
          raise RuntimeError("ratio node has inf entries")
  
      if torch.isinf(ratio_edge).any():
          raise RuntimeError("ratio edge has inf entries")
      if (
          self.conf_rl["reward_type"] == "imp"
          and self.conf_rl["no_baseline"]
      ):
          advantage_node = node_inputs_rewards
          advantage_edge = adj_inputs_rewards
      else:
          advantage_node = node_inputs_rewards - node_inputs_baseline
          advantage_edge = adj_inputs_rewards - adj_inputs_baseline
  
      surr1_node = ratio_node * advantage_node
      surr2_node = (
          torch.clamp(ratio_node, 1 - 0.2, 1 + 0.2) * advantage_node
      )
  
      surr1_edge = ratio_edge * advantage_edge
      surr2_edge = (
          torch.clamp(ratio_edge, 1 - 0.2, 1 + 0.2) * advantage_edge
      )
  
      if torch.isnan(surr1_node).any():
          raise RuntimeError("surr1 node has NaN entries")
      if torch.isnan(surr2_node).any():
          raise RuntimeError("surr2 node has NaN entries")
      if torch.isnan(surr1_edge).any():
          raise RuntimeError("surr1 edge has NaN entries")
      if torch.isnan(surr2_edge).any():
          raise RuntimeError("surr2 edge has NaN entries")
  
      return (
          -(
              (
                  torch.min(surr1_node, surr2_node).sum()
                  + torch.min(surr1_edge, surr2_edge).sum()
              )
              / (node_total_length + edge_total_length)
              - 1.0
          ),
          per_mol_reward,
          per_mol_property_score,
          reward_baseline_per_step,
      )
  
file_path: DIG/dig/ggraph/method/GraphDF/model/graphflow_con_rl.py
imports: import torch
  import torch.nn
  import numpy
  from rdkit import Chem
  from dig.ggraph.utils import check_chemical_validity, check_valency, calculate_min_plogp, reward_target_molecule_similarity
  from dig.ggraph.utils import convert_radical_electrons_to_hydrogens, steric_strain_filter, zinc_molecule_filter
  from .disgraphaf import DisGraphAF
label_desc: Use image node_id 3 for calling the GraphFlowModel_con_rl obj's underlying member method code with example usage: obj.reinforce_forward_constrained_optim(mol_xs, mol_adjs, mol_sizes, raw_smiles, bfs_perm_origin, in_baseline, cur_iter) and returns: unknown,  per_mol_reward,  per_mol_property_score,  reward_baseline_per_step