class_docstr: Actor Critic class stores both the actor and the critic models
  and it generates values and action for given sequences during the training
  of the actor.
  
  Attributes:
      actor (ActorModel): Actor model
      critic (CriticModel): Critic model
      debug (bool): enable prints for Debugging
      use_same_tokenizer (bool): if True the actor and critic use the same
          tokenizer
  
  Methods:
      forward: given a sequence returns action logits and values (used
          to evaluate the actor during training)
      generate: given a sequence returns action, action logits, values
          sequences and sequences masks (used to generate new sequences
          during acting phase)
class_docstr_tok: ["Actor","Critic","class","stores","both","the","actor","and","the","critic","models","and","it","generates","values","and","action","for","given","sequences","during","the","training","of","the","actor",".","Attributes",":","actor","(","ActorModel",")",":","Actor","model","critic","(","CriticModel",")",":","Critic","model","debug","(","bool",")",":","enable","prints","for","Debugging","use_same_tokenizer","(","bool",")",":","if","True","the","actor","and","critic","use","the","same","tokenizer","Methods",":","forward",":","given","a","sequence","returns","action","logits","and","values","(","used","to","evaluate","the","actor","during","training",")","generate",":","given","a","sequence","returns","action",",","action","logits",",","values","sequences","and","sequences","masks","(","used","to","generate","new","sequences","during","acting","phase",")"]
docstr: Generate actions, actions_logits, values and sequences from states
  
  Args:
      states_actor (torch.Tensor): States for the actor
      states_mask_actor (torch.Tensor): Mask for the states for the
          actor
      states_critic (torch.Tensor): States for the critic
  
  Returns:
      actions (torch.Tensor): Actions generated from the states
      actions_logits (torch.Tensor): Logits for the actions generated
          from the states (i.e. pi(y | x))
      values (torch.Tensor): Values generated by the critic model
          for the actions generated by the actor (i.e. V(x))
      sequences (torch.Tensor): Sequences generated from the states
          as [states, actions]
code: def generate(
      self,
      states_actor: torch.Tensor,
      states_mask_actor: torch.Tensor,
      states_critic: torch.Tensor,
  ) -> Tuple:
      """Generate actions, actions_logits, values and sequences from states
  
      Args:
          states_actor (torch.Tensor): States for the actor
          states_mask_actor (torch.Tensor): Mask for the states for the
              actor
          states_critic (torch.Tensor): States for the critic
  
      Returns:
          actions (torch.Tensor): Actions generated from the states
          actions_logits (torch.Tensor): Logits for the actions generated
              from the states (i.e. pi(y | x))
          values (torch.Tensor): Values generated by the critic model
              for the actions generated by the actor (i.e. V(x))
          sequences (torch.Tensor): Sequences generated from the states
              as [states, actions]
      """
  
      # generate action sequence from the actor
      actions, sequences_actor = self.actor.generate(
          states_actor, states_mask_actor
      )
  
      # create mask for the actor sequences
      sequences_mask_actor = (
          (sequences_actor != self.actor.tokenizer.pad_token_id)
          .to(sequences_actor.device)
          .long()
          .detach()
      )
  
      # get the length of the actions
      action_len_actor = actions.shape[1]
  
      # check if different encoding is needed for the critic
      if self.use_same_tokenizer:
          sequences_critic = sequences_actor
          sequences_mask_critic = sequences_mask_actor
          action_len_critic = action_len_actor
      else:
          encoded_critic = change_tokenization(
              sequences_actor,
              self.actor.tokenizer,
              self.critic.tokenizer,
          )
          # split the encoded_critic in tokens and maks
          sequences_critic = encoded_critic["input_ids"].to(
              sequences_actor.device,
          )
          sequences_mask_critic = (
              encoded_critic["attention_mask"]
              .to(sequences_actor.device)
              .long()
              .detach()
          )
  
          # compute len of actions for the critic tokenizer
          action_len_critic = states_critic.shape[1]
  
      # generate actions_logits and values
      actions_logits, values = self.forward(
          sequences_actor,
          sequences_mask_actor,
          sequences_critic,
          sequences_mask_critic,
          action_len_actor,
          action_len_critic,
      )
      if self.debug:
          print("ActorCritic.generate")
          print("actions shape", actions.shape)
          print("actions", actions)
          print("sequence shape", sequences_actor.shape)
          print("sequence", sequences_actor)
          print("actions_logits shape", actions_logits.shape)
          print("actions_logits", actions_logits)
          print("values shape", values.shape)
          print("values", values)
  
      return (
          actions,
          actions_logits,
          values,
          sequences_actor,
          sequences_mask_actor,
          sequences_critic,
          sequences_mask_critic,
          action_len_actor,
          action_len_critic,
      )
  
imports: import json
  import os
  import random
  from collections import deque, namedtuple
  import deepspeed
  import torch
  import torch.distributed
  from accelerate import Accelerator
  from beartype import beartype
  from beartype.typing import Deque, List, Tuple, Union
  from deepspeed.runtime.engine import DeepSpeedEngine
  from torch.utils.data import DataLoader, Dataset
  from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
  from chatllama.rlhf.actor import ActorModel
  from chatllama.rlhf.config import Config, ConfigActor, ConfigCritic, ConfigReward
  from chatllama.rlhf.model_list import hf_models
  from chatllama.rlhf.model_loader import ModelLoader
  from chatllama.rlhf.reward import RewardModel, CriticModel
  from chatllama.rlhf.utils import TrainingStats, ConversationLog
label_desc: Use image node_id 6 for calling the ActorCritic obj's underlying member method code with example usage: obj.generate(states_actor, states_mask_actor, states_critic) and returns: actions,  actions_logits,  values,  sequences_actor,  sequences_mask_actor,  sequences_critic,  sequences_mask_critic,  action_len_actor,  action_len_critic