class_docstr: Train and evaluate DQN on Atari.
class_docstr_tok: ["Train","and","evaluate","DQN","on","Atari","."]
docstr: Run a single step (or 2 steps on life loss) in the environment.
docstr_tok: ["Run","a","single","step","(","or","2","steps","on","life","loss",")","in","the","environment","."]
code: def _collect_step(self, time_step, metric_observers, train=False):
      """Run a single step (or 2 steps on life loss) in the environment."""
      if train:
          policy = self._collect_policy
      else:
          policy = self._eval_policy
  
      with self._action_timer:
          action_step = policy.action(time_step)
      with self._step_timer:
          next_time_step = self._env.step(action_step.action)
          traj = trajectory.from_transition(
              time_step, action_step, next_time_step
          )
  
      if next_time_step.is_last() and not self.game_over():
          traj = traj._replace(
              discount=np.array([1.0], dtype=np.float32)
          )
  
      if train:
          self._store_to_rb(traj)
  
      # When AtariPreprocessing.terminal_on_life_loss is True, we receive LAST
      # time_steps when lives are lost but the game is not over. In this mode, the
      # replay buffer and agent's policy must see the life loss as a LAST step
      # and the subsequent step as a FIRST step. However, we do not want to
      # actually terminate the episode and metrics should be computed as if all
      # steps were MID steps, since life loss is not actually a terminal event
      # (it is mostly a trick to make it easier to propagate rewards backwards by
      # shortening episode durations from the agent's perspective).
      if next_time_step.is_last() and not self.game_over():
          # Update metrics as if this is a mid-episode step.
          next_time_step = ts.transition(
              next_time_step.observation, next_time_step.reward
          )
          self._observe(
              metric_observers,
              trajectory.from_transition(
                  time_step, action_step, next_time_step
              ),
          )
  
          # Produce the next step as if this is the first step of an episode and
          # store to RB as such. The next_time_step will be a MID time step.
          reward = time_step.reward
          time_step = ts.restart(next_time_step.observation)
          with self._action_timer:
              action_step = policy.action(time_step)
          with self._step_timer:
              next_time_step = self._env.step(action_step.action)
          if train:
              self._store_to_rb(
                  trajectory.from_transition(
                      time_step, action_step, next_time_step
                  )
              )
  
          # Update metrics as if this is a mid-episode step.
          time_step = ts.transition(time_step.observation, reward)
          traj = trajectory.from_transition(
              time_step, action_step, next_time_step
          )
  
      self._observe(metric_observers, traj)
  
      return next_time_step
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from absl import app
  from absl import flags
  from absl import logging
  import gin
  import numpy
  import tensorflow
  from tf_agents.agents.categorical_dqn import categorical_dqn_agent
  from tf_agents.environments import batched_py_environment
  from tf_agents.environments import suite_atari
  from tf_agents.eval import metric_utils
  from tf_agents.metrics import py_metric
  from tf_agents.metrics import py_metrics
  from tf_agents.networks import categorical_q_network
  from tf_agents.networks import network
  from tf_agents.policies import epsilon_greedy_policy
  from tf_agents.policies import py_tf_policy
  from tf_agents.policies import random_py_policy
  from tf_agents.replay_buffers import py_hashed_replay_buffer
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import common
  from tf_agents.utils import timer
label_desc: Use image node_id 9 for calling the TrainEval obj's underlying member method code with example usage: obj._collect_step(time_step, metric_observers, train) and returns: next_time_step