class_docstr: Tests insertion and (weighted) sampling of the PrioritizedReplayBuffer.
class_docstr_tok: ["Tests","insertion","and","(","weighted",")","sampling","of","the","PrioritizedReplayBuffer","."]
docstr: Tests adding, sampling, and eviction of episodes.
docstr_tok: ["Tests","adding",",","sampling",",","and","eviction","of","episodes","."]
code: def test_episodes_unit(self):
      """Tests adding, sampling, and eviction of episodes."""
      # We do not test the mathematical correctness of our prioritization
      # here but rather if it works together with episode mode at all
      buffer = PrioritizedReplayBuffer(
          capacity=18, storage_unit="episodes"
      )
  
      batches = [
          SampleBatch(
              {
                  SampleBatch.T: [0, 1, 2, 3],
                  SampleBatch.ACTIONS: 4 * [np.random.choice([0, 1])],
                  SampleBatch.REWARDS: 4 * [np.random.rand()],
                  SampleBatch.TERMINATEDS: [False, False, False, True],
                  SampleBatch.TRUNCATEDS: [False, False, False, False],
                  SampleBatch.SEQ_LENS: [4],
                  SampleBatch.EPS_ID: 4 * [i],
              }
          )
          for i in range(3)
      ]
  
      # Add some batches with episodes of low priority
      for batch in batches:
          buffer.add(batch, weight=0.01)
  
      # Add two high priority episodes
      buffer.add(
          SampleBatch(
              {
                  SampleBatch.T: [0, 1, 0, 1],
                  SampleBatch.ACTIONS: 4 * [np.random.choice([0, 1])],
                  SampleBatch.REWARDS: 4 * [np.random.rand()],
                  SampleBatch.TERMINATEDS: [False, True, False, True],
                  SampleBatch.TRUNCATEDS: [False, False, False, True],
                  SampleBatch.SEQ_LENS: [2, 2],
                  SampleBatch.EPS_ID: [3, 3, 4, 4],
              }
          ),
          weight=1,
      )
  
      num_sampled_dict = {_id: 0 for _id in range(5)}
      num_samples = 200
      for i in range(num_samples):
          sample = buffer.sample(1, beta=self.beta)
          _id = sample[SampleBatch.EPS_ID][0]
          assert len(sample[SampleBatch.SEQ_LENS]) == 1
          num_sampled_dict[_id] += 1
  
      # All episodes, even though in different batches should be sampled
      # equally often
      assert np.allclose(
          np.array(list(num_sampled_dict.values())) / num_samples,
          [0, 0, 0, 0.5, 0.5],
          atol=0.1,
      )
  
      # Episode 6 is not entirely inside this batch, it should not be added
      # to the buffer
      buffer.add(
          SampleBatch(
              {
                  SampleBatch.T: [0, 1, 0, 1],
                  SampleBatch.ACTIONS: 4 * [np.random.choice([0, 1])],
                  SampleBatch.REWARDS: 4 * [np.random.rand()],
                  SampleBatch.TERMINATEDS: [False, True, False, False],
                  SampleBatch.TRUNCATEDS: [False, False, False, False],
                  SampleBatch.SEQ_LENS: [2, 2],
                  SampleBatch.EPS_ID: [5, 5, 6, 6],
              }
          ),
          weight=1,
      )
  
      num_sampled_dict = {_id: 0 for _id in range(7)}
      num_samples = 200
      for i in range(num_samples):
          sample = buffer.sample(1, beta=self.beta)
          _id = sample[SampleBatch.EPS_ID][0]
          assert len(sample[SampleBatch.SEQ_LENS]) == 1
          num_sampled_dict[_id] += 1
  
      # Episode 7 should be dropped for not ending inside the batch
      assert np.allclose(
          np.array(list(num_sampled_dict.values())) / num_samples,
          [0, 0, 0, 1 / 3, 1 / 3, 1 / 3, 0],
          atol=0.1,
      )
  
      # Add another batch to evict the first batch.
      buffer.add(
          SampleBatch(
              {
                  SampleBatch.T: [0, 1, 2, 3],
                  SampleBatch.ACTIONS: 4 * [np.random.choice([0, 1])],
                  SampleBatch.REWARDS: 4 * [np.random.rand()],
                  SampleBatch.TERMINATEDS: [False, False, False, True],
                  SampleBatch.TRUNCATEDS: [False, False, False, True],
                  SampleBatch.SEQ_LENS: [4],
                  SampleBatch.EPS_ID: 4 * [7],
              }
          ),
          weight=0.01,
      )
  
      # After adding 1 more batch, eviction has started with 24
      # timesteps added in total, 2 of which were discarded
      assert len(buffer) == 6
      assert buffer._num_timesteps_added == 4 * 6 - 2
      assert buffer._num_timesteps_added_wrap == 4
      assert buffer._next_idx == 1
      assert buffer._eviction_started is True
  
      num_sampled_dict = {_id: 0 for _id in range(8)}
      num_samples = 200
      for i in range(num_samples):
          sample = buffer.sample(1, beta=self.beta)
          _id = sample[SampleBatch.EPS_ID][0]
          assert len(sample[SampleBatch.SEQ_LENS]) == 1
          num_sampled_dict[_id] += 1
  
      assert np.allclose(
          np.array(list(num_sampled_dict.values())) / num_samples,
          [0, 0, 0, 1 / 3, 1 / 3, 1 / 3, 0, 0],
          atol=0.1,
      )
  
file_path: ray/python/ray/rllib/utils/replay_buffers/tests/test_prioritized_replay_buffer_replay_buffer_api.py
imports: from collections import Counter
  import numpy
  import unittest
  from ray.rllib.utils.replay_buffers.prioritized_replay_buffer import PrioritizedReplayBuffer
  from ray.rllib.policy.sample_batch import SampleBatch, MultiAgentBatch, concat_samples
  from ray.rllib.utils.test_utils import check
label_desc: Use image node_id 8 for calling the TestPrioritizedReplayBuffer obj's underlying member method code with example usage: obj.test_episodes_unit() without return types