class_docstr: A stateless agent interface.
class_docstr_tok: ["A","stateless","agent","interface","."]
docstr: For a given single-step, unbatched timestep, output the chosen action.
docstr_tok: ["For","a","given","single-step",",","unbatched","timestep",",","output","the","chosen","action","."]
code: def step(
      self,
      rng_key,
      params: hk.Params,
      timestep: dm_env.TimeStep,
      state: Nest,
  ) -> tuple[AgentOutput, Nest]:
      """For a given single-step, unbatched timestep, output the chosen action."""
      # Pad timestep, state to be [T, B, ...] and [B, ...] respectively.
      timestep = jax.tree_util.tree_map(
          lambda t: t[None, None, ...], timestep
      )
      state = jax.tree_util.tree_map(lambda t: t[None, ...], state)
  
      net_out, next_state = self._apply_fn(params, timestep, state)
      # Remove the padding from above.
      net_out = jax.tree_util.tree_map(
          lambda t: jnp.squeeze(t, axis=(0, 1)), net_out
      )
      next_state = jax.tree_util.tree_map(
          lambda t: jnp.squeeze(t, axis=0), next_state
      )
      # Sample an action and return.
      action = hk.multinomial(
          rng_key, net_out.policy_logits, num_samples=1
      )
      action = jnp.squeeze(action, axis=-1)
      return (
          AgentOutput(net_out.policy_logits, net_out.value, action),
          next_state,
      )
  
file_path: dm-haiku/examples/impala/agent.py
imports: import collections
  import functools
  from typing import Any, Callable, Optional
  import dm_env
  import haiku
  from examples.impala import util
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 4 for calling the Agent obj's underlying member method code with example usage: obj.step(rng_key, params, timestep, state) and returns: AgentOutput,  next_state