code: def _assert_is_segment(segment: SampleBatch, episode: SampleBatch):
      """Assert that the sampled segment is a segment of episode."""
      timesteps = segment[SampleBatch.T]
      masks = segment[SampleBatch.ATTENTION_MASKS] > 0.5
      seq_len = timesteps.shape[0]
      episode_segment = episode.slice(timesteps[0], timesteps[-1] + 1)
      assert np.allclose(
          segment[SampleBatch.OBS][masks],
          episode_segment[SampleBatch.OBS],
      )
      assert np.allclose(
          segment[SampleBatch.ACTIONS][masks],
          episode_segment[SampleBatch.ACTIONS],
      )
      assert np.allclose(
          segment[SampleBatch.RETURNS_TO_GO][:seq_len][masks],
          episode_segment[SampleBatch.RETURNS_TO_GO],
      )
  
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 a global function with example usage: _assert_is_segment(segment, episode) without return types