(self)
Use the RLTask self._task to create inputs to the policy model.
Use the RLTask self._task to create inputs to the policy model.
Use the RLTask self . _task to create inputs to the policy model .
def policy_batches_stream(self):
    """Use the RLTask self._task to create inputs to the policy model."""
    # For now TD-0 estimation of the value. TODO(pkozakowski): Support others?
    for np_trajectory in self._task.trajectory_batch_stream(
        self._policy_batch_size,
        epochs=self._replay_epochs,
        max_slice_length=self._max_slice_length,
    ):
      dist_inputs = self._get_dist_inputs(np_trajectory)
      (q_values, actions, act_log_probs) = self._run_value_model(
          np_trajectory.observation, dist_inputs)
      shapes.assert_same_shape(q_values, act_log_probs)

      # q_values shape: (batch_size, n_samples, length)
      if len(q_values.shape) != 3:
        raise ValueError('Q-values are expected to have shape [batch_size, ' +
                         'n_samples, length], got: %s' % str(q_values.shape))
      if q_values.shape[1] != self._q_value_n_samples:
        raise ValueError('Q-values dimension 1 should = n_samples, %d != %d'
                         % (q_values.shape[1], self._q_value_n_samples))
      if q_values.shape[0] != self._policy_batch_size:
        raise ValueError('Q-values dimension 0 should = policy batch size, ' +
                         '%d!=%d' %(q_values.shape[1], self._policy_batch_size))

      mask = np_trajectory.mask
      mask = np.reshape(mask, [mask.shape[0], 1] + list(mask.shape[1:]))
      mask = jnp.broadcast_to(mask, q_values.shape)
      shapes.assert_same_shape(mask, q_values)
      yield (np_trajectory.observation, actions, q_values, act_log_probs, mask)
def policy_batches_stream ( self ) : # For now TD-0 estimation of the value. TODO(pkozakowski): Support others? for np_trajectory in self . _task . trajectory_batch_stream ( self . _policy_batch_size , epochs = self . _replay_epochs , max_slice_length = self . _max_slice_length , ) : dist_inputs = self . _get_dist_inputs ( np_trajectory ) ( q_values , actions , act_log_probs ) = self . _run_value_model ( np_trajectory . observation , dist_inputs ) shapes . assert_same_shape ( q_values , act_log_probs ) # q_values shape: (batch_size, n_samples, length) if len ( q_values . shape ) != 3 : raise ValueError ( 'Q-values are expected to have shape [batch_size, ' + 'n_samples, length], got: %s' % str ( q_values . shape ) ) if q_values . shape [ 1 ] != self . _q_value_n_samples : raise ValueError ( 'Q-values dimension 1 should = n_samples, %d != %d' % ( q_values . shape [ 1 ] , self . _q_value_n_samples ) ) if q_values . shape [ 0 ] != self . _policy_batch_size : raise ValueError ( 'Q-values dimension 0 should = policy batch size, ' + '%d!=%d' % ( q_values . shape [ 1 ] , self . _policy_batch_size ) ) mask = np_trajectory . mask mask = np . reshape ( mask , [ mask . shape [ 0 ] , 1 ] + list ( mask . shape [ 1 : ] ) ) mask = jnp . broadcast_to ( mask , q_values . shape ) shapes . assert_same_shape ( mask , q_values ) yield ( np_trajectory . observation , actions , q_values , act_log_probs , mask )