code: def test_create_episode_dataset(self):
      d4rl_dataset = {
          "observations": [
              [1.0, 2.0],
              [3.0, 4.0],
              [5.0, 6.0],
              [7.0, 8.0],
              [9.0, 10.0],
              [11.0, 12.0],
              [13.0, 14.0],
          ],
          "actions": [[1.0], [2.0], [3.0], [4.0], [5.0], [6.0], [7.0]],
          "rewards": [[0.0], [1.0], [0.0], [1.0], [0.0], [0.0], [1.0]],
          "terminals": [False, True, False, True, False, False, True],
          "timeouts": [False, False, False, False, False, False, False],
          "infos/goal": [[0.0]] * 8,
      }
  
      episode_dict = create_episode_dataset(
          d4rl_dataset, exclude_timeouts=True
      )
      expected_dict = {
          "states": np.array(
              [
                  [1.0, 2.0],
                  [3.0, 4.0],
                  [5.0, 6.0],
                  [7.0, 8.0],
                  [9.0, 10.0],
                  [11.0, 12.0],
                  [13.0, 14.0],
              ]
          ),
          "actions": np.array(
              [[1.0], [2.0], [3.0], [4.0], [5.0], [6.0], [7.0]]
          ),
          "rewards": np.array(
              [[0.0], [1.0], [0.0], [1.0], [0.0], [0.0], [1.0]]
          ),
          "discounts": np.array([1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0]),
          "episode_start_index": np.array([0, 2, 4]),
      }
      self.assertDictEqual(episode_dict, expected_dict)
  
file_path: agents/tf_agents/examples/cql_sac/kumar20/dataset/dataset_utils_test.py
imports: import numpy
  import tensorflow
  from tf_agents.examples.cql_sac.kumar20.dataset.dataset_utils import create_collect_data_spec
  from tf_agents.examples.cql_sac.kumar20.dataset.dataset_utils import create_episode_dataset
  from tf_agents.specs.array_spec import ArraySpec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import test_utils
label_desc: Use image node_id 2 for calling the UtilsTest obj's underlying member method code with example usage: obj.test_create_episode_dataset() without return types