code: def _get_policy_info(
      builder_config: BuilderConfig, ds_config: DatasetConfig
  ) -> Dict[str, tfds.typing.FeatureConnectorArg]:
      """Builds the features dict of the policy weights.
  
      Args:
        builder_config: builder config of the dataset.
        ds_config: config containing the dataset specs.
  
      Returns:
        Dictionary with the policy-related features of this dataset.
      """
      float_type = builder_config.float_type
      # The policy dictionary contains the weights of the policy used to
      # generate the dataset.
      # See https://github.com/rail-berkeley/d4rl/wiki/Tasks#gym.
      policy_dict = {
          "fc0": {
              "bias": tfds.features.Tensor(
                  shape=(builder_config.policy_size,), dtype=float_type
              ),
              "weight": tfds.features.Tensor(
                  shape=(builder_config.policy_size, ds_config.obs_len),
                  dtype=float_type,
              ),
          },
          "fc1": {
              "bias": tfds.features.Tensor(
                  shape=(builder_config.policy_size,), dtype=float_type
              ),
              "weight": tfds.features.Tensor(
                  shape=(
                      builder_config.policy_size,
                      builder_config.policy_size,
                  ),
                  dtype=float_type,
              ),
          },
          "last_fc": {
              "bias": tfds.features.Tensor(
                  shape=(ds_config.action_len,), dtype=float_type
              ),
              "weight": tfds.features.Tensor(
                  shape=(
                      ds_config.action_len,
                      builder_config.policy_size,
                  ),
                  dtype=float_type,
              ),
          },
          "nonlinearity": tf.string,
          "output_distribution": tf.string,
      }
      if (
          builder_config.env == "adroit"
          and builder_config.name == "v1-cloned"
      ):
          # v1-cloned from d4rl_adroit uses a different policy shape
          policy_dict["fc0"]["weight"] = tfds.features.Tensor(
              shape=(ds_config.obs_len, builder_config.policy_size),
              dtype=float_type,
          )
          policy_dict["last_fc"]["weight"] = tfds.features.Tensor(
              shape=(builder_config.policy_size, ds_config.action_len),
              dtype=float_type,
          )
      if builder_config.has_policy_last_fc_log_std:
          policy_dict["last_fc_log_std"] = {
              "bias": tfds.features.Tensor(
                  shape=(ds_config.action_len,), dtype=float_type
              ),
              "weight": tfds.features.Tensor(
                  shape=(
                      ds_config.action_len,
                      builder_config.policy_size,
                  ),
                  dtype=float_type,
              ),
          }
  
      return policy_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 2 for calling a global function with example usage: _get_policy_info(builder_config, ds_config) and returns: policy_dict