class_docstr: Manages state and performs updates for IMPALA learner.
docstr: Compute vtrace-based actor-critic loss.
docstr_tok: ["Compute","vtrace-based","actor-critic","loss","."]
code: def _loss(
      self,
      theta: hk.Params,
      trajectories: util.Transition,
  ) -> tuple[jax.Array, dict[str, jax.Array]]:
      """Compute vtrace-based actor-critic loss."""
      initial_state = jax.tree_util.tree_map(
          lambda t: t[0], trajectories.agent_state
      )
      learner_outputs = self._agent.unroll(
          theta, trajectories.timestep, initial_state
      )
      v_t = learner_outputs.values[1:]
      # Remove bootstrap timestep from non-timesteps.
      _, actor_out, _ = jax.tree_util.tree_map(
          lambda t: t[:-1], trajectories
      )
      learner_outputs = jax.tree_util.tree_map(
          lambda t: t[:-1], learner_outputs
      )
      v_tm1 = learner_outputs.values
  
      # Get the discount, reward, step_type from the *next* timestep.
      timestep = jax.tree_util.tree_map(
          lambda t: t[1:], trajectories.timestep
      )
      discounts = timestep.discount * self._discount_factor
      rewards = timestep.reward
      if self._max_abs_reward > 0:
          rewards = jnp.clip(
              rewards, -self._max_abs_reward, self._max_abs_reward
          )
  
      # The step is uninteresting if we transitioned LAST -> FIRST.
      # timestep corresponds to the *next* time step, so we filter for FIRST.
      mask = jnp.not_equal(
          timestep.step_type, int(dm_env.StepType.FIRST)
      )
      mask = mask.astype(jnp.float32)
  
      rhos = rlax.categorical_importance_sampling_ratios(
          learner_outputs.policy_logits,
          actor_out.policy_logits,
          actor_out.action,
      )
      # vmap vtrace_td_error_and_advantage to take/return [T, B, ...].
      vtrace_td_error_and_advantage = jax.vmap(
          rlax.vtrace_td_error_and_advantage, in_axes=1, out_axes=1
      )
  
      vtrace_returns = vtrace_td_error_and_advantage(
          v_tm1, v_t, rewards, discounts, rhos
      )
      pg_advs = vtrace_returns.pg_advantage
      pg_loss = policy_gradient_loss(
          learner_outputs.policy_logits, actor_out.action, pg_advs, mask
      )
  
      baseline_loss = 0.5 * jnp.sum(
          jnp.square(vtrace_returns.errors) * mask
      )
      ent_loss = entropy_loss(learner_outputs.policy_logits, mask)
  
      total_loss = pg_loss
      total_loss += 0.5 * baseline_loss
      total_loss += 0.01 * ent_loss
  
      logs = {}
      logs["PG_loss"] = pg_loss
      logs["baseline_loss"] = baseline_loss
      logs["entropy_loss"] = ent_loss
      logs["total_loss"] = total_loss
      return total_loss, logs
  
imports: import functools
  import itertools
  import queue
  import threading
  import warnings
  import dm_env
  import haiku
  from examples.impala import agent
  from examples.impala import util
  import jax
  from jax.example_libraries import optimizers
  import jax.numpy
  import numpy
  import optax
  import rlax
label_desc: Use image node_id 2 for calling the Learner obj's underlying member method code with example usage: obj._loss(theta, trajectories) and returns: total_loss,  logs