(self)
Use the RLTask self._task to create inputs to the value model.
Use the RLTask self._task to create inputs to the value model.
Use the RLTask self . _task to create inputs to the value model .
def value_batches_stream(self):
    """Use the RLTask self._task to create inputs to the value model."""
    max_slice_length = self._max_slice_length
    min_slice_length = 1
    for np_trajectory in self._task.trajectory_batch_stream(
        self._value_batch_size,
        max_slice_length=max_slice_length,
        min_slice_length=min_slice_length,
        margin=0,
        epochs=self._replay_epochs,
    ):
      values_target = self._run_value_model(
          np_trajectory.observation, use_eval_model=True)
      if self._double_dqn:
        values = self._run_value_model(
            np_trajectory.observation, use_eval_model=False
        )
        index_max = np.argmax(values, axis=-1)
        ind_0, ind_1 = np.indices(index_max.shape)
        values_max = values_target[ind_0, ind_1, index_max]
      else:
        values_max = np.array(jnp.max(values_target, axis=-1))

      # The advantage_estimator returns
      #     gamma^n_steps * values_max(s_{i + n_steps}) + discounted_rewards
      #        - values_max(s_i)
      # hence we have to add values_max(s_i) in order to get n-step returns:
      #     gamma^n_steps * values_max(s_{i + n_steps}) + discounted_rewards
      # Notice, that in DQN the tensor values_max[:, :-self._margin]
      # is the same as values_max[:, :-1].
      n_step_returns = values_max[:, :-self._margin] + \
          self._advantage_estimator(
              rewards=np_trajectory.reward,
              returns=np_trajectory.return_,
              values=values_max,
              dones=np_trajectory.done,
              discount_mask=np_trajectory.env_info.discount_mask,
          )

      length = n_step_returns.shape[1]
      target_returns = n_step_returns[:, :length]
      inputs = np_trajectory.observation[:, :length, :]

      yield (
          # Inputs are observations
          # (batch, length, obs)
          inputs,
          # the max indices will be needed to compute the loss
          np_trajectory.action[:, :length],  # index_max,
          # Targets: computed returns.
          # target_returns, we expect here shapes such as
          # (batch, length, num_actions)
          target_returns / self._value_network_scale,
          # TODO(henrykm): mask has the shape (batch, max_slice_length)
          # that is it misses the action dimension; the preferred format
          # would be np_trajectory.mask[:, :length, :] but for now we pass:
          np.ones(shape=target_returns.shape),
      )
def value_batches_stream ( self ) : max_slice_length = self . _max_slice_length min_slice_length = 1 for np_trajectory in self . _task . trajectory_batch_stream ( self . _value_batch_size , max_slice_length = max_slice_length , min_slice_length = min_slice_length , margin = 0 , epochs = self . _replay_epochs , ) : values_target = self . _run_value_model ( np_trajectory . observation , use_eval_model = True ) if self . _double_dqn : values = self . _run_value_model ( np_trajectory . observation , use_eval_model = False ) index_max = np . argmax ( values , axis = - 1 ) ind_0 , ind_1 = np . indices ( index_max . shape ) values_max = values_target [ ind_0 , ind_1 , index_max ] else : values_max = np . array ( jnp . max ( values_target , axis = - 1 ) ) # The advantage_estimator returns #     gamma^n_steps * values_max(s_{i + n_steps}) + discounted_rewards #        - values_max(s_i) # hence we have to add values_max(s_i) in order to get n-step returns: #     gamma^n_steps * values_max(s_{i + n_steps}) + discounted_rewards # Notice, that in DQN the tensor values_max[:, :-self._margin] # is the same as values_max[:, :-1]. n_step_returns = values_max [ : , : - self . _margin ] + self . _advantage_estimator ( rewards = np_trajectory . reward , returns = np_trajectory . return_ , values = values_max , dones = np_trajectory . done , discount_mask = np_trajectory . env_info . discount_mask , ) length = n_step_returns . shape [ 1 ] target_returns = n_step_returns [ : , : length ] inputs = np_trajectory . observation [ : , : length , : ] yield ( # Inputs are observations # (batch, length, obs) inputs , # the max indices will be needed to compute the loss np_trajectory . action [ : , : length ] , # index_max, # Targets: computed returns. # target_returns, we expect here shapes such as # (batch, length, num_actions) target_returns / self . _value_network_scale , # TODO(henrykm): mask has the shape (batch, max_slice_length) # that is it misses the action dimension; the preferred format # would be np_trajectory.mask[:, :length, :] but for now we pass: np . ones ( shape = target_returns . shape ) , )