class_docstr: A driver that runs a TF policy in a TF environment.
class_docstr_tok: ["A","driver","that","runs","a","TF","policy","in","a","TF","environment","."]
docstr: Run policy in environment given initial time_step and policy_state.
  
  Args:
    time_step: The initial time_step.
    policy_state: The initial policy_state.
  
  Returns:
    A tuple (final time_step, final policy_state).
docstr_tok: ["Run","policy","in","environment","given","initial","time_step","and","policy_state",".","Args",":","time_step",":","The","initial","time_step",".","policy_state",":","The","initial","policy_state",".","Returns",":","A","tuple","(","final","time_step",",","final","policy_state",")","."]
code: def run(  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
      self,
      time_step: ts.TimeStep,
      policy_state: types.NestedTensor = (),
  ) -> Tuple[ts.TimeStep, types.NestedTensor]:
      """Run policy in environment given initial time_step and policy_state.
  
      Args:
        time_step: The initial time_step.
        policy_state: The initial policy_state.
  
      Returns:
        A tuple (final time_step, final policy_state).
      """
      num_steps = tf.constant(0.0)
      num_episodes = tf.constant(0.0)
  
      while (
          num_steps < self._max_steps
          and num_episodes < self._max_episodes
      ):
          action_step = self.policy.action(time_step, policy_state)
          next_time_step = self.env.step(action_step.action)
  
          traj = trajectory.from_transition(
              time_step, action_step, next_time_step
          )
          for observer in self._transition_observers:
              observer((time_step, action_step, next_time_step))
          for observer in self.observers:
              observer(traj)
  
          num_episodes += tf.math.reduce_sum(
              tf.cast(traj.is_boundary(), tf.float32)
          )
          num_steps += tf.math.reduce_sum(
              tf.cast(~traj.is_boundary(), tf.float32)
          )
  
          time_step = next_time_step
          policy_state = action_step.state
  
      return time_step, policy_state
  
file_path: agents/tf_agents/drivers/tf_driver.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Any, Callable, Optional, Sequence, Tuple
  import numpy
  import tensorflow
  from tf_agents.drivers import driver
  from tf_agents.environments import tf_environment
  from tf_agents.policies import tf_policy
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 2 for calling the TFDriver obj's underlying member method code with example usage: obj.run(time_step, policy_state) and returns: time_step,  policy_state