code: def sample_episode() -> rlds.Episode:
      """Returns a sample episode."""
      steps = {
          rlds.OBSERVATION: [
              [1, 1],
              [2, 2],
              [3, 3],
              [4, 4],
              [5, 5],
          ],
          rlds.ACTION: [[1], [2], [3], [4], [5]],
          rlds.REWARD: [1.0, 2.0, 3.0, 4.0, 5.0],
          rlds.DISCOUNT: [1, 1, 1, 1, 1],
          rlds.IS_FIRST: [True, False, False, False, False],
          rlds.IS_LAST: [False, False, False, False, True],
          rlds.IS_TERMINAL: [False, False, False, False, True],
      }
      return {rlds.STEPS: tf.data.Dataset.from_tensor_slices(steps)}
  
file_path: acme/acme/agents/jax/mbop/dataset_test.py
imports: from acme.agents.jax.mbop import dataset
  import rlds
  from rlds.transformations import transformations_testlib
  import tensorflow
  from absl.testing import absltest
label_desc: Use image node_id 1 for calling a global function with example usage: sample_episode() and returns: dict