code: def test_create_episode_dataset_exclude_timeout(self):
      d4rl_dataset = {
          "observations": [[1.0, 2.0], [3.0, 4.0]],
          "actions": [[1.0], [2.0]],
          "rewards": [[0.0], [0.0]],
          "terminals": [False, False],
          "timeouts": [False, True],
          "infos/goal": [[10.0, 10.0], [10.0, 10.0]],
      }
  
      episode_dict = create_episode_dataset(
          d4rl_dataset, exclude_timeouts=True
      )
  
      # Threw out timeout step.
      expected_dict = {
          "states": np.array([[1.0, 2.0]]),
          "actions": np.array([[1.0]]),
          "rewards": np.array([[0.0]]),
          "discounts": np.array([1.0]),
          "episode_start_index": np.array([0]),
      }
      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 3 for calling the UtilsTest obj's underlying member method code with example usage: obj.test_create_episode_dataset_exclude_timeout() without return types