code: def testGetBatchEpisodeIds(self):
      spec = [
          specs.TensorSpec([3], tf.float32, "action"),
          [
              specs.TensorSpec([5], tf.float32, "lidar"),
              specs.TensorSpec([3, 2], tf.float32, "camera"),
          ],
      ]
      replay_buffer = episodic_replay_buffer.EpisodicReplayBuffer(
          spec, capacity=5
      )
      episode_ids = [replay_buffer.create_episode_ids(num_episodes=3)]
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.evaluate(tf.compat.v1.local_variables_initializer())
      expected_ids = [[-1, -1, -1]]
  
      # batch_episode_ids should be updated the first time regardless of
      # begin_episode.
      episode_ids.append(
          replay_buffer._get_batch_episode_ids(
              episode_ids[-1], begin_episode=False
          )
      )
      expected_ids.append([0, 1, 2])
  
      # batch_episode_ids should not be updated when begin_episode is False.
      episode_ids.append(
          replay_buffer._get_batch_episode_ids(
              episode_ids[-1], begin_episode=False
          )
      )
      expected_ids.append([0, 1, 2])
  
      episode_ids.append(
          replay_buffer._get_batch_episode_ids(
              episode_ids[-1], begin_episode=[False, False, False]
          )
      )
      expected_ids.append([0, 1, 2])
  
      # batch_episode_ids should be updated only when begin_episode is True.
      episode_ids.append(
          replay_buffer._get_batch_episode_ids(
              episode_ids[-1], begin_episode=[True, False, False]
          )
      )
      expected_ids.append([3, 1, 2])
  
      episode_ids.append(
          replay_buffer._get_batch_episode_ids(
              episode_ids[-1], begin_episode=[False, True, False]
          )
      )
      expected_ids.append([3, 4, 2])
  
      episode_ids.append(
          replay_buffer._get_batch_episode_ids(
              episode_ids[-1], begin_episode=[False, True, True]
          )
      )
      expected_ids.append([3, 5, 6])
  
      self.assertAllEqual(expected_ids, self.evaluate(episode_ids))
  
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 7 for calling the EpisodicReplayBufferTest obj's underlying member method code with example usage: obj.testGetBatchEpisodeIds() without return types