def __init__(
      self, task, model_fn,
      expert_policy_fn=sharpened_network_policy,
      quantile=0.9,
      n_replay_epochs=10,
      n_train_steps_per_epoch=1000,
      filter_buffer_size=256,
      **kwargs
  ):
    """Initializes ExpertIteration.

    Args:
      task: Instance of trax.rl.task.RLTask.
      model_fn: Function (policy_distribution, mode) -> policy_model.
      expert_policy_fn: Function of the same signature as `network_policy`, to
        be used as an expert. The policy will be trained to mimic the expert on
        the "solved" trajectories.
      quantile: Quantile of best trajectories to be marked as "solved". They
        will be used to train the policy.
      n_replay_epochs: Number of last epochs to include in the replay buffer.
      n_train_steps_per_epoch: Number of policy training steps to run in each
        epoch.
      filter_buffer_size: Number of trajectories in the trajectory filter
        buffer, used to select the best trajectories based on the quantile.
      **kwargs: Keyword arguments passed to the superclass.
    """
    self._expert_policy_fn = expert_policy_fn
    self._quantile = quantile
    self._filter_buffer_size = filter_buffer_size
    super().__init__(
        task, model_fn,
        # Don't use a baseline - it's not useful in our weights.
        value_fn=(lambda batch: jnp.zeros_like(batch.return_)),
        # Don't weight trajectories - the training signal is provided by
        # filtering trajectories.
        weight_fn=jnp.ones_like,
        # Filter trajectories based on the quantile.
        trajectory_stream_preprocessing_fn=self._filter_trajectories,
        # Advantage normalization is a no-op here.
        advantage_normalization=False,
        n_replay_epochs=n_replay_epochs,
        n_train_steps_per_epoch=n_train_steps_per_epoch,
        **kwargs
    )
Initializes ExpertIteration.

    Args:
      task: Instance of trax.rl.task.RLTask.
      model_fn: Function (policy_distribution, mode) -> policy_model.
      expert_policy_fn: Function of the same signature as `network_policy`, to
        be used as an expert. The policy will be trained to mimic the expert on
        the "solved" trajectories.
      quantile: Quantile of best trajectories to be marked as "solved". They
        will be used to train the policy.
      n_replay_epochs: Number of last epochs to include in the replay buffer.
      n_train_steps_per_epoch: Number of policy training steps to run in each
        epoch.
      filter_buffer_size: Number of trajectories in the trajectory filter
        buffer, used to select the best trajectories based on the quantile.
      **kwargs: Keyword arguments passed to the superclass.
    
trax/rl/training.py