code: def create_episode_dataset(  # pytype: disable=annotation-type-mismatch  # numpy-scalars
      d4rl_dataset: Dict[str, D4RLListType],
      exclude_timeouts: bool,
      observation_dtype: np.dtype = np.float32,
  ) -> EpisodeDictType:
      """Create a dataset of episodes."""
      dataset = dict(
          states=[],
          actions=[],
          rewards=[],
          discounts=[],
          episode_start_index=[],
      )
  
      new_episode = True
      for i in range(len(d4rl_dataset["observations"])):
          if new_episode:
              dataset["episode_start_index"].append(
                  len(dataset["states"])
              )
          dataset["states"].append(d4rl_dataset["observations"][i])
          dataset["actions"].append(d4rl_dataset["actions"][i])
          dataset["rewards"].append(d4rl_dataset["rewards"][i])
          dataset["discounts"].append(
              1.0 - int(d4rl_dataset["terminals"][i])
          )
  
          is_terminal = d4rl_dataset["terminals"][i]
          is_timeout = d4rl_dataset["timeouts"][i]
          # Remove the last step if this is a timeout. Assumes timeouts occur
          # after at least 1 step in an episode.
          if exclude_timeouts and is_timeout and not is_terminal:
              for key in ["states", "actions", "rewards", "discounts"]:
                  dataset[key].pop()
          new_episode = is_terminal or is_timeout
  
      for key, value in dataset.items():
          dtype = observation_dtype if key == "states" else np.float32
          dataset[key] = np.asarray(value, dtype)
  
      return dataset
  
file_path: agents/tf_agents/examples/cql_sac/kumar20/dataset/dataset_utils.py
imports: from typing import Dict, List, Union
  import d4rl
  import numpy
  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
label_desc: Use image node_id 2 for calling a global function with example usage: create_episode_dataset(d4rl_dataset, exclude_timeouts, observation_dtype) and returns: dataset