code: def _get_step_metadata(
      builder_config: BuilderConfig, ds_config: DatasetConfig
  ) -> Dict[
      str,
      Union[np.dtype, Type[np.generic], tfds.features.FeatureConnector],
  ]:
      """Builds the features dict of the step metadata.
  
      Args:
        builder_config: builder config of the dataset.
        ds_config: config containing the dataset specs.
  
      Returns:
        Dictionary with the step metadata features of this dataset.
      """
      infos_dict = {}
      float_type = builder_config.step_metadata_type
      # Step metadata corresponds to state information.
      # See https://github.com/rail-berkeley/d4rl/wiki/Tasks#gym.
      for k in builder_config.step_metadata_keys:
          if k in _ACTION_LOG_PROBS:
              infos_dict[k] = float_type
          elif k == _QPOS:
              infos_dict[k] = tfds.features.Tensor(
                  shape=(ds_config.qpos_len,), dtype=float_type
              )
          elif k == _QVEL:
              infos_dict[k] = tfds.features.Tensor(
                  shape=(ds_config.qvel_len,), dtype=float_type
              )
          elif k == _GOAL:
              infos_dict[k] = tfds.features.Tensor(
                  shape=(ds_config.goal_len,), dtype=float_type
              )
          elif k == _ADROIT_BODY_POS:
              # We use the task name to get the names of the body position keys.
              adroit_keys = _ADROIT_BODY_POS_KEYS[ds_config.name]
              for ak, shape in adroit_keys.items():
                  infos_dict[ak] = tfds.features.Tensor(
                      shape=shape, dtype=float_type
                  )
          elif k in [_ACTION_MEAN, _ACTION_LOG_STD, _ACTION_LOGSTD]:
              infos_dict[k] = tfds.features.Tensor(
                  shape=(ds_config.action_len,), dtype=float_type
              )
          else:
              raise ValueError(f"Unknown key in the step metadata {k}")
      return infos_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 1 for calling a global function with example usage: _get_step_metadata(builder_config, ds_config) and returns: infos_dict