docstr: Tests the PPO loss function math.
  
  This test is overfitted to the old ModelV2 stack (e.g.
  policy.model.trainable_variables is not callable in the new stack)
  # TODO (Kourosh) we should create a new test for the new RLModule stack.
docstr_tok: ["Tests","the","PPO","loss","function","math",".","This","test","is","overfitted","to","the","old","ModelV2","stack","(","e.g",".","policy.model.trainable_variables","is","not","callable","in","the","new","stack",")","#","TODO","(","Kourosh",")","we","should","create","a","new","test","for","the","new","RLModule","stack","."]
code: def test_ppo_loss_function(self):
      """Tests the PPO loss function math.
  
      This test is overfitted to the old ModelV2 stack (e.g.
      policy.model.trainable_variables is not callable in the new stack)
      # TODO (Kourosh) we should create a new test for the new RLModule stack.
      """
      config = (
          ppo.PPOConfig()
          .experimental(_enable_new_api_stack=False)
          .environment("CartPole-v1")
          .rollouts(
              num_rollout_workers=0,
          )
          .training(
              gamma=0.99,
              model=dict(
                  fcnet_hiddens=[10],
                  fcnet_activation="linear",
                  vf_share_layers=True,
              ),
          )
      )
  
      for fw, sess in framework_iterator(config, session=True):
          algo = config.build()
          policy = algo.get_policy()
  
          # Check no free log std var by default.
          if fw == "torch":
              matching = [
                  v
                  for (n, v) in policy.model.named_parameters()
                  if "log_std" in n
              ]
          else:
              matching = [
                  v
                  for v in policy.model.trainable_variables()
                  if "log_std" in str(v)
              ]
          assert len(matching) == 0, matching
  
          # Post-process (calculate simple (non-GAE) advantages) and attach
          # to train_batch dict.
          # A = [0.99^2 * 0.5 + 0.99 * -1.0 + 1.0, 0.99 * 0.5 - 1.0, 0.5] =
          # [0.50005, -0.505, 0.5]
          train_batch = compute_gae_for_sample_batch(
              policy, CARTPOLE_FAKE_BATCH.copy()
          )
          if fw == "torch":
              train_batch = policy._lazy_tensor_dict(train_batch)
  
          # Check Advantage values.
          check(
              train_batch[Postprocessing.VALUE_TARGETS],
              [0.50005, -0.505, 0.5],
          )
  
          # Calculate actual PPO loss.
          if fw == "tf2":
              PPOTF2Policy.loss(
                  policy, policy.model, Categorical, train_batch
              )
          elif fw == "torch":
              PPOTorchPolicy.loss(
                  policy, policy.model, policy.dist_class, train_batch
              )
  
          vars = (
              policy.model.variables()
              if fw != "torch"
              else list(policy.model.parameters())
          )
          if fw == "tf":
              vars = policy.get_session().run(vars)
          expected_shared_out = fc(
              train_batch[SampleBatch.CUR_OBS],
              vars[0 if fw != "torch" else 2],
              vars[1 if fw != "torch" else 3],
              framework=fw,
          )
          expected_logits = fc(
              expected_shared_out,
              vars[2 if fw != "torch" else 0],
              vars[3 if fw != "torch" else 1],
              framework=fw,
          )
          expected_value_outs = fc(
              expected_shared_out, vars[4], vars[5], framework=fw
          )
  
          (
              kl,
              entropy,
              pg_loss,
              vf_loss,
              overall_loss,
          ) = self._ppo_loss_helper(
              policy,
              policy.model,
              Categorical if fw != "torch" else TorchCategorical,
              train_batch,
              expected_logits,
              expected_value_outs,
              sess=sess,
          )
          if sess:
              policy_sess = policy.get_session()
              k, e, pl, v, tl = policy_sess.run(
                  [
                      policy._mean_kl_loss,
                      policy._mean_entropy,
                      policy._mean_policy_loss,
                      policy._mean_vf_loss,
                      policy._total_loss,
                  ],
                  feed_dict=policy._get_loss_inputs_dict(
                      train_batch, shuffle=False
                  ),
              )
              check(k, kl)
              check(e, entropy)
              check(pl, np.mean(-pg_loss))
              check(v, np.mean(vf_loss), decimals=4)
              check(tl, overall_loss, decimals=4)
          elif fw == "torch":
              check(policy.model.tower_stats["mean_kl_loss"], kl)
              check(policy.model.tower_stats["mean_entropy"], entropy)
              check(
                  policy.model.tower_stats["mean_policy_loss"],
                  np.mean(-pg_loss),
              )
              check(
                  policy.model.tower_stats["mean_vf_loss"],
                  np.mean(vf_loss),
                  decimals=4,
              )
              check(
                  policy.model.tower_stats["total_loss"],
                  overall_loss,
                  decimals=4,
              )
          else:
              check(policy._mean_kl_loss, kl)
              check(policy._mean_entropy, entropy)
              check(policy._mean_policy_loss, np.mean(-pg_loss))
              check(policy._mean_vf_loss, np.mean(vf_loss), decimals=4)
              check(policy._total_loss, overall_loss, decimals=4)
          algo.stop()
  
imports: import unittest
  import numpy
  import ray
  from ray.rllib.algorithms.callbacks import DefaultCallbacks
  import ray.rllib.algorithms.ppo
  from ray.rllib.algorithms.ppo.ppo_tf_policy import PPOTF2Policy
  from ray.rllib.algorithms.ppo.ppo_torch_policy import PPOTorchPolicy
  from ray.rllib.evaluation.postprocessing import compute_gae_for_sample_batch, Postprocessing
  from ray.rllib.models.tf.tf_action_dist import Categorical
  from ray.rllib.models.torch.torch_modelv2 import TorchModelV2
  from ray.rllib.models.torch.torch_action_dist import TorchCategorical
  from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch
  from ray.rllib.utils.metrics.learner_info import LEARNER_INFO, LEARNER_STATS_KEY
  from ray.rllib.utils.numpy import fc
  from ray.rllib.utils.test_utils import check, check_compute_single_action, check_off_policyness, check_train_results, framework_iterator, check_inference_w_connectors
label_desc: Use image node_id 7 for calling the TestPPO obj's underlying member method code with example usage: obj.test_ppo_loss_function() without return types