class_docstr: The Direct Method estimator.
  
  Let s_t, a_t, and r_t be the state, action, and reward at timestep t.
  
  This method trains a Q-model for the evaluation policy \pi_e on behavior
  data generated by \pi_b. Currently, RLlib implements this using
  Fitted-Q Evaluation (FQE). You can also implement your own model
  and pass it in as `q_model_config = {"type": your_model_class, **your_kwargs}`.
  
  This estimator computes the expected return for \pi_e for an episode as:
  V^{\pi_e}(s_0) = \sum_{a \in A} \pi_e(a | s_0) Q(s_0, a)
  and returns the mean and standard deviation over episodes.
  
  For more information refer to https://arxiv.org/pdf/1911.06854.pdf
docstr: Calculates the Direct Method estimate on the given dataset.
  
  Note: This estimate works for only discrete action spaces for now.
  
  Args:
      dataset: Dataset to compute the estimate on. Each record in dataset should
          include the following columns: `obs`, `actions`, `action_prob` and
          `rewards`. The `obs` on each row shoud be a vector of D dimensions.
      n_parallelism: The number of parallel workers to use.
  
  Returns:
      Dictionary with the following keys:
          v_target: The estimated value of the target policy.
          v_behavior: The estimated value of the behavior policy.
          v_gain: The estimated gain of the target policy over the behavior
              policy.
          v_std: The standard deviation of the estimated value of the target.
docstr_tok: ["Calculates","the","Direct","Method","estimate","on","the","given","dataset",".","Note",":","This","estimate","works","for","only","discrete","action","spaces","for","now",".","Args",":","dataset",":","Dataset","to","compute","the","estimate","on",".","Each","record","in","dataset","should","include","the","following","columns",":","`","obs","`",",","`","actions","`",",","`","action_prob","`","and","`","rewards","`",".","The","`","obs","`","on","each","row","shoud","be","a","vector","of","D","dimensions",".","n_parallelism",":","The","number","of","parallel","workers","to","use",".","Returns",":","Dictionary","with","the","following","keys",":","v_target",":","The","estimated","value","of","the","target","policy",".","v_behavior",":","The","estimated","value","of","the","behavior","policy",".","v_gain",":","The","estimated","gain","of","the","target","policy","over","the","behavior","policy",".","v_std",":","The","standard","deviation","of","the","estimated","value","of","the","target","."]
code: def estimate_on_dataset(
      self, dataset: Dataset, *, n_parallelism: int = ...
  ) -> Dict[str, Any]:
      """Calculates the Direct Method estimate on the given dataset.
  
      Note: This estimate works for only discrete action spaces for now.
  
      Args:
          dataset: Dataset to compute the estimate on. Each record in dataset should
              include the following columns: `obs`, `actions`, `action_prob` and
              `rewards`. The `obs` on each row shoud be a vector of D dimensions.
          n_parallelism: The number of parallel workers to use.
  
      Returns:
          Dictionary with the following keys:
              v_target: The estimated value of the target policy.
              v_behavior: The estimated value of the behavior policy.
              v_gain: The estimated gain of the target policy over the behavior
                  policy.
              v_std: The standard deviation of the estimated value of the target.
      """
      # compute v_values
      batch_size = max(dataset.count() // n_parallelism, 1)
      updated_ds = dataset.map_batches(
          compute_q_and_v_values,
          batch_size=batch_size,
          batch_format="pandas",
          fn_kwargs={
              "model_class": self.model.__class__,
              "model_state": self.model.get_state(),
              "compute_q_values": False,
          },
      )
  
      v_behavior = updated_ds.mean("rewards")
      v_target = updated_ds.mean("v_values")
      v_gain_mean = v_target / v_behavior
      v_gain_ste = (
          updated_ds.std("v_values")
          / v_behavior
          / math.sqrt(dataset.count())
      )
  
      return {
          "v_behavior": v_behavior,
          "v_target": v_target,
          "v_gain_mean": v_gain_mean,
          "v_gain_ste": v_gain_ste,
      }
  
imports: import logging
  from typing import Dict, Any, Optional, List
  import math
  import numpy
  from ray.data import Dataset
  from ray.rllib.offline.estimators.off_policy_estimator import OffPolicyEstimator
  from ray.rllib.offline.offline_evaluation_utils import compute_q_and_v_values
  from ray.rllib.offline.offline_evaluator import OfflineEvaluator
  from ray.rllib.offline.estimators.fqe_torch_model import FQETorchModel
  from ray.rllib.policy import Policy
  from ray.rllib.policy.sample_batch import convert_ma_batch_to_sample_batch
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.utils.annotations import DeveloperAPI, override
  from ray.rllib.utils.typing import SampleBatchType
  from ray.rllib.utils.numpy import convert_to_numpy
label_desc: Use image node_id 6 for calling the DirectMethod obj's underlying member method code with example usage: obj.estimate_on_dataset(dataset) and returns: dict