code: def convert_rlds_to_trajectories(
      rlds_data: tf.data.Dataset, policy_info_fn: _PolicyFnType = None
  ) -> tf.data.Dataset:
      """Converts the RLDS data to a dataset of trajectories.
  
      Converts the rlds_data provided to a dataset of TF Agents trajectories by
      flattening and converting it into batches and then tuples of overlapping pairs
      of adjacent RLDS steps.
  
      An end step of first step type is padded to the RLDS data to ensure that the
      trajectory created using the last step of the last episode has a valid next
      step type.
  
        Input    ------------------------------------------->       Output
  
       RLDS Data --> Flatten -->   Overlapping Pairs     -->  TF Agents trajectories
      -----------    --------   ------------------------      ---------------------
       episode 1     step 1      (step 1, step 2)                Trajectory 1
        --------     --------   ------------------------      ---------------------
         step 1      step 2      (step 2, step 3)                Trajectory 2
         step 2      --------   ------------------------      ---------------------
         step 3      step 3      (step 3, step 4)                Trajectory 3
        --------     --------   ------------------------      ---------------------
      -----------    step 4      (step 4, step 5)                Trajectory 4
       episode 2     --------   ------------------------      ---------------------
        --------     step 5      (step 5, step 6)                Trajectory 5
         step 4      --------   ------------------------      ---------------------
         step 5      step 6      (step 6, padded end step)       Trajectory 6
         step 6      --------   ------------------------      ---------------------
        --------
  
      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.
        policy_info_fn: An optional function to create some policy.info to be used
          while generating TF-Agents trajectories.
  
      Returns:
        A dataset of type tf.data.Dataset, elements of which are TF Agents
        trajectories corresponding to the RLDS steps provided in rlds_data.
  
      Raises:
        ValueError: If no RLDS steps exist in rlds_data.
        ValueError: If step spec for rlds_data is not valid.
        InvalidArgumentError: If the RLDS dataset provided has episodes that:
          - Incorrectly end i.e. does not end in last step.
          - Incorrectly terminate i.e. a terminal step is not the last step.
          - Incorrectly begin i.e.  a last step in not followed by the first step.
            Please note that the last step of the last episode is taken care of
            in the function and user does not need to make sure that the last step
            of the last episode is followed by a first step.
      """
  
      def _pair_to_tuple(
          pair: Dict[str, tf.Tensor]
      ) -> Tuple[Dict[str, tf.Tensor], Dict[str, tf.Tensor]]:
          """Converts a batch of two adjacent RLDS steps to a Tuple of RLDS steps."""
          current_step = {}
          next_step = {}
          for key, value in pair.items():
              if isinstance(value, dict):
                  current_step[key] = value
                  next_step[key] = value
              else:
                  current_step[key] = value[0]
                  next_step[key] = value[1]
          return (current_step, next_step)
  
      def _get_step_type(step: Dict[str, tf.Tensor]) -> np.ndarray:
          if step[rlds.IS_FIRST]:
              return ts.StepType.FIRST
          elif step[rlds.IS_LAST]:
              return ts.StepType.LAST
          else:
              return ts.StepType.MID
  
      def _is_complete(step: Dict[str, tf.Tensor]) -> tf.Tensor:
          return step[rlds.IS_LAST] and step[rlds.IS_TERMINAL]
  
      def _get_discount(step: Dict[str, tf.Tensor]) -> tf.Tensor:
          """Returns 0 for complete episode, else returns the current discount."""
          return (
              tf.constant(0.0, dtype=tf.float32)
              if _is_complete(step)
              else step[rlds.DISCOUNT]
          )
  
      def _validate_step(
          current_step: Dict[str, tf.Tensor],
          next_step: Dict[str, tf.Tensor],
      ) -> None:
          """Validates a tuple of adjacent RLDS steps."""
  
          # This check validates any incorrectly terminated episodes in RLDS data.
          if current_step[rlds.IS_TERMINAL]:
              tf.Assert(
                  current_step[rlds.IS_LAST],
                  [
                      "Terminal step must be the last step of an episode."
                  ],
              )
  
          # This check validates any incorrectly ending episodes in the RLDS dataset.
          #
          # Please note that the last step of last episode will also be validated by
          # this check since we pad the RLDS data with a step of first step type.
          if (
              not current_step[rlds.IS_LAST]
              and not current_step[rlds.IS_FIRST]
          ):
              tf.Assert(
                  not next_step[rlds.IS_FIRST],
                  ["Mid step should not be followed by a first step."],
              )
  
          # This check validates any episodes that begin incorrectly in the RLDS data.
          if current_step[rlds.IS_LAST]:
              tf.Assert(
                  next_step[rlds.IS_FIRST],
                  [
                      "Last step of an episode must be followed by a first step."
                  ],
              )
  
      def _to_trajectory(
          current_step: Dict[str, tf.Tensor],
          next_step: Dict[str, tf.Tensor],
      ) -> trajectory.Trajectory:
          """Converts a tuple of adjacent RLDS steps to a trajectory."""
          _validate_step(current_step, next_step)
          step_type = _get_step_type(current_step)
          next_step_type = _get_step_type(next_step)
  
          policy_info = (
              policy_info_fn(current_step, next_step)
              if policy_info_fn
              else ()
          )
  
          # Represent single step episode with step_type=LAST and next_step_type=LAST.
          if (
              step_type == ts.StepType.FIRST
              and next_step_type == ts.StepType.FIRST
          ):
              step_type = ts.StepType.LAST
              next_step_type = ts.StepType.LAST
  
          return trajectory.Trajectory(
              step_type=step_type,
              observation=current_step[rlds.OBSERVATION],
              action=current_step[rlds.ACTION],
              policy_info=policy_info,
              next_step_type=next_step_type,
              reward=current_step[rlds.REWARD],
              discount=_get_discount(current_step),
          )
  
      _validate_rlds_episode_spec(rlds_data)
      rlds_data = rlds_data.flat_map(
          lambda episode: episode[rlds.STEPS]
      )
      _validate_rlds_step_spec(rlds_data.element_spec)
  
      # Pad a step of first step type at the end of the RLDS data to ensure that the
      # trajectory created using the last step of the last episode has a valid next
      # step type.
      rlds_data = rlds.transformations.concatenate(
          rlds_data, tf.data.Dataset.from_tensors({rlds.IS_FIRST: True})
      )
  
      # Batch the RLDS data as pair of adjacent steps with an overlap of one step
      # and then create tuples of those batches to finally to create TF Agents
      # trajectories with correct next step type.
      return (
          rlds.transformations.batch(
              dataset=rlds_data, size=2, shift=1, drop_remainder=True
          )
          .map(_pair_to_tuple)
          .map(_to_trajectory)
      )
  
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 5 for calling a global function with example usage: convert_rlds_to_trajectories(rlds_data, policy_info_fn) and returns: rlds, current_step,  next_step, unknown, unknown, trajectory, ts, ts, ts