code: def _load_test_rlds(dataset_name: str) -> tf.data.Dataset:
      count = 20
  
      def as_dataset(*args, **kwargs):
          del args, kwargs
          obs_dim = 29
          act_dim = 8
          steps = {
              rlds.OBSERVATION: [
                  [float(i)] * obs_dim for i in range(count + 1)
              ],
              rlds.ACTION: [
                  [float(i)] * act_dim for i in range(count + 1)
              ],
              rlds.REWARD: [float(j) for j in list(range(count))]
              + [0.0],
              rlds.DISCOUNT: [float(1)] * (count + 1),
              rlds.IS_FIRST: [i == 0 for i in range(count)] + [True],
              rlds.IS_LAST: [i == count - 1 for i in range(count + 1)],
              rlds.IS_TERMINAL: [False for i in range(count + 1)],
          }
          episode = {
              rlds.STEPS: [tf.data.Dataset.from_tensor_slices(steps)]
          }
          return tf.data.Dataset.from_tensor_slices(episode)
  
      with tfds.testing.mock_data(
          as_dataset_fn=as_dataset, num_examples=count
      ):
          return rlds.load(dataset_name)
  
file_path: agents/tf_agents/examples/cql_sac/kumar20/cql_sac_train_eval_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from gym import error
  import rlds
  import tensorflow
  import tensorflow_datasets
  from tf_agents.examples.cql_sac.kumar20 import cql_sac_train_eval
label_desc: Use image node_id 1 for calling a global function with example usage: _load_test_rlds(dataset_name) and returns: tf, rlds