docstr: Tests V-trace calculated from logits.
docstr_tok: ["Tests","V-trace","calculated","from","logits","."]
code: def test_vtrace_from_logits(self):
      """Tests V-trace calculated from logits."""
      seq_len = 5
      batch_size = 15
      num_actions = 3
      clip_rho_threshold = None  # No clipping.
      clip_pg_rho_threshold = None  # No clipping.
      space = Box(-1.0, 1.0, (seq_len, batch_size, num_actions))
      action_space = Box(
          0,
          num_actions - 1,
          (
              seq_len,
              batch_size,
          ),
          dtype=np.int32,
      )
      space_w_time = Box(
          -1.0,
          1.0,
          (
              seq_len,
              batch_size,
          ),
      )
      space_only_batch = Box(-1.0, 1.0, (batch_size,))
  
      for fw, sess in framework_iterator(
          frameworks=("torch", "tf"), session=True
      ):
          vtrace = vtrace_tf if fw != "torch" else vtrace_torch
  
          if fw == "tf":
              # Intentionally leaving shapes unspecified to test if V-trace
              # can deal with that.
              inputs_ = {
                  # T, B, NUM_ACTIONS
                  "behaviour_policy_logits": tf1.placeholder(
                      dtype=tf.float32, shape=[None, None, None]
                  ),
                  # T, B, NUM_ACTIONS
                  "target_policy_logits": tf1.placeholder(
                      dtype=tf.float32, shape=[None, None, None]
                  ),
                  "actions": tf1.placeholder(
                      dtype=tf.int32, shape=[None, None]
                  ),
                  "discounts": tf1.placeholder(
                      dtype=tf.float32, shape=[None, None]
                  ),
                  "rewards": tf1.placeholder(
                      dtype=tf.float32, shape=[None, None]
                  ),
                  "values": tf1.placeholder(
                      dtype=tf.float32, shape=[None, None]
                  ),
                  "bootstrap_value": tf1.placeholder(
                      dtype=tf.float32, shape=[None]
                  ),
              }
          else:
              inputs_ = {
                  # T, B, NUM_ACTIONS
                  "behaviour_policy_logits": space.sample(),
                  # T, B, NUM_ACTIONS
                  "target_policy_logits": space.sample(),
                  "actions": action_space.sample(),
                  "discounts": space_w_time.sample(),
                  "rewards": space_w_time.sample(),
                  "values": space_w_time.sample(),
                  "bootstrap_value": space_only_batch.sample(),
              }
          from_logits_output = vtrace.from_logits(
              clip_rho_threshold=clip_rho_threshold,
              clip_pg_rho_threshold=clip_pg_rho_threshold,
              **inputs_,
          )
  
          if fw != "torch":
              target_log_probs = (
                  vtrace.log_probs_from_logits_and_actions(
                      inputs_["target_policy_logits"],
                      inputs_["actions"],
                  )
              )
              behaviour_log_probs = (
                  vtrace.log_probs_from_logits_and_actions(
                      inputs_["behaviour_policy_logits"],
                      inputs_["actions"],
                  )
              )
          else:
              target_log_probs = (
                  vtrace.log_probs_from_logits_and_actions(
                      torch.from_numpy(inputs_["target_policy_logits"]),
                      torch.from_numpy(inputs_["actions"]),
                  )
              )
              behaviour_log_probs = (
                  vtrace.log_probs_from_logits_and_actions(
                      torch.from_numpy(
                          inputs_["behaviour_policy_logits"]
                      ),
                      torch.from_numpy(inputs_["actions"]),
                  )
              )
          log_rhos = target_log_probs - behaviour_log_probs
          ground_truth = (
              log_rhos,
              behaviour_log_probs,
              target_log_probs,
          )
  
          if sess:
              values = {
                  "behaviour_policy_logits": space.sample(),
                  "target_policy_logits": space.sample(),
                  "actions": action_space.sample(),
                  "discounts": space_w_time.sample(),
                  "rewards": space_w_time.sample(),
                  "values": space_w_time.sample() / batch_size,
                  "bootstrap_value": space_only_batch.sample() + 1.0,
              }
              feed_dict = {inputs_[k]: v for k, v in values.items()}
              from_logits_output = sess.run(
                  from_logits_output, feed_dict=feed_dict
              )
              (
                  log_rhos,
                  behaviour_log_probs,
                  target_log_probs,
              ) = sess.run(ground_truth, feed_dict=feed_dict)
  
              # Calculate V-trace using the ground truth logits.
              from_iw = vtrace.from_importance_weights(
                  log_rhos=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,
              )
              from_iw = sess.run(from_iw)
          else:
              from_iw = vtrace.from_importance_weights(
                  log_rhos=log_rhos,
                  discounts=inputs_["discounts"],
                  rewards=inputs_["rewards"],
                  values=inputs_["values"],
                  bootstrap_value=inputs_["bootstrap_value"],
                  clip_rho_threshold=clip_rho_threshold,
                  clip_pg_rho_threshold=clip_pg_rho_threshold,
              )
  
          check(from_iw.vs, from_logits_output.vs)
          check(from_iw.pg_advantages, from_logits_output.pg_advantages)
          check(
              behaviour_log_probs,
              from_logits_output.behaviour_action_log_probs,
          )
          check(
              target_log_probs,
              from_logits_output.target_action_log_probs,
          )
          check(log_rhos, from_logits_output.log_rhos)
  
file_path: ray/rllib/algorithms/impala/tests/test_vtrace.py
imports: from gymnasium.spaces import Box
  import numpy
  import unittest
  from ray.rllib.algorithms.impala import vtrace_tf
  from ray.rllib.algorithms.impala import vtrace_torch
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.numpy import softmax
  from ray.rllib.utils.test_utils import check, framework_iterator
label_desc: Use image node_id 2 for calling the VtraceTest obj's underlying member method code with example usage: obj.test_vtrace_from_logits() without return types