code: def generic_stack_multi_steps_test_helper(
      buffer_size, trajectory_lengths, stack_size, multi_steps
  ):
      batch = setup_buffer(
          buffer_size,
          trajectory_lengths,
          stack_size=stack_size,
          multi_steps=multi_steps,
      )
  
      # start with state, action, extra1 (these are single)
      expected = {k: [] for k in ["state", "action", "extra1"]}
      terminal_array = []
  
      actual_rb_index = stack_size - 1
      i = 0
      for traj_len in trajectory_lengths:
          traj_start = i
          for j in range(traj_len):
              cur = get_stacked_transition(i, stack_size, traj_start)
              for k in expected:
                  expected[k].append(cur[k])
              terminal_array.append(bool(j >= traj_len - multi_steps))
  
              actual_rb_index += 1
              i += 1
          actual_rb_index += stack_size - 1
      assert (
          actual_rb_index <= buffer_size
      ), f"{actual_rb_index} is larger than {buffer_size}"
  
      expected["terminal"] = np.expand_dims(terminal_array, axis=1)
      for k in expected:
          expected[k] = torch.tensor(expected[k])
      batch_size = expected["state"].shape[0]
  
      for k in expected:
          batch_val = getattr(batch, k)
          npt.assert_array_equal(
              batch_val,
              expected[k],
              err_msg=f"key {k}; expected {expected[k]}, got {batch_val}",
              verbose=True,
          )
  
      # now examine reward, next_state, next_action, next_extra1, which would be lists of size step
      expected = {
          k: []
          for k in [
              "reward",
              "next_state",
              "next_action",
              "next_extra1",
          ]
      }
      i = 0
      for traj_len in trajectory_lengths:
          traj_start = i
          for _ in range(traj_len):
              multistep_trans = {k: [] for k in expected}
  
              # rewards start at current
              traj_end = traj_start + traj_len
              for j in range(i, i + multi_steps):
                  if j < traj_end:
                      stacked_trans = get_stacked_transition(
                          j, stack_size, traj_start
                      )
                      multistep_trans["reward"].append(
                          stacked_trans["reward"]
                      )
  
              # next features start at current + 1
              for j in range(i + 1, i + multi_steps + 1):
                  if j <= traj_end:
                      stacked_trans = get_stacked_transition(
                          j, stack_size, traj_start
                      )
                      for k in [
                          "next_state",
                          "next_action",
                          "next_extra1",
                      ]:
                          stripped_k = k[len("next_") :]
                          multistep_trans[k].append(
                              stacked_trans[stripped_k]
                          )
  
              multistep_trans = {
                  k: torch.tensor(v) for k, v in multistep_trans.items()
              }
              for k in expected:
                  expected[k].append(multistep_trans[k])
  
              i += 1
  
      # validate the lengths
      for k in expected:
          batch_val = getattr(batch, k)
          assert isinstance(
              batch_val, list
          ), f"batch[{k}] has type {type(batch_val)}"
          assert len(batch_val) == batch_size
          for i in range(batch_size):
              assert isinstance(
                  batch_val[i], torch.Tensor
              ), f"batch[{k}][{i}] has type {type(batch_val[i])};\n{batch_val}"
              assert batch_val[i].shape[0] == batch.step[i], (
                  f"batch[{k}][{i}] {batch_val[i].shape} doesn't start "
                  f"with {batch.step[i]};\n{batch_val}"
              )
              # sanity check
              assert len(expected[k][i]) == batch.step[i], (
                  f"expected[{k}][{i}] {expected[k][i]} with len {len(expected[k][i])} should have len "
                  f"{batch.step[i]};\n{expected[k]}, {batch_val}"
              )
  
      for k in expected:
          batch_val = getattr(batch, k)
          for i in range(batch_size):
              # NOTE: the last transition for terminals is undefined
              A = batch_val[i]
              B = expected[k][i]
              if batch.terminal[i]:
                  A = A[:-1]
                  B = B[:-1]
  
              npt.assert_array_equal(A, B)
  
file_path: ReAgent/reagent/test/replay_memory/extra_replay_buffer_test.py
imports: import logging
  import numpy
  import numpy.testing
  import torch
  from reagent.replay_memory.circular_replay_buffer import ReplayBuffer
  from reagent.test.base.horizon_test_base import HorizonTestBase
label_desc: Use image node_id 5 for calling a global function with example usage: generic_stack_multi_steps_test_helper(buffer_size, trajectory_lengths, stack_size, multi_steps) without return types