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"]
docstr: Train the agent-critic model using RL:
  - for each batch of episodes, compute action logits and values
  - then compare action logits probs with memories one and values with
      rewards to compute the PPO loss and update the actor-critic model
code: def learn(self, memories: Deque[Memory]) -> None:
      """Train the agent-critic model using RL:
      - for each batch of episodes, compute action logits and values
      - then compare action logits probs with memories one and values with
          rewards to compute the PPO loss and update the actor-critic model
      """
      print("Start to Learn...")
  
      # get parameters
      epochs = self.config.trainer.epochs
      actor_eps_clip = self.config.trainer.actor_eps_clip
      critic_eps_clip = self.config.trainer.critic_eps_clip
      beta_s = self.config.trainer.beta_s
      batch_size = self.config.trainer.batch_size
      device = (
          torch.device(f"cuda:{dist.get_rank()}")
          if self.is_deepspeed_init
          else self.config.trainer.device
      )
  
      # create dataset from memories
      dataset = ExperienceDataset(memories, device)
      if self.is_deepspeed_init:
          engine = self.actor_model_engine or self.critic_model_engine
          dataloader = engine.deepspeed_io(dataset)
      else:
          dataloader = DataLoader(dataset, batch_size=batch_size)
  
      # initialize scheduler for actor
      actor_lr = self.config.trainer.actor_lr
      # This lr_scheduler is not available in deepspeed
      # see https://deepspeed.readthedocs.io/en/latest/schedulers.html
      if not self.is_deepspeed_init:
          self.actor_scheduler = CosineAnnealingWarmRestarts(
              self.actor_optimizer,
              T_0=len(dataset),
              eta_min=actor_lr * 0.1,
          )
  
      # initialize scheduler for critic
      critic_lr = self.config.trainer.critic_lr
      # This lr_scheduler is not available in deepspeed
      # see https://deepspeed.readthedocs.io/en/latest/schedulers.html
      if not self.is_deepspeed_init:
          self.critic_scheduler = CosineAnnealingWarmRestarts(
              self.critic_optimizer,
              T_0=len(dataset),
              eta_min=critic_lr * 0.1,
          )
  
      # initialize actor accelerate
      if self.config.actor.accelerate_enable is True:
          actor_accelerator = Accelerator()
          (
              actor_model,
              self.actor_optimizer,
              self.train_dataloader,
              self.actor_scheduler,
          ) = actor_accelerator.prepare(
              self.actorcritic.actor,
              self.actor_optimizer,
              self.train_dataloader,
              self.actor_scheduler,
          )
          self.actorcritic.actor = actor_model
  
      # initialize critic accelerate
      if self.config.critic.accelerate_enable is True:
          critic_accelerator = Accelerator()
          (
              critic_model,
              self.critic_optimizer,
              self.critic_scheduler,
          ) = critic_accelerator.prepare(
              self.actorcritic.critic,
              self.critic_optimizer,
              self.critic_scheduler,
          )
          self.actorcritic.critic = critic_model
  
      # train agent-critic
      self.actorcritic.train()
      for epoch in range(epochs):
          for k, batch in enumerate(dataloader):
              (
                  states_actor,
                  old_actions,
                  old_values,
                  rewards,
                  old_actions_log_probs,
                  sequences_actor,
                  sequences_mask_actor,
                  sequences_critic,
                  sequences_mask_critic,
                  action_len_actor,
                  action_len_critic,
              ) = [tensor.to(device) for tensor in batch]
  
              if self.debug:
                  print(
                      f"#########################################"
                      f" batch from memories {k} \n "
                      f"#########################################"
                      f"states_actor {states_actor.shape} \n"
                      f"old_actions {old_actions.shape} \n"
                      f"old_values {old_values.shape} \n"
                      f"rewards {rewards.shape} \n"
                      f"old_actions_log_probs "
                      f"{old_actions_log_probs.shape}\n"
                      f"sequences_actor {sequences_actor.shape} \n"
                      f"sequences_mask_actor "
                      f"{sequences_mask_actor.shape} \n"
                      f"sequences_critic {sequences_critic.shape} \n"
                      f"sequences_mask_critic "
                      f"{sequences_mask_critic.shape} \n"
                      f"action_len_actor {action_len_actor} \n"
                      f"action_len_critic {action_len_critic} \n"
                      f"#########################################"
                  )
  
              # get actor critic new probabilities and values
              actions_logits, values = self.actorcritic.forward(
                  sequences_actor,
                  sequences_mask_actor,
                  sequences_critic,
                  sequences_mask_critic,
                  action_len_actor.item(),
                  action_len_critic.item(),
              )
  
              # get action log prob
              actions_prob = (
                  torch.softmax(actions_logits, dim=-1)
                  .max(dim=-1)
                  .values
              )
              actions_log_prob = torch.log(actions_prob + self.eps)
  
              # compute entropy
              entropies = (actions_prob * actions_log_prob).sum(dim=-1)
  
              # compute KL divergence
              kl_div_loss = (
                  (
                      actions_prob
                      * (old_actions_log_probs - actions_log_prob)
                  )
                  .sum(dim=-1)
                  .mean()
              )
  
              # compute ratios
              ratios = (actions_log_prob - old_actions_log_probs).exp()
  
              # compute PPO loss
              if check_model_family(
                  self.config.actor, self.config.critic
              ):
                  # compute discounted rewards as in TRL
                  gamma = self.config.trainer.gamma_discounted
                  discounted_rewards = torch.zeros_like(old_values)
                  for i in range(discounted_rewards.shape[1]):
                      for j in range(i, discounted_rewards.shape[1]):
                          discounted_rewards[:, i] += (
                              gamma ** (j - i) * rewards[:, j]
                          )
  
                  advantages = (
                      discounted_rewards - old_values
                  )  # TRL has opposite sign for old values
                  advantages = (
                      advantages - advantages.mean(dim=-1)
                  ) / (advantages.std() + self.eps)
  
                  surr1 = advantages * ratios
              else:
                  advantages = rewards - old_values[:, -1]
                  surr1 = advantages * ratios
  
              surr2 = (
                  torch.clamp(
                      ratios, 1 - actor_eps_clip, 1 + actor_eps_clip
                  )
                  * advantages
              )
  
              policy_loss = (
                  -torch.min(surr1, surr2) - beta_s * entropies
              )
              policy_loss = policy_loss.mean()
              loss = policy_loss + kl_div_loss
  
              # check if loss item is NaN
              if torch.isnan(loss):
                  raise ValueError("Loss is nan")
  
              # update actor with loss
              if self.config.actor.deepspeed_enable:
                  self.actor_model_engine.backward(loss)
                  self.actor_model_engine.step()
              elif self.config.actor.accelerate_enable:
                  self.actor_optimizer.zero_grad()
                  actor_accelerator.backward(loss)
                  self.actor_optimizer.step()
                  self.actor_scheduler.step()
              else:
                  self.actor_optimizer.zero_grad()
                  loss.backward()
                  self.actor_optimizer.step()
                  self.actor_scheduler.step()
  
              # compute value loss
              # the loss is the distance between the rewards and the values
              # I want this distance to be small so that values are
              # representative of the rewards, for this reason i took the
              # maximum between the two.
              # The clip is limiting the slew-rate of values_loss_clipped
              value_loss_clipped = old_values + (
                  values - old_values
              ).clamp(-critic_eps_clip, critic_eps_clip)
              value_loss1 = (value_loss_clipped - rewards) ** 2
              value_loss2 = (values - rewards) ** 2
              value_loss = torch.max(value_loss1, value_loss2).mean()
  
              if torch.isnan(value_loss):
                  raise ValueError("Value loss is nan")
  
              # upate critic
              if self.config.critic.deepspeed_enable:
                  self.critic_model_engine.backward(value_loss)
                  self.critic_model_engine.step()
              elif self.config.critic.accelerate_enable:
                  self.critic_optimizer.zero_grad()
                  critic_accelerator.backward(loss)
                  self.critic_optimizer.step()
                  self.critic_scheduler.step()
              else:
                  self.critic_optimizer.zero_grad()
                  value_loss.backward()
                  self.critic_optimizer.step()
                  self.critic_scheduler.step()
  
              # append the losses to the training stats
              self.training_stats.training_loss.append(
                  loss.detach().cpu().item()
              )
              self.training_stats.value_loss.append(
                  value_loss.detach().cpu().item()
              )
  
              # print iteration info
              print(
                  f"Epoch {epoch+1}/{epochs}",
                  f"Step {k+1}/{int(len(dataloader) / batch_size)}",
                  f"Loss {loss.detach().cpu().item():.4f}",
                  f"Value Loss {value_loss.detach().cpu().item():.4f}",
              )
  
      self.actorcritic.eval()
      print("End Learning")
  
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 RLTrainer obj's underlying member method code with example usage: obj.learn(memories) without return types