(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 batches_stream(self):
    """Use the RLTask self._task to create inputs to the value model."""
    for np_trajectory in self._task.trajectory_batch_stream(
        self._batch_size, max_slice_length=self._max_slice_length, epochs=[-1]):
      # Insert an extra depth dimension, so the target shape is consistent with
      # the network output shape.
      yield (np_trajectory.observation,         # Inputs to the value model.
             np_trajectory.return_[:, :, None],
             np_trajectory.done[:, :, None],
             np_trajectory.reward[:, :, None],
             np_trajectory.action,
             jnp.zeros_like(np_trajectory.mask),
             np_trajectory.mask)
def batches_stream ( self ) : for np_trajectory in self . _task . trajectory_batch_stream ( self . _batch_size , max_slice_length = self . _max_slice_length , epochs = [ - 1 ] ) : # Insert an extra depth dimension, so the target shape is consistent with # the network output shape. yield ( np_trajectory . observation , # Inputs to the value model. np_trajectory . return_ [ : , : , None ] , np_trajectory . done [ : , : , None ] , np_trajectory . reward [ : , : , None ] , np_trajectory . action , jnp . zeros_like ( np_trajectory . mask ) , np_trajectory . mask )