(
      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.
Initializes ExpertIteration.
Initializes ExpertIteration .
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
    )
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 ) : 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 )