def __init__(self, task, model_fn, **kwargs):
    """Initializes PolicyGradient.

    Args:
      task: Instance of trax.rl.task.RLTask.
      model_fn: Function (policy_distribution, mode) -> policy_model.
      **kwargs: Keyword arguments passed to the superclass.
    """
    super().__init__(
        task, model_fn,
        # We're on-policy, so we can only use data from the last epoch.
        n_replay_epochs=1,
        # Each gradient computation needs a new data sample, so we do 1 step
        # per epoch.
        n_train_steps_per_epoch=1,
        # Very simple baseline: mean return across trajectories.
        value_fn=self._value_fn,
        # Weights are just advantages.
        weight_fn=(lambda x: x),
        # Normalize advantages, because this makes optimization nicer.
        advantage_normalization=True,
        **kwargs
    )
Initializes PolicyGradient.

    Args:
      task: Instance of trax.rl.task.RLTask.
      model_fn: Function (policy_distribution, mode) -> policy_model.
      **kwargs: Keyword arguments passed to the superclass.
    
trax/rl/training.py