class_docstr: Creates (s,a,r,s',a') transitions.
class_docstr_tok: ["Creates","(","s",",","a",",","r",",","s","'",",","a","'",")","transitions","."]
docstr: Convert a batch to transition tuple.
docstr_tok: ["Convert","a","batch","to","transition","tuple","."]
code: def create_transitions(self, batch, episode_return=None):
      """Convert a batch to transition tuple."""
      all_states = tf.squeeze(batch[rlds.OBSERVATION], axis=-1)
      all_states = tf.transpose(all_states, perm=[1, 2, 0])
      rewards = batch[rlds.REWARD][self.stack_size - 1 : -1]
      terminals = batch[rlds.IS_TERMINAL][
          self.stack_size : self.total_frames
      ]
      transitions = {
          "state": all_states[:, :, : self.stack_size],
          "action": batch[rlds.ACTION][self.stack_size - 1],
          "reward": tf.reduce_sum(rewards * self.cumulative_discount),
          "next_state": all_states[:, :, self.update_horizon :],
          "terminal": tf.reduce_any(terminals),
      }
      if episode_return is not None:
          transitions["episode_return"] = episode_return
      if self.next_action:
          transitions["next_action"] = batch[rlds.ACTION][
              self.total_frames - 1
          ]
      if self.return_to_go:
          transitions["return_to_go"] = batch["return_to_go"][
              self.stack_size - 1
          ]
      return transitions
  
file_path: dopamine/dopamine/labs/offline_rl/rlu_tfds/tfds_atari_utils.py
imports: import rlds
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling the BatchToTransition obj's underlying member method code with example usage: obj.create_transitions(batch, episode_return) and returns: transitions