code: def create_collect_data_spec(
      dataset_dict: EpisodeDictType, use_trajectories: bool = True
  ) -> Union[trajectory.Transition, trajectory.Trajectory]:
      """Create a spec that describes the data collected by agent.collect_policy."""
      reward = dataset_dict["rewards"][0]
      discount = dataset_dict["discounts"][0]
      observation = dataset_dict["states"][0]
      action = dataset_dict["actions"][0]
      step_type = np.asarray(0, dtype=np.int32)
  
      if use_trajectories:
          return trajectory.Trajectory(
              step_type=ArraySpec(
                  shape=step_type.shape, dtype=step_type.dtype
              ),
              observation=ArraySpec(
                  shape=observation.shape, dtype=observation.dtype
              ),
              action=ArraySpec(shape=action.shape, dtype=action.dtype),
              policy_info=(),
              next_step_type=ArraySpec(
                  shape=step_type.shape, dtype=step_type.dtype
              ),
              reward=ArraySpec(shape=reward.shape, dtype=reward.dtype),
              discount=ArraySpec(
                  shape=discount.shape, dtype=discount.dtype
              ),
          )
      else:
          time_step_spec = time_step.TimeStep(
              step_type=ArraySpec(
                  shape=step_type.shape, dtype=step_type.dtype
              ),
              reward=ArraySpec(shape=reward.shape, dtype=reward.dtype),
              discount=ArraySpec(
                  shape=discount.shape, dtype=discount.dtype
              ),
              observation=ArraySpec(
                  shape=observation.shape, dtype=observation.dtype
              ),
          )
          action_spec = policy_step.PolicyStep(
              action=ArraySpec(shape=action.shape, dtype=action.dtype),
              state=(),
              info=(),
          )
          return trajectory.Transition(
              time_step=time_step_spec,
              action_step=action_spec,
              next_time_step=time_step_spec,
          )
  
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 1 for calling a global function with example usage: create_collect_data_spec(dataset_dict, use_trajectories) and returns: trajectory, trajectory