code: def testMultipleAddBatch(self):
      spec = (
          specs.TensorSpec([3], tf.int32, "lidar"),
          specs.TensorSpec([], tf.bool, "begin_episode"),
      )
      replay_buffer = episodic_replay_buffer.EpisodicReplayBuffer(
          spec,
          capacity=3,
          begin_episode_fn=lambda value_and_begin: value_and_begin[1],
          end_episode_fn=lambda _: False,
      )
      stateful_replay_buffer = (
          episodic_replay_buffer.StatefulEpisodicReplayBuffer(
              replay_buffer, num_episodes=3
          )
      )
      episode_ids_var = stateful_replay_buffer.episode_ids
  
      values = np.ones(spec[0].shape.as_list(), dtype=np.int32)
      values = np.stack([values, 10 * values, 100 * values])
      # In this case all episodes are valid, so it will add all the values.
      # Add 1 to ep0, 10 to ep1, 100 to ep2
      new_episode_ids = stateful_replay_buffer.add_batch(
          (values, tf.constant([False, False, False]))
      )
      with tf.control_dependencies([new_episode_ids]):
          items = [[replay_buffer._get_episode(i) for i in [0, 1, 2]]]
      # In this case the second episode is invalid, so it won't add its values.
      # Add 1 to ep3 (was ep0), ., 100 to ep4 (was ep1).
      with tf.control_dependencies(items[-1][0]):
          new_episode_ids_2 = stateful_replay_buffer.add_batch(
              (values, tf.constant([True, False, True]))
          )
      with tf.control_dependencies([new_episode_ids_2]):
          items.append(
              [replay_buffer._get_episode(i) for i in [2, 3, 4]]
          )
      # In this case all episodes are valid, so it will add all the values.
      # Add 1 to ep3, 10 to ep5 (was ep2), 100 to ep4.
      with tf.control_dependencies(items[-1][0]):
          new_episode_ids_3 = stateful_replay_buffer.add_batch(
              (values, tf.constant([False, True, False]))
          )
      with tf.control_dependencies([new_episode_ids_3]):
          # End result: ep3 with [1, 1], ep4 with [100, 100], ep5 with [10].
          items.append(
              [replay_buffer._get_episode(i) for i in [3, 4, 5]]
          )
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.evaluate(tf.compat.v1.local_variables_initializer())
      (
          new_episode_ids,
          new_episode_ids_2,
          new_episode_ids_3,
          items,
      ) = self.evaluate(
          (new_episode_ids, new_episode_ids_2, new_episode_ids_3, items)
      )
      self.assertAllEqual([0, 1, 2], new_episode_ids)
      self.assertAllEqual([3, 1, 4], new_episode_ids_2)
      self.assertAllEqual([3, 5, 4], new_episode_ids_3)
      episode_ids_value = self.evaluate(episode_ids_var)
      self.assertAllEqual([3, 5, 4], episode_ids_value)
      self.assertEqual(
          5, self.evaluate(replay_buffer._get_last_episode_id())
      )
  
      # First add_batch: ep0, ep1, ep2.
      self.assertAllEqual(items[0][0][0], [values[0]])
      self.assertAllEqual(items[0][1][0], [values[1]])
      self.assertAllEqual(items[0][2][0], [values[2]])
      # Second add_batch: ep2, ep3, ep4.
      self.assertAllEqual(items[1][0][0], [values[2]])
      self.assertAllEqual(items[1][1][0], [values[0]])
      self.assertAllEqual(items[1][2][0], [values[2]])
      # Third add_batch: ep3, ep4, ep5.
      self.assertAllEqual(items[2][0][0], [values[0], values[0]])
      self.assertAllEqual(items[2][1][0], [values[2], values[2]])
      self.assertAllEqual(items[2][2][0], [values[1]])
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tf_agents import specs
  from tf_agents.replay_buffers import episodic_replay_buffer
  from tf_agents.utils import common
  from tf_agents.utils import test_utils
  from tensorflow.python.ops import list_ops
label_desc: Use image node_id 9 for calling the StatefulEpisodicReplayBufferTest obj's underlying member method code with example usage: obj.testMultipleAddBatch() without return types