class_docstr: A stateless agent interface.
class_docstr_tok: ["A","stateless","agent","interface","."]
docstr: Constructs an Agent object.
  
  Args:
    num_actions: Number of possible actions for the agent. Assumes a flat,
      discrete, 0-indexed action space.
    obs_spec: The observation spec of the environment.
    net_factory: A function from num_actions to a Haiku module representing
      the agent. This module should have an initial_state() function and an
      unroll function.
docstr_tok: ["Constructs","an","Agent","object",".","Args",":","num_actions",":","Number","of","possible","actions","for","the","agent",".","Assumes","a","flat",",","discrete",",","0-indexed","action","space",".","obs_spec",":","The","observation","spec","of","the","environment",".","net_factory",":","A","function","from","num_actions","to","a","Haiku","module","representing","the","agent",".","This","module","should","have","an","initial_state","(",")","function","and","an","unroll","function","."]
code: def __init__(
      self, num_actions: int, obs_spec: Nest, net_factory: NetFactory
  ):
      """Constructs an Agent object.
  
      Args:
        num_actions: Number of possible actions for the agent. Assumes a flat,
          discrete, 0-indexed action space.
        obs_spec: The observation spec of the environment.
        net_factory: A function from num_actions to a Haiku module representing
          the agent. This module should have an initial_state() function and an
          unroll function.
      """
      self._obs_spec = obs_spec
      net_factory = functools.partial(net_factory, num_actions)
      # Instantiate two hk.transforms() - one for getting the initial state of the
      # agent, another for actually initializing and running the agent.
      _, self._initial_state_apply_fn = hk.without_apply_rng(
          hk.transform(
              lambda batch_size: net_factory().initial_state(batch_size)
          )
      )
  
      self._init_fn, self._apply_fn = hk.without_apply_rng(
          hk.transform(
              lambda obs, state: net_factory().unroll(obs, 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 1 to create a new Agent object  with example: obj = Agent(num_actions, obs_spec, net_factory)