code: def _log_probs_from_logits_and_actions(batch_size):
      """Tests log_probs_from_logits_and_actions."""
      seq_len = 7
      num_actions = 3
  
      action_space = gym.spaces.Dict(
          {"action1": gym.spaces.Discrete(num_actions)}
      )
  
      policy_logits = convert_to_torch(
          _shaped_arange(seq_len, batch_size, num_actions) + 10,
          cast=None,
          device=None,
          in_place="try",
      )
      actions = convert_to_torch(
          np.random.randint(
              0, num_actions, size=(seq_len, batch_size), dtype=np.int32
          ),
          cast=None,
          device=None,
          in_place="try",
      )
  
      distribution_mapper = DistributionMapper(
          action_space=action_space, distribution_mapper_config={}
      )
  
      (
          action_log_probs_tensor,
          _,
      ) = impala_vtrace.log_probs_from_logits_and_actions_and_spaces(
          policy_logits=[{"action1": policy_logits}],
          actions=[{"action1": actions}],
          distribution_mapper=distribution_mapper,
      )
      action_log_probs_tensor = action_log_probs_tensor[0]["action1"]
      # Ground Truth
      # Using broadcasting to create a mask that indexes action logits
      action_index_mask = np.array(actions[..., None]) == np.arange(
          num_actions
      )
  
      def index_with_mask(array, mask):
          return array[mask].reshape(*array.shape[:-1])
  
      # Note: Normally log(softmax) is not a good idea because it's not
      # numerically stable. However, in this test we have well-behaved values.
      ground_truth_v = index_with_mask(
          np.log(_softmax(np.array(policy_logits))), action_index_mask
      )
  
      assert np.allclose(ground_truth_v, action_log_probs_tensor)
  
file_path: maze/maze/test/train/impala/test_vtrace.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gym
  import numpy
  from maze.distributions.distribution_mapper import DistributionMapper
  from maze.perception.perception_utils import convert_to_torch
  from maze.train.trainers.impala import impala_vtrace
label_desc: Use image node_id 4 for calling a global function with example usage: _log_probs_from_logits_and_actions(batch_size) and returns: array