code: def generate_valid_episodes() -> (
      Dict[str, Tuple[tf.data.Dataset, List[trajectory.Trajectory]]]
  ):
      """Get test data for valid RLDS datasets to be used across differrent tests.
  
      Returns:
        A dict representing all valid cases as keys and tuples of valid input RLDS
        and expected TF-Agents trajectories as values.
      """
      complete_steps = tf.data.Dataset.from_tensor_slices(
          {
              rlds_types.OBSERVATION: _OBSERVATIONS,
              rlds_types.ACTION: _ACTIONS,
              rlds_types.REWARD: _REWARDS,
              rlds_types.DISCOUNT: _DISCOUNTS,
              rlds_types.IS_TERMINAL: [False, False, True],
              rlds_types.IS_LAST: [False, False, True],
              rlds_types.IS_FIRST: [True, False, False],
          }
      )
      complete_episode_trajectories = [
          trajectory.Trajectory(0, [1.0, 2.0], [11.0], (), 1, 1.0, 1.0),
          trajectory.Trajectory(1, [3.0, 4.0], [21.0], (), 2, 2.0, 1.0),
          trajectory.Trajectory(2, [5.0, 6.0], [31.0], (), 0, 3.0, 0.0),
      ]
  
      truncated_steps = tf.data.Dataset.from_tensor_slices(
          {
              rlds_types.OBSERVATION: _OBSERVATIONS,
              rlds_types.ACTION: _ACTIONS,
              rlds_types.REWARD: _REWARDS,
              rlds_types.DISCOUNT: _DISCOUNTS,
              rlds_types.IS_TERMINAL: [False, False, False],
              rlds_types.IS_LAST: [False, False, True],
              rlds_types.IS_FIRST: [True, False, False],
          }
      )
      truncated_episode_trajectories = [
          trajectory.Trajectory(0, [1.0, 2.0], [11.0], (), 1, 1.0, 1.0),
          trajectory.Trajectory(1, [3.0, 4.0], [21.0], (), 2, 2.0, 1.0),
          trajectory.Trajectory(2, [5.0, 6.0], [31.0], (), 0, 3.0, 1.0),
      ]
  
      single_step = tf.data.Dataset.from_tensor_slices(
          {
              rlds_types.OBSERVATION: [[1.0, 2.0]],
              rlds_types.ACTION: [[11.0]],
              rlds_types.REWARD: [1.0],
              rlds_types.DISCOUNT: [1.0],
              rlds_types.IS_TERMINAL: [False],
              rlds_types.IS_LAST: [True],
              rlds_types.IS_FIRST: [True],
          }
      )
      single_step_episode_trajectories = [
          trajectory.Trajectory(2, [1.0, 2.0], [11.0], (), 2, 1.0, 1.0),
      ]
  
      extra_param_steps = tf.data.Dataset.from_tensor_slices(
          {
              rlds_types.OBSERVATION: [[1.0, 2.0]],
              rlds_types.ACTION: [[11.0]],
              rlds_types.REWARD: [1.0],
              rlds_types.DISCOUNT: [1.0],
              rlds_types.IS_TERMINAL: [False],
              rlds_types.IS_LAST: [True],
              rlds_types.IS_FIRST: [True],
              _EXTRA_PARAM: [True],
          }
      )
  
      return {
          "complete_episode": (
              tf.data.Dataset.from_tensor_slices(
                  {
                      rlds_types.STEPS: [complete_steps],
                  }
              ),
              complete_episode_trajectories,
          ),
          "truncated_episode": (
              tf.data.Dataset.from_tensor_slices(
                  {
                      rlds_types.STEPS: [truncated_steps],
                  }
              ),
              truncated_episode_trajectories,
          ),
          "single_step_episode": (
              tf.data.Dataset.from_tensor_slices(
                  {
                      rlds_types.STEPS: [single_step],
                  }
              ),
              single_step_episode_trajectories,
          ),
          "extra_param_episode": (
              tf.data.Dataset.from_tensor_slices(
                  {
                      rlds_types.STEPS: [extra_param_steps],
                  }
              ),
              single_step_episode_trajectories,
          ),
          "multiple_episodes": (
              tf.data.Dataset.from_tensor_slices(
                  {
                      rlds_types.STEPS: [
                          complete_steps,
                          single_step,
                          truncated_steps,
                      ],
                  }
              ),
              complete_episode_trajectories
              + single_step_episode_trajectories
              + truncated_episode_trajectories,
          ),
      }
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Dict, List, Tuple
  from absl.testing import parameterized
  import reverb
  from rlds import rlds_types
  import tensorflow
  from tf_agents.replay_buffers import reverb_replay_buffer
  from tf_agents.replay_buffers import reverb_utils
  from tf_agents.replay_buffers import rlds_to_reverb
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import trajectory
  from tf_agents.utils import test_utils
label_desc: Use image node_id 1 for calling a global function with example usage: generate_valid_episodes() and returns: dict