code: def test_data_split():
      def _extract_observation_values_from(
          imitation_samples: List[Tuple[Dict, Dict]]
      ):
          """Extract observation values from array of imitation samples of (obs, act) tuples"""
          return list(
              map(
                  lambda sample: sample[0][0]["observation"],
                  imitation_samples,
              )
          )
  
      dataset = InMemoryDataset(
          n_workers=1,
          conversion_env_factory=_env_factory,
          input_data=None,
          trajectory_processor=IdentityTrajectoryProcessor(),
          deserialize_in_main_thread=False,
      )
  
      # Fill dataset with two episodes with 5 usable steps each
      for _ in range(2):
          dataset.append(_mock_state_trajectory_record(6))
  
      assert dataset.trajectory_references[0] == range(0, 5)
      assert dataset.trajectory_references[1] == range(5, 10)
  
      train, valid = dataset.random_split([5, 5])
      assert len(train) == 5
      assert len(valid) == 5
      assert _extract_observation_values_from(train) == [0, 2, 4, 6, 8]
      assert _extract_observation_values_from(valid) == [0, 2, 4, 6, 8]
  
      # Add two more episodes with 5 usable steps each (now we have 4 in total)
      for _ in range(2):
          dataset.append(_mock_state_trajectory_record(6))
  
      # 50:50 split
      train, valid = dataset.random_split([10, 10])
      assert len(train) == 10
      assert len(valid) == 10
      assert (
          _extract_observation_values_from(train) == [0, 2, 4, 6, 8] * 2
      )
      assert (
          _extract_observation_values_from(valid) == [0, 2, 4, 6, 8] * 2
      )
  
      # 75:25 split
      train, valid = dataset.random_split([15, 5])
      assert len(train) == 15
      assert len(valid) == 5
      assert (
          _extract_observation_values_from(train) == [0, 2, 4, 6, 8] * 3
      )
      assert _extract_observation_values_from(valid) == [0, 2, 4, 6, 8]
  
      # 80:20 split: Not satisfiable with episodes of 4 * 5 steps, should fall back on 15:5 ratio
      train, valid = dataset.random_split([16, 4])
      assert len(train) == 15
      assert len(valid) == 5
      assert (
          _extract_observation_values_from(train) == [0, 2, 4, 6, 8] * 3
      )
      assert _extract_observation_values_from(valid) == [0, 2, 4, 6, 8]
  
      # 70:30 split: Not satisfiable with episodes of 4 * 5 steps, should fall back on 10:10 ratio
      train, valid = dataset.random_split([14, 6])
      assert len(train) == 10
      assert len(valid) == 10
      assert (
          _extract_observation_values_from(train) == [0, 2, 4, 6, 8] * 2
      )
      assert (
          _extract_observation_values_from(valid) == [0, 2, 4, 6, 8] * 2
      )
  
      # 50:25:25 split
      train, test, valid = dataset.random_split([10, 5, 5])
      assert len(train) == 10
      assert len(test) == 5
      assert len(valid) == 5
      assert (
          _extract_observation_values_from(train) == [0, 2, 4, 6, 8] * 2
      )
      assert (
          _extract_observation_values_from(test) == [0, 2, 4, 6, 8] * 1
      )
      assert (
          _extract_observation_values_from(valid) == [0, 2, 4, 6, 8] * 1
      )
  
imports: import pickle
  from typing import Any, Dict, Union, Tuple, List, Optional
  import gym
  import numpy
  from maze.core.env.maze_action import MazeActionType
  from maze.core.env.maze_state import MazeStateType
  from maze.core.env.structured_env import ActorID
  from maze.core.env.structured_env_spaces_mixin import StructuredEnvSpacesMixin
  from maze.core.log_events.step_event_log import StepEventLog
  from maze.core.trajectory_recording.datasets.trajectory_processor import IdentityTrajectoryProcessor
  from maze.core.trajectory_recording.records.spaces_record import SpacesRecord
  from maze.core.trajectory_recording.records.structured_spaces_record import StructuredSpacesRecord
  from maze.core.trajectory_recording.records.state_record import StateRecord
  from maze.core.trajectory_recording.records.trajectory_record import StateTrajectoryRecord, SpacesTrajectoryRecord
  from maze.core.wrappers.maze_gym_env_wrapper import make_gym_maze_env
  from maze.core.wrappers.wrapper import ObservationWrapper
  from maze.test.shared_test_utils.dummy_env.dummy_core_env import DummyCoreEnvironment
  from maze.test.shared_test_utils.dummy_env.dummy_maze_env import DummyEnvironment
  from maze.test.shared_test_utils.dummy_env.space_interfaces.action_conversion.double import DoubleActionConversion
  from maze.test.shared_test_utils.dummy_env.space_interfaces.observation_conversion.double import DoubleObservationConversion
  from maze.test.shared_test_utils.run_maze_utils import run_maze_job
  from maze.core.trajectory_recording.datasets.in_memory_dataset import InMemoryDataset
label_desc: Use image node_id 7 for calling a global function with example usage: test_data_split() and returns: list