class_docstr: A simple, feed-forward agent.
class_docstr_tok: ["A","simple",",","feed-forward","agent","."]
docstr: Computes a loss of trajs wrt params.
docstr_tok: ["Computes","a","loss","of","trajs","wrt","params","."]
code: def loss(self, params: hk.Params, trajs: Transition) -> jax.Array:
      """Computes a loss of trajs wrt params."""
      # Re-run the agent over the trajectories.
      # Due to https://github.com/google/jax/issues/1459, we use hk.BatchApply
      # instead of vmap.
      # BatchApply turns the input tensors from [T, B, ...] into [T*B, ...].
      # We `functools.partial` params in so it does not get transformed.
      net_curried = hk.BatchApply(functools.partial(self._net, params))
      learner_logits, baseline_with_bootstrap = net_curried(
          trajs.timestep
      )
  
      # Separate the bootstrap from the value estimates.
      baseline = baseline_with_bootstrap[:-1]
      baseline_tp1 = baseline_with_bootstrap[1:]
  
      # Remove bootstrap timestep from non-observations.
      _, actions, behavior_logits = jax.tree_util.tree_map(
          lambda t: t[:-1], trajs
      )
      learner_logits = learner_logits[:-1]
  
      # Shift step_type/reward/discount back by one, so that actions match the
      # timesteps caused by the action.
      timestep = jax.tree_util.tree_map(lambda t: t[1:], trajs.timestep)
      discount = timestep.discount * self._discount
      # The step is uninteresting if we transitioned LAST -> FIRST.
      mask = jnp.not_equal(
          timestep.step_type, int(dm_env.StepType.FIRST)
      )
      mask = mask.astype(jnp.float32)
  
      # Compute v-trace returns.
      vtrace_td_error_and_advantage = jax.vmap(
          rlax.vtrace_td_error_and_advantage, in_axes=1, out_axes=1
      )
      rhos = rlax.categorical_importance_sampling_ratios(
          learner_logits, behavior_logits, actions
      )
      vtrace_returns = vtrace_td_error_and_advantage(
          baseline, baseline_tp1, timestep.reward, discount, rhos
      )
  
      # Note that we use mean here, rather than sum as in canonical IMPALA.
      # Compute policy gradient loss.
      pg_advantage = jax.lax.stop_gradient(vtrace_returns.pg_advantage)
      tb_pg_loss_fn = jax.vmap(
          rlax.policy_gradient_loss, in_axes=1, out_axes=0
      )
      pg_loss = tb_pg_loss_fn(
          learner_logits, actions, pg_advantage, mask
      )
      pg_loss = jnp.mean(pg_loss)
  
      # Baseline loss.
      bl_loss = 0.5 * jnp.mean(jnp.square(vtrace_returns.errors) * mask)
  
      # Entropy regularization.
      ent_loss_fn = jax.vmap(rlax.entropy_loss, in_axes=1, out_axes=0)
      ent_loss = ent_loss_fn(learner_logits, mask)
      ent_loss = jnp.mean(ent_loss)
  
      total_loss = pg_loss + 0.5 * bl_loss + 0.01 * ent_loss
      return total_loss
  
imports: import functools
  import queue
  import threading
  from typing import Any, Callable, NamedTuple
  from absl import app
  from absl import logging
  from bsuite.environments import catch
  import dm_env
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  import rlax
label_desc: Use image node_id 3 for calling the Agent obj's underlying member method code with example usage: obj.loss(params, trajs) and returns: total_loss