Replay buffers#

Quick intro to replay buffers in RL#

In reinforcement learning (RL), a replay buffer stores and replays experiences that agents collect from interactions with the environment. In Python, you can implement a simple buffer as a list that you add elements to and later sample from. Off-policy learning algorithms use these buffers most. This makes intuitive sense because these algorithms can learn from experiences in the buffer that a previous version of the policy, or even a completely different behavior policy, produced.

Sampling strategy#

When you sample from a replay buffer, you choose which experiences to train your agent with. A straightforward strategy that works well for many algorithms is to pick samples uniformly at random. A more advanced strategy, which performs better in many cases, is Prioritized Experience Replay (PER). PER assigns each item in the buffer a scalar priority value that denotes its significance, or how much you expect to learn from it. PER samples experiences with a higher priority more often.

Eviction strategy#

A buffer has a limited capacity to hold experiences. As an algorithm runs, a buffer eventually reaches its capacity, and to make room for new experiences, it deletes, or evicts, older ones. Eviction generally happens on a first-in, first-out basis. For your algorithms, this means a buffer with a high capacity can learn from older samples, while a smaller buffer makes the learning process more on-policy. Buffers that implement reservoir sampling are an exception to this strategy.

Replay buffers in RLlib#

RLlib comes with a set of extendable replay buffers built in. All of them support the two basic methods add() and sample(). RLlib provides a base ReplayBuffer class that you build your own buffer from. Most algorithms require MultiAgentReplayBuffers so that they generalize to the multi-agent case. These buffers’ add() and sample() methods require a policy_id to handle experiences per policy. See the MultiAgentReplayBuffer for how it extends the base class. You can find buffer types and arguments to modify their behavior in RLlib’s default parameters, as part of the replay_buffer_config.

Basic usage#

When running an experiment, you rarely define your own replay buffer subclass. Instead, you configure existing buffers. The following example from RLlib’s examples section runs the R2D2 algorithm with PER, which R2D2 doesn’t use by default. The highlighted lines focus on the PER configuration.

Tip

Because PER is so common, most Q-learning algorithms support it. Their training iteration functions embed the required priority update step.

Warning

If your custom buffer requires extra interaction, you also have to change the training iteration function.

Specifying a buffer type works the same way as specifying an exploration type. The following example shows three ways to specify a type:

Apart from the type, you can specify the capacity and other parameters. These parameters are mostly constructor arguments for the buffer. They fall into three categories:

Parameters that define how algorithms interact with replay buffers. For example, worker_side_prioritization decides where to compute priorities.
Constructor arguments that instantiate the replay buffer. For example, capacity limits the buffer’s size.
Call arguments for underlying replay buffer methods. For example, the MultiAgentPrioritizedReplayBuffer uses prioritized_replay_beta to call the sample() method of every underlying PrioritizedReplayBuffer.

Tip

Most of the time, only the first two categories are of interest. The third is an advanced feature that supports use cases where a MultiAgentReplayBuffer instantiates underlying buffers that need constructor or default call arguments.

ReplayBuffer base class#

The base ReplayBuffer class only supports storing and replaying experiences in different StorageUnits. Add data to the buffer’s storage with the add() method, and replay it with the sample() method. Advanced buffer types add features while trying to retain compatibility through inheritance. The following example shows the most basic scheme of interaction with a ReplayBuffer.

# We choose fragments because it does not impose restrictions on our batch to be added
buffer = ReplayBuffer(capacity=2, storage_unit=StorageUnit.FRAGMENTS)
dummy_batch = SampleBatch({"a": [1], "b": [2]})
buffer.add(dummy_batch)
buffer.sample(2)
# Because elements can be sampled multiple times, we receive a concatenated version
# of dummy_batch `{a: [1, 1], b: [2, 2,]}`.

Build your own ReplayBuffer#

The following example implements a toy ReplayBuffer class and makes SimpleQ use it:

class LessSampledReplayBuffer(ReplayBuffer):
    @override(ReplayBuffer)
    def sample(
        self, num_items: int, evict_sampled_more_then: int = 30, **kwargs
    ) -> Optional[SampleBatchType]:
        """Evicts experiences that have been sampled > evict_sampled_more_then times."""
        idxes = [random.randint(0, len(self) - 1) for _ in range(num_items)]
        often_sampled_idxes = list(
            filter(lambda x: self._hit_count[x] >= evict_sampled_more_then, set(idxes))
        )

        sample = self._encode_sample(idxes)
        self._num_timesteps_sampled += sample.count

        for idx in often_sampled_idxes:
            del self._storage[idx]
            self._hit_count = np.append(
                self._hit_count[:idx], self._hit_count[idx + 1 :]
            )

        return sample


config = (
    DQNConfig()
    .api_stack(
        enable_env_runner_and_connector_v2=False, enable_rl_module_and_learner=False
    )
    .environment(env="CartPole-v1")
    .training(replay_buffer_config={"type": LessSampledReplayBuffer})
)

tune.Tuner(
    "DQN",
    param_space=config,
    run_config=tune.RunConfig(
        stop={"training_iteration": 1},
    ),
).fit()

For a full implementation, consider other methods such as get_state() and set_state(). For a more extensive example, see RLlib’s implementation of reservoir sampling, the ReservoirReplayBuffer.

Advanced usage#

In RLlib, all replay buffers implement the ReplayBuffer interface. They therefore support different StorageUnits whenever possible. A replay buffer’s storage_unit constructor argument defines how it stores experiences, and therefore the unit in which it samples them. When you later call the sample() method, num_items relates to that storage_unit.

The following example modifies the storage_unit and interacts with a custom buffer:

# This line will make our buffer store only complete episodes found in a batch
config.training(replay_buffer_config={"storage_unit": StorageUnit.EPISODES})

less_sampled_buffer = LessSampledReplayBuffer(**config.replay_buffer_config)

# Gather some random experiences
env = RandomEnv()
terminated = truncated = False
batch = SampleBatch({})
t = 0
while not terminated and not truncated:
    obs, reward, terminated, truncated, info = env.step([0, 0])
    # Note that in order for RLlib to find out about start and end of an episode,
    # "t" and "terminateds" have to properly mark an episode's trajectory
    one_step_batch = SampleBatch(
        {
            "obs": [obs],
            "t": [t],
            "reward": [reward],
            "terminateds": [terminated],
            "truncateds": [truncated],
        }
    )
    batch = concat_samples([batch, one_step_batch])
    t += 1

less_sampled_buffer.add(batch)
for i in range(10):
    assert len(less_sampled_buffer._storage) == 1
    less_sampled_buffer.sample(num_items=1, evict_sampled_more_then=9)

assert len(less_sampled_buffer._storage) == 0

As described earlier, RLlib’s MultiAgentReplayBuffers support modifying underlying replay buffers. The MultiAgentReplayBuffer stores experiences per policy in separate underlying replay buffers. Modify their behavior by specifying an underlying replay_buffer_config that works the same way as the parent’s config.

The following example creates a MultiAgentReplayBuffer with an alternative underlying ReplayBuffer. The MultiAgentReplayBuffer can stay the same. You only specify your own buffer along with a default call argument:

config = (
    DQNConfig()
    .api_stack(
        enable_env_runner_and_connector_v2=False, enable_rl_module_and_learner=False
    )
    .training(
        replay_buffer_config={
            "type": "MultiAgentReplayBuffer",
            "underlying_replay_buffer_config": {
                "type": LessSampledReplayBuffer,
                # We can specify the default call argument
                # for the sample method of the underlying buffer method here.
                "evict_sampled_more_then": 20,
            },
        }
    )
    .environment(env="CartPole-v1")
)

tune.Tuner(
    "DQN",
    param_space=config.to_dict(),
    run_config=tune.RunConfig(
        stop={"env_runners/episode_return_mean": 40, "training_iteration": 7},
    ),
).fit()