code: def do_test_log_likelihood(
      run,
      config,
      prev_a=None,
      continuous=False,
      layer_key=("fc", (0, 4), ("_hidden_layers.0.", "_logits.")),
      logp_func=None,
  ):
      config = config.copy(copy_frozen=False)
      # Run locally.
      config.num_rollout_workers = 0
      # Env setup.
      if continuous:
          config.env = "Pendulum-v1"
          obs_batch = preprocessed_obs_batch = np.array(
              [[0.0, 0.1, -0.1]]
          )
      else:
          config.env = "FrozenLake-v1"
          config.env_config = {"is_slippery": False, "map_name": "4x4"}
          obs_batch = np.array([0])
          # PG does not preprocess anymore by default.
          preprocessed_obs_batch = one_hot(obs_batch, depth=16)
  
      prev_r = None if prev_a is None else np.array(0.0)
  
      # Test against all frameworks.
      for fw in framework_iterator(config):
          algo = config.build()
  
          policy = algo.get_policy()
          vars = policy.get_weights()
          # Sample n actions, then roughly check their logp against their
          # counts.
          num_actions = 1000 if not continuous else 50
          actions = []
          for _ in range(num_actions):
              # Single action from single obs.
              actions.append(
                  algo.compute_single_action(
                      obs_batch[0],
                      prev_action=prev_a,
                      prev_reward=prev_r,
                      explore=True,
                      # Do not unsquash actions
                      # (remain in normalized [-1.0; 1.0] space).
                      unsquash_action=False,
                  )
              )
  
          # Test all taken actions for their log-likelihoods vs expected values.
          if continuous:
              for idx in range(num_actions):
                  a = actions[idx]
  
                  logp = policy.compute_log_likelihoods(
                      np.array([a]),
                      preprocessed_obs_batch,
                      prev_action_batch=np.array([prev_a])
                      if prev_a
                      else None,
                      prev_reward_batch=np.array([prev_r])
                      if prev_r
                      else None,
                      actions_normalized=True,
                      in_training=False,
                  )
  
                  # The expected logp computation logic is overfitted to the ModelV2
                  # stack and does not generalize to RLModule API.
                  if not config._enable_new_api_stack:
                      expected_logp = _get_expected_logp(
                          fw, vars, obs_batch, a, layer_key, logp_func
                      )
                      check(logp, expected_logp[0], rtol=0.2)
          # Test all available actions for their logp values.
          else:
              for a in [0, 1, 2, 3]:
                  count = actions.count(a)
                  expected_prob = count / num_actions
                  logp = policy.compute_log_likelihoods(
                      np.array([a]),
                      preprocessed_obs_batch,
                      prev_action_batch=np.array([prev_a])
                      if prev_a
                      else None,
                      prev_reward_batch=np.array([prev_r])
                      if prev_r
                      else None,
                      in_training=False,
                  )
  
                  if not config._enable_new_api_stack:
                      check(np.exp(logp), expected_prob, atol=0.2)
  
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 2 for calling a global function with example usage: do_test_log_likelihood(run, config, prev_a, continuous, layer_key, logp_func) without return types