class_docstr:     Base class for on-policy rl trainers, e.g. PPO.
  
  Args:
      strategy (Strategy):the strategy to use for training
      data_buffer (NaiveExperienceBuffer): the buffer to collect experiences
      sample_buffer (bool, defaults to False): whether to sample from buffer
      dataloader_pin_memory (bool, defaults to True): whether to pin memory for data loader
      callbacks (List[Callback], defaults to []): the callbacks to call during training process
class_docstr_tok: ["Base","class","for","on-policy","rl","trainers",",","e.g",".","PPO",".","Args",":","strategy","(","Strategy",")",":","the","strategy","to","use","for","training","data_buffer","(","NaiveExperienceBuffer",")",":","the","buffer","to","collect","experiences","sample_buffer","(","bool",",","defaults","to","False",")",":","whether","to","sample","from","buffer","dataloader_pin_memory","(","bool",",","defaults","to","True",")",":","whether","to","pin","memory","for","data","loader","callbacks","(","List","[","Callback","]",",","defaults","to","[","]",")",":","the","callbacks","to","call","during","training","process"]
docstr: The main training loop of on-policy rl trainers.
  
  Args:
      num_episodes (int): the number of episodes to train
      num_collect_steps (int): the number of collect steps per episode
      num_update_steps (int): the number of update steps per episode
docstr_tok: ["The","main","training","loop","of","on-policy","rl","trainers",".","Args",":","num_episodes","(","int",")",":","the","number","of","episodes","to","train","num_collect_steps","(","int",")",":","the","number","of","collect","steps","per","episode","num_update_steps","(","int",")",":","the","number","of","update","steps","per","episode"]
code: def fit(
      self,
      num_episodes: int,
      num_collect_steps: int,
      num_update_steps: int,
      *args,
      **kwargs,
  ):
      """
      The main training loop of on-policy rl trainers.
  
      Args:
          num_episodes (int): the number of episodes to train
          num_collect_steps (int): the number of collect steps per episode
          num_update_steps (int): the number of update steps per episode
      """
      self._before_fit(*args, **kwargs)
      with self._fit_ctx():
          for episode in tqdm.trange(
              num_episodes, desc="Episodes", disable=not is_rank_0()
          ):
              with self._episode_ctx(episode):
                  for collect_step in tqdm.trange(
                      num_collect_steps,
                      desc="Collect steps",
                      disable=not is_rank_0(),
                  ):
                      self._collect_phase(collect_step)
                  if not self.sample_buffer:
                      # HACK(cwher): according to the design of boost API, dataloader should also be boosted,
                      #  but it is impractical to adapt this pattern in RL training. Thus, I left dataloader unboosted.
                      #  I only call strategy.setup_dataloader() to setup dataloader.
                      self.dataloader = self.strategy.setup_dataloader(
                          self.data_buffer, self.dataloader_pin_memory
                      )
                  for update_step in tqdm.trange(
                      num_update_steps,
                      desc="Update steps",
                      disable=not is_rank_0(),
                  ):
                      self._update_phase(update_step)
                  # NOTE: this is for on-policy algorithms
                  self.data_buffer.clear()
  
imports: from abc import ABC, abstractmethod
  from contextlib import contextmanager
  from typing import List
  import torch.nn
  import tqdm
  from coati.experience_buffer import NaiveExperienceBuffer
  from coati.experience_maker import Experience
  from torch.optim import Optimizer
  from .callbacks import Callback
  from .strategies import Strategy
  from .utils import is_rank_0
label_desc: Use image node_id 15 for calling the OnPolicyTrainer obj's underlying member method code with example usage: obj.fit(num_episodes, num_collect_steps, num_update_steps) without return types