docstr: Estimate the feature importance of the policy given a dataset.
  
  For each feature in the dataset, the importance is computed by applying
  perturbations to each feature and computing the difference between the
  perturbed prediction and the reference prediction. The importance
  computation for each feature and each perturbation is repeated `self.repeat`
  times. If dataset is large the user can initialize the estimator with a
  `limit_fraction` to limit the dataset to a fraction of the original dataset.
  
  The dataset should include a column named `obs` where each row is a vector of D
  dimensions. The importance is computed for each dimension of the vector.
  
  Note (Implementation detail): The computation across features are distributed
  with ray workers since each feature is independent of each other.
  
  Args:
      dataset: the dataset to use for feature importance.
      n_parallelism: number of parallel workers to use for feature importance.
  
  Returns:
      A dict mapping each feature index string to its importance.
docstr_tok: ["Estimate","the","feature","importance","of","the","policy","given","a","dataset",".","For","each","feature","in","the","dataset",",","the","importance","is","computed","by","applying","perturbations","to","each","feature","and","computing","the","difference","between","the","perturbed","prediction","and","the","reference","prediction",".","The","importance","computation","for","each","feature","and","each","perturbation","is","repeated","`","self.repeat","`","times",".","If","dataset","is","large","the","user","can","initialize","the","estimator","with","a","`","limit_fraction","`","to","limit","the","dataset","to","a","fraction","of","the","original","dataset",".","The","dataset","should","include","a","column","named","`","obs","`","where","each","row","is","a","vector","of","D","dimensions",".","The","importance","is","computed","for","each","dimension","of","the","vector",".","Note","(","Implementation","detail",")",":","The","computation","across","features","are","distributed","with","ray","workers","since","each","feature","is","independent","of","each","other",".","Args",":","dataset",":","the","dataset","to","use","for","feature","importance",".","n_parallelism",":","number","of","parallel","workers","to","use","for","feature","importance",".","Returns",":","A","dict","mapping","each","feature","index","string","to","its","importance","."]
code: def estimate_on_dataset(
      self, dataset: Dataset, *, n_parallelism: int = ...
  ) -> Dict[str, Any]:
      """Estimate the feature importance of the policy given a dataset.
  
      For each feature in the dataset, the importance is computed by applying
      perturbations to each feature and computing the difference between the
      perturbed prediction and the reference prediction. The importance
      computation for each feature and each perturbation is repeated `self.repeat`
      times. If dataset is large the user can initialize the estimator with a
      `limit_fraction` to limit the dataset to a fraction of the original dataset.
  
      The dataset should include a column named `obs` where each row is a vector of D
      dimensions. The importance is computed for each dimension of the vector.
  
      Note (Implementation detail): The computation across features are distributed
      with ray workers since each feature is independent of each other.
  
      Args:
          dataset: the dataset to use for feature importance.
          n_parallelism: number of parallel workers to use for feature importance.
  
      Returns:
          A dict mapping each feature index string to its importance.
      """
  
      policy_state = self.policy.get_state()
      # step 1: limit the dataset to a few first rows
      ds = dataset.limit(int(self.limit_fraction * dataset.count()))
  
      # step 2: compute the reference actions
      bsize = max(1, ds.count() // n_parallelism)
      actions_ds = ds.map_batches(
          _compute_actions,
          batch_size=bsize,
          fn_kwargs={
              "output_key": "ref_actions",
              "policy_state": policy_state,
          },
      )
  
      # step 3: compute the feature importance
      n_features = ds.take(1)[0][SampleBatch.OBS].shape[-1]
      importance = np.zeros((self.repeat, n_features))
      for r in range(self.repeat):
          # shuffle the entire dataset
          shuffled_ds = actions_ds.random_shuffle()
          bsize_per_task = max(
              1, (shuffled_ds.count() * n_features) // n_parallelism
          )
  
          # for each index perturb the dataset and compute the feat importance score
          remote_fns = [
              get_feature_importance_on_index.remote(
                  dataset=shuffled_ds,
                  index=i,
                  perturb_fn=self.perturb_fn,
                  bsize=bsize_per_task,
                  policy_state=policy_state,
              )
              for i in range(n_features)
          ]
          ds_w_fi_scores = ray.get(remote_fns)
          importance[r] = np.array(
              [d.mean("delta") for d in ds_w_fi_scores]
          )
  
      importance = importance.mean(0)
      metrics = {
          f"feature_{i}": importance[i] for i in range(len(importance))
      }
  
      return metrics
  
file_path: ray/python/ray/rllib/offline/feature_importance.py
imports: import copy
  import numpy
  import pandas
  from typing import Callable, Dict, Any
  import ray
  from ray.data import Dataset
  from ray.rllib.policy import Policy
  from ray.rllib.policy.sample_batch import SampleBatch, convert_ma_batch_to_sample_batch
  from ray.rllib.utils.annotations import override, DeveloperAPI, ExperimentalAPI
  from ray.rllib.utils.typing import SampleBatchType
  from ray.rllib.offline.offline_evaluator import OfflineEvaluator
label_desc: Use image node_id 3 for calling the FeatureImportance obj's underlying member method code with example usage: obj.estimate_on_dataset(dataset) and returns: metrics