code: def train_policy_network(self, dataset, epochs):
      loss = []
      for j in range(epochs):
          loss = {
              "total_loss": [],
              "policy_losses": [],
              "unclipped_grads": [],
              "fetch_result": [],
          }
          # shuffle(dataset)
          for i in range(
              len(dataset)
              // self.ap.network_wrappers["actor"].batch_size
          ):
              batch = Batch(
                  dataset[
                      i
                      * self.ap.network_wrappers["actor"].batch_size : (
                          i + 1
                      )
                      * self.ap.network_wrappers["actor"].batch_size
                  ]
              )
  
              network_keys = self.ap.network_wrappers[
                  "actor"
              ].input_embedders_parameters.keys()
  
              advantages = batch.info("advantage")
              actions = batch.actions()
              if (
                  not isinstance(
                      self.spaces.action, DiscreteActionSpace
                  )
                  and len(actions.shape) == 1
              ):
                  actions = np.expand_dims(actions, -1)
  
              # get old policy probabilities and distribution
              old_policy = force_list(
                  self.networks["actor"].target_network.predict(
                      batch.states(network_keys)
                  )
              )
  
              # calculate gradients and apply on both the local policy network and on the global policy network
              fetches = [
                  self.networks["actor"]
                  .online_network.output_heads[0]
                  .kl_divergence,
                  self.networks["actor"]
                  .online_network.output_heads[0]
                  .entropy,
              ]
  
              inputs = copy.copy(batch.states(network_keys))
              inputs["output_0_0"] = actions
  
              # old_policy_distribution needs to be represented as a list, because in the event of discrete controls,
              # it has just a mean. otherwise, it has both a mean and standard deviation
              for input_index, input in enumerate(old_policy):
                  inputs["output_0_{}".format(input_index + 1)] = input
  
              (
                  total_loss,
                  policy_losses,
                  unclipped_grads,
                  fetch_result,
              ) = self.networks[
                  "actor"
              ].online_network.accumulate_gradients(
                  inputs, [advantages], additional_fetches=fetches
              )
  
              self.networks["actor"].apply_gradients_to_online_network()
              if isinstance(
                  self.ap.task_parameters, DistributedTaskParameters
              ):
                  self.networks[
                      "actor"
                  ].apply_gradients_to_global_network()
  
              self.networks[
                  "actor"
              ].online_network.reset_accumulated_gradients()
  
              loss["total_loss"].append(total_loss)
              loss["policy_losses"].append(policy_losses)
              loss["unclipped_grads"].append(unclipped_grads)
              loss["fetch_result"].append(fetch_result)
  
              self.unclipped_grads.add_sample(unclipped_grads)
  
          for key in loss.keys():
              loss[key] = np.mean(loss[key], 0)
  
          if (
              self.ap.network_wrappers[
                  "critic"
              ].learning_rate_decay_rate
              != 0
          ):
              curr_learning_rate = self.networks[
                  "critic"
              ].online_network.get_variable_value(self.ap.learning_rate)
              self.curr_learning_rate.add_sample(curr_learning_rate)
          else:
              curr_learning_rate = self.ap.network_wrappers[
                  "critic"
              ].learning_rate
  
          # log training parameters
          screen.log_dict(
              OrderedDict(
                  [
                      ("Surrogate loss", loss["policy_losses"][0]),
                      ("KL divergence", loss["fetch_result"][0]),
                      ("Entropy", loss["fetch_result"][1]),
                      ("training epoch", j),
                      ("learning_rate", curr_learning_rate),
                  ]
              ),
              prefix="Policy training",
          )
  
      self.total_kl_divergence_during_training_process = loss[
          "fetch_result"
      ][0]
      self.entropy.add_sample(loss["fetch_result"][1])
      self.kl_divergence.add_sample(loss["fetch_result"][0])
      return loss["total_loss"]
  
imports: import copy
  from collections import OrderedDict
  from typing import Union
  import numpy
  from rl_coach.agents.actor_critic_agent import ActorCriticAgent
  from rl_coach.agents.policy_optimization_agent import PolicyGradientRescaler
  from rl_coach.architectures.embedder_parameters import InputEmbedderParameters
  from rl_coach.architectures.head_parameters import PPOHeadParameters, VHeadParameters
  from rl_coach.architectures.middleware_parameters import FCMiddlewareParameters
  from rl_coach.base_parameters import AlgorithmParameters, NetworkParameters, AgentParameters, DistributedTaskParameters
  from rl_coach.core_types import EnvironmentSteps, Batch
  from rl_coach.exploration_policies.additive_noise import AdditiveNoiseParameters
  from rl_coach.exploration_policies.categorical import CategoricalParameters
  from rl_coach.logger import screen
  from rl_coach.memories.episodic.episodic_experience_replay import EpisodicExperienceReplayParameters
  from rl_coach.spaces import DiscreteActionSpace, BoxActionSpace
  from rl_coach.utils import force_list
label_desc: Use image node_id 6 for calling the PPOAgent obj's underlying member method code with example usage: obj.train_policy_network(dataset, epochs) and returns: loss