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 load_checkpoint(
      self,
  ) -> int:
      critic_episode = -1
      actor_episode = -1
  
      # check if there are some checkpoint for the critic
      print("Looking for checkpoints...")
      path = ModelLoader.check_model_path(
          config=self.config.critic,
          is_checkpoint=True,
          current_epoch=None,
      )
  
      # if there are checkpoint
      if path is not None:
          # load the critic checkpoint
          print("Loading ...")
          try:
              checkpoint = torch.load(path)
          except Exception:
              print(
                  "Checkpoint of critic corrupted!"
                  "Try to remove the last checkpoint."
                  "Now Starting from episode 0"
              )
              return 0
  
          # load checkpoint into model
          critic_episode = checkpoint["episode"]
          self.actorcritic.critic.load_state_dict(
              checkpoint["critic_state_dict"]
          )
          self.critic_optimizer.load_state_dict(
              checkpoint["critic_optim_state_dict"]
          )
  
      # check if there are checkpoints for the actor
      print("Looking for checkpoints...")
      path = ModelLoader.check_model_path(
          config=self.config,
          is_checkpoint=True,
          current_epoch=None,
      )
  
      # if there are some checkpoints
      if path is not None:
          # load the actor checkpoint
          print("Loading ...")
          try:
              checkpoint = torch.load(path)
          except Exception:
              print(
                  "Checkpoint of actor corrupted!"
                  "Try to remove the last checkpoint."
                  "Now Starting from episode 0"
              )
              return 0
  
          # load checkpoint into the model
          actor_episode = checkpoint["episode"]
          self.actorcritic.actor.load_state_dict(
              checkpoint["actor_state_dict"]
          )
          self.actor_optimizer.load_state_dict(
              checkpoint["actor_optim_state_dict"]
          )
          self.training_stats = checkpoint["training_stats"]
  
      # check if there are some discrepancies between the checkpoints
      if critic_episode == actor_episode:
          # all ok start from next episode
          return critic_episode + 1
      else:
          print(
              f"There are some discrepancies between the checkpoints"
              f"of actor and critic \nactor episode: {actor_episode}"
              f"\n critic episode: {critic_episode}\n"
          )
          return min(critic_episode, actor_episode) + 1
  
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 4 for calling the RLTrainer obj's underlying member method code with example usage: obj.load_checkpoint() and returns: unknown, unknown, int, int