docstr: Test sampling from a segmentation buffer.
docstr_tok: ["Test","sampling","from","a","segmentation","buffer","."]
code: def test_sample_basic(self):
      """Test sampling from a segmentation buffer."""
      for buffer_cls in (
          SegmentationBuffer,
          MultiAgentSegmentationBuffer,
      ):
          max_seq_len = 5
          max_ep_len = 15
          capacity = 4
          obs_dim = 10
          act_dim = 2
  
          buffer = buffer_cls(capacity, max_seq_len, max_ep_len)
  
          # generate batch and add to buffer
          episode_batches = []
          for i in range(8):
              episode_batches.append(
                  _generate_episode_batch(
                      max_ep_len, i, obs_dim, act_dim
                  )
              )
          batch = concat_samples(episode_batches)
          buffer.add(batch)
  
          # sample a few times and check shape
          for bs in range(10, 20):
              batch = _as_sample_batch(buffer.sample(bs))
              # check the keys exist
              _assert_sample_batch_keys(batch)
  
              # check the shapes
              self.assertEquals(
                  batch[SampleBatch.OBS].shape,
                  (bs, max_seq_len, obs_dim),
              )
              self.assertEquals(
                  batch[SampleBatch.ACTIONS].shape,
                  (bs, max_seq_len, act_dim),
              )
              self.assertEquals(
                  batch[SampleBatch.RETURNS_TO_GO].shape,
                  (
                      bs,
                      max_seq_len + 1,
                      1,
                  ),
              )
              self.assertEquals(
                  batch[SampleBatch.T].shape, (bs, max_seq_len)
              )
              self.assertEquals(
                  batch[SampleBatch.ATTENTION_MASKS].shape,
                  (bs, max_seq_len),
              )
  
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 4 for calling the TestSegmentationBuffer obj's underlying member method code with example usage: obj.test_sample_basic() without return types