code: def test_split_by_episode(self):
      s = SampleBatch(
          {
              "a": np.array([0, 1, 2, 3, 4, 5]),
              "eps_id": np.array([0, 0, 0, 0, 1, 1]),
              "terminateds": np.array([0, 0, 0, 1, 0, 1]),
          }
      )
      true_split = [np.array([0, 1, 2, 3]), np.array([4, 5])]
  
      # Check that splitting by EPS_ID works correctly
      eps_split = [b["a"] for b in s.split_by_episode()]
      check(true_split, eps_split)
  
      # Check that splitting by EPS_ID works correctly when explicitly specified
      eps_split = [b["a"] for b in s.split_by_episode(key="eps_id")]
      check(true_split, eps_split)
  
      # Check that splitting by DONES works correctly when explicitly specified
      eps_split = [b["a"] for b in s.split_by_episode(key="dones")]
      check(true_split, eps_split)
  
      # Check that splitting by DONES works correctly
      del s["eps_id"]
      terminateds_split = [b["a"] for b in s.split_by_episode()]
      check(true_split, terminateds_split)
  
      # Check that splitting without the EPS_ID or DONES key raise an error
      del s["terminateds"]
      with self.assertRaises(KeyError):
          s.split_by_episode()
  
      # Check that splitting with DONES always False returns the whole batch
      s["terminateds"] = np.array([0, 0, 0, 0, 0, 0])
      batch_split = [b["a"] for b in s.split_by_episode()]
      check(s["a"], batch_split[0])
  
imports: import copy
  import functools
  import os
  import unittest
  import numpy
  import torch
  import tree
  import ray
  from ray.rllib.models.repeated_values import RepeatedValues
  from ray.rllib.policy.sample_batch import SampleBatch, attempt_count_timesteps, concat_samples
  from ray.rllib.utils.compression import is_compressed
  from ray.rllib.utils.framework import try_import_torch
  from ray.rllib.utils.test_utils import check
  from ray.rllib.utils.torch_utils import convert_to_torch_tensor
label_desc: Use image node_id 11 for calling the TestSampleBatch obj's underlying member method code with example usage: obj.test_split_by_episode() without return types