ReplayBufferEnsemble¶

class torchrl.data.replay_buffers.ReplayBufferEnsemble(*rbs, storages: StorageEnsemble | None = None, samplers: SamplerEnsemble | None = None, writers: WriterEnsemble | None = None, transform: 'Transform' | None = None, batch_size: int | None = None, collate_fn: Callable | None = None, collate_fns: List[Callable] | None = None, p: Tensor = None, sample_from_all: bool = False, num_buffer_sampled: int | None = None, generator: torch.Generator | None = None, shared: bool = False, **kwargs)[source]¶

重播緩衝區的集合。

此類別允許一次從多個重播緩衝區讀取和取樣。它會自動組成儲存空間 (StorageEnsemble)、寫入器 (WriterEnsemble) 和取樣器 (SamplerEnsemble) 的集合。

備註

禁止直接寫入此類別，但可以對其進行索引以檢索巢狀緩衝區並擴充它。

有兩種不同的方式建構 ReplayBufferEnsemble：可以傳遞重播緩衝區的清單，或直接傳遞元件（儲存空間、寫入器和取樣器），就像其他重播緩衝區子類別一樣。

參數:

rbs (重播緩衝區實例的序列, 可選) – 要集合的重播緩衝區。
storages (StorageEnsemble, 可選) – 儲存空間的集合，如果未傳遞重播緩衝區。
samplers (SamplerEnsemble, 可選) – 取樣器的集合，如果未傳遞重播緩衝區。
writers (WriterEnsemble, 可選) – 寫入器的集合，如果未傳遞重播緩衝區。
transform (Transform, optional) – 若有傳入，這會是重播緩衝區集合的轉換。每個重播緩衝區的個別轉換會從其父重播緩衝區中取得，或直接寫入 StorageEnsemble 物件中。
batch_size (int, optional) – 取樣時使用的批次大小。
collate_fn (callable, optional) – 在每個個別的 collate_fn 被呼叫且資料被放入列表 (連同緩衝區 ID) 後，用來整理 (collate) 資料的函式。
collate_fns (list of callables, optional) – 每個巢狀重播緩衝區的 collate_fn。如果沒有提供，則從 ReplayBuffer 實例中取得。
p (list of float or Tensor, optional) – 一個浮點數列表，指示每個重播緩衝區的相對權重。如果明確地建立緩衝區，也可以傳遞給 torchrl.data.replay_buffers.samplers.SamplerEnsemble`。
sample_from_all (bool, optional) – 如果 True，則會從每個資料集取樣。這與 p 引數不相容。預設值為 False。如果明確地建立緩衝區，也可以傳遞給 torchrl.data.replay_buffers.samplers.SamplerEnsemble`。
num_buffer_sampled (int, optional) – 要取樣的緩衝區數量。如果 sample_from_all=True，這不會產生任何影響，因為它預設為緩衝區的數量。如果 sample_from_all=False，則會根據機率 p 對緩衝區進行取樣。如果明確地建立緩衝區，也可以傳遞給 torchrl.data.replay_buffers.samplers.SamplerEnsemble`。
generator (torch.Generator, optional) –
用於取樣的 generator。為重播緩衝區使用專用的 generator 可以更精細地控制種子 (seeding)，例如，保持全域種子不同，但對於分散式任務，RB 種子相同。預設值為 None (全域預設 generator)。

Warning

到目前為止，generator 對於轉換沒有任何影響。
shared (bool, optional) – 緩衝區是否將使用多進程 (multiprocessing) 進行共享。預設值為 False。

Examples

>>> from torchrl.envs import Compose, ToTensorImage, Resize, RenameTransform
>>> from torchrl.data import TensorDictReplayBuffer, ReplayBufferEnsemble, LazyMemmapStorage
>>> from tensordict import TensorDict
>>> import torch
>>> rb0 = TensorDictReplayBuffer(
...     storage=LazyMemmapStorage(10),
...     transform=Compose(
...         ToTensorImage(in_keys=["pixels", ("next", "pixels")]),
...         Resize(32, in_keys=["pixels", ("next", "pixels")]),
...         RenameTransform([("some", "key")], ["renamed"]),
...     ),
... )
>>> rb1 = TensorDictReplayBuffer(
...     storage=LazyMemmapStorage(10),
...     transform=Compose(
...         ToTensorImage(in_keys=["pixels", ("next", "pixels")]),
...         Resize(32, in_keys=["pixels", ("next", "pixels")]),
...         RenameTransform(["another_key"], ["renamed"]),
...     ),
... )
>>> rb = ReplayBufferEnsemble(
...     rb0,
...     rb1,
...     p=[0.5, 0.5],
...     transform=Resize(33, in_keys=["pixels"], out_keys=["pixels33"]),
... )
>>> print(rb)
ReplayBufferEnsemble(
    storages=StorageEnsemble(
        storages=(<torchrl.data.replay_buffers.storages.LazyMemmapStorage object at 0x13a2ef430>, <torchrl.data.replay_buffers.storages.LazyMemmapStorage object at 0x13a2f9310>),
        transforms=[Compose(
                ToTensorImage(keys=['pixels', ('next', 'pixels')]),
                Resize(w=32, h=32, interpolation=InterpolationMode.BILINEAR, keys=['pixels', ('next', 'pixels')]),
                RenameTransform(keys=[('some', 'key')])), Compose(
                ToTensorImage(keys=['pixels', ('next', 'pixels')]),
                Resize(w=32, h=32, interpolation=InterpolationMode.BILINEAR, keys=['pixels', ('next', 'pixels')]),
                RenameTransform(keys=['another_key']))]),
    samplers=SamplerEnsemble(
        samplers=(<torchrl.data.replay_buffers.samplers.RandomSampler object at 0x13a2f9220>, <torchrl.data.replay_buffers.samplers.RandomSampler object at 0x13a2f9f70>)),
    writers=WriterEnsemble(
        writers=(<torchrl.data.replay_buffers.writers.TensorDictRoundRobinWriter object at 0x13a2d9b50>, <torchrl.data.replay_buffers.writers.TensorDictRoundRobinWriter object at 0x13a2f95b0>)),
batch_size=None,
transform=Compose(
        Resize(w=33, h=33, interpolation=InterpolationMode.BILINEAR, keys=['pixels'])),
collate_fn=<built-in method stack of type object at 0x128648260>)
>>> data0 = TensorDict(
...     {
...         "pixels": torch.randint(255, (10, 244, 244, 3)),
...         ("next", "pixels"): torch.randint(255, (10, 244, 244, 3)),
...         ("some", "key"): torch.randn(10),
...     },
...     batch_size=[10],
... )
>>> data1 = TensorDict(
...     {
...         "pixels": torch.randint(255, (10, 64, 64, 3)),
...         ("next", "pixels"): torch.randint(255, (10, 64, 64, 3)),
...         "another_key": torch.randn(10),
...     },
...     batch_size=[10],
... )
>>> rb[0].extend(data0)
>>> rb[1].extend(data1)
>>> for _ in range(2):
...     sample = rb.sample(10)
...     assert sample["next", "pixels"].shape == torch.Size([2, 5, 3, 32, 32])
...     assert sample["pixels"].shape == torch.Size([2, 5, 3, 32, 32])
...     assert sample["pixels33"].shape == torch.Size([2, 5, 3, 33, 33])
...     assert sample["renamed"].shape == torch.Size([2, 5])

add(data: Any) → int¶

向重播緩衝區新增單個元素。

參數:: data (Any) – 要新增到重播緩衝區的資料
Returns:: 資料在重播緩衝區中的索引。

append_transform(transform: Transform, *, invert: bool = False) → ReplayBuffer¶

在末端附加轉換。

當呼叫 sample 時，轉換會依序應用。

參數:: transform (Transform) – 要附加的轉換
Keyword Arguments:: invert (bool, optional) – 如果 True，轉換將會反轉 (forward 呼叫將在寫入期間被呼叫，而 inverse 呼叫將在讀取期間被呼叫)。預設值為 False。

Example

>>> 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()

dump(*args, **kwargs)¶: 是 dumps() 的別名。

dumps(path)¶

將重播緩衝區儲存到磁碟上的指定路徑。

參數:: path (Path or str) – 要儲存重播緩衝區的路徑。

Examples

>>> 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()¶: 清空重播緩衝區並將游標重設為 0。

extend(data: Sequence) → Tensor¶

使用一個或多個包含在可迭代物件中的元素來擴充回放緩衝區。

如果存在，將會調用反向轉換 (inverse transforms)。

參數:: data (iterable) – 要新增至回放緩衝區的資料集合。
Returns:: 新增至回放緩衝區的資料索引。

Warning

當處理數值列表時，extend() 可能會有不明確的簽名。該列表應被解讀為 PyTree（在這種情況下，列表中的所有元素都將放入儲存 PyTree 中的一個切片中），或者是一個個新增數值的列表。為了解決這個問題，TorchRL 明確地區分了列表和元組：元組將被視為 PyTree，列表（在根層級）將被解釋為一次一個地添加到緩衝區的數值堆疊。對於 ListStorage 實例，只能提供未綁定的元素（沒有 PyTrees）。

insert_transform(index: int, transform: Transform, *, invert: bool = False) → ReplayBuffer¶

插入轉換 (transform)。

當調用 sample 時，轉換會依序執行。

參數:

index (int) – 插入轉換的位置。
transform (Transform) – 要附加的轉換

Keyword Arguments:

invert (bool, optional) – 如果 True，轉換將會反轉 (forward 呼叫將在寫入期間被呼叫，而 inverse 呼叫將在讀取期間被呼叫)。預設值為 False。

load(*args, **kwargs)¶: 是 loads() 的別名。

loads(path)¶

在給定的路徑加載回放緩衝區狀態。

緩衝區應具有匹配的組件，並使用 dumps() 儲存。

參數:: path (Path or str) – 回放緩衝區儲存的路徑。

有關更多資訊，請參閱 dumps()。

register_load_hook(hook: Callable[[Any], Any])¶: 為儲存註冊一個加載 hook。

備註

Hooks 目前在儲存回放緩衝區時不會被序列化：每次建立緩衝區時都必須手動重新初始化它們。

register_save_hook(hook: Callable[[Any], Any])¶: 為儲存註冊一個儲存 hook。

備註

Hooks 目前在儲存回放緩衝區時不會被序列化：每次建立緩衝區時都必須手動重新初始化它們。

sample(batch_size: Optional[int] = None, return_info: bool = False) → Any¶

從回放緩衝區取樣一批資料。

使用 Sampler 取樣索引，並從 Storage 檢索它們。

參數:

batch_size (int, optional) – 要收集的資料大小。如果未提供，此方法將取樣一個由 sampler 指示的 batch-size。
return_info (bool) – 是否傳回資訊。如果為 True，結果是一個元組 (data, info)。如果為 False，結果是資料。

Returns:

在回放緩衝區中選取的一批資料。如果 return_info 標記設定為 True，則包含此批次和資訊的元組。

property sampler¶

回放緩衝區的 sampler。

sampler 必須是 Sampler 的實例。

save(*args, **kwargs)¶: 是 dumps() 的別名。

set_sampler(sampler: Sampler)¶: 在回放緩衝區中設定一個新的 sampler 並傳回先前的 sampler。

set_storage(storage: Storage, collate_fn: Optional[Callable] = None)¶

在 replay buffer 中設定新的 storage，並回傳先前的 storage。

參數:

storage (Storage) – buffer 的新 storage。
collate_fn (callable, optional) – 若提供，collate_fn 會被設定為這個值。否則，它會被重置為預設值。

set_writer(writer: Writer)¶: 在 replay buffer 中設定新的 writer，並回傳先前的 writer。

property storage¶

replay buffer 的 storage。

storage 必須是 Storage 的一個實例。

property write_count¶: 透過 add 和 extend 寫入 buffer 的項目總數。

property writer¶

replay buffer 的 writer。

writer 必須是 Writer 的一個實例。

ReplayBufferEnsemble¶

文件

教學

資源