code: def test_desired_placement_of_rewards_and_actions(self):
      batch_size = 1
      bt = trajectory.BatchTrajectory(batch_size=batch_size)
  
      indices = np.arange(batch_size)
      (
          observations,
          _,
          _,
          _,
      ) = self.get_random_observations_rewards_actions_dones(
          batch_size=batch_size
      )
  
      # Have to call reset first.
      bt.reset(indices, observations)
  
      # Create some fake data for calling step.
      (
          new_observations,
          raw_rewards,
          actions,
          _,
      ) = self.get_random_observations_rewards_actions_dones(
          batch_size=batch_size
      )
      processed_rewards = raw_rewards.astype(np.int64)
      dones = np.full(batch_size, False)
  
      # Call step.
      bt.step(
          new_observations,
          raw_rewards,
          processed_rewards,
          dones,
          actions,
      )
  
      # Assert that nothing is done, since dones is False
      self.assertEqual(0, bt.num_completed_trajectories)
  
      # The only trajectory is active.
      self.assertEqual(batch_size, len(bt.trajectories))
      t = bt.trajectories[0]
      self.assertTrue(t.is_active)
      self.assertEqual(2, t.num_time_steps)
  
      ts = t.time_steps
  
      # Now assert on placements
  
      # i.e. the old observation/done is first and the new one comes later.
      self.assertAllEqual(observations[0], ts[0].observation)
      self.assertAllEqual(new_observations[0], ts[1].observation)
  
      self.assertEqual(False, ts[0].done)
      self.assertEqual(False, ts[1].done)
  
      # Similarly actions went to the first time-step.
      self.assertEqual(actions[0], ts[0].action)
      self.assertIsNone(ts[1].action)
  
      # However make sure reward went into the second time-step and not the first.
      self.assertNear(raw_rewards[0], ts[1].raw_reward, 1e-6)
      self.assertIsNone(ts[0].raw_reward)
  
      # Similarly with processed_rewards.
      self.assertEqual(processed_rewards[0], ts[1].processed_reward)
      self.assertIsNone(ts[0].processed_reward)
  
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 8 for calling the BatchTrajectoryTest obj's underlying member method code with example usage: obj.test_desired_placement_of_rewards_and_actions() without return types