def __init__(self, task,
               value_model=None,
               value_optimizer=None,
               value_lr_schedule=lr.multifactor,
               value_batch_size=64,
               value_train_steps_per_epoch=500,
               value_evals_per_epoch=1,
               value_eval_steps=1,
               n_shared_layers=0,
               added_policy_slice_length=0,
               n_replay_epochs=1,
               scale_value_targets=False,
               q_value=False,
               q_value_aggregate='logsumexp',
               q_value_temperature=1.0,
               q_value_n_samples=1,
               q_value_normalization=False,
               offline=False,
               **kwargs):  # Arguments of PolicyAgent come here.
    """Configures the actor-critic trainer.

    Args:
      task: `RLTask` instance to use.
      value_model: Model to use for the value function.
      value_optimizer: Optimizer to train the value model.
      value_lr_schedule: lr schedule for value model training.
      value_batch_size: Batch size for value model training.
      value_train_steps_per_epoch: Number of steps are we using to train the
          value model in each epoch.
      value_evals_per_epoch: Number of value trainer evaluations per RL epoch.
          Every evaluation, we also synchronize the weights of the target
          network.
      value_eval_steps: Number of value trainer steps per evaluation; only
          affects metric reporting.
      n_shared_layers: Number of layers to share between value and policy
          models.
      added_policy_slice_length: How much longer should slices of
          trajectories be for policy than for value training; this
          is useful for TD calculations and only affect the length
          of elements produced for policy batches; value batches
          have maximum length set by `max_slice_length` in `**kwargs`.
      n_replay_epochs: Number of last epochs to take into the replay buffer;
          only makes sense for off-policy algorithms.
      scale_value_targets: If `True`, scale value function targets by
          `1 / (1 - gamma)`.
      q_value: If `True`, use Q-values as baselines.
      q_value_aggregate: How to aggregate Q-values. Options: 'mean', 'max',
        'softmax', 'logsumexp'.
      q_value_temperature: Temperature parameter for the 'softmax' and
        'logsumexp' aggregation methods.
      q_value_n_samples: Number of samples to average over when calculating
          baselines based on Q-values.
      q_value_normalization: How to normalize Q-values before aggregation.
          Allowed values: 'std', 'abs', `None`. If `None`, don't normalize.
      offline: Whether to train in offline mode. This matters for some
        algorithms, e.g. QWR.
      **kwargs: Arguments for `PolicyAgent` superclass.
    """
    self._n_shared_layers = n_shared_layers
    self._value_batch_size = value_batch_size
    self._value_train_steps_per_epoch = value_train_steps_per_epoch
    self._value_evals_per_epoch = value_evals_per_epoch
    self._value_eval_steps = value_eval_steps

    # The 2 below will be initalized in super.__init__ anyway, but are needed
    # to construct value batches which are needed before PolicyAgent init
    # since policy input creation calls the value model -- hence this code.
    self._task = task
    self._max_slice_length = kwargs.get('max_slice_length', 1)
    self._added_policy_slice_length = added_policy_slice_length
    self._n_replay_epochs = n_replay_epochs
    task.set_n_replay_epochs(n_replay_epochs)

    if scale_value_targets:
      self._value_network_scale = 1 / (1 - self._task.gamma)
    else:
      self._value_network_scale = 1

    self._q_value = q_value
    self._q_value_aggregate = q_value_aggregate
    self._q_value_temperature = q_value_temperature
    self._q_value_n_samples = q_value_n_samples
    self._q_value_normalization = q_value_normalization

    is_discrete = isinstance(self._task.action_space, gym.spaces.Discrete)
    self._is_discrete = is_discrete
    self._vocab_size = None
    self._sample_all_discrete_actions = False
    if q_value and is_discrete:
      self._vocab_size = self.task.action_space.n
      # TODO(lukaszkaiser): the code below is specific to AWR, move it.
      # If n_samples = n_actions, we'll take them all in actor and reweight.
      if self._q_value_n_samples == self._vocab_size:
        # TODO(lukaszkaiser): set this explicitly once it's in AWR Trainer.
        self._sample_all_discrete_actions = True
    if offline and is_discrete:
      raise NotImplementedError(
          'Offline training is only supported for continuous action spaces for '
          'now.'
      )
    self._offline = offline

    if q_value:
      value_model = functools.partial(value_model,
                                      inject_actions=True,
                                      is_discrete=is_discrete,
                                      vocab_size=self._vocab_size)
    self._value_eval_model = value_model(mode='eval')
    self._value_eval_model.init(self._value_model_signature)
    self._value_eval_jit = tl.jit_forward(
        self._value_eval_model.pure_fn, fastmath.local_device_count(),
        do_mean=False)

    # Initialize policy training.
    super().__init__(task, **kwargs)

    # Initialize training of the value function.
    value_output_dir = kwargs.get('output_dir', None)
    if value_output_dir is not None:
      value_output_dir = os.path.join(value_output_dir, 'value')
      # If needed, create value_output_dir and missing parent directories.
      if not tf.io.gfile.isdir(value_output_dir):
        tf.io.gfile.makedirs(value_output_dir)
    self._value_inputs = data.inputs.Inputs(
        train_stream=lambda _: self.value_batches_stream())
    self._value_trainer = supervised.Trainer(
        model=value_model,
        optimizer=value_optimizer,
        lr_schedule=value_lr_schedule(),
        loss_fn=tl.L2Loss(),
        inputs=self._value_inputs,
        output_dir=value_output_dir,
        metrics={'value_loss': tl.L2Loss(),
                 'value_mean': self.value_mean})
Configures the actor-critic trainer.

    Args:
      task: `RLTask` instance to use.
      value_model: Model to use for the value function.
      value_optimizer: Optimizer to train the value model.
      value_lr_schedule: lr schedule for value model training.
      value_batch_size: Batch size for value model training.
      value_train_steps_per_epoch: Number of steps are we using to train the
          value model in each epoch.
      value_evals_per_epoch: Number of value trainer evaluations per RL epoch.
          Every evaluation, we also synchronize the weights of the target
          network.
      value_eval_steps: Number of value trainer steps per evaluation; only
          affects metric reporting.
      n_shared_layers: Number of layers to share between value and policy
          models.
      added_policy_slice_length: How much longer should slices of
          trajectories be for policy than for value training; this
          is useful for TD calculations and only affect the length
          of elements produced for policy batches; value batches
          have maximum length set by `max_slice_length` in `**kwargs`.
      n_replay_epochs: Number of last epochs to take into the replay buffer;
          only makes sense for off-policy algorithms.
      scale_value_targets: If `True`, scale value function targets by
          `1 / (1 - gamma)`.
      q_value: If `True`, use Q-values as baselines.
      q_value_aggregate: How to aggregate Q-values. Options: 'mean', 'max',
        'softmax', 'logsumexp'.
      q_value_temperature: Temperature parameter for the 'softmax' and
        'logsumexp' aggregation methods.
      q_value_n_samples: Number of samples to average over when calculating
          baselines based on Q-values.
      q_value_normalization: How to normalize Q-values before aggregation.
          Allowed values: 'std', 'abs', `None`. If `None`, don't normalize.
      offline: Whether to train in offline mode. This matters for some
        algorithms, e.g. QWR.
      **kwargs: Arguments for `PolicyAgent` superclass.
    
trax/rl/actor_critic.py