class_docstr: A driver that takes N episodes in an environment using a tf.while_loop.
  
  The while loop will run num_episodes in the environment, counting transitions
  that result in ending an episode.
  
  As environments run batched time_episodes, the counters for all batch elements
  are summed, and execution stops when the total exceeds num_episodes.
  
  This termination condition can be overridden in subclasses by implementing the
  self._loop_condition_fn() method.
class_docstr_tok: ["A","driver","that","takes","N","episodes","in","an","environment","using","a","tf.while_loop",".","The","while","loop","will","run","num_episodes","in","the","environment",",","counting","transitions","that","result","in","ending","an","episode",".","As","environments","run","batched","time_episodes",",","the","counters","for","all","batch","elements","are","summed",",","and","execution","stops","when","the","total","exceeds","num_episodes",".","This","termination","condition","can","be","overridden","in","subclasses","by","implementing","the","self._loop_condition_fn","(",")","method","."]
docstr: Returns a function with the driver's loop body ops.
docstr_tok: ["Returns","a","function","with","the","driver","'s","loop","body","ops","."]
code: def _loop_body_fn(self):
      """Returns a function with the driver's loop body ops."""
  
      def loop_body(counter, time_step, policy_state):
          """Runs a step in environment.
  
          While loop will call multiple times.
  
          Args:
            counter: Episode counters per batch index. Shape [batch_size].
            time_step: TimeStep tuple with elements shape [batch_size, ...].
            policy_state: Poicy state tensor shape [batch_size, policy_state_dim].
              Pass empty tuple for non-recurrent policies.
  
          Returns:
            loop_vars for next iteration of tf.while_loop.
          """
          action_step = self.policy.action(time_step, policy_state)
  
          # TODO(b/134487572): TF2 while_loop seems to either ignore
          # parallel_iterations or doesn't properly propagate control dependencies
          # from one step to the next. Without this dep, self.env.step() is called
          # in parallel.
          with tf.control_dependencies(tf.nest.flatten([time_step])):
              next_time_step = self.env.step(action_step.action)
  
          policy_state = action_step.state
  
          if self._is_bandit_env:
              # For Bandits we create episodes of length 1.
              # Since the `next_time_step` is always of type LAST we need to replace
              # the step type of the current `time_step` to FIRST.
              batch_size = tf.shape(input=time_step.discount)
              time_step = time_step._replace(
                  step_type=tf.fill(batch_size, ts.StepType.FIRST)
              )
  
          traj = trajectory.from_transition(
              time_step, action_step, next_time_step
          )
          observer_ops = [
              observer(traj) for observer in self._observers
          ]
          transition_observer_ops = [
              observer((time_step, action_step, next_time_step))
              for observer in self._transition_observers
          ]
          with tf.control_dependencies(
              [tf.group(observer_ops + transition_observer_ops)]
          ):
              (
                  time_step,
                  next_time_step,
                  policy_state,
              ) = tf.nest.map_structure(
                  tf.identity, (time_step, next_time_step, policy_state)
              )
  
          # While loop counter is only incremented for episode reset episodes.
          # For Bandits, this is every trajectory, for MDPs, this is at boundaries.
          if self._is_bandit_env:
              counter += tf.ones(batch_size, dtype=tf.int32)
          else:
              counter += tf.cast(traj.is_boundary(), dtype=tf.int32)
  
          return [counter, next_time_step, policy_state]
  
      return loop_body
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gin
  import tensorflow
  from tf_agents.bandits.environments import bandit_py_environment
  from tf_agents.bandits.environments import bandit_tf_environment
  from tf_agents.drivers import driver
  from tf_agents.environments import tf_py_environment
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 3 for calling the DynamicEpisodeDriver obj's underlying member method code with example usage: obj._loop_body_fn() and returns: loop_body, list