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: Given the whole sequences, use the actor forward to get the logits
      for each token in the sequence and the critic forward to get the
      values for each generation step.
  
  Args:
      sequences_actor (torch.Tensor): Sequences composed of
          [states, actions] for the actor
      sequence_mask_actor (torch.Tensor): Mask for the sequences
          of the actor
      sequences_critic (torch.Tensor): Sequences composed of
          [states, actions] for the critic
      sequences_mask_critic (torch.Tensor): Mask for the sequences
          of the critic
      action_len_actor (int): Length of the actions in the sequences
          for the actor
      action_len_critic (int): Length of the actions in the sequences
          for the critic
  
  Returns:
      action_logits (torch.Tensor): Logits for the actions in the
          sequences
      values (torch.Tensor): Values for the actions in the sequences
code: def forward(
      self,
      sequences_actor: torch.Tensor,
      sequences_mask_actor: torch.Tensor,
      sequences_critic: torch.Tensor,
      sequences_mask_critic: torch.Tensor,
      action_len_actor: int,
      action_len_critic: int,
  ) -> Tuple:
      """Given the whole sequences, use the actor forward to get the logits
          for each token in the sequence and the critic forward to get the
          values for each generation step.
  
      Args:
          sequences_actor (torch.Tensor): Sequences composed of
              [states, actions] for the actor
          sequence_mask_actor (torch.Tensor): Mask for the sequences
              of the actor
          sequences_critic (torch.Tensor): Sequences composed of
              [states, actions] for the critic
          sequences_mask_critic (torch.Tensor): Mask for the sequences
              of the critic
          action_len_actor (int): Length of the actions in the sequences
              for the actor
          action_len_critic (int): Length of the actions in the sequences
              for the critic
  
      Returns:
          action_logits (torch.Tensor): Logits for the actions in the
              sequences
          values (torch.Tensor): Values for the actions in the sequences
      """
  
      # use a single forward on the whole sequence
      # to get pi(y | x) and ignore predicted output
      actions_logits = self.actor.forward(
          sequences_actor, sequences_mask_actor
      )
  
      # use the critic forward to get the values for the actions
      values = self.critic.forward(
          sequences_critic, sequences_mask_critic
      )
  
      # return only logits and values for the actions taken
      real_actions_logits = actions_logits[:, -action_len_actor:, :]
      real_values = values[:, -action_len_critic:]
  
      if self.debug:
          print("ActorCritic.forward")
          print("action_len_actor", action_len_actor)
          print("action_len_critic", action_len_critic)
          print("sequences_actor.shape", sequences_actor.shape)
          print("sequences_actor", sequences_actor)
          print("sequences_critic.shape", sequences_critic.shape)
          print("sequences_critic", sequences_critic)
          print("real_action_logits.shape", actions_logits.shape)
          print("real_action_logits", actions_logits)
          print("real_values.shape", values.shape)
          print("real_values", values)
  
      return (
          real_actions_logits,
          real_values,
      )
  
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 5 for calling the ActorCritic obj's underlying member method code with example usage: obj.forward(sequences_actor, sequences_mask_actor, sequences_critic, sequences_mask_critic, action_len_actor, action_len_critic) and returns: real_actions_logits,  real_values