code: def setup_buffer(
      buffer_size, trajectory_lengths, stack_size=None, multi_steps=None
  ):
      """We will insert one trajectory into the RB."""
      stack_size = stack_size if stack_size is not None else 1
      update_horizon = multi_steps if multi_steps is not None else 1
      memory = ReplayBuffer(
          stack_size=stack_size,
          replay_capacity=buffer_size,
          batch_size=1,
          update_horizon=update_horizon,
          return_everything_as_stack=stack_size is not None,
          return_as_timeline_format=multi_steps is not None,
      )
  
      i = 0
      for traj_len in trajectory_lengths:
          for j in range(traj_len):
              trans = get_add_transition(i)
              terminal = bool(j == traj_len - 1)
              memory.add(
                  observation=trans["state"],
                  action=trans["action"],
                  reward=trans["reward"],
                  terminal=terminal,
                  extra1=trans["extra1"],
              )
              i += 1
      return memory.sample_all_valid_transitions()
  
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 3 for calling a global function with example usage: setup_buffer(buffer_size, trajectory_lengths, stack_size, multi_steps) and returns: memory