code: def features_dict():
      return tfds.features.FeaturesDict(
          {
              "steps": tfds.features.Dataset(
                  {
                      "observation": tfds.features.Image(
                          shape=(
                              84,
                              84,
                              1,
                          ),
                          dtype=np.uint8,
                          encoding_format="png",
                      ),
                      "action": np.int64,
                      "reward": tfds.features.Scalar(
                          dtype=np.float32,
                          doc=tfds.features.Documentation(
                              desc="Clipped reward.",
                              value_range="[-1, 1]",
                          ),
                      ),
                      "is_terminal": np.bool_,
                      "is_first": np.bool_,
                      "is_last": np.bool_,
                      "discount": np.float32,
                  }
              ),
              "checkpoint_id": np.int64,
              "episode_id": np.int64,
              "episode_return": tfds.features.Scalar(
                  dtype=np.float32,
                  doc=tfds.features.Documentation(
                      desc="Sum of the clipped rewards."
                  ),
              ),
          }
      )
  
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 8 for calling a global function with example usage: features_dict() and returns: tfds