docstr: Test that sampling from buffer of capacity > 1 works.
docstr_tok: ["Test","that","sampling","from","buffer","of","capacity",">","1","works","."]
code: def test_sample_capacity(self):
      """Test that sampling from buffer of capacity > 1 works."""
      for buffer_cls in (
          SegmentationBuffer,
          MultiAgentSegmentationBuffer,
      ):
          max_seq_len = 3
          max_ep_len = 10
          capacity = 100
          obs_dim = 1
          act_dim = 1
  
          buffer = buffer_cls(capacity, max_seq_len, max_ep_len)
  
          # Generate batch and add to buffer
          episode_batches = []
          for i in range(capacity):
              episode_batches.append(
                  _generate_episode_batch(
                      max_ep_len, i, obs_dim, act_dim
                  )
              )
          buffer.add(concat_samples(episode_batches))
  
          # Sample 100 times and check that samples are from at least 2 different
          # episodes. The [robability of all sampling from 1 episode by chance is
          # (1/100)**99 which is basically zero.
          batch = _as_sample_batch(buffer.sample(100))
          eps_ids = set()
          for i in range(100):
              # obs generated by _generate_episode_batch contains eps_id
              # use -1 because there might be front padding
              eps_id = int(batch[SampleBatch.OBS][i, -1, 0])
              eps_ids.add(eps_id)
  
          self.assertGreater(
              len(eps_ids),
              1,
              "buffer.sample is always returning the same episode.",
          )
  
imports: import unittest
  from typing import List, Union
  import numpy
  from rllib_dt.dt.segmentation_buffer import MultiAgentSegmentationBuffer, SegmentationBuffer
  import ray
  from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, MultiAgentBatch, SampleBatch, concat_samples
  from ray.rllib.utils import test_utils
  from ray.rllib.utils.framework import try_import_tf, try_import_torch
  from ray.rllib.utils.typing import PolicyID
label_desc: Use image node_id 6 for calling the TestSegmentationBuffer obj's underlying member method code with example usage: obj.test_sample_capacity() without return types