(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.
Initializes PolicyGradient.
Initializes PolicyGradient .
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
    )
def __init__ ( self , task , model_fn , * * kwargs ) : 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 )