code: def test_step(self):
      bt = trajectory.BatchTrajectory(batch_size=self.BATCH_SIZE)
  
      indices = np.arange(self.BATCH_SIZE)
      (
          observations,
          _,
          _,
          _,
      ) = self.get_random_observations_rewards_actions_dones()
  
      # Have to call reset first.
      bt.reset(indices, observations)
  
      # Create some fake data for calling step.
      (
          new_observations,
          raw_rewards,
          actions,
          dones,
      ) = self.get_random_observations_rewards_actions_dones()
      processed_rewards = raw_rewards.astype(np.int64)
  
      # Force mark the first one as done anyways, so that there is something to
      # test.
      dones[0] = True
  
      num_done = sum(dones)
      self.assertLessEqual(1, num_done)  # i.e. num_done is atleast 1.
  
      num_not_done = len(dones) - num_done
  
      # Finally call step.
      bt.step(
          new_observations,
          raw_rewards,
          processed_rewards,
          dones,
          actions,
      )
  
      # Expect to see `num_done` number of completed trajectories.
      self.assertEqual(num_done, bt.num_completed_trajectories)
  
      # Expect to see that the rest are marked as active.
      num_active = sum(t.is_active for t in bt.trajectories)
      self.assertEqual(num_not_done, num_active)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  import numpy
  from tensor2tensor.envs import time_step
  from tensor2tensor.envs import trajectory
  import tensorflow.compat.v1
  from tensorflow.compat.v1.io import gfile
label_desc: Use image node_id 7 for calling the BatchTrajectoryTest obj's underlying member method code with example usage: obj.test_step() without return types