code: def select_future_state(
      possible_states: List[torch.Tensor],
      q_values: torch.Tensor,
      N_s_a: torch.Tensor,
      repetitions: Dict[int, list],
      c_1: float = 1.25,
      c_2: float = 19652,
      return_idx: bool = False,
  ) -> torch.Tensor:
      """Select the future state maximizing the upper confidence bound."""
      # q_values (1, K, 1)
      pi = torch.tensor(
          [
              len(repetitions[i])
              for i in range(len(possible_states))
              if i in repetitions
          ]
      ).to(q_values.device)
      if pi.shape[0] != N_s_a.shape[1]:
          print(pi)
          print(pi.shape, q_values.shape, N_s_a.shape)
          pi = pi[: N_s_a.shape[1]]
      ucb = q_values.reshape(-1) + pi * torch.sqrt(
          torch.sum(N_s_a) / (1 + N_s_a)
      ) * (c_1 + torch.log((torch.sum(N_s_a) + c_2 + 1) / c_2))
      if return_idx:
          return ucb.argmax()
      return possible_states[ucb.argmax()]
  
file_path: nebuly/optimization/open_alpha_tensor/open_alpha_tensor/core/actors/stage.py
imports: from typing import Dict, List
  import torch
  from open_alpha_tensor.core.data.utils import get_scalars, map_action_to_triplet
  from open_alpha_tensor.core.modules.alpha_tensor import AlphaTensorModel
label_desc: Use image node_id 10 for calling a global function with example usage: select_future_state(possible_states, q_values, N_s_a, repetitions, c_1, c_2, return_idx) and returns: possible_states, ucb