ReplayBuffer¶
- class torchrl.data.ReplayBuffer(*, storage: Storage | Callable[[], Storage] | None = None, sampler: Sampler | Callable[[], Sampler] | None = None, writer: Writer | Callable[[], Writer] | None = None, collate_fn: Callable | None = None, pin_memory: bool = False, prefetch: int | None = None, transform: Transform | Callable | None = None, transform_factory: Callable[[], Transform | Callable] | None = None, batch_size: int | None = None, dim_extend: int | None = None, checkpointer: StorageCheckpointerBase | Callable[[], StorageCheckpointerBase] | None = None, generator: torch.Generator | None = None, shared: bool = False, compilable: bool | None = None)[source]¶
一個通用的、可組合的重放緩衝區類。
- 關鍵字引數:
storage (Storage, Callable[[], Storage], optional) – 要使用的儲存。如果傳入一個可呼叫物件,它將用作儲存的建構函式。如果未提供,則將建立一個預設的
ListStorage,其max_size為1_000。sampler (Sampler, Callable[[], Sampler], optional) – 要使用的取樣器。如果傳入一個可呼叫物件,它將用作取樣器的建構函式。如果未提供,將使用預設的
RandomSampler。writer (Writer, Callable[[], Writer], optional) – 要使用的寫入器。如果傳入一個可呼叫物件,它將用作寫入器的建構函式。如果未提供,將使用預設的
RoundRobinWriter。collate_fn (callable, optional) – 將樣本列表合併以形成 Tensor/輸出的 mini-batch。在使用 map-style 資料集的批次載入時使用。預設值將根據儲存型別確定。
pin_memory (bool) – 是否應對 rb 樣本呼叫 pin_memory()。
prefetch (int, optional) – 要透過多執行緒預取的下一個批次數。預設為 None(無預取)。
transform (Transform 或 Callable[[Any], Any], optional) – 在呼叫
sample()時要執行的變換。要連結變換,請使用Compose類。變換應與tensordict.TensorDict內容一起使用。如果重放緩衝區與 PyTree 結構一起使用(參見下面的示例),也可以傳入一個通用的可呼叫物件。與儲存、寫入器和取樣器不同,變換建構函式必須作為單獨的關鍵字引數transform_factory傳入,因為無法區分建構函式和變換。transform_factory (Callable[[], Callable], optional) – 變換的工廠。與
transform互斥。batch_size (int, optional) –
呼叫 sample() 時要使用的批次大小。
注意
批處理大小可以在構造時透過
batch_size引數指定,也可以在取樣時指定。當批處理大小在整個實驗中一致時,應優先選擇前者。如果批處理大小可能會更改,則可以將其傳遞給sample()方法。此選項與預取(因為這需要提前知道批處理大小)以及具有drop_last引數的取樣器不相容。dim_extend (int, optional) –
表示呼叫
extend()時要考慮用於擴充套件的維度。預設為storage.ndim-1。當使用dim_extend > 0時,建議在儲存例項化中使用ndim引數(如果該引數可用),以讓儲存知道資料是多維的,並在取樣過程中保持儲存容量和批處理大小的一致性。generator (torch.Generator, optional) –
要用於取樣的生成器。使用專用的生成器用於重放緩衝區可以對種子進行精細控制,例如在分散式作業中使全域性種子不同但 RB 種子相同。預設為
None(全域性預設生成器)。警告
目前,生成器對變換沒有影響。
shared (bool, optional) – 緩衝區是否將使用多程序共享。預設為
False。compilable (bool, optional) – 寫入器是否可編譯。如果為
True,則寫入器不能在多個程序之間共享。預設為False。
示例
>>> import torch >>> >>> from torchrl.data import ReplayBuffer, ListStorage >>> >>> torch.manual_seed(0) >>> rb = ReplayBuffer( ... storage=ListStorage(max_size=1000), ... batch_size=5, ... ) >>> # populate the replay buffer and get the item indices >>> data = range(10) >>> indices = rb.extend(data) >>> # sample will return as many elements as specified in the constructor >>> sample = rb.sample() >>> print(sample) tensor([4, 9, 3, 0, 3]) >>> # Passing the batch-size to the sample method overrides the one in the constructor >>> sample = rb.sample(batch_size=3) >>> print(sample) tensor([9, 7, 3]) >>> # one cans sample using the ``sample`` method or iterate over the buffer >>> for i, batch in enumerate(rb): ... print(i, batch) ... if i == 3: ... break 0 tensor([7, 3, 1, 6, 6]) 1 tensor([9, 8, 6, 6, 8]) 2 tensor([4, 3, 6, 9, 1]) 3 tensor([4, 4, 1, 9, 9])
重放緩衝區接受*任何*型別的資料。並非所有儲存型別都能正常工作,因為有些只接受數值資料,但預設的
ListStorage可以。示例
>>> torch.manual_seed(0) >>> buffer = ReplayBuffer(storage=ListStorage(100), collate_fn=lambda x: x) >>> indices = buffer.extend(["a", 1, None]) >>> buffer.sample(3) [None, 'a', None]
TensorStorage、LazyMemmapStorage和LazyTensorStorage也支援任何 PyTree 結構(PyTree 是由字典、列表或元組組成的任意深度的巢狀結構,其中葉子是張量),前提是它只包含張量資料。示例
>>> from torch.utils._pytree import tree_map >>> def transform(x): ... # Zeros all the data in the pytree ... return tree_map(lambda y: y * 0, x) >>> rb = ReplayBuffer(storage=LazyMemmapStorage(100), transform=transform) >>> data = { ... "a": torch.randn(3), ... "b": {"c": (torch.zeros(2), [torch.ones(1)])}, ... 30: -torch.ones(()), ... } >>> rb.add(data) >>> # The sample has a similar structure to the data (with a leading dimension of 10 for each tensor) >>> s = rb.sample(10) >>> # let's check that our transform did its job: >>> def assert0(x): >>> assert (x == 0).all() >>> tree_map(assert0, s)
- append_transform(transform: Transform, *, invert: bool = False) ReplayBuffer[source]¶
將變換附加到末尾。
呼叫 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¶
重放緩衝區的批次大小。
批處理大小可以透過在
sample()方法中設定 batch_size 引數來覆蓋。它定義了
sample()返回的樣本數量以及ReplayBuffer迭代器產生的樣本數量。
- dumps(path)[source]¶
將重放緩衝區儲存到指定路徑的磁碟上。
- 引數:
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)[source]¶
清空重放緩衝區並將遊標重置為 0。
- 引數:
empty_write_count (bool, optional) – 是否清空 write_count 屬性。預設為 True。
- extend(data: Sequence, *, update_priority: bool | None = None) torch.Tensor[source]¶
使用包含在可迭代物件中的一個或多個元素擴充套件經驗回放緩衝區。
如果存在,將呼叫逆向轉換。
- 引數:
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[source]¶
插入變換。
呼叫 sample 時按順序執行變換。
- 引數:
index (int) – 插入變換的位置。
transform (Transform) – 要附加的變換
- 關鍵字引數:
invert (bool, optional) – 如果為
True,則轉換將被反轉(寫入時呼叫正向呼叫,讀取時呼叫反向呼叫)。預設為False。
- loads(path)[source]¶
在給定路徑載入重放緩衝區狀態。
緩衝區應具有匹配的元件,並使用
dumps()儲存。- 引數:
path (Path 或 str) – 重放緩衝區儲存的路徑。
有關更多資訊,請參閱
dumps()。
- register_load_hook(hook: Callable[[Any], Any])[source]¶
為儲存註冊載入鉤子。
注意
鉤子目前不會在儲存重放緩衝區時序列化:每次建立緩衝區時都必須手動重新初始化它們。
- register_save_hook(hook: Callable[[Any], Any])[source]¶
為儲存註冊儲存鉤子。
注意
鉤子目前不會在儲存重放緩衝區時序列化:每次建立緩衝區時都必須手動重新初始化它們。
- sample(batch_size: int | None = None, return_info: bool = False) Any[source]¶
從重放緩衝區中取樣資料批次。
使用 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)[source]¶
在重放緩衝區中設定新的儲存並返回之前的儲存。
- 引數:
storage (Storage) – 緩衝區的新的儲存。
collate_fn (callable, optional) – 如果提供,collate_fn 將設定為此值。否則,它將被重置為預設值。
- property write_count: int¶
透過 add 和 extend 寫入緩衝區的總項數。