code: def _generate_steps(
      episode: Sequence[Any],
      step_fn: Callable[[Dict[str, Any]], Dict[str, Any]],
  ) -> Dict[str, Any]:
      """Constructs a dictionary of steps for the given episode.
  
      Args:
        episode: Sequence of steps (envlogger.StepData). Assumes that there is at
          least one step.
        step_fn: function that gets an RLDS step and returns the processed step.
  
      Returns:
        Nested dictionary with the steps of one episode. Steps
        are represented by a nested dictionary where each nested field is a list.
      """
      if len(episode) < 2:
          step = _build_last_step(episode[0], step_fn)
          return _to_nested_list(step)
      # We use episode[1] here because in the first steps some of the fields(like
      # the action, migt be None, so the shape is unknown).
      empty_step = _build_empty_step(episode[1], step_fn)
      steps = _empty_nested_list(empty_step)
  
      prev_step = None
      for step in episode:
          if prev_step is not None:
              rlds_step = _build_step(prev_step, step, step_fn)
              for k, v in rlds_step.items():
                  steps[k] = _append_nested(steps[k], v)
          prev_step = step
  
      if prev_step is not None:
          last_step = _build_last_step(prev_step, step_fn)
          for k, v in last_step.items():
              steps[k] = _append_nested(steps[k], v)
      return steps
  
file_path: datasets/tensorflow_datasets/rlds/envlogger_reader.py
imports: from typing import Any, Callable, Dict, Generator, List, Optional, Sequence, Tuple, Union
  import numpy
  from tensorflow_datasets.core.utils import tree_utils
label_desc: Use image node_id 8 for calling a global function with example usage: _generate_steps(episode, step_fn) and returns: steps, _to_nested_list