class_docstr: A simple, feed-forward agent.
class_docstr_tok: ["A","simple",",","feed-forward","agent","."]
docstr: Steps on a single observation.
docstr_tok: ["Steps","on","a","single","observation","."]
code: def step(
      self,
      params: hk.Params,
      rng: jax.Array,
      timestep: dm_env.TimeStep,
  ) -> tuple[jax.Array, jax.Array]:
      """Steps on a single observation."""
      timestep = jax.tree_util.tree_map(
          lambda t: jnp.expand_dims(t, 0), timestep
      )
      logits, _ = self._net(params, timestep)
      logits = jnp.squeeze(logits, axis=0)
      action = hk.multinomial(rng, logits, num_samples=1)
      action = jnp.squeeze(action, axis=-1)
      return action, logits
  
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 Agent obj's underlying member method code with example usage: obj.step(params, rng, timestep) and returns: action,  logits