code: def main(_):
      logging.set_verbosity(logging.INFO)
  
      d4rl_env = gym.make(FLAGS.env_name)
      d4rl_dataset = d4rl_env.get_dataset()
      root_dir = os.path.join(FLAGS.root_dir, FLAGS.env_name)
  
      dataset_dict = dataset_utils.create_episode_dataset(
          d4rl_dataset,
          FLAGS.exclude_timeouts,
          observation_dtype=d4rl_env.observation_space.dtype,
      )
      num_episodes = len(dataset_dict["episode_start_index"])
      logging.info(
          "Found %d episodes, %s total steps.",
          num_episodes,
          len(dataset_dict["states"]),
      )
  
      collect_data_spec = dataset_utils.create_collect_data_spec(
          dataset_dict, use_trajectories=FLAGS.use_trajectories
      )
      logging.info("Collect data spec %s", collect_data_spec)
  
      num_replicas = FLAGS.replicas or 1
      interval_size = num_episodes // num_replicas + 1
  
      # If FLAGS.replica_id is set, only run that section of the dataset.
      # This is useful if distributing the replicas on Borg.
      if FLAGS.replica_id is not None:
          file_name = "%s_%d.tfrecord" % (
              FLAGS.env_name,
              FLAGS.replica_id,
          )
          start_index = FLAGS.replica_id * interval_size
          end_index = min(
              (FLAGS.replica_id + 1) * interval_size, num_episodes
          )
          file_utils.write_samples_to_tfrecord(
              dataset_dict=dataset_dict,
              collect_data_spec=collect_data_spec,
              dataset_path=os.path.join(root_dir, file_name),
              start_episode=start_index,
              end_episode=end_index,
              use_trajectories=FLAGS.use_trajectories,
          )
      else:
          # Otherwise, parallelize with tf_agents.system.multiprocessing.
          jobs = []
          context = multiprocessing.get_context()
  
          for i in range(num_replicas):
              if num_replicas == 1:
                  file_name = "%s.tfrecord" % FLAGS.env_name
              else:
                  file_name = "%s_%d.tfrecord" % (FLAGS.env_name, i)
              dataset_path = os.path.join(root_dir, file_name)
              start_index = i * interval_size
              end_index = min((i + 1) * interval_size, num_episodes)
              kwargs = dict(
                  dataset_dict=dataset_dict,
                  collect_data_spec=collect_data_spec,
                  dataset_path=dataset_path,
                  start_episode=start_index,
                  end_episode=end_index,
                  use_trajectories=FLAGS.use_trajectories,
              )
              job = context.Process(  # pytype: disable=attribute-error  # re-none
                  target=file_utils.write_samples_to_tfrecord,
                  kwargs=kwargs,
              )
              job.start()
              jobs.append(job)
  
          for job in jobs:
              job.join()
  
file_path: agents/tf_agents/examples/cql_sac/kumar20/dataset/dataset_generator.py
imports: import functools
  import os
  from absl import app
  from absl import flags
  from absl import logging
  import d4rl
  import gym
  from tf_agents.examples.cql_sac.kumar20.dataset import dataset_utils
  from tf_agents.examples.cql_sac.kumar20.dataset import file_utils
  from tf_agents.system import system_multiprocessing
label_desc: Use image node_id 1 for calling a global function with example usage: main(_) without return types