code: def _vtrace_from_logits(batch_size):
      """Tests V-trace calculated from logits."""
      seq_len = 5
      num_actions = 3
      clip_rho_threshold = None  # No clipping.
      clip_pg_rho_threshold = None  # No clipping.
  
      # Intentionally leaving shapes unspecified to test if V-trace can
      # deal with that.
  
      values = {
          "behaviour_policy_logits": [
              {
                  "action1": convert_to_torch(
                      _shaped_arange(seq_len, batch_size, num_actions),
                      device=None,
                      cast=None,
                      in_place="try",
                  )
              }
          ],
          "target_policy_logits": [
              {
                  "action1": convert_to_torch(
                      _shaped_arange(seq_len, batch_size, num_actions),
                      device=None,
                      cast=None,
                      in_place="try",
                  )
              }
          ],
          "actions": [
              {
                  "action1": convert_to_torch(
                      np.random.randint(
                          0, num_actions - 1, size=(seq_len, batch_size)
                      ),
                      device=None,
                      cast=None,
                      in_place="try",
                  )
              }
          ],
          "discounts": convert_to_torch(
              np.array(  # T, B where B_i: [0.9 / (i+1)] * T
                  [
                      [0.9 / (b + 1) for b in range(batch_size)]
                      for _ in range(seq_len)
                  ]
              ),
              device=None,
              cast=None,
              in_place="try",
          ),
          "rewards": convert_to_torch(
              _shaped_arange(seq_len, batch_size),
              device=None,
              cast=None,
              in_place="try",
          ),
          "values": [
              convert_to_torch(
                  _shaped_arange(seq_len, batch_size) / batch_size,
                  device=None,
                  cast=None,
                  in_place="try",
              )
          ],
          "bootstrap_value": [
              convert_to_torch(
                  _shaped_arange(batch_size) + 1.0,
                  device=None,
                  cast=None,
                  in_place="try",
              )
          ],
      }
      action_space = {
          0: gym.spaces.Dict(
              {"action1": gym.spaces.Discrete(num_actions)}
          )
      }
      # initialize distribution mapper
      distribution_mapper = DistributionMapper(
          action_space=action_space[0], distribution_mapper_config={}
      )
  
      from_logits_output = impala_vtrace.from_logits(
          clip_rho_threshold=clip_rho_threshold,
          clip_pg_rho_threshold=clip_pg_rho_threshold,
          device=None,
          distribution_mapper=distribution_mapper,
          **values,
      )
  
      (
          target_log_probs,
          _,
      ) = impala_vtrace.log_probs_from_logits_and_actions_and_spaces(
          values["target_policy_logits"],
          values["actions"],
          distribution_mapper=distribution_mapper,
      )
      (
          behaviour_log_probs,
          _,
      ) = impala_vtrace.log_probs_from_logits_and_actions_and_spaces(
          values["behaviour_policy_logits"],
          values["actions"],
          distribution_mapper=distribution_mapper,
      )
      log_rhos = impala_vtrace.get_log_rhos(
          target_log_probs, behaviour_log_probs
      )
      (
          ground_truth_log_rhos,
          ground_truth_behaviour_action_log_probs,
          ground_truth_target_action_log_probs,
      ) = (log_rhos, behaviour_log_probs, target_log_probs)
  
      # Calculate V-trace using the ground truth logits.
      from_iw = impala_vtrace.from_importance_weights(
          log_rhos=ground_truth_log_rhos[0],
          discounts=values["discounts"],
          rewards=values["rewards"],
          values=values["values"][0],
          bootstrap_value=values["bootstrap_value"][0],
          clip_rho_threshold=clip_rho_threshold,
          clip_pg_rho_threshold=clip_pg_rho_threshold,
      )
  
      assert np.allclose(from_iw.vs, from_logits_output.vs[0])
      assert np.allclose(
          from_iw.pg_advantages, from_logits_output.pg_advantages[0]
      )
      assert np.allclose(
          ground_truth_behaviour_action_log_probs[0]["action1"],
          from_logits_output.behaviour_action_log_probs[0]["action1"],
      )
      assert np.allclose(
          ground_truth_target_action_log_probs[0]["action1"],
          from_logits_output.target_action_log_probs[0]["action1"],
      )
      assert np.allclose(
          ground_truth_log_rhos[0], from_logits_output.log_rhos[0]
      )
  
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 8 for calling a global function with example usage: _vtrace_from_logits(batch_size) without return types