code: def test_write_trajectories_to_tfrecord(
      self, mock_close, mock_call, mock_init
  ):
      mock_init.return_value = None
      mock_close.return_value = None
      episode_dict = {
          "states": np.array(
              [
                  [1.0, 2.0],
                  [3.0, 4.0],
                  [5.0, 6.0],
                  [7.0, 8.0],
                  [9.0, 10.0],
                  [11.0, 12.0],
                  [13.0, 14.0],
              ]
          ),
          "actions": np.array(
              [[1.0], [2.0], [3.0], [4.0], [5.0], [6.0], [7.0]]
          ),
          "rewards": np.array(
              [[0.0], [1.0], [0.0], [1.0], [0.0], [0.0], [1.0]]
          ),
          "discounts": np.array([1.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0]),
          "episode_start_index": np.array([0, 2, 4]),
      }
  
      first_trajectory = create_trajectory(
          state=np.array([5.0, 6.0]),
          action=np.array([3.0]),
          discount=np.array(1.0),
          reward=np.array([0.0]),
          step_type=np.array(0),
          next_step_type=np.array(2),
      )
  
      write_samples_to_tfrecord(
          episode_dict,
          first_trajectory,  # used as a dummy collect_data_spec
          dataset_path="dataset_path",
          start_episode=1,
          end_episode=3,
          use_trajectories=True,
      )
  
      expected_trajectories = [
          first_trajectory,
          create_trajectory(
              state=np.array([7.0, 8.0]),
              action=np.array([4.0]),
              discount=np.array(0.0),
              reward=np.array([1.0]),
              step_type=np.array(2),
              next_step_type=np.array(0),
          ),
          create_trajectory(
              state=np.array([9.0, 10.0]),
              action=np.array([5.0]),
              discount=np.array(1.0),
              reward=np.array([0.0]),
              step_type=np.array(0),
              next_step_type=np.array(1),
          ),
          create_trajectory(
              state=np.array([11.0, 12.0]),
              action=np.array([6.0]),
              discount=np.array(1.0),
              reward=np.array([0.0]),
              step_type=np.array(1),
              next_step_type=np.array(2),
          ),
          create_trajectory(
              state=np.array([13.0, 14.0]),
              action=np.array([7.0]),
              discount=np.array(0.0),
              reward=np.array([1.0]),
              step_type=np.array(2),
              next_step_type=np.array(0),
          ),
      ]
  
      # Check that the trajectories passed to the mock TFRecordObserver
      # match what we expect.
      num_trajectories = len(expected_trajectories)
      self.assertEqual(len(mock_call.call_args_list), num_trajectories)
      for i in range(num_trajectories):
          actual_traj = mock_call.call_args_list[i][0][0]
          expected_traj = expected_trajectories[i]
  
          self.assertAllEqual(
              actual_traj.step_type, expected_traj.step_type
          )
          self.assertAllEqual(
              actual_traj.observation, expected_traj.observation
          )
          self.assertAllEqual(actual_traj.action, expected_traj.action)
          self.assertAllEqual(
              actual_traj.policy_info, expected_traj.policy_info
          )
          self.assertAllEqual(
              actual_traj.next_step_type, expected_traj.next_step_type
          )
          self.assertAllEqual(actual_traj.reward, expected_traj.reward)
          self.assertAllEqual(
              actual_traj.discount, expected_traj.discount
          )
  
file_path: agents/tf_agents/examples/cql_sac/kumar20/dataset/file_utils_test.py
imports: import mock
  import numpy
  import tensorflow
  from tf_agents.examples.cql_sac.kumar20.dataset.file_utils import create_trajectory
  from tf_agents.examples.cql_sac.kumar20.dataset.file_utils import create_transition
  from tf_agents.examples.cql_sac.kumar20.dataset.file_utils import write_samples_to_tfrecord
  from tf_agents.utils import test_utils
label_desc: Use image node_id 2 for calling the FileUtilsTest obj's underlying member method code with example usage: obj.test_write_trajectories_to_tfrecord(mock_close, mock_call, mock_init) without return types