code: def test_episode_to_timestep_batch_drop_return_horizon(self):
      steps = {
          rlds.OBSERVATION: [[1], [2], [3], [4], [5], [6]],
          rlds.REWARD: [1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
      }
      episode = {rlds.STEPS: tf.data.Dataset.from_tensor_slices(steps)}
  
      batched = dataset_lib.episode_to_timestep_batch(
          episode,
          return_horizon=2,
          calculate_episode_return=True,
          drop_return_horizon=True,
      )
  
      # The two steps of the episode should be dropped. There will be 4 steps left
      # and since the return horizon is 2, only a single 3-batched step should be
      # emitted. The episode return should be the sum of the rewards of the first
      # 4 steps.
      expected_steps = {
          rlds.OBSERVATION: [[[1], [2], [3]]],
          rlds.REWARD: [[[1.0], [2.0], [3.0]]],
          dataset_lib.N_STEP_RETURN: [[[3.0], [5.0], [7.0]]],
          dataset_lib.EPISODE_RETURN: [[[10.0], [10.0], [10.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 4 for calling the DatasetTest obj's underlying member method code with example usage: obj.test_episode_to_timestep_batch_drop_return_horizon() without return types