class_docstr: Slim wrapper around an agent/optimizer pair.
class_docstr_tok: ["Slim","wrapper","around","an","agent\/optimizer","pair","."]
code: def update(
      self,
      params: hk.Params,
      opt_state: optax.OptState,
      trajs: Transition,
  ) -> tuple[hk.Params, optax.OptState]:
      g = jax.grad(self._agent.loss)(params, trajs)
      updates, new_opt_state = self._opt_update(g, opt_state)
      return optax.apply_updates(params, updates), new_opt_state
  
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 2 for calling the Learner obj's underlying member method code with example usage: obj.update(params, opt_state, trajs) and returns: optax,  new_opt_state