code: def _step_to_transition(
      rlds_step: rlds.BatchedStep,
  ) -> types.Transition:
      """Converts batched RLDS steps to batched transitions."""
      return types.Transition(
          observation=rlds_step[rlds.OBSERVATION],
          action=rlds_step[rlds.ACTION],
          reward=rlds_step[rlds.REWARD],
          discount=rlds_step[rlds.DISCOUNT],
          #  We provide next_observation if an algorithm needs it, however note that
          # it will only contain s_t and s_t+1, so will be one element short of all
          # other attributes (which contain s_t-1, s_t, s_t+1).
          next_observation=tree.map_structure(
              lambda x: x[1:], rlds_step[rlds.OBSERVATION]
          ),
          extras={
              N_STEP_RETURN: rlds_step[N_STEP_RETURN],
          },
      )
  
file_path: acme/acme/agents/jax/mbop/dataset.py
imports: import functools
  import itertools
  from typing import Iterator, Optional
  from acme import types
  from acme.jax import running_statistics
  import jax
  import jax.numpy
  import rlds
  import tensorflow
  import tree
label_desc: Use image node_id 5 for calling a global function with example usage: _step_to_transition(rlds_step) and returns: types