code: def _generate_episode_batch(ep_len, eps_id, obs_dim=8, act_dim=3):
      """Generate a batch containing one episode."""
      # These values are not actually correct as usual. But using eps_id
      # as the values allow us to identify them in the tests.
      batch = SampleBatch(
          {
              SampleBatch.OBS: np.full(
                  (ep_len, obs_dim), eps_id, dtype=np.float32
              ),
              SampleBatch.ACTIONS: np.full(
                  (ep_len, act_dim), eps_id + 100, dtype=np.float32
              ),
              SampleBatch.REWARDS: np.ones((ep_len,), dtype=np.float32),
              SampleBatch.RETURNS_TO_GO: np.arange(
                  ep_len, -1, -1, dtype=np.float32
              ).reshape((ep_len + 1, 1)),
              SampleBatch.EPS_ID: np.full(
                  (ep_len,), eps_id, dtype=np.int32
              ),
              SampleBatch.T: np.arange(ep_len, dtype=np.int32),
              SampleBatch.ATTENTION_MASKS: np.ones(
                  ep_len, dtype=np.float32
              ),
              SampleBatch.TERMINATEDS: np.array(
                  [False] * (ep_len - 1) + [True]
              ),
              SampleBatch.TRUNCATEDS: np.array([False] * ep_len),
          }
      )
      return 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 1 for calling a global function with example usage: _generate_episode_batch(ep_len, eps_id, obs_dim, act_dim) and returns: batch