class_docstr: A compact implementation of agents for Atari 100k.
docstr: Selects an action, and optionally records a transition and trains.
  
  If `reward` or `observation` is None, the agent's state will _not_ be
  updated and nothing will be written to the buffer. The user must call
  `log_transition` themselves in this case.
  
  Args:
    reward: Optional reward to log.
    observation: Optional observation to log. Must call `log_transition` later
      if not passed here.
  
  Returns:
    Selected action.
docstr_tok: ["Selects","an","action",",","and","optionally","records","a","transition","and","trains",".","If","`","reward","`","or","`","observation","`","is","None",",","the","agent","'s","state","will","_not_","be","updated","and","nothing","will","be","written","to","the","buffer",".","The","user","must","call","`","log_transition","`","themselves","in","this","case",".","Args",":","reward",":","Optional","reward","to","log",".","observation",":","Optional","observation","to","log",".","Must","call","`","log_transition","`","later","if","not","passed","here",".","Returns",":","Selected","action","."]
code: def step(self, reward=None, observation=None):
      """Selects an action, and optionally records a transition and trains.
  
      If `reward` or `observation` is None, the agent's state will _not_ be
      updated and nothing will be written to the buffer. The user must call
      `log_transition` themselves in this case.
  
      Args:
        reward: Optional reward to log.
        observation: Optional observation to log. Must call `log_transition` later
          if not passed here.
  
      Returns:
        Selected action.
      """
      if reward is not None and observation is not None:
          self._last_observation = self._observation
          self._record_observation(observation)
          if not self.eval_mode:
              self._store_transition(
                  self._last_observation, self.action, reward, False
              )
  
      if not self.eval_mode:
          self._train_step()
  
      state = self.preprocess_fn(self.state)
      self._rng, action = select_action(
          self.network_def,
          self.online_params,
          state,
          self._rng,
          self.num_actions,
          self.eval_mode,
          self.epsilon_eval,
          self.epsilon_train,
          self.epsilon_decay_period,
          self.training_steps,
          self.min_replay_history,
          self.epsilon_fn,
          self._support,
      )
      self.action = onp.asarray(action)
      return self.action
  
imports: import copy
  import functools
  from absl import logging
  from dopamine.jax import networks
  from dopamine.jax.agents.full_rainbow import full_rainbow_agent
  import gin
  import jax
  import jax.numpy
  import numpy
  import tensorflow
label_desc: Use image node_id 3 for calling the Atari100kRainbowAgent obj's underlying member method code with example usage: obj.step(reward, observation) and returns: self