动态加速强化学习的闪存缓冲库——Flashbax

2024-06-07 18:02:45作者：廉彬冶Miranda

项目简介 🌟

Flashbax 是一款专为强化学习（RL）设计的高效体验回放缓冲库，特别针对JAX环境进行了优化。它提供了一种简单且灵活的方式来管理和利用这些缓冲区，即使在完全编译的函数和训练循环中也不例外。

通过Flashbax，您可以轻松实现不同类型的缓冲区，包括平面缓冲区、轨迹缓冲区以及优先级版本的这两种缓冲区。无论是学术研究、工业应用还是个人项目，都能享受到这个库带来的便利性。

技术亮点 🔥

效能优化的缓冲区

所有Flashbax缓冲区都是基于轨迹缓冲区的专业变体，旨在优化内存使用和功能，适用于各种缓冲类型。

平面缓冲区

类比于DQN算法中的过渡缓冲区，平面缓冲区采用2周期（即 $s_{t}$ ， $s_{t+1}$ ），用于完整考虑过渡对。

项缓冲区

项缓冲区存储独立的项，例如（观测，动作，奖励，折扣，下一个观测）元组或整个episode，非常灵活。

轨迹缓冲区

轨迹缓冲区支持多步轨迹的采样，非常适合使用递归网络如R2D2的算法。

优先级缓冲区

平面上和轨迹上的优先级缓冲区允许基于用户定义的优先级进行采样，遵循PER论文的原理。

路径/平面队列

提供了队列数据结构，适合按先进先出顺序采样数据，在特定的On-Policy算法中有其用途。

安装与起步 📦

要将Flashbax集成到您的项目中，请按照以下步骤操作：

安装：首先使用pip安装Flashbax：
```
pip install flashbax
```

选择缓冲区：从各种缓冲选项中挑选，例如：

from flashbax import make_trajectory_buffer, make_prioritised_trajectory_buffer, make_flat_buffer, make_prioritised_flat_buffer, make_item_buffer, make_trajectory_queue

buffer = make_trajectory_buffer(...) 
# 或者
buffer = make_prioritised_trajectory_buffer(...) 
# 或者
buffer = make_flat_buffer(...) 
# 或者
buffer = make_prioritised_flat_buffer(...) 
# 或者
buffer = make_item_buffer(...) 
# 或者
buffer = make_trajectory_queue(...)

# 初始化
state = buffer.init(example_timestep)
# 添加数据
state = buffer.add(state, example_data)
# 采样数据
batch = buffer.sample(state, rng_key)

快速入门 🏃‍♂️

以下是一个使用平面缓冲区的简易示例，展示了如何使用纯函数来定义平面缓冲区的基本用法。请注意，这些纯函数与jax.pmap和jax.jit兼容，但此处为了简洁起见未使用它们：

import jax
import jax.numpy as jnp
from flashbax import make_flat_buffer

# 使用简单的配置实例化平面缓冲区。
# 返回的`buffer`只是一个容器，用于平面缓冲区所需的纯函数。
buffer = make_flat_buffer(max_length=32, min_length=2, sample_batch_size=1)

# 初始化缓冲区状态。
fake_timestep = {"obs": jnp.array([0, 0]), "reward": jnp.array(0.0)}
state = buffer.init(fake_timestep)

# 向缓冲区添加数据。
state = buffer.add(state, {"obs": jnp.array([1, 2]), "reward": jnp.array(3.0)})
print(buffer.can_sample(state))  # 输出：False，因为min_length尚未达到。

state = buffer.add(state, {"obs": jnp.array([4, 5]), "reward": jnp.array(6.0)})
print(buffer.can_sample(state))  # 输出：False，因为我们仍需2个transition（即3个时间步长）。

state = buffer.add(state, {"obs": jnp.array([7, 8]), "reward": jnp.array(9.0)})
print(buffer.can_sample(state))  # 输出：True！我们有了2个transition（3个时间步长）。

# 从缓冲区获取一个transition。
rng_key = jax.random.PRNGKey(0)  # 随机数源
batch = buffer.sample(state, rng_key)  # 采样

# 我们得到了一个transition！打印：obs = [[4 5]]，obs' = [[7 8]]
print(
    f"obs = {batch.experience.first['obs']}, obs' = {batch.experience.second['obs']}"
)