code: def get_features_dict(
      builder_config: BuilderConfig, ds_config: DatasetConfig
  ) -> Dict[str, tfds.typing.FeatureConnectorArg]:
      """Builds the features dict of a D4RL dataset.
  
      Args:
        builder_config: builder config of the Mujoco dataset.
        ds_config: config of the Mujoco dataset containing the specs.
  
      Returns:
        Dictionary with the features of this dataset.
      """
  
      float_type = builder_config.float_type
  
      steps_dict = {
          "observation": tfds.features.Tensor(
              shape=(ds_config.obs_len,), dtype=float_type
          ),
          "action": tfds.features.Tensor(
              shape=(ds_config.action_len,), dtype=float_type
          ),
          "reward": float_type,
          "is_terminal": tf.bool,
          "is_first": tf.bool,
          "is_last": tf.bool,
          "discount": float_type,
      }
      if (
          builder_config.env == "adroit"
          and builder_config.name == "v0-cloned"
      ):
          # D4RL adroit in the v0-cloned config, action uses a different float type
          # than the rest of the dataset.
          steps_dict["action"] = tfds.features.Tensor(
              shape=(ds_config.action_len,), dtype=np.float32
          )
  
      if builder_config.step_metadata_keys:
          steps_dict["infos"] = _get_step_metadata(
              builder_config, ds_config
          )
  
      episode_metadata = {}
      if builder_config.episode_metadata_keys:
          for k in builder_config.episode_metadata_keys:
              if k == _ITERATION:
                  episode_metadata[k] = tf.int32
              else:
                  episode_metadata[k] = tf.string
      if builder_config.has_policy_metadata:
          episode_metadata["policy"] = _get_policy_info(
              builder_config, ds_config
          )
  
      features_dict = {
          "steps": tfds.features.Dataset(steps_dict),
      }
      if episode_metadata:
          features_dict.update(episode_metadata)
  
      return features_dict
  
file_path: datasets/tensorflow_datasets/d4rl/dataset_builder.py
imports: from __future__ import annotations
  import dataclasses
  from typing import Any, Dict, FrozenSet, Optional, Type, Union
  import numpy
  from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow
  from tensorflow_datasets.d4rl import dataset_utils
  import tensorflow_datasets.public_api
label_desc: Use image node_id 3 for calling a global function with example usage: get_features_dict(builder_config, ds_config) and returns: features_dict