code: def _assert_sample_batch_keys(batch: SampleBatch):
      """Assert sampled batch has the requisite keys."""
      assert SampleBatch.OBS in batch
      assert SampleBatch.ACTIONS in batch
      assert SampleBatch.RETURNS_TO_GO in batch
      assert SampleBatch.T in batch
      assert SampleBatch.ATTENTION_MASKS in batch
  
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 2 for calling a global function with example usage: _assert_sample_batch_keys(batch) without return types