code: def create_trajectory_data_spec(
      rlds_data: tf.data.Dataset,
  ) -> trajectory.Trajectory:
      """Creates data spec for initializing Reverb server and Reverb Replay Buffer.
  
      Creates a data spec for the corresponding trajectory dataset that can be
      created using the rlds_data provided as input. This data spec is necessary for
      initializing Reverb server and Reverb Replay Buffer.
  
      Args:
        rlds_data: An RLDS dataset is tf.data.Dataset of RLDS episodes, where each
          episode contains a tf.data.Dataset of RLDS steps. An RLDS step is a
          dictionary of tensors containing is_first, is_last, observation, action,
          reward, is_terminal, and discount.
  
      Returns:
        A trajectory representing tensor specs for trajectory dataset that can be
        created using the rlds_data provided as input.
  
      Raises:
        ValueError: If no RLDS steps exist in rlds_data.
        ValueError: If step spec for rlds_data is not valid.
      """
      _validate_rlds_episode_spec(rlds_data)
      rlds_data = rlds_data.flat_map(
          lambda episode: episode[rlds.STEPS]
      )
      element_spec = rlds_data.element_spec
      _validate_rlds_step_spec(element_spec)
      time_step_spec = ts.time_step_spec(
          element_spec[rlds.OBSERVATION], element_spec[rlds.REWARD]
      )
      action_spec = policy_step.PolicyStep(element_spec[rlds.ACTION])
      return trajectory.from_transition(
          time_step_spec, action_spec, time_step_spec
      )
  
file_path: agents/tf_agents/replay_buffers/rlds_to_reverb.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Callable, Dict, List, Optional, Tuple, Union
  from absl import logging
  import numpy
  import rlds
  import tensorflow
  from tf_agents.replay_buffers import reverb_utils
  from tf_agents.specs import tensor_spec
  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 4 for calling a global function with example usage: create_trajectory_data_spec(rlds_data) and returns: trajectory