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 save_checkpoint(
      self,
      current_episode: int,
      max_episode: int,
  ) -> None:
      print(f"Saving checkpoint for episode {current_episode+1}..")
  
      # get the path to save the checkpoint for the critic
      model_folder, model_name, path = ModelLoader.get_model_path(
          config=self.config.critic,
          is_checkpoint=True,
          current_epoch=current_episode,
          max_epochs=max_episode,
          max_steps=0,
      )
  
      # if the checkpoint already exists remove it.
      # Deepspeed checkpoints are already directories and will be overwritten
      if os.path.exists(path) and not self.is_deepspeed_init:
          os.remove(path)
  
      # save the checkpoint
      actor_checkpoint_dict = {
          "episode": current_episode,
          "critic_state_dict": self.actorcritic.critic.state_dict(),
          "critic_optim_state_dict": self.critic_optimizer.state_dict(),
      }
  
      if self.config.actor.deepspeed_enable:
          # The model and optimizer state dicts are actually already saved
          # In the deepspeed model engine. But to make sure no depending
          # methods fail, the states are included in actor_checkpoint_dict.
          # ATTENTION: If you use deepspeed zero optimization, the client_state
          # will not be saved
          self.actor_model_engine.save_checkpoint(
              save_dir=path, client_state=actor_checkpoint_dict
          )
      else:
          torch.save(actor_checkpoint_dict, path)
  
      # get the path to save the checkpoint for the actor
      model_folder, model_name, path = ModelLoader.get_model_path(
          config=self.config,
          is_checkpoint=True,
          current_epoch=current_episode,
          max_epochs=max_episode,
          max_steps=0,
      )
  
      # if the checkpoint already exists remove it.
      # Deepspeed checkpoints are already directories and will be overwritten
      if os.path.exists(path) and not self.is_deepspeed_init:
          os.remove(path)
  
      # save the checkpoint
      critic_checkpoint_dict = {
          "episode": current_episode,
          "actor_state_dict": self.actorcritic.actor.state_dict(),
          "actor_optim_state_dict": self.actor_optimizer.state_dict(),
          "training_stats": self.training_stats,
      }
  
      if self.config.critic.deepspeed_enable:
          # The model and optimizer state dicts are actually already saved
          # In the deepspeed model engine. But to make sure no depending
          # methods fail, the states are included in critic_checkpoint_dict.
          # ATTENTION: If you use deepspeed zero optimization, the client_state
          # will not be saved
          self.critic_model_engine.save_checkpoint(
              save_dir=path, client_state=critic_checkpoint_dict
          )
      else:
          torch.save(critic_checkpoint_dict, path)
  
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 3 for calling the RLTrainer obj's underlying member method code with example usage: obj.save_checkpoint(current_episode, max_episode) without return types