class_docstr: Manages state and performs updates for IMPALA learner.
docstr: The actual update function.
docstr_tok: ["The","actual","update","function","."]
code: def update(self, params, opt_state, batch: util.Transition):
      """The actual update function."""
      (_, logs), grads = jax.value_and_grad(self._loss, has_aux=True)(
          params, batch
      )
  
      grad_norm_unclipped = optimizers.l2_norm(grads)
      updates, updated_opt_state = self._opt.update(grads, opt_state)
      params = optax.apply_updates(params, updates)
      weight_norm = optimizers.l2_norm(params)
      logs.update(
          {
              "grad_norm_unclipped": grad_norm_unclipped,
              "weight_norm": weight_norm,
          }
      )
      return params, updated_opt_state, 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 3 for calling the Learner obj's underlying member method code with example usage: obj.update(params, opt_state, batch) and returns: params,  updated_opt_state,  logs