code: def _batched_step_to_transition(
      step: rlds.BatchedStep,
  ) -> types.Transition:
      return types.Transition(
          observation=tf.nest.map_structure(
              lambda x: x[0], step[rlds.OBSERVATION]
          ),
          action=tf.nest.map_structure(
              lambda x: x[0], step[rlds.ACTION]
          ),
          reward=tf.nest.map_structure(
              lambda x: x[0], step[rlds.REWARD]
          ),
          discount=1.0
          - tf.cast(step[rlds.IS_TERMINAL][1], dtype=tf.float32),
          # If next step is terminal, then the observation may be arbitrary.
          next_observation=tf.nest.map_structure(
              lambda x: x[1], step[rlds.OBSERVATION]
          ),
      )
  
imports: import logging
  from typing import Any, Iterator, Optional, Tuple, Sequence
  from acme import specs
  from acme import types
  from flax import jax_utils
  import jax
  import jax.numpy
  import numpy
  import rlds
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling a global function with example usage: _batched_step_to_transition(step) and returns: types