google / trax / 4341

# coding=utf-8

# Copyright 2020 The Trax Authors.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

#     http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# Lint as: python3

  """Abstract class for RL agents, presenting the required API."""

               n_trajectories_per_epoch=None,

               n_interactions_per_epoch=None,

               n_eval_episodes=0,

               eval_steps=None,

               only_eval=False,

               output_dir=None,

               timestep_to_np=None):

    """Configures the Agent.

    Note that subclasses can have many more arguments, which will be configured

    using defaults and gin. But task and output_dir are passed explicitly.

    Args:

      task: RLTask instance, which defines the environment to train on.

      n_trajectories_per_epoch: How many new trajectories to collect in each

        epoch.

      n_interactions_per_epoch: How many interactions to collect in each epoch.

      n_eval_episodes: Number of episodes to play with policy at

        temperature 0 in each epoch -- used for evaluation only.

      eval_steps: an optional list of max_steps to use for evaluation

        (defaults to task.max_steps).

      only_eval: If set to True, then trajectories are collected only for

        for evaluation purposes, but they are not recorded.

      output_dir: Path telling where to save outputs such as checkpoints.

      timestep_to_np: Timestep-to-numpy function to override in the task.

    """

        'Exactly one of n_trajectories_per_epoch or n_interactions_per_epoch '

        'should be specified.'

    )

  def current_epoch(self):

    """Returns current step number in this training session."""

  def task(self):

    """Returns the task."""

  def avg_returns(self):

          'gin_config', jaxboard.markdownify_operative_config_str(config_str)

      )

                   task_file_name='trajectories.pkl'):

    """Save current epoch number and average returns to file."""

                  'avg_returns': self._avg_returns,

                  'avg_returns_temperature0': self._avg_returns_temperature0}

                     task_file_name='trajectories.pkl'):

    """Initialize epoch number and average returns from file."""

        self.policy,

        n_trajectories=self._n_trajectories_per_epoch,

        n_interactions=self._n_interactions_per_epoch,

        only_eval=self._only_eval,

        epoch_id=self._epoch

    )

    """Policy function that allows to play using this trainer.

    Args:

      trajectory: an instance of trax.rl.task.Trajectory

      temperature: temperature used to sample from the policy (default=1.0)

    Returns:

      a pair (action, dist_inputs) where action is the action taken and

      dist_inputs is the parameters of the policy distribution, that will later

      be used for training.

    """

    raise NotImplementedError

    """Trains this Agent for one epoch -- main RL logic goes here."""

    raise NotImplementedError

  def _open_summary_writer(self):

    """Opens the Jaxboard summary writer wrapped by a context manager.

    Yields:

      A Jaxboard summary writer wrapped in a GeneratorContextManager object.

      Elements of the lists correspond to the training and evaluation task

      directories created during initialization. If there is no output_dir

      provided, yields None.

    """

      finally:

    else:

    """Runs this loop for n epochs.

    Args:

      n_epochs: Stop training after completing n steps.

      n_epochs_is_total_epochs: if True, consider n_epochs as the total

        number of epochs to train, including previously trained ones

    """

            'RL training took %.2f seconds.' % (time.time() - cur_time))

              'Collecting %d episodes took %.2f seconds.'

              % (self._n_trajectories_per_epoch, time.time() - cur_time))

        else:

              'Collecting %d interactions took %.2f seconds.'

              % (self._n_interactions_per_epoch, time.time() - cur_time))

            'Average return in epoch %d was %.2f.' % (self._epoch, avg_return))

                lambda x: self.policy(x, temperature=0.0),

                n_trajectories=self._n_eval_episodes,

                max_steps=steps, only_eval=True)

                avg_return_temperature0

            )

                'Avg return with temperature 0 at %d steps in epoch %d was %.2f.'

                % (steps, self._epoch, avg_return_temperature0))

                    step=self._epoch)

                        self._avg_returns_temperature0[steps][-1],

                        step=self._epoch)

                    step=self._epoch)

                    (self.task.n_interactions() - cur_n_interactions)/ \

                    (time.time() - cur_time),

                    step=self._epoch)

                    step=self._epoch)

  """Agent that uses a deep learning model for policy.

  Many deep RL methods, such as policy gradient (REINFORCE) or actor-critic fall

  into this category, so a lot of classes will be subclasses of this one. But

  some methods only have a value or Q function, these are different.

  """

               policy_lr_schedule=lr.multifactor, policy_batch_size=64,

               policy_train_steps_per_epoch=500, policy_evals_per_epoch=1,

               policy_eval_steps=1, n_eval_episodes=0,

               only_eval=False, max_slice_length=1, output_dir=None, **kwargs):

    """Configures the policy trainer.

    Args:

      task: RLTask instance, which defines the environment to train on.

      policy_model: Trax layer, representing the policy model.

          functions and eval functions (a.k.a. metrics) are considered to be

          outside the core model, taking core model output and data labels as

          their two inputs.

      policy_optimizer: the optimizer to use to train the policy model.

      policy_lr_schedule: learning rate schedule to use to train the policy.

      policy_batch_size: batch size used to train the policy model.

      policy_train_steps_per_epoch: how long to train policy in each RL epoch.

      policy_evals_per_epoch: number of policy trainer evaluations per RL epoch

          - only affects metric reporting.

      policy_eval_steps: number of policy trainer steps per evaluation - only

          affects metric reporting.

      n_eval_episodes: number of episodes to play with policy at

        temperature 0 in each epoch -- used for evaluation only

      only_eval: If set to True, then trajectories are collected only for

        for evaluation purposes, but they are not recorded.

      max_slice_length: the maximum length of trajectory slices to use.

      output_dir: Path telling where to save outputs (evals and checkpoints).

      **kwargs: arguments for the superclass Agent.

    """

        task,

        n_eval_episodes=n_eval_episodes,

        output_dir=output_dir,

        **kwargs

    )

    # Inputs to the policy model are produced by self._policy_batches_stream.

        train_stream=lambda _: self.policy_batches_stream())

        policy_model,

        policy_distribution=self._policy_dist,

    )

    # This is the policy Trainer that will be used to train the policy model.

    # * inputs to the trainer come from self.policy_batches_stream

    # * outputs, targets and weights are passed to self.policy_loss

        model=policy_model,

        optimizer=policy_optimizer,

        lr_schedule=policy_lr_schedule(),

        loss_fn=self.policy_loss,

        inputs=self._policy_inputs,

        output_dir=output_dir,

        metrics=self.policy_metrics,

    )

        policy_model(mode='collect'), n_devices=1)

        policy_model(mode='eval'), n_devices=1)  # Not collecting stats

  def policy_loss(self):

    """Policy loss."""

  def policy_metrics(self):

    """Use self.task to create inputs to the policy model."""

    """Chooses an action to play after a trajectory."""

    # Add batch dimension to trajectory_np and run the model.

    # Pick element 0 from the batch (the only one), last (current) timestep.

    """Trains RL for one epoch."""

    # When restoring, calculate how many evals are remaining.

        self._policy_trainer.step,

        self._epoch,

        self._policy_train_steps_per_epoch,

        self._policy_evals_per_epoch)

          self._policy_train_steps_per_epoch // self._policy_evals_per_epoch,

          self._policy_eval_steps)

  """Helper function to calculate remaining evaluations for a trainer.

  Args:

    cur_step: current step of the supervised trainer

    epoch: current epoch of the RL trainer

    train_steps_per_epoch: supervised trainer steps per RL epoch

    evals_per_epoch: supervised trainer evals per RL epoch

  Returns:

    number of remaining evals to do this epoch

  Raises:

    ValueError if the provided numbers indicate a step mismatch

  """

                     % (cur_step, prev_steps))

                     % (done_steps_this_epoch, train_steps_per_eval))

  """Trains a policy model using policy gradient on the given RLTask."""

      self, task, model_fn,

      optimizer=adam.Adam,

      lr_schedule=lr.multifactor,

      batch_size=64,

      network_eval_at=None,

      n_eval_batches=1,

      max_slice_length=1,

      **kwargs

  ):

    """Initializes PolicyGradient.

    Args:

      task: Instance of trax.rl.task.RLTask.

      model_fn: Function (policy_distribution, mode) -> policy_model.

      optimizer: Optimizer for network training.

      lr_schedule: Learning rate schedule for network training.

      batch_size: Batch size for network training.

      network_eval_at: Function step -> bool indicating the training steps, when

        network evaluation should be performed.

      n_eval_batches: Number of batches to run during network evaluation.

      max_slice_length: The length of trajectory slices to run the network on.

      **kwargs: Keyword arguments passed to the superclass.

    """

        batch_size,

        epochs=[-1],

        max_slice_length=self._max_slice_length,

        sample_trajectories_uniformly=True,

    )

        trajectory_batch_stream,

        optimizer(),

        lr_schedule(),

        self._policy_dist,

        # Policy gradient uses the MC estimator. No need for margin - the MC

        # estimator only uses empirical returns.

        advantage_estimator=advantages.monte_carlo(task.gamma, margin=0),

        value_fn=self._value_fn,

    )

        model_fn,

        policy_distribution=self._policy_dist,

    )

    else:

    # Checkpoint every epoch. We do one step per epoch, so that's every step.

        model=model_fn(mode='train'),

        tasks=[train_task],

        eval_model=model_fn(mode='eval'),

        eval_tasks=[eval_task],

        output_dir=policy_output_dir,

        eval_at=network_eval_at,

        checkpoint_at=checkpoint_at,

    )

    # Validate the restored checkpoints. The number of network training steps

    # (self.loop.step) should be equal to the number of epochs (self._epoch),

    # because we do exactly one gradient step per epoch.

    # TODO(pkozakowski): Move this to the base class once all Agents use Loop.

          'The number of Loop steps must equal the number of Agent epochs, '

          'got {} and {}.'.format(self.loop.step, self._epoch)

      )

  def loop(self):

    """Loop exposed for testing."""

    """Policy function that allows to play using this agent."""

        collect_model=self._collect_model,

        policy_distribution=self._policy_dist,

        loop=self.loop,

        trajectory_np=trajectory_np,

        temperature=temperature,

    )

    """Trains RL for one epoch."""

    # Perform one gradient step per training epoch to ensure we stay on policy.

  def _value_fn(trajectory_batch):

    # Estimate the value of every state as the mean return across trajectories

    # and timesteps in a batch.

    collect_model, policy_distribution, loop, trajectory_np, temperature=1.0

):

  """Policy function powered by a neural network.

  Used to implement Agent.policy() in policy-based agents.

  Args:

    collect_model: the model used for collecting trajectories

    policy_distribution: an instance of trax.rl.distributions.Distribution

    loop: trax.supervised.training.Loop used to train the policy network

    trajectory_np: an instance of trax.rl.task.TrajectoryNp

    temperature: temperature used to sample from the policy (default=1.0)

  Returns:

    a pair (action, dist_inputs) where action is the action taken and

    dist_inputs is the parameters of the policy distribution, that will later

    be used for training.

  """

  else:

    # When evaluating (t != 1.0), use the evaluation model instead of the

    # collection model - some models accumulate normalization statistics

    # during data collection, and we don't want to do it in eval to avoid data

    # leakage.

  # Copying weights from loop.model should work, because the raw model's

  # weights should be updated automatically during training, but it doesn't.

  # TODO(pkozakowski): Debug.

  # Add batch dimension to trajectory_np and run the model.

      1, trajectory_np.observations.shape[0], policy_distribution.n_inputs

  )

  # Pick element 0 from the batch (the only one), last (current) timestep.

    # The result is composed of mutable numpy arrays. We copy them to avoid

    # accidental modification.

  """Trainer that uses a deep learning model for value function.

  Compute the loss using variants of the Bellman equation.

  """

               value_body=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,

               exploration_rate=functools.partial(

                   lr.multifactor,

                   factors='constant * decay_every',

                   constant=1.,  # pylint: disable=redefined-outer-name

                   decay_factor=0.99,

                   steps_per_decay=1,

                   minimum=0.1),

               n_eval_episodes=0,

               only_eval=False,

               n_replay_epochs=1,

               max_slice_length=1,

               sync_freq=1000,

               scale_value_targets=True,

               output_dir=None,

               **kwargs):

    """Configures the value trainer.

    Args:

      task: RLTask instance, which defines the environment to train on.

      value_body: Trax layer, representing the body of the value model.

          functions and eval functions (a.k.a. metrics) are considered to be

          outside the core model, taking core model output and data labels as

          their two inputs.

      value_optimizer: the optimizer to use to train the policy model.

      value_lr_schedule: learning rate schedule to use to train the policy.

      value_batch_size: batch size used to train the policy model.

      value_train_steps_per_epoch: how long to train policy in each RL epoch.

      value_evals_per_epoch: number of policy trainer evaluations per RL epoch

          - only affects metric reporting.

      value_eval_steps: number of policy trainer steps per evaluation - only

          affects metric reporting.

      exploration_rate: exploration rate schedule - used in the policy method.

      n_eval_episodes: number of episodes to play with policy at

        temperature 0 in each epoch -- used for evaluation only

      only_eval: If set to True, then trajectories are collected only for

        for evaluation purposes, but they are not recorded.

      n_replay_epochs: Number of last epochs to take into the replay buffer;

          only makes sense for off-policy algorithms.

      max_slice_length: the maximum length of trajectory slices to use; it is

          the second dimenions of the value network output:

          (batch, max_slice_length, number of actions)

          Higher max_slice_length implies that the network has to predict more

          values into the future.

      sync_freq: frequency when to synchronize the target

        network with the trained network. This is necessary for training the

        network on bootstrapped targets, e.g. using n-step returns.

      scale_value_targets: If `True`, scale value function targets by

          `1 / (1 - gamma)`. We are trying to fix the problem with very large

          returns in some games in a way which does not introduce an additional

          hyperparameters.

      output_dir: Path telling where to save outputs (evals and checkpoints).

      **kwargs: arguments for the superclass RLTrainer.

    """

        task,

        n_eval_episodes=n_eval_episodes,

        output_dir=output_dir,

        **kwargs

    )

    else:

        models.Quality,

        body=value_body,

        n_actions=self.task.action_space.n)

        self._value_eval_model.pure_fn, fastmath.device_count(), do_mean=False)

    # Inputs to the value model are produced by self._values_batches_stream.

        train_stream=lambda _: self.value_batches_stream())

    # This is the value Trainer that will be used to train the value model.

    # * inputs to the trainer come from self.value_batches_stream

    # * outputs, targets and weights are passed to self.value_loss

        model=value_model,

        optimizer=value_optimizer,

        lr_schedule=value_lr_schedule(),

        loss_fn=self.value_loss,

        inputs=self._inputs,

        output_dir=output_dir,

        metrics={'value_loss': self.value_loss,

                 'value_mean': self.value_mean,

                 'returns_mean': self.returns_mean}

    )

        value_model(mode='collect'), n_devices=1)

  def _value_model_signature(self):

        shape=(1, 1, self._task.action_space),

    )

    """Use self.task to create inputs to the policy model."""

    raise NotImplementedError

    """Chooses an action to play after a trajectory."""

    raise NotImplementedError

    """Trains RL for one epoch."""

    # Update the target value network.

    # When restoring, calculate how many evals are remaining.

        self._value_trainer.step,

        self._epoch,

        self._value_train_steps_per_epoch,

        self._value_evals_per_epoch)

          self._value_train_steps_per_epoch // self._value_evals_per_epoch,

          self._value_eval_steps)

          {'exploration_rate': self._exploration_rate(self._epoch)})

                                      self._value_trainer._train_sw, 'dqn')  # pylint: disable=protected-access

    # Update the target value network.

    # TODO(henrykm) a bit tricky if sync_at does not coincide with epochs

  def value_mean(self):

    """The mean value of actions selected by the behavioral policy."""

    raise NotImplementedError

  def returns_mean(self):

    """The mean value of actions selected by the behavioral policy."""

  r"""Trains a value model using DQN on the given RLTask.

  Notice that the algorithm and the parameters signficantly diverge from

  the original DQN paper. In particular we have separated learning and data

  collection.

  The Bellman loss is computed in the value_loss method. The formula takes

  the state-action values tensors Q and n-step returns R:

    .. math::

        L(s,a) = Q(s,a) - R(s,a)

  where R is computed in value_batches_stream. In the simplest case of the

  1-step returns we are getting

    .. math::

        L(s,a) = Q(s,a) - r(s,a) - gamma * \max_{a'} Q'(s',a')

  where s' is the state reached after taking action a in state s, Q' is

  the target network, gamma is the discount factor and the maximum is taken

  with respect to all actions avaliable in the state s'. The tensor Q' is

  updated using the sync_freq parameter.

  In code the maximum is visible in the policy method where we take

  sample = jnp.argmax(values). The epsilon-greedy policy is taking a random

  move with probability epsilon and oterhwise in state s it takes the

  action argmax_a Q(s,a).

  """

               task,

               advantage_estimator=advantages.monte_carlo,

               max_slice_length=1,

               smoothl1loss=True,

               double_dqn=False,

               **kwargs):

    # Our default choice of learning targets for DQN are n-step targets

    # implemented in the method td_k. We set the slice used for computation

    # of td_k to max_slice_length and we set the "margin" in td_k

    # to self._max_slice_length-1; in turn it implies that the shape of the

    # returned tensor of n-step targets is

    # values[:, :-(self.margin)] =  values[:, :1]

        gamma=task.gamma, margin=self._margin)

                              max_slice_length=max_slice_length,

                              **kwargs)

  def value_loss(self):

    """Value loss computed using smooth L1 loss or L2 loss."""

      # We calculate length using the shape of returns

      # and adequatly remove a superflous slice of values.

      # An analogous operation is done in value_batches_stream.

      else:

  def _replay_epochs(self):

    """Use the RLTask self._task to create inputs to the value model."""

        self._value_batch_size,

        max_slice_length=max_slice_length,

        min_slice_length=min_slice_length,

        margin=0,

        epochs=self._replay_epochs,

    ):

          np_trajectory.observations, use_eval_model=True)

            np_trajectory.observations, use_eval_model=False

        )

      else:

      # The advantage_estimator returns

      #     gamma^n_steps * values_max(s_{i + n_steps}) + discounted_rewards