PrioritizedReplayBuffer¶
- class torchrl.data.PrioritizedReplayBuffer(*, alpha: float, beta: float, eps: float = 1e-08, dtype: torch.dtype = torch.float32, storage: Storage | None = None, sampler: Sampler | None = None, collate_fn: Callable | None = None, pin_memory: bool = False, prefetch: int | None = None, transform: Transform | None = None, batch_size: int | None = None, dim_extend: int | None = None)[原始碼]¶
優先經驗回放緩衝區。
所有引數都是僅關鍵字引數。
源自“Schaul, T.; Quan, J.; Antonoglou, I.; and Silver, D. 2015. Prioritized experience replay.” (https://arxiv.org/abs/1511.05952)
- 引數:
alpha (
float) – 指數 α 決定了使用的優先順序程度,當 α = 0 時對應於均勻情況。beta (
float) – 重要性取樣負指數。eps (
float) – 新增到優先順序的 delta,以確保緩衝區不包含零優先順序。storage (Storage, 可選) – 要使用的儲存。如果未提供,將建立一個預設的
ListStorage,其max_size為1_000。sampler (Sampler, 可選) – 要使用的取樣器。如果未提供,將建立一個預設的
PrioritizedSampler,並附帶alpha、beta和eps。collate_fn (callable, optional) – 將樣本列表合併以形成 Tensor/輸出的 mini-batch。在使用 map-style 資料集的批次載入時使用。預設值將根據儲存型別確定。
pin_memory (bool) – 是否應對 rb 樣本呼叫 pin_memory()。
prefetch (int, optional) – 要透過多執行緒預取的下一個批次數。預設為 None(無預取)。
transform (Transform, 可選) – 呼叫 sample() 時執行的轉換。要連結轉換,請使用
Compose類。轉換應與tensordict.TensorDict內容一起使用。如果與其他結構一起使用,則應使用帶有"data"字首的鍵來編碼轉換,該鍵將用於從非 tensordict 內容構造 tensordict。batch_size (int, optional) –
呼叫 sample() 時要使用的批次大小。
注意
batch-size 可以在構造時透過
batch_size引數指定,也可以在取樣時指定。前者在整個實驗中 batch-size 保持一致時應優先使用。如果 batch-size 可能會發生變化,則可以將其傳遞給sample()方法。此選項與預取(因為這需要提前知道 batch-size)以及具有drop_last引數的取樣器不相容。dim_extend (int, optional) –
指示呼叫
extend()時要考慮的用於擴充套件的維度。預設為storage.ndim-1。在使用dim_extend > 0時,我們建議在儲存例項化時使用ndim引數(如果該引數可用),以告知儲存資料是多維的,並在取樣過程中保持儲存容量和 batch-size 的一致性。
注意
通用優先經驗回放緩衝區(即非 tensordict 後備的)需要將
sample()的return_info引數設定為True,才能訪問索引,從而更新優先順序。使用tensordict.TensorDict和相關的TensorDictPrioritizedReplayBuffer可以簡化此過程。示例
>>> import torch >>> >>> from torchrl.data import ListStorage, PrioritizedReplayBuffer >>> >>> torch.manual_seed(0) >>> >>> rb = PrioritizedReplayBuffer(alpha=0.7, beta=0.9, storage=ListStorage(10)) >>> data = range(10) >>> rb.extend(data) >>> sample = rb.sample(3) >>> print(sample) tensor([1, 0, 1]) >>> # get the info to find what the indices are >>> sample, info = rb.sample(5, return_info=True) >>> print(sample, info) tensor([2, 7, 4, 3, 5]) {'_weight': array([1., 1., 1., 1., 1.], dtype=float32), 'index': array([2, 7, 4, 3, 5])} >>> # update priority >>> priority = torch.ones(5) * 5 >>> rb.update_priority(info["index"], priority) >>> # and now a new sample, the weights should be updated >>> sample, info = rb.sample(5, return_info=True) >>> print(sample, info) tensor([2, 5, 2, 2, 5]) {'_weight': array([0.36278465, 0.36278465, 0.36278465, 0.36278465, 0.36278465], dtype=float32), 'index': array([2, 5, 2, 2, 5])}
- add(data: Any) int¶
將單個元素新增到重放緩衝區。
- 引數:
data (Any) – 要新增到重放緩衝區的資料
- 返回:
資料在重放緩衝區中的索引。
- append_transform(transform: Transform, *, invert: bool = False) ReplayBuffer¶
將變換附加到末尾。
呼叫 sample 時按順序應用變換。
- 引數:
transform (Transform) – 要附加的變換
- 關鍵字引數:
invert (bool, optional) – 如果為
True,則轉換將被反轉(寫入時呼叫正向呼叫,讀取時呼叫反向呼叫)。預設為False。
示例
>>> rb = ReplayBuffer(storage=LazyMemmapStorage(10), batch_size=4) >>> data = TensorDict({"a": torch.zeros(10)}, [10]) >>> def t(data): ... data += 1 ... return data >>> rb.append_transform(t, invert=True) >>> rb.extend(data) >>> assert (data == 1).all()
- classmethod as_remote(remote_config=None)¶
建立一個遠端 ray 類的例項。
- 引數:
cls (Python Class) – 要遠端例項化的類。
remote_config (dict) – 為該類保留的 CPU 核心數量。預設為 torchrl.collectors.distributed.ray.DEFAULT_REMOTE_CLASS_CONFIG。
- 返回:
一個建立 ray 遠端類例項的函式。
- property batch_size¶
重放緩衝區的批次大小。
batch_size 可以透過在
sample()方法中設定 batch_size 引數來覆蓋。它定義了
sample()返回的樣本數量以及ReplayBuffer迭代器產生的樣本數量。
- dumps(path)¶
將重放緩衝區儲存到指定路徑的磁碟上。
- 引數:
path (Path 或 str) – 儲存重放緩衝區的路徑。
示例
>>> import tempfile >>> import tqdm >>> from torchrl.data import LazyMemmapStorage, TensorDictReplayBuffer >>> from torchrl.data.replay_buffers.samplers import PrioritizedSampler, RandomSampler >>> import torch >>> from tensordict import TensorDict >>> # Build and populate the replay buffer >>> S = 1_000_000 >>> sampler = PrioritizedSampler(S, 1.1, 1.0) >>> # sampler = RandomSampler() >>> storage = LazyMemmapStorage(S) >>> rb = TensorDictReplayBuffer(storage=storage, sampler=sampler) >>> >>> for _ in tqdm.tqdm(range(100)): ... td = TensorDict({"obs": torch.randn(100, 3, 4), "next": {"obs": torch.randn(100, 3, 4)}, "td_error": torch.rand(100)}, [100]) ... rb.extend(td) ... sample = rb.sample(32) ... rb.update_tensordict_priority(sample) >>> # save and load the buffer >>> with tempfile.TemporaryDirectory() as tmpdir: ... rb.dumps(tmpdir) ... ... sampler = PrioritizedSampler(S, 1.1, 1.0) ... # sampler = RandomSampler() ... storage = LazyMemmapStorage(S) ... rb_load = TensorDictReplayBuffer(storage=storage, sampler=sampler) ... rb_load.loads(tmpdir) ... assert len(rb) == len(rb_load)
- empty(empty_write_count: bool = True)¶
清空重放緩衝區並將遊標重置為 0。
- 引數:
empty_write_count (bool, optional) – 是否清空 write_count 屬性。預設為 True。
- extend(data: Sequence, *, update_priority: bool | None = None) torch.Tensor¶
使用包含在可迭代物件中的一個或多個元素擴充套件經驗回放緩衝區。
如果存在,將呼叫逆向轉換。
- 引數:
data (iterable) – 要新增到經驗回放緩衝區的元素集合。
- 關鍵字引數:
update_priority (bool, optional) – 是否更新資料的優先順序。預設為 True。此類中無效。有關更多詳細資訊,請參閱
extend()。- 返回:
新增到經驗回放緩衝區的資料的索引。
警告
extend()在處理值列表時可能具有歧義的簽名,應將其解釋為 PyTree(在這種情況下,列表中的所有元素都將放入儲存中的 PyTree 片段)或要一次性新增的值列表。為了解決這個問題,TorchRL 區分了列表和元組:元組將被視為 PyTree,列表(在根級別)將被解釋為要一次性新增到緩衝區的值的堆疊。對於ListStorage例項,只能提供未繫結的元素(無 PyTrees)。
- insert_transform(index: int, transform: Transform, *, invert: bool = False) ReplayBuffer¶
插入變換。
呼叫 sample 時按順序執行變換。
- 引數:
index (int) – 插入變換的位置。
transform (Transform) – 要附加的變換
- 關鍵字引數:
invert (bool, optional) – 如果為
True,則轉換將被反轉(寫入時呼叫正向呼叫,讀取時呼叫反向呼叫)。預設為False。
- loads(path)¶
在給定路徑載入重放緩衝區狀態。
緩衝區應具有匹配的元件,並使用
dumps()儲存。- 引數:
path (Path 或 str) – 重放緩衝區儲存的路徑。
有關更多資訊,請參見
dumps()。
- next()¶
返回重放緩衝區的下一個項。
此方法用於在 __iter__ 不可用的情況下迭代重放緩衝區,例如
RayReplayBuffer。
- register_load_hook(hook: Callable[[Any], Any])¶
為儲存註冊載入鉤子。
注意
鉤子目前不會在儲存重放緩衝區時序列化:每次建立緩衝區時都必須手動重新初始化它們。
- register_save_hook(hook: Callable[[Any], Any])¶
為儲存註冊儲存鉤子。
注意
鉤子目前不會在儲存重放緩衝區時序列化:每次建立緩衝區時都必須手動重新初始化它們。
- sample(batch_size: int | None = None, return_info: bool = False) Any¶
從重放緩衝區中取樣資料批次。
使用 Sampler 取樣索引,並從 Storage 中檢索它們。
- 引數:
batch_size (int, optional) – 要收集的資料的大小。如果未提供,此方法將取樣由取樣器指示的批次大小。
return_info (bool) – 是否返回資訊。如果為 True,則結果為元組 (data, info)。如果為 False,則結果為資料。
- 返回:
從經驗回放緩衝區中選擇的資料批次。如果 return_info 標誌設定為 True,則返回包含此批次和資訊的元組。
- set_storage(storage: Storage, collate_fn: Callable | None = None)¶
在重放緩衝區中設定新的儲存並返回之前的儲存。
- 引數:
storage (Storage) – 緩衝區的新的儲存。
collate_fn (callable, optional) – 如果提供,collate_fn 將設定為此值。否則,它將被重置為預設值。
- property write_count: int¶
透過 add 和 extend 寫入緩衝區的總項數。