docstr: Tests V-trace calculated from logits.
docstr_tok: ["Tests","V-trace","calculated","from","logits","."]
code: def testVTraceFromLogits(self, 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.
      placeholders = {
          # T, B, NUM_ACTIONS
          "behaviour_policy_logits": tf.placeholder(
              dtype=tf.float32, shape=[None, None, None]
          ),
          # T, B, NUM_ACTIONS
          "target_policy_logits": tf.placeholder(
              dtype=tf.float32, shape=[None, None, None]
          ),
          "actions": tf.placeholder(dtype=tf.int32, shape=[None, None]),
          "discounts": tf.placeholder(
              dtype=tf.float32, shape=[None, None]
          ),
          "rewards": tf.placeholder(
              dtype=tf.float32, shape=[None, None]
          ),
          "values": tf.placeholder(
              dtype=tf.float32, shape=[None, None]
          ),
          "bootstrap_value": tf.placeholder(
              dtype=tf.float32, shape=[None]
          ),
      }
  
      from_logits_output = vtrace_ops.vtrace_from_logits(
          clip_rho_threshold=clip_rho_threshold,
          clip_pg_rho_threshold=clip_pg_rho_threshold,
          **placeholders,
      )
  
      target_log_probs = vtrace_ops.log_probs_from_logits_and_actions(
          placeholders["target_policy_logits"], placeholders["actions"]
      )
      behaviour_log_probs = (
          vtrace_ops.log_probs_from_logits_and_actions(
              placeholders["behaviour_policy_logits"],
              placeholders["actions"],
          )
      )
      log_rhos = target_log_probs - behaviour_log_probs
      ground_truth = (log_rhos, behaviour_log_probs, target_log_probs)
  
      values = {
          "behaviour_policy_logits": _shaped_arange(
              seq_len, batch_size, num_actions
          ),
          "target_policy_logits": _shaped_arange(
              seq_len, batch_size, num_actions
          ),
          "actions": np.random.randint(
              0, num_actions - 1, size=(seq_len, batch_size)
          ),
          "discounts": 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)
              ]
          ),
          "rewards": _shaped_arange(seq_len, batch_size),
          "values": _shaped_arange(seq_len, batch_size) / batch_size,
          "bootstrap_value": _shaped_arange(batch_size) + 1.0,  # B
      }
  
      feed_dict = {placeholders[k]: v for k, v in values.items()}
      with self.test_session() as session:
          from_logits_output_v = session.run(
              from_logits_output, feed_dict=feed_dict
          )
          (
              ground_truth_log_rhos,
              ground_truth_behaviour_action_log_probs,
              ground_truth_target_action_log_probs,
          ) = session.run(ground_truth, feed_dict=feed_dict)
  
      # Calculate V-trace using the ground truth logits.
      from_iw = vtrace_ops.vtrace_from_importance_weights(
          log_rhos=ground_truth_log_rhos,
          discounts=values["discounts"],
          rewards=values["rewards"],
          values=values["values"],
          bootstrap_value=values["bootstrap_value"],
          clip_rho_threshold=clip_rho_threshold,
          clip_pg_rho_threshold=clip_pg_rho_threshold,
      )
  
      with self.test_session() as session:
          from_iw_v = session.run(from_iw)
  
      self.assertAllClose(from_iw_v.vs, from_logits_output_v.vs)
      self.assertAllClose(
          from_iw_v.pg_advantages, from_logits_output_v.pg_advantages
      )
      self.assertAllClose(
          ground_truth_behaviour_action_log_probs,
          from_logits_output_v.behaviour_action_log_probs,
      )
      self.assertAllClose(
          ground_truth_target_action_log_probs,
          from_logits_output_v.target_action_log_probs,
      )
      self.assertAllClose(
          ground_truth_log_rhos, from_logits_output_v.log_rhos
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  from trfl import vtrace_ops
label_desc: Use image node_id 2 for calling the VtraceTest obj's underlying member method code with example usage: obj.testVTraceFromLogits(batch_size) without return types