code: def test_episode_to_timestep_batch_no_return_horizon(self):
      batched = dataset_lib.episode_to_timestep_batch(
          sample_episode(), return_horizon=1
      )
  
      expected_steps = {
          rlds.OBSERVATION: [
              [[1, 1], [2, 2], [3, 3]],
              [[2, 2], [3, 3], [4, 4]],
              [[3, 3], [4, 4], [5, 5]],
          ],
          rlds.ACTION: [
              [[1], [2], [3]],
              [[2], [3], [4]],
              [[3], [4], [5]],
          ],
          rlds.REWARD: [
              [[1.0], [2.0], [3.0]],
              [[2.0], [3.0], [4.0]],
              [[3.0], [4.0], [5.0]],
          ],
          rlds.DISCOUNT: [
              [[1], [1], [1]],
              [[1], [1], [1]],
              [[1], [1], [1]],
          ],
          rlds.IS_FIRST: [
              [[True], [False], [False]],
              [[False], [False], [False]],
              [[False], [False], [False]],
          ],
          rlds.IS_LAST: [
              [[False], [False], [False]],
              [[False], [False], [False]],
              [[False], [False], [True]],
          ],
          rlds.IS_TERMINAL: [
              [[False], [False], [False]],
              [[False], [False], [False]],
              [[False], [False], [True]],
          ],
          # n-step return should be equal to the rewards.
          dataset_lib.N_STEP_RETURN: [
              [[1.0], [2.0], [3.0]],
              [[2.0], [3.0], [4.0]],
              [[3.0], [4.0], [5.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 3 for calling the DatasetTest obj's underlying member method code with example usage: obj.test_episode_to_timestep_batch_no_return_horizon() without return types