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 train(
      self,
  ) -> None:
      print("Start RL Training")
  
      # initialize settings
      num_episodes = self.config.trainer.num_episodes
      max_timesteps = self.config.trainer.max_timesteps
      num_examples = self.config.trainer.num_examples
      update_timesteps = self.config.trainer.update_timesteps
      batch_size = self.config.trainer.batch_size
      checkpoint_steps = self.config.trainer.checkpoint_steps
      device = (
          torch.device(f"cuda:{dist.get_rank()}")
          if self.is_deepspeed_init
          else self.config.trainer.device
      )
  
      # number of elements that the memories should contain when learning
      number_of_memories_per_learn_iteration = (
          num_examples * update_timesteps
      )
  
      # the number of memories must be a multiple of the batch size
      assert (
          number_of_memories_per_learn_iteration % batch_size == 0
      ), "The number of memories must be a multiple of the batch size"
  
      # the total number of timesteps done in the train() are
      total_number_of_timesteps = num_episodes * max_timesteps
  
      # the total timesteps done should be a multiple of the update timesteps
      assert total_number_of_timesteps % update_timesteps == 0, (
          "The number of timesteps (num_episodes*max_timesteps)"
          "must be a multiple of the update_timesteps"
      )
  
      # initialize memories
      memories = deque([])
  
      # load checkpoint
      start_episode = self.load_checkpoint()
  
      # if it is a new training from the start clear the conversation log
      if start_episode == 0:
          self.conversation_log.clear()
  
      # initialize counters
      cnt_timesteps = 0
      cnt_learn_iter = 0
  
      # loop over episodes and timesteps
      self.actorcritic.eval()
      for episode in range(start_episode, num_episodes):
          for timestep in range(max_timesteps):
              # print the iteration info
              print(
                  f"Episode: {episode + 1}/{num_episodes}, "
                  f"Timestep: {timestep + 1}/{max_timesteps}",
                  f"Learning Cnt: {cnt_timesteps + 1}/{update_timesteps}",
              )
  
              # counter used to count timesteps into memory
              cnt_timesteps += 1
  
              # sample num_examples examples from  example dataset
              inputs = self.example_sampler.sample(num_examples)
  
              # tokenize examples for the actor
              tok_inputs_act = self.actorcritic.actor.tokenizer(
                  inputs,
                  padding=True,
                  return_tensors="pt",
                  truncation=True,
              )
  
              # states are [batch_size, seq_len_of_states]
              states_actor = tok_inputs_act["input_ids"].to(device)
              states_mask_actor = tok_inputs_act["attention_mask"].to(
                  device
              )
  
              # tokenize examples for the critic
              tok_inputs_crt = self.actorcritic.critic.tokenizer(
                  inputs,
                  padding=True,
                  return_tensors="pt",
                  truncation=True,
              )
  
              # states are [batch_size, seq_len_of_states]
              states_critic = tok_inputs_crt["input_ids"].to(device)
  
              # generate sequences of actions and values
              (
                  actions,
                  actions_logits,
                  values,
                  sequences_actor,
                  sequences_mask_actor,
                  sequences_critic,
                  sequences_mask_critic,
                  action_len_actor,
                  action_len_critic,
              ) = self.actorcritic.generate(
                  states_actor, states_mask_actor, states_critic
              )
  
              # compute action log probs
              action_prob = (
                  torch.softmax(actions_logits, dim=-1)
                  .max(dim=-1)
                  .values
              )
              actions_log_probs = torch.log(action_prob + self.eps)
  
              # get tokenized sequence for the reward models
              if self.use_same_tokenizer:
                  reward_sequence = sequences_actor
                  reward_mask = sequences_mask_actor
              elif check_model_family(
                  self.config.critic, self.config.reward
              ):
                  reward_sequence = sequences_critic
                  reward_mask = sequences_mask_critic
              else:
                  tokenized_responses = change_tokenization(
                      sequences_actor,
                      self.actorcritic.actor.tokenizer,
                      self.reward.tokenizer,
                  )
                  # get tokens and mask
                  reward_sequence = tokenized_responses["input_ids"].to(
                      device
                  )
                  reward_mask = tokenized_responses[
                      "attention_mask"
                  ].to(device)
  
              # compute rewards
              rewards = self.reward.forward(
                  reward_sequence,
                  reward_mask,
              )
  
              rewards = rewards[:, -action_len_critic:]
              reward = rewards[:, -1]
  
              # store memories of the episode / timestep
              for i in range(states_actor.shape[0]):
                  memories.append(
                      Memory(
                          states_actor[i, :].detach().cpu(),
                          actions[i, :].detach().cpu(),
                          values[i, :].detach().cpu(),
                          rewards[i, :].detach().cpu(),
                          actions_log_probs[i, :].detach().cpu(),
                          sequences_actor[i, :].detach().cpu(),
                          sequences_mask_actor[i, :].detach().cpu(),
                          sequences_critic[i, :].detach().cpu(),
                          sequences_mask_critic[i, :].detach().cpu(),
                          int(action_len_actor),
                          int(action_len_critic),
                      )
                  )
  
              # decode completions to be logged in the conversation log
              completions = [
                  self.actorcritic.actor.tokenizer.decode(action)
                  for action in actions
              ]
              # remove pad tokens from completions
              completions = [
                  c.replace(
                      self.actorcritic.actor.tokenizer.pad_token, ""
                  )
                  for c in completions
              ]
              # remove eos tokens from completions
              completions = [
                  c.replace(
                      self.actorcritic.actor.tokenizer.eos_token, ""
                  )
                  for c in completions
              ]
              # strange i need to force this?
              completions = [
                  c.replace("<pad>", "") for c in completions
              ]
  
              # log the memories in the conversation log
              for i in range(states_actor.shape[0]):
                  self.conversation_log.append(
                      inputs[i],
                      completions[i],
                      reward[i].detach().cpu().item(),
                      cnt_learn_iter,
                  )
  
              # learn from memories
              if (cnt_timesteps % update_timesteps == 0) and (
                  cnt_timesteps != 0
              ):
                  print("len memories", len(memories))
                  if not self.is_deepspeed_init or (
                      dist.get_rank() == 0
                  ):
                      self.conversation_log.save()
                  self.learn(memories)
                  mean_reward = sum(
                      [m.rewards[-1] for m in memories]
                  ) / len(memories)
                  print(f"Mean Reward: {mean_reward}")
                  memories.clear()
                  cnt_timesteps = 0
                  cnt_learn_iter += 1
                  if not self.is_deepspeed_init or (
                      dist.get_rank() == 0
                  ):
                      self.conversation_log.save()
  
          # save checkpoints
          if (episode % checkpoint_steps == 0) and (episode != 0):
              self.save_checkpoint(
                  current_episode=episode, max_episode=num_episodes
              )
              if not self.is_deepspeed_init or (dist.get_rank() == 0):
                  self.conversation_log.save()
  
      # save the models
      if self.is_deepspeed_init:
          self.actorcritic.save_deepspeed(
              self.actor_model_engine, self.config
          )
          self.actorcritic.save_deepspeed(
              self.critic_model_engine, self.config.critic
          )
      else:
          self.actorcritic.save()
      print("End RL Training")
  
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 RLTrainer obj's underlying member method code with example usage: obj.train() without return types