code: def test_truncate(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)
  
      # Take a few steps.
      ts = 5
      for _ in range(ts):
          (
              observations,
              rewards,
              actions,
              dones,
          ) = self.get_random_observations_rewards_actions_dones(
              batch_size=batch_size
          )
          dones[...] = False
          bt.step(observations, rewards, rewards, dones, actions)
  
      self.assertEqual(0, bt.num_completed_trajectories)
  
      num_to_keep = 2
      bt.truncate_trajectories(indices, num_to_keep=num_to_keep)
  
      self.assertEqual(batch_size, bt.num_completed_trajectories)
  
      # Assert they are all active.
      # Since the last `num_to_keep` observations were duplicated.
      self.assertTrue(all(t.is_active for t in bt.trajectories))
  
      orig_obs = bt.completed_trajectories[0].observations_np
      # + 1 because of the initial reset
      self.assertEqual(ts + 1, orig_obs.shape[0])
  
      trunc_obs = bt.trajectories[0].observations_np
      self.assertEqual(num_to_keep, trunc_obs.shape[0])
      self.assertEqual(num_to_keep, bt.trajectories[0].num_time_steps)
  
      # Test that the observations are the same.
      self.assertAllEqual(orig_obs[-num_to_keep:, ...], trunc_obs)
  
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 6 for calling the BatchTrajectoryTest obj's underlying member method code with example usage: obj.test_truncate() without return types