class_docstr: A minimal replay buffer used by Decision Transformer (DT)
  to process episodes into max_seq_len length segments and do shuffling.
  Stores MultiAgentSample.
class_docstr_tok: ["A","minimal","replay","buffer","used","by","Decision","Transformer","(","DT",")","to","process","episodes","into","max_seq_len","length","segments","and","do","shuffling",".","Stores","MultiAgentSample","."]
docstr: Sample segments from the buffer.
  
  Args:
      batch_size: number of segments to sample.
  
  Returns:
      MultiAgentBatch of segments with keys and shape {
          OBS: [batch_size, max_seq_len, obs_dim],
          ACTIONS: [batch_size, max_seq_len, act_dim],
          RETURNS_TO_GO: [batch_size, max_seq_len + 1, 1],
          T: [batch_size, max_seq_len],
          ATTENTION_MASKS: [batch_size, max_seq_len],
      }
docstr_tok: ["Sample","segments","from","the","buffer",".","Args",":","batch_size",":","number","of","segments","to","sample",".","Returns",":","MultiAgentBatch","of","segments","with","keys","and","shape","{","OBS",":","[","batch_size",",","max_seq_len",",","obs_dim","]",",","ACTIONS",":","[","batch_size",",","max_seq_len",",","act_dim","]",",","RETURNS_TO_GO",":","[","batch_size",",","max_seq_len","+","1",",","1","]",",","T",":","[","batch_size",",","max_seq_len","]",",","ATTENTION_MASKS",":","[","batch_size",",","max_seq_len","]",",","}"]
code: def sample(self, batch_size: int) -> MultiAgentBatch:
      """Sample segments from the buffer.
  
      Args:
          batch_size: number of segments to sample.
  
      Returns:
          MultiAgentBatch of segments with keys and shape {
              OBS: [batch_size, max_seq_len, obs_dim],
              ACTIONS: [batch_size, max_seq_len, act_dim],
              RETURNS_TO_GO: [batch_size, max_seq_len + 1, 1],
              T: [batch_size, max_seq_len],
              ATTENTION_MASKS: [batch_size, max_seq_len],
          }
      """
      samples = {}
      for policy_id, buffer in self.buffers.items():
          samples[policy_id] = buffer.sample(batch_size)
      return MultiAgentBatch(samples, batch_size)
  
file_path: ray/rllib_contrib/dt/src/rllib_dt/dt/segmentation_buffer.py
imports: import logging
  import random
  from collections import defaultdict
  from typing import List
  import numpy
  from ray.rllib.evaluation.postprocessing import discount_cumsum
  from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch, concat_samples
  from ray.rllib.utils.typing import SampleBatchType
label_desc: Use image node_id 3 for calling the MultiAgentSegmentationBuffer obj's underlying member method code with example usage: obj.sample(batch_size) and returns: MultiAgentBatch