code: def _get_episode(
      steps: Dict[str, Any],
      episode_metadata: Dict[str, Any],
      begin: int,
      end: int,
  ) -> Dict[str, Any]:
      """Builds a full episode dict.
  
      Args:
          steps: a dict with all steps in a dataset
          episode_metadata: dict with the episode metadata
          begin: defines a starting position of an episode
          end: defines an ending position of an episode
  
      Returns:
         A dict with data specific to one episode, already broken into steps.
      """
      # It's an initial step if the episode is empty.
      episode = {}
      for k in [
          "is_first",
          "is_last",
          "observation",
          "action",
          "reward",
          "discount",
      ]:
          episode[k] = steps[k][begin:end]
  
      episode["is_last"] = [False] * (end - begin)
      episode["is_terminal"] = [False] * (end - begin)
      if "infos" in steps.keys():
          episode["infos"] = {}
          for k in steps["infos"].keys():
              episode["infos"][k] = steps["infos"][k][begin:end]
  
      ends_in_terminal = steps["is_terminal"][end - 1]
      has_next_obs = "next_observation" in steps
  
      # If the episode ends in a terminal state, the discount of the previous step
      # is set to 0 (it's a transition to a terminal state).
      if ends_in_terminal:
          episode["discount"][-1] = 0.0
      if ends_in_terminal or has_next_obs:
          # In HDF5, the terminal bit is associated with the previous observation.
          # To comply with RLDS standard (see types.py), we propagate this information
          # to a next state. This matches the definition in RLDS. See types.py.
          # If there is information about a next observation, we also generate an
          # extra last step.
          episode["is_first"] = np.concatenate(
              (episode["is_first"], [False])
          )
          if has_next_obs:
              episode["observation"] = np.concatenate(
                  (
                      episode["observation"],
                      [steps["next_observation"][-1]],
                  )
              )
          else:
              # If the last observation is never recorded and to avoid discarding the
              # last transition to the terminal state, we create a dummy observation.
              # Since no solution is perfect, the design choice was to keep as much
              # information as possible and let the user decide to keep or ignore such
              # transitions.
              episode["observation"] = np.concatenate(
                  (
                      episode["observation"],
                      [np.zeros_like(steps["observation"][0])],
                  )
              )
          # Action/reward/discount are set to dummy values since not relevant.
          # When IS_LAST is set, any field coming temporally after the last
          # observation is invalid.
          episode["action"] = np.concatenate(
              (episode["action"], [np.zeros_like(steps["action"][0])])
          )
          episode["reward"] = np.concatenate(
              (episode["reward"], [np.zeros_like(steps["reward"][0])])
          )
          episode["discount"] = np.array(
              np.concatenate((episode["discount"], [0.0])),
              dtype=np.float32,
          )
  
          episode["is_terminal"] = np.concatenate(
              (episode["is_terminal"], [ends_in_terminal])
          )
          episode["is_last"] = np.concatenate(
              (episode["is_last"], [True])
          )
          if "infos" in steps.keys():
              for k in steps["infos"].keys():
                  episode["infos"][k] = np.concatenate(
                      (
                          episode["infos"][k],
                          [np.zeros_like(steps["infos"][k][0])],
                      )
                  )
      else:
          # Despite the fact that the last action and reward are valid in the
          # stored dataset (in the final transition, [obs, action, reward, next_obs],
          # only one step was stored as [obs, action, reward], instead of adding one
          # extra step with [next_obs, 0, 0]). We set IS_LAST=True so that
          # it is consistent with other typical datasets.
          episode["is_last"][-1] = True
      full_episode = {"steps": episode}
      if episode_metadata:
          full_episode.update(episode_metadata)
      return full_episode
  
file_path: datasets/tensorflow_datasets/d4rl/dataset_utils.py
imports: from typing import Any, Dict
  import h5py
  import numpy
  from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow
label_desc: Use image node_id 6 for calling a global function with example usage: _get_episode(steps, episode_metadata, begin, end) and returns: full_episode