class_docstr: Train the actor-critic model using RL
  
  Attributes:
      config (Config): Configuration of the trainer
      debug (bool): Debug mode
      actorcritic (ActorCritic): Actor-critic model
      actor_optim (torch.optim): Optimizer for the actor
      critic_optim (torch.optim): Optimizer for the critic
      actor_scheduler (torch.optim.lr_scheduler): Scheduler for the actor
      critic_scheduler (torch.optim.lr_scheduler): Scheduler for the critic
      reward (RewardModel): Reward model
      training_stats (TrainingStats): Class to store training stats
      conversation_log (ConversationLog): Class to store the conversation
      examples_sampler (ExamplesSampler): Class to sample examples
      eps (float): small epsilon to avoid division by zero
  
  Methods:
      train: the training loop that calls the learn function after generating
          the experiences.
      learn: Learn from a batch of experiences and update the actor and the
          critic model.
      load_checkpoint: Load the checkpoint of the actor-critic model
      save_checkpoint: Save the checkpoint of the actor-critic model
class_docstr_tok: ["Train","the","actor-critic","model","using","RL","Attributes",":","config","(","Config",")",":","Configuration","of","the","trainer","debug","(","bool",")",":","Debug","mode","actorcritic","(","ActorCritic",")",":","Actor-critic","model","actor_optim","(","torch.optim",")",":","Optimizer","for","the","actor","critic_optim","(","torch.optim",")",":","Optimizer","for","the","critic","actor_scheduler","(","torch.optim.lr_scheduler",")",":","Scheduler","for","the","actor","critic_scheduler","(","torch.optim.lr_scheduler",")",":","Scheduler","for","the","critic","reward","(","RewardModel",")",":","Reward","model","training_stats","(","TrainingStats",")",":","Class","to","store","training","stats","conversation_log","(","ConversationLog",")",":","Class","to","store","the","conversation","examples_sampler","(","ExamplesSampler",")",":","Class","to","sample","examples","eps","(","float",")",":","small","epsilon","to","avoid","division","by","zero","Methods",":","train",":","the","training","loop","that","calls","the","learn","function","after","generating","the","experiences",".","learn",":","Learn","from","a","batch","of","experiences","and","update","the","actor","and","the","critic","model",".","load_checkpoint",":","Load","the","checkpoint","of","the","actor-critic","model","save_checkpoint",":","Save","the","checkpoint","of","the","actor-critic","model"]
code: def __init__(
      self,
      config: Config,
  ) -> None:
      # save config
      self.config = config
  
      # set debug mode
      self.debug = config.trainer.debug
  
      # initialize agent-critic
      self.actorcritic = ActorCritic(config)
  
      # initialize actor optimizer
      self.actor_optimizer = torch.optim.Adam(
          self.actorcritic.actor.parameters(),
          lr=config.trainer.actor_lr,
      )
  
      # initialize critic optimizer
      self.critic_optimizer = torch.optim.Adam(
          self.actorcritic.critic.parameters(),
          lr=config.trainer.critic_lr,
      )
  
      # scheduler (defined in the learn() method (i need dataset size))
      self.actor_scheduler = None
      self.critic_scheduler = None
  
      # initialize reward model
      self.reward = RewardModel(config.reward)
  
      # initialize class to store training stats
      path = ModelLoader.get_training_stats_path(config)
      self.training_stats = TrainingStats(path)
      model_folder, _, _ = ModelLoader.get_model_path(
          config,
          is_checkpoint=True,
      )
      path = os.path.join(model_folder, "conversations_log.json")
      self.conversation_log = ConversationLog(path)
  
      # initialize examples sampler
      self.example_sampler = ExamplesSampler(
          config.trainer.examples_path
      )
  
      # check if actor and critic use the same tokenizer
      self.actorcritic.use_same_tokenizer = check_model_family(
          config.actor, config.critic
      )
  
      # check if actor and reward use the same tokenizer
      self.use_same_tokenizer = check_model_family(
          config.actor, config.reward
      )
  
      # eps
      self.eps = 1e-8
  
      # deepspeed initialization
      self.actor_model_engine = None
      self.critic_model_engine = None
      self.is_deepspeed_init = None
  
      if (
          self.config.actor.deepspeed_enable
          or self.config.critic.deepspeed_enable
          or self.config.critic.deepspeed_enable
      ):
          deepspeed.init_distributed("nccl")
          self.is_deepspeed_init = True
          os.environ["TOKENIZERS_PARALLELISM"] = "False"
  
      else:
          self.is_deepspeed_init = False
  
      if self.config.actor.deepspeed_enable:
          (
              self.actor_model_engine,
              self.actorcritic.actor,
              self.actor_optimizer,
          ) = self.initialize_deepspeed_model(
              config=self.config.actor, model=self.actorcritic.actor
          )
  
      if self.config.critic.deepspeed_enable:
          (
              self.critic_model_engine,
              self.actorcritic.critic,
              self.critic_optimizer,
          ) = self.initialize_deepspeed_model(
              config=self.config.critic, model=self.actorcritic.critic
          )
  
      if self.config.reward.deepspeed_enable:
          (
              _,
              self.reward,
              _,
          ) = self.initialize_deepspeed_model(
              config=self.config.reward, model=self.reward
          )
  
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 1 to create a new RLTrainer object  with example: obj = RLTrainer(config)