def make_replay_buffer(
    env,
    capacity: int = 1000000,
    rlds_logger_path: Optional[str] = None,
    type: str = "replay_buffer",
    image_keys: list = [],  # used only type=="memory_efficient_replay_buffer"
    preload_rlds_path: Optional[str] = None,
    preload_data_transform: Optional[callable] = None,
):
    """
    This is the high-level helper function to
    create a replay buffer for the given environment.

    Args:
    - env: gym or gymasium environment
    - capacity: capacity of the replay buffer
    - rlds_logger_path: path to save RLDS logs
    - type: support only for "replay_buffer" and "memory_efficient_replay_buffer"
    - image_keys: list of image keys, used only "memory_efficient_replay_buffer"
    - preload_rlds_path: path to preloaded RLDS trajectories
    - preload_data_transform: data transformation function for preloaded RLDS data
    """
    print("shape of observation space and action space")
    print(env.observation_space)
    print(env.action_space)

    # init logger for RLDS
    if rlds_logger_path:
        # from: https://github.com/rail-berkeley/oxe_envlogger
        from oxe_envlogger.rlds_logger import RLDSLogger

        rlds_logger = RLDSLogger(
            observation_space=env.observation_space,
            action_space=env.action_space,
            dataset_name="serl_rlds_dataset",
            directory=rlds_logger_path,
            max_episodes_per_file=5,  # TODO: arbitrary number
        )
    else:
        rlds_logger = None

    if type == "replay_buffer":
        replay_buffer = ReplayBufferDataStore(
            env.observation_space,
            env.action_space,
            capacity=capacity,
            rlds_logger=rlds_logger,
        )
    elif type == "memory_efficient_replay_buffer":
        replay_buffer = MemoryEfficientReplayBufferDataStore(
            env.observation_space,
            env.action_space,
            capacity=capacity,
            rlds_logger=rlds_logger,
            image_keys=image_keys,
        )
    else:
        raise ValueError(f"Unsupported replay_buffer_type: {type}")

    if preload_rlds_path:
        print(f" - Preloaded {preload_rlds_path} to replay buffer")
        dataset = tfds.builder_from_directory(preload_rlds_path).as_dataset(split="all")
        populate_datastore(
            replay_buffer,
            dataset,
            data_transform=preload_data_transform,
            type="with_dones",
        )
        print(f" - done populated {len(replay_buffer)} samples to replay buffer")

    return replay_buffer
This is the high-level helper function to
create a replay buffer for the given environment.

Args:
- env: gym or gymasium environment
- capacity: capacity of the replay buffer
- rlds_logger_path: path to save RLDS logs
- type: support only for "replay_buffer" and "memory_efficient_replay_buffer"
- image_keys: list of image keys, used only "memory_efficient_replay_buffer"
- preload_rlds_path: path to preloaded RLDS trajectories
- preload_data_transform: data transformation function for preloaded RLDS data
serl_launcher/serl_launcher/utils/launcher.py