code: def test_episode_to_timestep_batch_episode_return(self):
      batched = dataset_lib.episode_to_timestep_batch(
          sample_episode(),
          return_horizon=3,
          calculate_episode_return=True,
      )
  
      expected_steps = {
          rlds.OBSERVATION: [[[1, 1], [2, 2], [3, 3]]],
          rlds.ACTION: [[[1], [2], [3]]],
          rlds.REWARD: [[[1.0], [2.0], [3.0]]],
          rlds.DISCOUNT: [[[1], [1], [1]]],
          rlds.IS_FIRST: [[[True], [False], [False]]],
          rlds.IS_LAST: [[[False], [False], [False]]],
          rlds.IS_TERMINAL: [[[False], [False], [False]]],
          dataset_lib.N_STEP_RETURN: [[[6.0], [9.0], [12.0]]],
          # This should match to the sum of the rewards in the input.
          dataset_lib.EPISODE_RETURN: [[[15.0], [15.0], [15.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 2 for calling the DatasetTest obj's underlying member method code with example usage: obj.test_episode_to_timestep_batch_episode_return() without return types