original_docstring: 
          Input: s - tensor representing hidden state of task
                 a - integer representing which action is being performed (default = None) [OPTIONAL]
                 train - boolean representing if search is being used in training mode (default = False) [OPTIONAL]
          Description: Search the task action tree using upper confidence value
          Output: predicted value
          
docstring: tensor representing hidden state of task
  a - integer representing which action is being performed (default = None) [OPTIONAL]
  train - boolean representing if search is being used in training mode (default = False) [OPTIONAL]
  Description: Search the task action tree using upper confidence value
  Output: predicted value
docstring_tokens: tensor representing hidden state of task a - integer representing which action is being performed ( default = None ) [ OPTIONAL ] train - boolean representing if search is being used in training mode ( default = False ) [ OPTIONAL ] Description : Search the task action tree using upper confidence value Output : predicted value
code: def search(self, s, train = False, e_db = None, a = None):
          """
          Input: s - tensor representing hidden state of task
                 a - integer representing which action is being performed (default = None) [OPTIONAL]
                 train - boolean representing if search is being used in training mode (default = False) [OPTIONAL]
          Description: Search the task action tree using upper confidence value
          Output: predicted value
          """
          if a is not None:
              a_hash = deepcopy(a.reshape(1).item())
          else:
              a_hash = None
          s_hash = self.state_hash(s) #Create hash of state [sk-1] for game tree
          if (s_hash, a_hash) not in self.tree:
              self.tree[(s_hash, a_hash)] = self.Node() #Initialize new game tree node
          if a is not None and self.tree[(s_hash, a_hash)].S is None:
              with torch.no_grad():
                  d_k = self.g(s, a + 1) #backbone function
                  if e_db is not None:
                      d_k = torch.squeeze(d_k)
                      chunks = d_k.reshape(self.Cca.l, self.Cca.m, self.Cca.d)[:self.Cca.l - 1]
                      neighbours = ToolBox.get_kNN(chunks, e_db)
                      d_k = self.Cca(d_k, neighbours) #chunked cross-attention
                      d_k = d_k.reshape(1, d_k.size(0), d_k.size(1))
                  r_k = self.r(d_k) #reward function
                  s_k = self.s(d_k) #next state function
              s = s_k.reshape(s.size())
              self.tree[(s_hash, a_hash)].S = s
              self.tree[(s_hash, a_hash)].R = r_k.reshape(1).item()
          elif a is not None and self.tree[(s_hash, a_hash)].S is not None:
              s = self.tree[(s_hash, a_hash)].S
          sk_hash = self.state_hash(s) #Create hash of state [sk] for game tree
          if self.tree[(s_hash, a_hash)].N == 0:
              #EXPANSION ---
              self.expand_tree(s, s_hash, a_hash, sk_hash, noise = False)
          elif self.l < self.max_depth:
              a_k = torch.tensor(self.pUCT(sk_hash)) #Find best action to perform @ [sk]
              a_k = a_k.reshape((1,1))
              self.l += 1
              #BACKUP ---
              v_1  = self.search(s, a = a_k, e_db = e_db) #Go level deeper
              G = self.tree[(s_hash, a_hash)].R + (self.g_d * v_1)
              self.tree[(s_hash, a_hash)].Q = ((self.tree[(s_hash, a_hash)].N * self.tree[(s_hash, a_hash)].Q) + G) / (self.tree[(s_hash, a_hash)].N + 1) #Updated value
          if self.tree[(s_hash, a_hash)].Q < self.Q_min:
              self.Q_min = self.tree[(s_hash, a_hash)].Q
          if self.tree[(s_hash, a_hash)].Q > self.Q_max:
              self.Q_max = self.tree[(s_hash, a_hash)].Q
  
          self.tree[(s_hash, a_hash)].N += 1
          return self.tree[(s_hash, a_hash)].Q if self.single_player == True else -self.tree[(s_hash, a_hash)].Q
code_tokens: def search ( self , s , train = False , e_db = None , a = None ) : if a is not None : a_hash = deepcopy ( a . reshape ( 1 ) . item ( ) ) else : a_hash = None s_hash = self . state_hash ( s ) if ( s_hash , a_hash ) not in self . tree : self . tree [ ( s_hash , a_hash ) ] = self . Node ( ) if a is not None and self . tree [ ( s_hash , a_hash ) ] . S is None : with torch . no_grad ( ) : d_k = self . g ( s , a + 1 ) if e_db is not None : d_k = torch . squeeze ( d_k ) chunks = d_k . reshape ( self . Cca . l , self . Cca . m , self . Cca . d ) [ : self . Cca . l - 1 ] neighbours = ToolBox . get_kNN ( chunks , e_db ) d_k = self . Cca ( d_k , neighbours ) d_k = d_k . reshape ( 1 , d_k . size ( 0 ) , d_k . size ( 1 ) ) r_k = self . r ( d_k ) s_k = self . s ( d_k ) s = s_k . reshape ( s . size ( ) ) self . tree [ ( s_hash , a_hash ) ] . S = s self . tree [ ( s_hash , a_hash ) ] . R = r_k . reshape ( 1 ) . item ( ) elif a is not None and self . tree [ ( s_hash , a_hash ) ] . S is not None : s = self . tree [ ( s_hash , a_hash ) ] . S sk_hash = self . state_hash ( s ) if self . tree [ ( s_hash , a_hash ) ] . N == 0 : self . expand_tree ( s , s_hash , a_hash , sk_hash , noise = False ) elif self . l < self . max_depth : a_k = torch . tensor ( self . pUCT ( sk_hash ) ) a_k = a_k . reshape ( ( 1 , 1 ) ) self . l += 1 v_1 = self . search ( s , a = a_k , e_db = e_db ) G = self . tree [ ( s_hash , a_hash ) ] . R + ( self . g_d * v_1 ) self . tree [ ( s_hash , a_hash ) ] . Q = ( ( self . tree [ ( s_hash , a_hash ) ] . N * self . tree [ ( s_hash , a_hash ) ] . Q ) + G ) / ( self . tree [ ( s_hash , a_hash ) ] . N + 1 ) if self . tree [ ( s_hash , a_hash ) ] . Q < self . Q_min : self . Q_min = self . tree [ ( s_hash , a_hash ) ] . Q if self . tree [ ( s_hash , a_hash ) ] . Q > self . Q_max : self . Q_max = self . tree [ ( s_hash , a_hash ) ] . Q self . tree [ ( s_hash , a_hash ) ] . N += 1 return self . tree [ ( s_hash , a_hash ) ] . Q if self . single_player == True else - self . tree [ ( s_hash , a_hash ) ] . Q
short_docstring: Input: s - tensor representing hidden state of task
  a - integer representing which action is being performed (default = None) [OPTIONAL]
  train - boolean representing if search is being used in training mode (default = False) [OPTIONAL]
  Description: Search the task action tree using upper confidence value
  Output: predicted value
short_docstring_tokens: Input : s - tensor representing hidden state of task a - integer representing which action is being performed ( default = None ) [ OPTIONAL ] train - boolean representing if search is being used in training mode ( default = False ) [ OPTIONAL ] Description : Search the task action tree using upper confidence value Output : predicted value
comment: """
          Input: s - tensor representing hidden state of task
                 a - integer representing which action is being performed (default = None) [OPTIONAL]
                 train - boolean representing if search is being used in training mode (default = False) [OPTIONAL]
          Description: Search the task action tree using upper confidence value
          Output: predicted value
          """$SEP$#Create hash of state [sk-1] for game tree$SEP$#Initialize new game tree node$SEP$#backbone function$SEP$#chunked cross-attention$SEP$#reward function$SEP$#next state function$SEP$#Create hash of state [sk] for game tree$SEP$#EXPANSION ---$SEP$#Find best action to perform @ [sk]$SEP$#BACKUP ---$SEP$#Go level deeper$SEP$#Updated value