class_docstr: Trains a joint policy-and-value model using AWR.
class_docstr_tok: ["Trains","a","joint","policy-and-value","model","using","AWR","."]
docstr: Use the RLTask self._task to create inputs to the value model.
docstr_tok: ["Use","the","RLTask","self._task","to","create","inputs","to","the","value","model","."]
code: 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
      ):
          # 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
              ],  # Targets: regress to returns.
              np_trajectory.action,  # Policy targets: actions.
              np_trajectory.mask,
          )
  
imports: import functools
  from trax import data
  from trax import layers
  from trax import supervised
  from trax.fastmath import numpy
  from trax.fastmath import stop_gradient
  from trax.rl import actor_critic
  from trax.rl import distributions
  from trax.rl import rl_layers
  from trax.rl import training
  from trax.supervised import lr_schedules
label_desc: Use image node_id 2 for calling the AWRJoint obj's underlying member method code with example usage: obj.batches_stream() without return types