TensorDictMaxValueWriter

class torchrl.data.replay_buffers.TensorDictMaxValueWriter(rank_key=None, reduction: str = 'sum', **kwargs)

一個 Writer 類別，用於可組合的 replay buffers，基於某些排名鍵保留頂部元素。

參數:

• rank_key (str or str 的 tuple) – 用於對元素進行排名的鍵。預設為 ("next", "reward")。

• reduction (str) – 如果排名鍵有多個元素，則要使用的縮減方法。可以是 "max"、"min"、"mean"、"median" 或 "sum"。

範例

>>> import torch
>>> from tensordict import TensorDict
>>> from torchrl.data import LazyTensorStorage, TensorDictReplayBuffer, TensorDictMaxValueWriter
>>> from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
>>> rb = TensorDictReplayBuffer(
...     storage=LazyTensorStorage(1),
...     sampler=SamplerWithoutReplacement(),
...     batch_size=1,
...     writer=TensorDictMaxValueWriter(rank_key="key"),
... )
>>> td = TensorDict({
...     "key": torch.tensor(range(10)),
...     "obs": torch.tensor(range(10))
... }, batch_size=10)
>>> rb.extend(td)
>>> print(rb.sample().get("obs").item())
9
>>> td = TensorDict({
...     "key": torch.tensor(range(10, 20)),
...     "obs": torch.tensor(range(10, 20))
... }, batch_size=10)
>>> rb.extend(td)
>>> print(rb.sample().get("obs").item())
19
>>> td = TensorDict({
...     "key": torch.tensor(range(10)),
...     "obs": torch.tensor(range(10))
... }, batch_size=10)
>>> rb.extend(td)
>>> print(rb.sample().get("obs").item())
19

注意

此類別與具有多個維度的儲存體不相容。 這並不意味著禁止儲存 trajectories，而是儲存的 trajectories 必須基於每個 trajectory 儲存。 以下是一些類別的有效和無效用法的範例。 首先，一個平面 buffer，我們在其中儲存個別的轉換

>>> from torchrl.data import TensorStorage
>>> # Simplest use case: data comes in 1d and is stored as such
>>> data = TensorDict({
...     "obs": torch.zeros(10, 3),
...     "reward": torch.zeros(10, 1),
... }, batch_size=[10])
>>> rb = TensorDictReplayBuffer(
...     storage=LazyTensorStorage(max_size=100),
...     writer=TensorDictMaxValueWriter(rank_key="reward")
... )
>>> # We initialize the buffer: a total of 100 *transitions* can be stored
>>> rb.extend(data)
>>> # Samples 5 *transitions* at random
>>> sample = rb.sample(5)
>>> assert sample.shape == (5,)

其次，一個我們儲存 trajectories 的 buffer。 最大訊號會在每個批次中聚合（例如，每個 rollout 的獎勵都會被加總）

>>> # One can also store batches of data, each batch being a sub-trajectory
>>> env = ParallelEnv(2, lambda: GymEnv("Pendulum-v1"))
>>> # Get a batch of [2, 10] -- format is [Batch, Time]
>>> rollout = env.rollout(max_steps=10)
>>> rb = TensorDictReplayBuffer(
...     storage=LazyTensorStorage(max_size=100),
...     writer=TensorDictMaxValueWriter(rank_key="reward")
... )
>>> # We initialize the buffer: a total of 100 *trajectories* (!) can be stored
>>> rb.extend(rollout)
>>> # Sample 5 trajectories at random
>>> sample = rb.sample(5)
>>> assert sample.shape == (5, 10)

如果資料以批次形式傳入，但需要平面 buffer，我們可以簡單地在擴充 buffer 之前將資料扁平化

>>> rb = TensorDictReplayBuffer(
...     storage=LazyTensorStorage(max_size=100),
...     writer=TensorDictMaxValueWriter(rank_key="reward")
... )
>>> # We initialize the buffer: a total of 100 *transitions* can be stored
>>> rb.extend(rollout.reshape(-1))
>>> # Sample 5 trajectories at random
>>> sample = rb.sample(5)
>>> assert sample.shape == (5,)

不可能建立一個沿時間維度擴充的 buffer，這通常是使用具有 trajectories 批次的 buffers 的建議方式。 由於 trajectories 是重疊的，因此很難（如果不是不可能）聚合獎勵值並比較它們。 此建構函式無效（請注意 ndim 引數）

>>> rb = TensorDictReplayBuffer(
...     storage=LazyTensorStorage(max_size=100, ndim=2),  # Breaks!
...     writer=TensorDictMaxValueWriter(rank_key="reward")
... )

add(data: Any) → int | torch.Tensor

在適當的索引處插入單個資料元素，並傳回該索引。

傳遞到此模組的資料中的 rank_key 應該結構化為 []。 如果它有更多維度，它將使用 reduction 方法縮減為單個值。

extend(data: TensorDictBase) → None

在適當的索引位置插入一系列資料點。

傳遞到此模組的資料中的 rank_key 應構建為 [B]。如果它有更多維度，它將使用 reduction 方法縮減為單個值。

get_insert_index(data: Any) → int

返回應該插入資料的索引，如果資料不應該被插入，則返回 None。