code: def write_samples_to_tfrecord(
      dataset_dict: Dict[str, types.Array],
      collect_data_spec: trajectory.Transition,
      dataset_path: str,
      start_episode: int,
      end_episode: int,
      use_trajectories: bool = True,
  ) -> None:
      """Creates and writes samples to a TFRecord file."""
      tfrecord_observer = example_encoding_dataset.TFRecordObserver(
          dataset_path, collect_data_spec, py_mode=True
      )
      states = dataset_dict["states"]
      actions = dataset_dict["actions"]
      discounts = dataset_dict["discounts"]
      rewards = dataset_dict["rewards"]
      num_episodes = len(dataset_dict["episode_start_index"])
  
      for episode_i in range(start_episode, end_episode):
          episode_start_index = dataset_dict["episode_start_index"][
              episode_i
          ]
          # If this is the last episode, end at the final step.
          if episode_i == (num_episodes - 1):
              episode_end_index = len(states)
          else:
              # Otherwise, end before the next episode.
              episode_end_index = dataset_dict["episode_start_index"][
                  episode_i + 1
              ]
  
          for step_i in range(
              int(episode_start_index), int(episode_end_index)
          ):
              # Set step type.
              if step_i == episode_end_index - 1:
                  step_type = ts.StepType.LAST
              elif step_i == episode_start_index:
                  step_type = ts.StepType.FIRST
              else:
                  step_type = ts.StepType.MID
  
              # Set next state.
              # If at the last step in the episode, create a dummy next step.
              if step_type == ts.StepType.LAST:
                  next_state = np.zeros_like(states[step_i])
                  next_step_type = ts.StepType.FIRST
              else:
                  next_state = states[step_i + 1]
                  next_step_type = (
                      ts.StepType.LAST
                      if step_i == episode_end_index - 2
                      else ts.StepType.MID
                  )
  
              if use_trajectories:
                  sample = create_trajectory(
                      state=states[step_i],
                      action=actions[step_i],
                      discount=discounts[step_i],
                      reward=rewards[step_i],
                      step_type=step_type,
                      next_step_type=next_step_type,
                  )
              else:
                  sample = create_transition(
                      state=states[step_i],
                      action=actions[step_i],
                      next_state=next_state,
                      discount=discounts[step_i],
                      reward=rewards[step_i],
                      step_type=step_type,
                      next_step_type=next_step_type,
                  )
              tfrecord_observer(sample)
  
      tfrecord_observer.close()
      logging.info(
          "Wrote episodes [%d-%d] to %s",
          start_episode,
          end_episode,
          dataset_path,
      )
  
file_path: agents/tf_agents/examples/cql_sac/kumar20/dataset/file_utils.py
imports: from typing import Dict
  from absl import logging
  import d4rl
  import numpy
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.typing import types
  from tf_agents.utils import example_encoding_dataset
label_desc: Use image node_id 3 for calling a global function with example usage: write_samples_to_tfrecord(dataset_dict, collect_data_spec, dataset_path, start_episode, end_episode, use_trajectories) without return types