code: def test_episode_to_timestep_batch(self):
      batched = dataset_lib.episode_to_timestep_batch(
          sample_episode(), return_horizon=2
      )
  
      # Scalars should be expanded and the n-step return should be present. Each
      # element of a step should be a triplet containing the previous, current and
      # next values of the corresponding fields. Since the return horizon is 2 and
      # the number of steps in the episode is 5, there can be only 2 triplets for
      # time steps 1 and 2.
      expected_steps = {
          rlds.OBSERVATION: [
              [[1, 1], [2, 2], [3, 3]],
              [[2, 2], [3, 3], [4, 4]],
          ],
          rlds.ACTION: [
              [[1], [2], [3]],
              [[2], [3], [4]],
          ],
          rlds.REWARD: [
              [[1.0], [2.0], [3.0]],
              [[2.0], [3.0], [4.0]],
          ],
          rlds.DISCOUNT: [
              [[1], [1], [1]],
              [[1], [1], [1]],
          ],
          rlds.IS_FIRST: [
              [[True], [False], [False]],
              [[False], [False], [False]],
          ],
          rlds.IS_LAST: [
              [[False], [False], [False]],
              [[False], [False], [False]],
          ],
          rlds.IS_TERMINAL: [
              [[False], [False], [False]],
              [[False], [False], [False]],
          ],
          dataset_lib.N_STEP_RETURN: [
              [[3.0], [5.0], [7.0]],
              [[5.0], [7.0], [9.0]],
          ],
      }
  
      self.expect_equal_datasets(
          batched, tf.data.Dataset.from_tensor_slices(expected_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 the DatasetTest obj's underlying member method code with example usage: obj.test_episode_to_timestep_batch() without return types