class_docstr: Trains a joint policy-and-value model using actor-critic methods.
class_docstr_tok: ["Trains","a","joint","policy-and-value","model","using","actor-critic","methods","."]
docstr: Configures the joint trainer.
  
  Args:
    task: RLTask instance, which defines the environment to train on.
    joint_model: Trax layer, representing the joint policy and value model.
    optimizer: the optimizer to use to train the joint model.
    lr_schedule: learning rate schedule to use to train the joint model/.
    batch_size: batch size used to train the joint model.
    train_steps_per_epoch: how long to train the joint model in each RL epoch.
    supervised_evals_per_epoch: number of value trainer evaluations per RL
        epoch - only affects metric reporting.
    supervised_eval_steps: number of value trainer steps per evaluation -
        only affects metric reporting.
    n_trajectories_per_epoch: how many trajectories to collect per epoch.
    max_slice_length: the maximum length of trajectory slices to use.
    normalize_advantages: if True, then normalize advantages - currently
        implemented only in PPO.
    output_dir: Path telling where to save outputs (evals and checkpoints).
    n_replay_epochs: how many last epochs to take into the replay buffer;
         > 1 only makes sense for off-policy algorithms.
docstr_tok: ["Configures","the","joint","trainer",".","Args",":","task",":","RLTask","instance",",","which","defines","the","environment","to","train","on",".","joint_model",":","Trax","layer",",","representing","the","joint","policy","and","value","model",".","optimizer",":","the","optimizer","to","use","to","train","the","joint","model",".","lr_schedule",":","learning","rate","schedule","to","use","to","train","the","joint","model\/",".","batch_size",":","batch","size","used","to","train","the","joint","model",".","train_steps_per_epoch",":","how","long","to","train","the","joint","model","in","each","RL","epoch",".","supervised_evals_per_epoch",":","number","of","value","trainer","evaluations","per","RL","epoch","-","only","affects","metric","reporting",".","supervised_eval_steps",":","number","of","value","trainer","steps","per","evaluation","-","only","affects","metric","reporting",".","n_trajectories_per_epoch",":","how","many","trajectories","to","collect","per","epoch",".","max_slice_length",":","the","maximum","length","of","trajectory","slices","to","use",".","normalize_advantages",":","if","True",",","then","normalize","advantages","-","currently","implemented","only","in","PPO",".","output_dir",":","Path","telling","where","to","save","outputs","(","evals","and","checkpoints",")",".","n_replay_epochs",":","how","many","last","epochs","to","take","into","the","replay","buffer",";",">","1","only","makes","sense","for","off-policy","algorithms","."]
code: def __init__(
      self,
      task,
      joint_model=None,
      optimizer=None,
      lr_schedule=lr.multifactor,
      batch_size=64,
      train_steps_per_epoch=500,
      supervised_evals_per_epoch=1,
      supervised_eval_steps=1,
      n_trajectories_per_epoch=50,
      max_slice_length=1,
      normalize_advantages=True,
      output_dir=None,
      n_replay_epochs=1,
  ):
      """Configures the joint trainer.
  
      Args:
        task: RLTask instance, which defines the environment to train on.
        joint_model: Trax layer, representing the joint policy and value model.
        optimizer: the optimizer to use to train the joint model.
        lr_schedule: learning rate schedule to use to train the joint model/.
        batch_size: batch size used to train the joint model.
        train_steps_per_epoch: how long to train the joint model in each RL epoch.
        supervised_evals_per_epoch: number of value trainer evaluations per RL
            epoch - only affects metric reporting.
        supervised_eval_steps: number of value trainer steps per evaluation -
            only affects metric reporting.
        n_trajectories_per_epoch: how many trajectories to collect per epoch.
        max_slice_length: the maximum length of trajectory slices to use.
        normalize_advantages: if True, then normalize advantages - currently
            implemented only in PPO.
        output_dir: Path telling where to save outputs (evals and checkpoints).
        n_replay_epochs: how many last epochs to take into the replay buffer;
             > 1 only makes sense for off-policy algorithms.
      """
      super().__init__(
          task,
          n_trajectories_per_epoch=n_trajectories_per_epoch,
          output_dir=output_dir,
      )
      self._batch_size = batch_size
      self._train_steps_per_epoch = train_steps_per_epoch
      self._supervised_evals_per_epoch = supervised_evals_per_epoch
      self._supervised_eval_steps = supervised_eval_steps
      self._n_trajectories_per_epoch = n_trajectories_per_epoch
      self._max_slice_length = max_slice_length
      self._policy_dist = distributions.create_distribution(
          task.action_space
      )
      self._lr_schedule = lr_schedule()
      self._optimizer = optimizer
      self._normalize_advantages = normalize_advantages
      self._n_replay_epochs = n_replay_epochs
      self._task.set_n_replay_epochs(n_replay_epochs)
  
      # Inputs to the joint model are produced by self.batches_stream.
      self._inputs = data.inputs.Inputs(
          train_stream=lambda _: self.batches_stream()
      )
  
      self._joint_model = functools.partial(
          joint_model,
          policy_distribution=self._policy_dist,
      )
  
      # This is the joint Trainer that will be used to train the policy model.
      # * inputs to the trainer come from self.batches_stream
      # * outputs are passed to self._joint_loss
      self._trainer = supervised.Trainer(
          model=self._joint_model,
          optimizer=self._optimizer,
          lr_schedule=self._lr_schedule,
          loss_fn=self.joint_loss,
          inputs=self._inputs,
          output_dir=output_dir,
          metrics={
              "joint_loss": self.joint_loss,
              "advantage_mean": self.advantage_mean,
              "advantage_norm": self.advantage_norm,
              "value_loss": self.value_loss,
              "explained_variance": self.explained_variance,
              "log_probs_mean": self.log_probs_mean,
              "preferred_move": self.preferred_move,
          },
      )
      self._eval_model = tl.Accelerate(
          self._joint_model(mode="eval"), n_devices=1
      )
      example_batch = next(self.batches_stream())
      self._eval_model.init(example_batch)
  
imports: import functools
  from trax import data
  from trax import layers
  from trax import supervised
  from trax.fastmath import numpy
  from trax.fastmath import stop_gradient
  from trax.rl import actor_critic
  from trax.rl import distributions
  from trax.rl import rl_layers
  from trax.rl import training
  from trax.supervised import lr_schedules
label_desc: Use image node_id 1 to create a new ActorCriticJointAgent object from inherited base classes: rl_training with example: obj = ActorCriticJointAgent(task, joint_model, optimizer, lr_schedule, batch_size, train_steps_per_epoch, supervised_evals_per_epoch, supervised_eval_steps, n_trajectories_per_epoch, max_slice_length, normalize_advantages, output_dir, n_replay_epochs)