code: def test_collect_data_spec_trajectory(self):
      episode_dict = {
          "states": np.array(
              [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]],
              dtype=np.float32,
          ),
          "actions": np.array(
              [[1.0], [2.0], [3.0], [4.0]], dtype=np.float32
          ),
          "rewards": np.array(
              [[0.0], [1.0], [0.0], [1.0]], dtype=np.float32
          ),
          "discounts": np.array([1.0, 0.0, 1.0, 0.0], dtype=np.float32),
          "episode_start_index": np.array([0, 2], dtype=np.int32),
      }
  
      expected_spec = trajectory.Trajectory(
          step_type=ArraySpec(shape=[], dtype=np.int32),
          observation=ArraySpec(shape=[2], dtype=np.float32),
          action=ArraySpec(shape=[1], dtype=np.float32),
          policy_info=(),
          next_step_type=ArraySpec(shape=[], dtype=np.int32),
          reward=ArraySpec(shape=[1], dtype=np.float32),
          discount=ArraySpec(shape=[], dtype=np.float32),
      )
      actual_spec = create_collect_data_spec(
          episode_dict, use_trajectories=True
      )
      self.assertEqual(actual_spec, expected_spec)
  
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 7 for calling the UtilsTest obj's underlying member method code with example usage: obj.test_collect_data_spec_trajectory() without return types