google / trax / 83

pending completion

Build # 83

Build Type

Pull #534

travis-ci

Committed by

web-flow

Commit Message

<a href="https://github.com/google/trax/commit/<a class=hub.com/google/trax/commit/<a class="double-link" href="https://git"><a class=hub.com/google/trax/commit/<a class="double-link" href="https://git"><a class=hub.com/google/trax/commit/<a class="double-link" href="https://git"><a class=hub.com/google/trax/commit/be690807b88e39db74105740365c7e04fbdf8430">be690807b">&lt;a href=&quot;https://github.com/google/trax/commit/</a><a class="double-link" href="https://github.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/google/trax/commit/be690807b88e39db74105740365c7e04fbdf8430">be690807b</a><a href="https://github.com/google/trax/commit/be690807b88e39db74105740365c7e04fbdf8430">&quot;&gt;&amp;lt;a href=&amp;quot;https://github.com/google/trax/commit/&lt;/a&gt;&lt;a class=&quot;double-link&quot; href=&quot;https://github.com/google/trax/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/google/trax/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/google/trax/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/google/trax/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/google/trax/commit/be690807b88e39db74105740365c7e04fbdf8430&quot;&gt;be690807b&lt;/a&gt;&lt;a href=&quot;https://github.com/google/trax/commit/be690807b88e39db74105740365c7e04fbdf8430&quot;&gt;&amp;quot;&amp;gt;&amp;amp;lt;a href=&amp;amp;quot;https://github.com/google/trax/commit/&amp;lt;/a&amp;gt;&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/google/trax/commit/&amp;amp;lt;a class=&amp;amp;quot;double-link&amp;amp;quot; href=&amp;amp;quot;https://git&amp;quot;&amp;gt;&amp;amp... (continued)

Pull Request Pull Request #534: Report advantage metrics in AWR.

Run Details

23 of 23 new or added lines in 2 files covered. (100.0%)

2696 of 10815 relevant lines covered (24.93%)

0.25 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

24.61

/trax/rl/actor_critic.py

# 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
"""Classes for RL training in Trax."""

import os
import numpy as np
import tensorflow as tf

from trax import layers as tl
from trax import lr_schedules as lr
from trax import shapes
from trax import supervised
from trax.math import numpy as jnp
from trax.rl import advantages as rl_advantages
from trax.rl import training as rl_training


class ActorCriticTrainer(rl_training.PolicyTrainer):
  """Trains policy and value models using actor-critic methods.

  Attrs:
    on_policy (bool): Whether the algorithm is on-policy. Used in the data
      generators. Should be set in derived classes.
  """

  on_policy = None

  def __init__(self, task,
               value_model=None,
               value_optimizer=None,
               value_lr_schedule=lr.MultifactorSchedule,
               value_batch_size=64,
               value_train_steps_per_epoch=500,
               value_evals_per_epoch=1,
               value_eval_steps=1,
               n_shared_layers=0,
               added_policy_slice_length=0,
               n_replay_epochs=1,
               scale_value_targets=False,
               **kwargs):  # Arguments of PolicyTrainer come here.
    """Configures the actor-critic Trainer.

    Args:
      task: RLTask instance to use
      value_model: the model to use for the value function
      value_optimizer: the optimizer to train the value model
      value_lr_schedule: lr schedule for value model training
      value_batch_size: batch size for value model training
      value_train_steps_per_epoch: how many steps are we using to
        train the value model in each epoch
      value_evals_per_epoch: number of value trainer evaluations per RL
          epoch - only affects metric reporting.
      value_eval_steps: number of value trainer steps per evaluation -
          only affects metric reporting.
      n_shared_layers: how many layers to share between value and
        policy models
      added_policy_slice_length: how much longer should slices of
        trajectories be for policy than for value training; this
        is useful for TD calculations and only affect the length
        of elements produced for policy batches; value batches
        have maximum length set by max_slice_length in **kwargs
     n_replay_epochs: how many last epochs to take into the replay buffer;
        only makes sense for off-policy algorithms
     scale_value_targets: whether to scale targets for the value function by
        1 / (1 - gamma)
     **kwargs: arguments for PolicyTrainer super-class
    """
    self._n_shared_layers = n_shared_layers
    self._value_batch_size = value_batch_size
    self._value_train_steps_per_epoch = value_train_steps_per_epoch
    self._value_evals_per_epoch = value_evals_per_epoch
    self._value_eval_steps = value_eval_steps

    # The 2 below will be initalized in super.__init__ anyway, but are needed
    # to construct value batches which are needed before PolicyTrainer init
    # since policy input creation calls the value model -- hence this code.
    self._task = task
    self._max_slice_length = kwargs.get('max_slice_length', 1)
    self._added_policy_slice_length = added_policy_slice_length
    self._n_replay_epochs = n_replay_epochs

    if scale_value_targets:
      self._value_network_scale = 1 / (1 - self._task.gamma)
    else:
      self._value_network_scale = 1

    self._value_eval_model = value_model(mode='eval')
    self._value_eval_model.init(self._value_model_signature)

    # Initialize training of the value function.
    value_output_dir = kwargs.get('output_dir', None)
    if value_output_dir is not None:
      value_output_dir = os.path.join(value_output_dir, 'value')
      # If needed, create value_output_dir and missing parent directories.
      if not tf.io.gfile.isdir(value_output_dir):
        tf.io.gfile.makedirs(value_output_dir)
    self._value_inputs = supervised.Inputs(
        train_stream=lambda _: self.value_batches_stream())
    self._value_trainer = supervised.Trainer(
        model=value_model,
        optimizer=value_optimizer,
        lr_schedule=value_lr_schedule,
        loss_fn=tl.L2Loss(),
        inputs=self._value_inputs,
        output_dir=value_output_dir,
        metrics={'value_loss': tl.L2Loss()})

    # Initialize policy training.
    super(ActorCriticTrainer, self).__init__(task, **kwargs)

  @property
  def _value_model_signature(self):
    obs_sig = shapes.signature(self._task.observation_space)
    target_sig = mask_sig = shapes.ShapeDtype(
        shape=(1, 1, 1),
    )
    return (obs_sig.replace(shape=(1, 1) + obs_sig.shape), target_sig, mask_sig)

  @property
  def _replay_epochs(self):
    if self.on_policy:
      assert self._n_replay_epochs == 1, (
          'Non-unit replay buffer size only makes sense for off-policy '
          'algorithms.'
      )
    return [-(ep + 1) for ep in range(self._n_replay_epochs)]

  def value_batches_stream(self):
    """Use the RLTask self._task to create inputs to the value model."""
    max_slice_length = self._max_slice_length + self._added_policy_slice_length
    for np_trajectory in self._task.trajectory_batch_stream(
        self._value_batch_size,
        max_slice_length=max_slice_length,
        min_slice_length=(1 + self._added_policy_slice_length),
        epochs=self._replay_epochs,
    ):
      values = self._value_eval_model(
          np_trajectory.observations, n_accelerators=1
      ) * self._value_network_scale
      values = np.squeeze(values, axis=2)  # Remove the singleton depth dim.

      # TODO(pkozakowski): Add some shape assertions and docs.
      # Calculate targets based on the advantages over the target network - this
      # allows TD learning for value networks.
      advantages = self._advantage_estimator(
          np_trajectory.rewards, np_trajectory.returns, values,
          gamma=self._task.gamma,
          n_extra_steps=self._added_policy_slice_length,
      )
      length = advantages.shape[1]
      values = values[:, :length]
      target_returns = values + advantages

      # Insert an extra depth dimension, so the target shape is consistent with
      # the network output shape.
      yield (
          # Inputs: observations.
          np_trajectory.observations[:, :length],
          # Targets: computed returns.
          target_returns[:, :, None] / self._value_network_scale,
          # Mask to zero-out padding.
          np_trajectory.mask[:, :length, None],
      )

  def policy_inputs(self, trajectory, values):
    """Create inputs to policy model from a TrajectoryNp and values.

    Args:
      trajectory: a TrajectoryNp, the trajectory to create inputs from
      values: a numpy array: value function computed on trajectory

    Returns:
      a tuple of numpy arrays of the form (inputs, x1, x2, ...) that will be
      passed to the policy model; policy model will compute outputs from
      inputs and (outputs, x1, x2, ...) will be passed to self.policy_loss
      which should be overridden accordingly.
    """
    return NotImplementedError

  def policy_batches_stream(self):
    """Use the RLTask self._task to create inputs to the policy model."""
    # Maximum slice length for policy is max_slice_len + the added policy len.
    max_slice_length = self._max_slice_length + self._added_policy_slice_length
    for np_trajectory in self._task.trajectory_batch_stream(
        self._policy_batch_size,
        epochs=self._replay_epochs,
        max_slice_length=max_slice_length,
        include_final_state=False):
      values = self._value_eval_model(
          np_trajectory.observations, n_accelerators=1
      ) * self._value_network_scale
      values = np.squeeze(values, axis=2)  # Remove the singleton depth dim.
      if len(values.shape) != 2:
        raise ValueError('Values are expected to have shape ' +
                         '[batch_size, length], got: %s' % str(values.shape))
      if values.shape[0] != self._policy_batch_size:
        raise ValueError('Values first dimension should = policy batch size, ' +
                         '%d != %d' %(values.shape[0], self._policy_batch_size))
      yield self.policy_inputs(np_trajectory, values)

  def train_epoch(self):
    """Trains RL for one epoch."""
    # Copy policy state accumulated during data collection to the trainer.
    self._policy_trainer.model_state = self._policy_collect_model.state

    # Copy policy weights and state to value trainer.
    if self._n_shared_layers > 0:
      _copy_model_weights_and_state(
          0, self._n_shared_layers, self._policy_trainer, self._value_trainer
      )

    # Update the target value network.
    self._value_eval_model.weights = self._value_trainer.model_weights
    self._value_eval_model.state = self._value_trainer.model_state

    n_value_evals = rl_training.remaining_evals(
        self._value_trainer.step,
        self._epoch,
        self._value_train_steps_per_epoch,
        self._value_evals_per_epoch)
    for _ in range(n_value_evals):
      self._value_trainer.train_epoch(
          self._value_train_steps_per_epoch // self._value_evals_per_epoch,
          self._value_eval_steps,
      )
    # Copy value weights and state to policy trainer.
    if self._n_shared_layers > 0:
      _copy_model_weights_and_state(
          0, self._n_shared_layers, self._value_trainer, self._policy_trainer
      )
    n_policy_evals = rl_training.remaining_evals(
        self._policy_trainer.step,
        self._epoch,
        self._policy_train_steps_per_epoch,
        self._policy_evals_per_epoch)
    # Check if there was a restart after value training finishes and policy not.
    stopped_after_value = (n_value_evals == 0 and
                           n_policy_evals < self._policy_evals_per_epoch)
    should_copy_weights = self._n_shared_layers > 0 and not stopped_after_value
    if should_copy_weights:
      _copy_model_weights_and_state(
          0, self._n_shared_layers, self._value_trainer, self._policy_trainer
      )

    # Update the target value network.
    self._value_eval_model.weights = self._value_trainer.model_weights
    self._value_eval_model.state = self._value_trainer.model_state

    for _ in range(n_policy_evals):
      self._policy_trainer.train_epoch(
          self._policy_train_steps_per_epoch // self._policy_evals_per_epoch,
          self._policy_eval_steps,
      )

  def close(self):
    self._value_trainer.close()
    super().close()


def _copy_model_weights_and_state(  # pylint: disable=invalid-name
    start, end, from_trainer, to_trainer, copy_optimizer_slots=False
):
  """Copy model weights[start:end] from from_trainer to to_trainer."""
  from_weights = from_trainer.model_weights
  to_weights = to_trainer.model_weights
  shared_weights = from_weights[start:end]
  to_weights[start:end] = shared_weights
  to_trainer.model_weights = to_weights

  from_state = from_trainer.model_state
  to_state = to_trainer.model_state
  shared_state = from_state[start:end]
  to_state[start:end] = shared_state
  to_trainer.model_state = to_state

  if copy_optimizer_slots:
    # TODO(lukaszkaiser): make a nicer API in Trainer to support this.
    # Currently we use the hack below. Note [0] since that's the model w/o loss.
    # pylint: disable=protected-access
    from_slots = from_trainer._opt_state.slots[0][start:end]
    to_slots = to_trainer._opt_state.slots[0]
    # The lines below do to_slots[start:end] = from_slots, but on tuples.
    new_slots = to_slots[:start] + from_slots[start:end] + to_slots[end:]
    new_slots = tuple([new_slots] + list(to_trainer._opt_state.slots[1:]))
    to_trainer._opt_state = to_trainer._opt_state._replace(slots=new_slots)
    # pylint: enable=protected-access


### Implementations of common actor-critic algorithms.


class AdvantageBasedActorCriticTrainer(ActorCriticTrainer):
  """Base class for advantage-based actor-critic algorithms."""

  def __init__(
      self,
      task,
      advantage_estimator=rl_advantages.td_lambda,
      advantage_normalization=True,
      advantage_normalization_epsilon=1e-5,
      **kwargs
  ):
    self._advantage_estimator = advantage_estimator
    self._advantage_normalization = advantage_normalization
    self._advantage_normalization_epsilon = advantage_normalization_epsilon
    super(AdvantageBasedActorCriticTrainer, self).__init__(task, **kwargs)

  def policy_inputs(self, trajectory, values):
    """Create inputs to policy model from a TrajectoryNp and values."""
    # How much TD to use is determined by the added policy slice length,
    # as the policy batches need to be this much longer to calculate TD.
    advantages = self._advantage_estimator(
        trajectory.rewards, trajectory.returns, values,
        gamma=self._task.gamma,
        n_extra_steps=self._added_policy_slice_length,
    )
    # Observations should be the same length as advantages - so if we are
    # using n_extra_steps, we need to trim the length to match.
    obs = trajectory.observations[:, :advantages.shape[1]]
    act = trajectory.actions[:, :advantages.shape[1]]
    old_logps = trajectory.log_probs[:, :advantages.shape[1]]
    mask = trajectory.mask[:, :advantages.shape[1]]  # Mask to zero-out padding.
    # Shape checks to help debugging.
    if len(advantages.shape) != 2:
      raise ValueError('Advantages are expected to have shape ' +
                       '[batch_size, length], got: %s' % str(advantages.shape))
    if act.shape[0:2] != advantages.shape:
      raise ValueError('First 2 dimensions of actions should be the same as in '
                       'advantages, %s != %s' % (act.shape[0:2],
                                                 advantages.shape))
    if obs.shape[0:2] != advantages.shape:
      raise ValueError('First 2 dimensions of observations should be the same '
                       'as in advantages, %s != %s' % (obs.shape[0:2],
                                                       advantages.shape))
    if old_logps.shape != advantages.shape:
      raise ValueError('Old log-probs and advantages shapes should be the same'
                       ', %s != %s' % (old_logps.shape, advantages.shape))
    if mask.shape != advantages.shape:
      raise ValueError('Mask and advantages shapes should be the same'
                       ', %s != %s' % (mask.shape, advantages.shape))
    return (obs, act, advantages, old_logps, mask)

  @property
  def policy_loss_given_log_probs(self):
    """Policy loss given action log-probabilities."""
    raise NotImplementedError

  @property
  def policy_loss(self, **unused_kwargs):
    """Policy loss."""
    def NormalizeAdvantages():  # pylint: disable=invalid-name
      def normalize(adv):
        return (
            (adv - jnp.mean(adv)) /
            (jnp.std(adv) + self._advantage_normalization_epsilon)
        )
      return tl.Fn('NormalizeAdvantages', normalize)

    return tl.Serial(
        self._policy_dist.LogProb(),
        tl.Parallel(
            [],  # Distribution inputs.
            # Advantages.
            NormalizeAdvantages() if self._advantage_normalization else [],
            [],  # Old log probs.
            [],  # Mask.
        ),
        self.policy_loss_given_log_probs,
    )

  @property
  def policy_metrics(self):
    metrics = super(AdvantageBasedActorCriticTrainer, self).policy_metrics
    metrics.update({
        'advantage_mean': self.advantage_mean,
        'advantage_std': self.advantage_std,
    })
    return metrics

  @property
  def advantage_mean(self):
    return tl.Serial([
        # (dist_inputs, advantages, old_log_probs, mask)
        tl.Select([1]),  # Select just the advantages.
        tl.Fn('AdvantageMean', lambda x: jnp.mean(x)),  # pylint: disable=unnecessary-lambda
    ])

  @property
  def advantage_std(self):
    return tl.Serial([
        # (dist_inputs, advantages, old_log_probs, mask)
        tl.Select([1]),  # Select just the advantages.
        tl.Fn('AdvantageStd', lambda x: jnp.std(x)),  # pylint: disable=unnecessary-lambda
    ])


# A2C is one of the most basic actor-critic RL algorithms.
def A2CLoss():  # pylint: disable=invalid-name
  """Definition of the Advantage Actor Critic (A2C) loss."""
  def f(log_probs, advantages, old_log_probs, mask):
    del old_log_probs  # Not used in A2C.
    return -np.sum(log_probs * advantages * mask) / np.sum(mask)
  return tl.Fn('A2CLoss', f)


class A2CTrainer(AdvantageBasedActorCriticTrainer):
  """Trains policy and value models using the A2C algortithm."""

  on_policy = True

  @property
  def policy_loss_given_log_probs(self):
    """Policy loss."""
    return A2CLoss()  # pylint: disable=no-value-for-parameter


# PPO is a widely used actor-critic RL algorithm.
def PPOLoss(epsilon):  # pylint: disable=invalid-name
  """Definition of the Proximal Policy Optimization loss."""
  def f(new_log_probs, advantages, old_log_probs, mask):
    # Old log probs have an undesirable extra dimension which we remove here
    old_log_probs = old_log_probs.squeeze(axis=-1)

    # The ratio between new_probs and old_probs expressed
    # using log_probs and exponentaion
    probs_ratio = jnp.exp(new_log_probs - old_log_probs)
    unclipped_objective = probs_ratio * advantages
    clipped_objective = jnp.clip(probs_ratio,
                                 1 - epsilon,
                                 1 + epsilon) * advantages
    ppo_objective = jnp.minimum(unclipped_objective, clipped_objective)
    return -np.sum(ppo_objective * mask) / np.sum(mask)
  return tl.Fn('PPOLoss', f)


class PPOTrainer(AdvantageBasedActorCriticTrainer):
  """The Proximal Policy Optimization Algorithm aka PPO.

  Trains policy and value models using the PPO algortithm.
  """

  on_policy = True

  def __init__(self, task, epsilon=0.2, **kwargs):
    """Configures the PPO Trainer."""
    self._epsilon = epsilon
    super(PPOTrainer, self).__init__(task, **kwargs)

  @property
  def policy_loss_given_log_probs(self):
    """Policy loss."""
    return PPOLoss(epsilon=self._epsilon)  # pylint: disable=no-value-for-parameter


# AWR is an off-policy actor-critic RL algorithm.
def awr_weights(advantages, beta):
  return jnp.exp(advantages / beta)


def AWRLoss(beta, w_max):  # pylint: disable=invalid-name
  """Definition of the Advantage Weighted Regression (AWR) loss."""
  def f(log_probs, advantages, old_log_probs, mask):
    del old_log_probs  # Not used in AWR.
    weights = jnp.minimum(awr_weights(advantages, beta), w_max)
    return -np.sum(log_probs * weights * mask) / np.sum(mask)
  return tl.Fn('AWRLoss', f)


class AWRTrainer(AdvantageBasedActorCriticTrainer):
  """Trains policy and value models using AWR."""

  on_policy = False

  def __init__(self, task, beta=1.0, w_max=20.0, **kwargs):
    """Configures the AWR Trainer."""
    self._beta = beta
    self._w_max = w_max
    super(AWRTrainer, self).__init__(task, **kwargs)

  @property
  def policy_loss_given_log_probs(self):
    """Policy loss."""
    return AWRLoss(beta=self._beta, w_max=self._w_max)  # pylint: disable=no-value-for-parameter

  @property
  def policy_metrics(self):
    metrics = super(AWRTrainer, self).policy_metrics
    metrics.update({  # pylint: disable=g-complex-comprehension
        'awr_weight_' + name: self.awr_weight_stat(name, fn)
        for (name, fn) in [
            ('mean', jnp.mean),
            ('std', jnp.std),
            ('min', jnp.min),
            ('max', jnp.max),
        ]
    })
    return metrics

  def awr_weight_stat(self, stat_name, stat_fn):
    return tl.Serial([
        tl.Select([1]),  # Select just the advantages.
        tl.Fn(
            'AWRWeight' + stat_name.capitalize(),
            lambda x: stat_fn(awr_weights(x, self._beta)),
        ),
    ])

1	# coding=utf-8
2	# Copyright 2020 The Trax Authors.
3	#
4	# Licensed under the Apache License, Version 2.0 (the "License");
5	# you may not use this file except in compliance with the License.
6	# You may obtain a copy of the License at
7	#
8	# http://www.apache.org/licenses/LICENSE-2.0
9	#
10	# Unless required by applicable law or agreed to in writing, software
11	# distributed under the License is distributed on an "AS IS" BASIS,
12	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13	# See the License for the specific language governing permissions and
14	# limitations under the License.
15
16	# Lint as: python3
17	"""Classes for RL training in Trax."""	1✔
18
19	import os	1✔
20	import numpy as np	1✔
21	import tensorflow as tf	1✔
22
23	from trax import layers as tl	1✔
24	from trax import lr_schedules as lr	1✔
25	from trax import shapes	1✔
26	from trax import supervised	1✔
27	from trax.math import numpy as jnp	1✔
28	from trax.rl import advantages as rl_advantages	1✔
29	from trax.rl import training as rl_training	1✔
30
31
32	class ActorCriticTrainer(rl_training.PolicyTrainer):	1✔
33	"""Trains policy and value models using actor-critic methods.
34
35	Attrs:
36	on_policy (bool): Whether the algorithm is on-policy. Used in the data
37	generators. Should be set in derived classes.
38	"""
39
40	on_policy = None	1✔
41
42	def __init__(self, task,	1✔
43	value_model=None,
44	value_optimizer=None,
45	value_lr_schedule=lr.MultifactorSchedule,
46	value_batch_size=64,
47	value_train_steps_per_epoch=500,
48	value_evals_per_epoch=1,
49	value_eval_steps=1,
50	n_shared_layers=0,
51	added_policy_slice_length=0,
52	n_replay_epochs=1,
53	scale_value_targets=False,
54	**kwargs): # Arguments of PolicyTrainer come here.
55	"""Configures the actor-critic Trainer.
56
57	Args:
58	task: RLTask instance to use
59	value_model: the model to use for the value function
60	value_optimizer: the optimizer to train the value model
61	value_lr_schedule: lr schedule for value model training
62	value_batch_size: batch size for value model training
63	value_train_steps_per_epoch: how many steps are we using to
64	train the value model in each epoch
65	value_evals_per_epoch: number of value trainer evaluations per RL
66	epoch - only affects metric reporting.
67	value_eval_steps: number of value trainer steps per evaluation -
68	only affects metric reporting.
69	n_shared_layers: how many layers to share between value and
70	policy models
71	added_policy_slice_length: how much longer should slices of
72	trajectories be for policy than for value training; this
73	is useful for TD calculations and only affect the length
74	of elements produced for policy batches; value batches
75	have maximum length set by max_slice_length in **kwargs
76	n_replay_epochs: how many last epochs to take into the replay buffer;
77	only makes sense for off-policy algorithms
78	scale_value_targets: whether to scale targets for the value function by
79	1 / (1 - gamma)
80	**kwargs: arguments for PolicyTrainer super-class
81	"""
82	self._n_shared_layers = n_shared_layers	×
83	self._value_batch_size = value_batch_size	×
84	self._value_train_steps_per_epoch = value_train_steps_per_epoch	×
85	self._value_evals_per_epoch = value_evals_per_epoch	×
86	self._value_eval_steps = value_eval_steps	×
87
88	# The 2 below will be initalized in super.__init__ anyway, but are needed
89	# to construct value batches which are needed before PolicyTrainer init
90	# since policy input creation calls the value model -- hence this code.
91	self._task = task	×
92	self._max_slice_length = kwargs.get('max_slice_length', 1)	×
93	self._added_policy_slice_length = added_policy_slice_length	×
94	self._n_replay_epochs = n_replay_epochs	×
95
96	if scale_value_targets:	×
97	self._value_network_scale = 1 / (1 - self._task.gamma)	×
98	else:
99	self._value_network_scale = 1	×
100
101	self._value_eval_model = value_model(mode='eval')	×
102	self._value_eval_model.init(self._value_model_signature)	×
103
104	# Initialize training of the value function.
105	value_output_dir = kwargs.get('output_dir', None)	×
106	if value_output_dir is not None:	×
107	value_output_dir = os.path.join(value_output_dir, 'value')	×
108	# If needed, create value_output_dir and missing parent directories.
109	if not tf.io.gfile.isdir(value_output_dir):	×
110	tf.io.gfile.makedirs(value_output_dir)	×
111	self._value_inputs = supervised.Inputs(	×
112	train_stream=lambda _: self.value_batches_stream())
113	self._value_trainer = supervised.Trainer(	×
114	model=value_model,
115	optimizer=value_optimizer,
116	lr_schedule=value_lr_schedule,
117	loss_fn=tl.L2Loss(),
118	inputs=self._value_inputs,
119	output_dir=value_output_dir,
120	metrics={'value_loss': tl.L2Loss()})
121
122	# Initialize policy training.
123	super(ActorCriticTrainer, self).__init__(task, **kwargs)	×
124
125	@property	1✔
126	def _value_model_signature(self):
127	obs_sig = shapes.signature(self._task.observation_space)	×
128	target_sig = mask_sig = shapes.ShapeDtype(	×
129	shape=(1, 1, 1),
130	)
131	return (obs_sig.replace(shape=(1, 1) + obs_sig.shape), target_sig, mask_sig)	×
132
133	@property	1✔
134	def _replay_epochs(self):
135	if self.on_policy:	×
136	assert self._n_replay_epochs == 1, (	×
137	'Non-unit replay buffer size only makes sense for off-policy '
138	'algorithms.'
139	)
140	return [-(ep + 1) for ep in range(self._n_replay_epochs)]	×
141
142	def value_batches_stream(self):	1✔
143	"""Use the RLTask self._task to create inputs to the value model."""
144	max_slice_length = self._max_slice_length + self._added_policy_slice_length	×
145	for np_trajectory in self._task.trajectory_batch_stream(	×
146	self._value_batch_size,
147	max_slice_length=max_slice_length,
148	min_slice_length=(1 + self._added_policy_slice_length),
149	epochs=self._replay_epochs,
150	):
151	values = self._value_eval_model(	×
152	np_trajectory.observations, n_accelerators=1
153	) * self._value_network_scale
154	values = np.squeeze(values, axis=2) # Remove the singleton depth dim.	×
155
156	# TODO(pkozakowski): Add some shape assertions and docs.
157	# Calculate targets based on the advantages over the target network - this
158	# allows TD learning for value networks.
159	advantages = self._advantage_estimator(	×
160	np_trajectory.rewards, np_trajectory.returns, values,
161	gamma=self._task.gamma,
162	n_extra_steps=self._added_policy_slice_length,
163	)
164	length = advantages.shape[1]	×
165	values = values[:, :length]	×
166	target_returns = values + advantages	×
167
168	# Insert an extra depth dimension, so the target shape is consistent with
169	# the network output shape.
170	yield (	×
171	# Inputs: observations.
172	np_trajectory.observations[:, :length],
173	# Targets: computed returns.
174	target_returns[:, :, None] / self._value_network_scale,
175	# Mask to zero-out padding.
176	np_trajectory.mask[:, :length, None],
177	)
178
179	def policy_inputs(self, trajectory, values):	1✔
180	"""Create inputs to policy model from a TrajectoryNp and values.
181
182	Args:
183	trajectory: a TrajectoryNp, the trajectory to create inputs from
184	values: a numpy array: value function computed on trajectory
185
186	Returns:
187	a tuple of numpy arrays of the form (inputs, x1, x2, ...) that will be
188	passed to the policy model; policy model will compute outputs from
189	inputs and (outputs, x1, x2, ...) will be passed to self.policy_loss
190	which should be overridden accordingly.
191	"""
192	return NotImplementedError	×
193
194	def policy_batches_stream(self):	1✔
195	"""Use the RLTask self._task to create inputs to the policy model."""
196	# Maximum slice length for policy is max_slice_len + the added policy len.
197	max_slice_length = self._max_slice_length + self._added_policy_slice_length	×
198	for np_trajectory in self._task.trajectory_batch_stream(	×
199	self._policy_batch_size,
200	epochs=self._replay_epochs,
201	max_slice_length=max_slice_length,
202	include_final_state=False):
203	values = self._value_eval_model(	×
204	np_trajectory.observations, n_accelerators=1
205	) * self._value_network_scale
206	values = np.squeeze(values, axis=2) # Remove the singleton depth dim.	×
207	if len(values.shape) != 2:	×
208	raise ValueError('Values are expected to have shape ' +	×
209	'[batch_size, length], got: %s' % str(values.shape))
210	if values.shape[0] != self._policy_batch_size:	×
211	raise ValueError('Values first dimension should = policy batch size, ' +	×
212	'%d != %d' %(values.shape[0], self._policy_batch_size))
213	yield self.policy_inputs(np_trajectory, values)	×
214
215	def train_epoch(self):	1✔
216	"""Trains RL for one epoch."""
217	# Copy policy state accumulated during data collection to the trainer.
218	self._policy_trainer.model_state = self._policy_collect_model.state	×
219
220	# Copy policy weights and state to value trainer.
221	if self._n_shared_layers > 0:	×
222	_copy_model_weights_and_state(	×
223	0, self._n_shared_layers, self._policy_trainer, self._value_trainer
224	)
225
226	# Update the target value network.
227	self._value_eval_model.weights = self._value_trainer.model_weights	×
228	self._value_eval_model.state = self._value_trainer.model_state	×
229
230	n_value_evals = rl_training.remaining_evals(	×
231	self._value_trainer.step,
232	self._epoch,
233	self._value_train_steps_per_epoch,
234	self._value_evals_per_epoch)
235	for _ in range(n_value_evals):	×
236	self._value_trainer.train_epoch(	×
237	self._value_train_steps_per_epoch // self._value_evals_per_epoch,
238	self._value_eval_steps,
239	)
240	# Copy value weights and state to policy trainer.
241	if self._n_shared_layers > 0:	×
242	_copy_model_weights_and_state(	×
243	0, self._n_shared_layers, self._value_trainer, self._policy_trainer
244	)
245	n_policy_evals = rl_training.remaining_evals(	×
246	self._policy_trainer.step,
247	self._epoch,
248	self._policy_train_steps_per_epoch,
249	self._policy_evals_per_epoch)
250	# Check if there was a restart after value training finishes and policy not.
251	stopped_after_value = (n_value_evals == 0 and	×
252	n_policy_evals < self._policy_evals_per_epoch)
253	should_copy_weights = self._n_shared_layers > 0 and not stopped_after_value	×
254	if should_copy_weights:	×
255	_copy_model_weights_and_state(	×
256	0, self._n_shared_layers, self._value_trainer, self._policy_trainer
257	)
258
259	# Update the target value network.
260	self._value_eval_model.weights = self._value_trainer.model_weights	×
261	self._value_eval_model.state = self._value_trainer.model_state	×
262
263	for _ in range(n_policy_evals):	×
264	self._policy_trainer.train_epoch(	×
265	self._policy_train_steps_per_epoch // self._policy_evals_per_epoch,
266	self._policy_eval_steps,
267	)
268
269	def close(self):	1✔
270	self._value_trainer.close()	×
271	super().close()	×
272
273
274	def _copy_model_weights_and_state( # pylint: disable=invalid-name	1✔
275	start, end, from_trainer, to_trainer, copy_optimizer_slots=False
276	):
277	"""Copy model weights[start:end] from from_trainer to to_trainer."""
278	from_weights = from_trainer.model_weights	×
279	to_weights = to_trainer.model_weights	×
280	shared_weights = from_weights[start:end]	×
281	to_weights[start:end] = shared_weights	×
282	to_trainer.model_weights = to_weights	×
283
284	from_state = from_trainer.model_state	×
285	to_state = to_trainer.model_state	×
286	shared_state = from_state[start:end]	×
287	to_state[start:end] = shared_state	×
288	to_trainer.model_state = to_state	×
289
290	if copy_optimizer_slots:	×
291	# TODO(lukaszkaiser): make a nicer API in Trainer to support this.
292	# Currently we use the hack below. Note [0] since that's the model w/o loss.
293	# pylint: disable=protected-access
294	from_slots = from_trainer._opt_state.slots[0][start:end]	×
295	to_slots = to_trainer._opt_state.slots[0]	×
296	# The lines below do to_slots[start:end] = from_slots, but on tuples.
297	new_slots = to_slots[:start] + from_slots[start:end] + to_slots[end:]	×
298	new_slots = tuple([new_slots] + list(to_trainer._opt_state.slots[1:]))	×
299	to_trainer._opt_state = to_trainer._opt_state._replace(slots=new_slots)	×
300	# pylint: enable=protected-access
301
302
303	### Implementations of common actor-critic algorithms.
304
305
306	class AdvantageBasedActorCriticTrainer(ActorCriticTrainer):	1✔
307	"""Base class for advantage-based actor-critic algorithms."""
308
309	def __init__(	1✔
310	self,
311	task,
312	advantage_estimator=rl_advantages.td_lambda,
313	advantage_normalization=True,
314	advantage_normalization_epsilon=1e-5,
315	**kwargs
316	):
317	self._advantage_estimator = advantage_estimator	×
318	self._advantage_normalization = advantage_normalization	×
319	self._advantage_normalization_epsilon = advantage_normalization_epsilon	×
320	super(AdvantageBasedActorCriticTrainer, self).__init__(task, **kwargs)	×
321
322	def policy_inputs(self, trajectory, values):	1✔
323	"""Create inputs to policy model from a TrajectoryNp and values."""
324	# How much TD to use is determined by the added policy slice length,
325	# as the policy batches need to be this much longer to calculate TD.
326	advantages = self._advantage_estimator(	×
327	trajectory.rewards, trajectory.returns, values,
328	gamma=self._task.gamma,
329	n_extra_steps=self._added_policy_slice_length,
330	)
331	# Observations should be the same length as advantages - so if we are
332	# using n_extra_steps, we need to trim the length to match.
333	obs = trajectory.observations[:, :advantages.shape[1]]	×
334	act = trajectory.actions[:, :advantages.shape[1]]	×
335	old_logps = trajectory.log_probs[:, :advantages.shape[1]]	×
336	mask = trajectory.mask[:, :advantages.shape[1]] # Mask to zero-out padding.	×
337	# Shape checks to help debugging.
338	if len(advantages.shape) != 2:	×
339	raise ValueError('Advantages are expected to have shape ' +	×
340	'[batch_size, length], got: %s' % str(advantages.shape))
341	if act.shape[0:2] != advantages.shape:	×
342	raise ValueError('First 2 dimensions of actions should be the same as in '	×
343	'advantages, %s != %s' % (act.shape[0:2],
344	advantages.shape))
345	if obs.shape[0:2] != advantages.shape:	×
346	raise ValueError('First 2 dimensions of observations should be the same '	×
347	'as in advantages, %s != %s' % (obs.shape[0:2],
348	advantages.shape))
349	if old_logps.shape != advantages.shape:	×
350	raise ValueError('Old log-probs and advantages shapes should be the same'	×
351	', %s != %s' % (old_logps.shape, advantages.shape))
352	if mask.shape != advantages.shape:	×
353	raise ValueError('Mask and advantages shapes should be the same'	×
354	', %s != %s' % (mask.shape, advantages.shape))
355	return (obs, act, advantages, old_logps, mask)	×
356
357	@property	1✔
358	def policy_loss_given_log_probs(self):
359	"""Policy loss given action log-probabilities."""
360	raise NotImplementedError
361
362	@property	1✔
363	def policy_loss(self, **unused_kwargs):
364	"""Policy loss."""
365	def NormalizeAdvantages(): # pylint: disable=invalid-name	×
366	def normalize(adv):	×
367	return (	×
368	(adv - jnp.mean(adv)) /
369	(jnp.std(adv) + self._advantage_normalization_epsilon)
370	)
371	return tl.Fn('NormalizeAdvantages', normalize)	×
372
373	return tl.Serial(	×
374	self._policy_dist.LogProb(),
375	tl.Parallel(
376	[], # Distribution inputs.
377	# Advantages.
378	NormalizeAdvantages() if self._advantage_normalization else [],
379	[], # Old log probs.
380	[], # Mask.
381	),
382	self.policy_loss_given_log_probs,
383	)
384
385	@property	1✔
386	def policy_metrics(self):
387	metrics = super(AdvantageBasedActorCriticTrainer, self).policy_metrics	×
388	metrics.update({	×
389	'advantage_mean': self.advantage_mean,
390	'advantage_std': self.advantage_std,
391	})
392	return metrics	×
393
394	@property	1✔
395	def advantage_mean(self):
396	return tl.Serial([	×
397	# (dist_inputs, advantages, old_log_probs, mask)
398	tl.Select([1]), # Select just the advantages.
399	tl.Fn('AdvantageMean', lambda x: jnp.mean(x)), # pylint: disable=unnecessary-lambda
400	])
401
402	@property	1✔
403	def advantage_std(self):
404	return tl.Serial([	×
405	# (dist_inputs, advantages, old_log_probs, mask)
406	tl.Select([1]), # Select just the advantages.
407	tl.Fn('AdvantageStd', lambda x: jnp.std(x)), # pylint: disable=unnecessary-lambda
408	])
409
410
411	# A2C is one of the most basic actor-critic RL algorithms.
412	def A2CLoss(): # pylint: disable=invalid-name	1✔
413	"""Definition of the Advantage Actor Critic (A2C) loss."""
414	def f(log_probs, advantages, old_log_probs, mask):	×
415	del old_log_probs # Not used in A2C.	×
416	return -np.sum(log_probs * advantages * mask) / np.sum(mask)	×
417	return tl.Fn('A2CLoss', f)	×
418
419
420	class A2CTrainer(AdvantageBasedActorCriticTrainer):	1✔
421	"""Trains policy and value models using the A2C algortithm."""
422
423	on_policy = True	1✔
424
425	@property	1✔
426	def policy_loss_given_log_probs(self):
427	"""Policy loss."""
428	return A2CLoss() # pylint: disable=no-value-for-parameter	×
429
430
431	# PPO is a widely used actor-critic RL algorithm.
432	def PPOLoss(epsilon): # pylint: disable=invalid-name	1✔
433	"""Definition of the Proximal Policy Optimization loss."""
434	def f(new_log_probs, advantages, old_log_probs, mask):	×
435	# Old log probs have an undesirable extra dimension which we remove here
436	old_log_probs = old_log_probs.squeeze(axis=-1)	×
437
438	# The ratio between new_probs and old_probs expressed
439	# using log_probs and exponentaion
440	probs_ratio = jnp.exp(new_log_probs - old_log_probs)	×
441	unclipped_objective = probs_ratio * advantages	×
442	clipped_objective = jnp.clip(probs_ratio,	×
443	1 - epsilon,
444	1 + epsilon) * advantages
445	ppo_objective = jnp.minimum(unclipped_objective, clipped_objective)	×
446	return -np.sum(ppo_objective * mask) / np.sum(mask)	×
447	return tl.Fn('PPOLoss', f)	×
448
449
450	class PPOTrainer(AdvantageBasedActorCriticTrainer):	1✔
451	"""The Proximal Policy Optimization Algorithm aka PPO.
452
453	Trains policy and value models using the PPO algortithm.
454	"""
455
456	on_policy = True	1✔
457
458	def __init__(self, task, epsilon=0.2, **kwargs):	1✔
459	"""Configures the PPO Trainer."""
460	self._epsilon = epsilon	×
461	super(PPOTrainer, self).__init__(task, **kwargs)	×
462
463	@property	1✔
464	def policy_loss_given_log_probs(self):
465	"""Policy loss."""
466	return PPOLoss(epsilon=self._epsilon) # pylint: disable=no-value-for-parameter	×
467
468
469	# AWR is an off-policy actor-critic RL algorithm.
470	def awr_weights(advantages, beta):	1✔
471	return jnp.exp(advantages / beta)	×
472
473
474	def AWRLoss(beta, w_max): # pylint: disable=invalid-name	1✔
475	"""Definition of the Advantage Weighted Regression (AWR) loss."""
476	def f(log_probs, advantages, old_log_probs, mask):	×
477	del old_log_probs # Not used in AWR.	×
478	weights = jnp.minimum(awr_weights(advantages, beta), w_max)	×
479	return -np.sum(log_probs * weights * mask) / np.sum(mask)	×
480	return tl.Fn('AWRLoss', f)	×
481
482
483	class AWRTrainer(AdvantageBasedActorCriticTrainer):	1✔
484	"""Trains policy and value models using AWR."""
485
486	on_policy = False	1✔
487
488	def __init__(self, task, beta=1.0, w_max=20.0, **kwargs):	1✔
489	"""Configures the AWR Trainer."""
490	self._beta = beta	×
491	self._w_max = w_max	×
492	super(AWRTrainer, self).__init__(task, **kwargs)	×
493
494	@property	1✔
495	def policy_loss_given_log_probs(self):
496	"""Policy loss."""
497	return AWRLoss(beta=self._beta, w_max=self._w_max) # pylint: disable=no-value-for-parameter	×
498
499	@property	1✔
500	def policy_metrics(self):
501	metrics = super(AWRTrainer, self).policy_metrics	×
502	metrics.update({ # pylint: disable=g-complex-comprehension	×
503	'awr_weight_' + name: self.awr_weight_stat(name, fn)
504	for (name, fn) in [
505	('mean', jnp.mean),
506	('std', jnp.std),
507	('min', jnp.min),
508	('max', jnp.max),
509	]
510	})
511	return metrics	×
512
513	def awr_weight_stat(self, stat_name, stat_fn):	1✔
514	return tl.Serial([	×
515	tl.Select([1]), # Select just the advantages.
516	tl.Fn(
517	'AWRWeight' + stat_name.capitalize(),
518	lambda x: stat_fn(awr_weights(x, self._beta)),
519	),
520	])

google / trax / 83

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous