code: def _compute_adv_and_logps(
      self,
      model: ModelV2,
      dist_class: Type[TorchDistributionWrapper],
      train_batch: SampleBatch,
  ) -> None:
      # uses mean|max|expectation to compute estimate of advantages
      # continuous/discrete action spaces:
      # for max:
      #   A(s_t, a_t) = Q(s_t, a_t) - max_{a^j} Q(s_t, a^j)
      #   where a^j is m times sampled from the policy p(a | s_t)
      # for mean:
      #   A(s_t, a_t) = Q(s_t, a_t) - avg( Q(s_t, a^j) )
      #   where a^j is m times sampled from the policy p(a | s_t)
      # discrete action space and adv_type=expectation:
      #   A(s_t, a_t) = Q(s_t, a_t) - sum_j[Q(s_t, a^j) * pi(a^j)]
      advantage_type = self.config["advantage_type"]
      n_action_sample = self.config["n_action_sample"]
      batch_size = len(train_batch)
      out_t, _ = model(train_batch)
  
      # construct pi(s_t) and Q(s_t, a_t) for computing advantage actions
      pi_s_t = dist_class(model.get_policy_output(out_t), model)
      q_t = self._get_q_value(
          model, out_t, train_batch[SampleBatch.ACTIONS]
      )
  
      # compute the logp of the actions in the dataset (for computing actor's loss)
      action_logp = pi_s_t.dist.log_prob(
          train_batch[SampleBatch.ACTIONS]
      )
  
      # fix the shape if it's not canonical (i.e. shape[-1] != 1)
      if len(action_logp.shape) <= 1:
          action_logp.unsqueeze_(-1)
      train_batch[SampleBatch.ACTION_LOGP] = action_logp
  
      if advantage_type == "expectation":
          assert (
              self._is_action_discrete
          ), "Action space should be discrete when advantage_type = expectation."
          assert hasattr(
              self.model, "q_model"
          ), "CRR's ModelV2 should have q_model neural network in discrete \
                  action spaces"
          assert isinstance(
              pi_s_t.dist, torch.distributions.Categorical
          ), "The output of the policy should be a torch Categorical \
                  distribution."
  
          q_vals = self.model.q_model(out_t)
          if hasattr(self.model, "twin_q_model"):
              q_twins = self.model.twin_q_model(out_t)
              q_vals = torch.minimum(q_vals, q_twins)
  
          probs = pi_s_t.dist.probs
          v_t = (q_vals * probs).sum(-1, keepdims=True)
      else:
          policy_actions = pi_s_t.dist.sample(
              (n_action_sample,)
          )  # samples
  
          if self._is_action_discrete:
              flat_actions = policy_actions.reshape(-1)
          else:
              flat_actions = policy_actions.reshape(
                  -1, *self.action_space.shape
              )
  
          reshaped_s_t = train_batch[SampleBatch.OBS].view(
              1, batch_size, *self.observation_space.shape
          )
          reshaped_s_t = reshaped_s_t.expand(
              n_action_sample, batch_size, *self.observation_space.shape
          )
          flat_s_t = reshaped_s_t.reshape(
              -1, *self.observation_space.shape
          )
  
          input_v_t = SampleBatch(
              **{
                  SampleBatch.OBS: flat_s_t,
                  SampleBatch.ACTIONS: flat_actions,
              }
          )
          out_v_t, _ = model(input_v_t)
  
          flat_q_st_pi = self._get_q_value(model, out_v_t, flat_actions)
          reshaped_q_st_pi = flat_q_st_pi.reshape(-1, batch_size, 1)
  
          if advantage_type == "mean":
              v_t = reshaped_q_st_pi.mean(dim=0)
          elif advantage_type == "max":
              v_t, _ = reshaped_q_st_pi.max(dim=0)
          else:
              raise ValueError(
                  f"Invalid advantage type: {advantage_type}."
              )
  
      adv_t = q_t - v_t
      train_batch["advantages"] = adv_t
  
      # logging
      self.log("q_batch_avg", q_t.mean())
      self.log("q_batch_max", q_t.max())
      self.log("q_batch_min", q_t.min())
      self.log("v_batch_avg", v_t.mean())
      self.log("v_batch_max", v_t.max())
      self.log("v_batch_min", v_t.min())
      self.log("adv_batch_avg", adv_t.mean())
      self.log("adv_batch_max", adv_t.max())
      self.log("adv_batch_min", adv_t.min())
      self.log(
          "reward_batch_avg", train_batch[SampleBatch.REWARDS].mean()
      )
  
imports: from typing import Dict, List, Tuple, Type, Union, cast
  import gymnasium
  import numpy
  from rllib_crr.crr.crr_torch_model import CRRModel
  from ray.rllib.algorithms import AlgorithmConfig
  from ray.rllib.algorithms.ddpg.noop_model import TorchNoopModel
  from ray.rllib.models.catalog import ModelCatalog
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.models.torch.torch_action_dist import TorchCategorical, TorchDistributionWrapper, get_torch_categorical_class_with_temperature
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.torch_mixins import TargetNetworkMixin
  from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2
  from ray.rllib.utils.annotations import override
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.torch_utils import huber_loss, l2_loss
  from ray.rllib.utils.typing import TensorType, TrainerConfigDict
label_desc: Use image node_id 9 for calling the CRRTorchPolicy obj's underlying member method code with example usage: obj._compute_adv_and_logps(model, dist_class, train_batch) without return types