code: def loss_fn(
      params: flax.core.FrozenDict,
      apply_fn: Callable[..., Any],
      minibatch: tuple,
      clip_param: float,
      vf_coeff: float,
      entropy_coeff: float,
  ):
      """Evaluate the loss function.
  
      Compute loss as a sum of three components: the negative of the PPO clipped
      surrogate objective, the value function loss and the negative of the entropy
      bonus.
  
      Args:
        params: the parameters of the actor-critic model
        apply_fn: the actor-critic model's apply function
        minibatch: tuple of five elements forming one experience batch:
                   states: shape (batch_size, 84, 84, 4)
                   actions: shape (batch_size, 84, 84, 4)
                   old_log_probs: shape (batch_size,)
                   returns: shape (batch_size,)
                   advantages: shape (batch_size,)
        clip_param: the PPO clipping parameter used to clamp ratios in loss function
        vf_coeff: weighs value function loss in total loss
        entropy_coeff: weighs entropy bonus in the total loss
  
      Returns:
        loss: the PPO loss, scalar quantity
      """
      states, actions, old_log_probs, returns, advantages = minibatch
      log_probs, values = agent.policy_action(apply_fn, params, states)
      values = values[:, 0]  # Convert shapes: (batch, 1) to (batch, ).
      probs = jnp.exp(log_probs)
  
      value_loss = jnp.mean(jnp.square(returns - values), axis=0)
  
      entropy = jnp.sum(-probs * log_probs, axis=1).mean()
  
      log_probs_act_taken = jax.vmap(lambda lp, a: lp[a])(
          log_probs, actions
      )
      ratios = jnp.exp(log_probs_act_taken - old_log_probs)
      # Advantage normalization (following the OpenAI baselines).
      advantages = (advantages - advantages.mean()) / (
          advantages.std() + 1e-8
      )
      pg_loss = ratios * advantages
      clipped_loss = advantages * jax.lax.clamp(
          1.0 - clip_param, ratios, 1.0 + clip_param
      )
      ppo_loss = -jnp.mean(jnp.minimum(pg_loss, clipped_loss), axis=0)
  
      return ppo_loss + vf_coeff * value_loss - entropy_coeff * entropy
  
imports: import functools
  from typing import Any, Callable
  from absl import logging
  import flax
  from flax import linen
  import agent
  import models
  import test_episodes
  from flax.metrics import tensorboard
  from flax.training import checkpoints
  from flax.training import train_state
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  import optax
label_desc: Use image node_id 2 for calling a global function with example usage: loss_fn(params, apply_fn, minibatch, clip_param, vf_coeff, entropy_coeff) and returns: unknown