code: def test_observations_np(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)
  
      # Number of time-steps now looks like the following:
      # (1, 1, 1, 1, 1, 1, 1, 1, 1, 1)
      lengths = np.full((self.BATCH_SIZE,), 1)
  
      ts = 5
      for _ in range(ts):
          (
              observations,
              rewards,
              actions,
              dones,
          ) = self.get_random_observations_rewards_actions_dones()
          dones[...] = False
          bt.step(observations, rewards, rewards, dones, actions)
  
      # Number of time-steps now looks like the following:
      # (6, 6, 6, 6, 6, 6, 6, 6, 6, 6)
      lengths = lengths + ts
  
      # Now let's mark the first two as done.
      (
          observations,
          _,
          _,
          _,
      ) = self.get_random_observations_rewards_actions_dones(
          batch_size=2
      )
      bt.reset(np.array([0, 1]), observations)
  
      # Number of time-steps now looks like the following:
      # (1, 1, 6, 6, 6, 6, 6, 6, 6, 6)
      lengths[0] = lengths[1] = 1
  
      for _ in range(ts):
          (
              observations,
              rewards,
              actions,
              dones,
          ) = self.get_random_observations_rewards_actions_dones()
          dones[...] = False
          bt.step(observations, rewards, rewards, dones, actions)
  
      # Number of time-steps now looks like the following:
      # (6, 6, 11, 11, 11, 11, 11, 11, 11, 11)
      lengths = lengths + ts
  
      boundary = 20
      len_history_for_policy = 40
  
      padded_obs_np, padded_lengths = bt.observations_np(
          boundary=boundary,
          len_history_for_policy=len_history_for_policy,
      )
  
      # The lengths are what we expect them to be.
      self.assertAllEqual(lengths, padded_lengths)
  
      # The padded_observations are the shape we expect them to be.
      self.assertEqual(
          (self.BATCH_SIZE, boundary + 1) + self.OBSERVATION_SHAPE,
          padded_obs_np.shape,
      )
  
      # Let's now request the last n = [1, 2 * boundary) steps for the history.
      for len_history_for_policy in range(1, 2 * boundary):
          # The expected lengths will now be:
          truncated_lengths = [
              min(l, len_history_for_policy) for l in lengths
          ]
  
          padded_obs_np, padded_lengths = bt.observations_np(
              boundary=boundary,
              len_history_for_policy=len_history_for_policy,
          )
  
          self.assertAllEqual(truncated_lengths, padded_lengths)
  
          # This shouldn't change, since even if we request lengths > boundary + 1
          # there are no trajectories that long.
          self.assertEqual(
              (self.BATCH_SIZE, boundary + 1) + self.OBSERVATION_SHAPE,
              padded_obs_np.shape,
          )
  
      # Let's do 10 more steps (to go on the other side of the boundary.
      ts = 10
      for _ in range(ts):
          (
              observations,
              rewards,
              actions,
              dones,
          ) = self.get_random_observations_rewards_actions_dones()
          dones[...] = False
          bt.step(observations, rewards, rewards, dones, actions)
  
      # Number of time-steps now looks like the following:
      # (16, 16, 21, 21, 21, 21, 21, 21, 21, 21)
      lengths = lengths + ts
  
      len_history_for_policy = 40
      padded_obs_np, padded_lengths = bt.observations_np(
          boundary=boundary,
          len_history_for_policy=len_history_for_policy,
      )
  
      # The lengths are what we expect them to be.
      self.assertAllEqual(lengths, padded_lengths)
  
      # The padded_observations are the shape we expect them to be.
      self.assertEqual(
          (self.BATCH_SIZE, (2 * boundary) + 1)
          + self.OBSERVATION_SHAPE,
          padded_obs_np.shape,
      )
  
      # Test that the padding is the only part that is all 0s.
      # NOTE: There is almost 0 probability that the random observation is all 0s.
      zero_obs = np.full(self.OBSERVATION_SHAPE, 0.0)
      for b in range(self.BATCH_SIZE):
          # The first lengths[b] will be actual data, rest is 0s.
          for ts in range(lengths[b]):
              self.assertFalse(np.all(zero_obs == padded_obs_np[b][ts]))
  
          for ts in range(lengths[b], len(padded_obs_np[b])):
              self.assertAllEqual(zero_obs, padded_obs_np[b][ts])
  
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 9 for calling the BatchTrajectoryTest obj's underlying member method code with example usage: obj.test_observations_np() without return types