code: def _get_expected_logp(
      fw, vars, obs_batch, a, layer_key, logp_func=None
  ):
      """Get the expected logp for the given obs_batch and action.
  
      Args:
          fw: Framework ("tf" or "torch").
          vars: The ModelV2 weights.
          obs_batch: The observation batch.
          a: The action batch.
          layer_key: The layer key to use for the fc layers.
          logp_func: Optional custom logp function to use.
  
      Returns:
          The expected logp.
      """
      if fw != "torch":
          if isinstance(vars, list):
              expected_mean_logstd = fc(
                  fc(obs_batch, vars[layer_key[1][0]]),
                  vars[layer_key[1][1]],
              )
          else:
              expected_mean_logstd = fc(
                  fc(
                      obs_batch,
                      vars[
                          "default_policy/{}_1/kernel".format(
                              layer_key[0]
                          )
                      ],
                  ),
                  vars[
                      "default_policy/{}_out/kernel".format(
                          layer_key[0]
                      )
                  ],
              )
      else:
          expected_mean_logstd = fc(
              fc(
                  obs_batch,
                  vars["{}_model.0.weight".format(layer_key[2][0])],
                  framework=fw,
              ),
              vars["{}_model.0.weight".format(layer_key[2][1])],
              framework=fw,
          )
      mean, log_std = np.split(expected_mean_logstd, 2, axis=-1)
      if logp_func is None:
          expected_logp = np.log(norm.pdf(a, mean, np.exp(log_std)))
      else:
          expected_logp = logp_func(mean, log_std, a)
  
      return expected_logp
  
imports: import unittest
  import numpy
  from scipy.stats import norm
  import ray
  import ray.rllib.algorithms.dqn
  import ray.rllib.algorithms.ppo
  import ray.rllib.algorithms.sac
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.utils.numpy import MAX_LOG_NN_OUTPUT, MIN_LOG_NN_OUTPUT, fc, one_hot
  from ray.rllib.utils.test_utils import check, framework_iterator
label_desc: Use image node_id 1 for calling a global function with example usage: _get_expected_logp(fw, vars, obs_batch, a, layer_key, logp_func) and returns: expected_logp