class_docstr: Trainer for the AlphaTensor model. The trainer does not require an
  explicit loss since the loss is computed by the model itself. The trainer
  is responsible for both the training step and the acting one, storing
  acting performance in a buffer.
class_docstr_tok: ["Trainer","for","the","AlphaTensor","model",".","The","trainer","does","not","require","an","explicit","loss","since","the","loss","is","computed","by","the","model","itself",".","The","trainer","is","responsible","for","both","the","training","step","and","the","acting","one",",","storing","acting","performance","in","a","buffer","."]
docstr: Runs actors in parallel to generate multiple games starting from
  the same input tensor.
  
  Args:
      input_tensor (torch.Tensor): The input tensor used to generate the
      games.
      n_games (int): Number of games to generate / actors to be run in
      parallel.
      mc_n_sim (int): Number of simulations used in the Monte Carlo tree
      search.
      N_bar (int): N_bar parameter used to compute tau when improving
      the policy.
docstr_tok: ["Runs","actors","in","parallel","to","generate","multiple","games","starting","from","the","same","input","tensor",".","Args",":","input_tensor","(","torch.Tensor",")",":","The","input","tensor","used","to","generate","the","games",".","n_games","(","int",")",":","Number","of","games","to","generate","\/","actors","to","be","run","in","parallel",".","mc_n_sim","(","int",")",":","Number","of","simulations","used","in","the","Monte","Carlo","tree","search",".","N_bar","(","int",")",":","N_bar","parameter","used","to","compute","tau","when","improving","the","policy","."]
code: def act_step(
      self,
      input_tensor: torch.Tensor,
      n_games: int,
      mc_n_sim: int,
      N_bar: int,
  ):
      """Runs actors in parallel to generate multiple games starting from
      the same input tensor.
  
      Args:
          input_tensor (torch.Tensor): The input tensor used to generate the
          games.
          n_games (int): Number of games to generate / actors to be run in
          parallel.
          mc_n_sim (int): Number of simulations used in the Monte Carlo tree
          search.
          N_bar (int): N_bar parameter used to compute tau when improving
          the policy.
      """
      self.model.eval()
      best_reward = -1e10
      best_game = None
  
      if self.extra_devices:
          from joblib import Parallel, delayed
  
          # this means that there is an empty GPU available
          # thus we can use it to parallelize the acting step
          # use joblib to parallelize the acting step
          # we should use _single_act as a function to be parallelized
          extra_devices = (
              self.extra_devices * (n_games // len(self.extra_devices))
              + self.extra_devices[: n_games % len(self.extra_devices)]
          )
          self.model.to("cpu")
          input_tensor = input_tensor.to("cpu")
  
          print(f"Starting acting phase with {n_games} games")
          results = Parallel(n_jobs=len(self.extra_devices))(
              delayed(_single_act)(
                  actor_id,
                  self.model,
                  input_tensor,
                  extra_devices[actor_id],
                  mc_n_sim,
                  N_bar,
                  self.change_of_basis,
                  self.max_rank,
              )
              for actor_id in range(n_games)
          )
          self.model.to(self.device)
  
          for actor_id, states, policies, rewards in results:
              if rewards[-1] > best_reward:
                  print(f"New best actor! Actor: {actor_id}")
                  best_reward = rewards[-1]
                  best_game = (states, policies, rewards)
              self.dataset.add_game(states, policies, rewards)
              if self.data_augmentation:
                  states, policies = swap_data(states, policies)
                  self.dataset.add_game(states, policies, rewards)
          if best_game is not None:
              self.dataset.add_best_game(*best_game)
      else:
          for actor_id in range(n_games):
              input_tensor_cob = self.change_of_basis(input_tensor).to(
                  self.device
              )
              print(f"Running actor {actor_id} / {n_games}")
              states, policies, rewards = actor_prediction(
                  self.model,
                  input_tensor_cob,
                  self.max_rank,
                  mc_n_sim,
                  N_bar,
              )
              print(
                  f"Actor {actor_id} finished. Final reward: {rewards[-1]}"
              )
              if rewards[-1] > best_reward:
                  print("New best actor!")
                  best_reward = rewards[-1]
                  best_game = (states, policies, rewards)
              self.dataset.add_game(states, policies, rewards)
              if self.data_augmentation:
                  states, policies = swap_data(states, policies)
                  self.dataset.add_game(states, policies, rewards)
          if best_game is not None:
              self.dataset.add_best_game(*best_game)
  
file_path: nebuly/optimization/open_alpha_tensor/open_alpha_tensor/core/training.py
imports: from pathlib import Path
  from typing import Tuple, List
  import torch.optim
  import tqdm
  from torch.utils.data import DataLoader
  from open_alpha_tensor.config import BASE_CHECKPOINT_DATA_DIR, BASE_CHECKPOINT_DIR
  from open_alpha_tensor.core.actors.stage import actor_prediction
  from open_alpha_tensor.core.data.basis_change import ChangeOfBasis
  from open_alpha_tensor.core.data.dataset import TensorGameDataset
  from open_alpha_tensor.core.data.generation import f_prob_distribution
  from open_alpha_tensor.core.data.utils import map_action_to_triplet
  from open_alpha_tensor.core.modules.alpha_tensor import AlphaTensorModel
label_desc: Use image node_id 3 for calling the Trainer obj's underlying member method code with example usage: obj.act_step(input_tensor, n_games, mc_n_sim, N_bar) without return types