code: def _vtrace(batch_size):
      """Tests V-trace against ground truth data calculated in python."""
      seq_len = 5
  
      # Create log_rhos such that rho will span from near-zero to above the
      # clipping thresholds. In particular, calculate log_rhos in [-2.5, 2.5),
      # so that rho is in approx [0.08, 12.2).
      log_rhos = _shaped_arange(seq_len, batch_size) / (
          batch_size * seq_len
      )
      log_rhos = 5 * (log_rhos - 0.5)  # [0.0, 1.0) -> [-2.5, 2.5).
      values = {
          "log_rhos": log_rhos,
          # T, B where B_i: [0.9 / (i+1)] * T
          "discounts": np.array(
              [
                  [0.9 / (b + 1) for b in range(batch_size)]
                  for _ in range(seq_len)
              ]
          ),
          "rewards": _shaped_arange(seq_len, batch_size),
          "values": _shaped_arange(seq_len, batch_size) / batch_size,
          "bootstrap_value": _shaped_arange(batch_size) + 1.0,
          "clip_rho_threshold": 3.7,
          "clip_pg_rho_threshold": 2.2,
      }
  
      output = impala_vtrace.from_importance_weights(**values)
  
      ground_truth_v = _ground_truth_calculation(**values)
      for a, b in zip(ground_truth_v, output):
          assert np.allclose(a, b)
  
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 6 for calling a global function with example usage: _vtrace(batch_size) without return types