code: def atari_example_to_rlds(
      tf_example: tf.train.Example,
  ) -> Dict[str, Any]:
      """Generates an RLDS episode from an Atari TF Example.
  
      Args:
        tf_example: example from an Atari dataset.
  
      Returns:
        RLDS episode.
      """
  
      data = tf.io.parse_single_example(
          tf_example, _feature_description()
      )
      episode_length = tf.size(data["actions"])
      is_first = tf.concat(
          [[True], [False] * tf.ones(episode_length - 1)], axis=0
      )
      is_last = tf.concat(
          [[False] * tf.ones(episode_length - 1), [True]], axis=0
      )
  
      is_terminal = [False] * tf.ones_like(data["actions"])
      discounts = data["discounts"]
      if discounts[-1] == 0.0:
          is_terminal = tf.concat(
              [[False] * tf.ones(episode_length - 1, tf.int64), [True]],
              axis=0,
          )
          # If the episode ends in a terminal state, in the last step only the
          # observation has valid information (the terminal state).
          discounts = tf.concat([discounts[1:], [0.0]], axis=0)
      episode = {
          # Episode Metadata
          "episode_id": data["episode_idx"],
          "checkpoint_id": data["checkpoint_idx"],
          "episode_return": data["episode_return"],
          "steps": {
              "observation": data["observations"],
              "action": data["actions"],
              "reward": data["unclipped_rewards"],
              "discount": discounts,
              "is_first": is_first,
              "is_last": is_last,
              "is_terminal": is_terminal,
          },
      }
      return episode
  
file_path: datasets/tensorflow_datasets/rl_unplugged/atari_utils.py
imports: from __future__ import annotations
  import dataclasses
  from typing import Any, Dict
  import numpy
  from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow
  import tensorflow_datasets.public_api
label_desc: Use image node_id 6 for calling a global function with example usage: atari_example_to_rlds(tf_example) and returns: episode